⚡ electrical engineering

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

本文介绍了 JAGG（Jacobian-Aggregated Group Gradient），这是一种利用扩散模型轨迹的近线性变化来近似中间雅可比矩阵并聚合组梯度的方法，该方法在不显著损失质量的情况下，将扩散 Transformer 进行组相对策略优化（GRPO）训练的计算成本降低了约 2 倍。

原作者： Ruiyi Ding, Jie Li, He Kang, Ziyan Liu, Chengru Song, Yuan chen

发布于 2026-07-21

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Ruiyi Ding, Jie Li, He Kang, Ziyan Liu, Chengru Song, Yuan chen

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图教一位数字艺术家如何画出更好的画作。这位艺术家并不是人类，而是一个复杂的计算机程序，被称为“扩散模型”。把这个模型想象成一位雕塑家，他从一块充满噪声和静电感的黏土块开始，一步步地凿掉噪声，直到一座完美的雕像显现出来。为了教会这位雕塑家创作出人类真正喜欢的艺术品，我们使用了一种叫做“强化学习”的技术。这就像一场游戏，雕塑家尝试几种不同的凿刻动作，从一位评委（奖励模型）那里得到一个分数，然后根据这个分数学习如何在下次做得更好。

问题在于，这种学习过程极其昂贵。每当雕塑家进行一次微小的凿刻（一个“步骤”）时，计算机都必须运行一次大规模、高能耗的计算，以确定究竟该如何调整雕塑家的“大脑”。如果雕像需要 20 个步骤才能完成，那么对于每一次教导，计算机都要进行 20 次这种繁重的数学运算。这就像是在学习一段新舞蹈时，每做一个动作就要停下来向教练请求一次全面的评价，重写你的整个肌肉记忆，然后重新开始。对于高分辨率图像来说，这所花费的时间和电力多到让有效地训练这些模型几乎变得不可能。

这正是名为 JAGG 的新思路介入的地方。这项研究背后的研究人员提出了一个聪明的疑问：“我们真的需要在每一个动作之后都停下来并重新计算一切吗？”他们注意到，在雕塑过程初期那些混乱的阶段，变化发生得非常平滑且具有可预测性——几乎就像一条直线。如果你知道雕塑在几步之后从起点到了终点，你实际上可以推测出中间发生了什么，而无需为每一个步骤进行繁重的数学运算。

该论文引入了一种名为 JAGG（雅可比聚合组梯度，Jacobian-Aggregated Group Gradient）的方法。JAGG 不再在每一个步骤后都停下来计算教导，而是让模型连续执行一组步骤（例如连续做四个动作），并且只在这一组的开头和结尾进行两次繁重的“大脑更新”数学运算。对于中间的步骤，它使用了一个聪明的捷径：它假设起始点和终点之间的路径是一条直线，并以此填补空白。

研究人员在数学上证明了，如果雕塑过程是完美平滑且线性的，那么这个捷径实际上是完美的——它给出的答案与每次都进行繁重计算的结果完全一致。在现实世界中，尽管情况并非完全笔直，但他们发现这个捷径在图像生成的“噪声”早期阶段表现得异常出色。通过使用这种方法，他们能够将这些模型的训练时间缩短约 17% 到 18%（按每步计算）。

最棒的部分是？图像质量并没有因此受损。当他们在不同模型（如 Flux 和 QwenImage）上测试 JAGG 时，生成的图片与使用缓慢的传统方法所生成的图片一样好。事实上，在某些情况下，这个捷径甚至帮助模型学习得更加稳定。论文表明，通过在何时进行“重体力活”方面表现得更聪明一点，我们可以让这些强大的 AI 艺术家训练得更快、更便宜，且无需牺牲最终艺术品的精美程度。

技术摘要：JAGG (Jacobian-Aggregated Group Gradient，雅可比聚合组梯度)

问题陈述

群体相对策略优化（GRPO）已成为一种领先的强化学习（RL）算法，用于使生成模型符合人类偏好，并成功消除了大型语言模型（LLM）对独立价值网络的需求。然而，将 GRPO 扩展到扩散模型和流匹配模型（特别是 Diffusion Transformers 或 DiTs）会引入严重的计算瓶颈。

在基于扩散的 GRPO 中，生成过程被视为一个多步决策问题，策略是在包含 $W$ 个时间步的采样轨迹上进行优化的。标准训练要求在轨迹中的每一个时间步都进行一次完整的正向和反向传播，以通过高容量的 DiT 主干网络计算策略梯度。对于高分辨率文本生成图像（T2I）任务，其轨迹通常跨越数十个步骤，这导致了极高的训练成本。与 Token 步长较轻量级的语言模型不同，现代 DiT 的每一次反向传播都会产生等同于 Transformer 正向传播的全额成本，使得“每步反向传播”的要求成为了主要的训练开销。

虽然无训练推理加速方法（如 $\Delta$ -DiT、ScalingCache）已成功利用了 DiT 隐藏状态和速度预测沿去噪轨迹平滑且近乎线性变化的观察结果来跳过中间评估，但利用这种线性特性来降低 反向传播成本 在 RL 训练中是否可行，仍然是一个悬而未决的问题。

方法论：JAGG

