Flash-GRPO: Efficient Alignment for Video Diffusion via One-Step Policy Optimization
Paper: arXiv:2605.15980v1 Code: Shredded-Pork/Flash-GRPO Code reference:
main@42110137(2026-05-18) Project: shredded-pork.github.io/Flash-GRPO.github.io
1. Motivation (研究动机)
视频 Diffusion / Flow Matching 模型已经能生成高保真、时序一致的视频,但在审美质量、prompt adherence、物理合理性与 motion dynamics 上继续对齐人类偏好仍然依赖 RL。Flow-GRPO、DanceGRPO 等方法把 GRPO 迁移到视觉生成后,主要瓶颈变成训练成本:标准 GRPO 对每个 rollout 的整条 denoising trajectory 做 dense supervision,视频 latent 的时空维度很大,因此一次 14B 视频模型 alignment 往往需要数百 GPU days,研究迭代和部署都被成本限制。
现有效率化路线(如 Flow-GRPO-Fast、MixGRPO)通过 sliding window / 少量连续 timestep subsampling 降低反传成本,但论文指出这种做法不是“免费加速”:当窗口缩到 one-step 时,优化会严重不稳定,reward curve 可能 collapse,且很难达到 full-trajectory GRPO 的上限。Flash-GRPO 想解决的具体问题是:只在一个 timestep 上做 policy update,却不牺牲 full-trajectory alignment 的稳定性和质量。关键动机不是单纯“少算几个 timestep”,而是解释为什么 naive one-step GRPO 会失败。
作者归纳出两个根因:第一,组内 rollout 被分配到不同噪声水平时,reward variance 同时混入 policy quality 与 timestep difficulty,advantage 不再是干净的相对偏好信号;第二,SDE 离散化会在 policy gradient 中引入随 timestep 变化的缩放因子,导致不同 timestep 的梯度量级相差很大。若解决这两点,video RL alignment 可以从 full-trajectory dense training 转成更接近预训练单步 timestep training 的范式。

Figure 1 解读:左侧展示 Flash-GRPO 在 Motion、Aesthetic、Prompt Following 三类样例上的定性提升;右上 reward curve 显示它比 Flow-GRPO 更稳定地单调上升;右下给出核心效率结论:Flash-GRPO 在训练成本上达到 acceleration,同时取得更高 evaluation performance。这个图把本文问题定义得很清楚:目标不是牺牲质量换速度,而是在低计算预算下超过 full trajectory baseline。
2. Idea (核心思想)
核心 insight:one-step video GRPO 失败的主因不是“只看一个 timestep”本身,而是组内 credit assignment 被 timestep 难度污染以及不同 timestep 的 policy gradient 被 SDE scale 隐式加权。Flash-GRPO 因此把单步训练拆成两个互补校正:Iso-Temporal Grouping 保证同一 prompt 的 个 rollout 在同一 timestep 比较,只让 noise / generation quality 变化;
Temporal Gradient Rectification 则除掉时间相关的梯度缩放 ,让不同 timestep 的 prompt 对参数更新贡献均衡。与 Flow-GRPO-Fast / MixGRPO 的根本差异在于:这些效率方法主要做 timestep 子采样或滑窗,默认原始 GRPO 目标在少步上仍然可用;Flash-GRPO 则认为少步训练必须重新处理统计比较和梯度尺度,否则 one-step 会产生 confounded advantage 和 biased gradient magnitude。
它保留全局 batch 里的 temporal diversity,但在每个 prompt group 内固定 timestep,这相当于把“时刻难度”从组内相对优势估计里剥离出去。从实现角度看,Flash-GRPO 只在一个 ODE-to-SDE transition timestep 上采样 stochastic branch、记录 log probability 并反传;其它 timestep 使用 deterministic ODE 继续生成,以获得更高质量、更可靠的 final reward。这使 reward 仍然来自完整生成视频,而梯度只落在选中的 one-step transition 上。
3. Method (方法)
3.1 Overall framework

