HyperAIHyperAI

Command Palette

Search for a command to run...

扩散模型中的在线策略自蒸馏

摘要

强化学习能够使扩散模型与人类偏好及任务特定目标对齐,但终点奖励并未指明中间去噪预测应如何变化。我们提出 DiffusionOPSD,一种在线策略自蒸馏框架,将图像级奖励引导转化为采样查询处干净输出预测的显式目标。在每次外循环迭代中,一个冻结的行为策略生成轨迹并提供查询状态与锚点。奖励梯度围绕每个锚点构建有界的正负目标。可训练策略通过有限次拟合将这些目标作为解耦监督进行拟合,随后通过指数移动平均更新刷新行为策略。这一设计使我们能够分别衡量目标构建与有限实现。受控的同查询实验表明,更大的目标构建增益并不一定在单次拟合更新后转化为更大的实现增益。在 SD3.5-M 和步蒸馏的 Z-Image-Turbo 上,我们的方法在两种骨干网络和十个评估器的 20 个奖励匹配设置中,有 19 个取得了最佳的最终留出分数。它比最强的竞争方法性能提升最高达 44.0%,并在 SD3.5-M 上将训练 GPU 小时数相对于 DiffusionNFT 减少了 40%,在 Z-Image-Turbo 上减少了 63%。这些结果支持在线策略自蒸馏作为一种高效且可分析的扩散后训练方法,通过将图像级奖励引导转化为显式且持续刷新的中间监督,从而为更高效、可诊断的对齐开辟了路径。

一句话总结

来自 ByteDance Seed、新加坡国立大学、加州大学圣迭戈分校及其他机构的研究人员提出了 DiffusionOPSD,这是一种 on-policy 自蒸馏方法,将图像级奖励梯度转换为有界的正、负干净输出目标,通过有限拟合和 EMA 策略更新拟合这些目标,并在 SD3.5-M 和 Z-Image-Turbo 上的 20 个奖励匹配设置中的 19 个取得最佳 held-out 分数,同时将训练 GPU 小时数最多减少 63%63\%63%

核心贡献

  • DiffusionOPSD 是一种 on-policy 自蒸馏框架,将图像级奖励引导转换为采样去噪查询处的显式、有界正、负干净输出目标,将其作为分离式监督信号进行拟合,并随着行为策略演化而刷新目标。
  • 该框架将目标构造与模型拟合分开,受控的同查询实验表明,在单次拟合更新后,更大的目标构造收益未必转化为更大的实际收益。
  • 在 SD3.5-M 和 Z-Image-Turbo 上,DiffusionOPSD 在十个评估器的 20 个奖励匹配设置中的 19 个取得最佳最终 held-out 分数,相对最强竞争方法最多提升 44.0%,训练 GPU 小时数相对 DiffusionNFT 在 SD3.5-M 上减少 40%、在 Z-Image-Turbo 上减少 63%,并在联合多奖励训练中在全部三个优化奖励上均优于 DiffusionNFT。

引言

扩散模型和流模型广泛用于高保真图像生成,强化学习越来越多地被用于使它们与人类偏好、美学和任务特定奖励对齐。核心挑战在于,奖励仅在最终解码图像之后才会被观测到,而扩散策略通过许多中间去噪预测起作用,因此结果级反馈必须转换为有用的局部监督。先前方法依赖采样轨迹信用分配,在不解码完整 rollout 的情况下通过单个后期干净输出预测进行反向传播,或对最终端点重新加权,同时使期望的中间变化保持隐式。作者提出了 DiffusionOPSD,这是一种 on-policy 自蒸馏方法,在行为策略查询处从图像级奖励梯度构造有界正、负目标,将这些目标作为分离式中间监督进行拟合,并随着行为策略演化而重新构造这些目标。这将目标构造与有限模型拟合分开,并使两个阶段都能被独立衡量。

方法

作者提出了 DiffusionOPSD,这是一个将扩散奖励优化视为 on-policy 自蒸馏过程的框架。通过将图像级奖励梯度转换为用于中间预测的分离式目标,该方法在行为策略演化时不断重新构造这些目标。

有关架构和训练流程的全面概览,请参见下方的框架图。

该过程首先建立干净输出坐标系,用于解码和评估预测。设 vθv_{\theta}vθ 为可训练的 rectified-flow 速度场,voldv_{\mathrm{old}}vold 表示行为策略的速度场。对于给定提示 ccc、带噪隐变量 zσz_{\sigma}zσ 和噪声水平 σ\sigmaσ,查询元组定义为 s=(c,zσ,σ)s = (c, z_{\sigma}, \sigma)s=(c,zσ,σ)。在 rectified-flow 路径下,与速度预测对应的干净输出预测 yθ(s)y_{\theta}(s)yθ(s) 的计算方式为:

yθ(s)=zσσvθ(s)y_{\theta}(s) = z_{\sigma} - \sigma v_{\theta}(s)yθ(s)=zσσvθ(s)

