PaperlayerView exact revision on arXiv

Source-linked AI summary

LeanGRPO: Eliminating Redundant Recomputation in Diffusion RL

Sijie Wang, Zhiqiang Tan, Xinrui Yang, Shaohuai Shi

arXiv:2609.03528v1cs.LGcs.AIcs.AR

TL;DR

Trajectory-logprob diffusion RL often recomputes selected timesteps after rollout, although this is redundant in on-policy training with an unchanged backend, while retaining rollout graphs can create prohibitive memory overhead. LeanGRPO restructures data parallelism and introduces Retain and Reweight schedules that remove recomputation while managing different memory tradeoffs. Across multiple algorithms and backbones, it preserves reward improvement and reaches up to 1.83× end-to-end speedup.

  • Problem

    Trajectory-logprob diffusion RL retains a recompute-based rollout/update workflow, creating redundant update-stage computation and substantial memory pressure when gradient-enabled rollout graphs are retained.

  • Method

    LeanGRPO distributes samples for shared prompts across ranks and introduces Retain, which reuses rollout graphs, and Reweight, which performs provisional backward then corrects gradients before synchronization.

  • Results

    1.83× is the maximum end-to-end speedup reported across multiple diffusion RL algorithms and FLUX.1-dev, SD3.5, and Wan backbones while preserving reward improvement.

  • Takeaways & Limitations

    LeanGRPO provides an objective-preserving execution framework whose complementary schedules target different model scales, tensor sizes, and memory constraints.

  • Takeaways & Limitations

    Retain’s activation memory grows with selected timesteps, whereas Reweight adds one full gradient per rank and is preferable for larger video tensors.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Trajectory-logprob diffusion RL is effective but recomputes selected denoising steps during update, even when on-policy rollout and update use the same unchanged backend. LeanGRPO restructures execution and introduces Retain and Reweight schedules to remove this redundant recomputation while managing memory.

  • Motivation: Trajectory-logprob diffusion RL samples denoising trajectories without gradients, computes terminal rewards, then recomputes selected timesteps with gradients during update.This workflow is followed by representative methods including DDPO, DPOK, FlowGRPO, and DanceGRPO.
  • Motivation: Gradient-enabled rollout is mathematically feasible, but retaining all selected-timestep graphs and activations until terminal advantages are available can exhaust GPU memory.The overhead grows as the number of selected timesteps increases.
  • Framework: LeanGRPO restructures data parallelism so all GPUs process the same prompt while independently generating different samples, reducing per-GPU retained graph or activation state.The layout distributes samples for one prompt across ranks rather than assigning different prompts to different GPUs.
  • Schedules: LeanGRPO-Retain preserves selected-timestep rollout graphs and reuses them for backward after terminal advantages become available, eliminating update-stage recomputation.Its activation memory still grows with the number of selected timesteps.
  • Schedules: LeanGRPO-Reweight immediately backpropagates each selected timestep with a provisional advantage, releases activations, and corrects the provisional gradient after reward completion.It trades activation accumulation for retaining a full provisional gradient on each rank until correction.
  • Results: 1.83× is the maximum reported speedup over native training across diffusion RL algorithms and model backbones.The framework is integrated with multiple algorithms and backbones, with GPU profiling used to characterize schedule applicability.

2 RELATED WORK

LeanGRPO targets redundant policy evaluation in trajectory-logprob diffusion RL and is orthogonal to algorithmic improvements in reward assignment, timestep selection, and sampling efficiency.

  • Positioning: LeanGRPO preserves sampling distributions and policy objectives while changing when differentiable transition information is materialized and consumed.This makes it orthogonal to recent improvements in reward assignment, timestep selection, and sampling efficiency.
  • Positioning: Related distributed RL systems primarily optimize resource placement, generation backends, asynchronous execution, orchestration, and data movement rather than eliminating redundant policy evaluations.LeanGRPO addresses the redundant evaluation itself.

3 PRELIMINARIES