Figure 2 解读:左侧是 Iso-Temporal Grouping:每个 prompt group 只选一个 timestep 做 ODE-to-SDE transition,组内 rollout 共享该 timestep 但使用不同初始噪声;未选中的 timestep 用 deterministic ODE 走完整生成轨迹。右侧是 Temporal Gradient Rectification:SDE 离散化产生时间相关缩放 ,Flash-GRPO 用 抵消不同 timestep 的梯度量级差异。整套框架把 reward 评估保持在完整视频级别,把训练成本压到单步反传。
直觉上,GRPO 的 advantage 是“同一 prompt 下哪个样本更好”的相对比较。如果组内样本来自不同 timestep,reward 高低可能只是因为某个噪声水平更容易生成好视频,而不是 policy 的局部决策更优;这会让 advantage 成为带偏的 credit signal。Flash-GRPO 让组内样本共享 timestep,相当于控制住实验条件,只比较不同 stochastic rollout 的质量。随后再用 TGR 把不同 prompt group 的 timestep scale 拉平,避免某些 timestep 因为数值系数更大而支配整次更新。
3.2 Preliminary: Flow-GRPO objective and ODE-to-SDE
Flow-GRPO 的目标写为: 其中 full-trajectory objective 聚合所有 timestep 的 clipped policy ratio: 这里 , 是 advantage。瓶颈正来自 的 dense supervision。
Flow matching 的 deterministic ODE 为: 为得到 GRPO 所需探索,Flow-GRPO / DanceGRPO 使用与原 ODE marginal 等价的 SDE: \begin{split} \boldsymbol{x}_{t+\Delta t} = \boldsymbol{x}_t + \left[ \boldsymbol{v}_\theta(\boldsymbol{x}_t, t) + \frac{\sigma_t^2}{2t} \left(\boldsymbol{x}_t + (1-t)\boldsymbol{v}_\theta(\boldsymbol{x}_t, t)\right) \right] \Delta t + \sigma_t \sqrt{\Delta t} \, \boldsymbol{\epsilon}, \end{split} 其中 。Flash-GRPO 的关键是只在一个选定 timestep 使用这个 stochastic SDE transition,其它 timestep 仍用 ODE。
3.3 Iso-Temporal Grouping (ITG)
Naive one-step GRPO 会让同一 prompt group 内每个样本独立抽 timestep: 这个 baseline 把不同噪声水平的 intrinsic difficulty 混进 reward 均值。ITG 改成:对 batch 中 个 prompt ,每个 prompt 抽一个 timestep ,组内 个 rollout 共享 ,只改变噪声: \begin{split} &\mathcal{G}_k = \{ \boldsymbol{x}^{i}_{t_k} \mid i \in [1, G] \}. \end{split} 训练时只在 计算 policy gradient。这样每个 prompt group 内 advantage 只反映“同一 denoising condition 下哪个 rollout 更好”,而不同 prompt group 又覆盖不同 timestep,保留 batch-level temporal diversity。
3.4 Temporal Gradient Rectification (TGR)
在 selected timestep ,标准 RL objective 的梯度为: Euler-Maruyama 诱导的 Gaussian transition kernel 写成: 其中 mean 由 vector field 参数化: \begin{split} &\boldsymbol{\mu}_\theta(\boldsymbol{x}_t, t) = \\ &\boldsymbol{x}_t + \left[ \boldsymbol{v}_\theta(\boldsymbol{x}_t, t) + \frac{\sigma_t^2}{2t} \left(\boldsymbol{x}_t + (1-t)\boldsymbol{v}_\theta(\boldsymbol{x}_t, t)\right) \right] \Delta t. \end{split} 代入 score function 后得到: $ \begin{split} \nabla_\theta \log &p_\theta(\boldsymbol{x}{t-1} | \boldsymbol{x}t) \&= \nabla\theta \left( -\frac{|\boldsymbol{x}{t-1} - \boldsymbol{\mu}\theta(\boldsymbol{x}t,t)|^2}{2\sigma_t^2\Delta t} \right)\ &= \frac{\boldsymbol{x}{t-1}-\boldsymbol{\mu}\theta(\boldsymbol{x}t,t)}{\sigma_t^2\Delta t}\nabla\theta\boldsymbol{\mu}\theta(\boldsymbol{x}t,t) \ &= \frac{\sigma_t\sqrt{\Delta t}\boldsymbol{\epsilon}}{\sigma_t^2\Delta t}\nabla\theta\boldsymbol{\mu}\theta(\boldsymbol{x}t,t) \ &= \frac{\boldsymbol{\epsilon}}{\sigma_t\sqrt{\Delta t}} \cdot \Delta t\left(1+\frac{\sigma_t^2(1-t)}{2t}\right)\nabla\theta\boldsymbol{v}\theta(\boldsymbol{x}t,t) \ &= \underbrace{\left(\frac{\sqrt{\Delta t}}{\sigma_t}+\frac{\sigma_t\sqrt{\Delta t}(1-t)}{2t}\right)}{\lambda(t): \text{ Time-dependent Scaling}} \boldsymbol{\epsilon} \cdot \nabla\theta\boldsymbol{v}_\theta(\boldsymbol{x}_t,t).
\end{split} \lambda(t)=\frac{\sqrt{\Delta t}}{\sigma_t}+\frac{\sigma_t\sqrt{\Delta t}(1-t)}{2t} \mathcal{L}{\text{TGR}}(\theta) = \frac{1}{G} \sum{i=1}^G \frac{\hat{A}^i_t}{\lambda(t)} \cdot r_t^i(\theta). $ 这一步本质上把 sampler discretization scale 从 RL objective 里消掉,使“某个 timestep 数值系数大”不再等价于“这个 prompt 更该更新”。
3.5 Algorithm and PyTorch-style pseudocode
论文算法输入 prompt 、group size 、total timesteps 、reward models ;每轮随机采样 timestep ,在 用 SDE,其他 timestep 用 ODE,计算 的 mean/std 与 advantage,最后用 更新 。
def sample_iso_temporal_indices(prompts, max_selected_steps=10):
"""Released code behavior: one shared selected index per unique prompt."""
index_dict = {}
weights = [1] * max_selected_steps # code samples only the first 10 scheduler steps
for prompt in gather_object(prompts):
if prompt not in index_dict:
index_dict[prompt] = random.choices(range(len(weights)), weights=weights, k=1)[0]
broadcast(index_dict, src=0)
return torch.tensor([index_dict[p] for p in prompts], device="cuda")def rollout_with_one_sde_step(pipeline, prompt_embeds, neg_embeds, selected_index, cfg_scale=4.5):
"""Only selected_index is stochastic; all other steps are deterministic ODE."""
latents = pipeline.prepare_latents(...)
saved_latents, saved_log_probs = [], []
for i, t in enumerate(pipeline.scheduler.timesteps):
noise_pred = pipeline.transformer(latents, timestep=t, encoder_hidden_states=prompt_embeds)[0]
noise_uncond = pipeline.transformer(latents, timestep=t, encoder_hidden_states=neg_embeds)[0]
noise_pred = noise_uncond + cfg_scale * (noise_pred - noise_uncond)
deterministic = (i != selected_index)
if i == selected_index - 1:
saved_latents.append(latents)
latents, log_prob, mean, std = sde_step_with_logprob(
pipeline.scheduler, noise_pred.float(), t[None], latents.float(), determistic=deterministic
)
if i == selected_index:
saved_latents.append(latents)
saved_log_probs.append(log_prob)
return decode_video(latents), torch.stack(saved_latents), torch.stack(saved_log_probs)def compute_video_hpsv3_reward(videos, prompts, reward_server):
"""Released code sends every frame to an HPSv3 server and averages top 30%."""
scores = []
for video, prompt in zip(videos, prompts):
frame_scores = reward_server(images=list(video), prompts=[prompt] * len(video))
frame_scores = sorted(frame_scores, reverse=True)
k = int(len(frame_scores) * 0.3)
scores.append(sum(frame_scores[:k]) / k)
return torch.tensor(scores, device="cuda")def temporal_gradient_rectified_grpo_loss(log_prob, old_log_prob, advantages, timesteps):
"""Released code uses a hard-coded inverse-lambda table for the sampled Wan scheduler timesteps."""
inv_lambda = {
999: 7.4770, 982: 7.0414, 963: 6.6112, 944: 6.1867, 922: 5.7682,
899: 5.3559, 874: 4.9502, 847: 4.5513, 817: 4.1596, 785: 3.7754,
}
value = torch.tensor([inv_lambda[int(t)] for t in timesteps.flatten()], device=timesteps.device)
value = value.reshape(timesteps.shape)
w = 1.0 / distributed_mean(value.mean())
ratio = torch.exp(log_prob - old_log_prob)
unclipped = -w * value.squeeze() * advantages * ratio
clipped = -w * value.squeeze() * advantages * torch.clamp(ratio, 1.0 - 1e-3, 1.0 + 1e-3)
return torch.mean(torch.maximum(unclipped, clipped))Code reference:
main@42110137(2026-05-18) — pseudocode and mapping based on this commit
| Paper Concept | Source File | Key Class/Function |
|---|---|---|
| Wan2.1 Flash-GRPO launch | scripts/multi_node/train_wanx2_1_flash.sh | accelerate launch --num_machines 12 --num_processes 96 ... --config config/dgx.py:wan2_1_flash |
| Experiment config | config/dgx.py | wan2_1_flash() sets Wan2.1-T2V-1.3B, 20 train steps, 50 eval steps, CFG 4.5, 480×832×81 frames, clip_range=1e-3, LR 1e-4, bf16, HPSv3 reward |
| Iso-temporal timestep assignment | scripts/train_wanx2_1_flash.py | prompt-level index_dict, weights=[1]*10, broadcasted selected index |
| One-step SDE / deterministic ODE rollout | flow_grpo/diffusers_patch/wan2_1_pipeline_with_logprob_sample.py | wan_pipeline_with_logprob, sde_step_with_logprob; i == sampled_indices is stochastic, otherwise deterministic |
| Log-prob recomputation for policy ratio | scripts/train_wanx2_1_flash.py | compute_log_prob() calls sde_step_with_logprob(..., prev_sample=sample["next_latents"]) |
| TGR policy loss | scripts/train_wanx2_1_flash.py | hard-coded inverse-λ value_dict, normalized by distributed mean, PPO-style clipped loss |
| Reward computation | flow_grpo/rewards.py | video_hpsv3_remote() sends frames to HPSv3 server and averages top 30% frame scores |
论文公式与 released code 实现差异:论文把 和 写成连续 / 通用形式;当前 main@42110137 的 Wan2.1-1.3B 实现只在 scheduler 的前 10 个 index 中均匀采样,实际 timestep keys 为 999, 982, 963, 944, 922, 899, 874, 847, 817, 785,并用 hard-coded inverse-λ table 加一个 batch mean normalization w=1/mean(value_tensor)。论文还报告 14B 模型和 12 sampling steps,但公开仓库当前可见的 wan2_1_flash() 配置是 Wan2.1-T2V-1.3B、20 training sampling steps;未在该 commit 中看到单独 14B training config。
4. Experimental Setup (实验设置)
Datasets and models
论文遵循 DanceGRPO 设置,使用其 prompt dataset 训练,并 hold out 一个包含 300 prompts 的 distinct evaluation split。公开代码的 dataset/video/test.txt 正好包含 300 行;dataset/video/train.txt 包含 19,700 行,可作为 release 中实际训练 prompt split 的证据。基础模型使用 Wan2.1 family,实验覆盖 Wan2.1-T2V-1.3B 与 14B variants。
Baselines
主要比较对象包括:Flow-GRPO、Flow-GRPO-Fast1(Flow-GRPO-Fast 的 single-step setting,用于对齐 Flash-GRPO 的 one-step budget)、vanilla Wan2.1;VBench general quality 表中还列出 CogVideoX-2B、Hunyuan-Video、Wan2.1-T2V-1.3B。论文也讨论 MixGRPO 与 Flow-GRPO-Fast 的 sliding-window 思路作为现有效率化路线。
Evaluation metrics
- Visual Quality / HPSv3:使用 HPSv3 reward model;论文按 sampled frames 打分,并取 top 30% scoring frames 的平均值计算 reward,以缓解 temporal transition 中低分帧的影响。公开代码
video_hpsv3_remote()也对每个视频的 frame scores 排序后取int(l*0.3)个最高分求均值。 - Motion Quality / VideoAlign motion score:衡量 temporal coherence 与 motion dynamics;论文报告 Flash-GRPO 在 motion score 上比 Flow-GRPO-Fast1 更稳定。
- General Video Quality / VBench:按 VBench protocol,使用 enhanced prompts 与 negative prompts,评估 Aesthetic Quality、Imaging Quality、Subject Consistency、Object Class 等维度。
Training config
论文写明:1.3B model training 使用 20 sampling steps;14B model 使用 accelerated 12 sampling steps;evaluation held-out set inference 使用 50 sampling steps;CFG scale 固定为 4.5;GRPO clip ratio 为 0.001。公开 launch/config 对应如下:
Released Wan2.1-1.3B config
- Launch script:
scripts/multi_node/train_wanx2_1_flash.sh - Accelerate / hardware: Deepspeed ZeRO-2 config;
--num_machines 12、--num_processes 96,即 96 GPU processes。 - Config path:
config/dgx.py:wan2_1_flash - Model / shape:
pretrained.model="Wan2.1-T2V-1.3B-Diffusers";height480、width832、frames81。 - Sampling:
sample.num_steps=20、sample.eval_num_steps=50、sample.guidance_scale=4.5、sample.train_batch_size=1、sample.num_batches_per_epoch=2、sample.sample_time_per_prompt=1、sample.test_batch_size=2。 - Optimization:
train.batch_size=1、gradient_accumulation_steps=1、train.num_inner_epochs=1、train.timestep_fraction=0.51、train.beta=0、train.learning_rate=1e-4、train.clip_range=1e-3、train.sft=0.0、mixed precisionbf16、EMA enabled。 - Reward / tracking:
reward_fn={"videohpsv3": 1.0}、prompt_fn="general_ocr"、per_prompt_stat_tracking=True、sample.global_std=True。
5. Experimental Results (实验结果)
5.1 VBench general video quality
| Method | GPU Hours | Aesthetic Quality ↑ | Imaging Quality ↑ | Subject Consistency ↑ | Object Class ↑ |
|---|---|---|---|---|---|
| CogVideoX-2B | — | 61.07 | 62.37 | 96.52 | 86.48 |
| Hunyuan-Video | — | 60.36 | 67.56 | 97.37 | 86.10 |
| Wan2.1-T2V-1.3B | — | 65.46 | 66.79*/67.01 | 97.56 | 88.84*/88.81 |
| Flow-GRPO-Fast1 | 350 | 65.92 | 65.96 | 98.46 | 88.15 |
| Flow-GRPO | 350 | 65.79 | 68.60 | 97.28 | 87.92 |
| Flash-GRPO | 350 | 66.43 | 68.28 | 98.70 | 90.00 |
结论:同样 350 GPU hours 下,Flash-GRPO 在 Aesthetic Quality、Subject Consistency、Object Class 上取得表内最高值,并保持接近 Flow-GRPO 的 Imaging Quality(68.28 vs 68.60)。这支持作者的主张:single-step training 不必像 Flow-GRPO-Fast1 那样牺牲优化稳定性或生成质量。
5.2 Qualitative comparison

