用线性近似加速扩散模型强化学习训练，速度翻倍且质量几乎不变。
JAGG: Jacobian-Aggregated Group Gradient for Efficient GRPO Training of Diffusion Models
- 通过时间加权插值端点雅可比矩阵，近似中间步梯度。
- 将每组步骤的梯度聚合为两次反向传播，实现约2倍加速。
- 仅在速度线性时启用，适合高分辨率文本到图像生成训练。
组相对策略优化（GRPO）是一种强大的强化学习算法，用于将生成模型与人类偏好对齐。尽管在大语言模型中表现优异，将其扩展到扩散模型和流匹配模型时引入了严重的计算瓶颈：必须在采样轨迹的每个时间步对高容量DiT主干进行反向传播，导致高分辨率文本到图像（T2I）训练成本过高。训练无关的DiT推理加速方法（如Δ-DiT、ScalingCache）利用了DiT隐藏状态和速度预测沿轨迹平滑且近似线性的特性。我们提出，这种线性是否也能降低DiT强化学习训练的反向传播开销？答案是肯定的，为此我们提出了JAGG（雅可比聚合分组梯度），将每组W个连续步骤的完整Transformer反向传播从W次减少至2次。JAGG通过时间加权插值端点雅可比矩阵来近似中间步雅可比，再将每步的上游信号聚合为两个复合梯度，通过一次联合反向传播完成。我们证明该插值在速度关于(z,t)线性时是精确的，并设计了一个余弦相似度路由规则（jagg_frac）仅在假设成立时启用JAGG。在T2I基准测试上的实验表明，JAGG实现了约2倍的反向传播加速，且质量损失可忽略。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。