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

Group Relative Policy Optimization (GRPO) is a powerful reinforcement learning algorithm for aligning generative models with human preferences. While successful in large language models \cite{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., Δ-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 JAGG (Jacobian-Aggregated Group Gradient), 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 (jagg_frac) deploys JAGG only where the assumption holds. Experiments on T2I benchmarks show JAGG delivers \sim2\times backward speedup with negligible quality degradation.