Technical Summary: JAGG (Jacobian-Aggregated Group Gradient)

Problem Statement

Group Relative Policy Optimization (GRPO) has emerged as a leading reinforcement learning (RL) algorithm for aligning generative models with human preferences, successfully eliminating the need for a separate value network in Large Language Models (LLMs). However, extending GRPO to diffusion and flow matching models (specifically Diffusion Transformers or DiTs) introduces a severe computational bottleneck.

In diffusion-based GRPO, the generation process is treated as a multi-step decision problem where the policy is optimized over a sampling trajectory of $W$ timesteps. Standard training requires a full forward and backward pass through the high-capacity DiT backbone for every single timestep in the trajectory to compute the policy gradient. For high-resolution text-to-image (T2I) generation, where trajectories often span dozens of steps, this results in prohibitively expensive training costs. Unlike language models where token steps are lightweight, each backward pass through a modern DiT incurs the full cost of a transformer forward pass, making the per-step backward requirement the dominant training expense.

While training-free inference acceleration methods (e.g., $\Delta$-DiT, ScalingCache) have successfully exploited the observation that DiT hidden states and velocity predictions vary smoothly and nearly linearly along the denoising trajectory to skip intermediate evaluations, it remained an open question whether this same linearity could be leveraged to reduce the backward-pass cost in RL training.

Methodology: JAGG

The authors propose JAGG (Jacobian-Aggregated Group Gradient), a mathematically principled acceleration method that reduces the number of network backward passes from $W$ to 2 for a group of $W$ consecutive denoising steps.

Core Insight

The method relies on the insight that the Jacobian of the denoising velocity field $v_\theta(z, t)$ with respect to model parameters $\theta$ exhibits approximate linearity in the inputs $(z, t)$, particularly during the high-noise regime of the denoising process. This allows intermediate-step Jacobians to be accurately approximated via linear interpolation using only the endpoint Jacobians ($J_0$ and $J_{W-1}$).

Algorithmic Workflow

JAGG processes a group of $W$ timesteps through four phases:

1. Forward Passes: The model performs forward passes for all $W$ steps. However, gradient tracking is enabled only for the two endpoints (step 0 and step $W-1$). The intermediate $W-2$ steps are executed using torch.no_grad() to save memory and computation.
2. Upstream Gradient Computation: For each step $j$, a detached proxy of the model output is created. The per-step loss $L_j$ (derived from the PPO-style objective) is computed, and the upstream gradient $s_j = \partial L_j / \partial \hat{v}_j$ is calculated via automatic differentiation. Crucially, this step involves no transformer backward pass; the gradient flows only through the lightweight SDE step and the PPO objective.
3. Timestep-Weighted Aggregation: The upstream gradients $s_j$ are aggregated into two composite signals, $s_{base}$ and $s_{corr}$, using interpolation weights $\alpha_j$ derived from the timestep $t_j$:
   $$\alpha_j = \frac{t_0 - t_j}{t_0 - t_{W-1}}$$
   $$s_{base} = \sum_{j=0}^{W-1} (1 - \alpha_j) s_j, \quad s_{corr} = \sum_{j=0}^{W-1} \alpha_j s_j$$
   This aggregation effectively reconstructs the full gradient signal using the linear interpolation assumption.
4. Joint Backward: A single torch.autograd.backward call is issued on the two endpoint outputs $[\hat{v}_0, \hat{v}_{W-1}]$ with the aggregated gradients $[s_{base}, s_{corr}]$. This triggers exactly one FSDP all-gather per module, replacing the $W$ separate backward passes required by standard training.

Theoretical Guarantee

The authors provide a formal proof (Theorem 1) stating that if the velocity field $v_\theta(z, t)$ is strictly linear in $(z, t)$, the $t$-weighted Jacobian interpolation is exact with zero approximation error. Under these conditions, the two-endpoint backward pass recovers the exact group gradient.

Adaptive Routing ($jagg\_frac$)

Recognizing that the linearity assumption degrades in the low-noise (late) stages of denoising, the authors introduce a routing parameter $jagg\_frac$. This parameter applies JAGG only to the first fraction of denoising groups (where the high-noise linearity assumption holds, empirically verified via cosine similarity of gradients) and uses exact per-step backward passes for the remaining groups. For the T2I models tested, $jagg\_frac$ was set to 0.6.

Key Contributions

1. Bottleneck Identification: The paper formally identifies and characterizes the computational bottleneck in diffusion GRPO training caused by per-timestep backward passes across the sampling trajectory.
2. JAGG Framework: It proposes a Jacobian-aggregation framework that reduces backward passes from $W$ to 2 per group, grounded in a rigorous theoretical analysis of velocity Jacobian linearity.
3. Theoretical Proof: The authors prove that under linear velocity conditions, the proposed interpolation is exact, providing a solid theoretical foundation for the approximation.
4. Empirical Validation: Extensive experiments on multiple diffusion backbones (Flux, QwenImage) and RL algorithms (DanceGRPO, Flow-GRPO-Fast) demonstrate the method's generality and effectiveness.

Experimental Results

The authors evaluated JAGG on T2I benchmarks using HPSv2.1, PickScore, and CLIPScore as metrics.

• Efficiency: JAGG achieved a consistent wall-clock speedup of 17.3% to 18.2% per training step on standard DanceGRPO setups (reducing time from ~835s to ~691s on QwenImage and ~279s to ~228s on Flux). When applied to Flow-GRPO-Fast, the update stage alone saw speedups of 29.9% to 30.0%.
• Quality: The method delivered these speedups with negligible degradation in generation quality. In many cases, JAGG closely tracked the baseline reward curves and matched or slightly exceeded the baseline at convergence.
• Ablation: An ablation variant that performed backward passes on endpoints but without Jacobian interpolation (assigning full loss to endpoints independently) consistently underperformed. This confirms that the aggregation of intermediate gradient information via Jacobian interpolation is critical for stable training.
• Gradient Quality: Empirical analysis showed that in high-noise groups (0–2), the cosine similarity between the JAGG-aggregated gradient and the exact gradient was high (0.74–0.85), validating the linearity assumption. Similarity dropped in later groups, justifying the $jagg\_frac$ routing strategy.

Significance and Claims

The paper claims that JAGG represents a significant step toward making large-scale diffusion GRPO training practical. By exploiting the near-linearity of velocity Jacobians in the high-noise regime, the method reduces the computational burden of RL training without sacrificing generation quality.

The authors emphasize that JAGG is not merely a heuristic shortcut but is grounded in the mathematical properties of the denoising trajectory. The method is presented as orthogonal to existing inference acceleration techniques, specifically targeting the backward-pass bottleneck unique to diffusion RL. The results suggest that JAGG enables more widespread adoption of RL-based alignment for visual generative models by significantly lowering the training cost barrier.