稀疏注意力 - Shong

發展歷程#

起源#

稀疏注意力（Sparse Attention）是一種優化的注意力機制，它可以將一個查詢向量和一組鍵值對映射到一個輸出向量，但與單頭注意力和多頭注意力不同的是，它不會計算查詢向量和所有鍵向量的相似度，而是只計算查詢向量和部分鍵向量的相似度，從而減少計算量和內存消耗。稀疏注意力的概念最早出現在 2018 年的論文《Generating Long Sequences with Sparse Transformers》中，該論文提出了一種基於 Transformer 的長序列生成模型，其中使用了稀疏注意力來處理超過 8000 個單詞的文本序列。

實際上在訓練好的 Transformer 模型中，注意力矩陣往往是稀疏的，這意味著並不是每個令牌都需要關注其他所有令牌。有些令牌之間的相互作用可能對最終的輸出貢獻不大，可以被忽略。

稀疏的方式可以是固定的模式（如局部窗口）、基於內容的選擇（如與當前位置最相關的其他位置），或者是通過學習得到的模式。

根據確定稀疏連接的度量標準，方法分為兩類：基於位置的稀疏注意力和基於內容的稀疏注意力。

基於位置的稀疏注意力#

spa_att

Global Attention

設置了全局節點的概念，在這些全局節點上，一個節點能關注其他所有節點，換種理解方式也就是：這些節點充當了所有節點相互交流的中轉站，稀疏注意力的本質就是沒必要點對點的每個節點都去交流，有點像是 P2P 到 P2S 的概念。
Band Attention

考慮數據分布的局部性質，也就是類似於滑動窗口的概念，也就是一個節點只關注它周圍的節點，將注意力的交互限制在局部注意力中（Attention 好不容易有了的長距離感受野被你這麼一截斷就壞了）。
Dilated Attention

在 Band Attention 的基礎上，通過留白給自己設置一個更遠的交互節點，相當於是在滑動窗口上給自己設置一個步長間隔（截了一段發現效果不好又想着延長一下，不過感覺依然很拙劣）。
Random Attention

為每個查詢隨機抽樣一些邊來實現，感覺純隨機的中獎，效果應該不咋樣。
Block Attention

將輸入序列劃分為若干個不重疊的查詢塊（query blocks），並為每個查詢塊分配一個局部記憶塊（memory block），來實現對長序列的高效處理。（沒太懂這個玩意，這不就是 Attention 把長和寬都變小了一號嗎🤓）

基於內容的稀疏注意力#

最大內積搜索（MIPS）

為了高效地構建基於內容的稀疏圖，可以利用最大內積搜索（Maximum Inner Product Search，MIPS）問題的解決方案。MIPS 的目標是找到與查詢具有最大點積（dot product）的鍵，而不需要計算查詢與所有鍵之間的點積。

NSA (推理效率和訓練可行)#

起源#

長文本建模對於大模型來說是極其重要的，但是傳統的注意力機制是平方的運算複雜度，增加 context length 會增加大量的額外計算，長度增加兩倍，計算量會增加三倍 $(2n)^2-n^2=3n^2$

最先進的稀疏注意力分成了兩類，分別是 KV-cache eviction 和 block KV-cache 的選擇，採樣，聚類的方法。但是這兩類方法都沒有他們說得那麼好。

主要是以下幾個方面的問題: 局部稀疏，不適配 Attention，端到端訓練

局部稀疏：H2O 這種方法只在自回歸 decode 階段應用了稀疏矩陣，但是 prefill 中需要計算密集型預處理。MInference 則只在 prefill 時採用了稀疏注意力。這些方法都沒有在所有階段實現稀疏注意力，那麼對於 prefill 主導的工作比如書籍摘要，代碼補全或者 decode 主導的工作比如思維鏈上就會表現不好，也就是針對下游任務沒有一個統一的架構來實現端到端訓練。
不適配 Attention: 大部分稀疏矩陣考慮的是對 MHA 的稀疏，而對於 MQA 和 GQA 這樣的結構，會有不適配的情況，比如 Quest 方法，每一個注意力頭都有它獨立的 kv-cache，然而對於 MQA 和 GQA 這樣共用。
端到端訓練: 現在大多數稀疏矩陣都是針對的推理任務，需要一個針對訓練任務的稀疏矩陣，但是基於 dense 矩陣訓練的模型在稀疏推理下表現不佳，因為 20% 的注意力只能 cover70% 的 Attention Score。更有甚者像是 ClusterKV 和 MagicPIG 這樣的工作引入了不連續的計算圖，從而導致反向傳播無法正常執行。非連續的內存訪問阻止了對 FlashAttention 等快速注意技術的有效適應，這些技術依賴於連續的內存訪問和分塊計算來實現高吞吐量。

