Source-linked AI summary
DiffusionOPD: A Unified Perspective of On-Policy Distillation in Diffusion Models
Quanhao Li, Junqiu Yu, Kaixun Jiang, Yujie Wei, Zhen Xing, Pandeng Li, Ruihang Chu, Shiwei Zhang, Yu Liu, Zuxuan Wu
TL;DR
Diffusion models are typically optimized for single tasks, while multi-task RL faces interference, imbalance, cumbersome training, and forgetting. DiffusionOPD distills independently trained task teachers into one student using closed-form on-policy KL objectives, achieving consistently better training efficiency and final performance than multi-reward and cascade RL baselines, with state-of-the-art benchmark results.
Problem
Existing diffusion-model RL methods mainly optimize single tasks, while multi-task approaches face objective conflict, task imbalance, cumbersome staging, and catastrophic forgetting.
Method
DiffusionOPD independently trains task-specific teachers, then distills them into a unified student along student rollouts using a closed-form per-step KL objective for SDE and ODE samplers.
Results
DiffusionOPD consistently improves training efficiency and final performance over prior multi-task baselines, achieving state-of-the-art results on aesthetics, OCR, and GenEval.
Takeaways & Limitations
Closed-form KL optimization provides lower-variance training than PPO-style policy gradients while applying across both stochastic SDE and deterministic ODE samplers.
Takeaways & Limitations
The PPO-equivalence analysis assumes parameters remain fixed throughout each rollout, so the rollout policy equals the current student policy.
Abstract
from arXiv · showhide
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 Online Policy Distillation (OPD). DiffusionOPD first trains task-specific teachers independently, then distills their capabilities into a unified student along the student 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.
1 Introduction
DiffusionOPD addresses multi-task diffusion training challenges by separating single-task on-policy exploration from multi-task capability integration, distilling task-specific teachers into a unified student along its own rollouts. It extends on-policy distillation to diffusion models with a unified closed-form objective and reports stronger efficiency and performance than multi-task RL baselines.
- Motivation: Multi-task RL faces objective conflict and task-difficulty imbalance under joint optimization, while cascade RL requires cumbersome sequential stages and risks catastrophic forgetting.These limitations motivate decoupling exploration from capability integration.
- DiffusionOPD: DiffusionOPD trains task-specific teacher models independently, then distills their capabilities into a unified student along the student’s own rollout trajectories.This separates single-task on-policy exploration from multi-task capability integration.
- Diffusion OPD objective: The method lifts OPD from autoregressive token transitions to continuous-state diffusion denoising transitions, deriving a closed-form per-step KL objective for stochastic reverse-time SDE processes.Student and teacher define Gaussian one-step transition kernels with shared covariance at each denoising state.
- Optimization: Direct optimization of the closed-form KL avoids the Gaussian-noise score-function term in PPO-style surrogates, reducing gradient variance and extending naturally to deterministic ODE samplers through transition matching.The framework therefore provides a unified view across stochastic and deterministic diffusion samplers.
- Generality and evaluation: DiffusionOPD is a general on-policy distillation framework: teachers can supervise student-visited denoising states with a broad family of existing distillation objectives, and experiments surpass multi-task RL baselines in training efficiency and final performance.Ablations examine the distillation objective, loss formulation, and sampler noise level.
2 Related Works
Related work has applied reinforcement learning to improve diffusion-based text-to-image models under task-specific rewards. Diffusion distillation instead transfers teacher knowledge to students, with prior work largely emphasizing few-step inference through step distillation.
- Reinforcement Learning for Diffusion Models: Reinforcement learning improves diffusion-based text-to-image models under task-specific rewards, including aesthetic quality, text rendering accuracy, and compositional alignment.This line of work builds on advances in reinforcement learning and adapts them to diffusion generation.
- Diffusion Distillation: Diffusion distillation transfers knowledge from a teacher diffusion model to a student model, most often compressing many-step generation into few-step inference.Prior work primarily focuses on step distillation for more efficient inference.
- Diffusion Distillation: Trajectory distillation imitates intermediate denoising transitions or enforces consistency across timesteps.It is one of two broad categories described for existing diffusion-distillation approaches.
3 Method
DiffusionOPD extends on-policy distillation from token sequences to continuous-state diffusion processes, yielding closed-form per-step supervision along student rollouts. Its mean-matching objective is exact, lower variance than PPO-style optimization, and applies to both stochastic SDE and deterministic ODE samplers, within a two-stage multi-task training paradigm.
- Continuous-state OPD: OPD samples trajectories from the student and matches the teacher through closed-form per-step transition KLs, avoiding REINFORCE.The continuous-state lift preserves on-policy sampling and analytic per-step KL evaluation.
- Stochastic SDE regime: For same-covariance Gaussian transitions evaluated at identical student states, the exact reverse KL depends only on the student–teacher mean mismatch.The shared covariance is determined by the scheduler and global noise level, while the sample noise cancels analytically.
- Deterministic ODE regime: In deterministic ODE sampling, DiffusionOPD directly matches the teacher’s one-step transition targets along the student’s rollout trajectory.Unlike stochastic token distributions, deterministic diffusion updates uniquely determine the next state from the current latent.
- Optimization: Direct KL minimization and PPO-style optimization have the same expected gradient, but the closed-form KL removes PPO’s additional noise-proportional score-function variance.The direct loss uses pathwise backpropagation, whereas PPO remains unbiased but has nonzero gradient variance.
- Optimization: The closed-form KL supports both stochastic and deterministic samplers, whereas PPO-style optimization relies on a stochastic policy density and importance ratio.DiffusionOPD therefore uses one training principle across SDE and ODE sampling regimes.
- Two-stage multi-task training: DiffusionOPD first trains separate reward-specialized teachers, then distills them into one student using task-wise on-policy rollouts and accumulated losses.Independent teacher training avoids inter-task interference before multi-task integration.
4 Experiments
Experiments compare DiffusionOPD with multi-task RL, cascade RL, and alternative distillation methods, then examine its loss formulation and sampler noise. DiffusionOPD achieves stronger multi-domain performance, improved efficiency, and better optimization behavior than the evaluated alternatives.
- Experimental setup: The experiments use SD3.5-Medium at 512×512 with rule-based and model-based rewards, LoRA adaptation, and a 40-step first-order ODE sampler.The reward suite includes GenEval, OCR, PickScore, ClipScore, HPSv2.1, Aesthetics, ImageReward, and UnifiedReward.
- Multi-task comparisons: Single-task teachers specialize in their own domains, whereas multi-task RL improves task coverage but requires substantially longer training.GenEval, OCR, and Aesthetics teachers respectively excel at compositional alignment, text rendering, and aesthetic objectives, with limited transfer beyond those targets.
- Multi-task comparisons: DiffusionOPD achieves the best overall performance and superior visual quality while requiring less training time than multi-task RL to reach the same target score.It also attains a substantially higher performance ceiling than the multi-task RL baselines.
- Distillation methods: DiffusionOPD is compared with DMD, TDM, and SFT, using student on-policy rollouts for the distillation baselines except SFT’s teacher-generated samples.The comparison evaluates methods for transferring knowledge from multiple single-task teachers.
- Ablation studies: At identical noise level a = 0.7, the closed-form KL objective improves rewards faster and reaches a higher performance ceiling than PPO-style policy gradients.The sampler ablation further finds that lower noise consistently accelerates convergence and raises evaluation scores, with the ODE sampler up to five times more efficient than SDE at noise level=0.7.
5 Conclusion
DiffusionOPD introduces an on-policy distillation paradigm for multi-task diffusion-model training. It separates single-task exploration from multi-task capability integration and extends OPD to diffusion Markov chains with a closed-form per-step reverse-KL objective.
- DiffusionOPD is an on-policy distillation paradigm designed for multi-task training of diffusion models.
- By decoupling single-task exploration from multi-task capability integration, DiffusionOPD avoids joint multi-task RL’s optimization conflict.
- DiffusionOPD avoids cascade RL’s inefficiency and forgetting through its decoupled multi-task training approach.
- The framework extends OPD to diffusion Markov chains and yields a closed-form per-step reverse-KL objective.