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

What the math gives to ML

The paper contains a directly transferable computational principle: when a neural vector field changes nearly linearly along a diffusion trajectory, intermediate parameter Jacobians can be replaced by interpolation between two endpoint Jacobians. This converts a sum of W per-timestep backward passes into two endpoint backward passes by aggregating the corresponding upstream gradients with interpolation weights. The most useful neural-network application is a training-time backward-pass reduction for diffusion or flow-policy objectives, with a similarity gate that falls back to exact backpropagation when endpoint and interpolated predictions disagree. The method is especially attractive for GRPO because forward sampling is retained while the expensive DiT backward graph is evaluated only at selected trajectory endpoints.