Self Forcing - Bridging the Train-Test Gap in Autoregressive Video Diffusion

Huang et al. 2025

---

Hozy Summary

• Problem)
  • Previous methods denoise future frames based on ground-truth (GT) context frames.
  • This causes exposure bias.
    • i.e.) a model is trained exclusively on ground-truth context but must rely on its own imperfect predictions at inference time, resulting in a distributional mismatch that compounds errors as generation progresses.
    • Additionally, the error accumulates following the AR generation.
    • Consequently, the quality of the generated video degrades.
• Sol.)
  • Generate each frame conditioned on previously self-generated frames rather than GT ones.
    • How?)
      • Self Rollout
      • Supervision with Holistic Distribution Matching Loss
  • Computation
    • Few step diffusion distillation backbone
    • Gradient truncation strategy
  • Further accomplishments
    • Long video generation through rolling KV cache mechanism
• Implementation)
  • Main Training Algorithm) model.base.py
    • Pipeline)
      • For blocks \(i=1,\ldots,N\) (causality!)
        • Sample \(s\sim\mathcal U\{1,\ldots,K\}\) : randomly selected exit denoising-step index.
          • where \(\{t_1,\ldots,t_K\}\) : student’s few denoising timesteps
          • Sampled at each training iteration.
            • The default implementation samples a separate \(s^i,\quad i=1,\ldots,N\)
          • Gradient is enabled only at the selected exit step \(j=s\).
        • Self Rollout
          • \(\underbrace{x_{t_K}^i \sim \mathcal{N}(0, \mathbf{I})}_{t=t_K} \longrightarrow {\left\{\begin{array}{ccc} \nearrow & \overbrace{\hat{x}_0^i = G_\theta(x_{t_j}^i;\; t_j, \text{KV})}^{\text{Predict } x_0 \text{ cond. on } \hat{x}_0^{\lt i} \text{ in KV}} & \searrow \\ \nwarrow & \underbrace{x_{t_{j-1}^i} = \Psi(\hat{x}_0^i, \epsilon, t_{j-1})}_{\text{fwd(inject noise) to } t_{j-1}} & \swarrow \end{array} \right\}}_{j=K,\ldots,s} \longrightarrow \begin{cases} \overbrace{X_\theta\text{.append}(\hat{x}_0^i)}^{\text{Append generated block to the rollout video}} \quad (j=s \text{ only!}) \\ \underbrace{\text{KV}\leftarrow \overbrace{G_\theta^{\text{KV}}(\hat{x}_0^i)}^{t=0}}_{\text{Cache the } i \text{-th block}}\end{cases}\).
            • Desc.)
              • \(\{t_1,\ldots,t_K\}\) : student’s few denoising timesteps
              • Compute and append the K/V features of the generated \(i\)-th block.
    • Objectives)
      • Same DMD objective as the CausVid
        • \(\mathcal{L}_{\text{DMD}}\) : Distribution Matching Loss
        • \(\mathcal{L}_{\text{denoise}}\) : Denoising Loss
  • Inference Algorithm) pipeline.self_forcing_training.py

---

3. Setup & Model

3.1 Preliminaries: Autoregressive Video Diffusion Models

• Idea)
  • Combine AR chain-rule decomposition with denoising diffusion models for generation
• Def.)
  • \(x^{1:N} = \left( x^1, x^2, \ldots, x^N \right)\) where \(N\) is the length of the sequence
    • factorized using the chain rule of \(p\left( x^{1:N} \right) = \displaystyle\prod_{i=1}^N p\left( x^i \mid x^{\lt i} \right)\)
      • with each \(p\left( x^i \mid x^{\lt i} \right)\) is modeled using a diffusion process
  • Each frame \(x^i\) is corrupted by the forward process \(q_{t^i\mid0}\left( x_{t^i}^i \mid x_0^i \right)\)
    • s.t. \(x_{t^i}^i = \Psi\left( x^i, \epsilon^i, t^i \right) = \alpha_{t^i}x^i + \sigma_{t^i}\epsilon^i\) : the forward process
      • where
        • \(\alpha_{t^i}, \sigma_{t^i}\) : pre-defined noise schedule
        • \(t^i\in[0,1000]\) : a finite time horizon
        • \(\epsilon^i\sim\mathcal{N}(0,mathbf{I})\) : Gaussian noise
    • \(p_\theta\left( x^i \mid x^{\lt i} \right)\) is implicitly defined as \(f_{\theta, t_1}\circ f_{\theta, t_2}\circ \cdots \circ f_{\theta, t_T}\left( x_{t_T}^i \right)\)
      • where \(f_{\theta, t_j}\left( x_{t_j}^i \right) = \Psi\bigg( \underbrace{G_\theta\left( x_{t_j}^i, j_j, x^{\lt i} \right)}_{\text{output from network}}, \epsilon_{t_{j-1}}, t_{j-1} \bigg)\)
  • Optimization Target : Epsilon prediction
    • \(\mathcal{L}_\theta^{\text{DM}} = \mathbb{E}_{x^i, t^i, \epsilon^i} \left[ w_{t^i} \Vert \hat{\epsilon}_\theta^i - \epsilon^i \Vert_2^2 \right]\) s.t.
      • \(\hat{\epsilon}_\theta^i := G_\theta\left(x_{t^i}^i, t^i, x\right)\) : a neural network conditioned on the context \(c\)
      • \(w_{t^i}\) : the weighting function
