From Slow Bidirectional to Fast Autoregressive Video Diffusion Models (CausVid)

Yin et al. 2025

code

---

Hozy Summary

• Problem)
  • Existing video diffusion models with the DiT backbone struggle in interactive applications
    • Why?) Bidirectional Attention Dependencies
      • Desc.)
        • Existing video diffusion models typically use full bidirectional attention across video frames.
        • Thus, each frame can attend to the entire video, including future frames.
        • Hence, they are unsuitable for streaming autoregressive generation and interactive applications.
        • Also, full attention is too costly.
• Sol.)
  • Adapt a pretrained bidirectional DiT to an AR transformer
  • Extend DMD to videos to reduce latency
    • Student initialization on teacher’s ODE trajectories
    • Asymmetric distillation strategy
      • i.e.) Supervise a causal student model with a bidirectional teacher
      • Authors argue that using a bidirectional teacher avoids inheriting the error accumulation of a causal teacher during distillation.
• Method)
  1. Initialize the DMD-style model roles
     • \(s_{\text{data}}\) : pretrained bidirectional teacher score model
       • Bidirectional
     • \(G_\phi\) : few-step causal student generator
       • causal between blocks
       • bidirectional within blocks
     • \(s_{\text{gen},\xi}\) : auxiliary model estimating the generator/fake score
       • Bidirectional
  2. Further initialize \(G_\phi\) with ODE regression on teacher trajectories.
  3. Train asymmetric DMD.
     • cf.) Implementation details
       • Mostly follows the DMD2 implementation but differs in that…
         • The input state for \(G_\phi\) is obtained by adding noise to real video data.
           • Unlike original DMD, where the generator sample starts from z ~ N(0,I).
           • The resulting prediction \(\hat{x}_0\) is then used for the DMD objective.
         • GAN loss is not applied.
       • Diffusion Forcing style
         • Each chunk independently selects a timestep from the few-step schedule.
     • Main Training Algorithm)
       • Pipeline)
         • Divide a video into \(L\) chunks.
         • \(\{x_0^i\}_{i=1}^{L}\sim\mathcal D \stackrel{\Psi(\cdot,\epsilon^i,t^i)}{\longrightarrow} \underbrace{\{x_{t^i}^i\}_{i=1}^{L}}_{\text{noisy}} \longrightarrow \underbrace{\hat{x}_0 = G_\phi(\{x_{t^i}^i\}_{i=1}^{L}, \{t^i\}_{i=1}^{L})}_{\text{fake image}}\).
           • \(i\) : video chunk index.
           • cf.) \(t^i\) denotes the independent timestep for each chunk (Diffusion Forcing).
       • Objectives)
         • \(\mathcal{L}_{\text{DMD}}\) : Distribution Matching Loss
           • \(\hat{x}_0\stackrel{\Psi(\cdot, \epsilon, t_{\text{DMD}})}{\longrightarrow} \underbrace{\hat{x}_{t_{\text{DMD}}}}_{\text{noisy}} \longrightarrow \begin{cases} s_{\text{real}}(\hat{x}_{t_{\text{DMD}}}, t_{\text{DMD}}) \\ s_{\text{gen},\xi}(\hat{x}_{t_{\text{DMD}}}, t_{\text{DMD}}) \end{cases} \longrightarrow \nabla_\phi\;\mathcal{L}_{\text{DMD}} \longrightarrow G_\phi\).
         • \(\mathcal{L}_{\text{denoise}}\) : Denoising Loss
           • \(\hat{x}_0 \stackrel{\Psi(\cdot, \epsilon', t_{\text{denoise}})}{\longrightarrow} \underbrace{\hat{x}_{t_{\text{denoise}}}}_{\text{noisy}} \longrightarrow \begin{cases} \mu_{\text{gen}}^\xi(\hat{x}_{t_{\text{denoise}}}, t_{\text{denoise}}) \\ \text{stopgrad}(\hat{x}_0 = G_\phi(x_{t}, t)) \end{cases} \longrightarrow \nabla_\xi\;\mathcal{L}_{\text{denoise}} \longrightarrow s_{\text{gen},\xi}\).

---

4. Methods

4.1 AR

• Desc.)
  • 3D VAE Encoder
    • Compress video frames into shorter chunks of latent frames
  • DMD Framework
    • Settings)
      • \(s_{\text{data}}\) : score from many-step, bidirectional diffusion teacher model
      • \(G_\phi\) : Few-step AR diffusion student model
        • Initialization)
          • Parameters are initialized with \(s_{\text{data}}\)
          • Student initialization, i.e. \(G_\phi(x_t) \approx x_0\) from teacher’s ODE trajectory
      • \(s_{\text{gen}, \xi}\) : score from an auxiliary model initialized with the teacher model, parameterized with \(\xi\), and optimized with \(\mathcal{L}_{\text{denoise}} = \left\Vert \mu_{\text{gen}}^\phi(x_t, t) - x \right\Vert_2^2\)
        • where \(x = G_\phi(z)\), i.e. the fake image.
    • More details in 4.2 below
  • Latent Diffusion
    • Attention
      • Block(Chunk)-wise causal attention
      • Within each chunk, apply bidirectional self attention
        • Masking Implementation : \(M_{i,j} = \begin{cases} 1 & \text{if } \left\lfloor\frac{j}{k}\right\rfloor \le \left\lfloor\frac{i}{k}\right\rfloor \\ 0 & \text{otherwise} \end{cases}\)
          • \(i,j\) index the frames in the sequence
          • \(k\) is the chunk size

