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

Reinforcement Learning (RL) for diffusion models, specifically Group Relative Policy Optimization (GRPO), requires back-propagating gradients through the entire denoising trajectory to align models with human preferences. Because each denoising step involves a full transformer forward and backward pass, training high-resolution text-to-image models becomes prohibitively expensive. This paper investigates whether the observed smoothness and near-linearity of diffusion model activations can be leveraged to reduce the number of backward passes required during training.

The authors introduce JAGG (Jacobian-Aggregated Group Gradient), a technique that groups $W$ consecutive denoising steps and approximates the gradients for the intermediate steps. The core insight is that the Jacobian of the denoising velocity field with respect to model parameters is approximately linear in the input space during the early, high-noise denoising phase.

JAGG operates in four phases:

1. Forward Passes: Only the two endpoints of a group of $W$ steps are computed with gradient tracking; intermediate steps are computed without gradients.
2. Upstream Gradient Computation: Per-step losses are computed using detached proxies, avoiding full transformer backward passes for intermediate steps.
3. Timestep-Weighted Aggregation: The upstream gradients are aggregated into two composite signals using $t$-weighted interpolation.
4. Joint Backward: A single joint backward pass is performed on the two endpoints, effectively approximating the full gradient signal for the entire group.

To maintain stability, the authors implement a jagg_frac routing rule, which applies JAGG only to the initial, high-noise portion of the trajectory where the linearity assumption holds, while reverting to standard per-step backpropagation for the final, low-noise steps.

JAGG provides a mathematically grounded way to accelerate the training of diffusion models under RL frameworks. By reducing the number of full transformer backward passes from $W$ to 2 per group, it achieves a $\sim 2\times$ speedup in backward pass efficiency. This significantly lowers the barrier to entry for fine-tuning large-scale diffusion models on human preference data, enabling more frequent and efficient alignment iterations without sacrificing generation quality.