Source-linked AI summary
Causal-rCM: A Unified Teacher-Forcing and Self-Forcing Open Recipe for Autoregressive Diffusion Distillation in Streaming Video Generation and Interactive World Models
Kaiwen Zheng, Guande He, Min Zhao, Jintao Zhang, Huayu Chen, Jianfei Chen, Chen-Hsuan Lin, Ming-Yu Liu, Jun Zhu, Qianli Ma
TL;DR
Causal-rCM extends rCM to autoregressive diffusion by pairing teacher-forcing with consistency modeling and self-forcing with distribution matching, then implementing them as a sequential training recipe. It achieves state-of-the-art streaming video quality across multiple inference schedules and extends causal generation to an interactive Cosmos 3 world model.
Problem
Extending rCM's complementary consistency-model and distribution-matching objectives to autoregressive diffusion requires compatible causal training paradigms.
Method
Causal-rCM sequentially applies TF, TF-CM, and SF-DMD, supported by infrastructure that integrates causal training, continuous-time JVP-based consistency models, and large-scale parallel training.
Results
Causal-rCM achieves state-of-the-art streaming video quality with 4-step, 2-step, 2-step noisy-context, and 1-step inference schedules, while TF-sCM provides over 10× faster convergence than TF-dCM.
Takeaways & Limitations
TF-CM is the most reliable initialization choice when balancing VBench scores and qualitative detail, and Causal-rCM can convert Cosmos 3's vision stream into an interactive temporal-causal world model.
Takeaways & Limitations
Frame-wise text-to-video training with long rollout depth remains fragile because 4-step SF-DMD can develop camera drift during extended training, while the best pre-SF-DMD initialization does not always yield the best final model.
Abstract
from arXiv · showhide
Autoregressive video diffusion with causal diffusion transformers has emerged as a major paradigm for real-time streaming video generation and action-conditioned interactive world models. In this work, we extend rCM, an advanced diffusion distillation framework, to autoregressive video diffusion. The core philosophy of rCM lies in the complementarity between forward and reverse divergences, represented by consistency models (CMs) and distribution matching distillation (DMD), respectively, in diffusion distillation. This philosophy naturally carries over to the autoregressive setting, where teacher-forcing (TF) provides an offline, forward-divergence causal training paradigm, while self-forcing (SF) corresponds to an on-policy, reverse-divergence refinement. Our contributions are: (1) through extensive experiments, we show that teacher-forcing CM is currently the best complement to self-forcing DMD as an initialization strategy (2) we present the first implementation of teacher-forcing-based continuous-time CMs (e.g., sCM/MeanFlow) for autoregressive video diffusion, enabled by our custom-mask FlashAttention-2 JVP kernel, achieving 10$\times$ faster convergence compared to discrete-time CMs (dCMs) (3) we introduce Causal-rCM, a leading, unified, and scalable algorithm-infrastructure open recipe for diffusion distillation and causal training (4) we achieve state-of-the-art streaming video generation performance in both frame-wise and chunk-wise settings, using only synthetic data for training. Notably, our distilled 2-step causal Wan2.1-1.3B model achieves a VBench-T2V score of 84.63 with only 1 or 2 sampling steps. We further apply Causal-rCM to Cosmos 3, an advanced omnimodal world foundation model for physical AI with action-conditioned generation capability, enabling an interactive world model.
2. Background
Diffusion models learn to reverse a noise-perturbation trajectory, while consistency and distribution-matching objectives distill this process. Autoregressive video diffusion factorizes generation over temporal blocks and introduces causal training paradigms with complementary strengths and exposure-bias trade-offs.
- Diffusion models: Diffusion models perturb clean data with Gaussian noise and learn to reverse the resulting trajectory.Velocity, score, noise, and clean-data parameterizations are theoretically interconvertible.
- Diffusion distillation: Consistency models map noisy points at arbitrary times to points on the teacher trajectory, with continuous-time variants using tangent information and JVPs.MeanFlow combines continuous-time consistency modeling with consistency trajectory modeling under rectified flow.
- Diffusion distillation: DMD matches student and teacher distributions through reverse-KL distillation, using a fake-score network to estimate the student score.Student and critic phases alternate during training.
- Autoregressive video diffusion: Autoregressive video diffusion generates frames or chunks sequentially with causal attention while performing diffusion denoising within each temporal block.This factorization supports KV caching, streaming generation, and interactive generation.
- Causal training paradigms: Teacher-forcing is stable and parallelizable but conditions on clean history, whereas diffusion-forcing uses noisy histories without reproducing inference-time rollout errors.Self-forcing trains on its own autoregressive rollouts and pairs with reverse-type objectives such as DMD or GAN losses.
3. Causal-rCM: A Leading, Unified and Scalable Algorithm-Infrastructure Open Recipe for Diffusion Distillation and Causal Training
Causal-rCM combines teacher-forcing consistency distillation with self-forcing distribution matching in a staged recipe for autoregressive diffusion. It adds continuous-time causal consistency training and supports noisy-context and variable-step acceleration techniques.
- Algorithm: Causal-rCM pairs TF-CM’s offline forward-type consistency objective with SF-DMD’s on-policy reverse-type distribution matching.This preserves rCM’s forward-reverse complementarity in the autoregressive setting.
- Algorithm: The three-stage pipeline converts a bidirectional model with TF, distills it with TF-CM, and refines it on self-generated rollouts with SF-DMD.The recipe avoids ODE-pair knowledge distillation and GAN-style post-training.
- Continuous-time consistency: TF-sCM replaces finite-step consistency targets with continuous-time tangent targets computed through the TF-masked packed forward.The clean context is fixed while only the noisy branch follows the causal teacher velocity.
- Continuous-time consistency: Causal-rCM uses an RF-native sCM formulation because a TrigFlow-wrapped TF-sCM produced degraded causal results.The paper notes that analytically convertible schedules can still induce different normalized sCM objectives.
- Acceleration: Noisy context reduces effective per-chunk latency from N+1 to N function evaluations by reusing the last denoising KV states.The method can also improve long-horizon robustness, and the authors apply it in the final SF stage.
- Acceleration: Chunk-dependent schedules allocate extra denoising computation to the first chunk, while rollout-length cycling exposes different denoising intervals to SF-DMD supervision.A nominal 2-step model can use [4, 2, 2, ...].
Parallelisms
Causal-rCM scales video distillation through sharded model states, sequence-parallel attention, and selective activation recomputation. These components reduce memory demands for large causal video diffusion training.
- Model-state parallelism: FSDP2 partitions parameters, gradients, and optimizer states across data-parallel ranks, reducing per-GPU model-state memory.This permits student, teacher, fake-score, and EMA networks to coexist in the distillation pipeline.
- Context parallelism: Ulysses-style context parallelism shards the flattened spatiotemporal token sequence across devices and redistributes QKV for local attention.All-to-all operations restore the sequence representation after attention.
- Activation memory: Selective activation checkpointing recomputes selected attention and MLP regions while preserving lighter or stateful intermediates.This provides finer-grained recomputation control than region-based checkpointing.
KV Cache
Causal-rCM treats KV caching and execution state as explicit infrastructure for rollout, recomputation, and distributed attention. Detached historical context bounds SF-DMD memory while preserving causal semantics across chunks.
- Cache execution: The KV cache supports disabled, append, and readonly modes for packed training, committing generated chunks, and generating current chunks from prior context.Detached cached tensors prevent gradients from crossing previous chunks.
- Cache representation: Pre-RoPE and post-RoPE caching trade positional flexibility against direct reuse efficiency within the attention context.Pre-RoPE caching supports alternative positional treatment such as length extrapolation.
- Memory-efficient training: Replayed back-propagation reconstructs differentiable denoising from stored final-step inputs and detached causal state, reducing rollout memory.The rollout is first built without gradients before replaying selected computations.
- Compatibility design: Compatibility is implemented through execution semantics so cache state, checkpoint recomputation, JVPs, and distributed attention remain composable.Fresh causal state describes each forward while persistent historical K/V remains detached.
- JVP and distributed attention: Layer-level JVPs and tangent-aware Ulysses communication propagate primal and tangent tensors through sharded attention using a custom FlashAttention-2 JVP kernel.The same communication pattern is used for Q, K, V and their tangent counterparts.
- Distributed KV cache: Post-all-to-all KV caching stores head-sharded, full-sequence states in the layout consumed directly by local Ulysses attention.This avoids repeated conversion of old cache entries between global and local layouts.
4. Experiments
Experiments evaluate Causal-rCM across frame-wise and chunk-wise streaming generation, initialization strategies, continuous-time consistency models, and interactive Cosmos 3 control. Results show that optimal denoising depth depends on chunk span, while Causal-rCM supports strong streaming quality and action-conditioned trajectories.
- Main results: Causal-rCM achieves state-of-the-art streaming quality across frame-wise and chunk-wise settings with 4-step, 2-step, noisy-context, and 1-step schedules.The evaluation compares against bidirectional Wan2.1 and streaming baselines including Self-Forcing, LongLive, Causal Forcing, and AnyFlow.
- Sampling schedules: 1–2 denoising steps improve frame-wise stability, while 4 steps provide the best chunk-wise trajectory for motion and intra-chunk consistency.Frame-wise 4-step SF-DMD is prone to camera drift and trains stably for about 1k iterations, whereas shorter later-chunk rollouts remain stable for around 3k iterations.
- Sampling schedules: Noisy context reduces effective per-chunk cost from N+1 to N NFEs, with 1-step sampling preferred frame-wise and 2-step noisy-context sampling preferred chunk-wise.The difference follows from single-frame chunks having limited internal structure and three-frame chunks retaining useful spatiotemporal structure under residual noise.
- Consistency initialization: TF-sCM converges over 10× faster than TF-dCM, exceeding 81.8 VBench-T2V frame-wise and 83 chunk-wise within 1–2k iterations.TF-dCM remains lower after substantially longer training, including 10k iterations in the frame-wise comparison.
- Initialization ablations: TF-CM is the most reliable SF-DMD initialization overall, although TF-dCM can refine more stably frame-wise and DF/TF can score near 85 chunk-wise while losing fine details.Among TF-CM variants, TF-sCM slightly outperforms TF-dCM while requiring fewer SF-DMD iterations.
- Interactive world model: Interactive Cosmos 3 preserves causal temporal attention over vision supertokens and produces distinct future trajectories under left-turn, right-turn, and stay-forward controls.Actions are aligned to the next generated vision supertoken, so each action controls the transition between consecutive latent states.
5. Related Work
Related work situates Causal-rCM within JVP-based generative modeling, forward–reverse distillation, and autoregressive video diffusion. The paper connects these lines through unified CM/CTM recipe comparisons and causal teacher-forcing/self-forcing distillation.
- Differential information and JVPs: Causal-rCM extends JVP-based consistency modeling with DMD regularization to autoregressive video diffusion under clean causal contexts.The method applies teacher-forcing sCM as structured initialization for self-forcing DMD.
- Forward-reverse complementarity: Forward–reverse distillation recipes combine coverage-preserving forward components with quality- or reward-seeking reverse components.The paper frames CM and DMD methods within this broader family of few-step distillation approaches.
- Video and autoregressive diffusion distillation: Self-forcing addresses autoregressive exposure bias, while APT2 uses teacher-forcing consistency distillation to initialize self-forcing in causal architectures.The related-work discussion positions Causal-rCM alongside these autoregressive distillation methods.
- Recipe taxonomy: Table 6 organizes distillation methods as subsets of CM and CTM routes with bidirectional and causal variants.The table notes that discrete-time methods can serve as warmup stages for continuous-time JVP methods.
6. Limitations and Future
The paper identifies fragility in long frame-wise rollouts, misalignment between initialization quality and refinement stability, and lower ceilings from joint optimization. It proposes more systematic staged recipes and improved systems support as future directions.
- Limitations: Long frame-wise SF-DMD rollouts remain fragile because 4-step models can develop directional camera drift and cannot train for extended durations.The paper notes that action-conditioned settings may reduce this ambiguity by providing an explicit motion prior.
- Limitations: TF-sCM provides a stronger initialization than TF-dCM, but TF-dCM can achieve a higher final frame-wise peak through more stable long SF-DMD refinement.Initialization quality and refinement stability are therefore not fully aligned.
- Limitations: Joint optimization lowers the VBench ceiling in the causal setting, so the current recipe remains staged rather than fully joint.The paper attributes this possibility to a distribution gap between causal training stages.
- Future directions: Future work should improve custom-attention, JVP, and KV-cache kernels alongside runtime features such as torch.compile, CUDA Graphs, and NVFP4.These systems improvements are intended to reduce training and inference overhead.
A. Theoretical Analysis of TrigFlow-sCM and RF-sCM
TrigFlow-wrapped and direct RF-sCM share the same consistency map and zero-consistency condition after coordinate conversion, but generally optimize different normalized MSE objectives. The wrapper introduces no new RF-JVP structure, while tangent normalization and finite-precision evaluation create practical differences.
- TrigFlow wrapping and direct RF parameterization define the same consistency map after the coordinate change in Eqn. 21.
- The TrigFlow wrapper introduces no new RF-JVP structure; after input and output transforms, the same RF combination appears.
- The exact-arithmetic relation between TrigFlow and RF tangents differs by the factor b^2 = Z^-2.
- The two implementations share the same zero-consistency condition but are not generally the same normalized MSE objective.
- Tangent normalization breaks strict equivalence because the TrigFlow denominator depends on the stabilizing constant when c > 0.
- Finite-precision mixed-precision training can make the wrapped and direct implementations non-bitwise-equivalent, especially after normalized tangent targets amplify numerical differences.
B. FlashAttention-2 JVP Kernel with Custom Masks
The custom FlashAttention-2 JVP kernel computes tangent attention through the same sparse custom-mask schedule as the primal forward pass. This supports packed teacher-forcing and related causal training while avoiding dense-mask memory costs.
- TF-sCM evaluates packed clean-context and noisy-target tokens under a TF mask, requiring the JVP to use exactly the primal masked attention operator.
- Masked-out entries receive zero tangent contribution because the discrete mask has no tangent.
- The kernel computes attention tangents in the same online-softmax pass as primal FlashAttention, maintaining aligned normalization through shared rescaling.
- Custom masks are represented as query groups with admissible key ranges, covering dense attention and structured causal masks through query-key rectangles.
- The same sparse schedule computes the primal scores and tangent scores, ensuring the tangent corresponds to the masked attention computation.
- The implementation tiles queries and keys, processes allowed key ranges in parallel, and returns primal outputs, log-sum-exp values, and tangent outputs.