視頻生成太慢?英偉達、謝賽寧等發布TMD框架,實現70倍加速

視頻生成太慢?英偉達、謝賽寧等發布TMD框架,實現70倍加速

文章圖片

視頻生成太慢?英偉達、謝賽寧等發布TMD框架,實現70倍加速

文章圖片

視頻生成太慢?英偉達、謝賽寧等發布TMD框架,實現70倍加速

文章圖片

視頻生成太慢?英偉達、謝賽寧等發布TMD框架,實現70倍加速

文章圖片

視頻生成太慢?英偉達、謝賽寧等發布TMD框架,實現70倍加速

文章圖片


機器之心編輯部
近年來 , 大規模視頻擴散模型在視頻生成領域取得了顯著進展 。 然而 , 采樣效率低下仍然是這類模型的核心瓶頸 。
標準的擴散模型依賴多步去噪過程 , 通常需要數十甚至上百次迭代步驟 , 才能將噪聲逐步轉化為真實輸出 。 這種迭代特性帶來了高推理延遲和高計算成本 , 使得大型擴散模型難以應用于實時視頻生成、內容編輯、或為智能體訓練構建世界模型等交互式場景 。 因此 , 如何在不犧牲視覺質量的前提下加速擴散采樣 , 成為一個關鍵的開放性挑戰 。
現有的擴散蒸餾方法主要分為兩大類:一類是基于軌跡的蒸餾 , 包括知識蒸餾和一致性模型 , 它們直接回歸教師模型的去噪軌跡;另一類是基于分布的蒸餾 , 包括對抗蒸餾和變分分數蒸餾方法 , 它們對齊學生和教師的輸出分布 。 這些技術在圖像領域已經可以將采樣過程壓縮到僅需一到兩步 。
然而 , 將這些方法擴展到視頻擴散模型面臨獨特的挑戰 。 視頻具有高時空維度和復雜的幀間依賴關系 , 在蒸餾過程中很難同時保持全局運動連貫性和細粒度空間細節 。 而且 , 現有大多數方法將擴散網絡視為一個單一的整體映射 , 忽略了大型視頻擴散主干網絡中固有的層級結構和語義演進特性 。
針對上述挑戰 , NVIDIA 聯合 NYU 的謝賽寧等研究者提出了「轉移匹配蒸餾」(Transition Matching Distillation , 簡稱 TMD)框架 , 用于將大型視頻擴散模型蒸餾為高效的少步生成器 。

論文標題:Transition Matching Distillation for Fast Video Generation 論文鏈接:https://arxiv.org/pdf/2601.09881v1 項目鏈接:https://research.nvidia.com/labs/genair/tmd/TMD 的核心思想是:用一個緊湊的少步概率轉移過程來近似擴散模型的多步去噪過程 。 在這個框架中 , 每一個轉移步驟都捕捉了視頻樣本在相距較遠的噪聲水平之間的分布演變 , 從而使學生模型能夠以大的轉移步長來匹配教師模型的分布 。
實驗結果表明 , 在相當的推理預算下 , TMD 始終優于現有的蒸餾方法 , 能夠實現更好的視覺保真度和對提示詞的遵循度 。 作者蒸餾出的 14B 模型在 VBench 上的綜合得分為 84.24 , 接近一步生成(NFE=1.38) 。

下圖對比了教師模型(Wan2.1 14B)和 TMD 蒸餾模型在函數評估次數(NFE)上的差異 。 原始的 Wan2.1 14B 教師模型需要 100 次 NFE 才能生成一個視頻 , 在 VBench 上的得分是 86.2 。 而經過 TMD 蒸餾后 , 模型只需要 2.75 次 NFE 就能達到 84.6 的 VBench 得分 , 或者只需要 1.38 次 NFE 就能達到 84.2 的得分 。 換算下來 , TMD 實現了超過 70 倍的加速 , 而 VBench 得分僅下降約 2 分(從 86.2 降到 84.2) 。 這意味著在質量損失很小的情況下 , 視頻生成速度提升了一個數量級以上 。

