論文鏈接：https://arxiv.org/pdf/2509.01085

亮點直擊

• BSA——一種可訓(xùn)練的雙向動態(tài)稀疏注意力框架，該框架首次在視頻擴散訓(xùn)練中對全注意力機制中的查詢（Query）及鍵值對（Key-Value）進行正交稀疏化處理以加速訓(xùn)練過程。
• 為查詢塊和鍵值塊設(shè)計了不同的動態(tài)稀疏化策略，有效捕捉訓(xùn)練過程中的注意力變化特性，實現(xiàn)超越固定模式的自適應(yīng)標記選擇。
• 在Wan2.1-1.3B數(shù)據(jù)集表明：BSA可實現(xiàn)最高20倍的浮點運算量降低、17.7倍的訓(xùn)練加速以及6倍的推理加速，同時保持或超越全注意力機制的生成質(zhì)量。

結(jié)果

由于 DiT 模型采用Full Attention機制，計算量隨序列長度增加而呈二次方增長，計算復(fù)雜度為）（其中 L 為 token 序列長度）。這直接導(dǎo)致在訓(xùn)練與推理過程中的計算成本急劇攀升，嚴重制約了 DiT 模型在高分辨率長視頻生成任務(wù)中的實用性與效率，因此亟待針對性的優(yōu)化方案來解決這一核心限制。為了解決上述問題，提出了一種可訓(xùn)練的雙向動態(tài)稀疏注意力加速框架，首次對3D Full Attention中的Query和Key-Value 對分別進行動態(tài)稀疏化計算，同時設(shè)計了不同的動態(tài)稀疏化策略來提升訓(xùn)練、推理效率。

• 雙向Query-Key稀疏注意力：對于Query稀疏，通過對比token之間語義相似度來高效的選取Query內(nèi)部關(guān)鍵的query token，動態(tài)優(yōu)化query的稀疏性。對于Key-Value稀疏，只計算選取的關(guān)鍵KVBlock中的token。
• 動態(tài)稀疏注意力訓(xùn)練策略：分別針對KV block和Query block的動態(tài)稀疏性均設(shè)計了不同的動態(tài)策略。對于KV block稀疏，對不同的Query動態(tài)選擇對應(yīng)最關(guān)鍵的KV token，根據(jù)每一個訓(xùn)練step輸入的block之間的注意力分數(shù)，動態(tài)選擇關(guān)鍵 token 直至累積分數(shù)達到目標閾值p。 對于Query動態(tài)稀疏策略，分別針對時間、空間動態(tài)稀疏來選擇不同的block稀疏度。

大量實驗表明，該方法顯著加速了視頻擴散模型在不同長序列上的端到端訓(xùn)練速度，獲得了最大20倍的FLOPs減少和17.7倍的注意力訓(xùn)練加速，同時獲得了與Full Attention相當(dāng)甚至更好的生成質(zhì)量，除此之外，也可以在不降低推理質(zhì)量的情況下加速推理速度，在H100上將端到端的推理延遲從31s降低到5.2s ( 6.2x )。

問題與發(fā)現(xiàn)

解決的問題

視頻 DiT 在訓(xùn)練全分辨率、長序列數(shù)據(jù)時，大部分計算資源都耗費在注意力上，它可以消耗高達95 %的處理時間，且訓(xùn)練后的 DiT 在推理階段仍速度緩慢，這使得注意力計算成為視頻 DiT 縮放的首要瓶頸。為了改善這一狀況，近期很多工作提出了多種稀疏注意力機制。它們的核心思路是讓每個查詢Query僅與KV鍵值對的部分子集進行交互，以此來降低計算的復(fù)雜程度。它們只關(guān)注KV鍵值對中的部分冗余子集，卻忽略了Query查詢序列中同樣存在大量的冗余信息，這會導(dǎo)致大量的重復(fù)計算。除此之外，絕大多數(shù)稀疏注意力機制大多被設(shè)計成無需訓(xùn)練的形式。這些未經(jīng)過訓(xùn)練的方法通過直接截取部分KV子集來進行注意力計算，在實際訓(xùn)練中往往只能得出欠佳的結(jié)果。

關(guān)鍵發(fā)現(xiàn)

發(fā)現(xiàn)

為了設(shè)計高效的注意力訓(xùn)練框架，對當(dāng)前Full Attention的訓(xùn)練延遲進行了特異性分析，并揭示了以下兩個關(guān)鍵發(fā)現(xiàn)：

（1）Full Attention中的查詢Query和Key-Value序列均具有較大稀疏性而導(dǎo)致過多的計算浪費。

