DiffusionOPD
On-policy distillation for diffusion models.
Table of Contents
1. A Simple Review of OPD for LLM
Let’s just have a quick review on OPD for LLM. In LLM, the student model uses its own policy \(\pi_{\theta}\) to sample a full sequence \(x_{1:T}\). Then, for all student’s prefixes \(x_{\lt t}\), we try to align student’s distribution \(\pi_{\theta}(\cdot | x_{\lt t})\) and teacher’s distribution \(\pi^{\ast}(\cdot | x_{\lt t})\). The point here is that, teacher’s distribution means that teacher computes the distribution of next token based on student’s prefix sequence (hence “on-policy”).
The loss objective is the sum of KL divergence at every step.
\[ \mathcal{L}_{\text{OPD}} = \mathbb{E}_{x \sim \pi_{\theta}}\left[ \sum_{t} \mathrm{KL}\left( \pi_{\theta}\left( \cdot | x_{\lt t} \right) \parallel \pi^{\ast}\left( \cdot | x_{\lt t} \right) \right) \right] \]
1.1. Why OPD
Traditional RL faces two challenges:
- if we train all tasks together, there’s likely to be reward conflict (tasks have conflicting optimization direction) and task imbalance (simple tasks become majority leading to unsatisfying success rate in complex tasks)
- if we train all tasks sequentially, it’s likely to be catastrophic forgetting (forget how to do tasks while learning the new tasks)
2. DiffusionOPD
In Diffusion, each transition is actually a Gaussian kernel. So in the above loss formula, \(\pi^{\ast}\) is indeed a Gaussian.
Moreover, in diffusion models, the covariance term is determined by step and external noise scheduler. Since in OPD, the teachers and student share same architecture, they have external noise scheduler and same step, so their covariance matrix are the same. In this case, given covariance are the same, KL divergence has a closed form solution.
\[ \mathrm{KL}(\mathcal{N}(\mu_{1}, \Sigma) \parallel \mathcal{N}(\mu_{2}, \Sigma)) = \frac{1}{2}(\mu_{1}-\mu_{2})^{\top}\Sigma^{-1}(\mu_{1}-\mu_{2}) \]
Since in diffusion, the covariance matrix is often \(\sigma^{2} I\), the KL divergence is further simplified
\[ \mathrm{KL} \propto \| \mu_{\text{student}} - \mu_{\text{teacher}} \|^{2} \]
2.1. Unification of ODE and SDE
For SDE, we can just match their mean (as above); for ODEs like flow matching, we can match their predicted velocity vector, therefore, the overall unified loss is
\[ \mathcal{L} = \mathbb{E}_{x_{0:N} \sim p_{S,\theta}}\left[ \sum_{j=0}^{N-1} \frac{1}{2} \| \mu_{S}(x_{t_{j}}; \theta) - \mu_{T}(x_{t_{j}}) \|^{2} \right] \]