PyramidalWan:将预训练视频模型金字塔化以实现高效推理 PyramidalWan: On Making Pretrained Video Model Pyramidal for Efficient Inference
通过低成本微调将预训练视频扩散模型转化为金字塔模型,实现4.5倍FLOPs降低且无质量损失
前置知识
扩散模型(Diffusion Model)
一类生成模型,通过逐步向数据添加高斯噪声使其变为纯噪声(前向过程),然后训练一个神经网络学习从噪声逐步去噪恢复原始数据(反向过程)。视频扩散模型通常在潜空间(latent space)操作,先用VAE编码器将视频压缩为低维表示,再在潜空间中执行扩散过程,最后用解码器还原。推理时需要多步去噪(通常25-50步),因此计算成本很高。
本文的核心问题就是视频扩散模型推理计算量太大,理解扩散模型的基本原理是理解后续优化方案的前提。
流匹配(Flow Matching)
一种训练生成模型的框架,可以看作扩散模型的一种连续时间变体。它学习一个速度场(velocity field),将噪声分布沿直线路径(线性插值)平滑地传输到数据分布。给定干净信号 $x_0$ 和噪声 $\epsilon$,带噪信号 $x_\sigma = (1-\sigma)x_0 + \sigma\epsilon$,其中 $\sigma$ 为噪声水平。网络 $F(x_\sigma, \sigma)$ 学习预测 $dx_\sigma/d\sigma$。Wan2.1模型就采用流匹配目标训练。
本文的金字塔微调和蒸馏都基于流匹配框架,理解其数学形式才能看懂后续损失函数设计。
光谱自回归(Spectral Autoregression)
一个关于自然信号去噪过程的观测:在频域分解中,高斯噪声对所有频率分量的破坏速率基本一致,但自然信号(图像、视频)的高频分量幅度通常远低于低频分量。因此,在去噪初期,高频信息已经被噪声淹没,而低频信息仍有信噪比。这意味着高噪声水平时只需要在低分辨率处理,低噪声时才需要高分辨率——这就是金字塔采样的理论基础。
这是本文金字塔化策略的核心动机,理解了这一点才能明白为什么不同噪声水平可以对应不同分辨率。
分布匹配蒸馏(Distribution Matching Distillation, DMD)
一种将多步扩散模型压缩为少步模型的蒸馏方法。核心思想是训练一个学生模型在一步或几步内生成样本,使其分布与教师模型多步生成的分布匹配。具体实现中,学生模型预测干净信号后重新加噪,用一个fake score网络(基于教师初始化)提供梯度信号,通过反向KL散度引导学生输出匹配教师分布。相比一致性模型等方法,DMD在视觉质量上表现更好。
本文在金字塔模型基础上进一步研究步数蒸馏策略,DMD是其中效果最好的方法,理解其原理才能看懂金字塔DMD的变体设计。
金字塔补丁化流(Pyramidal Patchification Flow, PPF)
一种与PyramidalFlow不同的多分辨率推理策略。PyramidalFlow通过改变输入/输出的空间时间分辨率来降低计算量,而PPF保持原始分辨率不变,而是调整DiT中patchification层的卷积核大小:噪声较大时用大核(产生更少的token),噪声较小时用小核(产生更多token)。这样transformer块在不同阶段处理的token数量与PyramidalFlow相同,但无需在阶段间切换分辨率。
PPF是本文的对比方法之一,也是首次被证明可用于少步视频生成,理解它有助于评估两种策略的优劣。
研究动机
近年来视频扩散模型(如Wan2.1、HunyuanVideo、CogVideoX等)在生成质量上取得了显著进步,但这些能力的代价是极高的推理计算量。以本文使用的Wan2.1-1.3B模型为例,生成81帧448×832的视频需要50步去噪,总计算量高达2×12,592 TFLOPs(因为使用了classifier-free guidance,计算量翻倍)。这种计算开销严重限制了视频生成模型在实际应用中的部署,尤其是在延迟敏感的场景下。现有的加速手段主要包括步数蒸馏和架构优化(如注意力机制优化),但它们各自存在局限:步数蒸馏将50步压缩到2-4步时质量尚可,但进一步压缩到单步时质量会大幅下降(Wan-DMD单步的VBench总分从83.33骤降到79.45);架构优化则需要修改模型结构或注意力实现,对已有模型的兼容性有限。
本文的目标是本文的目标是提出一种高效的推理加速方案,将预训练的视频扩散模型转化为金字塔(pyramidal)版本,在保持原始模型视觉质量的前提下显著降低推理计算量。具体量化目标包括:(1)将推理FLOPs降低至少4倍(从2×12,592 TFLOPs降至约2×2,821 TFLOPs甚至更低);(2)保持VBench和VBench-2.0等自动指标与原始50步模型可比;(3)用户主观评价不显著劣于高计算量基线。此外,作者还希望在金字塔化基础上进一步通过步数蒸馏将推理成本压缩到极致(如267 TFLOPs),同时研究金字塔场景下不同的蒸馏策略。
与已有工作不同的是,此前的金字塔视频模型(如PyramidalFlow、TPDiff)都是从零开始训练的,在有限计算资源下训练导致生成质量不及当时的SOTA系统。本文抓住了一个关键洞察:现有的开源金字塔模型性能不佳,不是因为金字塔框架本身有问题,而是因为从零训练的模型质量不够。本文的独特角度是利用已有的高质量预训练模型(Wan2.1-1.3B),通过低成本的微调(仅5K迭代、2张H100 GPU)将其转化为金字塔版本,从而继承预训练模型的生成质量。同时,本文首次系统研究了金字塔模型中的步数蒸馏策略,包括使用原始非金字塔教师和金字塔教师两种DMD变体,以及对抗蒸馏方法。此外,作者还将PyramidalFlow的阶段间转换操作理论推广到基于正交变换的任意上采样函数(如小波变换),这在理论层面是有意义的贡献。
核心方法
本文的方法可以用一个简单的类比来理解:想象一个画家在画一幅画。传统扩散模型像是一直在最终尺寸的画布上作画,即使在最初只是涂底色的时候也是如此。金字塔模型的做法则是:先在一张小画布上画出粗略的轮廓和结构(对应低分辨率、高噪声的阶段),然后逐步放大画布、增加细节(中等分辨率、中等噪声),最后在完整尺寸的画布上精雕细琢(高分辨率、低噪声)。技术路线上,首先基于频谱分析确定各阶段的噪声水平边界(即每个分辨率适合处理什么范围的噪声),然后用金字塔流匹配损失微调预训练的Wan2.1-1.3B模型,使其学会在不同分辨率下处理不同噪声水平的输入。在推理时,从最低分辨率的纯噪声开始,逐步去噪并上采样,最终在原始分辨率得到干净的视频。此外,作者进一步研究了在此基础上的步数蒸馏,用DMD或对抗方法将每个阶段的步数压缩到1-2步。
本文最核心的创新在于证明了金字塔化可以通过微调而非从零训练来实现,且不会损失视觉质量。此前的工作(PyramidalFlow、TPDiff)都假设金字塔模型需要从头训练,这导致训练预算有限时质量不佳。本文发现,预训练模型已经学到了强大的视觉生成能力,金字塔化只需要教会模型「在不同分辨率下处理不同噪声水平」这一新技能即可。具体实现上,作者采用定义②的方式构建各阶段的干净信号:先在RGB空间降采样视频,再用VAE编码($x^{(i)}_0 = E(\text{Down}(V, i))$),而非在潜空间降采样。实验表明这种方式产生的视觉效果显著优于在潜空间降采样的定义①。另一个关键创新是将PyramidalFlow中基于平均池化和最近邻插值的阶段间转换操作,推广到基于任意正交变换(如Haar小波)的上采样函数,这为金字塔模型的灵活性提供了更大的设计空间。
方法步骤详情
整个方法包含以下几个关键步骤:(1)频谱分析与阶段划分:对训练数据集的VAE潜空间进行频谱分析,确定哪些噪声水平下高频信息已被充分破坏,据此设置三个阶段的边界噪声水平。阶段2(最低分辨率21×112×208)处理噪声最大的部分($\sigma \in [\sigma^{(2)}_c, \sigma^{(2)}_n]$),阶段1(41×224×416)处理中等噪声($\sigma \in [\sigma^{(1)}_c, \sigma^{(1)}_n]$),阶段0(81×448×832)处理低噪声到干净信号($\sigma \in [\sigma^{(0)}_c, 0]$)。(2)金字塔微调:使用预训练的Wan2.1-1.3B作为初始化,用金字塔流匹配损失 $\mathcal{L}_{pyr}$ 进行微调。在每个阶段,采样边界干净信号 $y^{(i)}_c$ 和 $y^{(i)}_n$,通过线性插值得到带噪输入 $x^{(i)}_\sigma$,网络学习预测 $dx^{(i)}_\sigma/d\sigma$。同时加入蒸馏损失 $\mathcal{L}_{dist}$ 将部分去噪的潜变量与教师预测对齐。(3)阶段间转换:在阶段 $i+1$ 的去噪完成后,通过 $R^\uparrow_N$ 操作上采样并添加去相关噪声,使其成为阶段 $i$ 的有效高噪声边界。这个操作满足分布等价性:$R^\uparrow_N N(y^{(i+1)}_c) \doteq y^{(i)}_n$。(4)步数蒸馏(可选):在金字塔微调的基础上,进一步用DMD或对抗蒸馏将每个阶段的步数压缩。DMD训练学生网络在一步内预测干净信号,用fake score网络提供梯度;对抗蒸馏则训练判别器区分生成样本和真实样本的特征。
技术新颖性
本文的技术新颖性体现在多个层面。首先,在应用层面,本文首次证明了预训练的视频扩散模型可以成功金字塔化,打破了此前「金字塔模型必须从零训练」的认知。这使得金字塔化成为一种即插即用的推理加速方案,无需承担昂贵的训练成本(5K迭代 vs PyramidalFlow的200K+迭代)。其次,在蒸馏层面,本文首次系统比较了金字塔模型中的多种步数蒸馏策略,包括使用原始非金字塔教师(DMD-OT)和金字塔教师(DMD-PT)的DMD,以及使用原始骨干(Adv-OD)和金字塔骨干(Adv-PD)的对抗蒸馏。特别是发现金字塔教师的DMD在实践中表现更好(视觉质量上DMD-PT*最优),而原始教师的DMD虽然指标略好但会出现颜色过饱和和卡通化伪影。第三,在理论层面,作者将PyramidalFlow中基于平均池化和最近邻插值的上/下采样操作推广到了基于任意正交变换(如小波变换)的一般框架,并证明了Haar小波操作可以作为这个框架的特例。这一推广使得金字塔模型可以使用更灵活的降采样策略,为未来的架构探索提供了理论基础。此外,本文首次展示了PPF框架可以成功用于少步视频生成(Wan-PPF-DMD),尽管在有限训练预算下其扩散训练效果不及PyramidalFlow。
实验结果
本文的实验结果可以分为几个层面来分析。首先是金字塔化的核心效果:PyramidalWan(20-20-10步调度)在VBench上获得82.83分,与原始Wan2.1-1.3B的50步(82.49分)相当甚至略优,同时FLOPs仅为后者的约22%(2×2,821 vs 2×12,592 TFLOPs)。在VBench-2.0上,PyramidalWan获得54.93分(原始模型56.02分),差距很小。特别值得注意的是,PyramidalWan在语义得分(Semantic score)上达到80.70,是所有模型中最高的,超过了原始模型的78.57,说明金字塔化甚至有助于语义理解。其次是步数蒸馏的效果:在非金字塔模型中,Wan-DMD(2步)获得83.28分,与50步模型相当,但单步推理时质量大幅下降至79.45分。金字塔模型在2-2-1调度下(阶段2和1各2步,阶段0仅1步)弥补了这一空缺,PyramidalWan-DMD-OT获得82.86分,PyramidalWan-DMD-PT*获得82.72分,均接近50步模型的水平。在VBench-2.0上,PyramidalWan-DMD-OT的55.36分是所有少步模型中最高的,超过了Wan-DMD(2步)的56.67分之外的所有模型。第三是延迟方面的实际收益:PyramidalWan在2-2-1调度下的单次去噪器前向传播总延迟为810毫秒(标准差<10ms),而Wan-PPF-DMD在相同调度下的总延迟为810毫秒(标准差<10ms),均远低于原始模型的多次前向传播延迟。对比Jenga等动态token减少方法,金字塔模型具有静态计算图的优势,更适合资源受限设备的部署。第四是用户研究结果:700份用户评价显示,PyramidalWan-DMD-PT*(2-2-1)与50步Wan模型和2步Wan-DMD相比,用户均未发现显著的质量差异(p值<0.001的检验均被拒绝),表明金字塔化的视觉质量得到了人类评估者的认可。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 视频生成质量 | VBench Total | 82.83(PyramidalWan 20-20-10) | 82.49(Wan2.1-1.3B 50步) | +0.34分,同时FLOPs降低78% |
| 视频生成质量 | VBench Total | 82.86(PyramidalWan-DMD-OT 2-2-1) | 83.33(Wan-DMD 4步) | 仅低0.47分,但FLOPs从534降至267 |
| 视频生成质量 | VBench-2.0 Total | 55.36(PyramidalWan-DMD-OT 2-2-1) | 56.02(Wan2.1-1.3B 50步) | 仅低0.66分,但推理速度快约47倍 |
| 视频生成质量 | VBench Semantic | 80.70(PyramidalWan 20-20-10) | 78.57(Wan2.1-1.3B 50步) | +2.13分,语义理解显著提升 |
| 推理效率 | FLOPs (TFLOPs) | 267(PyramidalWan-DMD-PT* 1-1-1) | 2×12,592(Wan2.1-1.3B 50步) | 降低约94倍 |
| 用户偏好 | 无显著偏好比例 | 29.1%(vs Wan 50步) | 41.7%偏好50步 | 用户未发现显著质量差异(p<0.001) |
局限与改进
本文存在几个值得注意的局限性。首先,在VBench-2.0的某些维度上,金字塔模型与非金字塔基线存在明显差距。特别是Creativity维度,PyramidalWan-DMD-PT*(2-2-1)仅获得34.81分,而Wan-DMD(4步)获得50.87分;Controllability维度上,金字塔模型为28.48分,而Wan-DMD(2步)为37.07分。这表明金字塔化和步数蒸馏的叠加可能在某些需要精细控制的任务上造成退化。其次,PPF(补丁化金字塔流)在视频生成上的表现令人失望:使用相同的训练预算和数据集,PPF扩散模型未能收敛,生成质量不令人满意。虽然DMD蒸馏可以在PPF扩散模型上成功应用(Wan-PPF-DMD),但基础扩散模型本身的质量问题限制了其上限。第三,使用原始教师的DMD蒸馏(DMD-OT)虽然获得了最好的定量指标,但经常出现颜色过饱和和卡通化伪影,这在用户研究中未被充分捕获(可能因为评估协议侧重整体质量而非细节保真度)。第四,金字塔模型在Dynamic Degree指标上表现不佳,PyramidalWan仅为44.44分(原始模型69.17分),说明生成视频的运动多样性有所下降。最后,作者也承认VBench-2.0指标与金字塔模型之间存在一定差距,这可能反映了当前评估体系对金字塔化带来的特定优势(如语义理解提升)和劣势(如运动动态性下降)的敏感度不同。
独立分析的弱点
从独立分析的角度,本文存在几个值得关注的弱点。第一,第一帧的特殊处理引入了架构复杂性。由于2D和3D降采样的缩放因子不同,第一帧的全局噪声水平与其他帧不同,作者不得不添加特殊的可学习嵌入向量来标记第一帧。这种设计虽然解决了兼容性问题,但增加了额外的超参数和潜在的不稳定性。更好的方向可能是重新设计时间维度的降采样策略,使其对所有帧一致。第二,DMD-PT与DMD-PT*之间的理论-实践悖论未被充分解释。DMD-PT*在理论上缺乏充分的理论支撑(将权重简化为 $\tilde{\beta}_1 = \tilde{\gamma}_1 = 1, \tilde{\beta}_2 = \tilde{\gamma}_2 = 0$),但实践中反而表现更好。作者将其归为未来工作,但这个问题对于理解和改进金字塔DMD至关重要。一个可能的方向是分析简化版本的隐式正则化效果。第三,训练数据依赖合成样本是一个潜在风险。作者使用Wan2.1-14B生成的80K合成视频进行训练,并发现合成数据优于真实视频样本。这引发了一个问题:蒸馏后的模型是否会继承教师模型的偏见和伪影,尤其是在长尾分布的场景下。改进方向可以是探索混合真实和合成数据的训练策略。第四,VBench-2.0上Creativity和Controllability维度的显著下降(34.81和28.48 vs 基线的50.87和36.36)暗示金字塔化可能限制了模型的创造性和可控性,这在需要高度创意或精确控制的应用场景(如广告制作、交互式创作)中可能是不可接受的。可以探索在这些维度上添加专门的损失项或正则化。
未来方向
基于本文的成果,可以延伸出多个有价值的研究方向。作者提出的方向包括:(1)进一步研究DMD-PT*简化版本为何在实践中表现更好,探究其隐式正则化机制;(2)探索金字塔化与其他推理加速技术(如注意力优化、token合并)的组合效果。基于本文的成果,还可以延伸出以下方向:(3)自适应阶段划分:当前的阶段边界(噪声水平阈值)是基于数据集的整体频谱分析确定的,对所有输入共享。可以根据每个具体视频的内容复杂度动态调整阶段划分,在复杂场景使用更多高分辨率步骤,在简单场景更多使用低分辨率。(4)金字塔化应用于更大模型:本文仅在1.3B参数的Wan2.1上验证,将其应用于14B甚至更大参数的模型(如Wan2.1-14B)可能带来更显著的效率收益,因为大模型的单步计算成本更高。(5)多模态输入的金字塔化:将金字塔框架扩展到文本-视频-音频多模态生成,在不同分辨率阶段使用不同粒度的条件信息。(6)金字塔化与检索增强:在低分辨率阶段利用检索到的参考图像或视频片段提供全局结构指导,在高分辨率阶段专注于细节生成。
复现评估
本文的可复现性总体较好。首先,所有模型基于开源的Wan2.1-1.3B,代码和权重均可获取。训练数据为80K合成视频(由Wan2.1-14B生成),作者提供了生成提示词的数据集来源(Lin et al. [31]的350K子集)。训练硬件需求适中:金字塔微调仅需2张H100 GPU训练5K迭代,这在学术研究环境下是可行的。步数蒸馏的硬件需求稍高,DMD-PT需要2张H100 GPU(5K迭代),而DMD-OT的教师微调需要更多资源。文中还详细说明了编译器设置(PyTorch compiler with max-autotune)、学习率、噪声水平选择等关键超参数。开源地址为 https://qualcomm-ai-research.github.io/PyramidalWan。需要注意的复现细节包括:(1)第一帧的特殊处理方式;(2)PPF训练时形状不一致需要三线性插值;(3)FlexAttention用于混合不同分辨率的批次训练。总体而言,一个熟悉扩散模型训练的团队应该能在1-2周内完成核心结果的复现。
论文图表