NSA 的報告中提到了主要要解決的是兩個問題：

一是與硬件聯合的推理優化，在 prefill 和 decode 兩個階段，將理論優化變為實際加速需要硬件友好型的算法，主要是在內存訪問和硬件瓶頸調度上。
二是訓練感知算法設計，支持端到端學習稀疏模式，避免傳統方法「先訓練後裁剪」的性能損失。

提到了解決的方法，主要分為三步：

compressed coarse-grained tokens(cmp)

將連續的 key/value 塊聚合為塊級別的表示，捕獲粗粒度的語義信息，減少計算負擔。

通俗來講就是把 kv 的多個維度融於一個維度，舉個例子: 1024 維的 kv 變為 64 維度。
selectively retained fine-grained tokens(slc)

有選擇地保留重要的 token，彌補壓縮可能帶來的信息損失。

通俗來講就是類似於 MIPS 的尋找最相關的 token 注意力，其餘 token 不值得關注。
sliding windows(win)

專門處理局部上下文的滑動窗口分支，解決局部模式可能主導學習過程的問題。

通俗來講就是上面的 Band Attention (壞了它還真有用 🤯)。

NSA

Demo#

我們可以舉一個簡單的例子來說明一下這個過程：

我們現在有的輸入是 $q_{:t},k_{:t},v_{:t}$ ，假設 $t=64$ ，然後假設按照長度為 8 進行 $Compress$ ，由於 $k,v$ 是對稱的，我們用 $k$ 舉例，即可把 $k$ 各分為 8 塊 $k_{1:8},k_{9:16},...,k_{57:64}$ ，經過壓縮後，也就是把 $k_{1:8}$ 變為和 $k_i$ 一樣 size 的向量塊，通俗來講就是把很多塊 $k$ 變為一個 $k$ 來減少 $k$ 所佔的顯存並加速計算，此時用原始的 $q$ 和壓縮後的 $k'_{:8},v'_{:8}$ 進行注意力分數計算得到壓縮注意力 $o_{64}^{cmp}$ 。

中間部分叫做 $Selection$ ，在壓縮時我們得到了壓縮後的 KV 塊 $k'_{:8},v'_{:8}$ ，此時計算最大的幾個注意力分數，我們這裡選擇 $Top-2$ 的，假設為 $k'_{3,7},v'_{3,7}$ ，也就是第三塊和第七塊，此時還原對應的選出來的壓縮塊，也就是將 $k'_{3}$ 擴展回到 $k_{17:24}$ 這樣去處理拿到我們所需要的 $kv$ 塊，然後計算得到選擇注意力 $o_{64}^{slc}$ 。

右邊是滑動窗口，在原 $k_{:64},v_{:64}$ 中選擇最近的 8 個 $k_{57:64},v_{57:64}$ 就可以得到滑動窗口注意力 $o_{64}^{win}$ 。

最後再用一個門控函數控制，即

O = g_{64}^{cmp}*o_{64}^{cmp}+g_{64}^{slc}*o_{64}^{slc}+g_{64}^{win}*o_{64}^{win}

分析一下節省的 $kv$ ，原本有 64 個 $kv$ ，我們的壓縮注意力用了 8 個 $kv$ ，選擇注意力用了 16 個 $kv$ ，滑動窗口注意力用了 8 個 $kv$ ，相當於現在一共只用了 32 個 $kv$ ，節省了一半的 $kv$ 顯存。

