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

• URL: http://arxiv.org/abs/2607.17572v3
• Date: Sat, 25 Jul 2026 08:59:31 GMT
• Title: JAGG: Jacobian-Aggregated Group Gradient for Efficient GRPO Training of Diffusion Models
• Abstract summary: We present textbfJAGG, which reduces full transformer backward passes from $W$ to $2$ per group of $W$ consecutive steps.<n>Experiments on T2I benchmarks show JAGG delivers $sim$times$ backward speedup with negligible quality degradation.
• Score: 21.274538255180314
• License: http://creativecommons.org/licenses/by/4.0/
• 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~\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 \emph{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 \emph{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 \textbf{JAGG} (\textbf{J}acobian-\textbf{A}ggregated \textbf{G}roup \textbf{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 \emph{exact} when the velocity is linear in $(z,t)$, and a cosine-similarity routing rule (\texttt{jagg\_frac}) deploys JAGG only where the assumption holds. Experiments on T2I benchmarks show JAGG delivers $\sim$2$\times$ backward speedup with negligible quality degradation. The code for this work can be accessed through https://github.com/SchumiDing/JAGG.

Related papers

• Concurrent Image Understanding and Generation: Self-Correcting Coupled Markov Jump Processes [70.61868608402723]
  We introduce $textbfSelf-Correcting Coupled Markov Jump Processes (SC-CMJP)<n>In conjunction with SC-CMJP, we introduce $textttCO_texttt2textttJump$, a novel training-free single-pass sampler for joint multimodal geneneration.<n>For training and evaluation purposes, we have created and will release three large-scale joint multimodal generation corpora.
  arXiv  Detail & Related papers  (2026-07-14T18:39:29Z)
• LoRIF: Low-Rank Influence Functions for Scalable Training Data Attribution [62.830878652285406]
  Training data attribution identifies which training examples most influenced a model's prediction.<n>LoRIF exploits low-rank structures of gradient to address both bottlenecks.<n>On models from 0.1B to 70B parameters trained on datasets with millions of examples, LoRIF achieves up to 20$times$ storage reduction and query-time speedup.
  arXiv  Detail & Related papers  (2026-01-29T16:18:34Z)
• Evolution Strategies at the Hyperscale [57.75314521465674]
  We introduce EGGROLL, an evolution strategies (ES) algorithm designed to scale backprop-free optimization to large population sizes.<n>ES is a set of powerful blackbox optimisation methods that can handle non-differentiable or noisy objectives.<n>EGGROLL overcomes these bottlenecks by generating random matrices $Ain mathbbRmtimes r, Bin mathbbRntimes r$ with $rll min(m,n)$ to form a low-rank matrix perturbation $A Btop$
  arXiv  Detail & Related papers  (2025-11-20T18:56:05Z)
• Multi-Timescale Gradient Sliding for Distributed Optimization [23.311605203774388]
  We propose two first-order methods for convex, non-smooth, distributed optimization problems.<n>Our MT-GS and AMT-GS can take advantage of similarities between (local) objectives to reduce the communication rounds.<n>These three desirable features are achieved through a block-decomposable primal-dual formulation.
  arXiv  Detail & Related papers  (2025-06-18T12:03:34Z)
• On Fitting Flow Models with Large Sinkhorn Couplings [34.004981672572484]
  Flow models transform data gradually from one modality (e.g. noise) onto another (e.g. images)<n>Recent works have proposed to sample mini-batches of $n$ source and $n$ target points and reorder them using an OT solver to form better pairs.<n>We show that flow models greatly benefit when fitted with large Sinkhorn couplings, with a low entropic regularization.
  arXiv  Detail & Related papers  (2025-06-05T19:19:01Z)
• MGDA Converges under Generalized Smoothness, Provably [27.87166415148172]
  Multi-objective optimization (MOO) is receiving more attention in various fields such as multi-task learning.<n>Recent works provide some effective algorithms with theoretical analysis but they are limited by the standard $L$-smooth or bounded-gradient assumptions.<n>We study a more general and realistic class of generalized $ell$-smooth loss functions, where $ell$ is a general non-decreasing function of gradient norm.
  arXiv  Detail & Related papers  (2024-05-29T18:36:59Z)
• Projection by Convolution: Optimal Sample Complexity for Reinforcement Learning in Continuous-Space MDPs [56.237917407785545]
  We consider the problem of learning an $varepsilon$-optimal policy in a general class of continuous-space Markov decision processes (MDPs) having smooth Bellman operators. Key to our solution is a novel projection technique based on ideas from harmonic analysis. Our result bridges the gap between two popular but conflicting perspectives on continuous-space MDPs.
  arXiv  Detail & Related papers  (2024-05-10T09:58:47Z)
• Near-Optimal Regret Bounds for Multi-batch Reinforcement Learning [54.806166861456035]
  We study the episodic reinforcement learning (RL) problem modeled by finite-horizon Markov Decision Processes (MDPs) with constraint on the number of batches. We design a computational efficient algorithm to achieve near-optimal regret of $tildeO(sqrtSAH3Kln (1/delta))$tildeO(cdot) hides logarithmic terms of $(S,A,H,K)$ in $K$ episodes. Our technical contribution are two-fold: 1) a near-optimal design scheme to explore
  arXiv  Detail & Related papers  (2022-10-15T09:22:22Z)
• Asynchronous Training Schemes in Distributed Learning with Time Delay [17.259708772713164]
  In the context of distributed deep learning, the issue of stale weights or gradients could result in poor algorithmic performance. In this paper, we propose a different approach to tackle the issue of stale weights or gradients. One practical variant of PC-ASGD is also proposed by adopting a condition to help with the determination of the tradeoff parameter.
  arXiv  Detail & Related papers  (2022-08-28T07:14:59Z)
• Adan: Adaptive Nesterov Momentum Algorithm for Faster Optimizing Deep Models [134.83964935755964]
  In deep learning, different kinds of deep networks typically need different extrapolations, which have to be chosen after multiple trials.<n>To relieve this issue and consistently improve the model training speed deep networks, we propose the ADAtive Nesterov momentum Transformer.
  arXiv  Detail & Related papers  (2022-08-13T16:04:39Z)
• Accelerated Gradient Tracking over Time-varying Graphs for Decentralized Optimization [59.65871549878937]
  We prove that the practical single loop accelerated gradient tracking needs $O(fracgamma1-sigma_gamma)2sqrtfracLepsilon)$.<n>Our convergence rates improve significantly over the ones of $O(frac1epsilon5/7)$ and $O(fracLmu)5/7frac1 (1-sigma)1.5logfrac1epsilon)$.
  arXiv  Detail & Related papers  (2021-04-06T15:34:14Z)
• Optimal Robust Linear Regression in Nearly Linear Time [97.11565882347772]
  We study the problem of high-dimensional robust linear regression where a learner is given access to $n$ samples from the generative model $Y = langle X,w* rangle + epsilon$ We propose estimators for this problem under two settings: (i) $X$ is L4-L2 hypercontractive, $mathbbE [XXtop]$ has bounded condition number and $epsilon$ has bounded variance and (ii) $X$ is sub-Gaussian with identity second moment and $epsilon$ is
  arXiv  Detail & Related papers  (2020-07-16T06:44:44Z)
• Training (Overparametrized) Neural Networks in Near-Linear Time [21.616949485102342]
  We show how to speed up the algorithm of [CGH+1] for training (mildly overetrized) ReparamLU networks. The centerpiece of our algorithm is to reformulate the Gauss-Newton as an $ell$-recondition.
  arXiv  Detail & Related papers  (2020-06-20T20:26:14Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.