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

TL;DR DiffusionOPSD 用在线自蒸馏把图像级奖励显式转为中间去噪目标，在 SD3.5-M 与 Z-Image-Turbo 上 20 项奖励匹配中 19 项最优，且训练 GPU 小时最高省 63%。

问题

问题背景

当前 diffusion model 后训练中，使用 RL 对齐人类偏好或任务目标成为研究热点，通常仅依赖图像级别奖励信号。

现有方法局限

• 端点奖励稀疏：奖励只针对最终生成图像计算，无法直接指导中间去噪步骤如何调整。基于策略梯度的 DDPO、DPOK 等方法虽能估计奖励梯度，但梯度方差高，且不能解析中间隐变量的影响。
• 目标构造与策略拟合耦合：一些方法（如 DiffusionNFT）将奖励梯度直接作为目标，但训练动态复杂，难以分离目标构造误差和拟合误差，导致调参困难、GPU 消耗大。
• 例如，DiffusionNFT 在 SD 3.5-M 上训练 GPU 小时较高，限制了快速迭代和故障诊断。

为什么难/重要

• 扩散模型采样是多步去噪过程，中间隐变量高维且相互依赖，将稀疏终端奖励传播到全部时间步是技术挑战。
• 若中间目标不合理，会破坏模型原有生成分布；若完全忽略中间指导，则对齐效率低。
• 业界需要更高效、可诊断的后训练方案，因为大规模扩散模型训练成本高，差分可解释性对调试和迭代至关重要。

行业类比

类似 LLM RLHF 中从 outcome reward 到 process reward 的演进，把图像级奖励扩展为逐步去噪目标，可让扩散模型对齐过程更透明、更省算力。

核心洞察

• DiffusionOPSD 的核心创新是将图像级奖励信号转化为对采样查询点上干净输出预测的显式监督目标，从而把扩散模型的策略更新与目标构造解耦。与直接使用奖励梯度更新参数或仅优化最终图像的方法不同，它通过为每个查询构造有界的正负目标，让策略拟合这些 detached 目标，并分离目标构造与有限拟合的影响，使诊断和归因成为可能。这种设计让研究者能独立分析目标构造质量和拟合实现效率，为后续优化提供了可解释的切入点。
• 实验揭示了一个反直觉现象：更大的目标构造增益不一定转化为单次有限拟合后的实际增益。这意味着在 on-policy 自蒸馏框架中，目标构造的改进可能被拟合过程的瓶颈（如优化步数、学习率、目标一致性）所抵消。这提示后续工作不应只追求更精确的目标，还需关注拟合效率与目标稳定性，为优化资源分配提供了新视角。
• DiffusionOPSD 在 SD 3.5-M 和 Z-Image-Turbo 两个骨干上，与十个评估器配合，在 20 个奖励匹配设置中有 19 个取得最佳 held-out 分数，且相对最强竞争方法提升最高达 44.0%，同时训练 GPU 小时比 DiffusionNFT 减少 40%-63%。这表明通过显式中间监督的自蒸馏路径，能同时兼顾对齐效果与计算效率，为扩散模型后训练提供了一种可扩展的替代方案。

方法

输入与策略初始化

输入为预训练扩散模型作为可训练策略、冻结的 behavior policy 用于采样、图像级奖励模型和提示词。行为策略生成去噪轨迹，在随机时间步 t 采样 query states，并记录对应的 clean image prediction 作为 anchor。

关键模块：目标构造与有限拟合

对每个 anchor，利用奖励模型计算 reward gradient，构造有界的 positive target（沿梯度方向）和 negative target（反向），并 detach 作为监督标签。可训练策略在相同 query states 上以这些 targets 为回归目标执行有限步拟合，损失可选用像素级或特征级距离。

策略刷新与外循环

每次外迭代完成后，通过 EMA 更新 behavior policy，使下一轮采样分布接近当前策略，保持 on-policy 特性。框架显式区分 target construction 与 target realization，便于归因分析。

输出

输出为对齐后的扩散模型策略，可直接用于推理。

差异点

与 DiffusionNFT 等直接利用 reward gradient 更新参数的方法不同，DiffusionOPSD 将奖励信号转化为持续刷新的显式中间监督目标，训练更高效且可诊断。

实验

实验设计