• 對于查詢Query來說，視頻幀之間及幀內(nèi) token 存在大量重復(fù)語義（如靜態(tài)背景、連續(xù)動作的相似幀）導(dǎo)致 token 數(shù)龐大。如圖2所示，F(xiàn)rame3和Frame12中的Full Attention中的query熱力圖呈現(xiàn)高度相似的表法，說明這些token提供相同的語義特征，對所有Query的序列token進行注意力查詢計算會導(dǎo)致嚴重的計算浪費。
• 對于KV鍵值對來說，token序列計算得到的注意力分數(shù)具有長尾效應(yīng)，只有部分關(guān)鍵KV子集于每個查詢Query具有強相關(guān)性，這一小部分計算顯著影響最終的輸出。因此，只需要計算小部分關(guān)鍵令牌就可以在不影響生成質(zhì)量的情況下顯著降低的計算成本。

（2）DiT中的注意力計算呈現(xiàn)動態(tài)稀疏性。動態(tài)稀疏性分別體現(xiàn)在Query和KV的時間、空間動態(tài)稀疏性。

• 空間動態(tài)稀疏：不同的 Query 所對應(yīng)的關(guān)鍵 KV 對子集本應(yīng)是動態(tài)變化的，如果采用固定的稀疏化策略，則無法適應(yīng)時空的動態(tài)稀疏，過選會造成計算冗余，漏選則會產(chǎn)生精度損失，因此需要設(shè)計動態(tài)稀疏策略來適配DiT中本身的動態(tài)稀疏性。
• 時間動態(tài)稀疏：隨著訓(xùn)練training step 的推進，稀疏度是隨時間變化的，前期注意力會獲取主要的全局信息，而后期注意力查詢則只關(guān)注于更高語義層次的局部特征，稀疏度隨著訓(xùn)練逐漸增大。

為了解決上述挑戰(zhàn)，提出了一種可訓(xùn)練的雙向動態(tài)稀疏注意力（BSA，Bidirectional Sparse Attention for Faster Video Diffusion Training）加速框架，首次對3D Full Attention中的Query和Key-Value 對分別進行動態(tài)稀疏化，同時設(shè)計了不同的動態(tài)稀疏化策略來提升訓(xùn)練、推理效率。

方法

1. Sparse Attention 回顧

2. 方法架構(gòu)

2.1 整體框架

主圖

如圖3 所示，方法框架主要分成三部分： （a）為注意力序列立方體劃分，將視頻 latent 劃分為時空立方體（Block），通過均值池化生成塊級表示來有效地篩選關(guān)鍵信息。 （b）提出的Query-Sparse方法，分別基于Query的語義冗余特征來高效的選取最優(yōu)query token，并根據(jù)時間空間動態(tài)稀疏性設(shè)計動態(tài)稀疏策略。 （c）提出的動態(tài)KV-sparse方法，對不同的Q選擇對應(yīng)最關(guān)鍵的KV token，動態(tài)選擇關(guān)鍵 token 直至累積分數(shù)達到目標閾值p，無需預(yù)設(shè)固定稀疏模式，適應(yīng)不同輸入內(nèi)容的稀疏需求。

2.2 立方體劃分

3. Query-Sparse

視頻數(shù)據(jù)本身具有多幀的時間相關(guān)性和每幀幀內(nèi)的空間相關(guān)性，因此存在時空信息冗余。實驗測試顯示在視頻擴散模型中，約 4% 的空間鄰近 token 貢獻了 80% 的注意力分數(shù)，可以去除冗余token的情況下實現(xiàn)無損性能。因此考慮到每個query查詢序列中也會存在很大的信息冗余（如靜態(tài)背景、連續(xù)動作的相似幀），主要的語義（如物體類別、動作趨勢）由少量關(guān)鍵 token 主導(dǎo)，丟棄相似語義的冗余 token 不會破壞整體語義結(jié)構(gòu)。

4. KV-Sparse

基于立方體劃分后的塊級表示，可以讓每個查詢Query僅與KV鍵值對的部分子集進行交互，以此來大量降低計算的復(fù)雜程度。但是如何確定每個查詢Query對應(yīng)的關(guān)鍵KV鍵值對子集是一個非常重要的問題。在實驗中發(fā)現(xiàn)，稀疏性在注意力塊之間和同一塊內(nèi)之間存在顯著差異，并且對于每一個query查詢對應(yīng)的關(guān)鍵kv對也是動態(tài)變化的，不應(yīng)該采用固定的top-k選擇方式來統(tǒng)一固定對每個query進行關(guān)鍵kv的選擇。

因此提出了基于統(tǒng)計閾值的動態(tài)KV-Sparse稀疏方法，分別針對每個Query選取動態(tài)的關(guān)鍵KV對，并通過輸入注意力分數(shù)的統(tǒng)計特性來計算得到動態(tài)的稀疏閾值來選取關(guān)鍵KV對，無需預(yù)設(shè)固定稀疏模式，適應(yīng)不同輸入內(nèi)容的稀疏需求。

