作者：Weili Nie等

解讀：AI生成未來

亮點直擊

一種新穎的視頻擴散蒸餾框架： 轉(zhuǎn)移匹配蒸餾（TMD），它將長去噪軌跡蒸餾成緊湊的少步概率轉(zhuǎn)移過程。

解耦的擴散主干設計： 將教師模型分解為語義主干和循環(huán)流頭，實現(xiàn)了具有靈活內(nèi)部流細化的分層蒸餾。

兩階段訓練策略： （1）轉(zhuǎn)移匹配適應，將流頭轉(zhuǎn)換為條件流映射；（2）分布匹配蒸餾，在每個轉(zhuǎn)移步驟中進行流頭展開。

全面的實證驗證： 證明了TMD在蒸餾Wan2.1 1.3B和14B T2V模型方面的有效性，在少步視頻生成中實現(xiàn)了速度和質(zhì)量之間的最先進權(quán)衡。

圖1。從 TMD 生成示例。使用我們的 TMD 方法（從 Wan2.1 14B T2V 中提�。┩ㄟ^兩個文本提示生成四幀 5s 480p 視頻，并具有兩個不同（有效）數(shù)量的功能評估 (NFE)效果對比

總結(jié)速覽

解決的問題

大型視頻擴散模型和流模型盡管在生成高質(zhì)量視頻方面取得了顯著成功，但其多步采樣過程效率低下，導致高推理延遲和計算成本。這使得它們難以應用于實時交互式場景，例如實時視頻生成、內(nèi)容編輯或用于代理訓練的世界建�！，F(xiàn)有方法在擴展到視頻擴散模型時面臨挑戰(zhàn)，因為視頻具有高時空維度和復雜的幀間依賴性，難以在蒸餾過程中同時保留全局運動連貫性和精細空間細節(jié)。

提出的方案

本文提出了一種新穎的框架——轉(zhuǎn)移匹配蒸餾（Transition Matching Distillation, TMD），旨在將大型視頻擴散模型蒸餾成高效的少步生成器。該方法的核心思想是將擴散模型漫長的多步去噪軌跡與一個緊湊的少步概率轉(zhuǎn)移過程進行匹配。

應用的技術(shù)

解耦架構(gòu)： 本文將原始的擴散主干網(wǎng)絡分解為兩個主要組件：一個主干網(wǎng)絡（包含大部分早期層），負責在每個外部轉(zhuǎn)移步驟中提取高級語義表示；以及一個輕量級流頭（由最后幾層組成），利用這些語義表示執(zhí)行多次內(nèi)部流更新以細化視覺細節(jié)。

兩階段訓練策略：

1.  轉(zhuǎn)移匹配預訓練： 在蒸餾之前，通過利用Mean-Flow的改編版本，對流頭進行預訓練，使其適應為一個能夠迭代細化特征的條件流映射。

2.  分布匹配蒸餾： 采用改進版的DMD2，將蒸餾問題表述為教師去噪過程與學生轉(zhuǎn)移過程之間的分布匹配問題。在訓練過程中，流頭會進行展開（rollout），以確保學生模型的概率轉(zhuǎn)移與教師的多步擴散分布對齊，從而捕獲語義演變和精細視覺細節(jié)。

達到的效果

在蒸餾Wan2.1 1.3B和14B文本到視頻模型上的廣泛實驗表明，TMD在生成速度和視覺質(zhì)量之間提供了靈活且強大的權(quán)衡。具體而言，在可比較的推理成本下，TMD在視覺保真度和提示遵循度方面始終優(yōu)于現(xiàn)有的蒸餾模型。例如，蒸餾后的14B模型在接近一步生成（NFE=1.38）的情況下，在VBench上取得了84.24的總分，顯著優(yōu)于所有其他一步蒸餾方法。

方法

