Temporal Difference Learning for Diffusion Models
Qizhen Ying, Yangchen Pan, Victor Prisacariu, Junfeng Wen
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 discreteand 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.Code is available at https://github.com/StephenYing/ Temporal_Difference_Learning_ for_Diffusion_Models .
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 3f217997-79dc-40ac-9b9e-ebcaa1d8aab3Builds on17
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 35,902 citations
- Denoising Diffusion Implicit ModelsJiaming Song, Chenlin Meng, Stefano ErmonICLR 2021 · 11,743 citations
- Improved Denoising Diffusion Probabilistic ModelsAlexander Quinn Nichol, Prafulla DhariwalICML 2021 · 5,234 citations
- Elucidating the Design Space of Diffusion-Based Generative ModelsTero Karras, Miika Aittala, Timo Aila, Samuli LaineNeurIPS 2022 · 3,959 citations
- DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 StepsCheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen et al.NeurIPS 2022 · 2,653 citations
Related papers
- Consistent Diffusion Models: Mitigating Sampling Drift by Learning to be ConsistentGiannis Daras, Yuval Dagan, Alex Dimakis, Constantinos DaskalakisNeurIPS 2023 · 79 citations
- Target-Driven Distillation: Consistency Distillation with Target Timestep Selection and Decoupled GuidanceCunzheng Wang, Ziyuan Guo, Yuxuan Duan, Huaxia Li et al.AAAI 2025 · 9 citations
- Simplifying, Stabilizing and Scaling Continuous-time Consistency ModelsCheng Lu, Yang SongICLR 2025
- Your Diffusion Model is Secretly a Noise Classifier and Benefits from Contrastive TrainingYunshu Wu, Yingtao Luo, Xianghao Kong, Vagelis Papalexakis et al.NeurIPS 2024 · 6 citations
- ReAlign: Text-to-Motion Generation via Step-Aware Reward-Guided AlignmentWanjiang Weng, Xiaofeng Tan, Junbo Wang, Guo-Sen Xie et al.AAAI 2026 · 6 citations