背景#

Attention#

對於一個新來的查詢 $q$ 需要查詢之前所有的 t 個 $kv$ 對

$o_t=Attn(q_t,k_{:t},v_{:t})=\sum^{t}_{i=1}\frac{\alpha_{t,i}*v_i}{\sum_{j=1}^{t}\alpha_{t,j}} \enspace \alpha_{t,i}=e^{\frac{q_t^T*k_i}{\sqrt{d}}}$

算術強度#

$T_{mem}=\frac{bytes}{BW_{mem}}$ ：訪問內存時間等於內存中訪問的字節數除以處理器的內存帶寬。

$T_{math}=\frac{ops}{BW_{math}}$ ：數學時間等於運算次數除以處理器的數學帶寬。

如果 $T_{math}>T_{mem}$ ，那麼該算法是受數學限制的。

上式可以被替換為 $\frac{ops}{bytes}>\frac{BW_{math}}{BW_{mem}}$ ，左邊是算法實現操作數與訪問字節數的比值，被稱為算法的算術強度，右邊是處理器的數學帶寬與內存帶寬的比值，被稱為字節比率。

如果算法的算術強度高於 GPU 的 字節比率，那麼該算法受算力限制的，也稱 math bound，即性能受算力 FLOPS 限制（算力受限 / 計算密集型算子）。
如果算法的算術強度低於 GPU 的 字節比率，則該算法受內存限制，也稱 memory bound，即性能受內存帶寬限制（內存受限 / 訪存密集型算子）。

應該儘可能讓算法 / 網絡層的算術強度高於 GPU 的字節比率，這樣才能充分利用 gpu 的算力。

在 prefill 階段，大量的 causal self-attention 所展現出的批量矩陣乘法展現出了高算術強度，即性能受到算力限制。在自回歸的 decode 階段，因為其每生成一個 token 都需要訪問之前所有的 kv-cache，則變為受到內存帶寬限制。而這種差異性會帶來優化方向的不一致，在 prefill 和 train 的時候減少計算複雜度，在 decode 階段減少內存訪問。

方法#

兩個方向：算法側設計和 kernal 優化。

整體概況#

將原始的 $k_{:t},v_{:t}$ 優化成更加 compact 和 information-dense 的 $\tilde{K_{t}},\tilde{V_{t}}$ ，其中這個變換是基於 $q_t$ 動態改變的，用公式表達就是：

\tilde{K_{t}}=f_K(q_t,k_{:t},v_{:t}) \enspace \enspace \enspace \tilde{V_{t}}=f_V(q_t,k_{:t},v_{:t})\enspace \enspace \enspace o_t^*=Attn(q_t,\tilde{K_{t}},\tilde{V_{t}})

對於函數映射，一共有三種方式，也就是上面所說的 cmp，slc，win 三種方式，通過一個門控因子來控制採用哪種映射，具體如下所示：

o_t^*=\sum_{c \in C}g_t^cAttn(q_t,\tilde{K_{t}},\tilde{V_{t}})\enspace \enspace \enspace \enspace \enspace \enspace g_t^c \in [0,1]

compressed coarse-grained tokens (壓縮)#

\tilde{K}_{t}^{cmp} = {f}_{t}^{cmp}(k_{:t})={\Phi(k_{id+1:id+l}|0\le i\le \lfloor \frac{t-l}{d}])}

其中 $l$ 是塊的長度， $d$ 是塊間滑動的步長， $\Phi$ 是一個可學習的 MLP，將塊中的鍵映射到一個壓縮鍵。文中提到通常來講這個 $d$ 是要小於 $l$ 的，來緩解信息碎片化。

原文說壓縮表示可以捕獲更粗粒度的高級語義信息，並減輕注意力的計算負擔，說是啥就是啥吧，實驗出來結果好就是王 (🤓)。

selectively retained fine-grained tokens (選擇)#

只使用粗粒度的上述 token 勢必是不行的，丟失了大量細粒度的信息，我們還需要細粒度的塊來幫助模型更好地理解。

塊的選擇：