在两个扩散骨干 SD 3.5-M 和 Z-Image-Turbo 上，与 DiffusionNFT 等基线在 reward-matched 设置下对比。外层迭代中，冻结的 behavior policy 生成轨迹并提供 query states 与 anchors；reward 梯度 围绕 anchor 构造有界正负目标；可训练策略通过 finite fitting 拟合这些 detached 监督，再用 EMA 更新 behavior policy。评估覆盖 10 个 evaluators，分离测量 target construction 与 finite realization。

关键发现

• 在 20 个 reward-matched 设置中，19 个 取得最佳最终 held-out 分数。
• 相对最强竞争方法最高提升 44.0%。
• 训练 GPU 小时相对 DiffusionNFT 减少 40%（SD 3.5-M）和 63%（Z-Image-Turbo）。
• 同 query 实验显示：larger target-construction gains 不必然转化为单次 fitting 后更大的 realized gains。

与基线对比的深度解读

与依赖 endpoint reward 间接梯度的 DiffusionNFT 相比，DiffusionOPSD 将图像级奖励转换为显式中间监督，通过 on-policy self-distillation 持续刷新目标，提升对齐效率与可诊断性。分离 target construction 与 finite fitting 的贡献，揭示了单纯增大目标构建增益的局限性，提示工程优化应关注拟合过程的实际收益。代码与项目页：GitHub | Project Page

行业影响

落地场景

DiffusionOPSD 适用于需要将扩散模型对齐到特定奖励信号的场景，例如：

• 电商商品图生成：以点击率、转化率或品牌视觉规范作为奖励，自动优化模特图与场景图，减少人工修图成本。
• 内容平台个性化封面生成：基于用户互动数据（点赞、停留时长）持续调整封面风格，提升内容分发效率。
• 广告创意生成：根据 A/B 测试反馈迭代优化素材，使其更符合目标受众审美。

商业价值

主要价值在降本与增效：

• 相比 DiffusionNFT，训练 GPU 小时减少 40%（SD 3.5-M）和 63%（Z-Image-Turbo），直接降低算力开支。
• 引入显式中间监督取代纯端点奖励，训练过程更可控、可诊断，减少调参成本。
• 在 19/20 个奖励匹配设置中取得最佳分数，质量提升更稳定，可转化为更高用户满意度和内容转化率。

与现有工作流集成

该方法可作为即插即用的后训练模块：

1. 准备奖励模型（如美学评分、CLIP 相似度）和未标注图像数据集。
2. 使用冻结的 行为策略 生成轨迹，计算奖励梯度并构造正负目标。
3. 对可训练策略进行 有限拟合（有限步监督学习），再用 EMA 刷新行为策略。

实现上可与 Diffusers、TRL 等库结合，目标构造与策略拟合解耦，无需复杂在线 RL 基础设施，便于在标准 GPU 集群上批量运行。建议在现有 LoRA 或 DreamBooth 微调之后接入，作为最终对齐步骤。项目页：DiffusionOPSD

局限

• **依赖可微奖励函数**：DiffusionOPSD 在目标构造阶段使用 reward gradients 生成正负 target，这要求奖励模型对中间 denoising prediction 可微分。对于黑盒奖励（如人类偏好 API、不可微的感知质量评估器），方法难以直接应用，需要额外近似或替代梯度估计，增加工程复杂度。与基于 DPO 的 pairwise 方法相比，后者无需显式梯度，对奖励形式更鲁棒。
• **模型与任务覆盖范围有限**：实验仅在 **SD 3.5-M** 与 **Z-Image-Turbo** 两个骨干上验证，虽然包含 10 个 evaluator 和 20 个 reward-matched settings，但未扩展到 **SDXL**、**Flux** 等主流模型，也未涉及视频扩散或 3D 生成任务。此外，评估依赖自动化指标，缺少真实人类偏好测试，存在 reward hacking 风险，泛化性和主观质量有待进一步确认。
• **超参数敏感性与实现复杂度**：方法引入多个新超参数，如 **target construction** 的边界范围、**finite fitting** 的迭代次数、EMA 系数等。虽然 controlled experiments 分离了目标构造与有限拟合的增益，但论文未系统分析超参数鲁棒性。实际部署时，为达到最优效果可能需要仔细调参，相比 DiffusionNFT 虽然训练时间减少，但调参复杂度可能部分抵消效率优势。