Source-linked AI summary
Flash-GRPO: Efficient Alignment for Video Diffusion via One-Step Policy Optimization
Xiaoxuan He, Siming Fu, Zeyue Xue, Weijie Wang, Ruizhe He, Yuming Li, Dacheng Yin, Shuai Dong, Haoyang Huang, Hongfa Wang, Nan Duan, Bohan Zhuang
TL;DR
Video RL alignment is expensive because full-trajectory GRPO requires dense supervision across long denoising trajectories, while existing subsampling methods compromise optimization. Flash-GRPO uses iso-temporal grouping and temporal gradient rectification for stable single-step training, achieving full-trajectory-level alignment across 1.3B to 14B models with substantially reduced cost.
Problem
Full-trajectory GRPO creates a computational bottleneck, while single-step methods suffer from timestep-confounded advantage estimation and inconsistent gradient magnitudes.
Method
Flash-GRPO combines iso-temporal grouping, which isolates timestep effects in advantage estimation, with temporal gradient rectification, which balances time-dependent gradient scaling.
Results
Flash-GRPO matches full-trajectory performance with stable optimization across video models ranging from 1.3B to 14B parameters while substantially reducing training costs.
Takeaways & Limitations
Flash-GRPO removes the efficiency-quality trade-off for practical and scalable high-quality video RL alignment.
Takeaways & Limitations
The stochastic formulation introduces time-dependent scaling factors that the method must account for during optimization.
Abstract
from arXiv · showhide
Group Relative Policy Optimization has emerged as essential for aligning video diffusion models with human preferences, but faces a critical computational bottleneck: training a 14B parametered model typically demands hundreds of GPU days per experiment. Existing efficiency methods reduce costs through sliding window subsampling training timesteps, but fundamentally compromise optimization, exhibiting severe instability and failing to reach full trajectory performance. We present Flash-GRPO, a single-step training framework that outperforms full trajectory training in alignment quality under low computational budgets while substantially improving training efficiency. Flash-GRPO addresses two critical challenges: iso-temporal grouping eliminates timestep-confounded variance by enforcing prompt-wise temporal consistency, decoupling policy performance from timestep difficulty; temporal gradient rectification neutralizes the time-dependent scaling factor that causes vastly inconsistent gradient magnitudes across timesteps. Experiments on 1.3B to 14B parameter models validate Flash-GRPO's effectiveness, demonstrating substantial training acceleration with consistent stability and state-of-the-art alignment quality.
1 Introduction
Flash-GRPO targets the computational and optimization barriers in video RL alignment by enabling stable single-step training while preserving full-trajectory performance. It combines iso-temporal grouping with temporal gradient rectification and demonstrates acceleration across model scales.
- 14B-parameter video alignment can require hundreds of GPU days per experiment, creating a scalability bottleneck.
- Sliding-window efficiency methods reduce computation but can destabilize training and fail to reach full-trajectory performance.
- Flash-GRPO uses one timestep per training while targeting full-trajectory performance.
- Iso-temporal grouping keeps rollouts for each prompt at one timestep, separating policy-induced variance from timestep-induced variance.
- Temporal gradient rectification balances optimization by addressing time-dependent gradient scaling across the diffusion trajectory.
- Experiments from 1.3B to 14B parameters show substantial acceleration, consistent stability, and alignment quality comparable to or exceeding relevant baselines.
2 Related Work
Related work establishes reinforcement learning as an important route for improving video generation quality, while highlighting the computational challenges that remain underexplored. Flash-GRPO addresses this gap with an efficient single-step framework.
- Video diffusion models generate high-fidelity, temporally coherent sequences with controllability, while human-preference alignment targets aesthetics, prompt adherence, and motion dynamics.
- Flow-GRPO and DanceGRPO extended GRPO to flow-matching models by converting deterministic ODE sampling into stochastic SDE formulations for exploration.
- Flash-GRPO's framework combines iso-temporal grouping with temporal gradient rectification to address efficiency and optimization issues.
- Computational challenges in video alignment remained largely unexplored, motivating an efficient single-step framework specifically for video diffusion models.
3 Preliminary
This section describes GRPO for flow matching and the ODE-to-SDE conversion needed for stochastic exploration. Full-trajectory objectives create the computational bottleneck that Flash-GRPO later removes.
- Flow-GRPO and DanceGRPO adapt GRPO to flow-matching models by maximizing group-relative advantages over generated trajectories.
- The full-trajectory objective aggregates clipped policy ratios across all T timesteps, requiring dense supervision.
- Deterministic ODE sampling precludes the exploration required for reinforcement learning.
- An equivalent SDE formulation injects controlled stochasticity while preserving the pretrained model's marginal distribution.
- The stochastic formulation introduces time-dependent scaling factors that later contribute to one-step gradient instability.
4 Method
Flash-GRPO enables efficient single-step GRPO alignment by addressing timestep-confounded advantage variance and time-dependent gradient-scale imbalance. Iso-temporal grouping stabilizes within-prompt comparisons, while temporal gradient rectification balances updates across diffusion timesteps.
- Iso-Temporal Grouping: Naive single-step GRPO mixes rewards from different noise levels, confounding policy quality with timestep difficulty and destabilizing advantage estimates.This undermines reliable policy optimization when samples within a prompt group use independently selected timesteps.
- Iso-Temporal Grouping: Iso-temporal grouping assigns one timestep to each prompt and shares it across that prompt’s rollouts while varying their initial noise.Different prompt groups retain temporal diversity, and only each group’s ODE-to-SDE transition timestep receives the policy gradient.
- Temporal Gradient Rectification: Time-dependent scaling factors modulate gradient magnitudes across timesteps, producing severe fluctuations unrelated to generation quality or reward strength.The imbalance arises from the discretization scheme and can cause catastrophic reward-curve collapses in baseline training.
- Temporal Gradient Rectification: Equation 10 identifies λ(t) as an implicit heterogeneous weighting factor that can vary by orders of magnitude across the diffusion trajectory.Prompts sampled at different timesteps therefore contribute inconsistent magnitudes to parameter updates.
- Temporal Gradient Rectification: Temporal Gradient Rectification rescales gradients by 1/λ(t), effectively setting λ(t) to 1 and decoupling optimization from sampler discretization.This makes prompts contribute equally to parameter updates regardless of trajectory position, improving stability and reward growth.
5 Experiment
Experiments on Wan2.1 models evaluate Flash-GRPO across VBench quality metrics, qualitative generation, ablations, reward curves, full-trajectory comparison, scalability, and motion quality. Across these settings, Flash-GRPO generally improves quality and stability while avoiding the optimization failures associated with naive single-step or sliding-window baselines.
- VBench Quality Metrics: Flash-GRPO achieves Aesthetic Quality 66.43 and Subject Consistency 98.70 on Wan2.1-T2V-1.3B after 350 GPU hours, outperforming Flow-GRPO-Fast1 and Flow-GRPO.These are the highest reported values for the two listed metrics in this comparison.
- VBench Quality Metrics: Flash-GRPO maintains Imaging Quality 68.28, compared with 65.96 for Flow-GRPO-Fast1 and 68.60 for full-trajectory Flow-GRPO.The result is reported alongside superior efficiency and is presented as avoiding the quality cost of naive subsampling.
- Visual Comparison: Qualitative comparisons show fewer flickering or unstable elements, smoother movements, more consistent facial expressions, and stronger spatial coherence with Flash-GRPO.The comparison covers savanna, animated panda, and cartoon animal scenes against vanilla Wan2.1.
- Ablation Study: Iso-temporal grouping reduces variance in credit assignment, while temporal gradient rectification removes gradient spikes caused by the time-dependent scaling factor λ(t).The ablation adds the two components incrementally over naive single-step training.
- Ablation Study: Without KL regularization, Flow-GRPO-Fast1 shows severe variance and optimization collapse, whereas Flash-GRPO achieves robust monotonic reward growth and stable convergence.The comparison is reported for HPSv3 reward curves on both training and evaluation.
- Analysis: Flash-GRPO reaches an evaluation peak reward of approximately 5.4 versus 5.1 for full-trajectory Flow-GRPO and maintains stable improvement during the compared schedule.The full-trajectory comparison is limited to the first half of the training schedule because of computational costs.
- Analysis: On the 14B Wan2.1 model, Flash-GRPO maintains consistent stability and monotonic growth, while Flow-GRPO exhibits slower growth.The paper presents this as evidence of effectiveness under low computational budgets at larger scale.
- Analysis: For Motion Quality, Flash-GRPO achieves a final score of approximately −0.28 compared with −0.34 for Flow-GRPO-Fast1.The method shows stable improvement on both training and evaluation sets, supporting improved temporal coherence.
6 Conclusion
Flash-GRPO enables stable single-step video RL alignment by addressing timestep-confounded advantage estimation and inconsistent time-dependent gradient scaling. Across 1.3B to 14B models, it substantially reduces training costs while preserving alignment quality comparable to full-trajectory methods.
- Flash-GRPO matches full-trajectory performance for video RL alignment through single-step training.
- Iso-temporal grouping prevents timestep mixing from confounding reward variance with timestep difficulty.
- Temporal gradient rectification balances update magnitudes across diffusion timesteps and stabilizes optimization.
- Experiments across 1.3B to 14B models validate the approach’s effectiveness and scalability while reducing training costs.
A More Experiments Comparison with Flow-GRPO-Fast1.
With KL regularization, Flash-GRPO converges faster and reaches a higher HPSv3 performance ceiling than Flow-GRPO-Fast1. Without KL regularization, it maintains stable gradients while the baseline exhibits severe instability.
- Approximately 5.35 HPSv3 versus 4.9 for Fast1 on the held-out set, with Flash-GRPO converging faster and reaching a higher ceiling.
- Flash-GRPO achieves faster convergence and higher performance on both training and evaluation curves than Flow-GRPO-Fast1.
- Without KL regularization, Flow-GRPO-Fast1 shows catastrophic gradient spikes and high variance, whereas Flash-GRPO maintains a consistently low, stable gradient norm.
- Temporal gradient rectification effectively regularizes the optimization landscape even without explicit KL penalties.
B Impact of Temporal Gradient Rectification.
Temporal gradient rectification produces a substantially more stable reward trajectory. Without rectification, the baseline experiences catastrophic reward collapse between 300 and 400 GPU hours.
- Between 300 and 400 GPU hours, the unrectified baseline suffers catastrophic reward collapse, while rectification yields a significantly more stable trajectory.
C Algorithm of Flash-GRPO.
Algorithm 1 samples one timestep for a prompt, performs the corresponding update within the denoising loop, computes group advantages, and optimizes the policy with the resulting loss.
- Algorithm 1 takes a prompt, group size, total timesteps, and reward models, and optimizes policy parameters θ.
- The algorithm randomly samples one timestep k for the prompt and applies the corresponding denoising update when t equals k.
- After rollout generation, it computes the mean and standard deviation of group rewards to form advantages.
- It minimizes the total loss LTGR(θ) by updating θ with a gradient step until convergence.
D More Qualitative Evaluation
Qualitative comparisons show that Flash-GRPO improves visual fidelity, motion smoothness, temporal coherence, and prompt adherence across 1.3B and 14B models and diverse content types.
- Flash-GRPO produces more coherent human motion and enhanced lighting and texture details on the 1.3B model.The waterfall scene shows improved foreground motion, while the cooking scene improves facial features, environment rendering, and watermelon-cutting action.
- On the 14B model, Flash-GRPO improves prompt following, depth of field, motion fluidity, semantic representation, and environmental detail across diverse scenes.Improvements appear in Japanese garden, bird, sailboat, cat, dog, and sword scenes.
- Flash-GRPO improves visual aesthetics, temporal coherence, and prompt adherence across different model scales and content types.
- Visualizations consistently show enhanced fidelity, smoother motion, stronger prompt following for complex prompts, and fewer visual artifacts than vanilla.