Large Scale Diffusion Distillation via Score-Regularized Continuous-Time Consistency

Motivation

• Large-scale sCM applicability remains unclear because JVP computation conflicts with BF16, FlashAttention, and context parallelism, while FID poorly captures strongly conditioned T2I and T2V attributes.These benchmarks emphasize fine-grained properties such as text rendering that weakly conditioned ImageNet evaluation does not capture.Source: §1 · p. 2

Infrastructure

• A FlashAttention-2 JVP kernel compatible with FSDP and context parallelism enables sCM exploration on models exceeding 10B parameters and high-dimensional video data.The infrastructure also supports investigating sCM scaling behavior and its fine-detail quality issues.Source: §1 · p. 2

Proposed method

• rCM integrates score distillation as a long-skip regularizer, pairing forward-divergence consistency training with reverse-divergence supervision on self-generated samples.The design combines offline teacher or data trajectories with on-policy student trajectories.Source: §1 · pp. 2–3

Results

• rCM matches or surpasses DMD2 on quality metrics while mitigating mode collapse and providing notable generation-diversity advantages.The framework requires no multi-stage training, GAN tuning, or extensive architecture and hyperparameter search.Source: §1 · p. 3

Future direction

• The approach is presented as promising for autoregressive video diffusion, where forward-divergence distillation could complement self-forcing and potentially improve diversity and motion dynamics.This extension is identified as future exploration rather than a validated result in the supplied passages.Source: §1 · p. 3

2.1 DIFFUSION MODELS

• Diffusion models perturb clean data with Gaussian noise along a trajectory and learn to reverse the process, with equivalent score, noise, clean-data, and velocity parameterizations.The velocity target is defined from derivatives of the signal and noise schedules.Source: §2.1 · p. 3
• Flow matching simplifies the probability-flow ODE to d x_t/dt = v_theta(x_t,t), while rectified flow uses α_t = 1 − t and σ_t = t.These schedules provide a concrete special case of the broader diffusion formulation.Source: §2.1 · p. 3

2.2 CONSISTENCY MODELS

• Consistency models learn a function mapping any teacher-trajectory point (x_t,t) to the initial point x_0 and enforce consistent outputs at adjacent timesteps.Their parameterization imposes the boundary condition f_theta(x,0) ≡ x.Source: §2.2 · pp. 3–4
• Continuous-time consistency models obtain the limit Δt → 0, using a tangent-alignment objective whose derivative can be computed with a Jacobian-vector product.sCM uses the TrigFlow schedule and preconditioning so its network predicts velocity directly.Source: §2.2 · p. 4

2.3 SCORE DISTILLATION

• Score distillation matches student and teacher distributions through diffused marginals and reverse divergences such as reverse KL and Fisher divergence.Because direct matching of clean high-dimensional distributions is difficult, the method operates on noise-perturbed samples.Source: §2.3 · p. 4
• An auxiliary fake score network approximates the intractable student score and is jointly optimized with the few-step generator through adversarial interplay.Both networks are commonly initialized from the teacher diffusion model.Source: §2.3 · p. 4

Implementation assumptions

• The formulation absorbs the data standard-deviation parameter into x_0 and, for simplicity, absorbs c_noise into F_theta.These are stated simplifications inherited from the sCM and EDM parameterizations.Source: §2.3 · p. 4

• sCM is evaluated on text-to-image and text-to-video tasks under different prompt types.Source: §3 · p. 5
• A simplified implementation retains the network structure because the target models use positional time embeddings, AdaLN, and QK normalization.Source: §3.1 · p. 5
• FlashAttention-2, FSDP, and context-parallel JVP support address the incompatibility of built-in JVP with large-scale training.The FlashAttention-2 kernel integrates JVP into self- and cross-attention, while layer-level tangent propagation supports FSDP and context parallelism distributes tangents with QKV.Source: §3.2 · pp. 5–6

3.3.1 EMPIRICAL OBSERVATION: QUALITY ISSUES