Trajectory-logprob diffusion RL models denoising as a multi-step policy, records transition log-probabilities during rollout, and recomputes selected transitions for a GRPO-style update. In the on-policy unchanged-policy setting, the current and rollout log-probabilities coincide, making recomputation redundant.

  • Trajectory-Logprob Diffusion RL: A diffusion or flow-matching model generates latent states z_T → z_T−1 → ··· → z_0, with z_0 decoded into an image or video.Each transition samples z_t−1 from the model’s conditional distribution given z_t and c.
  • Trajectory-Logprob Diffusion RL: Trajectory-logprob methods optimize the log-probabilities of sampled denoising transitions, recording per-step probabilities and latent pairs during rollout.The formulation treats the denoising process as a multi-step policy.
  • Training Objective: After rollout, terminal rewards are normalized within prompt groups to obtain advantages, and selected timesteps are recomputed under the update policy for differentiable GRPO ratios.The selected subset may be defined by a timestep fraction or an SDE window.
  • On-Policy Redundancy: In on-policy single-update training with the same backend and unchanged parameters, θ_upd = θ_roll and ρ_i,t = 1, so clipping is inactive.The rollout and update log-probabilities coincide in value.

4 METHODS

LeanGRPO enables gradient tracking during rollout and combines shared-prompt data parallelism with two memory-aware schedules. Retain delays backward while preserving graphs; Reweight consumes graphs immediately, corrects gradients later, and synchronizes only after correction.

  • Shared-Prompt Rollout: Both schedules compute selected denoising timesteps with gradient tracking during rollout, preserving either computation graphs and activations or provisional gradients.The differentiable rollout log-probability is distinct from the detached rollout record.
  • Shared-Prompt Rollout: Conventional data parallelism can require each rank to retain state for M samples, causing prohibitive memory consumption for both schedules.LeanGRPO instead distributes the M generations across R ranks so each rank retains Q = M/R samples.
  • LeanGRPO-Retain: LeanGRPO-Retain retains every selected-timestep graph and saved activation, then backpropagates after rewards are gathered and the trajectory advantage becomes available.The retained rollout information is reused directly for backward.
  • LeanGRPO-Retain: Retain’s activation memory grows approximately linearly with |S|, making it most suitable for small timestep fractions or moderate-resolution workloads.This memory limitation motivates Reweight.
  • LeanGRPO-Reweight: LeanGRPO-Reweight immediately consumes each selected-timestep graph through provisional backward, releases activations, and corrects the sample-local gradient after the terminal advantage is known.The corrected gradients are synchronized only after correction.
  • Within-Sample Correction: Because all selected timesteps in one trajectory share an advantage, scaling their accumulated provisional gradient equals scaling each timestep gradient before summation.Thus delayed correction preserves the resulting policy gradient.
  • Distributed Synchronization: Different samples require correction before distributed synchronization because a common scalar cannot recover independently advantage-weighted gradients after sample-wise decomposition is discarded.LeanGRPO-Reweight therefore disables synchronization during provisional backward and performs one reduce-scatter after correction.

5 EVALUATION

LeanGRPO improves end-to-end diffusion RL efficiency by eliminating update-stage recomputation, while exposing memory tradeoffs between Retain and Reweight across model and training settings.

  • End-to-end training efficiency: Speedup increases with the timestep fraction because recomputation occupies a larger share of update time.
  • End-to-end training efficiency: Under BF16 full fine-tuning, Reweight achieves 1.44×–1.81× speedups and Retain achieves 1.18×–1.29×.
  • Time breakdown: Native spends 402.45 seconds on recomputation, while Reweight reduces training time to 535.15 seconds and reaches a 1.81× end-to-end speedup.
  • Ablation study: Coalesced synchronization yields a 1.33× speedup, while Retain achieves a 1.29× speedup by eliminating recomputation.
  • Training convergence: Both LeanGRPO variants preserve Native’s overall reward improvement trend and reach target rewards faster on Wan2.1-1.3B and FLUX.1-dev.Retain is 1.14× faster on Wan2.1-1.3B and both variants are 1.46× faster on FLUX.1-dev.
  • GPU memory usage: Retain’s memory grows with retained graphs, whereas Reweight’s memory remains stable across selected timesteps in large-tensor video training.Retain reaches OOM at the sixth selected timestep, while Reweight uses more memory and achieves a 1.22× speedup.

