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

arXiv: https://arxiv.org/abs/2607.17572 · HTML: https://mlsanity.com/paper/2607.17572/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Reinforcement learning for reasoning, one of 17 directions in LLM reasoning.
- Builds on 40 of its 47 references that are in the map.
- First posted on arXiv 20 Jul 2026.

## Placed in

- LLM reasoning / [Reinforcement learning for reasoning](https://mlsanity.com/lr/reinforcement-learning-for-reasoning/)
- Diffusion models / [Diffusion and flow policies for robotics](https://mlsanity.com/diffusion/diffusion-and-flow-policies-for-robotics/)

## Builds on (most cited first, 30 of 40)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [CLIPScore: A Reference-free Evaluation Metric for Image Captioning](https://mlsanity.com/paper/2104.08718/) (2021)
- [Optimization Methods for Large-Scale Machine Learning](https://mlsanity.com/paper/1606.04838/) (2016)
- [Qwen-Image Technical Report](https://mlsanity.com/paper/2508.02324/) (2025)
- [FLUX.1 Kontext: Flow Matching for In-Context Image Generation and Editing in Latent Space](https://mlsanity.com/paper/2506.15742/) (2025)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Training Diffusion Models with Reinforcement Learning](https://mlsanity.com/paper/2305.13301/) (2023)
- [Pick-a-Pic: An Open Dataset of User Preferences for Text-to-Image Generation](https://mlsanity.com/paper/2305.01569/) (2023)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [Human Preference Score v2: A Solid Benchmark for Evaluating Human Preferences of Text-to-Image Synthesis](https://mlsanity.com/paper/2306.09341/) (2023)
- [Flow-GRPO: Training Flow Matching Models via Online RL](https://mlsanity.com/paper/2505.05470/) (2025)
- [DanceGRPO: Unleashing GRPO on Visual Generation](https://mlsanity.com/paper/2505.07818/) (2025)
- [DiffusionNFT: Online Diffusion Reinforcement with Forward Process](https://mlsanity.com/paper/2509.16117/) (2025)
- [A Survey of Reinforcement Learning for Large Reasoning Models](https://mlsanity.com/paper/2509.08827/) (2025)
- [MixGRPO: Unlocking Flow-based GRPO Efficiency with Mixed ODE-SDE](https://mlsanity.com/paper/2507.21802/) (2025)
- [Delta-DiT: A Training-Free Acceleration Method Tailored for Diffusion Transformers](https://mlsanity.com/paper/2406.01125/) (2024)
- [Real-Time Video Generation with Pyramid Attention Broadcast](https://mlsanity.com/paper/2408.12588/) (2024)
- [From Reusing to Forecasting: Accelerating Diffusion Models with TaylorSeers](https://mlsanity.com/paper/2503.06923/) (2025)
- [Diffusion Models for Reinforcement Learning: A Survey](https://mlsanity.com/paper/2311.01223/) (2023)
- [BranchGRPO: Stable and Efficient GRPO with Structured Branching in Diffusion Models](https://mlsanity.com/paper/2509.06040/) (2025)
- [GRPO-Guard: Mitigating Implicit Over-Optimization in Flow Matching via Regulated Clipping](https://mlsanity.com/paper/2510.22319/) (2025)
- [Improving Reasoning for Diffusion Language Models via Group Diffusion Policy Optimization](https://mlsanity.com/paper/2510.08554/) (2025)
- [AR-GRPO: Training Autoregressive Image Generation Models via Reinforcement Learning](https://mlsanity.com/paper/2508.06924/) (2025)
- [Rethinking the Design Space of Reinforcement Learning for Diffusion Models: On the Importance of Likelihood Estimation Beyond Loss Design](https://mlsanity.com/paper/2602.04663/) (2026)
- [Inpainting-Guided Policy Optimization for Diffusion Large Language Models](https://mlsanity.com/paper/2509.10396/) (2025)
- [Manifold-Aware Exploration for Reinforcement Learning in Video Generation](https://mlsanity.com/paper/2603.21872/) (2026)
- [Reinforcement Learning Meets Masked Generative Models: Mask-GRPO for Text-to-Image Generation](https://mlsanity.com/paper/2510.13418/) (2025)

## Abstract (arXiv)

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.