On-Policy Self-Distillation in Diffusion Models
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.CV
2026-08-25
DiffusionOPSD 用行为策略轨迹上的奖励梯度，构造有界的正负干净图目标，再当蒸馏监督去拟合。20 组设置里 19 组持有分数最高，相对最强对照最高高 44.0%，训练 GPU 时比 DiffusionNFT 少 40% 和 63%。
扩散和 flow 模型的奖励通常只出现在整张图生成完之后。策略真正动作的地方是中间每一步的去噪预测。轨迹级策略梯度用终点分数给整条路径打分，ReFL 一类方法把奖励反传过某一步干净图预测，DiffusionNFT 则把奖励选中的终点当成回归目标。中间这一步「该往哪改」始终是隐含的。
更麻烦的是，构造一个更好看的局部目标，和模型在有限步更新里能不能把这个目标学进去，不是一回事。同一个网络要服务很多噪声水平，拟合一个查询会牵动别的查询。把两段混在一个损失里，就没法判断失败出在目标还是出在拟合。
DiffusionOPSD 把后训练写成在线的 on-policy 自蒸馏。每个外循环里，冻结的行为策略采样轨迹，给出查询状态和干净图锚点。奖励对锚点求导，在锚点附近构造有界的正目标（沿奖励上升）和负目标（反向推开）。可训练策略用 stop-gradient 后的目标做有限步拟合，再靠 EMA 刷新行为策略，下一轮重新采样、重新造目标。
干净图预测来自 rectified-flow 速度场：给定噪声潜变量和噪声水平，速度可以唯一映射成一张可解码的干净图，于是图像级奖励可以直接标在中间预测上。正负分支的相对权重由组内归一化的终点奖励控制。因为造目标和改参数被切开，同一查询、同一后续噪声可以分别量「目标构造增益」和「拟合后实现增益」。
对照很直接。轨迹信用分配不造局部目标；直接反传要保住奖励计算图；终点监督不告诉当前中间预测怎么动。这里的目标跟行为策略走，行为策略一变，监督就重建。
主实验在 SD3.5-M 和步数蒸馏的 Z-Image-Turbo 上，十个评测器、奖励对齐的设置共 20 组，持有分数 19 组第一。相对最强对照的最大相对提升是 44.0%，对应 SD3.5-M 上奖励专用训练的 Pair 指标：0.465 对 DiffusionNFT 的 0.323。
共享检查点、300 次更新的 SD3.5-M 上，PickScore 为 25.51，高于同设置 DiffusionNFT 的 23.62 和 FlowGRPO 更长训练后的 23.50。奖励专用、100 次更新时，Z-Image-Turbo 的 PickScore 从基线 22.86 到 25.15，HPSv2.1 从 0.296 到 0.390。联合训练 PickScore、CLIPScore、HPSv2.1 的单一策略，三项都超过对应的 DiffusionNFT 策略。
| 骨干 | 方法 | 每 100 次更新 GPU 时 | 相对 NFT |
| SD3.5-M | DiffusionNFT | 47.2 | 1.00× |
| SD3.5-M | DiffusionOPSD | 28.2 | 0.60× |
| Z-Image-Turbo | DiffusionNFT | 405.8 | 1.00× |
| Z-Image-Turbo | DiffusionOPSD | 149.8 | 0.37× |
省时的关键是反向传播次数：NFT 在 SD3.5-M 上每个更新约 144 次扩散反向，OPSD 约 16 次。同查询对照里，构造增益更大的目标，拟合一步后实现增益可以更小，所以两段必须分开量。消融里，把奖励梯度方向换成随机、空操作或 rollout 残差，CLIPScore 从 0.3117 掉到 0.2311、0.2280、0.1456；只换查询来源到前向加噪，几乎不动（0.3103）。去掉负分支也几乎不掉（0.3137）。查询噪声到 0.90 或目标半径小到 0.02 才会明显变差。
扩散对齐的工程问题一直是：终点奖励很贵，中间步又没有标签。这篇给出一种可诊断的折中：用奖励梯度造局部目标，当蒸馏来拟合，并且每轮重造。若你已经在跑 DiffusionNFT 或 ReFL，值得看 GPU 时和持有分数能否一起动。
它没有宣称取代所有轨迹 RL。它宣称的是，中间监督可以显式化，而且造目标和学目标不该混成一笔糊涂账。
44.0% 是相对增益里的峰值，不是十个评测器的平均。主表大量数字来自单次训练曲线，不是多随机种子均值。同查询实验说明拟合缺口真实存在，但论文没有给出在完整后训练里自动闭合这个缺口的方法。评测奖励和解码器本身可被过拟合。Z-Image-Turbo 上 OPSD 的峰值显存到 61.5 GB，高于 NFT 的 49.9 GB，省的是时间不是卡数。筛选消融只用 20 条 DrawBench 提示、100 张图。