SGA-Flow-GRPO: Spatial Gradient-Guided Credit Assignment for Flow-GRPO
Organizations: Sun Yat-Sen University · Baidu · Beijing Institute of Technology · Tsinghua University ·
Abstract
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
| 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 | ||||||||
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
| 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 |
Diffusion Model Alignment
Diffusion model alignment steers generated content and internal representations toward prompts, human preferences, safety goals, and task constraints.
See all 185 papers