Figure 3 解读:奇数行是 vanilla Wan2.1,偶数行是 Flash-GRPO。三个案例分别强调 Motion、Aesthetic、Prompt Following:骑马序列中背景和人物运动更稳定;熊猫场景细节和视觉质量更好;卡通动物场景中红框标出的 unstable elements 被改善,prompt 中额外元素也更容易被遵循。
5.3 Ablation: ITG and TGR
| Method | Train Stability | Eval Reward |
|---|---|---|
| Wan2.1-1.3B | - | 4.67 |
| Naive Single-step | × | 4.64 |
| + ITG | × | 5.31 |
| + ITG + TGR (Full) | ✓ | 5.42 |
ITG 单独把 Eval Reward 从 4.64 提到 5.31,说明消除 timestep-confounded variance 对 reward 有显著帮助;但稳定性仍未解决。加入 TGR 后 Eval Reward 进一步到 5.42,并且 Train Stability 变为 ✓,说明梯度尺度校正确实对应优化稳定性的关键缺口。
Figure 4 解读:无 KL regularization 时,Flow-GRPO-Fast1 在 train / eval HPSv3 曲线上都出现 collapse;Flash-GRPO 保持稳定收敛。这张图主要验证 TGR 对 single-step setting 的必要性,因为它展示的不是最终数值小幅提升,而是 baseline 优化过程本身失控。
Figure 5 解读:与 full trajectory Flow-GRPO 相比,Flash-GRPO 在 train / eval HPSv3 上收敛更快、reward ceiling 更高。论文文字说明在 evaluation curve 上 Flash-GRPO peak reward approximately 5.4,而 Flow-GRPO approximately 5.1;这说明 full-trajectory 并非在低预算下天然更优,稳定的 one-step credit assignment 可以有更好的 sample / compute efficiency。
Figure 6 解读:Motion Quality 维度上,Flash-GRPO 仍表现出稳定提升;论文文字给出的 final score approximately ,Flow-GRPO-Fast1 approximately 。这说明改进不只是 HPSv3 / visual preference reward 上的过拟合,也反映到 temporal coherence 与 motion dynamics。
5.4 Appendix experiments
Figure 7 解读:加 KL regularization 后,Flow-GRPO-Fast1 不再彻底 collapse,但仍明显 plateau。论文文字给出 held-out set 上 Flash-GRPO approximately 5.35,Fast1 approximately 4.9,说明 KL 能缓解不稳定,却不能解决 timestep confounding 与 gradient scale imbalance 的根因。
Figure 8 解读:左图比较无 KL 时的 gradient norm,Flow-GRPO-Fast1 出现 catastrophic gradient spikes 和 high variance,Flash-GRPO 保持较低且稳定;右图显示 unrectified baseline 在 300–400 GPU hours 之间发生 reward collapse,而 TGR 后 reward 单调上升并有更高收敛上限。

