arXiv:2609.03528cs.LGcs.AI2026-09

提出新训练方法,让扩散强化学习无需重复计算,速度提升1.83倍。

LeanGRPO: Eliminating Redundant Recomputation in Diffusion RL

论文配图:LeanGRPO: Eliminating Redundant Recomputation in Diffusion RL
图 1 · 摘自论文原文
  • 重构数据并行布局,避免滚动和更新时重复计算梯度。
  • 两种无重计算调度方案,分别适配不同规模模型与输入大小。
  • 在FLUX.1-dev和Wan数据集上实现最高1.83倍加速,保持原优化目标。

扩散强化学习(Diffusion RL)在后训练图像与视频生成模型中取得显著进展。然而,多数扩散RL方法(如DanceGRPO和FlowGRPO)在轨迹采样后仍需对选定时间步重新计算并追踪梯度。在采用相同后端进行采样与更新的在线策略训练中,这一重计算在数学上是冗余的。直观上,采样与策略更新可共用前向主干网络以避免重复计算,但直接复用会带来较大的内存开销。为此,我们提出LeanGRPO,通过重构数据并行布局,并引入两种无重计算的轨迹-对数概率扩散强化学习训练策略:(1) LeanGRPO-Retain 在采样期间保留梯度追踪,直接复用生成的计算图与保存的激活值用于反向传播更新,无需重计算;(2) LeanGRPO-Reweight 同样在采样中启用梯度,但立即使用临时优势值对每个选定时步反向传播,延迟梯度同步,待轨迹完成后用真实优势值修正临时梯度。这两种策略针对不同模型规模与输入尺寸进行了优化。在FlowGRPO/DanceGRPO结合FLUX.1-dev与Wan数据集的实验中,LeanGRPO实现了最高达1.83倍的端到端加速,同时保持原始优化目标不变。

原文摘要 · Abstract (English)

Diffusion reinforcement learning (RL) has recently achieved significant success in post-training image and video generative models. However, most diffusion RL methods, including DanceGRPO and FlowGRPO, recompute selected timesteps with gradient tracking after rollout. Under on-policy training with the same backend for rollout and update, this recomputation is mathematically redundant. Intuitively, the rollout and policy update steps can reuse the same feed-forward backbone to avoid redundant computation, but doing so can incur a large memory overhead during rollout. To address the issue, we present LeanGRPO by restructuring the data-parallel layout and introducing two recompute-free training schedules for trajectory-logprob diffusion RL: (1) LeanGRPO-Retain enables gradient tracking during rollout and directly reuses the resulting computation graphs and saved activations for backward during update, requiring no recomputation; and (2) LeanGRPO-Reweight also enables gradients during rollout, but immediately backpropagates each selected step using a provisional advantage and delays gradient synchronization, then corrects the provisional gradients with the true advantage after the trajectory is completed. These schedules target different model scales and input sizes. Across FlowGRPO/DanceGRPO with FLUX.1-dev and Wan, LeanGRPO achieves up to 1.83x end-to-end speedup while preserving the original optimization objective.

扩散模型强化学习加速训练

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。