作者：Weili Nie等

解讀：AI生成未來

亮點(diǎn)直擊

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

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

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

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

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

總結(jié)速覽

解決的問題

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

提出的方案

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

應(yīng)用的技術(shù)

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

兩階段訓(xùn)練策略：

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

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

達(dá)到的效果

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

方法

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

第二階段：基于流頭部的蒸餾訓(xùn)練

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

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

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

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

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

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

實(shí)驗(yàn)

實(shí)驗(yàn)設(shè)置

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

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

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

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

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

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

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

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

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

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

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

結(jié)論

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

參考文獻(xiàn)

[1] Transition Matching Distillation for Fast Video Generation

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