Temporal Difference Learning for Diffusion Models
Summary
This paper introduces a temporal difference (TD) learning objective for diffusion models that enforces cross-time consistency along the denoising trajectory. It reformulates denoising as a reinforcement learning policy evaluation problem, showing significant improvements in sample quality (FID), especially for few-step samplers.
View Cached Full Text
Cached at: 06/16/26, 11:37 AM
# Temporal Difference Learning for Diffusion Models
Source: [https://arxiv.org/html/2606.15048](https://arxiv.org/html/2606.15048)
###### 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\.Code is available at[https://github\.com/StephenYing/Temporal\_Difference\_Learning\_for\_Diffusion\_Models](https://github.com/StephenYing/Temporal_Difference_Learning_for_Diffusion_Models)\.
Diffusion Models, Temporal Difference Learning
Figure 1:Overview of our algorithm\. By matching the model drift \(calculated from the posterior means\) with diffusion drift, we enforce consistency over denoising time index/noise level\.## 1Introduction
Diffusion models have become a standard tool for high\-fidelity generative modeling across images, audio, and beyond\(Ho et al\.,[2020](https://arxiv.org/html/2606.15048#bib.bib10); Song et al\.,[2021b](https://arxiv.org/html/2606.15048#bib.bib23); Karras et al\.,[2022](https://arxiv.org/html/2606.15048#bib.bib13)\)\. Despite impressive progress in*sampler*design \(e\.g\., probability\-flow ODE/DDIM\(Song et al\.,[2021a](https://arxiv.org/html/2606.15048#bib.bib22),[b](https://arxiv.org/html/2606.15048#bib.bib23)\), high\-order solvers such as DPM\-Solver\(Lu et al\.,[2022](https://arxiv.org/html/2606.15048#bib.bib15)\)and UniPC\(Zhao et al\.,[2023](https://arxiv.org/html/2606.15048#bib.bib32)\)\) and*training\-time accelerations*\(e\.g\., progressive distillation\(Salimans & Ho,[2022](https://arxiv.org/html/2606.15048#bib.bib21)\)and consistency\-style learning\(Song et al\.,[2023](https://arxiv.org/html/2606.15048#bib.bib24)\)\), the dominant training paradigm still optimizes*single\-time*reconstruction/noise\-prediction losses\.
Such single\-time objectives do not explicitly require that predictions made at different noise levels form a*time\-consistent*trajectory under the known forward corruption process\. The resulting cross\-time mismatch can accumulate along the denoising path and becomes particularly detrimental when the sampler uses few steps, i\.e\., a small number of function evaluations \(NFE\), where local errors have limited opportunity to average out\(Song et al\.,[2021b](https://arxiv.org/html/2606.15048#bib.bib23); Karras et al\.,[2022](https://arxiv.org/html/2606.15048#bib.bib13)\)\. This motivates treating diffusion training through the lens of sequential decision making, where predictions at different timesteps must be consistent over multiple steps rather than only locally accurate\.
Another line of recent work targets the same few\-step sampling bottleneck by modifying the learned transport or denoising map itself\. Shortcut models condition the network on both the current noise level and a desired step size, enabling the same model to take short or long denoising jumps at inference time\(Frans et al\.,[2025](https://arxiv.org/html/2606.15048#bib.bib7)\)\. MeanFlow instead learns an average velocity field for one\-step generative modeling, providing a flow\-based alternative to modeling only instantaneous velocities\(Geng et al\.,[2026](https://arxiv.org/html/2606.15048#bib.bib8)\)\. These approaches are closely related in motivation to fast sampling, but they primarily redesign the generator or its parameterization for one\-/few\-step generation\. Our goal is complementary: we keep the base diffusion or consistency\-training objective and add a TD regularizer that enforces cross\-time consistency of posterior\-mean drifts along the denoising trajectory\.
Recent years have seen a surge of works that formulate diffusion sampling as a multi\-step decision or control problem and apply reinforcement learning \(RL\) or control\-based methods to optimize non\-differentiable rewards\. DDPO casts denoising as a Markov decision process \(MDP\) and shows policy\-gradient updates can align text\-to\-image models with black\-box objectives such as aesthetics and compressibility\(Black et al\.,[2023](https://arxiv.org/html/2606.15048#bib.bib2)\)\. DPOK performs online RL with KL regularization to fine\-tune diffusion models from human\-trained reward functions, improving both alignment and fidelity\(Fan et al\.,[2023](https://arxiv.org/html/2606.15048#bib.bib5)\)\. Adjoint Matching formulates reward fine\-tuning for flow and diffusion models through stochastic optimal control, yielding a regression\-style objective for reward\-guided model improvement\(Domingo i Enrich et al\.,[2025](https://arxiv.org/html/2606.15048#bib.bib4)\)\.
Other algorithmic variants include LOOP, which analyzes the efficiency–performance trade\-off between REINFORCE and PPO and proposes a leave\-one\-out PPO scheme for diffusion fine\-tuning\(Gupta et al\.,[2025](https://arxiv.org/html/2606.15048#bib.bib9)\); and SEPO, which develops a policy\-gradient method for discrete diffusion with theoretical justification and strong results across discrete generative tasks\(Zekri & Boullé,[2025](https://arxiv.org/html/2606.15048#bib.bib30)\)\. Beyond images, RL fine\-tuning has been applied to text\-to\-speech diffusion models via a loss\-guided policy optimization objective \(DLPO\)\(Chen et al\.,[2024](https://arxiv.org/html/2606.15048#bib.bib3)\)\. Other directions include self\-play \(SPIN\-Diffusion\), where a model competes with its past checkpoints to iteratively improve under a reward signal\(Yuan et al\.,[2024](https://arxiv.org/html/2606.15048#bib.bib29)\), and forward\-process RL that integrates reinforcement signals into flow\-/score\-matching objectives for online fine\-tuning\(Zheng et al\.,[2025](https://arxiv.org/html/2606.15048#bib.bib33)\)\.
Finally, at the theoretical level, Temporal Difference Flows connect TD learning with flow\-based training, providing an RL interpretation of generative flows\(Farebrother et al\.,[2025](https://arxiv.org/html/2606.15048#bib.bib6)\)\. Earlier work on generative TD learning proposed theγ\\gamma\-model framework, which reinterprets TD updates as a generative modeling problem for infinite\-horizon prediction\(Janner et al\.,[2020](https://arxiv.org/html/2606.15048#bib.bib11)\)\. While these methods highlight the synergy between RL and generative models, they focus either on flow\-based models or on predictive state distributions\.
A complementary strand leverages diffusion*as the policy class*for decision making in RL: Diffuser denoises entire trajectories to plan behaviors\(Janner et al\.,[2022](https://arxiv.org/html/2606.15048#bib.bib12)\), Diffusion\-QL represents policies with conditional diffusion models for offline RL\(Wang et al\.,[2023](https://arxiv.org/html/2606.15048#bib.bib28)\), and hierarchical methods introduce subgoal\-conditioned diffusion for long\-horizon tasks\(Li et al\.,[2023](https://arxiv.org/html/2606.15048#bib.bib14)\)\. These approaches focus on maximizing environmental returns in external tasks rather than developing new training regimes for diffusion models\.
In contrast, our work performs*policy evaluation*over the denoising process itself\. We reformulate diffusion as a Markov reward process \(MRP\) and introduce a TD objective that enforces cross\-time consistency of predictions along the denoising path, unified across discrete\- and continuous\-time formulations\. Unlike prior consistency models \(CM\)\(Song et al\.,[2023](https://arxiv.org/html/2606.15048#bib.bib24)\)that require the reconstructions to agree over time, our method encourages the change in posterior means between two times to match the true diffusion drift, as illustrated in[Figure1](https://arxiv.org/html/2606.15048#S0.F1)\. Furthermore, to stabilize optimization across different time pairs, we propose a principled sample\-based loss reweighting scheme that equalizes loss scales\. Rather than learning a new one\-step generator or steering outputs through task\-specific rewards, our TD formulation serves as a general\-purpose training objective that improves fixed\-NFE generation by aligning the model’s internal temporal dynamics\.
The rest of the paper is organized as follows\. In Sec\.[2](https://arxiv.org/html/2606.15048#S2), we review background and notation, including a unified two\-time\-mean form\. Sec\.[3](https://arxiv.org/html/2606.15048#S3)presents our TD objective \(discrete derivation, unified form, design rules, etc\.\)\. We present the experiment results in Sec\.[4](https://arxiv.org/html/2606.15048#S4)and discuss our method in Sec\.[5](https://arxiv.org/html/2606.15048#S5)\. The appendices include implementation details, scheduler definitions, and additional derivations\.
## 2Background and Notation
Diffusion models corrupt data with a forward \(noising\) process and learn a denoiser or score function that inverts it\(Ho et al\.,[2020](https://arxiv.org/html/2606.15048#bib.bib10); Song et al\.,[2021b](https://arxiv.org/html/2606.15048#bib.bib23)\)\. Discrete\-time formulations such as DDPM\(Ho et al\.,[2020](https://arxiv.org/html/2606.15048#bib.bib10); Nichol & Dhariwal,[2021](https://arxiv.org/html/2606.15048#bib.bib19)\)and the deterministic DDIM\(Song et al\.,[2021a](https://arxiv.org/html/2606.15048#bib.bib22)\)provide simple training objectives and flexible sampling schedules, while continuous\-time formulations based on differential equations \(ODE/SDE\) unify these views and support principled SDE samplers\(Song et al\.,[2021b](https://arxiv.org/html/2606.15048#bib.bib23)\)\. To address the complex design space of diffusion models, EDM\(Karras et al\.,[2022](https://arxiv.org/html/2606.15048#bib.bib13)\)modularizes the framework, refines several design choices \(e\.g\., noise grids, preconditioning, loss weighting, etc\.\) and achieves strong empirical results with few\-step sampling\. In this work, we adopt these foundations and focus on enforcing*cross\-time consistency*at training time\.
Table 1:Unified two\-time posterior mean coefficients in𝝁τtrue\(𝒙t,𝒙0\)=At,τ𝒙0\+κt,τ𝒙t\{\\bm\{\\mu\}\}\_\{\\tau\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)=A\_\{t,\\tau\}\{\\bm\{x\}\}\_\{0\}\+\\kappa\_\{t,\\tau\}\{\\bm\{x\}\}\_\{t\}for DDPM\(Ho et al\.,[2020](https://arxiv.org/html/2606.15048#bib.bib10)\), DDIM\(Song et al\.,[2021a](https://arxiv.org/html/2606.15048#bib.bib22)\), VP/VE\-SDE\(Song et al\.,[2021b](https://arxiv.org/html/2606.15048#bib.bib23)\), EDM\(Karras et al\.,[2022](https://arxiv.org/html/2606.15048#bib.bib13)\), and CM\(Song et al\.,[2023](https://arxiv.org/html/2606.15048#bib.bib24)\)\. For DDPM/DDIM, we often use the adjacent\-stepτ=t−1\\tau=t\-1\(though the formula holds for anyτ<t\\tau<t\)\. For VP/VE/EDM/CM,\(t,τ\)\(t,\\tau\)are continuous time indices \(equivalently, noise levels\) selected by the sampler\. See Appendix[B](https://arxiv.org/html/2606.15048#A2)for definitions and derivations\.### 2\.1Unified Two\-Time Posterior Mean
Letttdenote the time index: a discrete timet∈\{0,…,T\}t\\in\\\{0,\\dots,T\\\}\(DDPM/DDIM\), or a continuous timet∈\[0,T\]t\\in\[0,T\]\(ODE/SDE\)\. Letτ<t\\tau<tbe an*earlier \(cleaner\)*time step\. Across families, the mean of the posteriorq\(𝒙τ\|𝒙t,𝒙0\)q\(\{\\bm\{x\}\}\_\{\\tau\}\|\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)admits the*same linear form*\(derived in Appendix[B](https://arxiv.org/html/2606.15048#A2)\)
𝝁τtrue\(𝒙t,𝒙0\)=At,τ𝒙0\+κt,τ𝒙t,\{\\bm\{\\mu\}\}\_\{\\tau\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)=A\_\{t,\\tau\}\\,\{\\bm\{x\}\}\_\{0\}\+\\kappa\_\{t,\\tau\}\\,\{\\bm\{x\}\}\_\{t\},\(1\)where𝒙0\{\\bm\{x\}\}\_\{0\}is a clean datum \(e\.g\., original image\) and𝒙t\{\\bm\{x\}\}\_\{t\}is a noisy sample at leveltt\. Family\-specific choices for\(At,τ,κt,τ\)\(A\_\{t,\\tau\},\\kappa\_\{t,\\tau\}\)are summarized in Table[1](https://arxiv.org/html/2606.15048#S2.T1)\. This linear form will allow us to define a surrogate mean by predicting𝒙0\{\\bm\{x\}\}\_\{0\}using a model𝒙θ,0\(𝒙t;t\)≈𝒙0\{\\bm\{x\}\}\_\{\\theta,0\}\(\{\\bm\{x\}\}\_\{t\};t\)\\approx\{\\bm\{x\}\}\_\{0\}parametrized byθ\\theta\(e\.g\., theDθD\_\{\\theta\}andfθf\_\{\\theta\}in the following paragraphs\)\. This will be the main strategy for our TD objective \(Sec\.[3](https://arxiv.org/html/2606.15048#S3)\)\.
#### EDM\-style preconditioning\.
EDM\(Karras et al\.,[2022](https://arxiv.org/html/2606.15048#bib.bib13)\)wraps a raw networkFθF\_\{\\theta\}in the preconditioned denoiser as111Note that EDM adopts a noise scheme instead of a time scheme, but we use time here as they are equivalent and it can facilitate later discussion\.
Dθ\(𝒙;t\)=cskipEDM\(t\)𝒙\+coutEDM\(t\)Fθ\(cin\(t\)𝒙;cnoise\(t\)\),D\_\{\\theta\}\(\{\\bm\{x\}\};t\)=c\_\{\\text\{skip\}\}^\{\\text\{EDM\}\}\(t\)\\,\{\\bm\{x\}\}\+c\_\{\\text\{out\}\}^\{\\text\{EDM\}\}\(t\)\\,F\_\{\\theta\}\\bigl\(c\_\{\\text\{in\}\}\(t\)\\,\{\\bm\{x\}\};c\_\{\\text\{noise\}\}\(t\)\\bigr\),\(2\)and trains this model with a weighted regression
ℒEDM=𝔼t,𝒙0,𝒙t\[w\(t\)‖Dθ\(𝒙t;t\)−𝒙0‖22\]\.\\mathcal\{L\}\_\{\\mathrm\{EDM\}\}=\\mathbb\{E\}\_\{t,\{\\bm\{x\}\}\_\{0\},\{\\bm\{x\}\}\_\{t\}\}\\Big\[w\(t\)\\,\\big\\\|D\_\{\\theta\}\(\{\\bm\{x\}\}\_\{t\};t\)\-\{\\bm\{x\}\}\_\{0\}\\big\\\|\_\{2\}^\{2\}\\Big\]\.\(3\)The specific choices of\{cskipEDM,coutEDM,cin,cnoise,w\(t\)\}\\\{c\_\{\\text\{skip\}\}^\{\\text\{EDM\}\},c\_\{\\text\{out\}\}^\{\\text\{EDM\}\},c\_\{\\text\{in\}\},c\_\{\\text\{noise\}\},w\(t\)\\\}equalize the effective variance across noise levels and improve optimization conditioning\(Karras et al\.,[2022](https://arxiv.org/html/2606.15048#bib.bib13)\)\. We will develop a weighting to achieve a similar effect later, and combine Equation \([3](https://arxiv.org/html/2606.15048#S2.E3)\) with our TD objective\.
#### Preconditioning for consistency models\.
Similarly to EDM, CM\(Song et al\.,[2023](https://arxiv.org/html/2606.15048#bib.bib24)\)parametrizes a prediction modelfθ\(𝒙t;t\)f\_\{\\theta\}\(\{\\bm\{x\}\}\_\{t\};t\)to predict𝒙ϵ\{\\bm\{x\}\}\_\{\\epsilon\}using anyt∈\[ϵ,T\]t\\in\[\\epsilon,T\]whereϵ\>0\\epsilon\>0is a small number\. To ensure that the boundary conditionfθ\(𝒙ϵ;ϵ\)=𝒙ϵf\_\{\\theta\}\(\{\\bm\{x\}\}\_\{\\epsilon\};\\epsilon\)=\{\\bm\{x\}\}\_\{\\epsilon\}is satisfied by construction, CM adopts a preconditioned form
fθ\(𝒙;t\)=cskipCM\(t\)𝒙\+coutCM\(t\)Fθ\(𝒙;t\),f\_\{\\theta\}\(\{\\bm\{x\}\};t\)=c\_\{\\text\{skip\}\}^\{\\text\{CM\}\}\(t\)\\,\{\\bm\{x\}\}\+c\_\{\\text\{out\}\}^\{\\text\{CM\}\}\(t\)\\,F\_\{\\theta\}\(\{\\bm\{x\}\};t\),\(4\)satisfyingcskipCM\(ϵ\)=1,coutCM\(ϵ\)=0c\_\{\\text\{skip\}\}^\{\\text\{CM\}\}\(\\epsilon\)=1,\\ c\_\{\\text\{out\}\}^\{\\text\{CM\}\}\(\\epsilon\)=0, which is a boundary\-shifted variant of EDM’s parameterization\. When training from scratch \(i\.e\., consistency training or CT\), it optimizes
ℒCT=𝔼\[w\(t\)‖fθ\(𝒙t\+1;t\+1\)−fθ′\(𝒙t,t\)‖22\],\\displaystyle\\mathcal\{L\}\_\{\\mathrm\{CT\}\}=\\mathbb\{E\}\\Big\[w\(t\)\\,\\big\\\|f\_\{\\theta\}\(\{\\bm\{x\}\}\_\{t\+1\};t\+1\)\-f\_\{\\theta^\{\\prime\}\}\(\{\\bm\{x\}\}\_\{t\},t\)\\big\\\|\_\{2\}^\{2\}\\Big\],\(5\)whereθ′\\theta^\{\\prime\}is updated via exponential moving average \(EMA\)\.
### 2\.2Temporal Difference Learning
In this work, we will recast training diffusion models as policy evaluation problems in reinforcement learning \(RL\)\(Pan et al\.,[2025](https://arxiv.org/html/2606.15048#bib.bib20)\)\. A finite\-horizon Markov Reward Process \(MRP\)\(Szepesvári,[2022](https://arxiv.org/html/2606.15048#bib.bib26)\)is characterized by a tuple\(𝒳,rt,Pt,T\)\(\{\\mathcal\{X\}\},r\_\{t\},P\_\{t\},T\)where𝒳\{\\mathcal\{X\}\}is the \(common\) state space,rt:𝒳×𝒳↦ℝr\_\{t\}:\{\\mathcal\{X\}\}\\times\{\\mathcal\{X\}\}\\mapsto\{\\mathbb\{R\}\}is the reward function at timett,Pt:𝒳↦Δ\(𝒳\)P\_\{t\}:\{\\mathcal\{X\}\}\\mapsto\\Delta\(\{\\mathcal\{X\}\}\)is the transition kernel at timettandTTis the length of the episode\. However, to match the time notation in diffusion models, here we let the MRP start from stept=Tt=Tand traverse*backward*tot=0t=0\. Specifically, when transitioning from𝒙t\{\\bm\{x\}\}\_\{t\}to𝒙t−1\{\\bm\{x\}\}\_\{t\-1\}according toPt\(⋅\|𝒙t\)P\_\{t\}\(\\cdot\|\{\\bm\{x\}\}\_\{t\}\), we receive a reward ofrt−1\(𝒙t,𝒙t−1\)r\_\{t\-1\}\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{t\-1\}\)\. Letgt:=∑i=1trt−ig\_\{t\}:=\\sum\_\{i=1\}^\{t\}r\_\{t\-i\}denote the MRP return collected from timettuntil the terminal time0, wherert−ir\_\{t\-i\}is short forrt−i\(xt−i\+1,xt−i\)r\_\{t\-i\}\(x\_\{t\-i\+1\},x\_\{t\-i\}\)\. The state value function is then defined as the expected return when starting from a given state:
vt\(𝒙\):=𝔼\[gt\|𝒙t=𝒙\]:=𝔼\[∑i=1trt−i\|𝒙t=𝒙\]\\displaystyle v\_\{t\}\(\{\\bm\{x\}\}\):=\\mathbb\{E\}\\left\[g\_\{t\}\\ \\middle\|\\ \{\\bm\{x\}\}\_\{t\}=\{\\bm\{x\}\}\\right\]:=\\mathbb\{E\}\\left\[\\sum\_\{i=1\}^\{t\}r\_\{t\-i\}\\ \\middle\|\\ \{\\bm\{x\}\}\_\{t\}=\{\\bm\{x\}\}\\right\]\(6\)where the expectation is taken overPiP\_\{i\}for0<i≤t0<i\\leq t\.
One can use temporal difference \(TD\) learning\(Sutton & Barto,[2018](https://arxiv.org/html/2606.15048#bib.bib25)\)to findvtv\_\{t\}\. Given some parametrized functionsvθ,tv\_\{\\theta,t\}, the TD error of the transition\(𝒙t,𝒙t−1\)\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{t\-1\}\)is
δt:=rt−1\+vθ,t−1\(𝒙t−1\)−vθ,t\(𝒙t\)\.\\displaystyle\\delta\_\{t\}:=r\_\{t\-1\}\+v\_\{\\theta,t\-1\}\(\{\\bm\{x\}\}\_\{t\-1\}\)\-v\_\{\\theta,t\}\(\{\\bm\{x\}\}\_\{t\}\)\.\(7\)By minimizing the TD error using samples from the MRP and semi\-gradient descent,vθ,tv\_\{\\theta,t\}will converge to the truevtv\_\{t\}\(Bertsekas & Tsitsiklis,[1996](https://arxiv.org/html/2606.15048#bib.bib1); Tsitsiklis & Van Roy,[1996](https://arxiv.org/html/2606.15048#bib.bib27)\)\.
## 3Temporal Difference Learning for Diffusion Models
This section develops TD training for diffusion models\. We first derive the TD loss on a discrete time index grid, then lift it to a unified formulation that applies to both discrete\- \(DDPM/DDIM\) and continuous\-time \(VP/VE/EDM\) families\. Whenever a posterior mean is needed, we*do not*restate family\-specific formulas – rather, we directly use the unified form in Equation \([1](https://arxiv.org/html/2606.15048#S2.E1)\)\.
### 3\.1Discrete Time: MRP Formulation for DDPM
#### MRP specifications on the time index grid\.
Following Sec\.[2\.2](https://arxiv.org/html/2606.15048#S2.SS2), we define an MRP, traversing backward in time fromt=Tt=Ttot=0t=0\. The construction is tied to a clean sample𝒙0\{\\bm\{x\}\}\_\{0\}: the rewards and posterior transitions below are defined conditionally on a clean𝒙0\{\\bm\{x\}\}\_\{0\}sampled from the data distribution\. In other words,𝒙0\{\\bm\{x\}\}\_\{0\}acts as an episode\-level context throughout the derivation, and we omit it when no confusion arises\. For training objectives,𝒙0∼qdata\{\\bm\{x\}\}\_\{0\}\\sim q\_\{\\rm data\}is a random variable, as shown later in Equation \([16](https://arxiv.org/html/2606.15048#S3.E16)\); within a sampled episode, the realized𝒙0\{\\bm\{x\}\}\_\{0\}is treated as fixed\. In our context of diffusion model training, rewards, returns, and values are all data\-space vector\-valued quantities, which corresponds to a multiple\-reward setting\. The MRP is specified as follows:
- •*State space:*𝒳\{\\mathcal\{X\}\}is the data space\. In the case of image generation,𝒳\{\\mathcal\{X\}\}is the set of images\.
- •*Reward function:*𝒓t−1:=𝒓t−1\(𝒙t,𝒙t−1\):=𝝁t−1true\(𝒙t,𝒙0\)−𝝁t−2true\(𝒙t−1,𝒙0\)\{\\bm\{r\}\}\_\{t\-1\}:=\{\\bm\{r\}\}\_\{t\-1\}\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{t\-1\}\):=\{\\bm\{\\mu\}\}\_\{t\-1\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)\-\{\\bm\{\\mu\}\}\_\{t\-2\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{t\-1\},\{\\bm\{x\}\}\_\{0\}\)is the posterior mean difference\. Note that𝒓t−1\{\\bm\{r\}\}\_\{t\-1\}is a vector in the data space, as opposed to a scalar reward common in reinforcement learning\. In the final step when transitioning from𝒙1\{\\bm\{x\}\}\_\{1\}to𝒙0\{\\bm\{x\}\}\_\{0\}, the reward is defined as𝒓0:=𝝁0true\(𝒙1,𝒙0\)−𝒙0=𝟎\{\\bm\{r\}\}\_\{0\}:=\{\\bm\{\\mu\}\}\_\{0\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{1\},\{\\bm\{x\}\}\_\{0\}\)\-\{\\bm\{x\}\}\_\{0\}=\{\\bm\{0\}\}\. The last equation is because𝝁0true\(𝒙1,𝒙0\)\{\\bm\{\\mu\}\}\_\{0\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{1\},\{\\bm\{x\}\}\_\{0\}\), the conditional mean, must be𝒙0\{\\bm\{x\}\}\_\{0\}*given*𝒙0\{\\bm\{x\}\}\_\{0\}\.
- •*Transition kernel:*Pt\(𝒙t−1∣𝒙t\):=q\(𝒙t−1\|𝒙t,𝒙0\)P\_\{t\}\(\{\\bm\{x\}\}\_\{t\-1\}\\mid\{\\bm\{x\}\}\_\{t\}\):=q\(\{\\bm\{x\}\}\_\{t\-1\}\|\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)is induced by the posterior of the predecessor in the diffusion process \(stochastic in𝒙t−1\{\\bm\{x\}\}\_\{t\-1\}through the forward coupling\)\.
#### Return, value, TD\(0\) andkk\-step return\.
In this diffusion MRP setup, the cumulative reward has a simple form\. Based on the definition of reward, the return𝒈t\{\\bm\{g\}\}\_\{t\}\(conditioned on𝒙0\{\\bm\{x\}\}\_\{0\}\) then equals the displacement of the posterior mean from the data:
𝒈t\|𝒙0=𝝁t−1true\(𝒙t,𝒙0\)−𝒙0,\{\\bm\{g\}\}\_\{t\}\\ \|\\ \{\\bm\{x\}\}\_\{0\}=\{\\bm\{\\mu\}\}\_\{t\-1\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)\-\{\\bm\{x\}\}\_\{0\},\(8\)It satisfies the usual return recursion, here written backward in diffusion time,𝒈t=𝒓t−1\+𝒈t−1\{\\bm\{g\}\}\_\{t\}=\{\\bm\{r\}\}\_\{t\-1\}\+\{\\bm\{g\}\}\_\{t\-1\}\. Moreover, it ensures that𝒈1=𝝁0true\(𝒙1,𝒙0\)−𝒙0=𝒙0−𝒙0=𝟎\{\\bm\{g\}\}\_\{1\}=\{\\bm\{\\mu\}\}\_\{0\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{1\},\{\\bm\{x\}\}\_\{0\}\)\\ \-\\ \{\\bm\{x\}\}\_\{0\}=\{\\bm\{x\}\}\_\{0\}\-\{\\bm\{x\}\}\_\{0\}=\{\\bm\{0\}\}\. Because the posterior mean in Equation \([8](https://arxiv.org/html/2606.15048#S3.E8)\) is determined once\(𝒙t,𝒙0\)\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)is known, the corresponding value function reduces to this conditional return:
𝒗t\(𝒙t\):=𝔼\[𝒈t∣𝒙t,𝒙0\]=𝝁t−1true\(𝒙t,𝒙0\)−𝒙0\.\{\\bm\{v\}\}\_\{t\}\(\{\\bm\{x\}\}\_\{t\}\):=\\mathbb\{E\}\[\\,\{\\bm\{g\}\}\_\{t\}\\mid\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\\,\]=\{\\bm\{\\mu\}\}\_\{t\-1\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)\-\{\\bm\{x\}\}\_\{0\}\.\(9\)
Here, the subscriptttin𝒗t\{\\bm\{v\}\}\_\{t\}denotes the diffusion\-time index of the quantity being estimated \(namely𝒈t\{\\bm\{g\}\}\_\{t\}\), while𝒙t\{\\bm\{x\}\}\_\{t\}is the conditioning state in the expectation\. We approximate𝒗t\{\\bm\{v\}\}\_\{t\}using the preconditioned denoiser, e\.g\., \([2](https://arxiv.org/html/2606.15048#S2.E2)\) or \([4](https://arxiv.org/html/2606.15048#S2.E4)\) earlier, as the model𝒙θ,0\{\\bm\{x\}\}\_\{\\theta,0\}:
𝒗t\(𝒙t\)\\displaystyle\{\\bm\{v\}\}\_\{t\}\(\{\\bm\{x\}\}\_\{t\}\)≈𝒗θ,t\(𝒙t\):=𝝁θ,t−1\(𝒙t\)−𝒙0\\displaystyle\\approx\{\\bm\{v\}\}\_\{\\theta,t\}\(\{\\bm\{x\}\}\_\{t\}\):=\{\\bm\{\\mu\}\}\_\{\\theta,t\-1\}\(\{\\bm\{x\}\}\_\{t\}\)\-\{\\bm\{x\}\}\_\{0\}\(10\)𝝁θ,t−1\(𝒙t\)\\displaystyle\{\\bm\{\\mu\}\}\_\{\\theta,t\-1\}\(\{\\bm\{x\}\}\_\{t\}\):=At,t−1𝒙θ,0\(𝒙t;t\)\+κt,t−1𝒙t\.\\displaystyle:=A\_\{t,t\-1\}\\,\{\\bm\{x\}\}\_\{\\theta,0\}\(\{\\bm\{x\}\}\_\{t\};t\)\+\\kappa\_\{t,t\-1\}\\,\{\\bm\{x\}\}\_\{t\}\.\(11\)
We learnθ\\thetaby constructing a bootstrap target with fixed parameters\. The “next”\-state’s value is estimated by
𝒗t−1\(𝒙t−1\)≈\\displaystyle\{\\bm\{v\}\}\_\{t\-1\}\(\{\\bm\{x\}\}\_\{t\-1\}\)\\\!\\approx𝒗θ′,t−1\(𝒙t−1\):=𝝁θ′,t−2\(𝒙t−1\)−𝒙0\\displaystyle\\ \{\\bm\{v\}\}\_\{\\theta^\{\\prime\},t\-1\}\(\{\\bm\{x\}\}\_\{t\-1\}\)\\\!:=\\\!\{\\bm\{\\mu\}\}\_\{\\theta^\{\\prime\},t\-2\}\(\{\\bm\{x\}\}\_\{t\-1\}\)\\\!\-\\\!\{\\bm\{x\}\}\_\{0\}\(12\)𝝁θ′,t−2\(𝒙t−1\)\\displaystyle\{\\bm\{\\mu\}\}\_\{\\theta^\{\\prime\},t\-2\}\(\{\\bm\{x\}\}\_\{t\-1\}\)=At−1,t−2𝒙θ′,0\(𝒙t−1;t−1\)\\displaystyle=A\_\{t\-1,t\-2\}\\,\{\\bm\{x\}\}\_\{\\theta^\{\\prime\},0\}\(\{\\bm\{x\}\}\_\{t\-1\};t\-1\)\+κt−1,t−2𝒙t−1,\\displaystyle\\qquad\+\\kappa\_\{t\-1,t\-2\}\\,\{\\bm\{x\}\}\_\{t\-1\},\(13\)whereθ′\\theta^\{\\prime\}is the fixed \(stop\-gradient\) target network’s parameters, updated using EMA\. By the definition of the reward, the bootstrap target is
𝒓t−1\+𝒗θ′,t−1\(𝒙t−1\)\\displaystyle\{\\bm\{r\}\}\_\{t\-1\}\+\{\\bm\{v\}\}\_\{\\theta^\{\\prime\},t\-1\}\(\{\\bm\{x\}\}\_\{t\-1\}\)=𝝁t−1true\(𝒙t,𝒙0\)−𝝁t−2true\(𝒙t−1,𝒙0\)\\displaystyle=\{\\bm\{\\mu\}\}\_\{t\-1\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)\-\{\\bm\{\\mu\}\}\_\{t\-2\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{t\-1\},\{\\bm\{x\}\}\_\{0\}\)\+𝝁θ′,t−2\(𝒙t−1\)−𝒙0\.\\displaystyle\\quad\+\{\\bm\{\\mu\}\}\_\{\\theta^\{\\prime\},t\-2\}\(\{\\bm\{x\}\}\_\{t\-1\}\)\-\{\\bm\{x\}\}\_\{0\}\.\(14\)Combining it with Equation \([10](https://arxiv.org/html/2606.15048#S3.E10)\) gives the TD error
𝜹t\\displaystyle\{\\bm\{\\delta\}\}\_\{t\}:=𝒓t−1\+𝒗θ′,t−1\(𝒙t−1\)−𝒗θ,t\(𝒙t\)\\displaystyle:=\{\\bm\{r\}\}\_\{t\-1\}\+\{\\bm\{v\}\}\_\{\\theta^\{\\prime\},t\-1\}\(\{\\bm\{x\}\}\_\{t\-1\}\)\-\{\\bm\{v\}\}\_\{\\theta,t\}\(\{\\bm\{x\}\}\_\{t\}\)=\[𝝁t−1true\(𝒙t,𝒙0\)−𝝁t−2true\(𝒙t−1,𝒙0\)\]⏟one\-step diffusion drift\\displaystyle=\\underbrace\{\[\{\\bm\{\\mu\}\}\_\{t\-1\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)\-\{\\bm\{\\mu\}\}\_\{t\-2\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{t\-1\},\{\\bm\{x\}\}\_\{0\}\)\]\}\_\{\\text\{one\-step diffusion drift\}\}−\[𝝁θ,t−1\(𝒙t\)−𝝁θ′,t−2\(𝒙t−1\)\]⏟one\-step model drift\.\\displaystyle\\quad\-\\underbrace\{\[\{\\bm\{\\mu\}\}\_\{\\theta,t\-1\}\(\{\\bm\{x\}\}\_\{t\}\)\-\{\\bm\{\\mu\}\}\_\{\\theta^\{\\prime\},t\-2\}\(\{\\bm\{x\}\}\_\{t\-1\}\)\]\}\_\{\\text\{one\-step model drift\}\}\.\(15\)The TD\(0\) objective is then
ℒTD\(0\):=𝔼t,𝒙0,𝒙t−1,𝒙t\[‖𝜹t‖22\],\\mathcal\{L\}\_\{\\mathrm\{TD\(0\)\}\}:=\\mathbb\{E\}\_\{t,\{\\bm\{x\}\}\_\{0\},\{\\bm\{x\}\}\_\{t\-1\},\{\\bm\{x\}\}\_\{t\}\}\\\!\\left\[\\\|\{\\bm\{\\delta\}\}\_\{t\}\\\|\_\{2\}^\{2\}\\right\],\(16\)wheret∼𝒰\{2,…,T\}t\\sim\{\\mathcal\{U\}\}\\\{2,\\ldots,T\\\},𝒙0∼qdata\{\\bm\{x\}\}\_\{0\}\\sim q\_\{\\mathrm\{data\}\},𝒙t−1∼q\(𝒙t−1∣𝒙0\)\{\\bm\{x\}\}\_\{t\-1\}\\sim q\(\{\\bm\{x\}\}\_\{t\-1\}\\mid\{\\bm\{x\}\}\_\{0\}\), and𝒙t∼q\(𝒙t∣𝒙t−1\)\{\\bm\{x\}\}\_\{t\}\\sim q\(\{\\bm\{x\}\}\_\{t\}\\mid\{\\bm\{x\}\}\_\{t\-1\}\)\.
Minimizing the TD error can be interpreted as aligning diffusion progression across time steps as shown in Equation \([15](https://arxiv.org/html/2606.15048#S3.E15)\)\. If the true mean has drifted in one step, the model should shift in the same way, thus enforcing*consistency*between time steps\. The same derivation applies to DDIM\(Song et al\.,[2021a](https://arxiv.org/html/2606.15048#bib.bib22)\)by substituting its\(At,t−1,κt,t−1\)\(A\_\{t,t\-1\},\\kappa\_\{t,t\-1\}\)from Table[1](https://arxiv.org/html/2606.15048#S2.T1)\.
We can also usekk\-step return as the bootstrap target\. Keep expanding Equation \([14](https://arxiv.org/html/2606.15048#S3.E14)\) overkksteps gives the followingkk\-step objective
ℒTD\(k\)\\displaystyle\\mathcal\{L\}^\{\(k\)\}\_\{\\mathrm\{TD\}\}:=𝔼\[∥\[𝝁t−1true\(𝒙t,𝒙0\)−𝝁t−k−1true\(𝒙t−k,𝒙0\)\]⏟k\-step diffusion drift\\displaystyle:=\\mathbb\{E\}\\Big\[\\big\\\|\\underbrace\{\[\{\\bm\{\\mu\}\}\_\{t\-1\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)\-\{\\bm\{\\mu\}\}\_\{t\-k\-1\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{t\-k\},\{\\bm\{x\}\}\_\{0\}\)\]\}\_\{\\text\{$k$\-step diffusion drift\}\}−\[𝝁θ,t−1\(𝒙t\)−𝝁θ′,t−k−1\(𝒙t−k\)\]⏟k\-step model drift∥22\]\.\\displaystyle\\qquad\-\\underbrace\{\[\{\\bm\{\\mu\}\}\_\{\\theta,t\-1\}\(\{\\bm\{x\}\}\_\{t\}\)\-\{\\bm\{\\mu\}\}\_\{\\theta^\{\\prime\},t\-k\-1\}\(\{\\bm\{x\}\}\_\{t\-k\}\)\]\}\_\{\\text\{$k$\-step model drift\}\}\\big\\\|\_\{2\}^\{2\}\\Big\]\.\(17\)
where the expectation is overt∼𝒰\{k\+1,T\},𝒙0∼qdata\(𝒙0\),𝒙t−k∼q\(𝒙t−k\|𝒙0\),𝒙t∼q\(𝒙t\|𝒙t−k\)t\\\!\\sim\\\!\{\\mathcal\{U\}\}\\\{k\+1,T\\\},\{\\bm\{x\}\}\_\{0\}\\\!\\sim\\\!q\_\{\\text\{data\}\}\(\{\\bm\{x\}\}\_\{0\}\),\{\\bm\{x\}\}\_\{t\-k\}\\\!\\sim\\\!q\(\{\\bm\{x\}\}\_\{t\-k\}\|\{\\bm\{x\}\}\_\{0\}\),\{\\bm\{x\}\}\_\{t\}\\\!\\sim\\\!q\(\{\\bm\{x\}\}\_\{t\}\|\{\\bm\{x\}\}\_\{t\-k\}\)\. The matching process is illustrated in[Figure1](https://arxiv.org/html/2606.15048#S0.F1)\.
#### Aligning model with posterior\.
Note that there is no need to bootstrap when at≤kt\\leq kis sampled because we do know the rest of the episode and the corresponding true return\. In this case, the TD loss actually reduces to posterior mean matching \(cf\. \([8](https://arxiv.org/html/2606.15048#S3.E8)\) and \([10](https://arxiv.org/html/2606.15048#S3.E10)\)\), similar to the DDPM loss
ℒDDPM=𝔼\[w\(t\)‖𝝁t−1true\(𝒙t,𝒙0\)−𝝁θ,t−1\(𝒙t\)‖22\],\\displaystyle\\mathcal\{L\}\_\{\\mathrm\{DDPM\}\}\\\!=\\\!\\mathbb\{E\}\\Big\[w\(t\)\\\|\{\\bm\{\\mu\}\}\_\{t\-1\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)\-\{\\bm\{\\mu\}\}\_\{\\theta,t\-1\}\(\{\\bm\{x\}\}\_\{t\}\)\\\|\_\{2\}^\{2\}\\Big\],\(18\)with a specific weightingw\(t\)w\(t\)defined by DDPM\. In practice, we combine the TD loss with the DDPM loss even whent\>kt\>kto facilitate training and speed up convergence, and the final objective is
ℒTD\+DDPM\(k\)=ℒTD\(k\)\+λℒDDPM,\\displaystyle\\mathcal\{L\}^\{\(k\)\}\_\{\\mathrm\{TD\+DDPM\}\}\\ =\\ \\mathcal\{L\}^\{\(k\)\}\_\{\\mathrm\{TD\}\}\\ \+\\ \\lambda\\ \\mathcal\{L\}\_\{\\mathrm\{DDPM\}\},\(19\)whereλ\>0\\lambda\>0is a hyper\-parameter\. In other words, if a timet\>kt\>kis sampled during training, the TD loss can be used and we optimizeℒTD\+DDPM\(k\)\\mathcal\{L\}^\{\(k\)\}\_\{\\mathrm\{TD\+DDPM\}\}\. Otherwise, we optimize\(1\+λ\)ℒDDPM\(1\+\\lambda\)\\mathcal\{L\}\_\{\\mathrm\{DDPM\}\}since the TD loss reduces to the DDPM loss in this case and the\(1\+λ\)\(1\+\\lambda\)factor maintains the scale of the objective\.
### 3\.2Discrete and Continuous Time: A Unified TD Objective
The derivations based on discrete time steps above can be easily extended to continuous\-time scenarios with ODE/SDE thanks to the unified mean \(Equation \([1](https://arxiv.org/html/2606.15048#S2.E1)\) from Sec\.[2\.1](https://arxiv.org/html/2606.15048#S2.SS1)\)\.
Instead of matching drifts that arekk\-step away in the discrete case, here we pick two time indicest,t′∈\[0,T\]t,t^\{\\prime\}\\in\[0,T\]\. Timett\(resp\.t′t^\{\\prime\}\) induces a true posterior mean for an earlier timeτ<t\\tau<t\(resp\.τ′<t′\\tau^\{\\prime\}<t^\{\\prime\}\)\. Their corresponding means can be expressed as
𝝁τtrue\(𝒙t,𝒙0\)\\displaystyle\{\\bm\{\\mu\}\}\_\{\\tau\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)=At,τ𝒙0\+κt,τ𝒙t\\displaystyle=A\_\{t,\\tau\}\\,\{\\bm\{x\}\}\_\{0\}\+\\kappa\_\{t,\\tau\}\\,\{\\bm\{x\}\}\_\{t\}\(20\)𝝁τ′true\(𝒙t′,𝒙0\)\\displaystyle\{\\bm\{\\mu\}\}\_\{\\tau^\{\\prime\}\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{t^\{\\prime\}\},\{\\bm\{x\}\}\_\{0\}\)=At′,τ′𝒙0\+κt′,τ′𝒙t′\.\\displaystyle=A\_\{t^\{\\prime\},\\tau^\{\\prime\}\}\\,\{\\bm\{x\}\}\_\{0\}\+\\kappa\_\{t^\{\\prime\},\\tau^\{\\prime\}\}\\,\{\\bm\{x\}\}\_\{t^\{\\prime\}\}\.\(21\)For the discrete\-time case \(Equation \([3\.1](https://arxiv.org/html/2606.15048#S3.Ex5)\)\),𝝁t−1true\(𝒙t,𝒙0\)\{\\bm\{\\mu\}\}\_\{t\-1\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)is the posterior mean in the previous time step \(hence the subscriptt−1t\-1\)\. In analogy,τ\\tauhere can be considered as the “previous time step” ofttin the continuous case\. In our experiments, we setτ′<t′<τ<t\\tau^\{\\prime\}<t^\{\\prime\}<\\tau<twith*span*k:=t−t′k:=t\-t^\{\\prime\}and*stride*Δ:=t′−τ′=t−τ<k\\Delta:=t^\{\\prime\}\-\\tau^\{\\prime\}=t\-\\tau<k, imitating the discrete case\. Accordingly, the model is defined as
𝝁θ,τ\(𝒙t\):=At,τ𝒙θ,0\(𝒙t;t\)\+κt,τ𝒙t,\\displaystyle\{\\bm\{\\mu\}\}\_\{\\theta,\\tau\}\(\{\\bm\{x\}\}\_\{t\}\):=A\_\{t,\\tau\}\\,\{\\bm\{x\}\}\_\{\\theta,0\}\(\{\\bm\{x\}\}\_\{t\};t\)\+\\kappa\_\{t,\\tau\}\\,\{\\bm\{x\}\}\_\{t\},\(22\)where𝒙θ,0\{\\bm\{x\}\}\_\{\\theta,0\}predicting𝒙0\{\\bm\{x\}\}\_\{0\}, instantiated asDθD\_\{\\theta\}for EDM \([2](https://arxiv.org/html/2606.15048#S2.E2)\) or \([4](https://arxiv.org/html/2606.15048#S2.E4)\) for CM parameterizations \(Sec\.[2\.1](https://arxiv.org/html/2606.15048#S2.SS1)\)\. The TD loss for continuous time reads
ℒTDcont\\displaystyle\\ \\mathcal\{L\}^\{\\mathrm\{cont\}\}\_\{\\mathrm\{TD\}\}=𝔼𝒙0,t,t′,𝒙t,𝒙t′\[∥\[𝝁τtrue\(𝒙t,𝒙0\)−𝝁τ′true\(𝒙t′,𝒙0\)\]\\displaystyle=\\mathbb\{E\}\_\{\{\\bm\{x\}\}\_\{0\},t,t^\{\\prime\},\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{t^\{\\prime\}\}\}\\Big\[\\big\\\|\\big\[\{\\bm\{\\mu\}\}\_\{\\tau\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)\-\{\\bm\{\\mu\}\}\_\{\\tau^\{\\prime\}\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{t^\{\\prime\}\},\{\\bm\{x\}\}\_\{0\}\)\\big\]−\[𝝁θ,τ\(𝒙t\)−𝝁θ′,τ′\(𝒙t′\)\]∥22\]\.\\displaystyle\\qquad\-\\big\[\{\\bm\{\\mu\}\}\_\{\\theta,\\tau\}\(\{\\bm\{x\}\}\_\{t\}\)\-\{\\bm\{\\mu\}\}\_\{\\theta^\{\\prime\},\\tau^\{\\prime\}\}\(\{\\bm\{x\}\}\_\{t^\{\\prime\}\}\)\\big\]\\big\\\|\_\{2\}^\{2\}\\Big\]\.\(23\)Inspired by prior work, it is preferable to include a weighting schemewTD\(t,t′\)w\_\{\\text\{TD\}\}\(t,t^\{\\prime\}\)for loss per sample within the expectation of Equation \([23](https://arxiv.org/html/2606.15048#S3.E23)\) to avoid drastic changes in gradient magnitudes across time steps\. The weighting scheme then depends on the parametrization of the prediction model and the training strategies of the base algorithm\. In the following, we will incorporate our TD objective into two widely used continuous\-time training paradigms, EDM and CT, and show specific recipes for stable training\.
### 3\.3Training Recipes: TD\+EDM and TD\+CT
#### Weighting for EDM parameterization\.
Here we expand the TD error \([23](https://arxiv.org/html/2606.15048#S3.E23)\) by substituting \([22](https://arxiv.org/html/2606.15048#S3.E22)\) into it
𝜹t,t′\\displaystyle\{\\bm\{\\delta\}\}\_\{t,t^\{\\prime\}\}:=\[𝝁τtrue\(𝒙t,𝒙0\)−𝝁τ′true\(𝒙t′,𝒙0\)\]\\displaystyle:=\\big\[\{\\bm\{\\mu\}\}\_\{\\tau\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)\-\{\\bm\{\\mu\}\}\_\{\\tau^\{\\prime\}\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{t^\{\\prime\}\},\{\\bm\{x\}\}\_\{0\}\)\\big\]−\[𝝁θ,τ\(𝒙t\)−𝝁θ′,τ′\(𝒙t′\)\]\\displaystyle\\qquad\-\\big\[\{\\bm\{\\mu\}\}\_\{\\theta,\\tau\}\(\{\\bm\{x\}\}\_\{t\}\)\-\{\\bm\{\\mu\}\}\_\{\\theta^\{\\prime\},\\tau^\{\\prime\}\}\(\{\\bm\{x\}\}\_\{t^\{\\prime\}\}\)\\big\]\(24\)=\[At,τ𝒙0−At′,τ′𝒙0\]\\displaystyle=\\big\[A\_\{t,\\tau\}\{\\bm\{x\}\}\_\{0\}\-A\_\{t^\{\\prime\},\\tau^\{\\prime\}\}\{\\bm\{x\}\}\_\{0\}\\big\]−\[At,τ𝒙θ,0\(𝒙t,t\)−At′,τ′𝒙θ′,0\(𝒙t′,t′\)\]\.\\displaystyle\\qquad\-\\big\[A\_\{t,\\tau\}\{\\bm\{x\}\}\_\{\\theta,0\}\(\{\\bm\{x\}\}\_\{t\},t\)\-A\_\{t^\{\\prime\},\\tau^\{\\prime\}\}\{\\bm\{x\}\}\_\{\\theta^\{\\prime\},0\}\(\{\\bm\{x\}\}\_\{t^\{\\prime\}\},t^\{\\prime\}\)\\big\]\.\(25\)When using the EDM preconditioned denoiserDθD\_\{\\theta\}\([2](https://arxiv.org/html/2606.15048#S2.E2)\) as𝒙θ,0\{\\bm\{x\}\}\_\{\\theta,0\}, the TD error can be decomposed further by using the raw networkFθF\_\{\\theta\}
𝜹t,t′\\displaystyle\{\\bm\{\\delta\}\}\_\{t,t^\{\\prime\}\}=\[At,τ𝒙0−At′,τ′𝒙0\]\\displaystyle=\\big\[A\_\{t,\\tau\}\{\\bm\{x\}\}\_\{0\}\-A\_\{t^\{\\prime\},\\tau^\{\\prime\}\}\{\\bm\{x\}\}\_\{0\}\\big\]−\[At,τcskipEDM\(t\)𝒙t−At′,τ′cskipEDM\(t′\)𝒙t′\]\\displaystyle\\quad\-\\big\[A\_\{t,\\tau\}c\_\{\\text\{skip\}\}^\{\\text\{EDM\}\}\(t\)\{\\bm\{x\}\}\_\{t\}\-A\_\{t^\{\\prime\},\\tau^\{\\prime\}\}c\_\{\\text\{skip\}\}^\{\\text\{EDM\}\}\(t^\{\\prime\}\)\{\\bm\{x\}\}\_\{t^\{\\prime\}\}\\big\]\(26\)−\[At,τcoutEDM\(t\)Fθ,t\(𝒙t\)−At′,τ′coutEDM\(t′\)Fθ′,t′\(𝒙t′\)\]\.\\displaystyle\\quad\-\\big\[A\_\{t,\\tau\}c\_\{\\text\{out\}\}^\{\\text\{EDM\}\}\(t\)F\_\{\\theta,t\}\(\{\\bm\{x\}\}\_\{t\}\)\-A\_\{t^\{\\prime\},\\tau^\{\\prime\}\}c\_\{\\text\{out\}\}^\{\\text\{EDM\}\}\(t^\{\\prime\}\)F\_\{\\theta^\{\\prime\},t^\{\\prime\}\}\(\{\\bm\{x\}\}\_\{t^\{\\prime\}\}\)\\big\]\.where we useFθ,t\(𝒙t\)=Fθ\(cin\(t\)𝒙t;cnoise\(t\)\)F\_\{\\theta,t\}\(\{\\bm\{x\}\}\_\{t\}\)=F\_\{\\theta\}\(c\_\{\\text\{in\}\}\(t\)\{\\bm\{x\}\}\_\{t\};c\_\{\\text\{noise\}\}\(t\)\)for short\. Following EDM, which balances the per\-sample losses for the raw modelFθF\_\{\\theta\}, here we also equalize the loss scale forFθF\_\{\\theta\}andFθ′F\_\{\\theta^\{\\prime\}\}\. Minimizing the TD error effectively minimizes the normalized errors𝒆tEDM,𝒆t′EDM\{\\bm\{e\}\}\_\{t\}^\{\\mathrm\{EDM\}\},\{\\bm\{e\}\}\_\{t^\{\\prime\}\}^\{\\mathrm\{EDM\}\}
𝒆tEDM:=𝒙0−cskipEDM\(t\)𝒙tcoutEDM\(t\)−Fθ,t\(𝒙t\),𝒆t′EDM:=𝒙0−cskipEDM\(t′\)𝒙t′coutEDM\(t′\)−Fθ′,t′\(𝒙t′\)\.\\begin\{split\}\{\\bm\{e\}\}\_\{t\}^\{\\mathrm\{EDM\}\}&:=\\frac\{\{\\bm\{x\}\}\_\{0\}\-c\_\{\\mathrm\{skip\}\}^\{\\text\{EDM\}\}\(t\)\\,\{\\bm\{x\}\}\_\{t\}\}\{c\_\{\\mathrm\{out\}\}^\{\\text\{EDM\}\}\(t\)\}\-F\_\{\\theta,t\}\(\{\\bm\{x\}\}\_\{t\}\),\\\\ \{\\bm\{e\}\}\_\{t^\{\\prime\}\}^\{\\mathrm\{EDM\}\}&:=\\frac\{\{\\bm\{x\}\}\_\{0\}\-c\_\{\\mathrm\{skip\}\}^\{\\text\{EDM\}\}\(t^\{\\prime\}\)\\,\{\\bm\{x\}\}\_\{t^\{\\prime\}\}\}\{c\_\{\\mathrm\{out\}\}^\{\\text\{EDM\}\}\(t^\{\\prime\}\)\}\-F\_\{\\theta^\{\\prime\},t^\{\\prime\}\}\(\{\\bm\{x\}\}\_\{t^\{\\prime\}\}\)\.\\end\{split\}\(27\)rescaled byAt,τcoutEDM\(t\)A\_\{t,\\tau\}c\_\{\\text\{out\}\}^\{\\text\{EDM\}\}\(t\)\(resp\.At′,τ′coutEDM\(t′\)A\_\{t^\{\\prime\},\\tau^\{\\prime\}\}c\_\{\\text\{out\}\}^\{\\text\{EDM\}\}\(t^\{\\prime\}\)\)\. That is,𝜹t,t′=ℬ𝒆t,t′EDM\{\\bm\{\\delta\}\}\_\{t,t^\{\\prime\}\}=\{\\mathcal\{B\}\}\\ \{\\bm\{e\}\}\_\{t,t^\{\\prime\}\}^\{\\mathrm\{EDM\}\}where
ℬ\\displaystyle\{\\mathcal\{B\}\}:=\[At,τcoutEDM\(t\)I,−At′,τ′coutEDM\(t′\)I\]∈ℝd×2d,\\displaystyle:=\\big\[\\,A\_\{t,\\tau\}c\_\{\\mathrm\{out\}\}^\{\\text\{EDM\}\}\(t\)\\,I,\\ \-A\_\{t^\{\\prime\},\\tau^\{\\prime\}\}c\_\{\\mathrm\{out\}\}^\{\\text\{EDM\}\}\(t^\{\\prime\}\)\\,I\\,\\big\]\\in\\mathbb\{R\}^\{d\\times 2d\},𝒆t,t′EDM\\displaystyle\{\\bm\{e\}\}\_\{t,t^\{\\prime\}\}^\{\\mathrm\{EDM\}\}:=\[𝒆tEDM𝒆t′EDM\]∈ℝ2d\.\\displaystyle:=\\begin\{bmatrix\}\{\\bm\{e\}\}\_\{t\}^\{\\mathrm\{EDM\}\}\\\\ \{\\bm\{e\}\}\_\{t^\{\\prime\}\}^\{\\mathrm\{EDM\}\}\\end\{bmatrix\}\\in\\mathbb\{R\}^\{2d\}\.\(28\)Then we can see that
‖𝜹t,t′‖22\\displaystyle\\\|\{\\bm\{\\delta\}\}\_\{t,t^\{\\prime\}\}\\\|\_\{2\}^\{2\}=‖ℬ𝒆t,t′‖22≤‖ℬ‖22‖𝒆t,t′EDM‖22\\displaystyle=\\\|\{\\mathcal\{B\}\}\\,\{\\bm\{e\}\}\_\{t,t^\{\\prime\}\}\\\|\_\{2\}^\{2\}\\leq\\\|\{\\mathcal\{B\}\}\\\|\_\{2\}^\{2\}\\,\\\|\{\\bm\{e\}\}\_\{t,t^\{\\prime\}\}^\{\\mathrm\{EDM\}\}\\\|\_\{2\}^\{2\}\(29\)=\(At,τ2coutEDM\(t\)2\+At′,τ′2coutEDM\(t′\)2\)‖𝒆t,t′EDM‖22\.\\displaystyle=\\Big\(A\_\{t,\\tau\}^\{2\}c\_\{\\mathrm\{out\}\}^\{\\text\{EDM\}\}\(t\)^\{2\}\+A\_\{t^\{\\prime\},\\tau^\{\\prime\}\}^\{2\}c\_\{\\mathrm\{out\}\}^\{\\text\{EDM\}\}\(t^\{\\prime\}\)^\{2\}\\Big\)\\,\\\|\{\\bm\{e\}\}\_\{t,t^\{\\prime\}\}^\{\\mathrm\{EDM\}\}\\\|\_\{2\}^\{2\}\.Since‖𝒆t,t′EDM‖22\\\|\{\\bm\{e\}\}\_\{t,t^\{\\prime\}\}^\{\\mathrm\{EDM\}\}\\\|^\{2\}\_\{2\}is the normalized error w\.r\.t\. the raw modelsFθ,Fθ′F\_\{\\theta\},F\_\{\\theta^\{\\prime\}\}, it would be helpful to set a uniform scale that is not affected by the choice of time indices, which leads to the following pairwise weighting:
wTDEDM\(t,t′\)=1At,τ2coutEDM\(t\)2\+At′,τ′2coutEDM\(t′\)2\.\\displaystyle w\_\{\\mathrm\{TD\}\}^\{\\mathrm\{EDM\}\}\(t,t^\{\\prime\}\)\\ =\\ \\frac\{1\}\{A\_\{t,\\tau\}^\{2\}c\_\{\\mathrm\{out\}\}^\{\\text\{EDM\}\}\(t\)^\{2\}\+A\_\{t^\{\\prime\},\\tau^\{\\prime\}\}^\{2\}c\_\{\\mathrm\{out\}\}^\{\\text\{EDM\}\}\(t^\{\\prime\}\)^\{2\}\}\.\(30\)ThenwTD\(t,t′\)‖𝜹t,t′‖22≤‖𝒆t,t′EDM‖22w\_\{\\mathrm\{TD\}\}\(t,t^\{\\prime\}\)\\\|\{\\bm\{\\delta\}\}\_\{t,t^\{\\prime\}\}\\\|\_\{2\}^\{2\}\\leq\\\|\{\\bm\{e\}\}\_\{t,t^\{\\prime\}\}^\{\\mathrm\{EDM\}\}\\\|\_\{2\}^\{2\}, and the weighted TD objective is
ℒwTD\+EDMcont=𝔼𝒙0,t,t′\[wTDEDM\(t,t′\)‖𝜹t,t′‖22\]\.\\displaystyle\\ \\mathcal\{L\}^\{\\mathrm\{cont\}\}\_\{\\mathrm\{wTD\+EDM\}\}=\\mathbb\{E\}\_\{\{\\bm\{x\}\}\_\{0\},t,t^\{\\prime\}\}\\Big\[w\_\{\\mathrm\{TD\}\}^\{\\mathrm\{EDM\}\}\(t,t^\{\\prime\}\)\\big\\\|\{\\bm\{\\delta\}\}\_\{t,t^\{\\prime\}\}\\big\\\|\_\{2\}^\{2\}\\Big\]\.\(31\)
#### Weighting for CM parameterization\.
Similarly, we can apply CM’s modelfθf\_\{\\theta\}\([4](https://arxiv.org/html/2606.15048#S2.E4)\) to Equation \([25](https://arxiv.org/html/2606.15048#S3.E25)\)\. Instead of balancing losses for the raw modelFθF\_\{\\theta\}, consistency training \(CT\) applies a uniform weightingw\(t\)≡1w\(t\)\\equiv 1in \([5](https://arxiv.org/html/2606.15048#S2.E5)\)\(Song et al\.,[2023](https://arxiv.org/html/2606.15048#bib.bib24)\), equalizing the losses forfθf\_\{\\theta\}directly\. Following this practice, we focus on the errors forfθf\_\{\\theta\}
𝒆tCT:=𝒙0−fθ\(𝒙t;t\),𝒆t′CT:=𝒙0−fθ′\(𝒙t′;t′\)\.\{\\bm\{e\}\}\_\{t\}^\{\\mathrm\{CT\}\}:=\{\\bm\{x\}\}\_\{0\}\-f\_\{\\theta\}\(\{\\bm\{x\}\}\_\{t\};t\),\\quad\{\\bm\{e\}\}\_\{t^\{\\prime\}\}^\{\\mathrm\{CT\}\}:=\{\\bm\{x\}\}\_\{0\}\-f\_\{\\theta^\{\\prime\}\}\(\{\\bm\{x\}\}\_\{t^\{\\prime\}\};t^\{\\prime\}\)\.\(32\)Substituting these into Equation \([25](https://arxiv.org/html/2606.15048#S3.E25)\) and following the same derivation, we bound the TD error as
‖𝜹t,t′‖22≤\(At,τ2\+At′,τ′2\)‖𝒆t,t′CT‖22\.\\\|\{\\bm\{\\delta\}\}\_\{t,t^\{\\prime\}\}\\\|\_\{2\}^\{2\}\\leq\\Big\(A\_\{t,\\tau\}^\{2\}\+A\_\{t^\{\\prime\},\\tau^\{\\prime\}\}^\{2\}\\Big\)\\,\\\|\{\\bm\{e\}\}\_\{t,t^\{\\prime\}\}^\{\\mathrm\{CT\}\}\\\|\_\{2\}^\{2\}\.\(33\)Hence, we have the CT\-specific pairwise weight:
wTDCT\(t,t′\):=1At,τ2\+At′,τ′2\.w\_\{\\mathrm\{TD\}\}^\{\\mathrm\{CT\}\}\(t,t^\{\\prime\}\):=\\frac\{1\}\{A\_\{t,\\tau\}^\{2\}\+A\_\{t^\{\\prime\},\\tau^\{\\prime\}\}^\{2\}\}\.\(34\)The general weighted TD objective is then
ℒwTD\+CTcont=𝔼𝒙0,t,t′\[wTDCT\(t,t′\)‖𝜹t,t′‖22\]\.\\mathcal\{L\}^\{\\mathrm\{cont\}\}\_\{\\mathrm\{wTD\+CT\}\}=\\mathbb\{E\}\_\{\{\\bm\{x\}\}\_\{0\},t,t^\{\\prime\}\}\[w^\{\\mathrm\{CT\}\}\_\{\\mathrm\{TD\}\}\(t,t^\{\\prime\}\)\\\|\{\\bm\{\\delta\}\}\_\{t,t^\{\\prime\}\}\\\|\_\{2\}^\{2\}\]\.\(35\)
#### Aligning model with posterior\.
During training,ttis sampled within\[0,T\]\[0,T\], which can makeτ′\\tau^\{\\prime\}invalid given fixed spankkand strideΔ\\Delta\. Specifically, whenτ′=t−k−Δ\\tau^\{\\prime\}=t\-k\-\\Deltafalls outside the valid time window \(or equivalently, when the corresponding noise level goes below the minimum noise of the sampler\), the TD loss reduces to a mean matching objective, similar to the discrete\-time case\. Therefore, for a base continuous\-time algorithm \(e\.g\., EDM or CT\), we can incorporate our TD loss into the original loss, and optimize an objective based on the sampled time\-step or noise\-level\. For EDM, our objective is
ℒTD\+EDM=ℒwTD\+EDMcont\+λℒEDM,\\displaystyle\{\\mathcal\{L\}\}\_\{\\mathrm\{TD\+EDM\}\}\\ =\\ \{\\mathcal\{L\}\}^\{\\mathrm\{cont\}\}\_\{\\mathrm\{wTD\+EDM\}\}\\ \+\\ \\lambda\{\\mathcal\{L\}\}\_\{\\mathrm\{EDM\}\},\(36\)whenτ′\\tau^\{\\prime\}is valid, and\(1\+λ\)ℒEDM\(1\+\\lambda\)\{\\mathcal\{L\}\}\_\{\\mathrm\{EDM\}\}otherwise\. Similarly, for CT, our objective is
ℒTD\+CT=ℒwTD\+CTcont\+λℒCT\.\\displaystyle\{\\mathcal\{L\}\}\_\{\\mathrm\{TD\+CT\}\}\\ =\\ \{\\mathcal\{L\}\}^\{\\mathrm\{cont\}\}\_\{\\mathrm\{wTD\+CT\}\}\\ \+\\ \\lambda\{\\mathcal\{L\}\}\_\{\\mathrm\{CT\}\}\.\(37\)whenτ′\\tau^\{\\prime\}is valid, and\(1\+λ\)ℒCT\(1\+\\lambda\)\{\\mathcal\{L\}\}\_\{\\mathrm\{CT\}\}otherwise\.
Our method is flexible in that, given a base continuous\-time algorithm, we can incorporate our TD loss \([23](https://arxiv.org/html/2606.15048#S3.E23)\) into the base lossℒbase\{\\mathcal\{L\}\}\_\{\\text\{base\}\}depending on the original algorithm and its model parametrization for predicting𝒙0\{\\bm\{x\}\}\_\{0\}\. The procedure is summarized in[Algorithm1](https://arxiv.org/html/2606.15048#alg1)\.
#### Noise\-grid index mapping and TD pairing\.
In implementation, we parameterize TD time indices using the sampler’s noise grid of EDM\(Karras et al\.,[2022](https://arxiv.org/html/2606.15048#bib.bib13)\):
σ\(i\)=\(σmax1/ρ\+iN−1\(σmin1/ρ−σmax1/ρ\)\)ρ,\\displaystyle\\sigma\(i\)=\\Big\(\\sigma\_\{\\max\}^\{1/\\rho\}\+\\frac\{i\}\{N\-1\}\\big\(\\sigma\_\{\\min\}^\{1/\\rho\}\-\\sigma\_\{\\max\}^\{1/\\rho\}\\big\)\\Big\)^\{\\rho\},\(38\)wherei∈\{0,…,N−1\}i\\in\\\{0,\\ldots,N\-1\\\},NNis the grid size,ρ\\rhois the schedule exponent, and\(σmin,σmax\)\(\\sigma\_\{\\min\},\\sigma\_\{\\max\}\)are the upper and lower bounds of the sampler noise level\. Given a noise levelσ\\sigmawithin\[σmin,σmax\]\[\\sigma\_\{\\min\},\\sigma\_\{\\max\}\], we identify its corresponding \(possibly non\-integer\) indexiifrom Equation \([38](https://arxiv.org/html/2606.15048#S3.E38)\) and treat it as the TD time indextt\. We then set222We uset′=t\+kt^\{\\prime\}=t\+kinstead oft′=t−kt^\{\\prime\}=t\-kbecause thent′t^\{\\prime\}will correspond to a smaller noise level, thus closer to the real data as common in the ODE/SDE formulations \(t=0t=0\)\. In noise grid largeriicorresponding to smaller noise level\.t′=t\+kt^\{\\prime\}=t\+kand defineτ=t\+Δ\\tau=t\+\\Deltaandτ′=τ\+k\\tau^\{\\prime\}=\\tau\+kto construct the TD pair in Equation \([23](https://arxiv.org/html/2606.15048#S3.E23)\)\. We apply TD only whenσ\(τ′\)∈\[σmin,σmax\]\\sigma\(\\tau^\{\\prime\}\)\\in\[\\sigma\_\{\\min\},\\sigma\_\{\\max\}\]; otherwise we fall back to the base loss mentioned above\.
Algorithm 1TD training \(general recipe\)Input: Preconditioned denoiser𝒙θ,0\{\\bm\{x\}\}\_\{\\theta,0\}\(e\.g\.,DθD\_\{\\theta\}for EDM orfθf\_\{\\theta\}for CT\); target networkθ′\\theta^\{\\prime\}; noise gridσ\(i\)\\sigma\(i\); spankk; strideΔ\\Delta; mixing coefficientλ\\lambda; base lossℒbase\{\\mathcal\{L\}\}\_\{\\text\{base\}\}\(EDM or CT\); pairwise TD weightwTD\(⋅,⋅\)w\_\{\\mathrm\{TD\}\}\(\\cdot,\\cdot\)\.
1:Sample
𝒙0∼qdata\(𝒙0\)\{\\bm\{x\}\}\_\{0\}\\\!\\sim\\\!q\_\{\\text\{data\}\}\(\{\\bm\{x\}\}\_\{0\}\)and select a grid/noise index
ttaccording to the base method\.
2:Compute the base loss
ℒbase\{\\mathcal\{L\}\}\_\{\\text\{base\}\}\(EDM or CT\) using
𝒙θ,0\{\\bm\{x\}\}\_\{\\theta,0\}
3:if
t≤N−1−k−Δt\\leq N\-1\-k\-\\Deltathen
4:Compute the TD error
𝜹t,t′\{\\bm\{\\delta\}\}\_\{t,t^\{\\prime\}\}\.
5:
ℒwTD=wTD\(t,t′\)‖𝜹t,t′‖22\{\\mathcal\{L\}\}\_\{\\mathrm\{wTD\}\}=w\_\{\\mathrm\{TD\}\}\(t,t^\{\\prime\}\)\\,\\\|\{\\bm\{\\delta\}\}\_\{t,t^\{\\prime\}\}\\\|\_\{2\}^\{2\}\(EDM: \([30](https://arxiv.org/html/2606.15048#S3.E30)\); CT: \([34](https://arxiv.org/html/2606.15048#S3.E34)\)\)
6:
ℒ=ℒwTD\+λℒbase\{\\mathcal\{L\}\}=\{\\mathcal\{L\}\}\_\{\\mathrm\{wTD\}\}\+\\lambda\{\\mathcal\{L\}\}\_\{\\text\{base\}\}
7:else
8:
ℒ=\(1\+λ\)ℒbase\{\\mathcal\{L\}\}=\(1\+\\lambda\)\{\\mathcal\{L\}\}\_\{\\text\{base\}\}
9:endif
10:Perform gradient descent on
ℒ\{\\mathcal\{L\}\}and update
θ′\\theta^\{\\prime\}using exponential moving average\.
## 4Empirical Study
We evaluate our TD objective with \(i\) EDM training \(TD\+EDM\) and \(ii\) Consistency model training \(TD\+CT\)\.
Experimental Setup\.In all experiments, we utilize the probability\-flow ODE with the Heun integrator, where the NFE isNFE=2×steps−1\\textit\{NFE\}=2\\times\\text\{steps\}\-1\. We report thelast\-15% FID\-50k, defined as the average FID\-50k over the last 15% of evaluation checkpoints to ensure a stable performance measure\. For TD\+EDM, unless otherwise stated, we use the default TD setup selected by the EDM ablations:Δ=0\.25,k=1,λ=0\.5\\Delta=0\.25,\\,k=1,\\,\\lambda=0\.5, and the TD pairwise weighting proposed in[Section3\.3](https://arxiv.org/html/2606.15048#S3.SS3)\. See App\.[A](https://arxiv.org/html/2606.15048#A1)for implementation details\.
### 4\.1Performance and Applicability
TD\+EDM\.We first apply our TD training with EDM and evaluate performance across different inference steps under the standard EDM configuration\(Karras et al\.,[2022](https://arxiv.org/html/2606.15048#bib.bib13)\)\(App\.[A\.1](https://arxiv.org/html/2606.15048#A1.SS1)\)\.[Table2](https://arxiv.org/html/2606.15048#S4.T2)reports results on three benchmarks: class\-conditional CIFAR\-10 \(32×3232\{\\times\}32\), AFHQv2 \(64×6464\{\\times\}64\), and FFHQ \(64×6464\{\\times\}64\)\. Across these datasets, TD\+EDM matches or improves upon the EDM baseline in the moderate few\-step regime \(12–18 steps\)\. On FFHQ \(and on AFHQv2 for most step budgets\) specifically, our method allows the model to maintain higher fidelity than the baseline, which relies solely on local denoising targets\. These results suggest that by aligning the model drift and diffusion drift, TD training makes the model more robust to the larger discretization intervals inherent in few\-step solvers\.
Table 2:Cross\-dataset comparison under the same TD\+EDM setting\. FID\-50k↓\\downarrow\(last\-15% average\)\.TD\+CT\.We further evaluate TD training within the consistency training \(CT\) framework, following the same CT baseline setup and evaluation setting in consistency models\(Song et al\.,[2023](https://arxiv.org/html/2606.15048#bib.bib24)\)\(App\.[A\.2](https://arxiv.org/html/2606.15048#A1.SS2)\)\. For TD\+CT, we also sweep the TD\-specific hyperparameters, i\.e\.,λ\\lambdaandΔ\\Delta, over the same ranges as the EDM ablations\. The best setting isλ=0\.5\\lambda=0\.5,Δ=0\.5\\Delta=0\.5\(one\-step stride=1/2=1/2\), withk=1k=1, which we use for both AFHQv2 \(64×6464\{\\times\}64\) and FFHQ \(64×6464\{\\times\}64\)\. As shown in[Table3](https://arxiv.org/html/2606.15048#S4.T3), TD\+CT improves one\-step FID on both benchmarks\. On AFHQv2, FID improves from 12\.97 to 12\.87, and on FFHQ, FID improves from 19\.45 to 15\.93\.
The advantage of our TD objective is more evident by the learning curves and the multi\-step sampling profiles\. As illustrated in[Figure2](https://arxiv.org/html/2606.15048#S4.F2), TD\+CT consistently exhibits better convergence and lower FID across different sampling steps for FFHQ throughout the training process, while remaining competitive or better on AFHQ\.
Table 3:TD\+CT vs\. CT baseline under one\-step sampling \(steps=1=1, NFE=1=1\)\. FID\-50k↓\\downarrow\(last\-15% average over evaluation checkpoints\)\.\(a\)AFHQv2 \(64×6464\\times 64\)
\(b\)FFHQ \(64×6464\\times 64\)
Figure 2:Learning curves \(FID\-50k vs\. training steps\) for TD\+CT and CT baseline\. Each plot shows 1\-step and 2\-step sampling results\.To investigate whether the learned consistency generalizes across various sampling densities, we evaluate the FID differenceΔFID\(s\):=FIDTD\+CT\(s\)−FIDCT\(s\)\\Delta\\text\{FID\}\(s\):=\\text\{FID\}\_\{\\text\{TD\+CT\}\}\(s\)\-\\text\{FID\}\_\{\\text\{CT\}\}\(s\)for steps∈\{1,2,3,4,6,8\}s\\in\\\{1,2,3,4,6,8\\\}\. As illustrated in[Figure3](https://arxiv.org/html/2606.15048#S4.F3), although minor fluctuations exist whereΔFID\\Delta\\text\{FID\}occasionally approaches or crosses the zero line \(notably for 4 and 8 steps during early or middle training phases\), the FID difference remains predominantly negative throughout the majority of the training process across all evaluated step counts\. This demonstrates that TD\+CT generally maintains a better performance profile over the baseline across different inference budgets\.
Figure 3:Relative improvement of TD\+CT over CT baseline on FFHQ \(64×6464\\times 64\) across different inference steps\. Negative values indicate that TD\+CT achieves lower FID than the baseline\.Computational Cost\.The additional computational cost of our method mainly arises from maintaining and updating the target networkθ′\\theta^\{\\prime\}\. For TD\+EDM, on class\-conditional CIFAR\-10 with the 55M\-parameter UNet trained for 200M images, the baseline training takes 50h 33m 30s \(45\.7 s/tick\), while training with TD takes 71h 36m 48s \(64\.7 s/tick\) on the same hardware\. CPU memory increases from 2\.29 GB to 2\.74 GB \(≈\\approx20%\), and GPU memory from 16\.58 GB to 17\.51 GB \(≈\\approx6%\)\. In contrast, for TD\+CT, the overhead is minimal because CT itself already utilizes a target network\. For 5,000 training steps, TD\+CT takes 95min \(AFHQv2\) and 92min \(FFHQ\), compared to 80min \(AFHQv2\) and 78min \(FFHQ\) for the CT baseline, respectively\. These indicate that our method can effectively improve generation quality without incurring significant overhead\.
### 4\.2Ablation Study
Given compute constraints, all ablations use CIFAR\-10 with a smaller UNet and report last\-15% FID\-50k \(mean±\\pmstd over 3 seeds; App\.[A\.3](https://arxiv.org/html/2606.15048#A1.SS3)\) using TD\+EDM\. Unless a factor is being swept, the default hyperparameters areΔ=0\.25,k=1,λ=0\.5\\Delta\{=\}0\.25,\\ k\{=\}1,\\ \\lambda\{=\}0\.5, with the weighted TD in[Equation30](https://arxiv.org/html/2606.15048#S3.E30)\.
Effect of pairwise TD weighting\.We first validate the necessity of our sample\-wise \(pairwise\) weightingwTDw\_\{\\mathrm\{TD\}\}in[Equation30](https://arxiv.org/html/2606.15048#S3.E30)by comparing it against a constant weight \(unweighted TD\)\.[Table4](https://arxiv.org/html/2606.15048#S4.T4)shows that weighting yields a clear improvement across step budgets\.
Table 4:Effects of pairwise weighting and sampling steps \(CIFAR\-10, small UNet\)\. FID\-50k↓\\downarrow\(last\-15% avg\.; mean±\\pmstd over 3 seeds\)\.Sensitivity to the regularization weightλ\\lambda\.To study the balance between TD and EDM losses, we sweepλ\\lambdawhile keepingΔ=0\.25\\Delta\{=\}0\.25,k=1k\{=\}1, and weighted TD fixed\.[Table5](https://arxiv.org/html/2606.15048#S4.T5)shows that the best performance is achieved with relatively smallλ\\lambda, and performance is stable across a low\-λ\\lambdarange\. Moreover, the improvement over EDM is consistent\.
Sensitivity to the strideΔ\\Delta\.We next investigate sensitivity to the strideΔ\\Deltaintroduced by our TD formulation\.[Table6](https://arxiv.org/html/2606.15048#S4.T6)shows that, under a fixed step budget, varyingΔ\\Deltatypically leads to only minor differences, suggesting robustness to this hyperparameter\.
Table 5:Ablation on weightλ\\lambda\. We varyλ∈\{0\.01,0\.5,1\.0,2\.0\}\\lambda\\in\\\{0\.01,0\.5,1\.0,2\.0\\\}underΔ=0\.25\\Delta\{=\}0\.25,k=1k\{=\}1\. CIFAR\-10 FID\-50k↓\\downarrow\(last\-15% avg\.; mean±\\pmstd over 3 seeds\)\.Table 6:Ablation on strideΔ∈\{1/2,1/3,1/4,1/5\}\\Delta\\\!\\in\\\!\\\{1/2,1/3,1/4,1/5\\\}\(weighted\)\. CIFAR\-10 FID\-50k↓\\downarrow\(last\-15% avg\.; mean±\\pmstd over 3 seeds\)\.
## 5Discussions
We have introduced a temporal difference \(TD\) learning framework for diffusion models that reformulates denoising across the time axis as a policy evaluation problem\. By casting the denoising trajectory as a Markov reward process \(MRP\), we derive a novel TD objective that explicitly penalizes inconsistencies in the model’s multi\-step progress, thereby enforcing cross\-time agreement\. A key theoretical and practical element of our work is the derivation of a principled, sample\-based loss reweighting scheme,wTDw\_\{\\mathrm\{TD\}\}, which ensures stable optimization by balancing the loss scales across heterogeneous time pairs\.
Insights from Empirical Results\.Our evaluations with TD\+EDM and TD\+CT provide empirical evidence for the benefits of temporal consistency training:
- •FID Improvements:Across multiple standard benchmarks, including CIFAR\-10, AFHQv2, and FFHQ, our TD objective generally matches or improves sample quality\. For TD\+EDM, we observe clear FID reductions in the 12–18 step range on most settings, demonstrating that enforcing cross\-time agreement can suppress discretization error accumulation in few\-step regimes\. For TD\+CT, we also observe improved one\-step performance on both AFHQv2 and FFHQ, with a particularly large gain on FFHQ \(Table[3](https://arxiv.org/html/2606.15048#S4.T3)\)\.
- •Robustness Across Sampling Budgets:A significant finding is the robustness of the TD\+CT model across a wide range of inference steps\. As shown in[Figure3](https://arxiv.org/html/2606.15048#S4.F3), the relative improvement \(Δ\\DeltaFID\) remains predominantly negative across all evaluated step counts \(1–8 steps\), indicating a broadly improved performance profile that is not tied to a single inference budget\.
- •Reliability of Hyperparameter Choices:Our ablation studies demonstrate that the TD framework is robust across various configurations\. While the pairwise weightingwTDw\_\{\\mathrm\{TD\}\}is essential for achieving substantial and stable gains, the performance remains consistently superior over the baseline across a broad range of regularization weightsλ\\lambdaand temporal stridesΔ\\Delta\. This confirms the reliability of our proposed default training scheme\.
Limitations\.While effective, our approach has several limitations\. First, compared with EDM\(Karras et al\.,[2022](https://arxiv.org/html/2606.15048#bib.bib13)\), the TD objective may introduce additional constant\-factor∼1\.5×\\sim\\\!1\.5\\timescomputational and memory overhead when a separate TD target network is maintained and updated\. However, this overhead is largely amortized in training pipelines that already rely on a target/teacher network \(e\.g\., consistency training with EMA teachers\), where our method can reuse the existing target network with little marginal cost\. Second, the current empirical scope is primarily focused on the generation at resolutions up to64×6464\\times 64using a probability\-flow ODE solver\. While this setting is consistent with prior algorithmic studies\(Ho et al\.,[2020](https://arxiv.org/html/2606.15048#bib.bib10); Karras et al\.,[2022](https://arxiv.org/html/2606.15048#bib.bib13)\), validation on higher\-resolution models, broader architectures, and specialized one\-/few\-step generators remains an important direction\.
Future work\.These findings motivate several promising research directions\. The computational cost could be mitigated by amortizing target\-network updates or employing lightweight adapter modules\. The framework also invites adaptive training curricula, such as progressively widening the step spankkor annealing the mixing coefficientλ\\lambdaduring training\(Karras et al\.,[2022](https://arxiv.org/html/2606.15048#bib.bib13)\)\. Related approaches, such as SDPO, have demonstrated the value of dense stepwise rewards\(Zhang et al\.,[2024](https://arxiv.org/html/2606.15048#bib.bib31)\), and noise\-correlation methods like ARTDiff\(Lu et al\.,[2024](https://arxiv.org/html/2606.15048#bib.bib16)\)could be combined with TD learning to further enhance temporal stability\. Finally, given its formulation as a drop\-in objective, we anticipate strong synergies with distillation techniques\(Salimans & Ho,[2022](https://arxiv.org/html/2606.15048#bib.bib21)\)and solver\-aware training policies\(Lu et al\.,[2022](https://arxiv.org/html/2606.15048#bib.bib15); Zhao et al\.,[2023](https://arxiv.org/html/2606.15048#bib.bib32)\)\.
## Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning\. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here\.
## Acknowledgement
Qizhen Ying and Yangchen Pan acknowledge the support from the Engineering and Physical Sciences Research Council \(EPSRC\) New Investigator Award under grant reference UKRI2775\. Junfeng Wen acknowledges the support of NSERC, RGPIN\-2024\-05357\. Qizhen Ying acknowledges the use of resources provided by the Isambard\-AI National AI Research Resource \(AIRR\)\. Isambard\-AI is operated by the University of Bristol and is funded by the UK Government’s Department for Science, Innovation and Technology \(DSIT\) via UK Research and Innovation; and the Science and Technology Facilities Council \[ST/AIRR/I\-A\-I/1023\]\(McIntosh\-Smith et al\.,[2024](https://arxiv.org/html/2606.15048#bib.bib18)\)\.
## References
- Bertsekas & Tsitsiklis \(1996\)Bertsekas, D\. and Tsitsiklis, J\.*Neuro\-Dynamic Programming*\.Athena Scientific, 1996\.
- Black et al\. \(2023\)Black, K\., Janner, M\., Du, Y\., Kostrikov, I\., and Levine, S\.Training diffusion models with reinforcement learning\.*arXiv preprint arXiv:2305\.13301*, 2023\.
- Chen et al\. \(2024\)Chen, J\., Byun, J\.\-S\., Elsner, M\., and Perrault, A\.Dlpo: Diffusion model loss\-guided reinforcement learning for fine\-tuning text\-to\-speech diffusion models\.*arXiv preprint arXiv:2405\.14632*, 2024\.
- Domingo i Enrich et al\. \(2025\)Domingo i Enrich, C\., Drozdzal, M\., Karrer, B\., and Chen, R\. T\.Adjoint matching: Fine\-tuning flow and diffusion generative models with memoryless stochastic optimal control\.In*International Conference on Learning Representations*, volume 2025, pp\. 53791–53846, 2025\.
- Fan et al\. \(2023\)Fan, Y\., Watkins, O\., Du, Y\., Liu, H\., Ryu, M\., Boutilier, C\., Abbeel, P\., Ghavamzadeh, M\., Lee, K\., and Lee, K\.Dpok: Reinforcement learning for fine\-tuning text\-to\-image diffusion models\.*Advances in Neural Information Processing Systems*, 2023\.
- Farebrother et al\. \(2025\)Farebrother, J\., Pirotta, M\., Tirinzoni, A\., Munos, R\., Lazaric, A\., and Touati, A\.Temporal difference flows\.*International Conference on Machine Learning*, 2025\.
- Frans et al\. \(2025\)Frans, K\., Hafner, D\., Levine, S\., and Abbeel, P\.One step diffusion via shortcut models\.In*International Conference on Learning Representations*, volume 2025, pp\. 34668–34684, 2025\.
- Geng et al\. \(2026\)Geng, Z\., Deng, M\., Bai, X\., Kolter, Z\., and He, K\.Mean flows for one\-step generative modeling\.*Advances in Neural Information Processing Systems*, 38:75460–75482, 2026\.
- Gupta et al\. \(2025\)Gupta, S\., Ahuja, C\., Lin, T\.\-Y\., Roy, S\. D\., Oosterhuis, H\., de Rijke, M\., and Shukla, S\. N\.A simple and effective reinforcement learning method for text\-to\-image diffusion fine\-tuning\.*arXiv preprint arXiv:2503\.00897*, 2025\.
- Ho et al\. \(2020\)Ho, J\., Jain, A\., and Abbeel, P\.Denoising diffusion probabilistic models\.*Advances in neural information processing systems*, 2020\.
- Janner et al\. \(2020\)Janner, M\., Mordatch, I\., and Levine, S\.Gamma\-models: Generative temporal difference learning for infinite\-horizon prediction\.*Advances in neural information processing systems*, 2020\.
- Janner et al\. \(2022\)Janner, M\., Du, Y\., Tenenbaum, J\. B\., and Levine, S\.Planning with diffusion for flexible behavior synthesis\.*International Conference on Machine Learning*, 2022\.
- Karras et al\. \(2022\)Karras, T\., Aittala, M\., Aila, T\., and Laine, S\.Elucidating the design space of diffusion\-based generative models\.*Advances in Neural Information Processing Systems*, 2022\.
- Li et al\. \(2023\)Li, W\., Wang, X\., Jin, B\., and Zha, H\.Hierarchical diffusion for offline decision making\.*International Conference on Machine Learning*, 2023\.
- Lu et al\. \(2022\)Lu, C\., Zhou, Y\., Bao, F\., Chen, J\., Li, C\., and Zhu, J\.Dpm\-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps\.*Advances in Neural Information Processing Systems*, 2022\.
- Lu et al\. \(2024\)Lu, K\., Yuxi, C\., Li, L\., Qin, D\., and Li, G\.Improve temporal consistency in diffusion models through noise correlations\.*Openreview*, 2024\.
- Luo \(2022\)Luo, C\.Understanding diffusion models: A unified perspective\.*arXiv preprint arXiv:2208\.11970*, 2022\.
- McIntosh\-Smith et al\. \(2024\)McIntosh\-Smith, S\., Alam, S\. R\., and Woods, C\.Isambard\-ai: a leadership class supercomputer optimised specifically for artificial intelligence\.*arXiv preprint arXiv:2410\.11199*, 2024\.
- Nichol & Dhariwal \(2021\)Nichol, A\. Q\. and Dhariwal, P\.Improved denoising diffusion probabilistic models\.*International Conference on Machine Learning*, 2021\.
- Pan et al\. \(2025\)Pan, Y\., Wen, J\., Xiao, C\., and Torr, P\. H\.An mrp formulation for supervised learning: Generalized temporal difference learning models\.*Journal of Artificial Intelligence Research*, 2025\.
- Salimans & Ho \(2022\)Salimans, T\. and Ho, J\.Progressive distillation for fast sampling of diffusion models\.*arXiv preprint arXiv:2202\.00512*, 2022\.
- Song et al\. \(2021a\)Song, J\., Meng, C\., and Ermon, S\.Denoising diffusion implicit models\.*International Conference on Learning Representations*, 2021a\.
- Song et al\. \(2021b\)Song, Y\., Sohl\-Dickstein, J\., Kingma, D\. P\., Kumar, A\., Ermon, S\., and Poole, B\.Score\-based generative modeling through stochastic differential equations\.*International Conference on Learning Representations*, 2021b\.
- Song et al\. \(2023\)Song, Y\., Dhariwal, P\., Chen, M\., and Sutskever, I\.Consistency models\.*International Conference on Machine Learning*, 2023\.
- Sutton & Barto \(2018\)Sutton, R\. and Barto, A\.*Reinforcement Learning, second edition: An Introduction*\.MIT Press, 2018\.
- Szepesvári \(2022\)Szepesvári, C\.*Algorithms for reinforcement learning*\.Springer nature, 2022\.
- Tsitsiklis & Van Roy \(1996\)Tsitsiklis, J\. and Van Roy, B\.Analysis of temporal\-diffference learning with function approximation\.*Advances in neural information processing systems*, 9, 1996\.
- Wang et al\. \(2023\)Wang, Z\., Hunt, J\. J\., and Zhou, M\.Diffusion policies as an expressive policy class for offline reinforcement learning\.*International Conference on Learning Representations*, 2023\.
- Yuan et al\. \(2024\)Yuan, H\., Chen, Z\., Ji, K\., and Gu, Q\.Self\-play fine\-tuning of diffusion models for text\-to\-image generation\.*Advances in Neural Information Processing Systems*, 2024\.
- Zekri & Boullé \(2025\)Zekri, O\. and Boullé, N\.Fine\-tuning discrete diffusion models with policy gradient methods\.*Advances in Neural Information Processing Systems*, 2025\.
- Zhang et al\. \(2024\)Zhang, Z\., Shen, L\., Zhang, S\., Ye, D\., Luo, Y\., Shi, M\., Du, B\., and Tao, D\.Aligning few\-step diffusion models with dense reward difference learning\.*arXiv preprint arXiv:2411\.11727*, 2024\.
- Zhao et al\. \(2023\)Zhao, W\., Bai, L\., Rao, Y\., Zhou, J\., and Lu, J\.Unipc: A unified predictor\-corrector framework for fast sampling of diffusion models\.*Advances in Neural Information Processing Systems*, 2023\.
- Zheng et al\. \(2025\)Zheng, K\., Chen, H\., Ye, H\., Wang, H\., Zhang, Q\., Jiang, K\., Su, H\., Ermon, S\., Zhu, J\., and Liu, M\.\-Y\.Diffusionnft: Online diffusion reinforcement with forward process\.*arXiv preprint arXiv:2509\.16117*, 2025\.
## Appendix AImplementation Details
This section provides full implementation settings for all experiments in Sec\.[4](https://arxiv.org/html/2606.15048#S4)\. We follow the same order as the main text: TD\+EDM \(Table[2](https://arxiv.org/html/2606.15048#S4.T2)\), TD\+CT \(Table[3](https://arxiv.org/html/2606.15048#S4.T3), Fig\.[2](https://arxiv.org/html/2606.15048#S4.F2)–[3](https://arxiv.org/html/2606.15048#S4.F3)\), and the ablation study \(Tables[4](https://arxiv.org/html/2606.15048#S4.T4)–[6](https://arxiv.org/html/2606.15048#S4.T6)\)\.
### A\.1TD\+EDM: Full Settings \(Table[2](https://arxiv.org/html/2606.15048#S4.T2)\)
Common EDM setup\.Unless stated otherwise, TD\+EDM follows the standard EDM configuration\(Karras et al\.,[2022](https://arxiv.org/html/2606.15048#bib.bib13)\)\(training pipeline, EDM preconditioning, noise sampling distribution, and the probability\-flow ODE sampler with Heun’s method\)\. We summarize the shared EDM constants in Table[7](https://arxiv.org/html/2606.15048#A1.T7)\.
Datasets & preprocessing\.Official train/test split; images normalized to\[−1,1\]\[\-1,1\]; AugmentPipe withp=0\.12p=0\.12\(xflip, yflip, scale, rotate, aniso, translate\), following the official PyTorch implementation ofKarras et al\. \([2022](https://arxiv.org/html/2606.15048#bib.bib13)\)\. In Table[2](https://arxiv.org/html/2606.15048#S4.T2), CIFAR\-10 experiments areclass\-conditional, while FFHQ and AFHQv2 areunconditional\.
Model architecture \(EDM default setting\)\.We use the EDM default SongUNet \(DDPM\+\+\) backbone with EDM preconditioning and𝒙^0\\hat\{\{\\bm\{x\}\}\}\_\{0\}parameterization\.
TD objective settings \(ours\)\.Default TD hyperparameters: strideΔ=0\.25\\Delta=0\.25\(index unit\), spank=1k=1, mixingλ=0\.5\\lambda=0\.5, and the TD pairwise weighting proposed in Sec\.[3\.3](https://arxiv.org/html/2606.15048#S3.SS3)\. Coupling: Markov increments\. EMA target network decayτ=0\.999\\tau=0\.999, updated every step\. Boundary mask nearσmin\\sigma\_\{\\min\}with valid index rangei∈\[0,N−1−k−Δ\]i\\in\[0,\\,N\-1\-k\-\\Delta\]\. Pair sampling: uniform over valid indices\.
Optimization and schedule \(EDM default setting\)\.Adam, learning rate2×10−42\{\\times\}10^\{\-4\}, betas\(0\.9,0\.999\)\(0\.9,0\.999\); no weight decay, no grad clip; constant LR, no warmup; batch128128per GPU, global512512, no accumulation; precision fp32 \(static scale11\); training length200200M images \(∼\\sim4000 ticks\) for all datasets in Table[2](https://arxiv.org/html/2606.15048#S4.T2)\.
Sampler & inference \(EDM default setting\)\.Probability\-flow ODE; Heun \(order 2\)\. We evaluate inference steps among\{9,12,15,18\}\\\{9,12,15,18\\\}, and reportNFE=2×steps−1\\textit\{NFE\}=2\\times\\text\{steps\}\-1\(i\.e\., NFE∈\{17,23,29,35\}\\in\\\{17,23,29,35\\\}\)\. We use the sameρ\\rho\-power grid\(σmin,σmax,ρ\)=\(2×10−3,80,7\)\(\\sigma\_\{\\min\},\\sigma\_\{\\max\},\\rho\)=\(2\{\\times\}10^\{\-3\},80,7\)\.
Evaluation protocol\.FID\-50k with deterministic evaluation seeds0–4999949999; we report thelast\-15%average over evaluation checkpoints \(as described in Sec\.[4\.1](https://arxiv.org/html/2606.15048#S4.SS1)\)\.
Computational resources:4×4\{\\times\}H100 80GB
Table 7:TD\+EDM: shared EDM constants\.Table 8:TD\+EDM: TD regularizer settings\.
### A\.2TD\+CT Experimental Details \(Table[3](https://arxiv.org/html/2606.15048#S4.T3)\)
Experiment setup\.TD\+CT follows the consistency training \(CT\) setup ofSong et al\. \([2023](https://arxiv.org/html/2606.15048#bib.bib24)\)for each dataset \(architecture, parameterization, schedules, solver, and evaluation protocol\), unless explicitly stated otherwise\.
Training and hyperparameters\.For both datasets, we use theadaptiveNNschedule inSong et al\. \([2023](https://arxiv.org/html/2606.15048#bib.bib24)\), where the number of discretization stepsN\(⋅\)N\(\\cdot\)increases during training\. Models are trained for a total of300,000 iterations\. We sweep the TD\-specific hyperparametersλ\\lambdaandΔ\\Deltafor TD\+CT, and selectΔ=0\.5\\Delta=0\.5,k=1k=1, andλ=0\.5\\lambda=0\.5based on the last\-15% one\-step FID\-50k, together with the CT\-specific pairwise TD weight and the same boundary masking rule as in Sec\.[3\.3](https://arxiv.org/html/2606.15048#S3.SS3)\.
Evaluation\.We evaluate every 5,000 training iterations\. We report one\-step FID\-50k \(steps=1=1, NFE=1=1\) using a single Heun update and plot the resulting learning curves in Fig\.[2](https://arxiv.org/html/2606.15048#S4.F2)\. For multi\-step analysis \(Fig\.[3](https://arxiv.org/html/2606.15048#S4.F3)\), we also evaluates∈\{1,2,3,4,6,8\}s\\in\\\{1,2,3,4,6,8\\\}at each evaluation point and reportΔFID\(s\)\\Delta\\text\{FID\}\(s\)as defined in the main text\.
Computational resources:4×4\{\\times\}H100 80GB
### A\.3Small\-model Ablations \(Tables[4](https://arxiv.org/html/2606.15048#S4.T4)–[6](https://arxiv.org/html/2606.15048#S4.T6)\)
Scope\.All ablations use CIFAR\-10 with a smaller UNet and report last\-15% FID\-50k \(mean±\\pmstd over 3 seeds\), as described in Sec\.[4\.2](https://arxiv.org/html/2606.15048#S4.SS2)\. Unless a factor is being swept, we anchor hyperparameters atΔ=0\.25\\Delta=0\.25,k=1k=1,λ=0\.5\\lambda=0\.5with the*weighted*TD in Eq\.[30](https://arxiv.org/html/2606.15048#S3.E30)\.
Datasets & preprocessing\.CIFAR\-10; official split; resolution32×3232\{\\times\}32for train/eval; normalized to\[−1,1\]\[\-1,1\]; AugmentPipe withp=0\.12p=0\.12\(xflip, yflip, scale, rotate, aniso, translate\)\.Notice:in contrast to Table[2](https://arxiv.org/html/2606.15048#S4.T2)\(conditional CIFAR\-10 for TD\+EDM\), all ablations on CIFAR\-10 are unconditional\.
Model architecture\.SongUNet \(DDPM\+\+\),∼\\sim13M parameters; base channelsch=64ch\{=\}64, multipliers\[1,2,2\]\[1,2,2\];22resblocks/level; attention at1616with11head, head\_dim6464; GroupNorm \+ SiLU; dropout0\.100\.10; time/noise embed dims256256\(positional\); output parameterization𝒙^0\\hat\{\{\\bm\{x\}\}\}\_\{0\}\(EDM\)\.
EDM constants & TD settings\.We use the same EDM constants as in Table[7](https://arxiv.org/html/2606.15048#A1.T7)\. Unless swept, we use the same TD settings as in Table[8](https://arxiv.org/html/2606.15048#A1.T8)\. The ablated ranges ofΔ\\Delta,λ\\lambda, and weighting are summarized in Table[9](https://arxiv.org/html/2606.15048#A1.T9)\.
Optimization and training length\.Follow the EDM default optimizer setup \(Table[2](https://arxiv.org/html/2606.15048#S4.T2)\) with Adam \(LR2×10−42\{\\times\}10^\{\-4\}, betas\(0\.9,0\.999\)\(0\.9,0\.999\)\), no weight decay, no grad clip, constant LR, no warmup\. We train each configuration for120120M images \(∼\\sim2400 ticks\)\.
Sampler & evaluation\.Probability\-flow ODE \+ Heun \(order 2\)\. We evaluate steps among\{12,15,18\}\\\{12,15,18\\\}and reportNFE=2×steps−1\\textit\{NFE\}=2\\times\\text\{steps\}\-1\(i\.e\., NFE∈\{23,29,35\}\\in\\\{23,29,35\\\}\)\. FID\-50k with deterministic evaluation seeds0–4999949999; last\-15% average\.
Computational resources:single NVIDIA RTX4090\.
Table 9:Ablation ranges for TD hyperparameters on CIFAR\-10 \(small UNet\)\.
## Appendix BForward Processes and Two\-Time Means
This appendix consolidates the forward/noising processes considered in the paper and rewrites their two\-time posterior means in the unified linear form
𝝁τtrue\(𝒙t,𝒙0\)=At,τ𝒙0\+κt,τ𝒙t,τ<t,\{\\bm\{\\mu\}\}\_\{\\tau\}^\{\\text\{true\}\}\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)=A\_\{t,\\tau\}\\,\{\\bm\{x\}\}\_\{0\}\+\\kappa\_\{t,\\tau\}\\,\{\\bm\{x\}\}\_\{t\},\\qquad\\tau<t,\(39\)where𝒙0∈ℝd\{\\bm\{x\}\}\_\{0\}\\in\\mathbb\{R\}^\{d\}is the clean datum and𝒙t\{\\bm\{x\}\}\_\{t\}is a noisy observation at leveltton the*same*time/noise axis asτ\\tau\. For discrete\-time models \(DDPM/DDIM\),t,τ∈\{0,…,T\}t,\\tau\\\!\\in\\\!\\\{0,\\dots,T\\\}and we typically takeτ=t−1\\tau=t\{\-\}1on the native grid\. For continuous\-time models \(VP/VE/EDM\),t,τ∈\[0,T\]t,\\tau\\\!\\in\\\!\[0,T\]\(or a monotone reparameterization such asσ\\sigma\) chosen by the sampler\.
#### Notation\.
For DDPM/DDIM, setαt=1−βt\\alpha\_\{t\}=1\-\\beta\_\{t\}andα¯t=∏s=1tαs\\bar\{\\alpha\}\_\{t\}=\\prod\_\{s=1\}^\{t\}\\alpha\_\{s\}\. For VP\-SDE, letα\(t\)=exp\(−12∫0tβ\(s\)𝑑s\)\\alpha\(t\)=\\exp\\\!\\big\(\-\\tfrac\{1\}\{2\}\\\!\\int\_\{0\}^\{t\}\\beta\(s\)\\,ds\\big\)andσ\(t\)=1−α\(t\)2\\sigma\(t\)=\\sqrt\{1\-\\alpha\(t\)^\{2\}\}\. For VE/EDM,σ\(⋅\)\\sigma\(\\cdot\)denotes the noise scale with𝒙t=𝒙0\+σ\(t\)ϵ\{\\bm\{x\}\}\_\{t\}=\{\\bm\{x\}\}\_\{0\}\+\\sigma\(t\)\\,\{\\bm\{\\epsilon\}\},ϵ∼𝒩\(𝟎,I\)\{\\bm\{\\epsilon\}\}\\sim\\mathcal\{N\}\(\{\\bm\{0\}\},I\)\. All random vectors are inℝd\\mathbb\{R\}^\{d\}, andIIis the identity matrix\.
### B\.1DDPM \(discrete time\)
#### Forward one\-step transition\.
The DDPM forward transition and the marginal w\.r\.t\.𝒙0\{\\bm\{x\}\}\_\{0\}are \(Luo \([2022](https://arxiv.org/html/2606.15048#bib.bib17), Eq\.\(31\); Eq\.\(69\)\-\(70\)\)\):
q\(𝒙t∣𝒙t−1\)=𝒩\(αt𝒙t−1,\(1−αt\)I\),\(DDPM\-Forward\)q\(\{\\bm\{x\}\}\_\{t\}\\mid\{\\bm\{x\}\}\_\{t\-1\}\)=\\mathcal\{N\}\\\!\\big\(\\sqrt\{\\alpha\_\{t\}\}\\,\{\\bm\{x\}\}\_\{t\-1\},\\ \(1\-\\alpha\_\{t\}\)I\\big\),\\quad\\text\{\(DDPM\-Forward\)\}\(40\)q\(𝒙t∣𝒙0\)=𝒩\(α¯t𝒙0,\(1−α¯t\)I\),\(DDPM\-Marginal\)q\(\{\\bm\{x\}\}\_\{t\}\\mid\{\\bm\{x\}\}\_\{0\}\)=\\mathcal\{N\}\\\!\\big\(\\sqrt\{\\bar\{\\alpha\}\_\{t\}\}\\,\{\\bm\{x\}\}\_\{0\},\\ \(1\-\\bar\{\\alpha\}\_\{t\}\)I\\big\),\\quad\\text\{\(DDPM\-Marginal\)\}\(41\)with the equivalent reparameterization
𝒙t=α¯t𝒙0\+1−α¯tϵ,ϵ∼𝒩\(𝟎,I\)\.\(DDPM\-Reparam\)\{\\bm\{x\}\}\_\{t\}=\\sqrt\{\\bar\{\\alpha\}\_\{t\}\}\\,\{\\bm\{x\}\}\_\{0\}\+\\sqrt\{1\-\\bar\{\\alpha\}\_\{t\}\}\\,\{\\bm\{\\epsilon\}\},\\quad\{\\bm\{\\epsilon\}\}\\sim\\mathcal\{N\}\(\{\\bm\{0\}\},I\)\.\\quad\\text\{\(DDPM\-Reparam\)\}\(42\)
#### Posterior towardt−1t\{\-\}1\(mean and variance\)\.
By linear\-Gaussian conditioning,
q\(𝒙t−1∣𝒙t,𝒙0\)\\displaystyle q\(\{\\bm\{x\}\}\_\{t\-1\}\\mid\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)=𝒩\(𝝁t−1true\(𝒙t,𝒙0\),β~tI\),\(DDPM\-Posterior\)\\displaystyle=\\mathcal\{N\}\\\!\\big\(\{\\bm\{\\mu\}\}^\{\\text\{true\}\}\_\{t\-1\}\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\),\\ \\tilde\{\\beta\}\_\{t\}\\,I\\big\),\\quad\\text\{\(DDPM\-Posterior\)\}\(43\)where𝝁t−1true\(𝒙t,𝒙0\)\\displaystyle\\text\{where \}\{\\bm\{\\mu\}\}^\{\\text\{true\}\}\_\{t\-1\}\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)=α¯t−1\(1−αt\)1−α¯t𝒙0\+αt\(1−α¯t−1\)1−α¯t𝒙t,\(DDPM\-PosteriorMean\)\\displaystyle=\\frac\{\\sqrt\{\\bar\{\\alpha\}\_\{t\-1\}\}\\,\(1\-\\alpha\_\{t\}\)\}\{1\-\\bar\{\\alpha\}\_\{t\}\}\\,\{\\bm\{x\}\}\_\{0\}\+\\frac\{\\sqrt\{\\alpha\_\{t\}\}\\,\(1\-\\bar\{\\alpha\}\_\{t\-1\}\)\}\{1\-\\bar\{\\alpha\}\_\{t\}\}\\,\{\\bm\{x\}\}\_\{t\},\\quad\\text\{\(DDPM\-PosteriorMean\)\}\(44\)β~t\\displaystyle\\tilde\{\\beta\}\_\{t\}=\(1−αt\)\(1−α¯t−1\)1−α¯t\.\(DDPM\-PosteriorVar\)\\displaystyle=\\frac\{\(1\-\\alpha\_\{t\}\)\(1\-\\bar\{\\alpha\}\_\{t\-1\}\)\}\{1\-\\bar\{\\alpha\}\_\{t\}\}\.\\quad\\text\{\(DDPM\-PosteriorVar\)\}\(45\)Equations \([44](https://arxiv.org/html/2606.15048#A2.E44)\)\-\([45](https://arxiv.org/html/2606.15048#A2.E45)\) fromLuo \([2022](https://arxiv.org/html/2606.15048#bib.bib17), Eq\.\(84\)\-\(85\), p\.12\)\.
#### SNR identity\.
We will also refer to the signal\-to\-noise ratio
SNR\(t\)=α¯t1−α¯t,\(DDPM\-SNR\)\\mathrm\{SNR\}\(t\)=\\frac\{\\bar\{\\alpha\}\_\{t\}\}\{1\-\\bar\{\\alpha\}\_\{t\}\},\\qquad\\text\{\(DDPM\-SNR\)\}\(46\)as used to simplify weighting expressions \(Luo \([2022](https://arxiv.org/html/2606.15048#bib.bib17), Eq\.\(109\), p\.14\)\)\.
#### Two\-time mean in unified form\.
Takingτ=t−1\\tau=t\{\-\}1, the DDPM two\-time mean𝝁t−1true\(𝒙t,𝒙0\)\{\\bm\{\\mu\}\}^\{\\text\{true\}\}\_\{t\-1\}\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)equals Equation \([44](https://arxiv.org/html/2606.15048#A2.E44)\), i\.e\., the unified form equation \([39](https://arxiv.org/html/2606.15048#A2.E39)\) with
At,τ=α¯t−1\(1−αt\)1−α¯t,κt,τ=αt\(1−α¯t−1\)1−α¯t\.A\_\{t,\\tau\}=\\frac\{\\sqrt\{\\bar\{\\alpha\}\_\{t\-1\}\}\\,\(1\-\\alpha\_\{t\}\)\}\{1\-\\bar\{\\alpha\}\_\{t\}\},\\qquad\\kappa\_\{t,\\tau\}=\\frac\{\\sqrt\{\\alpha\_\{t\}\}\\,\(1\-\\bar\{\\alpha\}\_\{t\-1\}\)\}\{1\-\\bar\{\\alpha\}\_\{t\}\}\.
### B\.2DDIM: Two\-Time Mean and Parameters \(short\)
Use the same notation as DDPM, letαt=1−βt\\alpha\_\{t\}=1\-\\beta\_\{t\}andα¯t=∏s=1tαs\\bar\{\\alpha\}\_\{t\}=\\prod\_\{s=1\}^\{t\}\\alpha\_\{s\}\. DDIM specifies a non\-Markovian reverse conditional whose mean depends on\(𝒙t,𝒙0\)\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)Song et al\. \([2021a](https://arxiv.org/html/2606.15048#bib.bib22), Eq\. 7\)\):
qσ\(𝒙t−1∣𝒙t,𝒙0\)=𝒩\(α¯t−1𝒙0\+1−α¯t−1−σt21−α¯t\(𝒙t−α¯t𝒙0\)⏟mean,σt2I\)\.q\_\{\\sigma\}\\\!\\left\(\{\\bm\{x\}\}\_\{t\-1\}\\mid\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\\right\)=\\mathcal\{N\}\\\!\\Bigl\(\\underbrace\{\\sqrt\{\\bar\{\\alpha\}\_\{t\-1\}\}\\,\{\\bm\{x\}\}\_\{0\}\+\\sqrt\{\\tfrac\{1\-\\bar\{\\alpha\}\_\{t\-1\}\-\\sigma\_\{t\}^\{2\}\}\{\\,1\-\\bar\{\\alpha\}\_\{t\}\\,\}\}\\,\\bigl\(\{\\bm\{x\}\}\_\{t\}\-\\sqrt\{\\bar\{\\alpha\}\_\{t\}\}\\,\{\\bm\{x\}\}\_\{0\}\\bigr\)\}\_\{\\text\{mean\}\},\\ \\sigma\_\{t\}^\{2\}I\\Bigr\)\.\(47\)
Hence in our unified linear form, we have
κt,t−1=1−α¯t−1−σt21−α¯t,At,t−1=α¯t−1−κt,t−1α¯t\.\\kappa\_\{t,t\-1\}=\\sqrt\{\\frac\{1\-\\bar\{\\alpha\}\_\{t\-1\}\-\\sigma\_\{t\}^\{2\}\}\{\\,1\-\\bar\{\\alpha\}\_\{t\}\\,\}\},\\qquad A\_\{t,t\-1\}=\\sqrt\{\\bar\{\\alpha\}\_\{t\-1\}\}\-\\kappa\_\{t,t\-1\}\\sqrt\{\\bar\{\\alpha\}\_\{t\}\}\.\(48\)
#### Stochasticity and theη\\eta\-parameterization \(DDIM Eq\. \(16\)\)\.
A convenient schedule for the reverse variance is
σt\(η\)=η1−α¯t−11−α¯t1−α¯tα¯t−1,η∈\[0,1\]\.\\sigma\_\{t\}\(\\eta\)=\\eta\\,\\sqrt\{\\frac\{1\-\\bar\{\\alpha\}\_\{t\-1\}\}\{\\,1\-\\bar\{\\alpha\}\_\{t\}\\,\}\}\\;\\sqrt\{\\,1\-\\frac\{\\bar\{\\alpha\}\_\{t\}\}\{\\bar\{\\alpha\}\_\{t\-1\}\}\\,\},\\qquad\\eta\\in\[0,1\]\.\(49\)Special cases: \(i\)η=0⇒σt=0\\eta\{=\}0\\Rightarrow\\sigma\_\{t\}\{=\}0gives the deterministic DDIM, whereκt,t−1=1−α¯t−11−α¯t\\kappa\_\{t,t\-1\}=\\sqrt\{\\tfrac\{1\-\\bar\{\\alpha\}\_\{t\-1\}\}\{1\-\\bar\{\\alpha\}\_\{t\}\}\}andAt,t−1=α¯t−1−1−α¯t−11−α¯tα¯tA\_\{t,t\-1\}=\\sqrt\{\\bar\{\\alpha\}\_\{t\-1\}\}\-\\sqrt\{\\tfrac\{1\-\\bar\{\\alpha\}\_\{t\-1\}\}\{1\-\\bar\{\\alpha\}\_\{t\}\}\}\\sqrt\{\\bar\{\\alpha\}\_\{t\}\}; \(ii\)η=1\\eta\{=\}1recovers the DDPM variance choice at steptt\.
### B\.3VP–SDE \(variance preserving\)
#### Forward dynamics and one\-time marginals\.
The VP SDE is
d𝒙t=−12β\(t\)𝒙tdt\+β\(t\)d𝒘t,t∈\[0,1\],\\,\\mathrm\{d\}\{\\bm\{x\}\}\_\{t\}\\;=\\;\-\\tfrac\{1\}\{2\}\\,\\beta\(t\)\\,\{\\bm\{x\}\}\_\{t\}\\,\\,\\mathrm\{d\}t\\;\+\\;\\sqrt\{\\beta\(t\)\}\\,\\,\\mathrm\{d\}\{\\bm\{w\}\}\_\{t\},\\qquad t\\in\[0,1\],\(VP\-SDE\)as stated in Eq\. \(11\) ofSong et al\. \([2021b](https://arxiv.org/html/2606.15048#bib.bib23)\)\. Let
α\(t\):=exp\(−12∫0tβ\(s\)𝑑s\),σ\(t\):=1−α\(t\)2\.\\alpha\(t\)\\;:=\\;\\exp\\Bigl\(\-\\tfrac\{1\}\{2\}\\\!\\int\_\{0\}^\{t\}\\beta\(s\)\\,ds\\Bigr\),\\qquad\\sigma\(t\)\\;:=\\;\\sqrt\{1\-\\alpha\(t\)^\{2\}\}\.\(50\)Solving Equation \([VP\-SDE](https://arxiv.org/html/2606.15048#A2.Ex14)\) gives the Gaussian marginal
p0t\(𝒙t∣𝒙0\)=𝒩\(α\(t\)𝒙0,\(1−α\(t\)2\)I\),p\_\{0t\}\(\{\\bm\{x\}\}\_\{t\}\\mid\{\\bm\{x\}\}\_\{0\}\)\\;=\\;\\mathcal\{N\}\\\!\\bigl\(\\alpha\(t\)\\,\{\\bm\{x\}\}\_\{0\},\\ \(1\-\\alpha\(t\)^\{2\}\)I\\bigr\),\(VP\-kernel\)which appears as the VP case of Eq\. \(29\) inSong et al\. \([2021b](https://arxiv.org/html/2606.15048#bib.bib23)\)\.
#### Two\-time conditional mean\.
Fix0≤τ<t≤10\\leq\\tau<t\\leq 1\. From the linear solution of equation \([VP\-SDE](https://arxiv.org/html/2606.15048#A2.Ex14)\), we can write
𝒙t=α\(t\)𝒙0\+α\(t\)∫0tβ\(s\)α\(s\)d𝒘s,𝒙τ=α\(τ\)𝒙0\+α\(τ\)∫0τβ\(s\)α\(s\)d𝒘s\.\{\\bm\{x\}\}\_\{t\}\\;=\\;\\alpha\(t\)\\,\{\\bm\{x\}\}\_\{0\}\\;\+\\;\\alpha\(t\)\\\!\\int\_\{0\}^\{t\}\\frac\{\\sqrt\{\\beta\(s\)\}\}\{\\alpha\(s\)\}\\,\\,\\mathrm\{d\}\{\\bm\{w\}\}\_\{s\},\\qquad\{\\bm\{x\}\}\_\{\\tau\}\\;=\\;\\alpha\(\\tau\)\\,\{\\bm\{x\}\}\_\{0\}\\;\+\\;\\alpha\(\\tau\)\\\!\\int\_\{0\}^\{\\tau\}\\frac\{\\sqrt\{\\beta\(s\)\}\}\{\\alpha\(s\)\}\\,\\,\\mathrm\{d\}\{\\bm\{w\}\}\_\{s\}\.Conditioned on𝒙0\{\\bm\{x\}\}\_\{0\},\(𝒙τ,𝒙t\)\(\{\\bm\{x\}\}\_\{\\tau\},\{\\bm\{x\}\}\_\{t\}\)is jointly Gaussian with
𝔼\[𝒙τ∣𝒙0\]=α\(τ\)𝒙0,𝔼\[𝒙t∣𝒙0\]=α\(t\)𝒙0,\\displaystyle\\mathbb\{E\}\[\{\\bm\{x\}\}\_\{\\tau\}\\\!\\mid\{\\bm\{x\}\}\_\{0\}\]=\\alpha\(\\tau\)\{\\bm\{x\}\}\_\{0\},\\quad\\mathbb\{E\}\[\{\\bm\{x\}\}\_\{t\}\\\!\\mid\{\\bm\{x\}\}\_\{0\}\]=\\alpha\(t\)\{\\bm\{x\}\}\_\{0\},Cov\(𝒙τ∣𝒙0\)=\(1−α\(τ\)2\)I,Cov\(𝒙t∣𝒙0\)=\(1−α\(t\)2\)I,\\displaystyle\\mathrm\{Cov\}\(\{\\bm\{x\}\}\_\{\\tau\}\\\!\\mid\{\\bm\{x\}\}\_\{0\}\)=\(1\-\\alpha\(\\tau\)^\{2\}\)I,\\quad\\mathrm\{Cov\}\(\{\\bm\{x\}\}\_\{t\}\\\!\\mid\{\\bm\{x\}\}\_\{0\}\)=\(1\-\\alpha\(t\)^\{2\}\)I,Cov\(𝒙τ,𝒙t∣𝒙0\)=α\(t\)α\(τ\)\(1−α\(τ\)2\)I,\\displaystyle\\mathrm\{Cov\}\(\{\\bm\{x\}\}\_\{\\tau\},\{\\bm\{x\}\}\_\{t\}\\\!\\mid\{\\bm\{x\}\}\_\{0\}\)=\\frac\{\\alpha\(t\)\}\{\\alpha\(\\tau\)\}\\bigl\(1\-\\alpha\(\\tau\)^\{2\}\\bigr\)I,where cross\-covariance follows Itô isometry:Cov\(𝒙τ,𝒙t∣𝒙0\)=α\(t\)α\(τ\)∫0τα\(s\)−2β\(s\)𝑑s=α\(t\)α\(τ\)\(1−α\(τ\)2\)I\\mathrm\{Cov\}\(\{\\bm\{x\}\}\_\{\\tau\},\{\\bm\{x\}\}\_\{t\}\\\!\\mid\{\\bm\{x\}\}\_\{0\}\)=\\alpha\(t\)\\alpha\(\\tau\)\\int\_\{0\}^\{\\tau\}\\alpha\(s\)^\{\-2\}\\beta\(s\)\\,ds=\\tfrac\{\\alpha\(t\)\}\{\\alpha\(\\tau\)\}\\bigl\(1\-\\alpha\(\\tau\)^\{2\}\\bigr\)I\. Applying the Gaussian conditioning formula—for jointly Gaussian vectors\(y1,y2\)\(y\_\{1\},y\_\{2\}\), the conditional mean is𝔼\[y1∣y2\]=𝝁1\+Σ12Σ22−1\(y2−𝝁2\)\\mathbb\{E\}\[y\_\{1\}\\mid y\_\{2\}\]=\{\\bm\{\\mu\}\}\_\{1\}\+\\Sigma\_\{12\}\\Sigma\_\{22\}^\{\-1\}\(y\_\{2\}\-\{\\bm\{\\mu\}\}\_\{2\}\)—then yields
𝝁τtrue\(𝒙t,𝒙0\)=α\(τ\)𝒙0\+\[α\(t\)α\(τ\)⋅1−α\(τ\)21−α\(t\)2\]⏟κt,τ\(𝒙t−α\(t\)𝒙0\)\.\{\\bm\{\\mu\}\}^\{\\text\{true\}\}\_\{\\tau\}\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)=\\alpha\(\\tau\)\\,\{\\bm\{x\}\}\_\{0\}\+\\underbrace\{\\left\[\\frac\{\\alpha\(t\)\}\{\\alpha\(\\tau\)\}\\cdot\\frac\{1\-\\alpha\(\\tau\)^\{2\}\}\{\\,1\-\\alpha\(t\)^\{2\}\\,\}\\right\]\}\_\{\\displaystyle\\kappa\_\{t,\\tau\}\}\\bigl\(\{\\bm\{x\}\}\_\{t\}\-\\alpha\(t\)\\,\{\\bm\{x\}\}\_\{0\}\\bigr\)\.\(VP\-TwoTime\)Equivalently, in the unified linear form𝝁τtrue\(𝒙t,𝒙0\)=At,τx0\+κt,τxt\{\\bm\{\\mu\}\}^\{\\text\{true\}\}\_\{\\tau\}\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)=A\_\{t,\\tau\}x\_\{0\}\+\\kappa\_\{t,\\tau\}x\_\{t\},
At,τ=α\(τ\)−κt,τα\(t\),κt,τ=α\(t\)α\(τ\)⋅1−α\(τ\)21−α\(t\)2\.A\_\{t,\\tau\}=\\alpha\(\\tau\)\-\\kappa\_\{t,\\tau\}\\,\\alpha\(t\),\\qquad\\kappa\_\{t,\\tau\}=\\frac\{\\alpha\(t\)\}\{\\alpha\(\\tau\)\}\\cdot\\frac\{1\-\\alpha\(\\tau\)^\{2\}\}\{\\,1\-\\alpha\(t\)^\{2\}\\,\}\.\(51\)
Remark\.With the commonly used linear scheduleβ\(t\)=β¯min\+t\(β¯max−β¯min\)\\beta\(t\)=\\bar\{\\beta\}\_\{\\min\}\+t\(\\bar\{\\beta\}\_\{\\max\}\-\\bar\{\\beta\}\_\{\\min\}\), the correspondingp0t\(𝒙t∣𝒙0\)p\_\{0t\}\(\{\\bm\{x\}\}\_\{t\}\\\!\\mid\{\\bm\{x\}\}\_\{0\}\)is given explicitly in Eqs\. \(32\)–\(33\) ofSong et al\. \([2021b](https://arxiv.org/html/2606.15048#bib.bib23)\)\.
### B\.4VE\-SDE and EDM \(variance exploding &σ\\sigma\-parameterization\)
#### Forward dynamics and one\-time marginals\.
The VE SDE reads
d𝒙t=ddtσ\(t\)2d𝒘t,σ\(0\)=0,\\,\\mathrm\{d\}\{\\bm\{x\}\}\_\{t\}\\;=\\;\\sqrt\{\\tfrac\{d\}\{dt\}\\,\\sigma\(t\)^\{2\}\}\\;\\,\\mathrm\{d\}\{\\bm\{w\}\}\_\{t\},\\qquad\\sigma\(0\)=0,\(VE\-SDE\)see Eq\. \(9\) inSong et al\. \([2021b](https://arxiv.org/html/2606.15048#bib.bib23)\)\. It induces the additive\-noise marginal
p0t\(𝒙t∣𝒙0\)=𝒩\(𝒙0,σ\(t\)2I\),p\_\{0t\}\(\{\\bm\{x\}\}\_\{t\}\\mid\{\\bm\{x\}\}\_\{0\}\)\\;=\\;\\mathcal\{N\}\\\!\\bigl\(\{\\bm\{x\}\}\_\{0\},\\ \\sigma\(t\)^\{2\}I\\bigr\),\(VE\-kernel\)the VE case of Eq\. \(29\) inSong et al\. \([2021b](https://arxiv.org/html/2606.15048#bib.bib23)\)\.𝒙t=𝒙τ\+σ\(t\)2−σ\(τ\)2𝒛\{\\bm\{x\}\}\_\{t\}=\{\\bm\{x\}\}\_\{\\tau\}\+\\sqrt\{\\sigma\(t\)^\{2\}\-\\sigma\(\\tau\)^\{2\}\}\\,\{\\bm\{z\}\}with𝒛∼𝒩\(𝟎,I\)\{\\bm\{z\}\}\\sim\\mathcal\{N\}\(\{\\bm\{0\}\},I\)independent of𝒙τ\{\\bm\{x\}\}\_\{\\tau\}\.
#### Two\-time conditional mean for VE\.
With the above coupling, conditioned on𝒙0\{\\bm\{x\}\}\_\{0\},
Cov\(𝒙t∣𝒙0\)=σ\(t\)2I,Cov\(𝒙τ∣𝒙0\)=σ\(τ\)2I,Cov\(𝒙τ,𝒙t∣𝒙0\)=σ\(τ\)2I\.\\mathrm\{Cov\}\(\{\\bm\{x\}\}\_\{t\}\\\!\\mid\{\\bm\{x\}\}\_\{0\}\)=\\sigma\(t\)^\{2\}I,\\qquad\\mathrm\{Cov\}\(\{\\bm\{x\}\}\_\{\\tau\}\\\!\\mid\{\\bm\{x\}\}\_\{0\}\)=\\sigma\(\\tau\)^\{2\}I,\\qquad\\mathrm\{Cov\}\(\{\\bm\{x\}\}\_\{\\tau\},\{\\bm\{x\}\}\_\{t\}\\\!\\mid\{\\bm\{x\}\}\_\{0\}\)=\\sigma\(\\tau\)^\{2\}I\.Thus forτ\>0\\tau\>0andt\>τt\>\\tau,
𝝁τtrue\(𝒙t,𝒙0\)=𝒙0\+σ\(τ\)2σ\(t\)2⏟κt,τ\(𝒙t−𝒙0\)=\(1−σ\(τ\)2σ\(t\)2\)𝒙0\+σ\(τ\)2σ\(t\)2𝒙t\.\{\\bm\{\\mu\}\}^\{\\text\{true\}\}\_\{\\tau\}\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)\\;=\\;\{\\bm\{x\}\}\_\{0\}\\;\+\\;\\underbrace\{\\frac\{\\sigma\(\\tau\)^\{2\}\}\{\\sigma\(t\)^\{2\}\}\}\_\{\\displaystyle\\kappa\_\{t,\\tau\}\}\\bigl\(\{\\bm\{x\}\}\_\{t\}\-\{\\bm\{x\}\}\_\{0\}\\bigr\)\\;=\\;\\Bigl\(1\-\\frac\{\\sigma\(\\tau\)^\{2\}\}\{\\sigma\(t\)^\{2\}\}\\Bigr\)\\,\{\\bm\{x\}\}\_\{0\}\\;\+\\;\\frac\{\\sigma\(\\tau\)^\{2\}\}\{\\sigma\(t\)^\{2\}\}\\,\{\\bm\{x\}\}\_\{t\}\.\(VE\-TwoTime\)Hence the unified coefficients areAt,τ=1−σ\(τ\)2σ\(t\)2A\_\{t,\\tau\}=1\-\\frac\{\\sigma\(\\tau\)^\{2\}\}\{\\sigma\(t\)^\{2\}\}andκt,τ=σ\(τ\)2σ\(t\)2\\kappa\_\{t,\\tau\}=\\frac\{\\sigma\(\\tau\)^\{2\}\}\{\\sigma\(t\)^\{2\}\}\.
#### EDM uses the same corruption and thus the same two\-time mean\.
EDM\(Karras et al\.,[2022](https://arxiv.org/html/2606.15048#bib.bib13)\)parameterizes “time” directly by the noise scaleσ\\sigmaand corrupts clean data additively:𝒙=𝒙0\+𝒏,𝒏∼𝒩\(𝟎,σ2I\)\{\\bm\{x\}\}=\{\\bm\{x\}\}\_\{0\}\+\{\\bm\{n\}\},\\ \{\\bm\{n\}\}\\sim\\mathcal\{N\}\(\{\\bm\{0\}\},\\sigma^\{2\}I\)\. Under the same Markov\-increments coupling as in VE, the joint\(𝒙τ,𝒙t\)∣𝒙0\(\{\\bm\{x\}\}\_\{\\tau\},\{\\bm\{x\}\}\_\{t\}\)\\mid\{\\bm\{x\}\}\_\{0\}is Gaussian with the same covariances as above, so the two\-time mean of EDM coincides with VE’s formula\.
#### Shared\-noise coupling for CT\.
Consistency training uses the same noise across times:
𝒙τ=𝒙0\+σ\(τ\)ϵ,𝒙t=𝒙0\+σ\(t\)ϵ,ϵ∼𝒩\(𝟎,I\)\.\{\\bm\{x\}\}\_\{\\tau\}=\{\\bm\{x\}\}\_\{0\}\+\\sigma\(\\tau\)\\,\{\\bm\{\\epsilon\}\},\\qquad\{\\bm\{x\}\}\_\{t\}=\{\\bm\{x\}\}\_\{0\}\+\\sigma\(t\)\\,\{\\bm\{\\epsilon\}\},\\qquad\{\\bm\{\\epsilon\}\}\\sim\\mathcal\{N\}\(\{\\bm\{0\}\},I\)\.In this case𝒙τ−𝒙0=σ\(τ\)σ\(t\)\(𝒙t−𝒙0\)\{\\bm\{x\}\}\_\{\\tau\}\-\{\\bm\{x\}\}\_\{0\}=\\frac\{\\sigma\(\\tau\)\}\{\\sigma\(t\)\}\(\{\\bm\{x\}\}\_\{t\}\-\{\\bm\{x\}\}\_\{0\}\)deterministically, so the posterior collapses and the mean is
𝝁τtrue\(𝒙t,𝒙0\)=𝒙0\+σ\(τ\)σ\(t\)⏟κt,τ\(𝒙t−𝒙0\)=\(1−σ\(τ\)σ\(t\)\)𝒙0\+σ\(τ\)σ\(t\)𝒙t\.\{\\bm\{\\mu\}\}^\{\\text\{true\}\}\_\{\\tau\}\(\{\\bm\{x\}\}\_\{t\},\{\\bm\{x\}\}\_\{0\}\)=\{\\bm\{x\}\}\_\{0\}\+\\underbrace\{\\frac\{\\sigma\(\\tau\)\}\{\\sigma\(t\)\}\}\_\{\\displaystyle\\kappa\_\{t,\\tau\}\}\\bigl\(\{\\bm\{x\}\}\_\{t\}\-\{\\bm\{x\}\}\_\{0\}\\bigr\)=\\Bigl\(1\-\\frac\{\\sigma\(\\tau\)\}\{\\sigma\(t\)\}\\Bigr\)\{\\bm\{x\}\}\_\{0\}\+\\frac\{\\sigma\(\\tau\)\}\{\\sigma\(t\)\}\{\\bm\{x\}\}\_\{t\}\.\(CT\-TwoTime\)HenceAt,τ=1−σ\(τ\)σ\(t\)A\_\{t,\\tau\}=1\-\\frac\{\\sigma\(\\tau\)\}\{\\sigma\(t\)\}andκt,τ=σ\(τ\)σ\(t\)\\kappa\_\{t,\\tau\}=\\frac\{\\sigma\(\\tau\)\}\{\\sigma\(t\)\}\.Similar Articles
Differencing the Diffusion Trajectory toward Uncertain Components for Time Series Forecasting
This paper proposes DiffDiff, a diffusion framework for probabilistic time series forecasting that embeds predictability asymmetry into the diffusion trajectory, outperforming six diffusion baselines on seven benchmarks across four prediction horizons.
GDSD: Reinforcement Learning as Guided Denoiser Self-Distillation for Diffusion Language Models
GDSD proposes a reinforcement learning method that directly distills denoisers from advantage-guided self-teachers for diffusion language models, avoiding biases from ELBO-based likelihood surrogates. It achieves up to +19.6% accuracy improvements on planning, math, and coding benchmarks over prior state-of-the-art methods.
Subliminal Clocks: Latent Time Modelling in Diffusion Language Models
This paper demonstrates that diffusion language models (DLMs) internally represent denoising progress as a latent timestep signal in their residual streams, which can be probed and steered to modulate model confidence and entropy.
Efficient Diffusion LLMs via Temporal-Spatial Parallel Decoding and Confidence Extrapolation
This paper introduces Temporal-Spatial Parallel Decoding (TSPD) and Confidence Extrapolation (CE) to accelerate inference in diffusion-based large language models by dynamically deciding when tokens have converged and forecasting logit trends, reducing unnecessary denoising steps while preserving output quality.
Drifting Objectives for Refining Discrete Diffusion Language Models
This paper introduces TokenDrift, a drifting objective that refines discrete diffusion language models by lifting categorical predictions to a continuous semantic space for anti-symmetric drifting, significantly improving generation quality under a fixed number of denoising steps.