题目:On-Policy Self-Distillation in Diffusion Models
论文地址:https://arxiv.org/abs/2608.24646
代码地址:https://diffusionopsd.github.io
创新点
•提出面向扩散模型偏好对齐的 DiffusionOPSD 在线自蒸馏范式,突破传统强化学习仅能在生成终点提供奖励的局限,将图像层面的奖励信号转化为扩散中间去噪步骤上干净输出预测的显式监督目标,无需依赖独立外部教师模型与价值网络arXiv。
•设计目标构建与模型参数更新解耦的迭代训练流程,使用冻结的行为策略生成采样轨迹获取锚点,依靠奖励梯度在锚点周边构建有界正负优化目标,将目标作为梯度断开的监督信号供可训练策略拟合,每轮更新后借助 EMA 指数移动平均刷新行为策略,还可分别评估目标构造带来的理论增益与参数更新后的实际落地增益,便于分析训练过程。
方法
本文将扩散模型的偏好对齐任务转化为在线自蒸馏优化过程,不再只依靠图像最终输出的奖励信号,而是把图像层面的奖励转化为扩散采样中间去噪步骤的显式监督目标,训练循环中先利用冻结的行为策略生成完整采样轨迹并得到基准锚点,借助奖励梯度在锚点附近构造有界正向目标与负向排斥目标,将这些目标设置为梯度断开的独立监督信号用来更新可训练策略,每一轮参数更新之后使用 EMA 指数移动平均更新行为策略以继续迭代生成新轨迹,整个流程把目标构建环节和模型参数更新环节解耦,能够分开评估目标本身带来的奖励增益和模型更新后的实际性能增益,该框架不需要额外的教师模型或价值网络,可直接在模型自身生成的样本轨迹上完成偏好对齐,并且能够适配少步蒸馏的快速扩散主干模型,在保持模型原有高速生成能力的同时完成偏好对齐,降低训练阶段的算力开销。
训练集与测试集的质量收益曲线
本图包含左右两张子图,横轴代表归一化累计 GPU 训练算力,纵轴为奖励收益,曲线对比 DiffusionOPSD 与 ReFL、FlowGRPO、DiffusionNFT 几种基线方法,阴影区域代表统计误差范围,帕累托前沿线用于表征最优性能边界,左子图展示训练集上的原生训练收益,右子图展示在未参与训练的留存测试集上的泛化收益,实验结果在十组奖励设置与两种主干模型上取平均,可以看出 DiffusionOPSD 的奖励收益上升速度最快,并且最终能够达到最高的收益水平,相比其他基线方法,在同等算力消耗下获得更高的模型对齐收益,同时在测试集上也保持更好的泛化表现。
奖励到策略的不同范式对比
本图依次展示了四种扩散模型奖励对齐范式,从左至右分别为轨迹信用分配、晚状态奖励反向传播、端点监督以及本文所提的显式中间目标方法,前三种传统方式大多只依靠最终生成图像的奖励信号,或是间接把终点奖励回传给中间去噪步骤,而 DiffusionOPSD 方法会在有限拟合之前,针对扩散采样的中间状态构建显式的正负目标,在每一个去噪步骤上直接建立监督信号,不再只依赖图像生成终点的奖励反馈,直观体现出本文方案和过往奖励对齐思路的核心差异。
DiffusionOPSD 方法整体框架图
本图完整展示 DiffusionOPSD 在线自蒸馏的整套流程,输入文本提示词送入冻结的行为策略,从高斯噪声开始生成扩散去噪轨迹并得到查询状态与行为锚点,在局部奖励曲面中,以锚点为中心划定有界邻域,沿着奖励上升与下降方向分别构建正向目标与负向目标,将查询状态送入可训练场网络预测输出,对目标执行梯度截断操作,借助有限拟合模块分别拟合正负残差,再通过加权组合构建损失函数用于参数更新,一轮训练结束后利用 EMA 指数移动平均更新行为策略,循环迭代持续优化模型,整套流程把奖励目标构造和策略参数拟合相互解耦,实现扩散模型中间去噪步骤的显式监督对齐。
实验
该表格汇总了 SD3.5-M 与 Z-Image-Turbo 两类主干模型在留存评测集上的各项评估指标,指标数值越高代表对齐效果越好,区分了奖励专属独立权重与通用权重两种设置,采用开源评估器与内部评估器多维度打分,同时纳入 SDXL、FLUX.1-dev 等基础模型以及 FlowGRPO、DiffusionNFT、ReFL 等多种对齐基线方法进行对照,加粗和下划线标记同组内最优与次优结果,部分两行基线采用两阶段在线蒸馏方案,表格展示不同更新步数下各方法的性能,DiffusionOPSD 在多种主干模型与评测指标上均取得优异结果,验证了该在线自蒸馏方案在扩散模型偏好对齐任务上的有效性。

