Stepwise Credit Assignment for GRPO on Flow-Matching Models

• Early denoising steps shape composition and layout, while late steps resolve details and textures, so uniform final-image credit ignores the trajectory’s temporal structure.A high-reward final image can also reinforce poor early decisions that later steps corrected.Source: §1 · pp. 1–2
• Stepwise-Flow-GRPO uses intermediate reward estimates and stepwise reward gains to credit each denoising step for its contribution.The method uses Tweedie’s formula to estimate clean images from intermediate noisy states and calculate rewards.Source: §1 · p. 2
• A DDIM-inspired SDE produces higher-quality, less noisy samples while matching the Flow-ODE Fokker-Planck marginal in the Taylor limit.The authors report that the combined approach improves sample efficiency and convergence over standard Flow-GRPO.Source: §1 · p. 2

• GRPO estimates and standardizes rewards through environment simulation, avoiding PPO’s requirement for a fine-grained critic model.Modern generative-model RL commonly uses policy gradients with KL regularization [38].Source: §2 · p. 2
• Stepwise gains provide fine-grained credit assignment analogous to a learned PPO critic, without a separate critic model; the authors motivate the approach using guarantees for monotone submodular rewards.They describe their gain-based transformation as the first application to GRPO and flow models, extending beyond the on-policy setting of Kveton et al. to off-policy group-relative advantages.Source: §2 · p. 2
• Unlike TempFlow-GRPO [16]’s hand-designed noise-level weighting and Granular-GRPO [58]’s first-half step selection, this method optimizes data-driven telescoping reward gains without manual scheduling or step selection.Source: §2 · pp. 2–3
• Stepwise-Flow-GRPO introduces step-wise rewards throughout diffusion, offering fine-grained guidance and superior performance in image-text alignment and generation quality, whereas Flow-GRPO [26] rewards only the final generation step.Source: §2 · p. 2

• Rectified flow interpolates between data and Gaussian noise, and generation integrates the learned velocity ODE backward from noise to produce samples.The interpolation is given as x_t = (1 − t)x_0 + t x_1 for t in [0, 1].Source: §3 · p. 3
• Flow-GRPO converts the deterministic flow ODE into an SDE with matched marginals, providing stochasticity for reinforcement-learning exploration while preserving the original flow marginals.The SDE uses a noise schedule and Brownian motion and is solved with Euler–Maruyama discretization.Source: §3 · p. 3
• GRPO generates groups of trajectories for each prompt and uses group-relative advantages based on image rewards to optimize the policy.The objective uses propensity-ratio clipping and a KL penalty to constrain policy updates and anchor the learned policy to its initialization [38].Source: §3 · p. 3
• When optimization is on-policy, the importance ratio is 1, reducing GRPO to KL-regularized policy gradients regardless of the clipping parameter.The paper indexes discrete RL steps with t in {T, T − 1, …, 1, 0}, distinct from continuous SDE time.Source: §3 · p. 3

• Low-frequency structure has higher signal-to-noise ratio than high-frequency structure, so denoising recovers composition and layout before details and textures.Uniform credit treats steps determining object layout like steps sharpening an edge.Source: §4 · p. 4
• When later steps correct an early mistake, final-image reward can still reinforce the mistaken step; intermediate reward trajectories can dip and recover.Source: §4 · p. 4

5. Stepwise-Flow-GRPO

• Stepwise-Flow-GRPO estimates intermediate rewards from predicted clean images and assigns advantages according to per-step reward improvements.The method uses Tweedie estimates, then computes gains rather than optimizing intermediate rewards directly.Source: §5 · p. 4, §5.1 · p. 4, §5.2 · pp. 4–5

5.1. Stepwise Rewards

• Tweedie estimates predict clean images from noisy states; five forward-Euler ODE substeps provide a strong reward signal with minimal added cost.A single substep reduces to the one-step Tweedie estimate, while the estimates from different states can be computed in parallel.Source: §5.1 · p. 4

5.2. Stepwise Policy Optimization

• Stepwise gains reward reward-improving steps and penalize reward-decreasing ones; their telescoping sum links local updates to final reward improvement.Joint normalization preserves larger early gains instead of inflating noise at later steps, where improvements are smaller.Source: §5.2 · p. 5

5.3. Connection to Submodular Optimization

• The method connects gain-based policy optimization to adaptive submodular optimization, but the paper does not formally verify submodularity of its reward functions.The authors describe their setting as extending prior work to off-policy optimization and group-relative advantages.Source: §5.3 · p. 5

5.4. Design Variations

• Among the tested alternatives, the gain formulation presented in Section 5.2 performs best in the authors’ experiments.The alternatives include an EMA baseline, GAE, and an ODE-based formulation using successive Tweedie estimates.Source: §5.4 · p. 5

5.5. Improved Sampling via DDIM-Style Updates

