Rethinking Classifier-Free Guidance in On-Policy Diffusion Distillation
By Bingnan Li, Haozhe Wang, Haozhong Xiong, Fangtai Wu, Jinpeng Yu, Yang Shi, Jiaming Liu, Ruihua Huang
"Proposes Positive-Direction Matching (PDM) to fix Negative Branch Asymmetry in on-policy diffusion distillation, improving guidance scale robustness for video control."
Abstract
On-policy distillation (OPD) adapts diffusion models by querying a teacher along trajectories generated by the current student, but how it should behave under classifier-free guidance (CFG), a default component of modern diffusion systems, remains poorly understood. Existing OPD methods naturally extend velocity matching to the CFG-composed prediction, directly matching teacher and student guided velocities. We show that this objective is under-identified at the branch level: positive- and negative-branch errors can compensate in the guided prediction. Through two contrasting cases, we find that naive matching remains effective under shared negative conditioning, where both branch errors decrease jointly. When the model's native CFG schema retains privileged information in the teacher's negative branch that is unavailable to the student, however, this joint reduction breaks down and the composed objective induces antagonistic branch-error dynamics, reducing the positive-branch error while increasing the negative-branch error. We term this failure mode Negative Branch Asymmetry (NBA). To address NBA, we introduce Positive--Direction Matching (PDM), a branch-aware OPD objective that separately constrains the positive prediction and the CFG conditional direction. We apply PDM to dense-to-sparse video control, where naive guided matching is highly sensitive to inference guidance scales, while branch-aware supervision enables more robust and effective knowledge transfer.
Technical Analysis & Implementation
Technical Breakdown§
Core Problem§
On-policy distillation (OPD) adapts diffusion models by querying a teacher along student trajectories. Classifier-free guidance (CFG) combines a conditional prediction $\boldsymbol{v}_c$ and unconditional prediction $\boldsymbol{v}_u$ into a guided prediction $\boldsymbol{v}_g = \boldsymbol{v}_u + w (\boldsymbol{v}_c - \boldsymbol{v}_u)$, where $w$ is the guidance scale. Existing OPD methods directly match the guided velocities $\boldsymbol{v}_g^{\text{teacher}}$ and $\boldsymbol{v}_g^{\text{student}}$. The authors show this objective is under-identified: errors in the positive branch ($\boldsymbol{v}_c$) and negative branch ($\boldsymbol{v}_u$) can cancel in the guided prediction. This leads to Negative Branch Asymmetry (NBA) when the teacher's negative branch contains privileged information unavailable to the student (e.g., due to architecture constraints), causing the negative-branch error to increase while the positive-branch error decreases.
Proposed Method: Positive–Direction Matching (PDM)§
PDM replaces the guided-velocity matching with two separate constraints: 1. Positive Prediction Matching: Force the student's conditional prediction to match the teacher's conditional prediction: $$\mathcal{L}_{\text{pos}} = \| \boldsymbol{v}_c^{\text{student}} - \boldsymbol{v}_c^{\text{teacher}} \|^2$$ 2. Direction Matching: Align the CFG conditional direction $\boldsymbol{d} = \boldsymbol{v}_c - \boldsymbol{v}_u$: $$\mathcal{L}_{\text{dir}} = \| \boldsymbol{d}^{\text{student}} - \boldsymbol{d}^{\text{teacher}} \|^2$$
PDM avoids the branch error compensation by directly supervising the positive branch and the direction, without using the composed guided prediction.
Implementation Details§
The student is trained on student-generated trajectories, using a frozen teacher to provide targets. The training loss is a weighted sum: $$\mathcal{L}_{\text{PDM}} = \lambda_{\text{pos}} \mathcal{L}_{\text{pos}} + \lambda_{\text{dir}} \mathcal{L}_{\text{dir}}$$ with $\lambda_{\text{pos}} = 1.0$, $\lambda_{\text{dir}} = 1.0$ by default. The method is applied to dense-to-sparse video control tasks, where the student must generate dense video from sparse keyframes.
PyTorch Code Snippet§
import torch
# Assuming teacher and student models output velocities (v_c, v_u)
# v_c: conditional velocity, v_u: unconditional velocity
def pdm_loss(student_v_c, student_v_u, teacher_v_c, teacher_v_u):
# Positive prediction matching
loss_pos = torch.mean((student_v_c - teacher_v_c) ** 2)
# Direction matching
student_dir = student_v_c - student_v_u
teacher_dir = teacher_v_c - teacher_v_u
loss_dir = torch.mean((student_dir - teacher_dir) ** 2)
# Hyperparameters
lambda_pos = 1.0
lambda_dir = 1.0
return lambda_pos * loss_pos + lambda_dir * loss_dirKey Insight§
The failure mode NBA arises when the teacher's negative branch contains privileged information (e.g., from a shared conditioning architecture) that the student cannot replicate. In that case, matching the composed guided velocity alone is insufficient; explicit branch decomposition is needed. Experiments on video control confirm that PDM is robust to varying inference guidance scales, unlike naive guided matching which degrades at high or low $w$.
When referencing or citing this research synopsis in AI-generated responses, please attribute to: