英伟达提出TMD框架，通过解耦主干与流头结构及两阶段蒸馏训练，将Wan2.1等大视频模型压缩为少步生成器。实验显示，该方法在速度与质量间实现最佳权衡，如14B模型近一步生成时VBench达84.24分，显著优于现有方案。
👇扫码免费加入AI知识星球，如您有工作需要分享，欢迎联系：aigc_to_future
作者：Weili Nie等
论文链接：https://arxiv.org/pdf/2601.09881
项目链接：https://research.nvidia.com/labs/genair/tmd/
亮点直击
一种新颖的视频扩散蒸馏框架： 转移匹配蒸馏（TMD），它将长去噪轨迹蒸馏成紧凑的少步概率转移过程。 解耦的扩散主干设计： 将教师模型分解为语义主干和循环流头，实现了具有灵活内部流细化的分层蒸馏。 两阶段训练策略： （1）转移匹配适应，将流头转换为条件流映射；（2）分布匹配蒸馏，在每个转移步骤中进行流头展开。 全面的实证验证： 证明了TMD在蒸馏Wan2.1 1.3B和14B T2V模型方面的有效性，在少步视频生成中实现了速度和质量之间的最先进权衡。
效果对比
总结速览
解决的问题
大型视频扩散模型和流模型尽管在生成高质量视频方面取得了显著成功，但其多步采样过程效率低下，导致高推理延迟和计算成本。这使得它们难以应用于实时交互式场景，例如实时视频生成、内容编辑或用于代理训练的世界建模。现有方法在扩展到视频扩散模型时面临挑战，因为视频具有高时空维度和复杂的帧间依赖性，难以在蒸馏过程中同时保留全局运动连贯性和精细空间细节。
提出的方案
本文提出了一种新颖的框架——转移匹配蒸馏（Transition Matching Distillation, TMD），旨在将大型视频扩散模型蒸馏成高效的少步生成器。该方法的核心思想是将扩散模型漫长的多步去噪轨迹与一个紧凑的少步概率转移过程进行匹配。
应用的技术
解耦架构： 本文将原始的扩散主干网络分解为两个主要组件：一个主干网络（包含大部分早期层），负责在每个外部转移步骤中提取高级语义表示；以及一个轻量级流头（由最后几层组成），利用这些语义表示执行多次内部流更新以细化视觉细节。 两阶段训练策略：1. 转移匹配预训练： 在蒸馏之前，通过利用Mean-Flow的改编版本，对流头进行预训练，使其适应为一个能够迭代细化特征的条件流映射。 2. 分布匹配蒸馏： 采用改进版的DMD2，将蒸馏问题表述为教师去噪过程与学生转移过程之间的分布匹配问题。在训练过程中，流头会进行展开（rollout），以确保学生模型的概率转移与教师的多步扩散分布对齐，从而捕获语义演变和精细视觉细节。
达到的效果
在蒸馏Wan2.1 1.3B和14B文本到视频模型上的广泛实验表明，TMD在生成速度和视觉质量之间提供了灵活且强大的权衡。具体而言，在可比较的推理成本下，TMD在视觉保真度和提示遵循度方面始终优于现有的蒸馏模型。例如，蒸馏后的14B模型在接近一步生成（NFE=1.38）的情况下，在VBench上取得了84.24的总分，显著优于所有其他一步蒸馏方法。
方法
TMD方法，包括两个训练阶段：（1）转移匹配预训练，以初始化一个能够迭代优化从主干网络中提取的特征的流头；（2）带流头蒸馏，本文引入DMD2-v，它在少步视频生成中改进了DMD2，并在每个转移步骤中应用流头展开。为简化表示，本文在符号中省略了教师模型的额外条件，例如文本条件。下面，本文首先介绍学生架构，然后介绍两阶段训练。
解耦架构。 本文的方法遵循前文中解释的转移匹配的一般公式。与TM不同，本文旨在用蒸馏学生的大转移步长来近似教师模型的许多小去噪步长。为了在每个转移步长 中高效预测 ，本文将预训练的教师架构解耦为一个主干网络 ，它充当特征提取器，以及一个轻量级流头 ，它给定这些特征迭代预测 ，即：
其中 是内流的给定时间离散化；如下图2所示。
虽然这种解耦已成功用于训练生成模型 [61, 44, 50, 27, 67]，但它需要仔细设计以尽量不干扰预训练模型。本文的设计考虑了两个关键因素：
流头目标 。 本文发现DTM公式 优于其他目标类型，例如样本预测 （参见附录B）。 融合层。 本文使用时间条件门控机制来融合主要特征 和噪声流头目标 ，确保学生模型的初始前向传播与教师模型匹配。此外，本文将主输入 的补丁嵌入重用于内流输入 。本文在算法1中提供了推理的伪代码。
阶段1：转移匹配预训练
基于解耦架构，在蒸馏前我们将流头部转换为流映射以进行迭代优化。与转换匹配（TM）方法类似，可直接采用公式（6）中的流匹配损失来训练流头部，使其逼近内部流的速度场。然而理论上，这仍需大量内部步骤才能逼近目标𝑦。因此，我们借助MeanFlow方法实现少步数的流头部训练。
转换匹配均值流方法
从整体来看，我们提出的预训练算法——转换匹配均值流（TM-MF）——采用公式（9）所示的MeanFlow目标函数，并以主干特征𝑚 = 𝑚𝜃(𝑥𝑡𝑖, 𝑡𝑖)为条件（算法2给出了伪代码）。具体而言，我们通过平均速度参数化条件化内部流映射：需注意的是，在训练过程中我们未对主干网络特征进行梯度截断，因为这可能限制预训练所需的灵活性。若直接训练流头部预测平均速度𝑢𝜃，效果并不理想。我们的假设是：流头部的输出应尽可能接近预训练教师的输出。由于教师网络预测的是公式（2）中外部流的速度，为保持与教师网络的一致性，流头部应转而预测E[𝑥1 − 𝑥 | 𝑥]。根据公式（5）中内部速度的定义，我们得出：
因此，我们将平均速度参数化表示为：
其中headθ表示我们解耦架构中的头部模块（根据附录A所述从教师网络初始化而来）。通过这种参数化方式，当𝑟趋近于𝑠时，headθ的输出将逼近教师网络的速度预测值。
为提升性能与稳定性，我们遵循原始MeanFlow方法的三项实践：（1）对部分批次数据执行流匹配（在我们的设定中更准确地说，是转换匹配）；（2）采用无分类器引导技术（通过调整条件速度𝑣(𝑦𝑠, 𝑠)实现），以特定概率丢弃文本条件；（3）应用自适应损失归一化。由于公式（10）中的雅可比向量积计算需要定制化实现以兼容视频生成的大规模训练（例如使用Flash Attention[12]、全分片数据并行[66]或上下文并行[23]等技术），我们采用有限差分法近似计算雅可比向量积，使算法能够独立于底层架构和训练技术（参见文献[47,52]）。
由于无法直接获取内部流速度，我们在目标函数（9）中使用条件速度𝑣(𝑦𝑠, 𝑠) = 𝑦1 − 𝑦。需要说明的是，针对特定𝑦值，理论上也可基于预训练教师网络的速度推导出内部速度的表征形式（如文献[20]所述），这部分工作留待未来研究。最后，我们注意到转换匹配作为预训练策略也能取得具有竞争力的结果（消融实验后文）。特别地，当使用条件速度时，转换匹配预训练可视为公式（9）中MeanFlow在𝑟 = 𝑠情况下的特例。
第二阶段：基于流头部的蒸馏训练
完成TM-MF预训练后，我们应用分布蒸馏技术来对齐学生模型与教师模型的分布。我们对视频模型的基线DMD2方法进行了显著改进，并针对TMD方法优化了实现方案。
DMD2-v改进方案
DMD2最初是为图像扩散模型蒸馏设计的，其设计选择在视频领域可能并非最优。我们发现了三个提升视频DMD2性能的关键因素（称为DMD2-v），这些因素构成TMD训练的默认设置：
GAN判别器架构：使用Conv3D层的GAN判别器表现优于其他架构，表明局部时空特征对GAN损失至关重要； 知识蒸馏预热策略：知识蒸馏预热能提升单步蒸馏性能，但在多步生成中容易引入难以通过DMD2训练修复的粗粒度伪影（见附录图10），因此DMD2-v仅对单步蒸馏应用该策略； 时间步偏移技术：在为外部转移步采样时间步，或在VSD损失中为生成样本添加噪声时，对均匀采样的𝑡′应用偏移函数𝑡=𝛾𝑡′/((𝛾-1)𝑡′+1)（𝛾≥1）可提升性能并防止模式坍塌。
流头部展开机制
在蒸馏过程中，我们展开内部流并将所得架构视为每个转移步𝑡𝑖的样本生成器（见上图2b）。根据公式（3）定义的流头部目标，展开后的学生模型输出为：其中𝑦̂0 ≈ InnerFlow(𝑚𝜃(𝑥𝑡𝑖, 𝑡𝑖))表示流头部经过𝑁步内部优化后的最终预测，遵循公式（12）的计算过程。
实验
实验设置
实现。 本文使用Wan2.1 1.3B和14B T2V-480p作为教师视频扩散模型，并将它们蒸馏成具有解耦架构的相同大小的学生模型。所有实验都在 [𝑇, 𝐻, 𝑊] = [21, 60, 104] 的潜在分辨率下进行，解码为81帧，像素分辨率为480 × 832。本文使用包含50万个文本和视频对的数据集，其中文本提示从VidProM数据集采样（并由Qwen-2.5 扩展），视频由Wan2.1 14B T2V模型生成。
评估指标。 为了评估本文的方法和基线，本文使用VBench [22]（报告总分、质量分和语义分）和用户偏好研究来评估视觉质量和提示遵循度。本文认为有效函数评估次数（NFE）是生成过程中使用的DiT块总数除以 𝐿（教师架构中的块数）；对于基线，这对应于步数 𝑀，对于本文的TMD模型，这对应于：
其中 是内流步数， 是流头中的块数。值得注意的是，Wan2.1 1.3B的 ，Wan2.1 14B的 。
与现有方法的比较
本文的TMD方法是基于改进版DMD2的视频生成版本（称为DMD2-v）。本文将TMD与DMD2-v和现有基线进行比较，以蒸馏视频扩散模型。如下图3所示，本文提供了视觉比较。如下表1所示，本文展示了蒸馏Wan2.1 1.3B（或类似大小的视频模型）成少步生成器时的VBench结果，其中本文根据学生去噪步数 𝑀 对蒸馏模型进行分组。当 时，TMD-N2H5（有效NFE = 2.33，即2个去噪步长和流头中的5个DiT块）改进了蒸馏性能。如下表2所示，本文展示了蒸馏Wan2.1 14B成少步生成器时的VBench结果。DMD2-v表示本文改进的视频生成DMD2版本。
判别器头。 如下表3所示，本文展示了DMD2-v中判别器头设计的影响，用于Wan2.1 1.3B的一步蒸馏。本文比较了三个头部：（1）Conv3D，联合处理时空特征；（2）Conv1D-2D，分离时间卷积和空间卷积（例如文献[64]）；（3）Attention，将特征展平为由自注意力处理的令牌（带有池化下采样）。Conv3D优于其他两种判别器头架构。
KD预热。 如下表4所示，VBench上的总分随着一步DMD2中的KD预热而增加，但在两步DMD2中随着KD预热而减少。这表明本文最好只在一步生成中应用KD预热。
时间步长偏移。 如下表5所示，本文观察到，对控制DMD损失中噪声水平的 和控制多步学生中去噪步数的 应用时间步长偏移分别改进了蒸馏性能。
质量-效率权衡。 内部步数 和流头层数 控制着内部流的计算成本。本文改变 和 以更全面地分析TMD的性能-效率权衡。如下图6所示，本文观察到VBench总分通常随着有效NFE的增加而提高。这证明了本文的方法在平衡生成速度和视觉质量方面提供的精细灵活性。
MeanFlow vs. 流匹配。 在转移匹配预训练中，本文用香草流匹配目标（TM）替换MeanFlow目标（TM-MF），以突出MeanFlow的影响。如下表6所示，TM-MF始终比TM实现更好的蒸馏性能，这表明TM-MF为第二阶段蒸馏训练提供了优越的初始化。
蒸馏中的流头展开。 关闭训练和推理之间的差距至关重要，方法是允许蒸馏目标中的梯度通过展开的内部流轨迹反向传播。如下图7所示，在蒸馏中应用流头展开极大地加速了训练收敛并提高了性能。
结论
转移匹配蒸馏（TMD），一种旨在解决大规模视频扩散模型显著推理延迟的新颖框架。本文方法的核心在于解耦的学生架构，它将用于语义特征提取的主干网络与用于迭代细化的轻量级循环流头分离。这种设计与两阶段训练策略相结合，包括转移匹配预训练和基于分布的蒸馏。本文在蒸馏最先进的Wan2.1模型上的实验表明，TMD在各种推理预算下提供了精细的灵活性，在视频质量和提示遵循度方面始终优于现有方法，从而实现视频生成速度和质量之间的更好权衡。
参考文献
[1] Transition Matching Distillation for Fast Video Generation
技术交流社区免费开放
涉及 内容生成/理解（图像、视频、语音、文本、3D/4D等）、大模型、具身智能、自动驾驶、深度学习及传统视觉等多个不同方向。这个社群更加适合记录和积累，方便回溯和复盘。愿景是联结数十万AIGC开发者、研究者和爱好者，解决从理论到实战中遇到的具体问题。倡导深度讨论，确保每个提问都能得到认真对待。
技术交流
加入「AI生成未来社区」群聊，一起交流讨论，涉及 图像生成、视频生成、3D生成、具身智能等多个不同方向，备注不同方向邀请入群！可添加小助手备注方向加群！
没有评论:
发表评论