扩散模型中的在策略自蒸馏
On-Policy Self-Distillation in Diffusion Models
August 25, 2026
作者: Wei Zhou, Xiongwei Zhu, Lingdong Kong, Bo Chen, Lei Zhang, Yongyuan Liang, Xiaoxia Hou, Ye Tian, Xian Sun, Yingshuo Wang, Linfeng Li, Shengqiong Wu, Leigang Qu, Feng Li, Wei Liu, Julian McAuley, Tat-Seng Chua
cs.AI
摘要
强化学习可以将扩散模型与人类偏好及任务特定目标对齐,但端点奖励并未指明中间去噪预测应如何变化。我们提出 DiffusionOPSD,作为一种同策略自蒸馏框架,将图像级奖励引导转换为采样查询处干净输出预测的显式目标。在每次外部迭代中,冻结的行为策略生成轨迹,并提供查询状态与锚点。奖励梯度在每个锚点周围构造有界正负目标。可训练策略通过有限次拟合将这些目标作为解耦监督进行拟合,随后指数移动平均更新刷新行为策略。这一设置使我们能够分别衡量目标构建与有限实现。受控的同查询实验表明,在单次拟合更新后,较大的目标构建收益并不一定会转化为较大的实际收益。在 SD 3.5-M 和步进蒸馏的 Z-Image-Turbo 上,我们的方法在两个骨干网络和十个评估器所构成的 20 个奖励匹配设置中的 19 个中取得了最佳最终留出分数。相比最强竞争方法,其性能提升最高达 44.0%;相对于 DiffusionNFT,在 SD 3.5-M 上训练 GPU 小时数减少 40%,在 Z-Image-Turbo 上减少 63%。这些结果支持同策略自蒸馏作为扩散模型后训练的一种高效且可分析的方法,通过将图像级奖励引导转换为显式且持续刷新的中间监督,从而为更高效且可诊断的对齐开辟了道路。
English
Reinforcement learning can align diffusion models with human preferences and task-specific objectives, but endpoint rewards do not specify how an intermediate denoising prediction should change. We introduce DiffusionOPSD as an on-policy self-distillation framework that converts image-level reward guidance into explicit targets for clean-output predictions at sampled queries. At each outer iteration, a frozen behavior policy generates trajectories and supplies query states and anchors. Reward gradients construct bounded positive and negative targets around each anchor. The trainable policy fits these targets as detached supervision through finite fitting before an exponential moving average update refreshes the behavior policy. This setup lets us measure target construction and finite realization separately. Controlled same-query experiments show that larger target-construction gains do not necessarily translate into larger realized gains after a single fitting update. Across SD 3.5-M and the step-distilled Z-Image-Turbo, our approach achieves the best final held-out scores in 19 of 20 reward-matched settings across two backbones and ten evaluators. It outperforms the strongest competing method by up to 44.0% and reduces training GPU-hours relative to DiffusionNFT by 40% on SD 3.5-M and 63% on Z-Image-Turbo. These results support on-policy self-distillation as an efficient and analyzable approach to diffusion post-training by converting image-level reward guidance into explicit and continually refreshed intermediate supervision, thereby opening a path toward more efficient and diagnosable alignment.