AnyFlow Flow Map Distillation Any-Step Sampling On-Policy Training

让视频 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。

4 到 32 NFEs单个模型支持不同测试采样预算。
1.3B 到 14B覆盖双向和 causal 视频 diffusion 架构。
84.41 VBenchAnyFlow-FAR 14B causal T2V 在 32 NFEs 达到总分 84.41。
87.87 I2V14B I2V 在 4 NFEs 已达到接近 Wan2.1-I2V-14B 50x2 NFEs 的分数。

摘要与定位

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。

核心替换 从 endpoint 到 flow map

模型学习任意时间点之间的转移,而不是只学到 z_0 的映射。

核心收益 统一 few-step 与 multi-step

同一模型既能 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。

test-time scaling
AnyFlow 的 test-time scaling 定量比较:相比 consistency-based 方法,AnyFlow 在少步数下表现强,并能随采样步数增加继续提升。 源文件:source/figures/src/teaser_test_scale.pdf
teaser frame 0 teaser frame 5 teaser frame 10 teaser frame 15 teaser frame 20
AnyFlow test-time scaling 的定性帧示例。这里直接取自源码中的动画帧序列。 源文件:source/figures/videosrc/teaser/*.jpg

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。

paradigm comparison
范式对比:consistency 只学 endpoint mapping;AnyFlow 学任意时间对之间的 flow map transition。 源文件:source/figures/src/paradigm_comparison.pdf

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 两个性质。

PF-ODE 与 flow map
d z_t / dt = v(z_t, t)
Phi_{r<-t}(z_t) = z_r, Phi_{q<-r} o Phi_{r<-t} = Phi_{q<-t}

这个公式约束模型遵循 ODE 轨迹上的状态转移。composition 性质是后面 Flow Map Backward Simulation 能用 shortcut transitions 替代完整 rollout 的理论抓手,减少显式模拟长轨迹的成本。

神经 flow map model
f_theta(z_t, t, r) ≈ z_r, 1 ≥ t > r ≥ 0

与只预测 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 用有限差分近似。

MeanFlow 目标
L(theta) = E[ || u_theta(z_t,r,t) - sg(u_tgt) ||_2^2 ]
u_tgt = v(z_t,t) - (t-r) * d u_theta(z_t,r,t) / dt

它约束 averaged velocity 接近目标速度,用于训练 flow map 的转移能力。误差减少的是任意时间区间 transition 的拟合误差。

有限差分替代 JVP
d u / dt ≈ [u(z_{t+Delta t}, r, t+Delta t) - u(z_{t-Delta t}, r, t-Delta t)] / (2 Delta t)

该式用两次 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。

DMD 梯度
grad_theta L_DMD = -E[(s_real(z_t,t)-s_fake(z_t,t)) * partial f_theta(z) / partial theta]

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。

AnyFlow pipeline
AnyFlow pipeline:先做 forward flow map training,再用 flow map backward simulation 做 on-policy distillation。 源文件:source/figures/src/pipeline.pdf

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。

插值 timestep conditioning
g * emb(t) + (1-g) * emb'(r), g = 0.25

这个设计在边界 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。

guidance-fused prediction
u = (1/g) * (u_c - (1-g) * sg(u_null))

该式让模型在训练中吸收 CFG 效果,用于减少测试时额外 conditional/unconditional 双分支计算。

Adaptive Loss Reweighting

post-training 中 t=r 的 instantaneous velocity field 已经较可靠。AnyFlow 用边界 case 的平均回归损失作为 baseline,动态缩放 t != r 的 loss。

自适应 loss 权重
w_{t,r} = 1 if t=r; w_{t,r} = mu_{t=r} / (||Delta||_2^2 + c) if t != r

它约束非边界时间对的 loss 尺度与已优化好的边界场对齐,减少 gradient instability,同时保留原始 instantaneous flow field。

interpolated embedding
timestep embedding norm 对比:interpolated timestep conditioning 比 zero-init 更稳定。 源文件:source/figures/src/interpolated_embedding.pdf
interpolated visualization
zero-init timestep conditioning 与 interpolated timestep conditioning 的可视化比较。 源文件:source/figures/src/interpolated_visualization.pdf

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。

flow map composition 近似
f_theta(z_t,t,q) ≈ f_theta(f_theta(z_t,t,r), r, q), t > r > q

该式约束 shortcut transition 与分段 transition 一致。它用于 backward simulation,将长 Euler rollout 分解成可训练的短/长区间 transition,减少完整轨迹显式模拟成本,并让同一模型学习不同 inference budgets。

backward simulation comparison
Backward simulation 范式对比:consistency 依赖 consistency trajectories;flow map backward simulation 通过 shortcut segments 高效模拟 Euler sampling trajectories。 源文件:source/figures/src/comp_simulation.pdf
ablation backward frame 0 ablation backward frame 5 ablation backward frame 10 ablation backward frame 15 ablation backward frame 20
Backward simulation 定性消融帧:flow map backward simulation 轨迹更直,能在保持 few-step 质量的同时支持 test-time scaling。 源文件:source/figures/videosrc/ablation_anyflow_simulation/*.jpg

应用到 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,平衡运动建模与吞吐。

downstream pipeline
AnyFlow 下游 fine-tuning pipeline:相比 Self-Forcing pretrained causal models,AnyFlow 更容易继续训练到新数据集。 源文件:source/figures/src/downstream_pipeline.pdf
downstream results
AnyFlow 下游 fine-tuning 可视化:在 robotics 和 driving 等专门领域继续训练后,可改善身份保持与轨迹准确性。 源文件:source/figures/src/downstream_results.pdf

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。

模型参数分辨率NFEsQualitySemanticTotal
rCM-Wan2.1-T2V-14B14B480x832485.4776.7283.73
AnyFlow-Wan2.1-T2V-14B14B480x832485.7077.3884.04
AnyFlow-Wan2.1-T2V-14B14B480x8323285.7677.4484.10
Krea-Realtime-Wan2.1-14B14B480x832484.8077.0783.25
AnyFlow-FAR-Wan2.1-14B14B480x832485.8276.9784.05
AnyFlow-FAR-Wan2.1-14B14B480x8323286.1277.5584.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。

模型参数分辨率NFEsQualityI2VTotal
Wan2.1-I2V-14B14B480x83250x280.3095.1287.71
FastVideo-CausalWan2.2-A14B-Preview14B480x832878.8294.8186.82
AnyFlow-FAR-Wan2.1-14B14B480x832480.3995.3587.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 保持首帧一致性和顺滑运动转移。

causal t2v frame 0 causal t2v frame 5 causal t2v frame 10 causal t2v frame 15 causal t2v frame 20
14B causal T2V 定性比较帧。源文件:source/figures/videosrc/t2v14b_causal/*.jpg
bidirectional t2v frame 0 bidirectional t2v frame 5 bidirectional t2v frame 10 bidirectional t2v frame 15 bidirectional t2v frame 20
14B bidirectional T2V 定性比较帧。源文件:source/figures/videosrc/t2v14b_bidirectional/*.jpg
i2v frame 0 i2v frame 5 i2v frame 10 i2v frame 15 i2v frame 20
14B I2V 定性比较帧。源文件:source/figures/videosrc/i2v14b/*.jpg

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 两种设置中整体最强。

方法NFEsBidirectional OverallCausal Overall
Flow Matching Training4x274.6476.80
Flow Matching Training32x283.8383.50
Flow Map Training481.7580.48
Flow Map Training3283.4083.13
Consistency ODE-Init + Consistency Backward Simulation4 / 3282.96 / 79.8082.49 / 79.64
Flow Map Training + Flow Map Backward Simulation4 / 3283.48 / 83.9683.54 / 83.96

Time Sampler

uniform weighting 的 32-step scaling 最差,可能因为它与预训练使用的 logit-normal time weighting 不匹配。给高噪声区域更多概率质量能获得更好整体表现,因此作者固定 w(t)=Beta(2,1.5)。

time sampler
timestep density f(t)=p(t)w(t) 的可视化。 源文件:source/figures/src/time_sampler.pdf

On-Policy 与 ODE-Init 消融

onpolicy frame 0 onpolicy frame 5 onpolicy frame 10 onpolicy frame 15 onpolicy frame 20
On-policy distillation 定性消融:forward flow map training alone 仍受离散化误差和 exposure bias 影响;on-policy flow map distillation 在模型 self-rollouts 上优化 reverse divergence。 源文件:source/figures/videosrc/ablation_onpolicy/*.jpg
odeinit frame 0 odeinit frame 5 odeinit frame 10 odeinit frame 15 odeinit frame 20
ODE-Init 与 AnyFlow 在 causal video generation 上的定性消融。 源文件:source/figures/videosrc/ablation_odeinit/*.jpg

训练成本

Flow Map Backward Simulation 在模拟 4 步时比 consistency backward simulation 略贵,因为需要通过完整 rollout chain 反传;但当模拟 16 步时,由于 learned shortcut transitions,训练成本明显更低。

方法Causal 成本Bidirectional 成本
Standard Flow-Matching5.8 s/iter9.7 s/iter
+ Guidance-Fused + Differential Derivation10.4 s/iter16.8 s/iter
Consistency Backward Simulation (4 Steps)45.941.8
Flow Map Backward Simulation (4 Steps)53.1 (+15.7%)51.2 (+22.5%)
Consistency Backward Simulation (16 Steps)93.896.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 中裁剪图片。