• The DDIM-style SDE preserves stochasticity for policy gradients while producing higher-quality samples for reward evaluation.Its update matches rectified-flow marginals using α_t = (1 − t)^2 and β_t = t^2; setting σ_t = 0 recovers the deterministic flow ODE.Source: §5.5 · p. 6
• The DDIM-style update significantly improves visual quality while preserving the ability to optimize with GRPO.It is designed to produce higher-quality reward-model samples while retaining policy-gradient stochasticity.Source: §5.5 · p. 6

6. Experiments

• Experiments use PickScore and GenEval, evaluate three pretrained reward models on GenEval, and include OCR text rendering evaluation.The main setup uses SD3.5-Medium, 10 denoising steps, cfg=1.0, and batch size 16; separate final-quality models use cfg=4.5 and batch size 8.Source: §6 · pp. 6–7

6.1. Main Results

• Stepwise-Flow-GRPO achieves higher reward per training iteration across reward functions, with faster convergence and superior final performance in 3 of 4 settings.It also matches or exceeds Flow-GRPO on final GenEval accuracy.Source: Figure 4 · p. 7, §6.1 · pp. 7–8, Table 1 · p. 8
• Despite added intermediate-denoising computation, Stepwise-Flow-GRPO converges faster in wall-clock time and reaches target rewards with less total computation.It achieves visibly superior performance in 3 of 4 settings.Source: Figure 5 · p. 7, §6.1 · p. 8
• With both methods using the DDIM-inspired SDE, stepwise credit assignment retains its sample-efficiency advantage; the two modifications are reported as complementary.The improved SDE substantially improves both methods.Source: Figure 6 · p. 7, §6.1 · p. 8
• On GenEval, Stepwise-Flow-GRPO matches Flow-GRPO at cfg=1.0 and is equal or better across all subcategories at cfg=4.5.The reported qualitative comparison highlights spatial reasoning, attribute binding, and counting.Source: Figure 3 · p. 6, §6.1 · p. 8, Table 1 · p. 8

6.2. Ablation Studies

• Using T′=5 ODE substeps balances reward-estimate quality and efficiency, while joint normalization preserves larger early gains and converges faster than per-step normalization.With 1–3 substeps estimates are noisy; 6–10 add computational cost with negligible quality gains.Source: §6.2 · p. 8

7. Conclusions and Future Work

• Stepwise-Flow-GRPO assigns credit to denoising steps using gain-based advantages from intermediate reward estimates, alongside a DDIM-inspired SDE to improve reward-signal quality.The authors identify per-prompt gain variance, adaptive step weighting, and self-correcting diffusion as future directions.Source: §7 · p. 8

A. Design Variations

• The standard gain formulation performs best in practice among the explored alternatives for variance reduction and temporal credit propagation.The alternatives are described in the following subsections.Source: A. Design Variations · p. 12

A.1. Exponential Moving Average Baseline

• An exponential moving average of past gains was explored as a per-step baseline to reduce variance while preserving the direction of reward improvement.The baseline uses a decay rate of α = 0.99, and centered gains are normalized and converted to advantages as in the standard formulation.Source: §A.1 · p. 12

A.2. Generalized Advantage Estimation (GAE)

• GAE was applied to gains to account for a step’s possible downstream impact, with γ = 0.95 and GAE values replacing raw gains in the advantage calculation.This was motivated by the possibility that an early compositional decision could enable larger gains during later refinement.Source: §A.2 · p. 12

A.3. ODE-Based Progressive Distillation

• ODE-based progressive distillation matches successive Tweedie estimates and scales the loss by gain-based advantages, avoiding the SDE’s required stochasticity.The authors hypothesized it could improve sample quality and gradient cleanliness, but found it less effective than the standard gain formulation in practice.Source: §A.3 · p. 12, A. Design Variations · p. 12

A.4. Experimental Results

• No tested variation significantly improves on the standard gain formulation, which the authors retain because it preserves diffusion’s natural temporal structure.The alternatives include EMA centering, GAE, and progressive distillation; the latter is less stable than the SDE formulation.Source: §A.4 · pp. 12–13, Figure 7 · p. 13

B. Improved SDE

• The DDIM-inspired SDE interpolates between deterministic ODE sampling and stochastic exploration while preserving the flow’s variance structure.Its noise scale controls stochasticity: σt = 0 recovers the deterministic ODE, while σt > 0 enables RL exploration through controlled noise injection.Source: B. Improved SDE · p. 13

B.1. Variance-Preserving Property

• With Var(x̂1) = 1, the DDIM SDE exactly matches the expected marginal variance, whereas Flow-GRPO’s SDE inflates variance across steps.The excess variance accumulates and produces noisier trajectories that confound reward evaluation.Source: §B.1 · p. 13

B.2. Adaptive Noise Schedule

