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.
Figures & tables
Figure 1: Training convergence under different CFG scales. (a) CFG=1.0 . (b) CFG=4.5 . SGA-Flow-GRPO converges comparably to DiffusionNFT and substantially faster than Flow-GRPO.
Figure 2: Qualitative comparison of spatial credit normalization strategies. Our post-processing better preserves active regions and boundaries than Min-Max normalization (with 0.15 quantile truncation) and its negated variant, yielding spatially consistent credits before and after patchification.
Figure 3: Overview of SGA-Flow-GRPO. Given sampled denoising trajectories, our approach decomposes transition-level log-likelihoods into token-wise components to construct fine-grained importance ratios. Meanwhile, reward gradients provide spatial credit signals to modulate the group-relative advantage, enabling differentiated policy updates across spatial tokens.
Model
Step
Overall ↑
Single Obj. ↑
Two Obj. ↑
Counting ↑
Colors ↑
Position ↑
Attr. Binding ↑
Off-the-Shelf Models
SD3.5-M (w/o CFG)
-
0.27
0.70
0.20
0.18
0.45
0.04
0.07
SD3.5-M
-
0.63
0.98
0.78
0.50
0.81
0.24
0.52
FLUX.1 Dev
-
0.66
0.98
0.81
0.74
0.79
0.22
0.45
GPT-4o
-
0.84
0.99
0.92
0.85
0.92
0.75
0.61
Post-training Methods
Table 1: Evaluation Results on Geneval.
Figure 4: Qualitative Comparison on the GenEval Benchmark.
Figure 5: Ablation of log-likelihood tokenization and spatial credit assignment. Under CFG=1.0 and CFG=4.5 , we report the training convergence curves and the GenEval scores on the test set.
Figure 6: Performance comparison between Multi-reward and GenEval-only optimization under CFG=4.5 (top) and CFG=1.0 (bottom).
Figure 7: Ablation of temperature τ in the reward-gradient spatial map under CFG=1.0 (left) and CFG=4.5 (right), evaluated by GenEval score on the test set.
Figure 8: Ablation study on the clipping range ϵ . Subfigures (a), (b), and (c) show the training reward curves with clipping ranges ϵ=0.1 , 0.2 , and 0.3 , respectively. Subfigure (d) compares the corresponding GenEval scores on the test set.
BS
PickScore
ImageReward
Ori.
Diff.
Ori.
Diff.
1
0.024
0.044
0.019
0.047
4
0.059
0.109
0.023
0.059
8
0.093
0.188
0.042
0.081
16
0.162
0.343
0.075
0.158
Appendix
Table 2: Runtime comparison of reward computation. All results are reported in seconds on an NVIDIA A100 80GB GPU.
Figure 9: PickScore comparison on Pick-a-Pic. We compare our method with DiffusionNFT and Flow-GRPO under CFG=1.0 and CFG=4.5 . All methods are trained for 600 steps, and test-set PickScore is evaluated every 120 training steps.
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.
This paper focuses on the alignment of flow matching models with human preferences. A promising way is fine-tuning by directly backpropagating reward gradients through the differentiable generation process of flow matching. However, backpropagating through long trajectories results in prohibitive memory costs and gradient explosion. Therefore, direct-gradient methods struggle to update early generation steps, which are crucial for determining the global structure of the final image. To address this issue, we introduce LeapAlign, a fine-tuning method that reduces computational cost and enables direct gradient propagation from reward to early generation steps. Specifically, we shorten the long trajectory into only two steps by designing two consecutive leaps, each skipping multiple ODE sampling steps and predicting future latents in a single step. By randomizing the start and end timesteps of the leaps, LeapAlign leads to efficient and stable model updates at any generation step. To better use such shortened trajectories, we assign higher training weights to those that are more consistent with the long generation path. To further enhance gradient stability, we reduce the weights of gradient terms with large magnitude, instead of completely removing them as done in previous works. When fine-tuning the Flux model, LeapAlign consistently outperforms state-of-the-art GRPO-based and direct-gradient methods across various metrics, achieving superior image quality and image-text alignment.
Zhanhao Liang, Tao Yang, Jie Wu +2
1The Australian National University · 2ByteDance Seed
Deploying GRPO on Flow Matching models has proven effective for text-to-image generation. However, existing paradigms typically propagate an outcome-based reward to all preceding denoising steps without distinguishing the local effect of each step. Moreover, current group-wise ranking mainly compares trajectories at matched timesteps and ignores within-trajectory dependencies, where certain early denoising actions can affect later states via delayed, implicit interactions. We propose TurningPoint-GRPO (TP-GRPO), a GRPO framework that alleviates step-wise reward sparsity and explicitly models long-term effects within the denoising trajectory. TP-GRPO makes two key innovations: (i) it replaces outcome-based rewards with step-level incremental rewards, providing a dense, step-aware learning signal that better isolates each denoising action's "pure" effect, and (ii) it identifies turning points-steps that flip the local reward trend and make subsequent reward evolution consistent with the overall trajectory trend-and assigns these actions an aggregated long-term reward to capture their delayed impact. Turning points are detected solely via sign changes in incremental rewards, making TP-GRPO efficient and hyperparameter-free. Extensive experiments also demonstrate that TP-GRPO exploits reward signals more effectively and consistently improves generation. Demo code is available at https://github.com/YunzeTong/TurningPoint-GRPO.
Yunze Tong, Mushui Liu, Canyu Zhao +7
1Zhejiang University, Hangzhou, China · 2Taobao & Tmall Group of Alibaba, Hangzhou, China · 3Tsinghua University, Beijing, China