Search papers, labs, and topics across Lattice.
LeanGRPO addresses the inefficiencies in diffusion reinforcement learning (RL) by eliminating redundant recomputation during the rollout and update phases. By restructuring the data-parallel layout and introducing two innovative training schedules鈥擫eanGRPO-Retain and LeanGRPO-Reweight鈥攖he method allows for gradient tracking without the need for recomputation, significantly reducing memory overhead. The results demonstrate up to 1.83x end-to-end speedup compared to existing methods while maintaining the original optimization objectives.
LeanGRPO achieves up to 1.83x speedup in diffusion RL without sacrificing optimization quality by eliminating redundant computations.
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.
