視頻生成模型太貴太慢怎么辦？

普林斯頓大學和Meta聯合推出的新框架LinGen，以MATE線性復雜度塊取代傳統自注意力，將視頻生成從像素數的平方復雜度壓到線性復雜度，使單張GPU就能在分鐘級長度下生成高質量視頻，大幅提高了模型的可擴展性和生成效率。

實驗結果表明，LinGen在視頻質量上優于DiT（勝率達75.6%），并且最高可減少15×（11.5×）FLOPs（延遲）。此外，自動指標和人工評估均顯示，LinGen-4B在視頻質量上與最先進模型相當（分別以50.5%、52.1%、49.1%的勝率優于Gen-3、Luma Labs和Kling）。

方法：線性復雜度的MATE模塊

LinGen維持Diffusion Transformer（DiT）中的其他結構不變，而將其計算瓶頸——平方復雜度的自注意力模塊替換為線性復雜度的MATE模塊，它由MA分支和TE分支組成。

其中，MA分支包含一個雙向的Mamba2模塊。

Mamba2作為State Space Model（SSM）的變體，善于處理超長的token序列，同時又對硬件非常友好，可以使用attention的各種硬件加速核，如xformers，FlashAttention等。但是Mamba系列模型在語言任務上的優秀表現難以直接遷移到大型視覺任務上，生成的高分辨率視頻往往一致性很差、質量不高。

一些特殊的scan方法嘗試解決這一問題，如Zigzag scan，Hilbert scan，但它們都要求對序列做復雜的順序變換，而這個操作對硬件極其不友好。在處理高分辨率、長視頻時，會帶來顯著的額外延遲。

針對于此，LinGen提出Rotary Major Scan（RMS），相鄰層中四種scan方式交替切換。

以上圖的方式為例，W，H和T分別在展開時有第一、第二和第三優先級，通過交換展開的優先級，就可以實現不同的scan方式。

相比于已有方法，該方法最大的好處是對硬件非常友好、可以通過簡單的tensor reshaping實現，因此也幾乎沒有額外開銷，同時還把scan后原相鄰token的平均距離降到了和已有特殊scan方式相同的水平。

然而，所有這些特殊的scan方式仍然不足以完全解決Mamba的臨近信息丟失問題，因為在模型的任意一層中，只會有一種scan方式被應用，如果不考慮跨層交流，大量臨近信息在單層中依舊有損失。

針對于此，LinGen在TE分支中應用了TEmporal Swin Attention（TESA）：它是一種特殊的3D window attention，窗口范圍在不同層中會滑動，每一個窗口都很小，并且窗口大小不隨視頻分辨率和長度（即3D tensor的大小）的變化而變化。

這是因為TESA僅用來處理最臨近的信息，這一固定的窗口大小也使得TESA實現了相對3D tensor中token數的線性復雜度。

作為額外的補充，LinGen還在MA分支中引入了review tokens。它被用以增強視頻中極長程的一致性，例如在60秒視頻的結尾復現視頻前幾秒消失的人。它把待處理video tensor的概覽提前寫入Mamba的hidden state memory中，為后續的視頻處理提供幫助。

評估：遠超基線，對標SOTA

從人類評測和模型自動評測兩個角度將LinGen與已有的先進視頻生成模型、以及DiT baseline進行比較。

無論是人類評測的結果，還是在VBench上的自動評測的結果，都顯示LinGen與先進的商業模型Kling、Runway Gen-3生成的視頻質量接近，并且遠勝于OpenSora v1.2。

可以看到，在FLOPs方面，當生成17秒、34秒和68秒長度的512p視頻時，LinGen-4B相對于DiT-4B分別實現了5×、8×和15×的加速；

在延遲方面，當在單個H100上生成512p和768p的17秒視頻時，LinGen-4B相對于DiT-4B分別實現了2.0×和3.6×的加速；

當生成17秒、34秒和68秒長度的512p視頻時，LinGen-4B相對于DiT-4B分別實現了2.0×、3.9×和11.5×的延遲加速。

這說明LinGen具有線性復雜度，可以在單卡上實現分鐘級視頻生成，速度遠快于DiT。與相同大小的DiT相比，LinGen可實現推理速度11倍以上的提升。

另外，LinGen和相同大小、在相同數據集上以相同training recipe訓練的DiT baseline相比，在視頻質量和文字-視頻一致性上取得全面領先。相比起DiT，LinGen可以更快地適應更長的token序列。

通常認為自注意力模塊的線性替代是對完整自注意力的近似，雖然在速度上有顯著優勢，但在模型性能上往往略遜一籌，而LinGen打破了這個慣有的看法。

在整個預訓練過程中，模型從低分辨率圖像生成開始，學習低分辨率視頻生成，再不斷增加所生成視頻的分辨率和長度，所處理的token數增長了上千倍。

而在從少token數的任務遷移到多token數的任務時，LinGen的適應性遠強于DiT（a圖中是從256x256分辨率視頻生成遷移到512x512分辨率視頻生成任務時的loss curve），這可能是受益于Mamba對于長序列的高適應性，這一特征已經在語言任務上被觀察到。

為了進一步驗證這里推理，選取這一預訓練階段的早期checkpoint進行比較，發現LinGen比DiT的win rate優勢變得更加顯著。這暗示了雖然LinGen在任務遷移的早期能大幅領先DiT，但是這種優勢隨著預訓練的進行，在不斷減小。

盡管如此，在訓練資源有限的情況下，LinGen在預訓練的極長一段時間內仍舊能對DiT保持優勢。

項目主頁：https://lineargen.github.io/

論文鏈接：https://arxiv.org/abs/2412.09856

項目代碼：https://github.com/jha-lab/LinGen