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 ConceptSource FileKey Class/Function
Wan2.1 Flash-GRPO launchscripts/multi_node/train_wanx2_1_flash.shaccelerate launch --num_machines 12 --num_processes 96 ... --config config/dgx.py:wan2_1_flash
Experiment configconfig/dgx.pywan2_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 assignmentscripts/train_wanx2_1_flash.pyprompt-level index_dict, weights=[1]*10, broadcasted selected index
One-step SDE / deterministic ODE rolloutflow_grpo/diffusers_patch/wan2_1_pipeline_with_logprob_sample.pywan_pipeline_with_logprob, sde_step_with_logprob; i == sampled_indices is stochastic, otherwise deterministic
Log-prob recomputation for policy ratioscripts/train_wanx2_1_flash.pycompute_log_prob() calls sde_step_with_logprob(..., prev_sample=sample["next_latents"])
TGR policy lossscripts/train_wanx2_1_flash.pyhard-coded inverse-λ value_dict, normalized by distributed mean, PPO-style clipped loss
Reward computationflow_grpo/rewards.pyvideo_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";height 480、width 832、frames 81
  • Sampling: sample.num_steps=20sample.eval_num_steps=50sample.guidance_scale=4.5sample.train_batch_size=1sample.num_batches_per_epoch=2sample.sample_time_per_prompt=1sample.test_batch_size=2
  • Optimization: train.batch_size=1gradient_accumulation_steps=1train.num_inner_epochs=1train.timestep_fraction=0.51train.beta=0train.learning_rate=1e-4train.clip_range=1e-3train.sft=0.0、mixed precision bf16、EMA enabled。
  • Reward / tracking: reward_fn={"videohpsv3": 1.0}prompt_fn="general_ocr"per_prompt_stat_tracking=Truesample.global_std=True

5. Experimental Results (实验结果)

5.1 VBench general video quality

MethodGPU HoursAesthetic Quality ↑Imaging Quality ↑Subject Consistency ↑Object Class ↑
CogVideoX-2B61.0762.3796.5286.48
Hunyuan-Video60.3667.5697.3786.10
Wan2.1-T2V-1.3B65.4666.79*/67.0197.5688.84*/88.81
Flow-GRPO-Fast135065.9265.9698.4688.15
Flow-GRPO35065.7968.6097.2887.92
Flash-GRPO35066.4368.2898.7090.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

MethodTrain StabilityEval 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 表现。