Aligning Few-Step Diffusion Models with Dense Reward Difference Learning

📄 arXiv: 2411.11727v1 📥 PDF

作者: Ziyi Zhang, Li Shen, Sen Zhang, Deheng Ye, Yong Luo, Miaojing Shi, Bo Du, Dacheng Tao

分类: cs.LG, cs.CV

发布日期: 2024-11-18

🔗 代码/项目: GITHUB

💡 一句话要点

提出SDPO，通过密集奖励差异学习对齐少步扩散模型，提升步泛化能力

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 扩散模型 强化学习 奖励学习 步泛化 少步扩散模型

📋 核心要点

  1. 现有方法在对齐少步扩散模型时，依赖最终步的稀疏奖励，导致步泛化能力差，不同去噪步长下性能不稳定。
  2. SDPO通过引入每步的密集奖励反馈，学习配对样本间的奖励差异，实现少步扩散模型的步进式优化，提升对齐一致性。
  3. SDPO采用在线强化学习框架，有效利用密集奖励的步进式粒度，实验表明其在不同步长配置下优于现有方法。

📝 摘要（中文）

将扩散模型与下游目标对齐对于其实际应用至关重要。然而，标准对齐方法在直接应用于少步扩散模型时，常常难以进行步泛化，导致不同去噪步长场景下的性能不一致。为了解决这个问题，我们提出了一种新颖的对齐方法，即步进式扩散策略优化（SDPO），专门为少步扩散模型定制。与先前依赖于每个去噪轨迹的最后一步的单一稀疏奖励进行轨迹级优化的方法不同，SDPO在每个中间步骤中都加入了密集的奖励反馈。通过学习配对样本之间密集奖励的差异，SDPO促进了少步扩散模型的步进式优化，确保了所有去噪步骤之间的一致对齐。为了促进稳定和高效的训练，SDPO引入了一个在线强化学习框架，该框架具有几个新颖的策略，旨在有效地利用密集奖励的步进式粒度。实验结果表明，SDPO在基于奖励的对齐方面始终优于先前的各种步长配置方法，突显了其强大的步泛化能力。代码可在https://github.com/ZiyiZhang27/sdpo 获取。

🔬 方法详解

问题定义：现有的扩散模型对齐方法，特别是应用于少步扩散模型时，通常依赖于最终步骤的稀疏奖励信号进行优化。这种方法忽略了中间步骤的信息，导致模型在不同去噪步骤上的性能不一致，即步泛化能力差。因此，如何有效地利用少步扩散模型中的每一步信息，实现更稳定和高效的对齐，是一个亟待解决的问题。

核心思路：SDPO的核心思路是引入密集奖励差异学习，即在每个去噪步骤中都提供奖励反馈，并学习配对样本之间奖励的差异。通过这种方式，模型可以逐步优化，而不仅仅依赖于最终结果。这种设计使得模型能够更好地理解每一步骤对最终结果的影响，从而提高步泛化能力。

技术框架：SDPO采用在线强化学习框架，包含以下主要模块：1) 扩散模型作为环境；2) 策略网络，用于生成去噪步骤；3) 奖励函数，用于评估每个步骤的质量；4) 优化器，用于更新策略网络。该框架通过不断与环境交互，收集经验数据，并利用这些数据来改进策略网络。

关键创新：SDPO的关键创新在于引入了密集奖励差异学习，并将其与在线强化学习框架相结合。与传统的稀疏奖励方法相比，SDPO能够更有效地利用少步扩散模型中的信息，从而提高对齐性能和步泛化能力。此外，该框架还设计了专门的策略来有效利用密集奖励的步进式粒度。

关键设计：SDPO的关键设计包括：1) 奖励函数的选择，需要能够准确反映每个步骤的质量；2) 策略网络的结构，需要能够有效地生成去噪步骤；3) 优化算法的选择，需要能够稳定和高效地更新策略网络。此外，如何有效地利用在线强化学习框架中的探索-利用平衡也是一个重要的设计考虑。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，SDPO在基于奖励的对齐任务中，相比于现有方法，在各种步长配置下均取得了显著的性能提升，验证了其强大的步泛化能力。具体性能数据未在摘要中给出，但强调了SDPO在不同步数设置下的优越性，表明其对步数具有鲁棒性。

🎯 应用场景

SDPO方法可应用于图像生成、文本生成等领域，尤其适用于对生成速度有较高要求的场景。通过提升少步扩散模型的性能，SDPO能够加速生成过程，同时保持生成质量。该研究对于推动扩散模型在实际应用中的普及具有重要意义，例如在实时图像编辑、快速原型设计等方面具有潜在的应用价值。

📄 摘要（原文）

Aligning diffusion models with downstream objectives is essential for their practical applications. However, standard alignment methods often struggle with step generalization when directly applied to few-step diffusion models, leading to inconsistent performance across different denoising step scenarios. To address this, we introduce Stepwise Diffusion Policy Optimization (SDPO), a novel alignment method tailored for few-step diffusion models. Unlike prior approaches that rely on a single sparse reward from only the final step of each denoising trajectory for trajectory-level optimization, SDPO incorporates dense reward feedback at every intermediate step. By learning the differences in dense rewards between paired samples, SDPO facilitates stepwise optimization of few-step diffusion models, ensuring consistent alignment across all denoising steps. To promote stable and efficient training, SDPO introduces an online reinforcement learning framework featuring several novel strategies designed to effectively exploit the stepwise granularity of dense rewards. Experimental results demonstrate that SDPO consistently outperforms prior methods in reward-based alignment across diverse step configurations, underscoring its robust step generalization capabilities. Code is avaliable at https://github.com/ZiyiZhang27/sdpo.