DiffusionOPD: A Unified Perspective of On-Policy Distillation in Diffusion Models

Abstract

Reinforcement learning has emerged as a powerful tool for improving diffusion-based text-to-image models, but existing methods are largely limited to single-task optimization. Extending RL to multiple tasks is challenging: joint optimization suffers from cross-task interference and imbalance, while cascade RL is cumbersome and prone to catastrophic forgetting. We propose DiffusionOPD, a new multi-task training paradigm for diffusion models based on On-Policy Distillation (OPD). DiffusionOPD first trains task-specific teachers independently, then distills their capabilities into a unified student along the student's own rollout trajectories. This decouples single-task exploration from multi-task integration and avoids the optimization burden of solving all tasks jointly from scratch.

Theoretically, we lift the OPD framework from discrete tokens to continuous-state Markov processes, deriving a closed-form per-step KL objective that unifies both stochastic SDE and deterministic ODE refinement via mean-matching. We formally and empirically demonstrate that this analytic gradient provides lower variance and better generality compared to conventional PPO-style policy gradients. Extensive experiments show that DiffusionOPD consistently surpasses both multi-reward RL and cascade RL baselines in training efficiency and final performance, while achieving state-of-the-art results on all evaluated benchmarks.

Method

Lifting OPD from LLMs to Continuous-State Markov Chains

On-Policy Distillation (OPD), originally proposed for autoregressive language models, lets the student generate a full trajectory from its own policy and trains it to match the teacher on the prefixes it visits. We generalize this idea to any discrete-time Markov chain in which the student and teacher share the same state space and transition-kernel structure:

$$\mathcal{L}_{\text{OPD}}(\theta) \;=\; \mathbb{E}_{x_{0:N}\sim p_{S}}\!\left[\sum_{j=0}^{N-1}\;\mathrm{KL}\!\left(p_S(\cdot\mid x_{t_j})\,\big\|\,p_T(\cdot\mid x_{t_j})\right)\right].$$

Per-Step Gaussian Transitions and a Closed-Form KL

For a flow-matching model with reverse-time SDE discretization (Euler–Maruyama), each denoising step induces a Gaussian one-step kernel $\;p_{S}(x_{t_{j+1}}\!\mid\!x_{t_j})=\mathcal{N}(\mu_S(x_{t_j}),\bar\sigma_j^2 \mathbf{I})\;$ and similarly for the teacher. Crucially, the per-step covariance $\bar\sigma_j^2\mathbf{I}$ depends only on the scheduler, so it is identical for student and teacher. The reverse KL therefore admits a closed form:

$$\mathrm{KL}\!\left(p_S\,\|\,p_T\right) \;=\; \frac{\lVert\mu_S(x_{t_j};\theta)-\mu_T(x_{t_j})\rVert_2^{\,2}}{2\,\bar\sigma_j^{\,2}}.$$

Plugging this into the OPD objective gives the diffusion-domain DiffusionOPD loss — an analytic, per-step mean-matching objective optimized by direct backpropagation:

$$\mathcal{L}^{\text{diffusion}}_{\text{OPD}}(\theta) \;=\; \mathbb{E}_{x_{0:N}\sim p_{S,\theta}}\!\left[\sum_{j=0}^{N-1}\frac{\lVert\mu_S(x_{t_j};\theta)-\mu_T(x_{t_j})\rVert_2^{\,2}}{2\,\bar\sigma_j^{\,2}}\right].$$

A Unified View: SDE and ODE in One Objective

In the deterministic ODE regime, the student and teacher each induce a unique next-state target $\mu_S(x_{t_j};\theta)$ and $\mu_T(x_{t_j})$, so distribution matching collapses to pointwise transition matching — a clean squared-$L_2$ loss:

$$\mathcal{L}^{\text{diffusion-ODE}}_{\text{OPD}}(\theta) \;=\; \mathbb{E}_{x_{0:N}\sim p_{S,\theta}}\!\left[\sum_{j=0}^{N-1}\frac{1}{2}\lVert\mu_S(x_{t_j};\theta)-\mu_T(x_{t_j})\rVert_2^{\,2}\right].$$

The closed-form KL and the deterministic $L_2$ loss are two faces of the same training principle, providing a unified view of on-policy distillation across SDE and ODE samplers.

Two-Stage Training Recipe

BibTeX

If you find DiffusionOPD useful, please consider citing:

@article{li2026diffusionopd,
  title         = {DiffusionOPD: A Unified Perspective of On-Policy Distillation in Diffusion Models},
  author        = {Li, Quanhao and Yu, Junqiu and Jiang, Kaixun and Wei, Yujie and
                   Xing, Zhen and Li, Pandeng and Chu, Ruihang and Zhang, Shiwei and
                   Liu, Yu and Wu, Zuxuan},
  journal       = {arXiv preprint arXiv:2605.15055},
  year          = {2026},
  eprint        = {2605.15055},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG}
}