Command Palette
Search for a command to run...
面向快速图像与视频生成的并行解码蒸馏
面向快速图像与视频生成的并行解码蒸馏
Neta Shaul Chao Liu Arash Vahdat Julius Berner
摘要
视频扩散或流匹配模型的生成过程因缓慢的迭代采样而计算开销高昂。当前最先进的加速方法严重依赖变分分数蒸馏(VSD)和对抗损失,将扩散模型蒸馏为少步生成器。尽管这些训练损失能实现高质量的视频生成,但它们优化难度极大且容易出现模式坍塌,导致视频多样性下降和运动感缺失。本文提出并行解码蒸馏(PDD),一种简洁且可扩展的、基于轨迹的蒸馏方法,用于扩散模型和流匹配模型的快速推理。我们的架构和训练流程兼容任意预训练模型,并支持使用不同数量的函数评估次数(NFE)进行采样。PDD 通过在每次网络评估中预测多个去噪步骤来加速生成。其核心思想是学习平均速度的表示,而无需借助 JVP 或有限差分近似来回归其导数。在 LTX-2.3 文本到视频/音频、Wan 14B 文本到视频和 Qwen-Image 文本到图像任务上,我们的方法以 4-8 NFE 取得了最先进的性能。此外,PDD 在生成视频的多样性方面也展现出显著提升。
一句话总结
并行解码蒸馏(PDD)由 NVIDIA 和魏茨曼科学研究所提出,是一种基于轨迹的蒸馏方法,通过每次网络评估预测多个去噪步骤来加速扩散和流匹配模型,学习无需导数回归的平均速度表示,并在 LTX-2.3、Wan 14B 和 Qwen-Image 上以 4–8 NFE 达到最先进性能,同时显著提升生成视频的多样性。
核心贡献
- 并行解码蒸馏(PDD)是一种基于轨迹的蒸馏方法,训练一个并行解码器在一次网络评估中预测多个连续去噪区间的平均速度,避免了 VSD、对抗损失、JVP 和有限差分近似。
- PDD 与任何预训练的扩散或流模型兼容,支持可变 NFE,并在 LTX-2.3 文本到视频/音频、Wan 14B 文本到视频、Qwen-Image 文本到图像上以 4–8 NFE 达到最先进的少步性能。
- PDD 相比基于分布的基线显著提升了生成多样性和运动,成为首个用于少步高分辨率视频生成的纯轨迹蒸馏方法。
引言
扩散和流模型已成为高质量文本到图像和视频生成的骨干,但其迭代采样过程通常需要数百次网络评估,对交互式编辑和世界建模等实时应用造成无法接受的延迟。基于轨迹的蒸馏方法可以加速图像生成,但在大型视频模型上因质量下降和依赖雅可比向量积或有限差分的昂贵训练而遇到困难。基于分布的替代方法虽然更易扩展到视频,但经常出现模式坍塌、运动多样性降低以及对超参数敏感的问题。作者引入并行解码蒸馏(PDD),这是一种基于轨迹的方法,不将多个步骤压缩为一步,而是训练一个并行解码器在一次前向传递中预测多个连续时间区间的平均速度。该设计产生了一个简单的回归目标,避免了对抗损失、雅可比向量积和多阶段训练,并且在推理时自然支持可变数量的函数评估。该方法在大规模视频和图像模型上实现了最先进的少步生成,同时相比基于分布的基线保持了更好的样本多样性。
方法
作者首先建立了生成流模型的基础,将流匹配和扩散模型视为由速度场 vt 和源分布 p0 定义的确定性过程。流过程 (Xt)0≤t≤1 遵循常微分方程:
dtdXt=vt(Xt),X0∼p0.目标是学习一个模型 vt,将来自源分布 p0 的样本映射到目标分布 p1。从训练好的流模型采样涉及通过将时间区间离散为 N 个更小的区间并顺序近似积分来数值求解该 ODE,通常使用欧拉或龙格-库塔求解器。
为了加速这种本质上的迭代采样过程,作者提出了并行解码蒸馏(PDD)。不是逐步近似积分步骤,而是学习一个并行解码模型,在单次网络评估中预测多个积分步骤。
如上图所示,采样轨迹被离散为 N 个区间,这些区间被分组为大小为 L 的块。对于时间步 tn 的状态 Xn,并行解码器 uˉnθ(⋅∣Xn) 被训练在一次前向传递中预测下一块 {n,…,n+L−1} 中所有区间的平均速度。一个块内的并行化过程定义为:
Xˉk+1=Xˉk+(tk+1−tk)uˉnθ(k∣Xn),对于 k∈{n,…,n+L−1},初始条件为 Xˉn=Xn。至关重要的是,并行化过程仅依赖于初始状态 Xn,因此整个块可以通过单次评估并行解码器来模拟。在生成过程中,作者应用块步规则同时推进 L 个区间:
Xˉn+L=Xn+k=n∑n+L−1(tk+1−tk)uˉnθ(k∣Xn).通过重复此规则 N/L 次,模型获得一个干净样本。
训练过程如上图所示。并行解码器通过回归到预训练教师流模型的平均速度的龙格-库塔近似来训练。作者采用在线策略训练,估计教师模型在学生输出上的平均速度。训练目标,即 PD 损失,定义为:
LPD(θ)=E[∥uˉnθ(k∣Xn)−uk(sg(Xˉk))∥2],其中 uk 是用单步欧拉或中点方法近似的教师平均速度,sg(⋅) 表示停止梯度算子。如右图所示,学生预测块的平均速度,从而产生块内状态 Xˉk。随机选择其中之一,并将学生的输出速度与教师在该区间内的平均速度匹配。由于近似平均速度仅需 1-2 次教师模型评估,该公式使损失在大规模下也可计算。
架构设计如上图所示。并行解码器使用与预训练流模型相同的骨干架构 Htθ,但修改了最终的输出层。作者不采用预测瞬时速度的单个线性层,而是学习 N 个线性层 W0θ,…,WN−1θ,分别对应网格中的每个时间步。并行解码器的架构为:
uˉnθ(k∣xn)=WkθHtnθ(xn).这种设计允许从预训练流模型的最后一层初始化,并使模型能够预测任意块大小,而无需引入第二个时间坐标。
在推理过程中出现了一个重要的优化,称为层融合。训练需要所有不同的学生输出方向来估计损失,而生成只需加权平均方向以跳过 L 个区间。作者定义了一个融合线性层:
Wn:n+Lθ=k=n∑n+L−1ΔkWkθ,Δk=tn+L−tntk+1−tk.因此,在推理时,共享骨干学习区间内平均速度的表示,作者通过仅在内存中保留每个块的一个融合线性层来避免扩大的最终层带来的额外计算。
实验
实验在类条件 ImageNet 生成、使用 Qwen-Image 的文本到图像合成以及使用 Wan2.1 和 LTX-2.3 的文本到视频生成上评估 PDD,证实了其广泛的适用性。PDD 始终支持灵活的少步推理,在相对于专用基线取得有竞争力或更优质量的同时,更好地保留了多样性和运动。该方法可有效扩展到大型多模态模型,并且可以在无数据的情况下训练,仅需数百次迭代即可匹配数据依赖蒸馏的性能。
PDD 蒸馏具有可变 NFE 的流图轨迹,避免雅可比向量积或有限差分,并使用融合线性头,简化了训练和推理。该设计在大规模基准上取得了有竞争力的视频质量、更高的运动性和更好的多样性,同时即使在少量训练迭代下也表现良好并防止模式坍塌。与欧拉/拉格朗日流图和 Pi-Flow 不同,PDD 支持可变 NFE,且不需要 JVP 或有限差分计算。其推理头是一个融合线性模块,消除了单独的高斯混合或线性策略头的额外成本和复杂性。在 Wan2.1 和 LTX-2.3 等文本到视频模型上,PDD 仅需 4 到 8 步和最少的训练即可达到顶级视频质量、更高的运动性和更大的多样性,通常匹配或超越基线。PDD 保留了样本多样性并紧密遵循教师分布,而 DMD 导致模式坍塌和多样性降低。
PDD 在极少的推理步骤下实现了强大的文本到图像性能。欧拉变体在 2 次函数评估时达到 DPG-Bench 评分 88.04 和 GenEval 0.86,而中点变体进一步将 OneIG-EN 提升至 0.516、DPG-Bench 提升至 88.10。与 DMD2 等蒸馏基线相比,PDD 保持了有竞争力的质量,同时避免了一步方法中观察到的多样性损失。PDD-Euler(2 NFE)在 DPG-Bench 上达到 88.04,在 GenEval 上达到 0.86,优于 TwinFlow 基线并接近教师的 88.30。中点变体将 OneIG-EN 从 0.508 提高到 0.516,DPG-Bench 提高到 88.10,与教师的 GenEval 0.86 匹配。Pi-Flow 以 4 NFE 达到 0.533 OneIG-EN 和 88.11 DPG-Bench,而 DMD2(lightning-step4-v2)得分为 0.524 和 88.25,显示 PDD 仅用一半评估次数即可达到相当的质量。论文指出 DMD2 遭受多样性坍塌,而 PDD 在保持有竞争力的同时更好地保留了生成多样性和教师保真度。
针对 Qwen-Image 的 PDD 蒸馏比 DMD2 更好地保留了多样性,同时质量保持竞争力。在 2 NFE 下,PDD 变体几乎匹配教师的多样性,而 DMD2 取得了最高的质量分数,但遭受严重的模式坍塌,多样性非常低。PDD-Euler 和 PDD-Midpoint 在 2 NFE 时达到 OneIG 多样性 0.197 和 0.177,接近教师的 0.200。DMD2 在 4 NFE 时达到最高的 HPSv2(32.34)和 PickScore(22.98),但多样性坍塌至 0.095。PDD-Midpoint 的 PickScore 为 22.66,与最好的少步方法相当,同时提供比 DMD2 和 TwinFlow 显著更高的多样性。
在 ImageNet-256 上使用 SiT-XL+REPA 教师的单次 NFE 下,FreeFlow 取得了最低的 FID(1.45)。使用欧拉和中点离散化的 PDD 模型分别达到 FID 分数 2.73 和 2.69,均优于 Pi-Flow 的 2.85。中点变体比欧拉版本略有改进。FreeFlow 以 1.45 的 FID 领先,而两个 PDD 变体的分数均高于 2.6。PDD-Midpoint(2.69)略优于 PDD-Euler(2.73)。在此单步设置下,两个 PDD 模型均超越 Pi-Flow(2.85)。
在 Wan2.1 1.3B 模型上,采用中点离散化的 PDD 获得了最高的 VBench 总体和质评分数,同时比 AnyFlow 和 DMD2 保留了更多多样性,接近教师的多样性。DMD2 取得了最高的质评分数,但语义分数最低且多样性明显下降,与模式坍塌一致。对于 14B 模型,经短时间训练的 PDD 检查点给出了最佳视频质量,并在总体指标上排名第二,多样性同样高于基线。PDD-Midpoint(4 NFE)在 1.3B 模型上获得了最佳的总体 VBench 分数(84.94)和质评分数(86.45),超越所有基线。PDD 变体比 AnyFlow 和 DMD2 保留了更高的多样性:PDD-Midpoint 的 V-JEPA 2 余弦距离为 0.1032,而 AnyFlow 为 0.0704,DMD2 为 0.0833。DMD2(4 NFE)取得了最高的质评分数(86.14),但语义分数最低(78.87)且多样性降低,体现了模式坍塌。教师模型(UniPC, 50x2 NFE)显示出最高的多样性(V-JEPA 2 余弦 0.1254),PDD 比其他少步方法更接近这一教师多样性。在 14B 模型上,PDD(short)检查点提供了最佳视频质量,并在总体 VBench 中排名第二,同时仍保持比基线更高的多样性分数。
实验评估了 PDD 在文本到视频(Wan2.1、LTX-2.3)和文本到图像(多种教师)任务上的蒸馏,验证了其通过避免雅可比向量积并使用融合线性头简化训练的能力。在各种设置下,PDD 以少步推理(2-4 NFE)实现了有竞争力或更优的质量和运动,同时始终保留接近教师的样本多样性,并防止了 DMD2 中观察到的模式坍塌。它匹配或超越了 Pi-Flow 和 AnyFlow 等基线,并且即使在最少的训练迭代下也表现良好。