让视频 diffusion 蒸馏模型真正支持任意步数采样
这篇论文提出 AnyFlow,一个基于 flow map 的 any-step 视频 diffusion 蒸馏框架。它的核心动机很锋利:现有 consistency-distilled 视频模型在 4 步这类少步数下很强,但测试时增加采样步数后,质量常常不升反降。AnyFlow 试图保留少步高效性,同时恢复 ODE sampling “步数越多、质量越好”的 test-time scaling。
方法上,AnyFlow 不再只学习 endpoint consistency mapping,即从 z_t 直接到 z_0,而是学习任意时间区间的 flow-map transition,即从 z_t 到 z_r。再配合 Flow Map Backward Simulation,它可以做高效 on-policy distillation,减少 few-step 采样的离散化误差和 causal generation 的 exposure bias。
摘要与定位
Few-step video generation 已被 consistency distillation 显著推进,但这类模型经常在测试时分配更多采样步数后性能下降,限制了它们作为 any-step video diffusion 的有效性。
作者认为,问题来自 consistency distillation 将原始 probability-flow ODE 轨迹替换为 consistency-sampling 轨迹,从而削弱了 ODE sampling 原本有利的 test-time scaling。为解决这一点,AnyFlow 把蒸馏目标从 endpoint consistency mapping z_t -> z_0 改为任意时间区间上的 flow-map transition learning z_t -> z_r,并提出 Flow Map Backward Simulation,把完整 Euler rollout 分解成 shortcut flow-map transitions。
consistency-based 方法会反复 endpoint projection 和 re-noising,轨迹可能偏离目标 PF-ODE。
模型学习任意时间点之间的转移,而不是只学到 z_0 的映射。
同一模型既能 4 NFEs 快速生成,也能在 32 NFEs 下继续提升质量。
1. 引言
大规模视频 diffusion 模型已经取得很高生成质量,但多数 pipeline 仍绑定固定推理步数。实际使用中,用户往往需要灵活生成:快速预览时希望低延迟,最终交付时希望更高保真。因此,模型需要具备 any-step 能力,即无需重新训练,就能在测试时用不同采样预算交换延迟和质量。
现有 few-step 视频蒸馏方法主要建立在 consistency models 上。它们在极小采样预算下有效,但 any-step 行为不稳。代表性方法如双向模型中的 rCM、causal 模型中的 Self-Forcing,常常随着采样步数增加而退化。结构原因在于,consistency formulation 强调到 z_0 的 fixed-point mapping,多步采样时反复 re-noising 中间状态,会引入累积偏差,使轨迹离目标 PF-ODE path 越来越远,而不是被逐步 refinement。
AnyFlow 包含两个互补阶段。第一阶段是 improved forward flow map training,把预训练视频 diffusion 模型转成 flow map model,为 any-step generation 提供强初始化。第二阶段是 on-policy flow map distillation,用 teacher guidance 在模型自身 rollouts 上优化 reverse divergence,以减少 test-time errors。
2. 背景与预备知识
Flow Map Formulation
令 z_t 表示 probability-flow ODE (PF-ODE) 在连续时间 t in [0,1] 的 latent state,速度场为 v(z_t,t)。精确 flow map Phi_{r<-t} 是把状态从时间 t 传输到时间 r 的算子,满足 identity 和 composition 两个性质。
这个公式约束模型遵循 ODE 轨迹上的状态转移。composition 性质是后面 Flow Map Backward Simulation 能用 shortcut transitions 替代完整 rollout 的理论抓手,减少显式模拟长轨迹的成本。
与只预测 z_0 的 endpoint mapping 相比,它能处理任意时间对,因此天然支持可变 step size 和 any-step inference。
MeanFlow 与有限差分
MeanFlow 通过平均 transport velocity 参数化区间 [r,t] 上的转移:f_theta(z_t,t,r)=z_t-(t-r)u_theta(z_t,r,t)。原始目标包含导数项,需要 Jacobian-vector product (JVP),在 FSDP 下难以扩展,因此 Transition Model 用有限差分近似。
它约束 averaged velocity 接近目标速度,用于训练 flow map 的转移能力。误差减少的是任意时间区间 transition 的拟合误差。
该式用两次 forward pass 近似导数,避免昂贵 JVP,使 flow map training 更适合大规模视频模型和 FSDP。
On-Policy Distillation
On-policy distillation 通过在 student 自己生成的 rollout states 上训练,并由强 teacher 指导,缓解 train-test mismatch。diffusion 中常用 Distribution Matching Distillation (DMD),包含 backward simulation 和 distribution-matching objective。
s_real 与 s_fake 分别是真实分布和 student 生成分布的 score。该梯度用于把 student rollout 分布拉向 teacher/data 分布,减少 few-step 采样误差和 causal generation 的 exposure bias。
3. 方法:AnyFlow Pipeline
AnyFlow 的目标是在保留 distilled video diffusion 少步高效性的同时,让质量随着采样步数增加而提升。它把蒸馏目标从 fixed-point endpoint mapping 转移到 flow map transition learning:模型学习 f_theta:(z_t,t,r) -> z_r,而不是只预测 z_0。
Stage 1:Flow Map Training
第一阶段在 teacher-synthesized data 上微调预训练视频 diffusion 模型,学习稳定 transition operator f_theta:(z_t,t,r)->z_r。作者基于 MeanFlow 目标和有限差分公式,并加入四个设计:interpolated timestep conditioning、time sampler、guidance-fused training、adaptive loss reweighting。
Interpolated Timestep Conditioning
flow map model 需要新增 timestep r。直接使用 emb(t)+emb'(t-r) 或 zero-initialized projection 会导致 embedding norm 过大,产生过饱和图像。AnyFlow 改用插值 timestep conditioning。
这个设计在边界 t=r 时能退化到预训练 timestep embedding,约束新增 r 条件不破坏预训练尺度。它减少 embedding norm 漂移和过饱和生成结果。
Guidance-Fused Training
作者将 classifier-free guidance 融入 flow map training,从而测试时可以省略 CFG,加快推理。与 MeanFlow 把 CFG 目标融合到 target velocity 不同,AnyFlow 将其融合到预测中,以更好对齐预训练 diffusion 模型中的 guidance scale。
该式让模型在训练中吸收 CFG 效果,用于减少测试时额外 conditional/unconditional 双分支计算。
Adaptive Loss Reweighting
post-training 中 t=r 的 instantaneous velocity field 已经较可靠。AnyFlow 用边界 case 的平均回归损失作为 baseline,动态缩放 t != r 的 loss。
它约束非边界时间对的 loss 尺度与已优化好的边界场对齐,减少 gradient instability,同时保留原始 instantaneous flow field。
4. On-Policy Flow Map Distillation
Forward flow map training 提供强初始化,但仍存在 test-time errors:低 NFE 下的 discretization error,以及 causal generation 中的 exposure bias。为修正 rollout drift,AnyFlow 引入 teacher-guided on-policy distillation,并用 DMD 实现 reverse-divergence supervision。
为什么不能直接用 consistency backward simulation?
consistency backward simulation 可以方便地到达 z_0,因此 endpoint supervision 很直接。但 consistency sampler 多步采样时需要额外 re-noising,这些状态会偏离 base PF-ODE trajectory,导致难以泛化到 simulation step 之外的步数。并且,若要模拟任意步数,完整 rollout 的训练成本很高。
Flow Map Backward Simulation
AnyFlow 利用 flow map 的任意时间对转移能力,把长 rollout 轨迹 shortcut。对目标采样预算 N,先沿采样轨迹选中间 timestep t,再设下一步 r,使 t-r=T/N。这样完整 T->0 轨迹被分为 T->t、t->r、r->0 三段,首尾两段由 learned flow map shortcut 处理,中间 t->r 是目标 transition。
该式约束 shortcut transition 与分段 transition 一致。它用于 backward simulation,将长 Euler rollout 分解成可训练的短/长区间 transition,减少完整轨迹显式模拟成本,并让同一模型学习不同 inference budgets。
应用到 causal video diffusion
对 causal video diffusion,AnyFlow 采用 FAR 训练管线:用 context compression 和 asymmetric patchify kernels 高效编码视频。它保留三个 full-token chunks,其余 chunks 用更大 kernel 压缩,以降低 teacher-forcing 训练成本和采样 KV cache 大小。为联合支持 I2V 和 T2V,模型使用 non-uniform chunk partition:第一 chunk size 为 1,以精确处理首帧条件;后续 chunks size 为 3,平衡运动建模与吞吐。
5. 实验
实现细节
AnyFlow 基于 Wan2.1 和 Diffusers 实现。训练数据是由 Wan2.1-T2V-14B 生成的 256K prompt-video pairs,每个样本最多 81 帧,分辨率 480 x 832。训练分两阶段:Stage 1 用 AdamW 优化 forward flow map objective,1.3B 模型全局 batch size 32、训练 6000 iterations,14B 模型 batch size 16、训练 4000 iterations;Stage 2 用 learning rate 2e-6 做 800 iterations 的 on-policy distillation。两阶段都使用 rank 256 的 LoRA。
T2V 定量结果
AnyFlow 在少步数下表现强,并能随采样预算增加继续提升。14B bidirectional backbone 上,AnyFlow-Wan2.1-T2V-14B 在 4 NFEs 得到 84.04,超过 rCM-Wan2.1-T2V-14B 的 83.73;14B causal backbone 上,AnyFlow-FAR-Wan2.1-14B 在 4 NFEs 达到 84.05,在 32 NFEs 进一步提升到 84.41。
| 模型 | 参数 | 分辨率 | NFEs | Quality | Semantic | Total |
|---|---|---|---|---|---|---|
| rCM-Wan2.1-T2V-14B | 14B | 480x832 | 4 | 85.47 | 76.72 | 83.73 |
| AnyFlow-Wan2.1-T2V-14B | 14B | 480x832 | 4 | 85.70 | 77.38 | 84.04 |
| AnyFlow-Wan2.1-T2V-14B | 14B | 480x832 | 32 | 85.76 | 77.44 | 84.10 |
| Krea-Realtime-Wan2.1-14B | 14B | 480x832 | 4 | 84.80 | 77.07 | 83.25 |
| AnyFlow-FAR-Wan2.1-14B | 14B | 480x832 | 4 | 85.82 | 76.97 | 84.05 |
| AnyFlow-FAR-Wan2.1-14B | 14B | 480x832 | 32 | 86.12 | 77.55 | 84.41 |
I2V 定量结果
在 VBench-I2V 上,AnyFlow-FAR-Wan2.1-14B 只用 4 NFEs 达到总分 87.87,接近 Wan2.1-I2V-14B 使用 50x2 NFEs 的 87.71,并超过 FastVideo-CausalWan2.2-A14B-Preview 的 86.82。
| 模型 | 参数 | 分辨率 | NFEs | Quality | I2V | Total |
|---|---|---|---|---|---|---|
| Wan2.1-I2V-14B | 14B | 480x832 | 50x2 | 80.30 | 95.12 | 87.71 |
| FastVideo-CausalWan2.2-A14B-Preview | 14B | 480x832 | 8 | 78.82 | 94.81 | 86.82 |
| AnyFlow-FAR-Wan2.1-14B | 14B | 480x832 | 4 | 80.39 | 95.35 | 87.87 |
定性比较
在 causal T2V 例子中,FastVideo-Wan2.2-A14B-Preview 细节模糊,LightX2V-Wan2.1-14B-CausVid 有 flickering artifacts,Krea-Realtime-14B 在部分 clip 中运动不真实。相比之下,AnyFlow-FAR-Wan2.1-14B 在 4 NFEs 下给出更好的视觉质量和运动。I2V 中,AnyFlow 复用 causal generator,通过 non-uniform chunk partition 保持首帧一致性和顺滑运动转移。
6. 消融研究
关键设计消融
表格显示,flow map training 比 flow matching training 和 consistency ODE-init 提供更强 few-step 初始化;但 forward training alone 仍有 test-time errors,on-policy distillation 能进一步缓解。Flow Map Backward Simulation 在 few-step 和 multi-step 两种设置中整体最强。
| 方法 | NFEs | Bidirectional Overall | Causal Overall |
|---|---|---|---|
| Flow Matching Training | 4x2 | 74.64 | 76.80 |
| Flow Matching Training | 32x2 | 83.83 | 83.50 |
| Flow Map Training | 4 | 81.75 | 80.48 |
| Flow Map Training | 32 | 83.40 | 83.13 |
| Consistency ODE-Init + Consistency Backward Simulation | 4 / 32 | 82.96 / 79.80 | 82.49 / 79.64 |
| Flow Map Training + Flow Map Backward Simulation | 4 / 32 | 83.48 / 83.96 | 83.54 / 83.96 |
Time Sampler
uniform weighting 的 32-step scaling 最差,可能因为它与预训练使用的 logit-normal time weighting 不匹配。给高噪声区域更多概率质量能获得更好整体表现,因此作者固定 w(t)=Beta(2,1.5)。
On-Policy 与 ODE-Init 消融
训练成本
Flow Map Backward Simulation 在模拟 4 步时比 consistency backward simulation 略贵,因为需要通过完整 rollout chain 反传;但当模拟 16 步时,由于 learned shortcut transitions,训练成本明显更低。
| 方法 | Causal 成本 | Bidirectional 成本 |
|---|---|---|
| Standard Flow-Matching | 5.8 s/iter | 9.7 s/iter |
| + Guidance-Fused + Differential Derivation | 10.4 s/iter | 16.8 s/iter |
| Consistency Backward Simulation (4 Steps) | 45.9 | 41.8 |
| Flow Map Backward Simulation (4 Steps) | 53.1 (+15.7%) | 51.2 (+22.5%) |
| Consistency Backward Simulation (16 Steps) | 93.8 | 96.6 |
| Flow Map Backward Simulation (16 Steps) | 53.1 (-43.4%) | 51.2 (-47.0%) |
7. 结论与局限
AnyFlow 是第一个基于 two-time flow map formulation 的 any-step 视频 diffusion 蒸馏框架。通过学习任意时间对之间的 transition,AnyFlow 让单个模型支持广泛采样预算。它结合 improved forward flow map training 和用于 on-policy distillation 的 flow map backward simulation,减少采样中的 discretization error 和 exposure bias。实验显示,在双向和 causal 架构、1.3B 到 14B 模型规模上,AnyFlow 在 few-step regime 中匹配或超过强 consistency-based baseline,并能随着采样预算增加继续提升。
主要局限是:flow map training 依赖外部数据集。即使用合成数据,训练分布仍可能与 base model 原始训练分布不同,从而带来轻微 distribution shift,例如纹理更平滑。作者认为,如果能用 base model 预训练时的同源数据应用 AnyFlow,这个问题可缓解。
未来工作包括更有效、更稳定的 forward flow map training 策略,以及把 AnyFlow 扩展到 autoregressive long-video generation,配合专门长视频训练。
阅读路径
快速读这篇可以按四步走:先看摘要和引言,理解为什么 consistency distillation 不具备稳定 any-step scaling;再看 Preliminary 的 flow map 形式化,抓住 composition property;然后重点看 Method 的 interpolated timestep conditioning、adaptive reweighting 和 Flow Map Backward Simulation;最后看实验表,确认 AnyFlow 在 4 NFEs 和 32 NFEs 都能维持或提升性能。
本 HTML 的插图均来自 arXiv 源码包中的独立 figure/source 文件。PDF 图来自 `source/figures/src/*.pdf` 的单图转换;动态示例来自 `source/figures/videosrc/.../*.jpg` 的源码帧序列。没有从整篇论文 PDF 中裁剪图片。