• Self Forcing Set-ups
  • DiT with text conditioning on latent space encoded by a causal 3D VAE
  • AR chain-rule decomposition implemented with causal attention

Concept) Teacher Forcing (TF)

• Idea)
  • Train the model to predict the next token conditioned on ground-truth tokens
  • In video diffusion, TF denoises each frame using clean, ground-truth context frames

Concept) Diffusion Forcing (DF)

• Idea)
  • Train the model on videos with noise levels independently sampled for each frame
  • In video diffusion, TF denoises each frame based on noisy context frames

3.2 Autoregressive Diffusion Post-Training via Self-Rollout

• Setup)
  • Student ODE Initialization Protocol from CausVid
    • Recap)
      • Optimize \(\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\)
• Idea)
  • Sample a batch of videos \(\left\{x^{1:N}\right\} \sim p\left( x^{1:N} \right) = \displaystyle\prod_{i=1}^N p\left( x^i \mid x^{\lt i} \right)\)
    • conditioned on self-generated outputs including both…
      • clean context frames in the past
      • noisy frames at the current time step
    • How?)
      • At each denoising step \(t_j\) and the frame index \(i\), the model denoises an intermediate noisy frame \(x_{t_j}^i\) by…
        1. Condition on previous clean frames \(x^{\lt i}\)
        2. Obtain previous timestep noisy frame \(x_{t_{j-1}}^i\) through the forward process \(\Psi\) and inject a Gaussian noise with lower noise level
        3. Perform few-step diffusion process
• Rollout Implementation)
  • At timestep \(t_j\),
    1. Suppose we have a intermediate denoising latent of \(x_{t_j}^i\) given from the previous \(t=t_{j+1}\) timestep.
       • cf.) When \(t=T,\quad x_{T}^i\sim\mathcal{N}(0,\mathbf{I})\)
    2. Sample \(\hat{x}_0^i \leftarrow G_\phi\left( x_{t_j}^i;\; t_j \right)\)
       • i.e.) Denoise to \(t=0\) using the generator.
    3. Inject noise to \(\hat{x}_0^i\) to back to the level of \(t=t_{j-1}\) using the forward process \(\Psi\).
       • i.e.) \(x_{t_{j-1}}^i \leftarrow \Psi\left( \hat{x}_0^i, \epsilon, t_{j-1} \right)\) where \(\epsilon\sim\mathcal{N}(0,\mathbf{I})\)
  • Repeat this for \(t=T,\ldots,s\) where \(s\sim \mathcal{U}(1,T)\) is a sampled denoising timestep.
    • At each epoch, it uses the \(s\)-th step output as the final output.
      • Why doing this?)
        • To ensure all intermediate denoising steps receive supervision signals

• Other technical details
  • Employ KV caching during training as well.
  • Gradient truncation
    • i.e.) Limit the backpropagation to only the final denoising step of each frame
  • Detach the gradients of the previous frames from the current frame by restricting gradient flow into KV cache embeddings.

3.3 Holistic Distribution Matching Loss

• Goal)
  • Utilize distillation method to enhance the quality of AR video generation
• Method)
  • Let
    • \(p_\theta\left( x^{1:N} \right)\) : the distribution of generated videos
    • \(p_\text{data}\left( x^{1:N} \right)\) : the distribution of real videos
  • Inject noise to both distributions as \(p_{(\cdot), t}\left( x_t^{1:N} \right) = \displaystyle\int q_{t\mid0}\left( x_t^{1:N}\mid x^{1:N} \right) p_{(\cdot)}\left(x^{1:N}\right) \text{d} x^{1:N}\)
  • Match the whole video (not individual frames) \(p_{\theta, t}\left( x_t^{1:N} \right)\) and \(p_{\text{data}, t}\left( x_t^{1:N} \right)\) using…
    1. Distribution Matching Distillation (DMD) : \(\mathbb{E}_t\left[ \mathcal{D}_{\text{KL}}\left( p_{\theta, t} \Vert p_{\text{data}, t} \right) \right]\)
    2. Score Identity Distillation (SiD) : \(\mathbb{E}_{t, p_{\theta, t}}\left[ \Vert \nabla\log p_{\theta, t} - \nabla \log p_{\text{data}, t} \Vert^2 \right]\)
    3. GAN : Minimizing JS-divergence

3.4 Long Video Generation with Rolling KV Cache

• How?)
  • Maintain a fixed-size KV cache for most recent \(L\) frames.
    • When full, remove the oldest ones, and push new ones.
  • This cause severe flickering artifacts specifically for the first latent frame.
    • Why?)
      • During training, the model is provided with clean first frames, while this is not available during the rolling KV cache procedure.
    • Sol.)
      • During training, restrict the attention window so the model cannot attend to the first chunk when denoising the final chunk