JAGG: 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~citeshao2024deepseekmathpushinglimitsmathematical, 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 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 (textttjaggfrac) deploys JAGG only where the assumption holds. Experiments on T2I benchmarks show JAGG delivers ~2× backward speedup with negligible quality degradation. The code for this work can be accessed through https://github.com/SchumiDing/JAGG.