PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Kaiwen Zheng, Yuji Wang, Qianli Ma, Huayu Chen, Jintao Zhang, Yogesh Balaji, Jianfei Chen, Ming-Yu Liu, Jun Zhu, Qinsheng Zhang

arXiv:2510.08431v3cs.CVcs.LG

TL;DR

Large-scale application of continuous-time consistency models is constrained by JVP infrastructure and limited evaluation of fine-grained image and video quality. The paper scales JVP-based distillation with compatible kernels and proposes rCM, which combines consistency and score distillation; it reports competitive quality, stronger diversity than DMD2, and up to 50× acceleration.

  • Problem

    The applicability of continuous-time consistency models to large-scale text-to-image and text-to-video diffusion remains unclear because JVP infrastructure is challenging and FID-based evaluations miss fine-grained attributes.

  • Method

    The paper develops parallelism-compatible FlashAttention-2 JVP infrastructure and proposes rCM, integrating forward-divergence consistency distillation with reverse-divergence score distillation.

  • Results

    rCM matches or surpasses DMD2 on quality metrics while mitigating mode collapse and offering diversity advantages, with distilled models achieving up to 50× acceleration over teacher models.

  • Takeaways & Limitations

    Combining forward- and reverse-divergence principles provides a practical framework for large-scale diffusion distillation with competitive quality and diversity.

  • Takeaways & Limitations

    Pure sCM exhibits fine-detail and temporal-consistency failures, while extension of rCM to autoregressive video diffusion remains future exploration.

Abstract

from arXiv · show

Although continuous-time consistency models (e.g., sCM, MeanFlow) are theoretically principled and empirically powerful for fast academic-scale diffusion, its applicability to large-scale text-to-image and video tasks remains unclear due to infrastructure challenges in Jacobian-vector product (JVP) computation and the limitations of evaluation benchmarks like FID. This work represents the first effort to scale up continuous-time consistency to general application-level image and video diffusion models, and to make JVP-based distillation effective at large scale. We first develop a parallelism-compatible FlashAttention-2 JVP kernel, enabling sCM training on models with over 10 billion parameters and high-dimensional video tasks. Our investigation reveals fundamental quality limitations of sCM in fine-detail generation, which we attribute to error accumulation and the "mode-covering" nature of its forward-divergence objective. To remedy this, we propose the score-regularized continuous-time consistency model (rCM), which incorporates score distillation as a long-skip regularizer. This integration complements sCM with the "mode-seeking" reverse divergence, effectively improving visual quality while maintaining high generation diversity. Validated on large-scale models (Cosmos-Predict2, Wan2.1) up to 14B parameters and 5-second videos, rCM generally matches the state-of-the-art distillation method DMD2 on quality metrics while mitigating mode collapse and offering notable advantages in diversity, all without GAN tuning or extensive hyperparameter searches. The distilled models generate high-fidelity samples in only $1\sim4$ steps, accelerating diffusion sampling by $15\times\sim50\times$. These results position rCM as a practical and theoretically grounded framework for advancing large-scale diffusion distillation. Code is available at https://github.com/NVlabs/rcm.

1 INTRODUCTION

The paper scales continuous-time consistency distillation to large application-level image and video models, exposing sCM quality limitations and introducing rCM to combine complementary divergence objectives. Across large-scale tasks, rCM preserves diversity while improving quality without complex engineering requirements.

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

2 BACKGROUND

The background introduces diffusion trajectories, consistency models, continuous-time consistency training, and score distillation as complementary approaches to learning fast generative samplers. It explains their parameterizations, objectives, and practical implementation assumptions.

  • 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.
  • 2.1 DIFFUSION MODELS: 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.
  • 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.
  • 2.2 CONSISTENCY MODELS: 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.
  • 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.
  • 2.3 SCORE DISTILLATION: 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.
  • 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.

3 SCALING UP SCM

The section scales sCM to application-level image and video tasks by simplifying its implementation and building JVP infrastructure compatible with large-model training. It also finds that pure sCM develops fine-detail and temporal-consistency distortions linked to error accumulation.

  • sCM is evaluated on text-to-image and text-to-video tasks under different prompt types.
  • A simplified implementation retains the network structure because the target models use positional time embeddings, AdaLN, and QK normalization.
  • 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.
  • 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.
  • 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.

4 SCORE-REGULARIZED CONTINUOUS-TIME CONSISTENCY MODELS

rCM repairs sCM quality limitations by adding score-based regularization to long-skip consistency, combining forward and reverse divergence signals. The method uses student rollouts for DMD training and includes stability strategies for time-derivative computation.

  • 4.1 QUALITY REPAIR WITH SCORE REGULARIZATION: rCM adds score-based regularization to long-skip consistency, complementing sCM with reverse divergence.
  • 4.1 QUALITY REPAIR WITH SCORE REGULARIZATION: 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.
  • 4.1 QUALITY REPAIR WITH SCORE REGULARIZATION: 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.
  • 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.
  • 4.2 STABLE TIME DERIVATIVE CALCULATION: 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.

5 EXPERIMENTS

Experiments evaluate rCM on large text-to-image and video models using established quality, diversity, and efficiency measures. Across these settings, rCM preserves quality in few-step generation, matches or surpasses DMD2 on reported quality metrics, and retains stronger diversity.

  • 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.
  • 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.
  • 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.
  • Performance: 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.
  • 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.
  • 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.
  • 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.

6 CONCLUSION

The paper presents rCM as a scalable framework combining consistency and score distillation for large image and video diffusion models. Its infrastructure supports JVP-based training at large scale, while the combined divergences address sCM quality limitations and preserve diversity.

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

C.2 NETWORK RESTRUCTURING

The network restructuring makes layer implementations support both ordinary forward execution and JVP-mode execution. This design is illustrated for RMSNorm and delegates tangent computation through torch.func.jvp.

  • 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.
  • Layer-specific JVP computation is delegated to forward jvp using torch.func.jvp, while original forward logic and parameter initialization remain unchanged.
  • Figure 8 contrasts the original RMSNorm implementation with its JVP-enabled restructuring.

D EXPERIMENT DETAILS

The experiments use smoothed student parameters for evaluation, stability-oriented optimization settings, repeated prompts, and timestep schedules that expose a quality–diversity trade-off.

  • 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.
  • GenEval repeats its 553 test prompts four times, while VBench uses GPT-4o–augmented prompts.
  • σ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.
  • Table 4 summarizes training and evaluation configurations, with T denoting the number of latent video frames.

E MORE RESULTS

The 4-step sCM video results demonstrate scalability and improvement over discrete-time consistency models, but remain lower quality than DMD.

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

F.1 CONTINUOUS-TIME CONSISTENCY TRAJECTORY MODELS

Continuous-time consistency trajectory models extend consistency models with a target time condition for finer ODE-trajectory transitions, but they underperform sCM in basic text-to-image distillation.

  • 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.
  • F.1 CONTINUOUS-TIME CONSISTENCY TRAJECTORY MODELS: 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.
  • F.1 CONTINUOUS-TIME CONSISTENCY TRAJECTORY MODELS: 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.
  • F.1 CONTINUOUS-TIME CONSISTENCY TRAJECTORY MODELS: 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.

F.2 ANALYSIS OF JVP ERRORS

Large-scale training requires BF16 to avoid FP16 overflow, but JVP computation remains substantially more numerically sensitive than ordinary network outputs.

  • 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.
  • F.2 ANALYSIS OF JVP ERRORS: 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_θ−.
  • F.2 ANALYSIS OF JVP ERRORS: 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.
Loading 2510.08431v3…