TMD方法，包括兩個訓練階段：（1）轉(zhuǎn)移匹配預訓練，以初始化一個能夠迭代優(yōu)化從主干網(wǎng)絡中提取的特征的流頭；（2）帶流頭蒸餾，本文引入DMD2-v，它在少步視頻生成中改進了DMD2，并在每個轉(zhuǎn)移步驟中應用流頭展開。為簡化表示，本文在符號中省略了教師模型的額外條件，例如文本條件。下面，本文首先介紹學生架構(gòu)，然后介紹兩階段訓練。

解耦架構(gòu)。 本文的方法遵循前文中解釋的轉(zhuǎn)移匹配的一般公式。與TM不同，本文旨在用蒸餾學生的大轉(zhuǎn)移步長來近似教師模型的許多小去噪步長。為了在每個轉(zhuǎn)移步長  中高效預測 ，本文將預訓練的教師架構(gòu)解耦為一個主干網(wǎng)絡 ，它充當特征提取器，以及一個輕量級流頭 ，它給定這些特征迭代預測 ，即：

其中  是內(nèi)流的給定時間離散化；如下圖2所示。

雖然這種解耦已成功用于訓練生成模型 [61, 44, 50, 27, 67]，但它需要仔細設計以盡量不干擾預訓練模型。本文的設計考慮了兩個關(guān)鍵因素：

流頭目標 。 本文發(fā)現(xiàn)DTM公式  優(yōu)于其他目標類型，例如樣本預測 （參見附錄B）。

融合層。 本文使用時間條件門控機制來融合主要特征  和噪聲流頭目標 ，確保學生模型的初始前向傳播與教師模型匹配。此外，本文將主輸入  的補丁嵌入重用于內(nèi)流輸入 。

本文在算法1中提供了推理的偽代碼。

階段1：轉(zhuǎn)移匹配預訓練

基于解耦架構(gòu)，在蒸餾前我們將流頭部轉(zhuǎn)換為流映射以進行迭代優(yōu)化。與轉(zhuǎn)換匹配（TM）方法類似，可直接采用公式（6）中的流匹配損失來訓練流頭部，使其逼近內(nèi)部流的速度場。然而理論上，這仍需大量內(nèi)部步驟才能逼近目標。因此，我們借助MeanFlow方法實現(xiàn)少步數(shù)的流頭部訓練。

轉(zhuǎn)換匹配均值流方法從整體來看，我們提出的預訓練算法——轉(zhuǎn)換匹配均值流（TM-MF）——采用公式（9）所示的MeanFlow目標函數(shù)，并以主干特征 = 為條件（算法2給出了偽代碼）。具體而言，我們通過平均速度參數(shù)化條件化內(nèi)部流映射：

需注意的是，在訓練過程中我們未對主干網(wǎng)絡特征進行梯度截斷，因為這可能限制預訓練所需的靈活性。若直接訓練流頭部預測平均速度，效果并不理想。我們的假設是：流頭部的輸出應盡可能接近預訓練教師的輸出。由于教師網(wǎng)絡預測的是公式（2）中外部流的速度，為保持與教師網(wǎng)絡的一致性，流頭部應轉(zhuǎn)而預測E。根據(jù)公式（5）中內(nèi)部速度的定義，我們得出：

因此，我們將平均速度參數(shù)化表示為：

其中headθ表示我們解耦架構(gòu)中的頭部模塊（根據(jù)附錄A所述從教師網(wǎng)絡初始化而來）。通過這種參數(shù)化方式，當趨近于時，headθ的輸出將逼近教師網(wǎng)絡的速度預測值。

為提升性能與穩(wěn)定性，我們遵循原始MeanFlow方法的三項實踐：（1）對部分批次數(shù)據(jù)執(zhí)行流匹配（在我們的設定中更準確地說，是轉(zhuǎn)換匹配）；（2）采用無分類器引導技術(shù)（通過調(diào)整條件速度實現(xiàn)），以特定概率丟棄文本條件；（3）應用自適應損失歸一化。由于公式（10）中的雅可比向量積計算需要定制化實現(xiàn)以兼容視頻生成的大規(guī)模訓練（例如使用Flash Attention[12]、全分片數(shù)據(jù)并行[66]或上下文并行[23]等技術(shù)），我們采用有限差分法近似計算雅可比向量積，使算法能夠獨立于底層架構(gòu)和訓練技術(shù)（參見文獻[47,52]）。

由于無法直接獲取內(nèi)部流速度，我們在目標函數(shù)（9）中使用條件速度。需要說明的是，針對特定值，理論上也可基于預訓練教師網(wǎng)絡的速度推導出內(nèi)部速度的表征形式（如文獻[20]所述），這部分工作留待未來研究。最后，我們注意到轉(zhuǎn)換匹配作為預訓練策略也能取得具有競爭力的結(jié)果（消融實驗后文）。特別地，當使用條件速度時，轉(zhuǎn)換匹配預訓練可視為公式（9）中MeanFlow在情況下的特例。

第二階段：基于流頭部的蒸餾訓練

完成TM-MF預訓練后，我們應用分布蒸餾技術(shù)來對齊學生模型與教師模型的分布。我們對視頻模型的基線DMD2方法進行了顯著改進，并針對TMD方法優(yōu)化了實現(xiàn)方案。

DMD2-v改進方案DMD2最初是為圖像擴散模型蒸餾設計的，其設計選擇在視頻領(lǐng)域可能并非最優(yōu)。我們發(fā)現(xiàn)了三個提升視頻DMD2性能的關(guān)鍵因素（稱為DMD2-v），這些因素構(gòu)成TMD訓練的默認設置：

GAN判別器架構(gòu)：使用Conv3D層的GAN判別器表現(xiàn)優(yōu)于其他架構(gòu)，表明局部時空特征對GAN損失至關(guān)重要；

知識蒸餾預熱策略：知識蒸餾預熱能提升單步蒸餾性能，但在多步生成中容易引入難以通過DMD2訓練修復的粗粒度偽影（見附錄圖10），因此DMD2-v僅對單步蒸餾應用該策略；

時間步偏移技術(shù)：在為外部轉(zhuǎn)移步采樣時間步，或在VSD損失中為生成樣本添加噪聲時，對均勻采樣的應用偏移函數(shù)可提升性能并防止模式坍塌。

流頭部展開機制在蒸餾過程中，我們展開內(nèi)部流并將所得架構(gòu)視為每個轉(zhuǎn)移步的樣本生成器（見上圖2b）。根據(jù)公式（3）定義的流頭部目標，展開后的學生模型輸出為：其中表示流頭部經(jīng)過步內(nèi)部優(yōu)化后的最終預測，遵循公式（12）的計算過程。

實驗

實驗設置

實現(xiàn)。 本文使用Wan2.1 1.3B和14B T2V-480p作為教師視頻擴散模型，并將它們蒸餾成具有解耦架構(gòu)的相同大小的學生模型。所有實驗都在 = [21, 60, 104] 的潛在分辨率下進行，解碼為81幀，像素分辨率為480 × 832。本文使用包含50萬個文本和視頻對的數(shù)據(jù)集，其中文本提示從VidProM數(shù)據(jù)集采樣（并由Qwen-2.5 擴展），視頻由Wan2.1 14B T2V模型生成。

評估指標。 為了評估本文的方法和基線，本文使用VBench [22]（報告總分、質(zhì)量分和語義分）和用戶偏好研究來評估視覺質(zhì)量和提示遵循度。本文認為有效函數(shù)評估次數(shù)（NFE）是生成過程中使用的DiT塊總數(shù)除以 （教師架構(gòu)中的塊數(shù)）；對于基線，這對應于步數(shù) ，對于本文的TMD模型，這對應于：

其中  是內(nèi)流步數(shù)， 是流頭中的塊數(shù)。值得注意的是，Wan2.1 1.3B的 ，Wan2.1 14B的 。

