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

AI summary

Overview

Research area: Computer Vision / generative modeling — distillation of diffusion models for fast text-to-image (T2I) and text-to-video (T2V) generation.

Technical level: Advanced. The paper builds on continuous-time consistency models, Jacobian-vector products (JVPs), score distillation, divergence families, and large-scale distributed training infrastructure (FSDP2, Ulysses context parallelism, FlashAttention-2).

Scope: The paper scales continuous-time consistency (sCM) to application-level image and video diffusion models of up to 14B parameters, diagnoses the resulting quality failures, and proposes score-regularized continuous-time consistency (rCM) to repair them.

What This Paper Is About

Continuous-time consistency models such as sCM are theoretically clean and work well on academic-scale diffusion benchmarks, but it is unclear whether they can distill large, strongly conditioned T2I and T2V models. The obstacle is twofold: JVP computation is hard to make compatible with large-scale training infrastructure (BF16, FlashAttention, FSDP, context parallelism), and FID-style benchmarks do not capture the fine-grained quality that these tasks demand. The paper's goal is to make JVP-based distillation work at scale and to fix the visual-quality problems that pure sCM exhibits there.

Key Contributions

1. First scaling of continuous-time consistency to application-level diffusion models. The authors develop a parallelism-compatible FlashAttention-2 JVP kernel (a Triton kernel integrating JVP into the FlashAttention-2 forward pass, supporting self- and cross-attention), make JVP work with FSDP by refactoring networks to expose per-layer forward/tangent functions, and extend JVP to Ulysses context parallelism by distributing QKV tangents. This enables sCM training on models with over 10 billion parameters and high-dimensional video.

2. A diagnosis of sCM's quality failures. Pure sCM generates sharp images but produces distortions where high accuracy or temporal consistency is needed — small text rendering in T2I, and blurry textures and unstable object geometry (e.g., object interpenetration) across frames in T2V. The authors attribute this to error accumulation and to the "mode-covering" character of sCM's forward-divergence objective, and show the issues are not resolved simply by scaling model size.

3. rCM: score-regularized continuous-time consistency. rCM adds score distillation (specifically a memory-efficient DMD loss, rather than SiD) as a long-skip regularizer on top of the sCM objective: L_rCM = L_sCM + λ·L_DMD, with λ = 0.01 found to generalize across the tested models and tasks. This pairs a "mode-seeking" reverse divergence with sCM's forward divergence. rCM needs no multi-stage training, GAN tuning, or extensive architecture/hyperparameter search.

4. Plug-in stabilization of the time derivative. A "semi-continuous time" scheme computes the spatial term (∇F)·F_teacher exactly via JVP while approximating ∂_t F with a finite difference (Δt = 10⁻⁴), stable for 2B-scale T2I models. For 10B+ models and video, the authors use "high-precision time": native continuous-time JVP with FP32 enforced for all time-embedding layers via torch.amp.autocast.

Main Findings

• sCM degrades on fine detail. Distillation with pure sCM is close to the teacher on typical prompts but shows quality degradation on challenging cases such as small text rendering; for video, artifacts are notable across prompts. Scaling model size does not fix this.

• Error accumulation explains the failure. The learning target of sCM contains a first-order self-feedback term ∂F/∂t weighted by sin(t) alongside the zeroth-order teacher signal weighted by cos(t). As cos(t)/sin(t) → 0 at large t, teacher supervision vanishes and JVP dominates; the term is numerically fragile under BF16.

• rCM matches or beats DMD2 on quality while improving diversity. rCM generally matches the state-of-the-art distillation method DMD2 on quality metrics while mitigating mode collapse, with notable advantages in generation diversity — all without GAN tuning.

• T2I (GenEval, Cosmos-Predict2, 1360×768). Teacher scores at 35×2 NFE: 0.6B = 0.81, 2B = 0.83, 14B = 0.84. Cosmos-Predict2 + DMD2 at 4 NFE: 0.6B = 0.77, 2B = 0.80. Cosmos-Predict2 + rCM at 4 NFE: 0.6B = 0.79, 2B = 0.81, 14B = 0.83; at 2 NFE: 0.6B = 0.78, 2B = 0.82, 14B = 0.81; at 1 NFE: 0.6B = 0.78, 2B = 0.81, 14B = 0.82. Figure 5 shows rCM rendering fine text details such as "Casio G-Shock", "11:44 AM", and "Thursday, March 22nd".

• T2V (VBench, Wan2.1 480p, 832×480×81). Teacher at 50×2 NFE: 1.3B total score 83.02, 14B 83.58. DMD2 at 4 NFE (1.3B): 84.56. rCM at 4 NFE: 1.3B = 84.43, 14B = 84.92; at 2 NFE: 1.3B = 84.09, 14B = 85.05; at 1 NFE: 1.3B = 82.65, 14B = 83.02.

• T2V (VBench, Cosmos 720p, 1280×704×93). Cosmos-Predict2 TI2V 2B teacher at 35×2 NFE: T2V score 83.03, I2V score 88.6. Cosmos-Predict2 TI2V + rCM at 4 NFE: T2V 84.40, I2V 88.2.

• Throughput (frames per second, batch size 1 on a single H100, covering diffusion sampling and VAE decoding). Wan2.1 1.3B: teacher 0.72, rCM 14.6 at 4 NFE, 23.0 at 2 NFE, 32.3 at 1 NFE. Wan2.1 14B: teacher 0.18, rCM 4.5 at 4 NFE, 8.3 at 2 NFE, 14.4 at 1 NFE. Cosmos-Predict2 2B: teacher 0.32, rCM 4.6 at 4 NFE.

• Efficiency claim. The distilled models generate high-fidelity samples in only 1~4 steps, accelerating diffusion sampling by 15×~50×.