6 CONCLUSION

LeanGRPO is an objective-preserving framework that removes redundant update-stage recomputation in on-policy trajectory-logprob diffusion RL while retaining the original policy gradient.

  • Conclusion: LeanGRPO eliminates redundant update-stage recomputation in on-policy trajectory-logprob diffusion RL.
  • Conclusion: Across multiple algorithms and FLUX.1-dev, SD3.5, and Wan backbones, LeanGRPO preserves reward improvement and achieves end-to-end speedups of up to 1.83×.
  • Gradient equivalence: Shared-prompt parallelism preserves every term in the native gradient while changing where each sample is evaluated.
  • Gradient equivalence: Under unchanged-policy and common-backend assumptions, gradient-enabled rollout log-probabilities match update-stage values and parameter gradients.
  • Delayed advantage correction: LeanGRPO-Reweight uses provisional gradients before advantages are available, with the provisional advantage set to one.

B.2 EQUIVALENCE OF DELAYED ADVANTAGE CORRECTION

Delayed advantage correction reconstructs the native gradient by applying each trajectory’s terminal advantage after provisional per-timestep gradients have been accumulated.

  • Delayed advantage correction: All selected timesteps in a trajectory receive the same terminal advantage, so provisional backward accumulates an advantage-independent gradient.
  • Delayed advantage correction: Once the terminal advantage is available, local correction applies it to the provisional gradient.
  • Delayed advantage correction: Summing corrected gradients over prompt-generation pairs recovers the native global gradient.
  • Delayed advantage correction: The equivalence is unchanged by common batch or timestep normalization factors.

B.3 CORRECTNESS OF DELAYED REDUCE-SCATTER

LeanGRPO-Reweight preserves exact distributed gradients by delaying synchronization until each rank applies its sample-specific advantage. This changes communication order without changing the optimizer gradient.

  • Premature synchronization discards sample-wise gradient terms needed for distinct advantage weighting.
  • Delaying communication preserves the shard of the independently advantage-weighted prompt-group gradient.The subsequent reduce-scatter operates on the locally corrected aggregate.

C WHY POST-HOC VECTOR CORRECTION CANNOT RECOVER THE EXACT GRADIENT

Post-hoc correction cannot generally recover the exact advantage-weighted gradient after synchronization has mixed sample-specific gradients. Exactness therefore requires preserving sample-specific information until advantage weighting.

  • No common scalar, element-wise vector, matrix, or nonlinear postprocessing function can generally reconstruct the exact target from the mixed gradient alone.An element-wise construction can also be circular or impossible when a mixed-gradient coordinate is zero while the corrected coordinate is nonzero.
  • Synchronization retains only one linear combination of sample gradients, while correction requires a different advantage-weighted combination.
  • When advantages differ, distinct sample-gradient decompositions can produce the same mixed gradient but different corrected gradients.
  • Exact correction requires retaining sample-specific gradient information or an already advantage-weighted sum before synchronization.LeanGRPO-Reweight preserves exactness by delaying synchronization until local provisional gradients have been advantage-weighted.

D WHY POLICY-GRADIENT RL USES UPDATE-STAGE RECOMPUTATION

Trajectory-logprob diffusion RL delays terminal advantages until the full trajectory and reward comparison are complete, creating a tension between immediate backward and activation memory. Native methods therefore re-forward selected transitions after rollout, while LeanGRPO removes that separate update-stage evaluation without necessarily removing checkpoint recomputation.

  • Terminal advantages become available only after the full trajectory is sampled, decoded, rewarded, and compared within the prompt group.
  • Native training re-evaluates selected transitions after rollout to reconstruct the differentiable path required for backward.Under the on-policy common-backend setting, this re-forward reproduces the rollout prediction numerically.
  • Activation checkpointing is an internal backward memory–computation trade-off, distinct from update-stage transition re-forwarding.
  • LeanGRPO eliminates separate update-stage transition re-forwarding but may still recompute checkpointed blocks during backward.This distinction defines “recompute-free” as avoiding a second update-stage evaluation of the sampled transition.