• Pure sCM produces fine-detail failures in challenging T2I prompts and blurry, geometrically unstable artifacts across T2V frames.Scaling model size alone does not solve the T2I issues, including small-text rendering distortions.Source: §3.3.1 · p. 6

3.3.2 THEORETICAL ANALYSIS: ERROR ACCUMULATION

• At large times, vanishing teacher supervision leaves training dynamics dominated by numerically fragile JVP self-feedback, allowing errors to accumulate.The JVP signal is especially fragile under BF16 precision, while teacher supervision vanishes as sin(t) approaches zero.Source: §3.3.2 · p. 6

4.1 QUALITY REPAIR WITH SCORE REGULARIZATION

• rCM adds score-based regularization to long-skip consistency, complementing sCM with reverse divergence.Source: §4.1 · p. 6, Figure 4 · p. 7
• The rCM objective is L_rCM(θ) = L_sCM(θ) + λL_DMD(θ), with λ = 0.01 generalized across the evaluated models and tasks.The DMD term uses a fake score-network denoiser, a time distribution p_D, and stop-gradient operations.Source: §4.1 · p. 7
• DMD loss and fake-score training require student-generated samples x_0 ∼ p_θ.The rollout strategy randomly selects simulation-step counts and backpropagates DMD loss only through the final step toward t_N → 0.Source: §4.1 · p. 7

4.2 STABLE TIME DERIVATIVE CALCULATION

• During rCM training, the time derivative combines JVP and partial-time-derivative terms to prevent sudden collapse after long training.The partial-time derivative is unstable because of oscillatory trigonometric time embeddings.Source: §4.2 · p. 7
• A finite-difference approximation computes the time derivative for stable 2B-scale T2I training, whereas 10B+ models and video tasks use full JVP with FP32 time embeddings.The finite-difference approach is sensitive to ∆t in larger models and video settings.Source: §4.2 · p. 7

Models, Tasks and Datasets

• rCM is evaluated on Cosmos-Predict2 T2I models from 0.6B to 14B parameters and Wan2.1 T2V models at 1.3B and 14B.Wan distillation uses curated data supplemented with synthetic data generated by Wan2.1 T2V 14B.Source: §5.1 · pp. 7–8

Evaluation Metrics

• The evaluation uses GenEval for compositional T2I prompts, VBench for video quality, and NFE plus FPS to quantify inference efficiency.Video FPS is measured with batch size 1 on a single H100 and includes diffusion sampling and VAE decoding.Source: §5.1 · p. 8

Performance

• The 14B Cosmos-Predict2 model achieves a state-of-the-art GenEval overall score of 0.83 in 4 steps.rCM closely approaches the teacher’s performance and benefits from scaling on Cosmos-Predict2.Source: §5.2 · p. 8
• Wan2.1 14B rCM reaches a total VBench score of 85, surpassing the 480p Wan teacher.Similar behavior is reported for Cosmos-Predict2 at 720p and on the I2V task.Source: §5.2 · p. 8, Table 2 · p. 9, Table 3 · p. 10

Comparison with DMD2

• rCM matches or surpasses DMD2 in GenEval and VBench quality while retaining a diversity advantage, especially for video generation.The comparison describes DMD2 as tending toward collapsed generations with similar object positions and orientations.Source: §5.2 · pp. 9–10

Generation with Fewer Steps

• rCM produces reasonable T2I samples across 1–4 steps, while 1-step T2V outputs are blurry and show a marked VBench decline.For challenging T2I prompts, fewer steps particularly impair detailed text rendering.Source: §5.2 · p. 10

Ablation Study on λ

• Increasing λ improves quality but reduces diversity, whereas smaller λ values favor diversity; λ = 0.01 preserves good quality at the tested granularity.The study searches λ ∈{1, 0.1, 0.01, 0.001} on Wan2.1-1.3B for 10k iterations.Source: §5.2 · p. 10, Figure 7 · p. 10

