LeanGRPO: Eliminating Redundant Recomputation in Diffusion RL

• URL: http://arxiv.org/abs/2609.03528v1
• Date: Thu, 03 Sep 2026 08:24:56 GMT
• Title: LeanGRPO: Eliminating Redundant Recomputation in Diffusion RL
• Abstract summary: Diffusion reinforcement learning (RL) has recently achieved significant success in post-training image and video generative models.<n>Most RL methods, including DanceGRPO and FlowGRPO, recompute selected timesteps with gradient tracking after rollout.<n>We present LeanGRPO by restructuring the data-parallel layout and introducing two recompute-free training schedules for trajectory-logprob computation RL.
• Score: 18.67514175891835
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• ReFPO: Reflow Regularization for Flow Matching Policy Gradients [58.32178725687043]
  We present a simple online RL method that adds explicit Reflow regularization to Flow Matching Policy Gradients.<n>We experimentally demonstrate that ReFPO improves average performance and discretization across GridWorld, MuJoCo Playground, and high-dimensional Humanoid Control tasks.
  arXiv  Detail & Related papers  (2026-06-19T04:23:10Z)
• Reversal Q-Learning [57.43956630070019]
  Reversal Q-learning is an off-policy RL algorithm that trains a flow policy based on prior data.<n>We show that RQL leads to the best average offline RL performance compared to state-of-the-art flow-based offline RL algorithms.
  arXiv  Detail & Related papers  (2026-06-16T05:56:10Z)
• Learn Where Outcomes Diverge: Efficient VLA RL via Probabilistic Chunk Masking [5.238545250784642]
  This paper presents Probabilistic Chunk Masking (PCM), a drop-in modification to GRPO that allocates gradient to a small, probabilistically selected subset of chunks per trajectory.<n>On three LIBERO benchmarks, PCM matches the final success rate of standard GRPO while achieving 2.38 times wall-clock speedup, 4.8 times faster gradient updates, and 60% lower peak activation memory.
  arXiv  Detail & Related papers  (2026-05-15T16:33:59Z)
• Off-Policy Value-Based Reinforcement Learning for Large Language Models [25.962820072445222]
  ReVal is a Bellman-update-based method that combines stepwise signals capturing internal consistency with trajectory-level signals derived from outcome verification.<n>On DeepSeek-R1-Distill-1.5B, ReVal improves training efficiency and achieves improvement of 2.7% in AIME24 and 4.5% in out-of-domain benchmark GPQA over GRPO.
  arXiv  Detail & Related papers  (2026-03-24T15:55:02Z)
• DyJR: Preserving Diversity in Reinforcement Learning with Verifiable Rewards via Dynamic Jensen-Shannon Replay [57.80564154223355]
  Existing experience replay methods address this by reusing accurate samples for direct policy updates.<n>We argue that historical data should prioritize sustaining diversity rather than simply reinforcing accuracy.<n>We propose Dynamic Jensen-Shannon Replay (LEPJR), a simple yet effective regularization framework.
  arXiv  Detail & Related papers  (2026-03-17T06:20:56Z)
• Parallel Diffusion Solver via Residual Dirichlet Policy Optimization [88.7827307535107]
  Diffusion models (DMs) have achieved state-of-the-art generative performance but suffer from high sampling latency due to their sequential denoising nature.<n>Existing solver-based acceleration methods often face significant image quality degradation under a low-dimensional budget.<n>We propose the Ensemble Parallel Direction solver (dubbed as EPD-EPr), a novel ODE solver that mitigates these errors by incorporating multiple gradient parallel evaluations in each step.
  arXiv  Detail & Related papers  (2025-12-28T05:48:55Z)
• CoPRIS: Efficient and Stable Reinforcement Learning via Concurrency-Controlled Partial Rollout with Importance Sampling [11.252930904797]
  We propose Concurrency- Controlled Partial Rollout with Importance Sampling (CoPRIS)<n>CoPRIS mitigates long-tail inefficiencies by maintaining a fixed number of concurrent rollouts, early-terminating once sufficient samples are collected, and reusing unfinished trajectories in subsequent rollouts.<n>Experiments show that CoPRIS achieves up to 1.94x faster training while maintaining comparable or superior performance to synchronous RL systems.
  arXiv  Detail & Related papers  (2025-11-05T11:39:32Z)
• BranchGRPO: Stable and Efficient GRPO with Structured Branching in Diffusion Models [57.304411396229035]
  We present BranchGRPO, a method that restructures the rollout process into a branching tree.<n>On HPDv2.1 image alignment, BranchGRPO improves alignment scores by up to textbf16% over DanceGRPO.<n>A hybrid variant, BranchGRPO-Mix, further accelerates training to 4.7x faster than DanceGRPO without degrading alignment.
  arXiv  Detail & Related papers  (2025-09-07T12:53:06Z)
• Shuffle-R1: Efficient RL framework for Multimodal Large Language Models via Data-centric Dynamic Shuffle [65.14124923451077]
  Reinforcement learning (RL) has emerged as an effective post-training paradigm for enhancing the reasoning capabilities of multimodal large language model (MLLM)<n>However, current RL pipelines often suffer from training inefficiencies caused by two underexplored issues: Advantage Collapsing and Rollout Silencing.<n>We propose Shuffle-R1, a simple yet principled framework that improves RL fine-tuning efficiency by dynamically restructuring trajectory sampling and batch composition.
  arXiv  Detail & Related papers  (2025-08-07T17:53:47Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.