Temporal Difference Learning for Diffusion Models

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 cross-time consistency can degrade performance, especially for few-step samplers. We introduce a temporal difference (TD) objective that penalizes inconsistency of the model’s \emph{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.

Modern AI systems can create realistic images by starting from random noise and gradually cleaning it up. However, their standard training objectives mainly teach the model to make accurate predictions at individual noise levels, rather than directly encouraging predictions at different stages to agree with each other along the whole generation path. As a result, small inconsistencies can accumulate, especially when we want the system to generate images quickly using only a few steps. In this paper, we introduce a new training method that encourages the model’s predictions at different stages of image creation to stay consistent with one another. Our idea is inspired by how learning systems can improve by comparing progress across time, rather than judging each step in isolation. The method can be added to several existing diffusion model training pipelines. In experiments on standard image-generation benchmarks, our approach often improves image quality, especially when using fewer generation steps. This suggests that encouraging consistency across the generation process can help diffusion models produce better images under fast sampling.