作者提出了 JAGG (Jacobian-Aggregated Group Gradient)，这是一种具有数学原理支持的加速方法，它能将一组 $W$ 个连续去噪步骤的反向传播次数从 $W$ 次减少到 2 次。

核心洞察

该方法依赖于这样一个洞察：去噪速度场 $v_\theta(z, t)$ 对模型参数 $\theta$ 的雅可比矩阵（Jacobian）在输入 $(z, t)$ 方面表现出近似线性，特别是在去噪过程的高噪声阶段。这使得中间步骤的雅可比矩阵可以通过仅使用端点雅可比矩阵（ $J_0$ 和 $J_{W-1}$ ）进行线性插值来得到准确近似。

算法工作流

JAGG 通过四个阶段处理一组 $W$ 个时间步：

正向传播： 模型执行所有 $W$ 步的正向传播。然而，仅对两个端点（第 0 步和第 $W-1$ 步）启用梯度追踪。中间的 $W-2$ 步使用 torch.no_grad() 执行，以节省内存和计算。
上游梯度计算： 对于每个步骤 $j$ ，创建一个模型输出的脱离（detached）代理。计算每步损失 $L_j$ （源自 PPO 式目标函数），并通过自动微分计算上游梯度 $s_j = \partial L_j / \partial \hat{v}_j$ 。至关重要的是，这一步不涉及 Transformer 反向传播；梯度仅通过轻量级的 SDE 步骤和 PPO 目标函数流动。
时间步加权聚合： 使用源自时间步 $t_j$ 的插值权重 $\alpha_j$ 将上游梯度 $s_j$ 聚合为两个复合信号 $s_{base}$ 和 $s_{corr}$ ：
$\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$
这种聚合有效地利用线性插值假设重建了完整的梯度信号。
联合反向传播： 对两个端点输出 $[\hat{v}_0, \hat{v}_{W-1}]$ 发起一次单一的 torch.autograd.backward 调用，并配合聚合后的梯度 $[s_{base}, s_{corr}]$ 。这会触发每个模块恰好一次 FSDP all-gather，取代了标准训练所需的 $W$ 次独立反向传播。

理论保证

作者提供了一个形式化证明（定理 1），指出如果速度场 $v_\theta(z, t)$ 对 $(z, t)$ 是严格线性的，则 $t$ 加权雅可比插值是精确的，近似误差为零。在此条件下，两端点反向传播可以恢复出精确的组梯度。

自适应路由 ( $jagg\_frac$ )

考虑到在线性度下降的低噪声（后期）去噪阶段，作者引入了一个路由参数 $jagg\_frac$ 。该参数仅将 JAGG 应用于前一部分去噪组（即高噪声线性假设成立的阶段，已通过梯度的余弦相似度进行了经验验证），并对剩余组使用精确的逐步反向传播。对于测试的 T2I 模型， $jagg\_frac$ 被设置为 0.6。

核心贡献

瓶颈识别： 本文正式识别并表征了由于跨采样轨迹的逐时间步反向传播导致的扩散 GRPO 训练中的计算瓶颈。
JAGG 框架： 提出了一种雅可比聚合框架，基于速度雅可比线性的严谨理论分析，将每组的反向传播次数从 $W$ 减少到 2。
理论证明： 作者证明了在速度线性条件下，所提出的插值是精确的，为该近似提供了坚实的理论基础。
经验验证： 在多个扩散主干网络（Flux, QwenImage）和 RL 算法（DanceGRPO, Flow-GRPO-Fast）上的广泛实验证明了该方法的通用性和有效性。

实验结果

作者使用 HPSv2.1、PickScore 和 CLIPScore 作为指标，在 T2I 基准测试中评估了 JAGG。

效率： 在标准的 DanceGRPO 设置下，JAGG 在每个训练步实现了 17.3% 至 18.2% 的稳定墙钟时间加速（将 QwenImage 的时间从 ~835s 降至 ~691s，将 Flux 从 ~279s 降至 ~228s）。当应用于 Flow-GRPO-Fast 时，仅更新阶段的加速比达到了 29.9% 至 30.0%。
质量： 该方法在实现这些加速的同时，生成的质量几乎没有退化。在许多情况下，JAGG 紧密追踪基准奖励曲线，并在收敛时达到或略微超过基准水平。
消融实验： 一个执行端点反向传播但未使用雅可比插值的消融变体（将完整损失独立分配给端点）的表现始终较差。这证实了通过雅可比插值聚合中间梯度信息对于稳定训练至关重要。
梯度质量： 经验分析显示，在高噪声组（0–2）中，JAGG 聚合梯度与精确梯度之间的余弦相似度较高（0.74–0.85），验证了线性假设。在随后的组中，相似度下降，从而证明了 $jagg\_frac$ 路由策略的合理性。

意义与主张

论文声称 JAGG 是迈向大规模扩散 GRPO 训练实用化的重要一步。通过利用高噪声阶段速度雅可比的近线性特性，该方法在不牺牲生成质量的前提下降低了 RL 训练的计算负担。

作者强调，JAGG 不仅仅是一个启发式的捷径，而是植根于去噪轨迹的数学属性。该方法被呈现为与现有推理加速技术正交的方法，专门针对扩散 RL 特有的反向传播瓶颈。研究结果表明，JAGG 通过显著降低训练成本壁垒，使视觉生成模型的 RL 对齐应用更加广泛。