Stepwise Credit Assignment for GRPO on Flow-Matching Models

Yash Savani; Branislav Kveton; Yuchen Liu; Yilin Wang; Jing Shi; Subhojyoti Mukherjee; Nikos Vlassis; Krishna Kumar Singh

cs.LGMar 30, 2026

Stepwise Credit Assignment for GRPO on Flow-Matching Models

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

Organizations: Carnegie Mellon University · Adobe Research

Abstract

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.

Figures & tables

Figure 2: Gain magnitudes across steps. Mean absolute gain

\E_{i}[|g_{t}^{i}|]

measured on 256 GenEval prompts using PickScore. Early steps show larger gains, indicating that compositional decisions drive most reward improvement.

Panel of Figure 2 — Algorithm 1 Stepwise-Flow-GRPO

Figure 3: Qualitative results . We compare our Stepwise-Flow-GRPO with Flow-GRPO and observe better spatial reasoning, attribute binding, and counting performance.

Panel of Figure 4 — Figure 4: Sample efficiency across reward functions. Stepwise-Flow-GRPO consistently outperforms Flow-GRPO in reward per training step across all settings, achieving both faster convergence and superior final performance in 3 out of 4 settings.

Panel of Figure 5 — Figure 5: Wall-clock efficiency matches sample efficiency gains. Reward versus wall-clock time for the same settings as Fig. 4 . Despite additional computational cost for intermediate denoising, Stepwise-Flow-GRPO converges faster in wall-clock time, achieving visibly superior performance in 3 out of 4 settings.

Panel of Figure 6 — Figure 6: Stepwise credit assignment remains effective with improved SDE. Reward versus training step when both methods use the DDIM-inspired SDE from Sec. 5.5 . Stepwise-Flow-GRPO retains its sample efficiency advantage, demonstrating that the improvements in credit assignment and sampling are complementary.

Pretrained Models
Experiment	Overall	Single Obj.	Two Objs.	Counting	Colors	Position	Attr. Binding
SD3.5-M (cfg=1.0)	0.28	0.71	0.23	0.15	0.45	0.05	0.08
SD3.5-M (cfg=4.5)	0.63	0.98	0.78	0.50	0.81	0.24	0.52
Fine-tuned Models
SD3.5-M+Flow-GRPO (cfg=1.0 PickScore)	0.60	0.96	0.73	0.67	0.67	0.21	0.35
SD3.5-M+Ours (cfg=1.0 PickScore)	0.60	0.96	0.75	0.67	0.67	0.21	0.34

Table 1: Final model quality on GenEval. Compositional generation performance for models trained with PickScore reward. Both methods substantially improve over the base model, with our method matching Flow-GRPO at cfg=1.0 and outperforming it across most categories at cfg=4.5, particularly in counting and spatial positioning.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: Design variation comparison. Reward vs. training iteration for different formulations of stepwise credit assignment on GenEval with PickScore reward. The standard gain formulation from the main paper matches all alternatives, demonstrating that preserving the natural temporal structure of diffusion gains is the most effective credit assignment.

Figure 8: Improved SDE produces cleaner images. Qualitative comparison between Flow-GRPO SDE (middle) and our DDIM-inspired SDE (right) from Sec. 5.5 . The improved formulation generates images that are much less noisy while maintaining stochasticity for policy gradients.

Figure 9: Joint normalization preserves temporal structure and accelerates convergence. Reward vs. training iteration comparing joint normalization (global mean/std across all steps and trajectories) against per-step normalization (separate mean/std for each step).

Figure 10: Stepwise-Flow-GRPO is robust to the number of denoising substeps. Reward vs. training iteration (top) and GPU Hours (bottom) for different numbers of substeps

T^{\prime}\in\{2,5,8\}

used to compute

\hat{x}_{0}(t)

. All settings achieve similar convergence speed and final performance, validating that our method is not sensitive to this hyperparameter choice.

Method (Reward)	Generation (s)	Reward eval (s)
Ours (PickScore)	24.4 $\pm$ 0.0	2.8 $\pm$ 0.0
Flow-GRPO (PickScore)	13.7 $\pm$ 0.9	1.4 $\pm$ 1.3
Ours (ImageReward)	29.8 $\pm$ 0.2	8.2 $\pm$ 0.2
Flow-GRPO (ImageReward)	7.9 $\pm$ 0.0	0.7 $\pm$ 0.0
Ours (UnifiedReward)	136.1 $\pm$ 5.8	114.6 $\pm$ 5.8
Flow-GRPO (UnifiedReward)	56.0 $\pm$ 30.9	48.8 $\pm$ 30.9

Appendix

Table 2: Per-iteration timing breakdown in seconds per training iteration, averaged over multiple runs.

Figure 11: OCR text rendering results. Reward vs. training iteration (top) and wall-clock time (bottom) using combined OCR+PickScore reward (80% OCR, 20% PickScore). Flow-GRPO diverges after 500 steps while Stepwise-Flow-GRPO continues improving, demonstrating superior stability and final performance on compositional text rendering.

Figure 12: Extended GenEval training. GenEval overall score vs. wall-clock time for 400 GPU hour runs. Stepwise-Flow-GRPO achieves 0.87, substantially outperforming Flow-GRPO (0.72) and approaching state-of-the-art autoregressive models. The widening performance gap demonstrates that stepwise credit assignment provides increasing benefits at high performance levels.

Figure 13: UnifiedReward training. Reward vs. wall-clock time for 60 GPU hour run on GenEval using UnifiedReward-7b-v1.5. Stepwise-Flow-GRPO trains stably while Flow-GRPO diverges with this reward function, demonstrating superior robustness with complex VLM-based rewards.