该预测通过局部奖励 R~(y,c)=R(D(y),c)\widetilde{R}(y, c) = R(D(y), c)R(y,c)=R(D(y),c) 进行评估,其中 DDD 为隐变量解码器。

在每个外层迭代中,冻结的行为策略收集轨迹,以提供查询状态和锚点。从轨迹中选择低噪声查询状态 zqz_qzq,行为锚点 y0y_0y0 计算为该查询处行为策略的干净输出预测。来自所收集轨迹的端点奖励确定一个按组归一化的拟合权重 ω\omegaω。该权重将每个奖励在其提示组内居中,同时使用完整 rollout 批次的全局标准差,确保所有提示组处于同一尺度。

为了指导优化,作者从行为锚点 y0y_0y0 出发构造有界正目标和负目标。作者应用稳定化的归一化奖励梯度步骤,迭代更新正、负目标。每次更新步骤都被投影回信任域球,以确保位移保持在半径 ρy02\rho \|y_0\|_2ρy02 以内。最终的分离式正目标 yˉ+\bar{y}_+yˉ+ 旨在提高奖励,而负目标 yˉ\bar{y}_-yˉ 则作为排斥性参考。

随后使用有限拟合过程优化可训练速度场。模型围绕锚点 y0y_0y0 形成正、负拟合分支,以编码对正目标的吸引和对负目标的排斥。采用分离式自适应归一化器来缩放各分支损失。总体目标 LOPSD\mathcal{L}_{\mathrm{OPSD}}LOPSD 定义为两个分支平方拟合误差的加权和,其中权重来自先前计算的拟合权重 ω\omegaω。关键的是,查询、锚点、权重和目标在有限拟合阶段保持固定,从而将目标构造与参数更新完全分离。

最后,通过刷新行为策略完成在线训练循环。在对可训练策略应用有限次优化器更新后,行为策略通过指数移动平均进行更新。更新后的行为策略随后为下一外层迭代生成新轨迹,提供新的查询状态和端点奖励。这种轨迹收集、目标构造、有限拟合和行为策略刷新的持续循环构成了系统自蒸馏循环的核心。

实验

实验在 SD3.5-M 和步数蒸馏的 Z-Image-Turbo 上评估 DiffusionOPSD,并与 FlowGRPO、ReFL 和 DiffusionNFT 进行对比,衡量各奖励的 held-out 质量、训练成本和优化稳定性。DiffusionOPSD 在 20 个奖励匹配设置中的 19 个取得最佳最终分数,训练成本降低 40% 至 63%,并将奖励增益持续到最终更新,人类标注者对其输出的偏好超过所有基线。消融实验确认,目标构造中的奖励梯度方向至关重要,而大多数实现选择和查询状态来源影响较小,并且该方法在原生少步采样器和多奖励训练中仍然有效。

DiffusionOPSD 在标准骨干网络和原生少步蒸馏模型上的 20 个奖励匹配设置中的 19 个取得最佳 held-out 评估器分数,在 HPSv3 和 VLM-Pairwise 上的增益最高达到 44%。它将改进保持到最终训练更新,保留了运行内奖励增益的 98%,而竞争方法通常趋于平台期或下降。单一联合训练策略在三个奖励上保持接近专家的质量,比联合基线高 8% 到 14%,并达到或超过单独专家。DiffusionOPSD 在 SD3.5-M 和 Z-Image-Turbo 上的 20 个奖励匹配比较中的 19 个领先,在 HPSv3 和 VLM-Pairwise 上的相对增益最高达到 44%。在 Z-Image-Turbo 上,DiffusionOPSD 在 ImageReward 上比最强基线高出 30.7%,在 VLM-Pairwise 上高出 14.6%,而 DiffusionNFT 在十个目标中的八个上低于未适配模型。DiffusionOPSD 在最终更新时保留了其运行内奖励改进的 98%,相比之下 FlowGRPO 为 83%,DiffusionNFT 为 90%。在三个奖励上训练的单一 DiffusionOPSD 策略达到 PickScore 专家增益的 113%,在统计上与 HPSv2.1 专家相当,并在全部三个指标上超过联合 DiffusionNFT 策略 8% 到 14%。

DiffusionOPSD 在奖励匹配的图像生成任务上使用标准骨干网络和原生少步蒸馏模型进行评估,比较 held-out 评估器分数和训练稳定性。它在几乎所有设置中取得领先性能,并将改进持续到最终训练更新,而竞争方法通常趋于平台期或下降。单一联合训练策略在多个奖励上保持接近专家的质量,优于联合基线,并达到或超过单独专家。


用 AI 构建 AI

从创意到上线——通过免费 AI 协同编码、开箱即用的环境和最优惠的 GPU 价格,加速您的 AI 开发。

AI 协同编码
开箱即用的 GPU
最优定价

HyperAI Newsletters

订阅我们的最新资讯
我们会在北京时间 每周一的上午九点 向您的邮箱投递本周内的最新更新
邮件发送服务由 MailChimp 提供