JAGG: Jacobian-Aggregated Group Gradient for Efficient GRPO Training of Diffusion Models

arXiv:2607.17572v2 Announce Type: replace Abstract: Group Relative Policy Optimization (GRPO) is a powerful reinforcement learning algorithm for aligning generative models with human preferences. While successful in large language models~ite{shao2024deepseekmathpushinglimitsmathematical}, its extension to diffusion and flow matching models introduces a severe computational bottleneck: gradients must be back-propagated through the high-capacity DiT backbone at every timestep of the sampling trajectory, making high-resolution text-to-image (T2I) training prohibitively expensive. Training-free DiT inference acceleration methods (e.g., Delta-DiT, ScalingCache) exploit the fact that DiT hidden states and velocity predictions vary smoothly and nearly linearly along the trajectory. We ask whether the same linearity can reduce the backward-pass cost of DiT RL training, and answer affirmatively with extbf{JAGG} (extbf{J}acobian-extbf{A}ggregated extbf{G}roup extbf{G}radient), which reduces full transformer backward passes from W to 2 per group of W consecutive steps. JAGG approximates intermediate-step Jacobians via t-weighted interpolation of the endpoint Jacobians, then aggregates per-step upstream signals into two composite gradients applied through a single joint backward pass. We prove this interpolation is exact when the velocity is linear in (z,t), and a cosine-similarity routing rule (exttt{jagg_frac}) deploys JAGG only where the assumption holds. Experiments on T2I benchmarks show JAGG delivers sim2imes backward speedup with negligible quality degradation.

Source: arXiv cs.LG | 2026-07-23