Pretrained Models
Model	Overall	Single Obj.	Two Objs.	Counting	Colors	Position	Attr. Binding
SD3.5-M (cfg=1.0)	0.28	0.71	0.23	0.15	0.45	0.05	0.08
SD3.5-M (cfg=4.5)	0.63	0.98	0.78	0.50	0.81	0.24	0.52
Standard Training Duration
Flow-GRPO (cfg=1.0, PickScore)	0.60	0.96	0.73	0.67	0.67	0.21	0.35
Ours (cfg=1.0, PickScore)	0.60	0.96	0.75	0.67	0.67	0.21	0.34

Appendix

Table 3: Complete GenEval results. After 400 GPU hours, our method achieves 0.87 overall, substantially outperforming Flow-GRPO (0.72) and approaching GPT-4o (0.84). Flow-GRPO extended results from [ 26 ] ; reference models from respective papers.

Figure 14: Qualitative comparison across training objectives. Generated images from GenEval prompts using base SD3.5-M (left), GenEval reward training (middle), and UnifiedReward training (right). While GenEval reward training improves prompt adherence and benchmark scores, UnifiedReward training produces higher overall visual quality and more photorealistic images.

Figure 15: Extended qualitative comparison (page 1). Each row shows a single prompt with the base model generation (step 0), followed by Flow-GRPO and Stepwise-Flow-GRPO (Ours) at training steps 60 and 120, both trained with PickScore reward. The gap is most visible at step 60, where our method demonstrates better counting, more plausible compositions, and higher image quality.

Figure 16: Extended qualitative comparison (page 2). Same layout as Fig. 15 . Stepwise-Flow-GRPO consistently produces more accurate object counts, better spatial arrangements, and higher overall quality than Flow-GRPO across training.

Figure 17: Extended qualitative comparison (page 3). Same layout as Fig. 15 . Our method maintains qualitative advantages throughout training, with the largest visible differences at step 60.

RL for Diffusion Models

Reinforcement learning for diffusion models adapts denoising and sampling to rewards, preferences, and task outcomes.

See all 179 papers

SGA-Flow-GRPO: Spatial Gradient-Guided Credit Assignment for Flow-GRPO
Sep 26, 2026Yunkai Yang, Yudong Zhang, Xinying Chen +5Diffusion Model Alignment RL for Diffusion Models
Sun Yat-Sen University · Baidu · Beijing Institute of Technology +2
Alleviating Sparse Rewards by Modeling Step-Wise and Long-Term Sampling Effects in Flow-Based GRPO
Feb 6, 2026Yunze Tong, Mushui Liu, Canyu Zhao +7Flow Matching Temporal Credit Assignment
1Zhejiang University, Hangzhou, China · 2Taobao & Tmall Group of Alibaba, Hangzhou, China · 3Tsinghua University, Beijing, China
$λ$ -Controlled GRPO: Turning Flow-Matching Ratio Instability into a Budgeted Resource
Sep 18, 2026Yufeng Wang, Parivesh Priye, Meeshawn Marathe +1Flow Matching T2I Generation
Stony Brook University · Georgia Institute of Technology · Rivian and Volkswagen Group Technologies

Sep 26, 2026J/K move · Enter open · S save

SGA-Flow-GRPO: Spatial Gradient-Guided Credit Assignment for Flow-GRPO

Yunkai Yang, Yudong Zhang, Xinying Chen, Bin Luo+4

Sun Yat-Sen University · Baidu · Beijing Institute of Technology+2

Diffusion Model Alignment RL for Diffusion Models Credit Assignment T2I Generation Diffusion Model Fine-Tuning Diffusion Transformer

Open paper AskPDF arXiv

Reinforcement Learning (RL) has proven effective in aligning flow-based generative models with human preferences. Recently, Flow-GRPO has emerged as an efficient critic-free paradigm by calculating advantages over sampled candidate trajectories. However, standard Flow-GRPO applies a uniform scalar advantage across both temporal denoising steps and spatial latent dimensions, without explicitly accounting for the spatial structure of generated images, which may lead to sub-optimal policy updates. To address this, we propose a novel gradient-guided spatial credit assignment framework tailored for Diffusion Transformers (DiTs). We first reformulate the transition-level log-likelihood in Flow-GRPO into a token-wise representation natively aligned with DiT patch architectures, constructing spatially fine-grained importance sampling ratios. To allocate localized credit without rigid, boundary-sensitive segmentation heuristics, we introduce a continuous spatial credit map derived from reward gradients. Crucially, we employ an outlier-robust normalization scheme based on Median Absolute Deviation (MAD) coupled with temperature scaling, effectively eliminating gradient noise while highlighting functional prompt-aligned regions. Extensive evaluations on GenEval show that our approach delivers SOTA alignment quality, achieving a convergence rate comparable to top-tier methods like DiffusionNFT while substantially improving upon Flow-GRPO-based methods in alignment performance.

Abstract

RL for Diffusion Models

Explore similar work

SGA-Flow-GRPO: Spatial Gradient-Guided Credit Assignment for Flow-GRPO

Alleviating Sparse Rewards by Modeling Step-Wise and Long-Term Sampling Effects in Flow-Based GRPO

λλλ-Controlled GRPO: Turning Flow-Matching Ratio Instability into a Budgeted Resource

$λ$ -Controlled GRPO: Turning Flow-Matching Ratio Instability into a Budgeted Resource