AnyFlow:基于在线策略流图蒸馏的任意步数视频扩散模型 AnyFlow: Any-Step Video Diffusion Model with On-Policy Flow Map Distillation
首个基于流图的任意步数视频扩散蒸馏框架
前置知识
概率流常微分方程 (PF-ODE)
扩散模型的前向加噪与反向去噪可被建模为沿一条连续概率流常微分方程 (Probability-Flow ODE) 的轨迹,轨迹两端分别为纯噪声 $z_T$ 与干净数据 $z_0$,中间任意时刻 $t\in[0,1]$ 对应不同程度的加噪潜变量 $z_t$。形式上,$\frac{dz_t}{dt} = v(z_t, t)$,其中 $v$ 为速度场,由训练好的扩散模型隐式定义。PF-ODE 是本文 flow map 形式化的基础理论载体。
任何蒸馏或加速采样的方法都必须在 PF-ODE 之上重新组织轨迹的离散化方式,理解 PF-ODE 才能明白 consistency model 与 flow map model 的本质差异。
一致性模型 (Consistency Models)
Consistency Models 学习一个把任意时刻 $z_t$ 直接映射到干净数据 $z_0$ 的算子 $f_\theta(z_t, t) \approx z_0$,从而把多步 ODE 求解压缩为少量前向传播。代表方法包括 CM、sCM、rCM 以及视频领域的 Self-Forcing。在多步采样时,需要在中间状态反复加噪后再去噪,逐步偏离原始 PF-ODE 轨迹。
本文明确指出当前视频扩散蒸馏几乎全部建立在 consistency model 之上,并指出一致性框架在多步采样时性能反而退化的结构性缺陷。
流图模型 (Flow Map Models)
流图把一致性模型的固定点映射 $z_t \to z_0$ 推广为任意两点间的转移算子 $f_\theta(z_t, t, r) \approx z_r$,其中 $1 \ge t \ge r \ge 0$。它满足恒等性 $f_\theta(z,t,t) = z$ 与可组合性 $f_\theta(z_t,t,q) \approx f_\theta(f_\theta(z_t,t,r), r, q)$。当 $r=0$ 时退化为 consistency model,当 $t=r$ 时退化为流匹配速度场。
这是 AnyFlow 取代 consistency model 的核心数学对象,可组合性是后续 flow map backward simulation 能够 shortcut 长轨迹的关键。
MeanFlow
MeanFlow 是当前主流的 flow map 训练算法,参数化区间 $[r, t]$ 上的平均传输速度 $u_\theta(z_t, r, t)$,目标函数为 $\mathcal{L}(\theta) = \mathbb{E}[\|u_\theta(z_t, r, t) - \mathrm{sg}(u_{tgt})\|_2^2]$,其中 $u_{tgt} = v(z_t, t) - (t-r)\frac{du_\theta}{dt}$。它使用 JVP 近似(差分推导方程)来避免对全模型反向传播,时间复杂度可控。
AnyFlow 的第一阶段直接采用 MeanFlow 目标,但针对视频后训练引入了时间条件、引导融合、损失重加权等改进。
分布匹配蒸馏 (DMD)
Distribution Matching Distillation 是一种典型的在线策略蒸馏范式:让学生模型先自回归采样得到 $\hat{z}_0$,再在 $t$ 时刻重新加噪,用真实分数 $s_{real}$ 与伪分数 $s_{fake}$ 之差反向传播 KL 梯度。梯度形式为 $\nabla_\theta \mathcal{L}_{DMD} = -\mathbb{E}_{t,z}[(s_{real}(z_t, t) - s_{fake}(z_t, t))\frac{\partial f_\theta(z)}{\partial \theta}]$。
DMD 是 AnyFlow 第二阶段 on-policy distillation 的核心,但传统 DMD 的 backward simulation 依赖 consistency 框架,本文需要重新设计为 flow map 版本。
无分类器引导 (CFG)
Classifier-Free Guidance 通过在条件 $c$ 与无条件 $\emptyset$ 预测之间插值来增强生成效果,公式为 $u = g\cdot u_c - (g-1)\cdot u_\emptyset$,其中 $g$ 为引导强度。本文在训练阶段就把 CFG 融合进预测中,从而在推理时无需再调用两次模型。
本文提出的 Guidance-Fused Training 与传统 MeanFlow 不同,它把 CFG 融合到预测而非目标速度场,是工程上的关键创新。
双向 vs 因果视频扩散
双向 (bidirectional) 模型在生成每一帧时可以看到全部未来与过去帧,因果 (causal) 模型则只能看到过去帧,更接近自回归长视频生成。AnyFlow 同时在两类架构上验证,其中因果模型还需要在 rollouts 中维护 KV-cache 以支持长程推理。
理解这两种架构的差异有助于把握 AnyFlow 引入 FAR (chunk-based 因果训练) 的目的——在 I2V/T2V/V2V 统一框架下大幅降低 teacher-forcing 成本。
VBench / VBench-I2V
VBench 是文本生成视频任务的事实标准评估套件,将生成质量拆为 16 个细粒度维度,汇总为 Quality 与 Semantic 两个总分。VBench-I2V 在此基础上加入 I2V score 以评估条件图一致性。本文 T2V 全部使用官方 augmented prompts 重新评估以保证公平。
理解评估指标才能读懂 Tables 3 与 4 中 Total / Quality / Semantic / I2V 各列的含义与对比逻辑。
研究动机
现有视频扩散蒸馏方法几乎全部建立在 consistency models 之上,代表性方法如 rCM(双向)与 Self-Forcing(因果)虽然在 4 NFE 极低预算下表现亮眼,但随着采样步数增加质量反而下降。论文 Fig. 1 清晰展示了这一现象:rCM 在 32 步时 VBench 分数从 4 步的 83.73 跌到 75.0 附近,Self-Forcing 同样出现退化。结构原因在于一致性采样需要在每一步重新对中间状态加噪再执行去噪,多次迭代后轨迹逐渐偏离教师模型的 PF-ODE,暴露出 cumulative bias 与 exposure bias 两类误差。Table 1 进一步对比了 APT1&2、rCM、Self-Forcing 等主流蒸馏方法,发现它们无一例外都受限于 consistency 范式,无法在少量步数与多步数之间实现平滑过渡,这一缺陷在需要灵活调度推理预算的实际部署中尤为突出。
本文的目标是本文的目标是设计一个真正支持 any-step 推理的视频扩散蒸馏框架:单个模型既能以 4 NFE 生成可用预,又能以 32 NFE 输出更高保真度结果,覆盖从 2 NFE 到 32 NFE 的全谱系采样预算。具体而言,AnyFlow 试图在 1.3B 到 14B 参数量级、双向与因果两类架构上,同时匹配或超越现有 consistency-based 蒸馏方法的少步性能,并恢复 ODE 采样应有的 test-time scaling 行为,最终还要支持在下游数据集上继续微调。
与已有工作不同的是,现有工作要么停留在 consistency 框架内修补缺陷(如 rCM、Self-Forcing),要么以流图替代 consistency 但仅停留在前向训练阶段(如 MeanFlow、Score Distillation 的早期工作)。本文的独特切入角度是首次把 flow map 与 on-policy distillation 真正结合,并提出专为 flow map 设计的 backward simulation——利用流图可组合性将完整 Euler 轨迹拆分为 shortcut 段,既避免 consistency 加噪再投影带来的漂移,又显著降低长步数下的模拟成本。Table 1 显示这是首个同时支持因果 + 双向 + any-step 的视频扩散蒸馏框架,区别于同期 TMD(仅双向、通过架构共享提高效率)。
核心方法
AnyFlow 的整体思路可以类比为'双阶段接力':第一阶段利用教师模型合成的 256K prompt-video 数据对预训练视频扩散模型进行 forward flow map 训练,把它改造成一个稳定的双时间步转移算子 $f_\theta(z_t, t, r) \approx z_r$,这一步只做离线学习、不涉及学生模型自身 rollout;第二阶段在第一阶段的检查点基础上,联合优化 forward flow map 目标与 DMD-style on-policy 目标,关键技术是 flow map backward simulation——把从 $z_T$ 到 $z_0$ 的完整 Euler rollout 拆解为 $T\to t$、$t\to r$、$r\to 0$ 三段,其中两端的 shortcut 由 $f_\theta$ 直接跳到目标时间,中间段才是真正需要回传梯度的子区间,从而使 on-policy 蒸馏在不同推理步数预算下都能保持可控计算量。整条流水线在 Wan2.1 主干网络上以 LoRA rank=256 进行参数高效微调,最终在双向 1.3B/14B 与因果 1.3B/14B 上同时落地。
AnyFlow 与已有方法的本质区别在于'用什么做蒸馏目标、用什么做反向模拟'两点同时革新。蒸馏目标方面,传统 consistency 把学习目标限定在 $z_t\to z_0$ 的端点映射,本文改为任意 $(t, r)$ 对的转移,使模型天然支持任意步长离散化;反向模拟方面,传统 consistency backward simulation 需要在每步把中间状态重新加噪并执行端点投影(导致 trajectory drift),而本文利用 flow map 的可组合性 $f_\theta(z_t, t, q) \approx f_\theta(f_\theta(z_t, t, r), r, q)$ 把长轨迹 shortcut 为三段 $T\to t$、$t\to r$、$r\to 0$,使得任意 $N$ 步推理预算都可以用恒定的 3 次前向 + 3 次反向传播来模拟,从而既消除 discretization error 又降低训练成本(16 步时相对 consistency 模拟降低 43.4%-47.0%)。
方法步骤详情
AnyFlow 训练分为两个串行阶段。阶段一(Algorithm 1:Forward Flow Map Training)逐批采样 (x, c, t, r),其中 $t, r$ 由时间采样器(最终采用 Beta(2, 1.5) 权重)生成,$z = (1-t)x + t\epsilon$,然后在 interpolated timestep condition $g\cdot\mathrm{emb}(t) + (1-g)\cdot\mathrm{emb}'(r)$($g=0.25$)下让模型预测 $u_\theta(z, t, r)$,再通过差分推导方程 $du/dt \approx [u(z_+, t_+, r) - u(z_-, t_-, r)]/(2\epsilon)$ 计算目标速度 $u_{tgt}$,最终用 adaptive reweighted loss $\mathcal{L} = w_{t,r}\|u_\theta - \mathrm{sg}(u_{tgt})\|^2$ 优化,其中 $w_{t,r} = 1$ 当 $t=r$、$w_{t,r} = \mu_{t=r}/(\|\Delta\|_2^2 + c)$ 当 $t\neq r$,$\mu_{t=r}$ 为边界样本的指数移动平均损失。阶段二(Algorithm 2:On-Policy Flow Map Distillation)每轮先随机采样推理步数 $s\in[1, S]$ 与梯度步 $t\in[1, s]$,令 $r = t - T/s$、$v = \epsilon - x$,从 $z_T$ 出发依次执行 $z_t = f_\theta(z_T, T, t)$、$z_r = f_\theta(z_t, t, r)$、$z_0 = f_\theta(z_r, r, 0)$(前两步为 shortcut,最后一步为目标梯度区间),把 $z_0$ 与 $x$ 一起送入 DMD 的 real-fake score 差分计算 KL 梯度,再把梯度沿整条 shortcut 链回传到 $f_\theta$;同时为防止遗忘,与阶段一相同的 flow map 目标被联合优化。双向模型直接套用上述两阶段;因果模型额外采用 FAR 框架,用非均匀 chunk 划分(首块 1 帧 + 后续每块 3 帧)联合支持 T2V/I2V,并缓存 rollout 期间的 KV 以压缩内存。
技术新颖性
技术新颖性体现在三个层面。第一,公式层面:首次把 MeanFlow 风格的 flow map 训练迁移到视频扩散后训练,并通过 interpolated timestep conditioning($g\cdot\mathrm{emb}(t) + (1-g)\cdot\mathrm{emb}'(r)$)解决 TMD 零初始化方案导致的嵌入范数爆炸(Fig. 11 显示 2K 步后范数稳定),避免过饱和生成。第二,模拟层面:flow map backward simulation 是本文最具突破性的设计,它把'模拟任意步数 Euler 轨迹'的问题从全链模拟降为常量 3 段 shortcut,巧妙利用了流图可组合性 $f_\theta(z_t, t, q) \approx f_\theta(f_\theta(z_t, t, r), r, q)$,使 16 步蒸馏相对 consistency 模拟提速近一半(Table 5)。第三,框架层面:首次实现'因果 + 双向 + any-step'三合一视频扩散蒸馏,并通过 FAR 非均匀 chunk 把 I2V 与 T2V 统一到同一模型,区别于 TMD(仅双向且通过架构侧效率提升)。同时 AnyFlow 保留教师模型的瞬时流场,因而可继续在下游机器人/自动驾驶/游戏数据上微调(Fig. 7-8),这是 Self-Forcing 蒸馏模型难以做到的。
实验结果
AnyFlow 在三类实验上同时取得 SOTA 表现。Table 2 的消融实验显示,在 14B 双向设置下,仅做 forward flow map training 已能拿到 81.75/83.40(4/32 NFE)总分,优于 Flow Matching 训练(74.64/83.83)与 Consistency ODE-Init(80.44/82.86);叠加 on-policy distillation 后,flow map backward simulation 在 4 NFE 取得 83.48、32 NFE 取得 83.96,相比 consistency backward simulation 的 83.55/82.96 消除了多步退化,验证了核心设计。Table 3 的主对比显示,14B 双向 AnyFlow 在 4 NFE 即达到 84.04,超过 rCM-14B 的 83.73(同样 4 NFE),且 32 NFE 进一步涨到 84.10;14B 因果 AnyFlow-FAR 在 4 NFE 拿到 84.05、32 NFE 涨到 84.41,全面超越 Krea-Realtime-14B(83.25@4 NFE)、LightX2V-WAN2.1-14B-CausVid(83.42@9 NFE)与 FastVideo-CausalWan2.2-A14B(83.12@8 NFE)。1.3B 规模同样保持优势:AnyFlow-Wan2.1-T2V-1.3B 在 4 NFE 取得 83.48、32 NFE 涨至 83.96,相对 rCM-1.3B 的 82.51 提升约 1 分;AnyFlow-FAR-1.3B 在 4 NFE 拿到 83.54、32 NFE 涨至 83.96,超过 Self-Forcing 的 83.39。Table 4 的 I2V 实验中,AnyFlow-FAR-14B 仅用 4 NFE 即取得 VBench-I2V 总分 87.87(Quality 80.39、I2V 95.35),与 Wan2.1-I2V-14B 在 50×2 NFE 下的 87.71 持平且 I2V 分数更高,并显著超过 FastVideo-CausalWan2.2-A14B-Preview 的 86.82@8 NFE,意味着 AnyFlow 能在 12.5× 的加速比下保持同等质量。Table 5 的训练成本显示,16 步模拟下 AnyFlow 相对 consistency backward simulation 在因果模型上减少 43.4% 时间(53.1 vs 93.8 s/iter)、双向模型减少 47.0%(51.2 vs 96.6 s/iter),而 4 步时仅增加 15.7%/22.5%,证实了 shortcut 设计的实用价值。Fig. 10 进一步消融了 $w(t)$ 权重策略,Beta(2, 1.5) 在双架构下均取得最佳综合表现。Fig. 11 则从嵌入范数角度验证了 interpolated timestep conditioning 抑制了零初始化方案带来的过饱和现象。Fig. 6 与 Fig. 12 的定性比较显示,AnyFlow 在 32 NFE 时仍能持续提升运动质量与纹理清晰度,而 consistency baseline 在多步后会出现过平滑与运动失真。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| Text-to-Video (14B 双向) | VBench Total Score (4 NFE) | 84.04 (AnyFlow-Wan2.1-T2V-14B) | 83.73 (rCM-Wan2.1-T2V-14B) | +0.31 VBench Total,相同 NFE 下反超 rCM 14B |
| Text-to-Video (14B 因果) | VBench Total Score (4 NFE) | 84.05 (AnyFlow-FAR-Wan2.1-14B) | 83.25 (Krea-Realtime-Wan2.1-14B) | +0.80 VBench Total,并在 32 NFE 进一步提升到 84.41 |
| Image-to-Video (14B 因果) | VBench-I2V Total Score (4 NFE) | 87.87 (AnyFlow-FAR-Wan2.1-14B) | 87.71 (Wan2.1-I2V-14B, 50×2 NFE) | +0.16 分,但 NFE 仅为对方的 1/25,加速比约 12.5× |
| Text-to-Video (1.3B 双向) | VBench Total Score (4 NFE) | 83.48 (AnyFlow-Wan2.1-T2V-1.3B) | 82.51 (rCM-Wan2.1-T2V-1.3B) | +0.97 分,32 NFE 时进一步达到 83.96 |
| Text-to-Video (1.3B 因果) | VBench Total Score (4 NFE) | 83.54 (AnyFlow-FAR-Wan2.1-1.3B) | 83.39 (Self-Forcing) | +0.15 分,且在 32 NFE 时涨到 83.96,体现真正 any-step 能力 |
| 训练效率 (16 步 on-policy) | 秒/迭代(8×H100,1.3B) | 53.1 (因果) / 51.2 (双向) | 93.8 (因果) / 96.6 (双向) consistency backward simulation | 分别减少 43.4% 与 47.0% 训练时间 |
局限与改进
作者明确承认两点局限。其一,AnyFlow 仍依赖外部数据集(包括 256K 合成 prompt-video)做 flow map 训练,即便使用教师模型合成数据,训练分布与原始教师预训练数据之间仍存在分布漂移,可能导致部分结果出现偏柔和的纹理风格,作者建议未来使用与教师完全相同的原始数据。其二,第一阶段 forward flow map 训练对噪声水平敏感(Fig. 10 表明不同 $w(t)$ 策略对 32 NFE 性能影响 0.2 分),时间采样器仍需在 uniform/normal/Beta 中手工挑选,并未达到 fully self-adaptive。从独立观察角度,论文对'为何 Beta(2, 1.5) 优于 logit-normal'缺乏理论解释;Table 3 中 1.3B 因果模型在 4 NFE 与 Self-Forcing 相比仅领先 0.15 分,绝对差距很小;且所有实验都建立在 Wan2.1 主干之上,未在 HunyuanVideo、CogVideoX 等其他基座上验证,方法的可迁移性需要额外实验确认。此外,论文对双向模型在 4 NFE 取得 84.04 vs 32 NFE 84.10 的差异也说明真正最大化的 any-step 上限并未在双向设置上充分释放,可能受限于 Wan2.1-14B 教师本身的 capacity。
独立分析的弱点
从独立视角看,AnyFlow 仍存在以下可改进之处。第一,第一阶段对合成数据质量的依赖过强:教师模型本身在 50×2 NFE 时的 VBench Total 为 83.74(14B 双向),但流图训练后的模型在 4 NFE 已经超过该值(84.04),暗示 forward flow map 训练实际承担了相当一部分泛化压力;若教师本身存在视觉伪影,学生很可能继承此类缺陷,因此阶段一的数据治理是隐性瓶颈。第二,guidance-fused training 把 CFG 嵌入到训练目标中($g=0.25$),但推理时如何选择 effective CFG 强度论文未充分讨论,意味着实际部署时仍需重新搜索 inference-time guidance,可能限制'即插即用'的便利性。第三,flow map backward simulation 的 shortcut 策略固定为 $T\to t$、$t\to r$、$r\to 0$ 三段,对于极低 NFE(如 1-2 步)情形,shortcut 段本身可能与目标 $t\to r$ 段跨度相当,优势会缩小;论文也未明确给出 NFE 下界。第四,因果模型依赖 FAR 框架与非均匀 chunk,这意味着把 AnyFlow 推广到其他非 FAR 风格的因果基座时需要重新调整 chunk 划分与 KV-cache 策略,方法的架构耦合性较强。第五,论文只汇报 VBench 数值,未提供人工评估或 FVD/IS 等更全面的指标,对美学与运动自然度的差异缺乏定量刻画。
未来方向
作者已指明两条未来方向。其一是开发更稳定、自适应的 forward flow map 训练策略,减少对 $w(t)$ 等超参的敏感性,使方法在更宽 NFE 区间内稳健。其二是把 AnyFlow 扩展到 autoregressive 长视频生成,依赖其保留瞬时流场的特性做继续预训练/微调。从成果可延伸的角度看,flow map backward simulation 的 shortcut 思想完全可以推广到 3D/4D 视频、音频-视频联合扩散乃至科学计算中的 ODE 求解加速;flow map 的可组合性也为模型并行化(不同时间区间分配到不同设备)打开新可能;on-policy distillation 中 DMD 损失与 GAN/对比损失的混合也是值得探索的方向。考虑到 AnyFlow 在保留 base model 风格上的优势,把它作为 personalized video diffusion 或可控生成(camera motion、character consistency)的统一底座也具有现实价值。
复现评估
论文在可复现性上表现良好:作者承诺代码与模型权重将发布在 https://github.com/NVLabs/AnyFlow,并基于开源的 Diffusers 框架实现,主干网络为公开的 Wan2.1 系列。训练数据明确说明为 256K prompt-video 对,由 Wan2.1-T2V-14B 合成,单样本最多 81 帧、480×832 分辨率。超参数全部给出:阶段一用 AdamW 学习率 $5\times 10^{-5}$,全局 batch size 32(1.3B)/16(14B),分别训练 6000/4000 步;阶段二学习率 $2\times 10^{-6}$ 训练 800 步;两阶段均使用 LoRA rank=256 节省显存。评估采用官方 VBench augmented prompts 重新跑分,保证与基线公平对比。算力门槛方面,1.3B 阶段一单次迭代约 10.4 秒(8×H100),阶段二 16 步模拟约 53 秒/迭代,14B 模型训练需更多节点。复现主要难点包括:Wan2.1 14B 模型的合规获取、256K 合成数据的复现成本、阶段二同时维护 student + teacher + DMD 判别器三组网络所需的显存,以及 interpolated timestep condition、adaptive reweighting 等新增超参的联合调优,整体复现难度偏高。
论文图表