實驗

基于Wan2.1-1.3B模型架構(gòu)進行T2V任務(wù)的模型訓(xùn)練，重新初始化進行training from scatch，所有的模型訓(xùn)練均訓(xùn)練至完全收斂，以保證公平比較。

Loss比較

如圖4所示，Sparse Attention與Full Attention基線的預(yù)訓(xùn)練損失曲線相重合，均表現(xiàn)出穩(wěn)定且平滑的下降趨勢，并且大部分優(yōu)于Full Attention 模型。

loss

Efficiency和Quality比較

如表1 所示，在2個不同的分辨率上對Sparse Attention 和Full Attention 進行from strach訓(xùn)練，分別為61 × 448 × 832,23K令牌)的原始分辨率，和擴展的更長token長度( 157x768x1280 , 153K令牌)。進行Sparse Attention和Full Attention在效率和生成質(zhì)量上的對比。

• 在原始序列長度（23k tokens）下，Sparse Attention比Full Attention的獲得了12.85倍的加速比，并且實現(xiàn)了93%的稀疏度，F(xiàn)LOPs為Full Attention的7%。除此之外，在加速的同時，BSA體現(xiàn)出了強大的生成質(zhì)量，它在Vbench的4個一致性測量指標上優(yōu)于Full Attention，尤其是在背景一致性上。這說明了Sparse Attention 可以在較短序列長度上也可以實現(xiàn)較大的加速訓(xùn)練，同時也可以達到更好的生成效果。
• 在更長的序列長度（153k tokens）下，Sparse Attention在加速比和生成質(zhì)量的優(yōu)勢上更加明顯。具體來說，BSA與Full Attention模型訓(xùn)練相比，獲得了17.79倍的加速比，稀疏度可以達到95%，F(xiàn)LOPs計算也可以達到Full Attention的5%。并且它在生成質(zhì)量上相對于Full Attention的提升幅度也更大，尤其是文本一致性和背景一致性。這種改進主要是源于對于更長的序列長度，那么模型訓(xùn)練時Attention計算的占比也更多，由此可以達到的稀疏度和加速比都會隨之增大。

Training on Longer Sequences 在不同序列長度上的對比

為了評估BSA在不同序列長度上的訓(xùn)練加速效果，分別在5種不同序列長度上進行訓(xùn)練加速比測試。所有的模型訓(xùn)練設(shè)置均保持一致來保證訓(xùn)練的公平性，結(jié)果如圖6所示。詳細地說，分別測試了23k、44k、59k、117k、153k序列長度，加速比隨著序列長度的增加逐漸增大。當(dāng)序列長度為最小的23k的時候，加速比也可以達到12.85x，當(dāng)序列長度增加為其2倍的44k的時候，加速比可以增加至14.72x。對于當(dāng)前測試的最長的序列長度153k時，最大加速比可以達到17.79倍，由此說明對于更長的序列長度，Sparse Attention可以更有效地縮短模型訓(xùn)練的時間。

speed

Sparse Adaptation 稀疏度討論

為了探究稀疏度與訓(xùn)練Loss和計算量之間的關(guān)系，還測試了不同稀疏度下的驗證損失Validation Loss和計算量FLOPs的實驗，如圖7所示。模型的稀疏度與Query-sparse中的保留token比例r和KV-sparse中的動態(tài)閾值p（動態(tài)閾值通過每一次計算得到的注意力分數(shù)來選取的k個關(guān)鍵值得到）相關(guān)，并且也存在trade-off的權(quán)衡。當(dāng)sparsity為0時，代表的是Full Attention的訓(xùn)練結(jié)果。從圖7中可以發(fā)現(xiàn)，當(dāng)Sparse Attention的稀疏度在0-0.93時，validation loss與Full Attention的Validation loss幾乎沒有區(qū)別，并且FLOPs隨著稀疏度的增加而下降。但是當(dāng)Sparse Attention的稀疏度超過0.95，雖然計算量FLOPs仍在減少，但是validation loss卻變得很大，這說明在這個稀疏度下無法實現(xiàn)無損的生成質(zhì)量。而當(dāng)稀疏度為0.93附近時，是一個最優(yōu)的結(jié)果，即既可以實現(xiàn)無損甚至更好的生成效果，還可以減少13x的計算量FLOPs。

Qualitative Results 定性實驗結(jié)果

如圖5所示，展示了4個分別在不同序列長度上的生成視頻不同幀下的T2V生成結(jié)果，分別包括不同幀數(shù)下較低分辨率（448??832）和高分辨率（782??1280）。如圖中4個不同的例子展示所示，所提出的Sparse Attention生成的視頻與Full attention相比可以達到無損的效果。

vis

Comparison with Other Sparse Attentions 與其他SparseAttention方法對比

