Stepwise Credit Assignment for GRPO on Flow-Matching Models

📄 arXiv: 2603.28718v1 📥 PDF

作者: Yash Savani, Branislav Kveton, Yuchen Liu, Yilin Wang, Jing Shi, Subhojyoti Mukherjee, Nikos Vlassis, Krishna Kumar Singh

分类: cs.LG, cs.AI, cs.CV

发布日期: 2026-03-30

备注: Accepted to the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2026 Project page: https://stepwiseflowgrpo.com

💡 一句话要点

提出Stepwise-Flow-GRPO，为Flow模型生成过程中的每一步骤分配合适的奖励。

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

关键词: Flow模型 强化学习 信用分配 扩散模型 图像生成

📋 核心要点

  1. Flow-GRPO在Flow模型上应用强化学习，但其均匀信用分配忽略了扩散过程的时序性，导致次优中间步骤可能被错误奖励。
  2. Stepwise-Flow-GRPO基于每一步骤的奖励改进分配信用，利用Tweedie公式估计中间奖励，并引入基于增益的优势函数。
  3. 实验结果表明，Stepwise-Flow-GRPO具有更高的样本效率和更快的收敛速度，并通过DDIM启发的SDE提高了奖励质量。

📝 摘要（中文）

Flow-GRPO成功地将强化学习应用于Flow模型，但它对所有步骤采用统一的信用分配。这忽略了扩散生成的时序结构：早期步骤决定构图和内容（低频结构），而后期步骤解决细节和纹理（高频细节）。此外，仅基于最终图像分配统一的信用可能会无意中奖励次优的中间步骤，尤其是在扩散轨迹中稍后纠正错误时。我们提出了Stepwise-Flow-GRPO，它基于每个步骤的奖励改进来分配信用。通过利用Tweedie公式获得中间奖励估计并引入基于增益的优势函数，我们的方法实现了卓越的样本效率和更快的收敛速度。我们还引入了一种受DDIM启发的SDE，该SDE提高了奖励质量，同时保留了策略梯度的随机性。

🔬 方法详解

问题定义：Flow-GRPO在训练Flow模型时，对扩散过程中的每一步骤都给予相同的奖励权重，这与扩散过程的实际情况不符。扩散过程早期步骤主要负责生成图像的整体结构和内容，而后期步骤则负责完善细节和纹理。这种均匀的奖励分配方式可能导致早期步骤的次优操作被后期步骤的修正所掩盖，从而影响模型的训练效果。现有方法无法有效区分不同步骤对最终生成结果的贡献，导致训练效率低下。

核心思路：Stepwise-Flow-GRPO的核心思路是根据每个步骤对最终奖励的贡献程度来分配信用。具体来说，它不再对所有步骤采用统一的奖励，而是根据每一步骤的奖励改进情况来动态调整奖励权重。通过这种方式，模型可以更准确地评估每个步骤的优劣，从而更好地学习生成策略。

技术框架：Stepwise-Flow-GRPO的整体框架包括以下几个主要模块：1) Flow模型：用于生成图像的扩散模型。2) 奖励函数：用于评估生成图像质量的函数。3) 信用分配模块：根据每一步骤的奖励改进情况来分配信用。4) 策略梯度优化器：用于更新Flow模型的参数，使其能够生成更高质量的图像。该框架通过强化学习的方式，不断优化Flow模型的生成策略，使其能够生成更符合要求的图像。

关键创新：Stepwise-Flow-GRPO的关键创新在于其动态信用分配机制。与传统的均匀信用分配方法不同，Stepwise-Flow-GRPO能够根据每一步骤的实际贡献来分配奖励，从而更准确地评估每个步骤的优劣。此外，该方法还引入了Tweedie公式来估计中间奖励，并采用了基于增益的优势函数，进一步提高了训练效率和模型性能。

关键设计：Stepwise-Flow-GRPO的关键设计包括：1) 使用Tweedie公式来估计中间奖励，从而更准确地评估每个步骤的贡献。2) 引入基于增益的优势函数，以减少方差并提高训练稳定性。3) 采用DDIM启发的SDE，以提高奖励质量，同时保留策略梯度的随机性。这些设计共同作用，使得Stepwise-Flow-GRPO能够更有效地训练Flow模型，并生成更高质量的图像。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

Stepwise-Flow-GRPO通过动态信用分配机制，实现了更高的样本效率和更快的收敛速度。实验结果表明，该方法在图像生成任务上取得了显著的性能提升，能够生成更高质量、更逼真的图像。与传统的Flow-GRPO相比，Stepwise-Flow-GRPO能够更快地达到相同的性能水平，并且能够生成更高质量的图像。

🎯 应用场景

Stepwise-Flow-GRPO可应用于图像生成、图像编辑、图像修复等领域。通过更有效地训练Flow模型，该方法可以生成更高质量、更逼真的图像，从而提升相关应用的用户体验。此外，该方法还可以应用于其他类型的生成模型，例如文本生成、音频生成等，具有广泛的应用前景。

📄 摘要（原文）

Flow-GRPO successfully applies reinforcement learning to flow models, but uses uniform credit assignment across all steps. This ignores the temporal structure of diffusion generation: early steps determine composition and content (low-frequency structure), while late steps resolve details and textures (high-frequency details). Moreover, assigning uniform credit based solely on the final image can inadvertently reward suboptimal intermediate steps, especially when errors are corrected later in the diffusion trajectory. We propose Stepwise-Flow-GRPO, which assigns credit based on each step's reward improvement. By leveraging Tweedie's formula to obtain intermediate reward estimates and introducing gain-based advantages, our method achieves superior sample efficiency and faster convergence. We also introduce a DDIM-inspired SDE that improves reward quality while preserving stochasticity for policy gradients.