• Comparison figure. Figure 1 shows 5 random video samples from 4-step sCM, DMD2, SiD, and rCM on Wan2.1 1.3B; rCM resolves sCM's quality issues while showing clear superiority to DMD2/SiD in generation diversity, with highly similar object position/orientation/motion to the teacher and sCM. The authors also report observing no clear advantage for SiD in T2I and T2V tasks, which is why they adopt the more memory-efficient DMD.

Methodology in Plain English

The authors start from sCM, which trains a student to move a noisy sample to clean data in one shot by matching the tangent of its own output along the teacher's ODE trajectory. That tangent is computed with a JVP, which PyTorch's built-in torch.func.jvp cannot handle under the sharding and attention kernels used for large models. So they write a custom Triton FlashAttention-2 kernel that carries tangents through the same block-wise tiling, refactor each network layer to accept and emit tangents so FSDP sharding still works, and distribute tangents across GPUs in the same all-to-all pattern that Ulysses context parallelism uses for QKV.

They also simplify sCM so it can be applied to an existing teacher without retraining: any noise schedule is mapped to the TrigFlow schedule by matching signal-to-noise ratio (σ/α = tan(t)), and a "wrapped" teacher is constructed in FP64. The sCM loss uses tangent normalization with c = 0.1, and the weighting cos(t) is absorbed into the JVP.

Having built this infrastructure, they apply it to real T2I and T2V models and find the quality problems described above. Their fix is to add a second, complementary loss: a DMD-style score-distillation term that supervises the student on its own generated samples using a fake score network, which is updated jointly in the same flow-matching manner as DMD2, with both student and fake score initialized from teacher weights. Because the student is a consistency model, rollouts alternate reverse denoising and forward noising from pure noise; the number of simulation steps N is drawn randomly from [1, N_max] and only the final step t_N → 0 is backpropagated. Instead of DMD2's fixed timesteps, they draw each t̂_n ∼ p_D and set t_n = min(t̂_n, t_{n-1}) to keep the sequence monotonically decreasing.

Experiments distill Cosmos-Predict2 T2I models (0.6B, 2B, 14B) and Wan2.1 T2V models (1.3B, 14B), using curated data plus synthetic data generated by Wan2.1 T2V 14B for the Wan distillation, with FSDP2, Ulysses CP, and selective activation checkpointing, full-parameter tuning without LoRA, and classifier-free guidance distilled into the student. Evaluation uses GenEval for T2I and VBench for video, with NFE and FPS as efficiency measures.

Why This Matters

Impact on research. The paper reframes diffusion distillation through the lens of divergence direction: forward divergence (used in pre-training and consistency models) is mode-covering and tends to lower sample quality, while reverse divergence (used in post-training and score distillation) is mode-seeking and tends to collapse modes. rCM shows empirically that combining the two at scale yields quality comparable to or better than DMD2 with improved diversity, and it removes the need for GAN tuning or multi-stage pipelines. It also provides an open infrastructure recipe (FlashAttention-2 JVP kernel, FSDP and context-parallel compatibility) that the community can reuse for any JVP-based distillation.

Real-world applications (from the tasks evaluated and discussed):

• Fast text-to-image generation, including images with fine-grained rendered text, on models of 0.6B to 14B parameters.
• Fast text-to-video generation up to 5-second videos, with the temporal-consistency and motion quality gains that VBench measures.
• Image-to-video generation, evaluated on Cosmos-Predict2 TI2V at 720p with an I2V score of 88.2 for the 4-step rCM model.
• Interactive world models via autoregressive video diffusion, which the authors discuss as a promising extension direction.

Industry relevance. The paper is a collaboration between Tsinghua (BNRist Center, THU-Bosch ML Center, AI Institute) and NVIDIA, and reports throughput on a single H100 (e.g., 14.6 FPS versus 0.72 FPS for the 1.3B teacher, and 4.5 FPS versus 0.18 FPS for the 14B teacher on Wan). Reducing a 50×2-NFE teacher to 1~4 steps with comparable or better VBench scores directly targets the inference cost that limits deployment of large video models.

Future Directions

• Forward-divergence distillation for autoregressive video. The authors suggest that for autoregressive video diffusion, where Self-Forcing can be viewed as a reverse-KL-style DMD tailored to a bidirectional teacher and a causal student, forward-divergence distillation via teacher forcing with a causal teacher could complement self-forcing and improve diversity and motion dynamics.
• Extending rCM beyond the tested model families and tasks. The paper validates on Cosmos-Predict2 T2I (0.6B, 2B, 14B) and Wan2.1 T2V (1.3B, 14B); generalization to other schedules, architectures, and modalities is left open.
• Understanding and controlling the λ balance. The authors find λ = 0.01 generalizes across their models and tasks, but the paper does not report a systematic study of how this weight should be chosen in general — an open question it raises implicitly.
• Improving the numerical behavior of the JVP term. The stabilization techniques (semi-continuous time with Δt = 10⁻⁴, and FP32 time embeddings for 10B+ models and video) are presented as remedies for a fragile first-order self-feedback signal; a principled alternative is not reported.

Target Audience

Researchers and engineers working on diffusion model acceleration, consistency models, and score/adversarial distillation; practitioners who need few-step text-to-image and text-to-video generation at billion-parameter scale; and infrastructure engineers interested in making JVP-based training compatible with FlashAttention, FSDP, and context parallelism. Readers need familiarity with diffusion ODEs, score functions, and one of forward/reverse KL divergences to follow the derivations, though the high-level argument (mode-covering versus mode-seeking, and why combining them helps) is accessible to anyone who knows modern generative modeling.

Authors’ abstract

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.