Aligning Few-Step Diffusion Models with Dense Reward Difference Learning
作者: Ziyi Zhang, Li Shen, Sen Zhang, Deheng Ye, Yong Luo, Miaojing Shi, Dongjing Shan, Bo Du, Dacheng Tao
分类: cs.LG, cs.CV
发布日期: 2026-02-28
💡 一句话要点
提出SDPO，通过密集奖励差异学习对齐少步扩散模型与下游目标
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 少步扩散模型 强化学习 密集奖励学习 奖励差异学习 策略优化
📋 核心要点
- 现有强化学习方法在少步扩散模型中存在局限，难以实现与下游目标的有效对齐。
- SDPO通过双状态轨迹采样和密集奖励差异学习，为少步扩散模型提供更有效的强化学习框架。
- 实验证明，SDPO在多种任务和设置下，显著提升了少步扩散模型与奖励的对齐效果。
📝 摘要（中文）
少步扩散模型能够高效地合成高分辨率图像，但由于现有强化学习（RL）方法在低步长、有限状态空间和次优样本质量方面的局限性，难以与特定的下游目标对齐。为了解决这个问题，我们提出了一种为少步扩散模型量身定制的新型RL框架——步进式扩散策略优化（SDPO）。SDPO引入了一种双状态轨迹采样机制，跟踪每一步的噪声状态和预测的干净状态，以提供密集的奖励反馈，并实现低方差的混合步长优化。为了进一步提高效率，我们开发了一种基于潜在相似性的密集奖励预测策略，以最大限度地减少昂贵的密集奖励查询。利用这些密集奖励，SDPO优化了一个密集奖励差异学习目标，从而能够进行更频繁和更精细的策略更新。额外的改进，包括逐步优势估计、时间重要性加权和步进式梯度更新，进一步增强了长期依赖性、低步长优先级和梯度稳定性。我们的实验表明，SDPO在各种少步设置和任务中始终提供卓越的奖励对齐结果。
🔬 方法详解
问题定义：论文旨在解决少步扩散模型难以与特定下游目标对齐的问题。现有强化学习方法在少步扩散模型中面临挑战，包括：状态空间有限、样本质量不高、难以提供有效的奖励信号等，导致策略优化困难。
核心思路：论文的核心思路是利用密集奖励差异学习来指导少步扩散模型的训练。通过在每一步都提供奖励信号，并学习奖励的差异，可以更有效地进行策略优化，克服了传统强化学习方法在稀疏奖励环境下的不足。同时，通过双状态轨迹采样，可以更准确地估计奖励，降低方差。
技术框架：SDPO框架主要包含以下几个模块：1) 双状态轨迹采样：同时跟踪噪声状态和预测的干净状态。2) 密集奖励预测：利用潜在相似性来预测密集奖励，减少计算开销。3) 密集奖励差异学习：优化奖励差异，实现更精细的策略更新。4) 策略优化：利用强化学习算法（如PPO）更新扩散模型的参数。
关键创新：论文的关键创新在于：1) 提出了双状态轨迹采样机制，为密集奖励学习提供了更准确的状态信息。2) 开发了基于潜在相似性的密集奖励预测策略，降低了计算复杂度。3) 引入了密集奖励差异学习目标，实现了更有效的策略优化。
关键设计：在双状态轨迹采样中，同时保存噪声状态和去噪后的状态，用于计算奖励和优势函数。密集奖励预测器可以使用神经网络进行训练，输入为潜在空间中的状态表示，输出为预测的奖励值。密集奖励差异学习目标可以采用均方误差损失函数，鼓励模型学习奖励的差异。在策略优化过程中，采用了步进式优势估计、时间重要性加权和步进式梯度更新等技巧，以提高训练的稳定性和效率。
🖼️ 关键图片
📊 实验亮点
实验结果表明，SDPO在各种少步设置和任务中均取得了显著的性能提升。例如，在图像生成任务中，SDPO能够生成更高质量、与目标更对齐的图像，并且所需的步数更少。与基线方法相比，SDPO在奖励对齐方面取得了显著的进步，证明了其在少步扩散模型强化学习方面的有效性。
🎯 应用场景
该研究成果可应用于图像生成、图像编辑、视频生成等领域，尤其是在需要快速生成高质量图像的场景下。例如，可以用于游戏开发中快速生成游戏素材，或者在设计领域快速生成设计方案。此外，该方法还可以应用于机器人控制领域，通过学习奖励函数，使机器人能够更快地完成特定任务。
📄 摘要（原文）
Few-step diffusion models enable efficient high-resolution image synthesis but struggle to align with specific downstream objectives due to limitations of existing reinforcement learning (RL) methods in low-step regimes with limited state spaces and suboptimal sample quality. To address this, we propose Stepwise Diffusion Policy Optimization (SDPO), a novel RL framework tailored for few-step diffusion models. SDPO introduces a dual-state trajectory sampling mechanism, tracking both noisy and predicted clean states at each step to provide dense reward feedback and enable low-variance, mixed-step optimization. For further efficiency, we develop a latent similarity-based dense reward prediction strategy to minimize costly dense reward queries. Leveraging these dense rewards, SDPO optimizes a dense reward difference learning objective that enables more frequent and granular policy updates. Additional refinements, including stepwise advantage estimates, temporal importance weighting, and step-shuffled gradient updates, further enhance long-term dependency, low-step priority, and gradient stability. Our experiments demonstrate that SDPO consistently delivers superior reward-aligned results across diverse few-step settings and tasks. Code is available atthis https URL.