Figure 9 解读:Wan1.3B + HPSv3 rewards 的 qualitative comparison。waterfall / cliff 场景中,Flash-GRPO 在红框 foreground region 上有更连贯的人体运动;cooking / watermelon 场景中,人物面部、厨房环境、切西瓜动作和光照纹理更真实。

Figure 10 解读:Wan14B qualitative comparison 中的 Japanese garden / wooden spoon prompt。Flash-GRPO 更好地保持 foreground object 与景深效果,说明方法在更大模型上仍能改善 prompt following 与审美细节。

Figure 11 解读:Wan14B 的 bird、sailboat、zebra / giraffe 场景展示了更流畅的 motion 和更丰富的环境细节。作者用这些样例说明 Flash-GRPO 的收益跨 prompt 类型存在,而不是只针对单一视觉风格。

Figure 12 解读:Wan14B 的 cat、dog chasing、hand-held sword CG 场景中,Flash-GRPO 改善了 motion aesthetics 与 prompt following。结合 Figure 9–12,作者的总体结论是 Flash-GRPO 在 1.3B 与 14B 两个尺度上都能提升 visual fidelity、motion smoothness 与 complex prompt adherence。
5.5 Limitations and risks
论文没有系统报告失败案例、reward hacking 或 human evaluation;主要证据来自 HPSv3、VideoAlign motion score、VBench 与 qualitative samples。公开代码还存在 engineering caveat:TGR inverse-λ 是 Wan scheduler 的 hard-coded table,且当前 release 只清楚给出 1.3B / 20-step / 96-GPU launch,若迁移到其它 sampler、timestep schedule 或 14B 配置,需要重新计算对应系数并验证 stability。
总体结论:Flash-GRPO 把 video GRPO 的效率瓶颈从 full-trajectory backprop 降到 single-step update,并通过 ITG + TGR 解决 naive one-step 的两个根因;在 350 GPU hours VBench general quality、HPSv3 reward、motion score 和 qualitative samples 上都显示出比 Flow-GRPO-Fast1 更稳定、比 full trajectory Flow-GRPO 更高效的 alignment 表现。