訂閱
糾錯
加入自媒體

視頻生成效率革命!英偉達TMD蒸餾框架:Wan2.1模型實時化,速度和質量最佳均衡!

2026-01-19 15:31
AI生成未來
關注

作者:Weili Nie等

解讀:AI生成未來

亮點直擊

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

解耦的擴散主干設計: 將教師模型分解為語義主干和循環流頭,實現了具有靈活內部流細化的分層蒸餾。

兩階段訓練策略: (1)轉移匹配適應,將流頭轉換為條件流映射;(2)分布匹配蒸餾,在每個轉移步驟中進行流頭展開。

全面的實證驗證: 證明了TMD在蒸餾Wan2.1 1.3B和14B T2V模型方面的有效性,在少步視頻生成中實現了速度和質量之間的最先進權衡。

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

總結速覽

解決的問題

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

提出的方案

本文提出了一種新穎的框架——轉移匹配蒸餾(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。根據公式(5)中內部速度的定義,我們得出:

因此,我們將平均速度參數化表示為:

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

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

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

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

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

DMD2-v改進方案DMD2最初是為圖像擴散模型蒸餾設計的,其設計選擇在視頻領域可能并非最優。我們發現了三個提升視頻DMD2性能的關鍵因素(稱為DMD2-v),這些因素構成TMD訓練的默認設置:

GAN判別器架構:使用Conv3D層的GAN判別器表現優于其他架構,表明局部時空特征對GAN損失至關重要;

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

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

流頭部展開機制在蒸餾過程中,我們展開內部流并將所得架構視為每個轉移步的樣本生成器(見上圖2b)。根據公式(3)定義的流頭部目標,展開后的學生模型輸出為:其中表示流頭部經過步內部優化后的最終預測,遵循公式(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

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

聲明: 本文由入駐維科號的作者撰寫,觀點僅代表作者本人,不代表OFweek立場。如有侵權或其他問題,請聯系舉報。

發表評論

0條評論,0人參與

請輸入評論內容...

請輸入評論/評論長度6~500個字

您提交的評論過于頻繁,請輸入驗證碼繼續

暫無評論

暫無評論

    人工智能 獵頭職位 更多
    掃碼關注公眾號
    OFweek人工智能網
    獲取更多精彩內容
    文章糾錯
    x
    *文字標題:
    *糾錯內容:
    聯系郵箱:
    *驗 證 碼:

    粵公網安備 44030502002758號