sota

如表2所示，與最相關(guān)的基于訓(xùn)練的稀疏注意方法（如MoBA和VSA）進行了詳細的比較。BSA在加速比方面比MoBA和VSA都有明顯的優(yōu)勢，對于23k序列長度，可以達到12.85x的attention加速，但是目前training-based最優(yōu)的VSA僅可以實現(xiàn)4.5x的attention加速比。并且與這些稀疏注意力方法相比，也提供了更好的生成質(zhì)量。

Ablation Study

為了探究Query-sparse和KV-sparse對加速效果和生成質(zhì)量的影響，分別對其進行了詳盡的消融實驗，如表3所示。采取Full Attention為基線在表2的第5行，總體的方法展示在最后一行，并且分別在第1-4行來計算Query-sparse及其window窗口、KV-sparse及其統(tǒng)計動態(tài)閾值對加速效果和生成質(zhì)量的影響。

Query-Sparse

• Original Query-sparse：在沒有進行KV-sparse的基礎(chǔ)上，通過表2的第1行可以發(fā)現(xiàn)，當(dāng)保持prune rate為0.5時，可以達到無損的驗證結(jié)果，在effciency方面，并且可以實現(xiàn)1.96x的加速比，減少50%的計算量。除此之外，在Vbench上的測試結(jié)果也都優(yōu)于Full Attention。
• Query-sparse with window size selection：還測試了采用window size來根據(jù)多個center token來選取有效token的方法。這說明了with window block selection可以更好地選取包含有效語義的tokens，而不會被冗余token干擾。

KV-Sparse

• Original KV-sparse：在沒有Query-sparse的基礎(chǔ)上，基于閾值的KV-sparse可以實現(xiàn)0.86的稀疏度和6.05x倍的訓(xùn)練加速，還節(jié)省了將近8.6倍的計算量。除此之外，總體生成效果與Full Attention相比還是可以達到無損的結(jié)果。
• KV-sparse with stastic dynamic threshold：還測試了加上動態(tài)統(tǒng)計閾值的KV-sparse。從表2中的結(jié)果可以驗證，這種基于統(tǒng)計信息的動態(tài)閾值可以在相同validation loss的基礎(chǔ)上實現(xiàn)更高的稀疏度，并且在生成質(zhì)量相當(dāng)?shù)那闆r下實現(xiàn)更高的訓(xùn)練加速比和更少的計算量FLOPs。

Query-Sparse + KV-Sparse

如表2的最后一行顯示，結(jié)合了Query-Sparse 和KV-Sparse的方法在相當(dāng)?shù)膙alidation loss和生成質(zhì)量的情況下實現(xiàn)了最大的稀疏度0.93和最大的加速比12.85倍。這得益于Query-Sparse 和KV-Sparse是可以正交實現(xiàn)的，兩者達到的稀疏效果可以進行疊加，達到最優(yōu)的加速效果，并且不會損害生成質(zhì)量，驗證了稀疏注意力的有效性。并且需要強調(diào)的是，稀疏方法所增加的計算量很小，幾乎可以忽略不計，這也顯示了Sparse Attention方法的高效性。

結(jié)論

視頻擴散Transformer（DiT）模型在生成質(zhì)量方面表現(xiàn)優(yōu)異，但在生成高分辨率長視頻時遇到了主要的計算瓶頸。Full Attention的二次復(fù)雜度會增加訓(xùn)練/推理成本。 為了克服這一限制，提出了一個雙向稀疏注意（BSA）框架，用于更快的視頻DiT訓(xùn)練，這是第一個提出雙向Query-KV動態(tài)稀疏化的框架，從而提高了訓(xùn)練和推理效率。完全關(guān)注效率低下源于兩個關(guān)鍵挑戰(zhàn)：由于查詢和鍵值對固有的稀疏性而導(dǎo)致的過度計算，以及由于固定的稀疏模式無法利用DiT的動態(tài)關(guān)注而導(dǎo)致的冗余計算 。BSA通過兩個關(guān)鍵組件來解決這些問題，查詢稀疏性通過語義相似度和動態(tài)時空訓(xùn)練策略選擇信息量最大的查詢令牌來優(yōu)化，而KV稀疏性通過計算統(tǒng)計動態(tài)閾值并僅保留關(guān)鍵KV塊進行計算來實現(xiàn)。 大量實驗表明，BSA顯著加速了長序列的DiT訓(xùn)練，將FLOPs降低了20倍，實現(xiàn)了17.79倍的注意力訓(xùn)練速度，同時保持甚至超過了完Full Attention的生成質(zhì)量。

本文轉(zhuǎn)自AI生成未來 ，作者：AI生成未來

原文鏈接:??https://mp.weixin.qq.com/s/m3cMf_o8nQMGhY18tZ7i0g??