與現(xiàn)有方法的比較

本文的TMD方法是基于改進版DMD2的視頻生成版本（稱為DMD2-v）。本文將TMD與DMD2-v和現(xiàn)有基線進行比較，以蒸餾視頻擴散模型。如下圖3所示，本文提供了視覺比較。如下表1所示，本文展示了蒸餾Wan2.1 1.3B（或類似大小的視頻模型）成少步生成器時的VBench結(jié)果，其中本文根據(jù)學生去噪步數(shù)  對蒸餾模型進行分組。當  時，TMD-N2H5（有效NFE = 2.33，即2個去噪步長和流頭中的5個DiT塊）改進了蒸餾性能。如下表2所示，本文展示了蒸餾Wan2.1 14B成少步生成器時的VBench結(jié)果。DMD2-v表示本文改進的視頻生成DMD2版本。

判別器頭。 如下表3所示，本文展示了DMD2-v中判別器頭設計的影響，用于Wan2.1 1.3B的一步蒸餾。本文比較了三個頭部：（1）Conv3D，聯(lián)合處理時空特征；（2）Conv1D-2D，分離時間卷積和空間卷積（例如文獻[64]）；（3）Attention，將特征展平為由自注意力處理的令牌（帶有池化下采樣）。Conv3D優(yōu)于其他兩種判別器頭架構(gòu)。

KD預熱。 如下表4所示，VBench上的總分隨著一步DMD2中的KD預熱而增加，但在兩步DMD2中隨著KD預熱而減少。這表明本文最好只在一步生成中應用KD預熱。

時間步長偏移。 如下表5所示，本文觀察到，對控制DMD損失中噪聲水平的  和控制多步學生中去噪步數(shù)的  應用時間步長偏移分別改進了蒸餾性能。

質(zhì)量-效率權(quán)衡。 內(nèi)部步數(shù)  和流頭層數(shù)  控制著內(nèi)部流的計算成本。本文改變  和  以更全面地分析TMD的性能-效率權(quán)衡。如下圖6所示，本文觀察到VBench總分通常隨著有效NFE的增加而提高。這證明了本文的方法在平衡生成速度和視覺質(zhì)量方面提供的精細靈活性。

MeanFlow vs. 流匹配。 在轉(zhuǎn)移匹配預訓練中，本文用香草流匹配目標（TM）替換MeanFlow目標（TM-MF），以突出MeanFlow的影響。如下表6所示，TM-MF始終比TM實現(xiàn)更好的蒸餾性能，這表明TM-MF為第二階段蒸餾訓練提供了優(yōu)越的初始化。

蒸餾中的流頭展開。 關(guān)閉訓練和推理之間的差距至關(guān)重要，方法是允許蒸餾目標中的梯度通過展開的內(nèi)部流軌跡反向傳播。如下圖7所示，在蒸餾中應用流頭展開極大地加速了訓練收斂并提高了性能。

結(jié)論

轉(zhuǎn)移匹配蒸餾（TMD），一種旨在解決大規(guī)模視頻擴散模型顯著推理延遲的新穎框架。本文方法的核心在于解耦的學生架構(gòu)，它將用于語義特征提取的主干網(wǎng)絡與用于迭代細化的輕量級循環(huán)流頭分離。這種設計與兩階段訓練策略相結(jié)合，包括轉(zhuǎn)移匹配預訓練和基于分布的蒸餾。本文在蒸餾最先進的Wan2.1模型上的實驗表明，TMD在各種推理預算下提供了精細的靈活性，在視頻質(zhì)量和提示遵循度方面始終優(yōu)于現(xiàn)有方法，從而實現(xiàn)視頻生成速度和質(zhì)量之間的更好權(quán)衡。

參考文獻

[1] Transition Matching Distillation for Fast Video Generation

       原文標題 : 視頻生成效率革命！英偉達TMD蒸餾框架：Wan2.1模型實時化，速度和質(zhì)量最佳均衡！