近日,南京理工大學本科校友、美國東北大學沈軒博士生和所在團隊提出一種全新的加速視頻擴散模型的方法,名為 Draft Attention。
圖 | 沈軒(來源:沈軒)
這一方法的核心創新在于:研究團隊首先將隱藏空間中的 tokens 重塑為“視頻”的結構,即按照時間、高度、寬度 [t, h, w] 的格式進行 reshape。隨后對 reshaped 的 query 和 key 進行下采樣,采用平均池化操作構建一個低分辨率的注意力圖。這一草圖注意力圖能夠高效捕捉視頻中的關鍵區域,同時識別出存在冗余的部分。
接著,研究團隊利用該草圖注意力圖生成稀疏掩碼,來引導完整分辨率下的稀疏注意力計算,僅保留重要的注意力連接。通過這種方式他們在幾乎不影響生成質量的情況下,顯著降低了注意力模塊的計算開銷。
研究團隊的理論分析也從數學上證明了:這種“草圖”注意力圖與原始全分辨率注意力圖之間的誤差是可控的,并且引入的稀疏性誤差也是有界的。
實驗結果進一步驗證了方法的有效性——Draft Attention 在保持生成質量的同時,在 H100 和 A100 GPU 上分別實現了最高 1.75 倍和 2 倍的端到端加速效果。
簡單來說,本次工作不僅提出了一種無需重新訓練、即插即用的加速框架,也為高效視頻生成提供了新的思路。
這項成果具有廣泛的應用前景,尤其是在對生成速度和計算資源要求較高的場景中會非常有價值。比如,在未來一到兩年內,研究團隊可以預見它被用于提升視頻生成平臺的效率,降低內容創作者使用 AI 生成高質量視頻的門檻。此外,它還可以用于多模態大模型中的視頻生成模塊,加速模型響應,提升用戶交互體驗。
往更長遠看,隨著邊緣計算能力的提升,研究團隊的方法有望被部署到本地終端設備,比如手機、XR 眼鏡等,從而實現高質量視頻的本地快速生成。這對于虛擬現實、游戲制作、數字人等領域都具有重要意義。
同時,草圖注意力作為一種輕量、高效的稀疏計算機制,也可能被擴展到超清圖像生成、視頻編輯、三維建模等其他視覺生成任務中。總的來說,這項工作為大規模、高質量視覺生成的高效實現提供了一種新的思路。
據介紹,沈軒和朋友韓晨夏(https://cxhan.com/)在一次關于 AI 加速的討論中,注意到當前的視頻生成任務,尤其是高分辨率視頻的生成,存在極高的時間成本。即便是在頂級 GPU 如 A100 和 H100 上,生成一個僅幾秒鐘的視頻也要花費幾十分鐘。研究團隊意識到,這一領域在生成效率方面仍有很大的優化空間,于是決定深入探索視頻生成加速的可行性。
后來,他們關注到其他團隊的一篇新論文 Sparse VideoGen (SVG)(https://arxiv.org/abs/2502.01776)。這項工作通過在空間和時間維度上分別構建稀疏注意力算子,在一定程度上緩解了視頻生成中的冗余計算問題。
然而,研究團隊在深入分析后發現了該方法的一些局限性:SVG 所采用的稀疏模式是預定義的,且只支持兩種稀疏策略,這種固定的設計容易在高稀疏度下損失視頻生成質量,限制了其適應性和效果。
隨后,沈軒開始動手實踐,基于騰訊開源的視頻生成模型 HunyuanVideo(https://github.com/Tencent-Hunyuan/HunyuanVideo),嘗試在其中引入 Max Pooling 機制,以 20% 的稀疏度生成視頻。初步結果出人意料地好,這一實驗結果讓他們意識到:視頻擴散過程中的注意力計算存在大量冗余,遠比研究團隊最初預想的要多。
為了進一步提升計算效率,研究團隊嘗試將草圖注意力圖引導下的稀疏模式對齊到模塊化稀疏注意力結構中,并發現在滿足 pool_h × pool_w = block_size 的條件下,可以更好地與現有的高效注意力框架兼容。
然而,新的挑戰也隨之而來:草圖注意力生成的稀疏掩碼在原始空間中是離散分布的,難以直接高效執行。為此,研究團隊設計了一種重排策略,將這些稀疏塊聚集成連續的內存布局,使得注意力計算可以在 GPU 上高效執行。
下圖直觀地展示了這一過程:草圖注意力圖(Draft Map)所對應的模塊化稀疏注意力是離散的(下圖左邊的 Full Map),只有引入適當的重排(Reorder)才能夠使得這些稀疏模塊聚集在一起(下圖右邊 Reordered Full Map)。
(來源:arXiv)
在驗證整體思路可行后,沈軒繼續深入開展實驗,進一步優化了設計細節。他們發現,在高稀疏度的設定下平均池化相比最大池化在生成質量上表現更佳,尤其是在保持背景一致性和主體細節方面更為穩定。
隨后,沈軒采用了 90% 稀疏度生成了一系列視頻,結果表明不僅生成速度顯著提升,而且視頻質量幾乎沒有明顯下降。為了更全面地評估方法性能,他們引入了多項圖像和視頻評估指標,包括 PSNR、SSIM、LPIPS 以及視頻生成領域的權威基準 VBench。
(來源:arXiv)
對比結果顯示,在相同稀疏率下,研究團隊的方法在生成質量上明顯優于 SVG。此外,他們還對草圖注意力圖的計算開銷進行了量化評估。結果表明,盡管他們在前期引入了 draft attention 的計算過程,但整體額外開銷極小,相較于整個視頻生成過程幾乎可以忽略不計。這進一步證明了本次方法在加速效率和生成質量之間實現了良好平衡。
下一步,研究團隊打算繼續將量化(Quantization)引入到 Draft Attention 當中來進一步提高視頻生成速度。
參考資料:
https://arxiv.org/pdf/2505.14708
運營/排版:何晨龍
特別聲明:以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布,本平臺僅提供信息存儲服務。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.