LeanGRPO: 消除扩散强化学习中的冗余重计算

# 消除扩散强化学习中的冗余重计算
来源: https://arxiv.org/html/2609.03528
谭志强 杨心瑞 史少帅*
* 作者单位：哈尔滨工业大学（深圳）计算机科学与技术学院
** 邮箱：[email protected]，[email protected]，[[email protected]](mailto:)，[email protected]
*** 代码: github.com/coderwayne3025/LeanGRPO (https://github.com/coderwayne30525/LeanGRPO)

###### 摘要
扩散强化学习（RL）最近在后训练图像和视频生成模型方面取得了显著成功。然而，大多数扩散RL方法（包括DanceGRPO和FlowGRPO）在 rollout 之后，会对选定的时间步进行带梯度追踪的重计算。在使用相同后端进行 rollout 和更新的 on-policy 训练中，这种重计算在数学上是冗余的。直观地说，rollout 和策略更新步骤可以共享相同的前向传播主干以避免冗余计算，但这可能会在 rollout 期间带来巨大的内存开销。为了解决这个问题，我们提出了 LeanGRPO，通过重构数据并行布局并引入两种无需重计算的训练调度方案来处理轨迹对数概率扩散 RL：(1) LeanGRPO-Retain 在 rollout 期间启用梯度追踪，并直接重用所得计算图和保存的激活值用于更新时的反向传播，无需任何重计算；(2) LeanGRPO-Reweight 在 rollout 期间也启用梯度，但它立即使用临时优势值反向传播每个选定步骤，并延迟梯度同步，然后在轨迹完成后使用真实优势值校正临时梯度。这些调度方案针对不同的模型规模和输入尺寸。在 FLUX.1-dev 和 Wan 上的 FlowGRPO/DanceGRPO 测试中，LeanGRPO 在保持原始优化目标不变的情况下，实现了高达 1.83 倍的端到端加速。

## 1 引言
强化学习最近已成为将扩散和流匹配生成模型与人类偏好及特定任务奖励对齐的重要工具（Ho 等人, 2020 (https://arxiv.org/html/2609.03528#bib.bib13); Song 等人, 2021 (https://arxiv.org/html/2609.03528#bib.bib42); Lipman 等人, 2023 (https://arxiv.org/html/2609.03528#bib.bib26)）。其中一种突出的方法族将去噪过程建模为策略轨迹，并通过转换对数概率来优化终端奖励。尽管这种方法有效，但这种轨迹对数概率公式在计算上仍然昂贵。代表性方法包括 DDPO（Black 等人, 2024 (https://arxiv.org/html/2609.03528#bib.bib2)）、DPOK（Fan 等人, 2023 (https://arxiv.org/html/2609.03528#bib.bib7)）、FlowGRPO（Liu 等人, 2025b (https://arxiv.org/html/2609.03528#bib.bib28)）和 DanceGRPO（Xue 等人, 2025b (https://arxiv.org/html/2609.03528#bib.bib54)），它们遵循图 1 (https://arxiv.org/html/2609.03528#S1.F1)(a) 所示的工作流程。模型首先进行不带梯度的去噪轨迹采样，以生成终端图像或视频并计算奖励。在奖励和组优势值可用后，更新阶段会重计算选定的去噪时间步，以重建反向传播所需的可微路径。当策略在 rollout 和更新之间保持不变且两个阶段使用相同的后端时，这种重计算在数学上是冗余的，因为相同的去噪预测在 rollout 期间已经计算过了。如附录 G (https://arxiv.org/html/2609.03528#A7) 详细所述，我们发现大多数公开可用的扩散 RL 实现都存在这个问题。

最直接的解决方案是在 rollout 期间启用梯度追踪。然而，在实践中，即使在使用相同后端的 on-policy 训练中，大多数现有的扩散 RL 方法仍然保留基于重计算的设计。这不仅仅是工程上的疏忽：终端优势值只有在 rollout 完成和奖励聚合后才能获得，而保留所有选定时间步的计算图及其保存的激活值直到那时，随着选定时间步数量的增加，可能会导致难以承受的内存消耗。因此，消除重计算在原理上是可行的，但核心挑战是控制由此产生的 GPU 内存开销。
参照标题图 1：原生扩散 RL 训练与无重计算的 LeanGRPO 对比。这里，$c_b$ 表示第 $b$ 个条件提示，$\epsilon_{b,m}$ 表示其第 $m$ 次生成独立采样的初始高斯潜在噪声。

因此，我们引入了 LeanGRPO，一个目标保留的系统框架，它从现有的轨迹对数概率扩散 RL 算法中消除了冗余的更新阶段重计算。我们首先重构了传统 GRPO 使用的数据并行（DP）布局（Li 等人, 2020 (https://arxiv.org/html/2609.03528#bib.bib22)）。不是将不同的提示分配给不同的 GPU 并在本地生成多个样本，而是所有 GPU 处理相同的提示，同时独立生成不同的样本。这避免了每个 GPU 需要保留多个样本的计算图和保存的激活值。基于这种共享提示布局，我们提出了两种互补的无重计算调度方案：LeanGRPO-Retain 在终端优势值可用后，保留选定时间步的 rollout 图并直接重用它们进行反向传播（图 1 (https://arxiv.org/html/2609.03528#S1.F1)(b)）。它通过简单的训练调度消除了更新阶段的重计算，但其激活内存随选定时间步数量增长。而 LeanGRPO-Reweight 在 rollout 期间通过临时反向传播消耗每个选定时间步的计算图，实现一个临时梯度，并释放其保存的激活值。在采样完成且终端奖励可用后，使用相应的优势值校正临时梯度（图 1 (https://arxiv.org/html/2609.03528#S1.F1)(c)）。通过立即消耗每个选定时间步的计算图及其保存的激活值，该调度避免了由其累积引起的内存增长。内存消耗转移到在每个等级上保留一个完整的临时梯度，直到优势值校正恢复出精确算术下相同策略梯度。

我们的贡献如下：
1. 1. 我们将更新阶段重计算识别为轨迹对数概率扩散 RL 的主要系统瓶颈：它在策略不变的情况下是数学冗余的，但又是恢复被丢弃梯度路径所必需的。
2. 2. 我们引入了 LeanGRPO，一个与大多数开源轨迹对数概率扩散 RL 算法兼容的框架。它提供了两种互补的无重计算训练调度方案：LeanGRPO-Retain 和 LeanGRPO-Reweight。
3. 3. 我们将 LeanGRPO 与多种扩散 RL 算法和模型主干网络集成。在不同模型规模和 RL 工作流上的广泛实验表明，LeanGRPO 相对于原生训练实现了高达 1.83 倍的加速。GPU 内存分析进一步描述了两种调度方案的适用场景，并为实际部署提供了指导。

## 2 相关工作
近期的工作改进了奖励分配（Huang 等人, 2026 (https://arxiv.org/html/2609.03528#bib.bib16)）、时间步选择（He 等人, 2026 (https://arxiv.org/html/2609.03528#bib.bib12)）和采样效率（Li 等人, 2026a (https://arxiv.org/html/2609.03528#bib.bib20)），但通常保留两阶段执行模式，即 rollout 不带梯度执行，选定的转换在更新阶段重计算。LeanGRPO 与这些算法改进是正交的：它保留其采样分布和策略目标，同时改变了可微转换信息被具体化和消耗的时间。分布式 RL 系统，如 VeRL-Omni（Huang 等人, 2026 (https://arxiv.org/html/2609.03528#bib.bib16)）和 BiDiRL（Tan 等人, 2026 (https://arxiv.org/html/2609.03528#bib.bib43)），优化了资源放置、专用生成后端和异步执行。它们主要改进跨阶段的协调和数据移动，而不是消除冗余的策略评估。

## 3 预备知识
轨迹对数概率扩散 RL。我们考虑针对扩散和流匹配生成模型（Ho 等人, 2020 (https://arxiv.org/html/2609.03528#bib.bib13); Song 等人, 2021 (https://arxiv.org/html/2609.03528#bib.bib42); Lipman 等人, 2023 (https://arxiv.org/html/2609.03528#bib.bib26)）的强化学习。给定一个条件输入 $c$，例如文本提示，模型生成一个潜在状态轨迹 $z_T \rightarrow z_{T-1} \rightarrow \cdots \rightarrow z_0$，其中 $z_T$ 从先验分布中采样，$z_0$ 被解码为最终图像或视频。在每个去噪步骤中，下一个潜在状态从模型的条件转换分布中采样：$z_{t-1} \sim p_{\theta}(\cdot \mid z_t, c)$，其中 $\theta$ 表示模型参数。轨迹对数概率方法将去噪过程视为多步策略，并优化采样去噪转换的对数概率（Black 等人, 2024 (https://arxiv.org/html/2609.03528#bib.bib2); Fan 等人, 2023 (https://arxiv.org/html/2609.03528#bib.bib7); Liu 等人, 2025b (https://arxiv.org/html/2609.03528#bib.bib28); Xue 等人, 2025b (https://arxiv.org/html/2609.03528#bib.bib54)）。令 $i$ 索引一个采样轨迹，$t$ 索引其去噪时间步。在 rollout 期间，实现会记录每步转换对数概率 $\ell^{\mathrm{roll}}_{i,t} = \log p_{\theta_{\mathrm{roll}}}(z^i_{t-1} \mid z^i_t, c_i), \quad t=1,\ldots,T$ (1)，以及相应的潜在状态对 $(z^i_t, z^i_{t-1})$。Rollout 之后，计算终端奖励并在每个提示组内进行归一化以获得优势值 $A_i$（Shao 等人, 2024 (https://arxiv.org/html/2609.03528#bib.bib40)）。然后，在选定的时间步子集 $\mathcal{S}$ 上执行训练，该子集可以通过时间步比例或随机微分方程（SDE）窗口选择（Xue 等人, 2025b (https://arxiv.org/html/2609.03528#bib.bib54); Liu 等人, 2025b (https://arxiv.org/html/2609.03528#bib.bib28)）。在训练期间，每个选定时间步在更新策略下被重计算，以获得用于 GRPO 比率的可微对数概率。令 $\ell^{\mathrm{upd}}_{i,t} = \log p_{\theta_{\mathrm{upd}}}(z^i_{t-1} \mid z^i_t, c_i)$ (2)，其中 $\ell^{\mathrm{upd}}_{i,t}$ 是通过重计算选定转换获得的可微当前对数概率。在 PPO/GRPO 目标中常用的概率比率（Schulman 等人, 2017 (https://arxiv.org/html/2609.03528#bib.bib39); Shao 等人, 2024 (https://arxiv.org/html/2609.03528#bib.bib40)）为 $\rho_{i,t} = \frac{p_{\theta_{\mathrm{upd}}}(z^i_{t-1} \mid z^i_t, c_i)}{p_{\theta_{\mathrm{roll}}}(z^i_{t-1} \mid z^i_t, c_i)} = \exp\left(\ell^{\mathrm{upd}}_{i,t} - \ell^{\mathrm{roll}}_{i,t}\right)$ (3)，而带有裁剪的 GRPO/PPO 风格损失可以写为 $\mathcal{L}_{\mathrm{clip}} = \sum_{i} \sum_{t \in \mathcal{S}} \max\left(-A_i \rho_{i,t}, -A_i \operatorname{clip}(\rho_{i,t}, 1-\epsilon, 1+\epsilon)\right)$ (4)。这里 $\epsilon$ 表示用于限制概率比率的裁剪范围。在本文分析的 on-policy 单次更新设置中，rollout 和更新使用相同的后端，并且在优化器步骤之前策略参数未改变。因此，$\theta_{\mathrm{upd}} = \theta_{\mathrm{roll}}$。在这个 on-policy 点，$\ell^{\mathrm{upd}}_{i,t}$ 和 $\ell^{\mathrm{roll}}_{i,t}$ 在数值上一致，得到 $\rho_{i,t}=1$。由于 $\rho_{i,t}=1$ 位于裁剪区间内，裁剪不活跃。此外，$\ell^{\mathrm{roll}}_{i,t}$ 是在无梯度模式下记录的，因此满足 $\nabla_{\theta} \ell^{\mathrm{roll}}_{i,t} = 0$。利用 $\nabla_{\theta} \rho_{i,t} = \rho_{i,t} \nabla_{\theta} (\ell^{\mathrm{upd}}_{i,t} - \ell^{\mathrm{roll}}_{i,t})$，策略梯度变为 $G^{\mathrm{native}} = \nabla_{\theta} \mathcal{L}_{\mathrm{clip}} = -\sum_{i} \sum_{t \in \mathcal{S}} A_i \nabla_{\theta} \rho_{i,t} = -\sum_{i} \sum_{t \in \mathcal{S}} A_i \nabla_{\theta} \ell^{\mathrm{upd}}_{i,t} = \sum_{i} G_i^{\mathrm{native}}$ (5)。

## 4 方法
### 4.1 共享提示梯度启用的 Rollout
两种 LeanGRPO 变体都在 rollout 期间对选定的去噪时间步进行带梯度追踪的评估。在终端优势值可用之前，LeanGRPO-Retain 保留每个样本的计算图和保存的激活值，而 LeanGRPO-Reweight 保留每个样本的临时梯度。我们将梯度启用 rollout 期间产生的可微对数概率表示为 $\ell^{\mathrm{grad-roll}}_{i,t}$，以区别于第 3 节 (https://arxiv.org/html/2609.03528#S3) 中定义的 detached rollout 记录 $\ell^{\mathrm{roll}}_{i,t}$。在传统的数据并行布局下，每个等级为他们自己的每个提示生成 $M$ 个样本。对于 LeanGRPO-Retain，这意味着需要为每个提示保留 $M$ 个不同样本的计算图和保存的激活值；对于 LeanGRPO-Reweight，则意味着需要维护每个提示 $M$ 个不同样本的临时梯度。两者都可能导致每个 GPU 难以承受的内存消耗。相反，我们将 $M$ 次生成分布在 $R$ 个等级上，每个等级并行采样相同的提示。当 $Q = M/R \geq 1$（为方便起见，我们在此假设 $Q$ 是整数）时，每个等级仅保留 $Q$ 个样本，即传统数据并行所需激活值或临时梯度的 $1/R$。当一个提示的所有生成完成后，LeanGRPO 在所有等级上收集 $M$ 个奖励以形成完整的提示组，并计算相应的优势值。与采样所有提示后再训练的传统扩散 RL 不同，LeanGRPO 在完成每个提示组后立即开始训练。这避免了累积多个提示组带来的内存压力。当 $Q < 1$ 且 $P = R/M$ 时，我们将工作者分为 $P$ 组，每组 $M$ 个等级。每个工作者组内的等级为同一提示生成不同的样本，而不同的工作者组并行处理不同的提示。在这两种情况下，这种布局在等级间分摊了可微状态，大大减轻了每个等级的内存压力。这种重分配在梯度上与原生数据并行等效，如附录 A (https://arxiv.org/html/2609.03528#A1) 所证明。

### 4.2 LeanGRPO-Retain：带图保留的反向传播
LeanGRPO-Retain 直接保留 rollout 期间产生的可微信息，如图 1 (https://arxiv.org/html/2609.03528#S1.F1)(b) 所示。对于提示 $b$ 的第 $m \in \{1,\ldots,M\}$ 次生成，每个选定时间步的去噪前向传播……