cs.CVSep 26, 2026

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

Authors: Yunkai Yang, Yudong Zhang, Xinying Chen, Bin Luo, Jienan Lyu, Kunquan Zhang, Weitao Wan, Runmin Dong

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

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Explore similar work

CardsList
  1. Stepwise Credit Assignment for GRPO on Flow-Matching Models

    Mar 30, 2026Yash Savani, Branislav Kveton, Yuchen Liu +5RL for Diffusion ModelsStep-Level Credit Assignment in RL

  2. LeapAlign: Post-Training Flow Matching Models at Any Generation Step by Building Two-Step Trajectories

    Apr 16, 2026Zhanhao Liang, Tao Yang, Jie Wu +2Diffusion Model AlignmentFlow Matching

  3. 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 MatchingTemporal Credit Assignment