Temporal Difference Learning for Diffusion Models

2026-06-12

AI Summary

This work addresses the degradation in generation quality of diffusion models under few-step sampling, which stems from insufficient cross-temporal consistency in the denoising trajectory. The authors propose the first formulation of the diffusion process as a Markov reward process, reframing denoising as a policy evaluation problem in reinforcement learning. To enforce consistency across multi-step denoising paths, they introduce a temporal difference (TD) objective that penalizes trajectory inconsistencies and incorporate a sample reweighting strategy to stabilize training. The resulting method is universally applicable to both discrete- and continuous-time diffusion models and achieves substantial improvements in generation quality—measured by FID—under limited sampling steps, demonstrating particularly pronounced advantages in low-compute regimes.

Abstract

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 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.

Problem

• diffusion models
• temporal consistency
• denoising trajectory
• few-step sampling
• sample quality

Innovation

• Temporal Difference Learning
• Diffusion Models
• Cross-time Consistency
• Reinforcement Learning
• Few-step Sampling