基於硬件友好的考慮和注意力分數的固定分布。這一步對於在現代 GPU 上實現高效的計算至關重要。現代 GPU 在連續塊訪問上的吞吐量遠勝於基於隨機索引的讀取，同時塊計算也可以最大效率利用 GPU 的 tensor core。注意力分數通常表現出空間連續性，這表明相鄰的鍵往往具有相似的重要性水平，這是 DS 後面做實驗發現的，淺色區域表示較高的關注值，如圖所示。

block

重要注意力分數計算：

計算所有的注意力分數顯然是一個開銷很大的事，但是我們可以通過計算前一步壓縮後的注意力來降低這個開銷。

p_{t}^{cmp}=softmax(q_t^T*\tilde{K}_{t}^{cmp})

但上面這個只是基於壓縮的注意力分數，普遍意義上，我們需要的選擇塊長度定義為 $l'$ ，當 $l'=l=d$ 時， $p_{t}^{cmp}=p_{t}^{slc}$ ，對於分塊不一致的情況，給定 $l \le l',\enspace\enspace d|l,\enspace\enspace d|l'$ ，那麼

p_{t}^{slc}=\sum_{m=0}^{\frac{l'}{d}-1}\sum_{m=0}^{\frac{l}{d}-1}p_{t}^{cmp}[\frac{l'}{d}j-m-n]

但其實 $p_{t}^{tls}$ 和 $p_{t}^{cmp}$ 是不一樣的，在 NSA 方案裡面，讓他們一致。在 GQA 和 MQA 中，對於不同的且共享同樣的 KV 值的 Q head ，他們的重要注意力分數是相同的，就是把所有的注意力分數相加作為這個 KV 的注意力分數，這一步直接可以節省大量內存。