• rCM combines forward-divergence consistency distillation with reverse-divergence score distillation for large image and video models.The paper frames the two signals as complementary supervision on forward offline and reverse self-generated data paths.Source: §6 · p. 10
• The distilled models achieve competitive text-to-image results in one step and text-to-video results in two steps, reaching up to 50× acceleration over teacher models.Source: §6 · p. 10
• The rCM algorithm initializes student and fake-score models from the teacher, then alternates generator and critic-style updates using consistency and flow-matching losses.The supplied algorithm passages specify teacher, student, fake-score, time-distribution, and simulation-step inputs, followed by student and fake-score updates.Source: §C.1 · p. 17
• JVP computation is integrated into FlashAttention-2-style streaming attention, computing outputs and tangents without materializing the full attention matrix.The procedure operates on blocked query, key, and value sequences and updates output and tangent states iteratively.Source: §C.1 · pp. 16, 18

• The JVP base class supports standard forward execution and JVP-mode execution by returning primals together with their tangents.When withT=True, each primal and tangent is wrapped in a TensorWithT tuple.Source: §C.2 · p. 18
• Layer-specific JVP computation is delegated to forward jvp using torch.func.jvp, while original forward logic and parameter initialization remain unchanged.Source: §C.2 · p. 19
• Figure 8 contrasts the original RMSNorm implementation with its JVP-enabled restructuring.Source: Figure 8 · p. 19, §C.2 · p. 19

• Power EMA student parameters are used for evaluation, with AdamW applied to both student and fake-score optimizers.The configurations use β1 = 0, β2 = 0.999, weight decay 0.01, and no gradient clipping; the paper reports disabling clipping as crucial for rCM stability.Source: D EXPERIMENT DETAILS · p. 19
• GenEval repeats its 553 test prompts four times, while VBench uses GPT-4o–augmented prompts.Source: D EXPERIMENT DETAILS · p. 19
• σmax governs the quality–diversity trade-off, with [arctan(σmax), 1.3, 1.0, 0.6] used for 4-step sampling.For fewer than four steps, the schedule uses its first k entries; σmax = 80 is used for high-diversity visualizations.Source: D EXPERIMENT DETAILS · p. 19
• Table 4 summarizes training and evaluation configurations, with T denoting the number of latent video frames.Source: Table 4 · p. 20

• 4-step sCM video results show that the infrastructure enables scalable training and outperform the discrete-time CM counterpart, while quality remains limited versus DMD.Source: Figure 9 · p. 20

F.1 CONTINUOUS-TIME CONSISTENCY TRAJECTORY MODELS

• sCTM adds a time condition s < t and learns transitions from x_t to x_s along the probability-flow ODE.The function f_θ maps (x_t, t, s) to x_s using DDIM-derived preconditioning coefficients.Source: §F.1 · p. 20
• The continuous-time sCTM objective changes sCM coefficients while excluding s from the JVP with respect to t.Because s is independent of t, it remains uninvolved in the JVP computation.Source: §F.1 · p. 20
• Under s = t and w(t, t) = 1, the objective recovers flow matching; MeanFlow and AYF correspond to sCTM under rectified flow.Other positive weighting functions yield an equivalent objective with a gradient proportional to flow matching.Source: §F.1 · p. 21
• sCTM underperforms sCM in both quality and diversity on basic text-to-image distillation tasks.The authors suggest that learning arbitrary mappings along the ODE trajectory is more demanding than learning only the mapping to the initial point.Source: §F.1 · p. 21

F.2 ANALYSIS OF JVP ERRORS

• BF16 precision is required for neural-network computation in large-model training to avoid FP16 overflow.This precision requirement motivates evaluating numerical errors in the JVP under BF16.Source: §F.2 · p. 21
• JVP computation incurs substantially larger numerical errors under BF16 than the zeroth-order network output.The comparison concerns the JVP term dF_θ/dt versus F_θ−.Source: §F.2 · p. 21, Figure 11 · p. 22
• The error analysis compares BF16 and FP32 computations across 100 uniformly sampled timesteps using Cosmos-Predict2 text-to-image models with 0.6B and 2B parameters.Only network precision is changed; wrapping conversions remain in FP64.Source: §F.2 · p. 21