• The adaptive noise schedule retains predicted noise structure and downweights exploration at earlier, more sensitive steps to make sensitivity-weighted exploration more uniform.The proposed schedule sets σt = ηt√(1 − t), with η = 0 corresponding to deterministic sampling and η = 1 to full replacement.Source: §B.2 · pp. 13–14

B.3. RL Objective

• The method applies the Flow-GRPO GRPO objective with the DDIM SDE as its policy, retaining clipping, KL regularization, and group-relative advantages.Cleaner intermediate samples support more accurate reward evaluation, particularly for vision-based reward models sensitive to image quality.Source: §B.3 · p. 14

C.1. Gain Normalization Strategy

• Joint normalization significantly outperforms per-step normalization in convergence speed, although both eventually reach similar final rewards.Joint normalization preserves relative gain magnitudes across steps; per-step normalization equalizes step importance regardless of natural gain magnitude.Source: §C.1 · p. 15

C.2. Number of Denoising Substeps

• T′ ∈ {2, 5, 8} achieves similar performance in training iterations and wall-clock time, making the method robust to this hyperparameter.The authors select T′ = 5 as the default and note practitioners can adjust it for computational constraints without significantly affecting final performance.Source: §C.2 · p. 15

D. Runtime Analysis

• Intermediate estimates require T′ = 5 ODE substeps per state and reward-model evaluation at each step, adding computational cost.Reward-evaluation overhead varies by model: ImageReward takes 8.2s versus 0.7s, while UnifiedReward takes 114.6s versus 48.8s.Source: D. Runtime Analysis · p. 15
• Despite 1.8–2.4× more generation time per iteration, Stepwise-Flow-GRPO converges 20–40% faster overall in wall-clock time across settings.It requires 2–3× fewer training iterations to reach target rewards, more than compensating for per-iteration overhead.Source: D. Runtime Analysis · pp. 15–16

E.1. OCR Text Rendering

• On OCR text rendering with an 80% OCR and 20% PickScore reward, Flow-GRPO diverges after about 500 steps while Stepwise-Flow-GRPO continues improving to a higher plateau.The authors characterize the task as compositional, requiring readable text alongside visual quality.Source: §E.1 · p. 16, Figure 11 · p. 16

E.2. Extended GenEval Training

• After 400 GPU hours on GenEval, Stepwise-Flow-GRPO scores 0.87 overall, versus 0.72 for Flow-GRPO and 0.84 for GPT-4o.It also reports 0.89 on counting, 0.73 on spatial positioning, and 0.80 on attribute binding.Source: §E.2 · p. 16

E.3. UnifiedReward Training

• After 60 GPU hours, Stepwise-Flow-GRPO achieves a 0.74 GenEval score, with smooth, stable training curves; the authors also report consistent efficiency advantages throughout extended training.Source: §E.3 · p. 16
• 0.87 GenEval overall after 400 GPU hours, versus Flow-GRPO’s 0.72.The figure describes the performance gap as widening during extended training.Source: Figure 12 · p. 17

E.4. Qualitative Results

• Qualitative comparisons show gains in object counting, real-world dynamics, and image quality, most pronounced at training step 60.The examples include improved prompt adherence and more plausible compositions.Source: §E.4 · p. 17

E.3. UnifiedReward Training

• Flow-GRPO with UnifiedReward diverged, while Stepwise-Flow-GRPO trained stably throughout.The authors note that this robustness is particularly valuable for large VLM-based rewards, where gradient noise can be substantial.Source: §E.3 · pp. 16–17, Figure 13 · p. 17

F. Comparison with Concurrent Work

• Stepwise-Flow-GRPO optimizes telescoping gains for marginal improvement, while TempFlow-GRPO and Granular-GRPO use normalized final rewards with trajectory branching.The related methods emphasize early steps through hand-designed noise-level weighting or optimization of only the first 8 of 16 steps; this method’s gains decrease naturally as t →0.Source: F. Comparison with Concurrent Work · p. 17
• 0.87 GenEval overall after 400 GPU hours, versus Flow-GRPO’s 0.72 and GPT-4o’s 0.84.Table 3 presents complete GenEval results; Flow-GRPO’s extended results and reference-model scores are from their respective papers.Source: Table 3 · p. 18
• Approximately 1.5× faster per epoch than TempFlow-GRPO’s released code at equivalent batch sizes.The comparison uses the authors’ observed runtime.Source: F. Comparison with Concurrent Work · p. 18
• Qualitative comparisons report more accurate counts, better spatial arrangements, and higher image quality than Flow-GRPO across training.The visible differences are reported as largest at step 60; a separate comparison says UnifiedReward training yields higher visual quality and photorealism than GenEval-reward training.Source: Figure 14 · p. 19, Figure 15 · p. 20, Figure 16 · p. 21