Technical Summary: Flash-GRPO

Problem Statement

Video diffusion models have achieved significant progress in generating realistic content, yet aligning these models with human preferences (aesthetic quality, prompt adherence, physical plausibility) remains a critical challenge. While Reinforcement Learning (RL), specifically Group Relative Policy Optimization (GRPO), has emerged as a dominant paradigm for this alignment, it faces a severe computational bottleneck. Standard GRPO requires backpropagating gradients through full spatiotemporal denoising trajectories, demanding hundreds of GPU days for large-scale models (e.g., 14B parameters).

Existing efficiency methods attempt to mitigate this cost by subsampling timesteps (e.g., sliding window approaches). However, the authors identify that these naive subsampling strategies fundamentally compromise optimization. They exhibit severe training instability and fail to reach the performance ceiling of full-trajectory training. The paper attributes this failure to two root causes:

1. Timestep-Confounded Advantage Estimation: Mixing timesteps within advantage groups entangles reward variance with the intrinsic difficulty of different noise levels, obscuring the true policy signal.
2. Time-Dependent Gradient Scaling: The policy gradient inherently contains a time-dependent scaling factor derived from Stochastic Differential Equation (SDE) discretization. This factor varies by orders of magnitude across the diffusion trajectory, causing early timesteps to dominate parameter updates regardless of their actual importance, leading to optimization imbalance.

Methodology: Flash-GRPO

The authors propose Flash-GRPO, a single-step training framework designed to achieve full-trajectory performance while utilizing only one timestep per training rollout. The method addresses the identified instabilities through two core mechanisms:

1. Iso-Temporal Grouping

To eliminate timestep-confounded variance, Flash-GRPO enforces temporal consistency within prompt groups.

• Mechanism: For a given prompt, all $G$ rollouts in a group share the exact same timestep $t_k$ but are initialized with different Gaussian noise.
• Process: The group performs a single-step ODE-to-SDE transition at the assigned timestep $t_k$ to enable exploration and gradient computation. All other timesteps in the trajectory use deterministic Ordinary Differential Equation (ODE) sampling to ensure high-quality generations and accurate reward signals.
• Effect: This factorizes the advantage computation, isolating policy-induced variance from timestep-induced variance. It ensures that relative performance comparisons occur under identical denoising conditions, decoupling policy performance from timestep difficulty.

2. Temporal Gradient Rectification

To resolve gradient scale heterogeneity, the authors derive the explicit policy gradient for the reverse generation process.

• Derivation: The gradient is shown to be scaled by a time-dependent coefficient $\lambda(t) = \frac{\sqrt{\Delta t}}{\sigma_t} + \frac{\sigma_t\sqrt{\Delta t}(1-t)}{2t}$. In a single-step setting with diverse timesteps, this causes vastly inconsistent gradient magnitudes.
• Mechanism: Flash-GRPO introduces Temporal Gradient Rectification, which explicitly normalizes the gradient by $1/\lambda(t)$.
• Effect: This forces $\lambda(t) \to 1$ for all timesteps, ensuring uniform contribution to parameter updates regardless of the timestep's position in the diffusion trajectory. This eliminates discretization-induced bias and stabilizes optimization dynamics.

Key Contributions

1. Identification of Instability Sources: The paper provides theoretical derivations and empirical validation for two specific causes of instability in single-step video GRPO: timestep-confounded advantage estimation and time-dependent gradient scaling.
2. Flash-GRPO Framework: A principled single-step training framework combining iso-temporal grouping for precise advantage estimation and temporal gradient rectification for balanced optimization.
3. Scalability and Efficiency: Validation across models ranging from 1.3B to 14B parameters, demonstrating that the method eliminates the traditional efficiency-quality trade-off, achieving state-of-the-art alignment quality with substantially reduced computational costs.

Experimental Results

The authors evaluated Flash-GRPO on Wan2.1 video models (1.3B and 14B parameters) using the HPSv3 reward model and VBench metrics.

• Training Efficiency: Flash-GRPO achieves a 6× acceleration in training cost compared to standard full-trajectory methods.
• Stability: Unlike naive single-step methods (e.g., Flow-GRPO-Fast1) which suffer from catastrophic optimization collapse and reward variance, Flash-GRPO maintains stable, monotonic reward growth.
• Performance:
  • On the 1.3B model, Flash-GRPO achieved the highest Aesthetic Quality (66.43) and Subject Consistency (98.70) among compared methods, outperforming both Flow-GRPO and Flow-GRPO-Fast1.
  • On the 14B model, Flash-GRPO demonstrated consistent stability and monotonic growth, whereas Flow-GRPO exhibited slower convergence and instability.
  • Qualitative Improvements: Visual comparisons show enhanced temporal dynamics (reduced flickering), improved visual fidelity (richer details), and better prompt adherence compared to vanilla baselines.

Significance and Claims

The paper claims that Flash-GRPO resolves the fundamental trade-off between computational efficiency and alignment quality in video diffusion models. By addressing the specific mathematical artifacts of single-step optimization (confounded variance and gradient scaling), the method enables high-quality video RL alignment to become both practical and scalable. The authors assert that their approach allows models to reach the performance ceiling of full-trajectory training while operating under low computational budgets, making it a robust alternative for large-scale video model alignment.