p_t^{slc'}=\sum^{H}_{h=1}p_t^{slc,(h)}

選擇最大的 k 個注意力分數：

選擇最大的 k 個注意力分數，注意到這裡我們的選出的是壓縮塊，也就是 $p_t^{cls}[i]$ ， $rank$ 表示排名。

I_t=\{i|rank(p_t^{slc}[i])\le n\}

基於壓縮塊再復原最開始的所有 $k$ 塊。

\tilde{K}_{t}^{slc}=Cat[\{k_{il':(i=1)l'|i\in T_t}\}]

sliding windows (滑動窗口)#

在注意力機制中，局部模式通常適應得更快，並且可以主導學習過程，這可能會阻止模型從前兩種 kv 中有效學習。為了解決這個問題，引入了一個專用的滑動窗口分支，它顯式地處理原始上下文，允許其他分支（壓縮和選擇）專注於學習它們各自的功能。

\tilde{K}_{t}^{win}=k_{t-w:t} \enspace \enspace \enspace \enspace \enspace \tilde{V}_{t}^{win}=v_{t-w:t}

三個分支分別提供了獨立的鍵和值。這種架構設計通過防止局部和全局之間的梯度干擾來實現穩定的學習，同時引入最小的開銷。獲取 $\tilde{K}_{t}^{cmp}, \tilde{V}_{t}^{cmp},\tilde{K}_{t}^{slc}, \tilde{V}_{t}^{slc},\tilde{K}_{t}^{win}, \tilde{V}_{t}^{win}$ 六個 KV 值後，採用 gate 門控的方式從中獲取結果。

Kernel Design#

要在 train 和 prefill 階段實現 FlashAttention 級別的加速，利用 Triton 實現了和硬件對齊的稀疏矩陣。

在壓縮階段和滑動窗口都可以很好利用 FlashAttention 進行優化，因此這裡提到的 kernel 優化主要是針對選擇階段所產生的離散注意力序列的計算。

針對 GQA 和 MQA 進行優化，如果我們遵循 FlashAttention 的策略，將時間連續的查詢塊加載到 SRAM 中，這將導致內存訪問效率低下，因為塊內的查詢可能需要不相交的 KV 塊，為了解決這個問題，將 GQA 中所有共享相同 kv 塊的查詢頭一起加載到 SRAM 中。

組中心數據加載：

在內循環中，加載有 $h$ 個頭的查詢 $Q \in \mathbb R^{[h, d_k]}$ ，然後找到剛才他屬於的被壓縮的 index $I_t$ 。

共享 KV：

在內循環中，根據 $I_t$ 選入此時需要的 $K \in \mathbb R^{[B_k, d_k]} ,\enspace V \in \mathbb R^{[B_k, d_v]}$ ，其中 $B_k$ 是滿足 $B_k|l'$ 最小的 kernel 塊大小。

Kernel

以上的一個綠色塊代表了一個 q 和一段 kv 計算。這裡我們注意到，當 t 增大時，由於我們選擇的 KV 塊恆為小於等於 3 塊，那麼越長 NSA 的加速越明顯。

MoBA (大道至簡，即插即用)#

背景#

對於稀疏性，不僅提到了注意力分數的稀疏性，還提到了與記憶存儲相關的腦區觀察到的稀疏連接特性（感覺可以投 ACL🤓）。

傳統的稀疏注意力有兩大缺陷，一是採用預定義的結構，基於一些特定的任務，泛化性很差，二是動態選取 Token 進行稀疏注意力訓練的方式，這個方法一般對於訓練階段毫無幫助。

與 NSA 要解決的推理速度和可訓練兩個問題差不多，Moba 解決的問題也是加速推理以及可訓練，塊注意力混合機制 (MoBA,Mixture of Block Attention)，將專家混合 (MoE) 從 MLP 遷移到 Attention。

核心方法是針對每個 $q$ 動態選擇相關歷史 $kv$ 塊的功能。

方法#

MoBA 核心的方法是 block partitioning（塊分區）和 selection strategy（選擇策略）（聽著是不是和 NSA 的壓縮塊然後選擇很像🤔）。

總體#

MoBA 的方法論很簡單， $o=Attn(q_{:N},k_{:N},v_{:N})$ ，假設長度為 $N$ ，然後將長度為 $N$ 的輸入分為 $n$ 個小塊，其中定義塊的大小為 $B=\frac{N}{n}$ ，此時定義一個索引，這個索引主要是為了後續選擇塊。

I_i=[(i-1)*B+1,i*B]

然後計算 $q$ 在 $k_{:I_1},k_{I_1:I_2},...,k_{I_{n-1}:I_n}$ 上的 $Score$ ，選出最大的 $k$ 個最大的 $k$ 塊，注意這裡的 $Score$ 的計算方式， $Score = <q,mean\_pool(k_{I_i:I_{i+1}})>$ ，外層尖括號表示內積，mean_pool表示平均值，相當於計算在 $i$ 塊上的平均 $q*k$ 值。

MoBA

然後 MoBA 提到了在自回歸語言模型中保持因果關係很重要，需要確保 $q$ 無法 Route 到任何未來的 $kv$ 塊上。其中一個比較特殊一點的情況是，將 “current block” 定義為包含查詢 token 本身的 block，到當前 $kv$ 塊的 Route 也可能違反因果關係，因為整個 $kv$ 塊的平均池可能會無意中包含來自未來 $kv$ 塊的信息。為了解決這個問題，我們強制要求每個 Token 都必須 Route 到其各自的當前 $kv$ 塊，並在當前 $kv$ 塊 Attention 運算期間應用因果掩碼。

最終思考#

MoBA 和 NSA 的核心不一樣在哪裡：

MoBA 幹的是將 $kv$ 分塊，然後選擇更小的塊進行計算，NSA 幹的是壓縮後選擇小塊進行計算再加一個滑動窗口。這個計算的核心邏輯就不一樣，MoBA 選擇是通過內積的 topk 來選擇的，是不需要梯度參與的，而 NSA 的選擇其實是會有梯度回傳來修正的。
NSA 幹的是取 KV block 的細粒度，MoBA 幹的是讓不同的 query head 能接觸到不同的 $kv$ 塊，側重點不一樣且兩者都不能幹對方幹的事。

筆者問題：

MoBA 裡面的分成小塊後計算 Attn 分數可以使用 FlashAttention，那為什麼 NSA 裡面選擇完小塊後不能這麼使用呢？