【視頻生成太慢?英偉達、謝賽寧等發布TMD框架,實現70倍加速】在該項目網站 , 我們可以看到視頻生成結果以及與基線方法的效果對比 。
方法詳解
為了實現高效蒸餾 , 研究團隊提出了一種解耦的學生模型架構 , 將原始的擴散主干網絡分解為兩個組件:第一個是「主干網絡」 , 包含大部分早期層 , 負責在每個外部轉移步驟提取高層語義表征;第二個是 Flow Head , 由最后幾層組成 , 它以主干網絡提取的表征為條件 , 通過多個內部流更新來細化細粒度的視覺細節 。
這種層級化的分解設計帶來了顯著優勢 。 主干網絡可以與 Flow Head 共享表征 , 而 Flow Head 則在每個外部轉移步驟內執行若干輕量級的內部細化步驟 , 為平衡采樣效率和視覺保真度提供了靈活的機制 。 舉例來說 , 如果從一個 30 層的 DiT 模型中選取最后 5 層作為 Flow Head , 并展開 2 個內部步驟 , 額外計算開銷不到 17% 。
TMD 采用兩階段訓練策略 。 第一階段是「轉移匹配預訓練」 , 研究團隊借鑒了 MeanFlow 的思想 , 將 Flow Head 訓練成一個條件流映射 , 使其能夠進行迭代細化 。 這一階段的關鍵在于讓 Flow Head 的輸出保持與預訓練教師模型的輸出對齊 。 第二階段是「帶 Flow Head 的分布匹配蒸餾」 , 研究團隊改進了 DMD2 方法使其適應視頻生成場景(稱為 DMD2-v) , 并在每個轉移步驟中展開 Flow Head 進行訓練 。 通過展開 Flow Head , 學生模型的概率轉移與教師模型的多步擴散分布得以對齊 , 同時捕捉語義演變和細粒度視覺細節 。

在 DMD2-v 的改進中 , 研究團隊識別出三個對視頻蒸餾至關重要的因素:首先是 GAN 判別器架構 , 使用 Conv3D 層優于其他架構 , 這表明局部化的時空特征對 GAN 損失很重要;其次是知識蒸餾預熱策略 , 它在單步蒸餾中有幫助 , 但在多步生成中反而會引入難以修復的粗粒度偽影;第三是時間步偏移 , 對采樣時間步應用偏移函數可以改善性能并防止模式崩潰 。
實驗結果
研究團隊在 Wan2.1 的 1.3B 和 14B 兩個文本到視頻模型上驗證了 TMD 的效果 , 生成的是 5 秒、480p 分辨率、81 幀的視頻 。 評估采用了 VBench 基準測試和用戶偏好研究兩種方式 。
在 Wan2.1 1.3B 模型的蒸餾實驗中 , 當使用 2 個外部去噪步驟時 , TMD-N2H5(即 2 個去噪步驟配合 5 層 Flow Head)以 2.33 的有效 NFE(函數評估次數)取得了 84.68 的 VBench 總分 , 超越了所有其他蒸餾模型 , 包括 NFE 為 4 的最強基線 rCM(總分 84.43) 。 當只使用 1 個外部步驟時 , TMD-N2H5 以 1.17 的有效 NFE 取得了 83.80 的總分 , 同樣優于所有其他單步蒸餾方法 , 大幅縮小了與兩步蒸餾方法之間的差距 。

在更大的 Wan2.1 14B 模型上 , TMD 的優勢更加明顯 。 當使用 1 個外部步驟時 , TMD-N4H5 以 1.38 的有效 NFE 取得了 84.24 的總分 , 顯著超越了所有其他單步蒸餾方法 。 相比單步 rCM 的 83.02 分 , TMD 提升了 1.22 分 , 而額外的推理成本幾乎可以忽略不計 。

此外 , TMD 還消除了單步 DMD2-v 所需的計算昂貴的知識蒸餾預熱步驟 。
除了自動化指標 , 研究團隊還進行了用戶偏好研究 。 他們從 VBench 中隨機抽取 60 個具有挑戰性的提示 , 針對每個提示用不同種子生成 5 個視頻 , 讓評估者在視覺質量和提示對齊兩個維度上進行盲測對比 。 結果顯示 , 無論是在單步還是兩步生成設置下 , 用戶都一致更偏好 TMD 生成的視頻 。 在兩步生成中 , TMD 在視覺質量上的勝率為 63.3% , 在提示對齊上的勝率高達 71.9% 。 這一結果表明 , Flow Head 的迭代細化機制對于提升提示遵循能力有顯著幫助 。


研究團隊還進行了詳細的消融實驗來驗證各設計選擇的有效性 。 在預訓練方法上 , 使用 MeanFlow 目標(TM-MF)一致優于普通的流匹配目標(TM) 。 在蒸餾階段 , 啟用 Flow Head 展開可以帶來更快的訓練收斂和更好的最終性能 。 此外 , 通過調整內部步數和 Flow Head 層數 , TMD 可以實現細粒度的質量 - 效率權衡 , VBench 總分隨有效 NFE 的增加而穩步提升 。

    推薦閱讀