Temporal difference learning for diffusion models

Conference item

Diffusion models are typically trained with objectives that focus on local denoising targets at individual time steps (or adjacent pairs), which do not enforce consistency between predictions along the denoising trajectory. This lack of crosstime consistency can degrade performance, especially for few-step samplers. We introduce a temporal difference (TD) objective that penalizes inconsistency of the model’s multi-step progress along the denoising path. By reformulating the diffusion process as a Markov reward process and casting denoising as a policy evaluation problem in reinforcement learning, we derive a unified TD approach that applies to both discrete- and continuous-time diffusion formulations. We further propose a principled sample-based reweighting method that stabilizes training. Empirically, we show that using our TD training can significantly improve sample quality measured by FID, with stronger advantages when the number of sampling steps is small, highlighting its practical utility under low-computation-budget scenarios. We provide ablation studies to justify our design choices, including pairwise loss reweighting, regularization weight, and one-step stride. Overall, our TD approach can be a general drop-in that enforces cross-time consistency and improves generation quality across different diffusion generative models.

Authors: Ying, Q; Pan, Y; Prisacariu, VA; Wen, J

• Publication status: Accepted
• Peer review status: Peer reviewed
• Publisher: Proceedings of Machine Learning Research
• Acceptance date: 2026-04-30
• Event title: 43rd International Conference on Machine Learning (ICML 2026)
• Event location: Seoul, South Korea
• Event website: https://icml.cc/Conferences/2026
• Event start date: 2026-07-06
• Event end date: 2026-07-11
• Language: English