4.2 Bidirectional to Causal Generator Distillation

• Desc.)
  • \(\left\{ x_t^i \right\}_{i=1}^L\) : \(N\) noisy input frames divided in to \(L\) chunks where \(i=1,\ldots,L\)
  • Diffusion Forcing
    • Each chunk \(\left\{ x_t^i \right\}\) has its own noise timestep \(t^i\sim[0,999]\)
  • Asymmetric Distillation
    • How?)
      • Employ bidirectional attention in the teacher model
      • Constrain the student model to causal attention
    • Why doing this?)
      • Causal diffusion models underperform bidirectional ones.
        • cf.) Causal Diffusion
          • Convert pretrained bidirectional DiT into causal diffusion teacher
          • Fine-tune with denoising loss
      • AR teacher’s error accumulation inherits to AR student.
    • Loss)
      \(\begin{aligned} \nabla_\phi\mathcal{L}_{\text{DMD}} &= \triangleq \mathbb{E}_t \left( \nabla_\phi \text{KL}(p_{\text{gen}, t} \Vert p_{\text{data}, t}) \right) \\ &\approx -\mathbb{E}_t \left( \int s_{\text{data}}(\Psi(G_\phi(\epsilon), t), t) - s_{\text{gen}, \xi}(\Psi(G_\phi(\epsilon), t), t)\frac{\text{d}G_\phi(\epsilon)}{\text{d}\phi} \text{d}\epsilon \right) \end{aligned}\)

4.3 Student Initialization

• Idea)
  • Recall that DMD initialized \(G_\phi\) with \(\mu_{\text{base}}\) parameters.
  • Authors argue that “directly training the causal student model using the DMD loss can be unstable due to architectural differences.”
  • Thus, they further adopt the strategy of training \(G_\phi\) with regression using the pair of \((x_t, x_0)\) from \(s_{\text{data}}\).
• How?)
  1. Sample a \(L\)-length noise inputs from the standard Gaussian \(\mathcal{N}(0,\mathbf{I})\).
  2. Simulate the reverse diffusion process with an ODE solver using the pre-trained bidirectional teacher model.
     • i.e.) Sample \(\left\{ x_t^i \mid i=1,\ldots,L;\; t=T,\ldots,0 \right\}\)
  3. Select a subset of \(t\) values that match those used in the student generator \(G_\phi\).
  4. Optimize \(G_\phi\) using the loss of \(\mathcal{L}_{\text{init}} = \mathbb{E}_{x, t^i}\left\Vert G_\phi\left( \left\{ x_{t^i}^i \right\}_{i=1}^N , \left\{ t^i \right\}_{i=1}^N \right) - \left\{ x_{0}^i \right\}_{i=1}^N \right\Vert\)

4.4 KV Caching during Inference

---

Code

Run Order

1. Bidirectional DMD Training : causvid/train_distillation.py
2. ODE Dataset Generation : causvid/models/wan/generate_ode_pairs.py
3. Causal ODE Pretraining : causvid/train_ode.py
4. Causal DMD Training : causvid/train_distillation.py