E NUMERICAL ALIGNMENT WITH NATIVE RECOMPUTATION

The numerical evaluation compares LeanGRPO with Native under controlled fixed trajectories and real-reward training. Retain can match Native bitwise in controlled settings, while Reweight remains closely aligned but exhibits finite-precision differences tied to delayed advantage scaling.

  • Identical transitions, rewards, and advantages isolate the gradient execution schedule in the controlled FLUX.1-dev comparison.
  • Reweight’s remaining numerical difference primarily arises because Native scales the loss before BF16 backward, whereas Reweight scales a formed gradient after the advantage arrives.Coalesced reduction order contributes a smaller additional effect, and reproducing Native’s exact rounding would require retaining per-sample backward operands or recomputing them.
  • Retain can reproduce Native bitwise when trajectory inputs, kernels, hook order, timestep order, and distributed transaction boundaries are controlled.
  • Reweight preserves the same gradient in exact arithmetic but remains closely rather than bitwise aligned under finite-precision backward.
  • Retain is more memory-sensitive to selected timesteps and high-resolution workloads, whereas Reweight’s full-gradient storage is largely unaffected by those factors.Larger generation counts and micro-batches create additional pressure for Reweight because it retains one full gradient per sample.

DIFFUSION RL IMPLEMENTATIONS

The audit identifies how widely LeanGRPO’s target transition recomputation and graph-colocation conditions occur across diffusion RL implementations. It also delineates the schedules, optimizer-step assumptions, and backend boundaries under which LeanGRPO applies.

  • Audit protocol: The audit counts a method only when publicly inspectable end-to-end code supports online reward optimization for continuous image or video diffusion or flow-matching models.Paper-only methods, incomplete training paths, offline preference-only methods, and infrastructure ports without distinct algorithms are excluded.
  • Operational definition: LeanGRPO’s operational target is rollout without a retained differentiable policy graph followed by update-time recomputation of the same sampled transition.An arbitrary update-time forward is insufficient; the recomputation must form an advantage-weighted transition-policy gradient or equivalent parameterization.
  • Audit coverage: 23 of 32 audited implementations contain the target transition recomputation, including DDPO, DPOK, FlowGRPO, and DanceGRPO.The remaining 9 optimize other objectives and lack the update-stage transition recomputation targeted by LeanGRPO.
  • Audit coverage: 22 of 23 target implementations are graph-colocated for LeanGRPO’s first on-policy update, representing 95.7% of the target class and 68.8% of the audited set.B2-DiffuRL has the relevant recomputation but separates sampling and training into different programs, requiring process restructuring.
  • Schedule assumptions: Exact graph reuse depends on the optimizer-step boundary: gradients accumulated before the next optimizer step remain on-policy, whereas later updates require current-policy recomputation.LeanGRPO fully supports single-policy-snapshot schedules and provides a hybrid for replay or multi-epoch methods.
  • Schedule assumptions: Repeated same-batch updates are a specialized off-policy regime requiring safeguards, since Flow-DPPO reports plateauing or degradation for some two-step Flow-GRPO and Flow-CPS settings.LeanGRPO can eliminate recomputation for the first on-policy update, while subsequent current-policy recomputation remains necessary.
  • Scope boundaries: LeanGRPO primarily targets shared rollout-update execution paths with an unchanged policy before the optimizer step.Different backends break graph colocation, and methods such as DiffusionNFT fall outside scope because they do not reconstruct sampled reverse transition log-probabilities.
  • Future work: Users currently select the execution schedule manually through hyperparameters, leaving automatic hardware- and workload-aware schedule selection for future work.The paper identifies automatic schedule selection as a main future direction.
Loading 2609.03528v1…