LeanGRPO: Eliminating Redundant Recomputation in Diffusion RL
Sijie Wang, Zhiqiang Tan, Xinrui Yang and Shaohuai Shi at Harbin Institute of Technology Shenzhen remove the recomputation step that DanceGRPO and FlowGRPO perform after rollout, which is mathematically redundant when rollout and update share a backend under on-policy training.
Ask this paper
The redundancy identified: most diffusion RL methods recompute selected timesteps with gradient tracking after the rollout finishes, even though the same feed-forward backbone already produced them.
Why nobody had removed it: reusing the rollout computation graph naively incurs large memory overhead during rollout, so the fix required restructuring the data-parallel layout.
Two recompute-free schedules for different scales: LeanGRPO-Retain tracks gradients during rollout and reuses the saved graphs and activations directly in the backward pass; LeanGRPO-Reweight backpropagates each selected step immediately with a provisional advantage, delays gradient synchronization, and corrects with the true advantage once the trajectory completes.
Measured speedup: up to 1.83x end-to-end on FlowGRPO and DanceGRPO with FLUX.1-dev and Wan, with quality preserved.
Code released at github.com/coderwayne3025/LeanGRPO.
Abstract
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.