稀疏注意力 - Shong

発展の経緯#

起源#

スパースアテンション（Sparse Attention）は、クエリベクトルと一組のキー・バリューペアを出力ベクトルにマッピングする最適化されたアテンションメカニズムですが、シングルヘッドアテンションやマルチヘッドアテンションとは異なり、クエリベクトルとすべてのキーの類似度を計算するのではなく、クエリベクトルと一部のキーの類似度のみを計算することで、計算量とメモリ消費を削減します。スパースアテンションの概念は、2018 年の論文『Generating Long Sequences with Sparse Transformers』に初めて登場し、この論文では、8000 語を超えるテキストシーケンスを処理するためにスパースアテンションを使用した Transformer ベースの長シーケンス生成モデルが提案されました。

実際に訓練された Transformer モデルでは、アテンションマトリックスはしばしばスパースであることが多く、これはすべてのトークンが他のすべてのトークンに注意を払う必要がないことを意味します。トークン間の相互作用の中には、最終的な出力に対する寄与が小さく、無視できるものもあります。

スパースな方法は、固定パターン（例えば、ローカルウィンドウ）、コンテンツに基づく選択（例えば、現在の位置に最も関連する他の位置）、または学習によって得られたパターンである可能性があります。

スパース接続の度合いを決定する基準に基づいて、方法は二つのカテゴリに分けられます：位置に基づくスパースアテンションとコンテンツに基づくスパースアテンション。

位置に基づくスパースアテンション#

spa_att

グローバルアテンション

グローバルノードの概念を設定し、これらのグローバルノード上で、1 つのノードが他のすべてのノードに注意を払うことができる。別の理解方法としては、これらのノードがすべてのノードの相互交流の中継点として機能することを意味し、スパースアテンションの本質は、ポイントツーポイントで各ノードが交流する必要がないということです。P2P から P2S の概念に似ています。
バンドアテンション

データ分布の局所的な性質を考慮し、スライディングウィンドウの概念に似ており、1 つのノードが周囲のノードにのみ注意を払うことを意味し、アテンションの相互作用をローカルアテンションに制限します（アテンションが辛うじて得た長距離の受容野が、あなたのこのカットで壊れてしまいます）。
ダイレーテッドアテンション

バンドアテンションの基礎の上に、より遠くの相互作用ノードを設定するために余白を残し、スライディングウィンドウにステップサイズの間隔を設定することに相当します（カットした部分が効果が悪いことがわかり、延長しようとしていますが、依然として非常に拙劣に感じます）。
ランダムアテンション

各クエリに対してランダムにいくつかのエッジをサンプリングして実現します。純粋にランダムな当選のように感じ、効果はあまり良くないでしょう。
ブロックアテンション

入力シーケンスをいくつかの重複しないクエリブロックに分割し、各クエリブロックにローカルメモリブロックを割り当てて、長シーケンスの効率的な処理を実現します。（このものがよくわからないのですが、これはアテンションが長さと幅を 1 つ小さくしただけではないですか🤓）

コンテンツに基づくスパースアテンション#

最大内積検索（MIPS）

コンテンツに基づくスパースグラフを効率的に構築するために、最大内積検索（Maximum Inner Product Search、MIPS）問題の解決策を利用できます。MIPS の目標は、クエリと最大の点積を持つキーを見つけることであり、クエリとすべてのキー間の点積を計算する必要はありません。

NSA（推論効率と訓練の実行可能性）#

起源#

長文モデリングは大規模モデルにとって非常に重要ですが、従来のアテンションメカニズムは平方の計算複雑度を持ち、コンテキストの長さを増やすと大量の追加計算が必要になります。長さが 2 倍になると、計算量は 3 倍になります $(2n)^2-n^2=3n^2$

最先端のスパースアテンションは、KV キャッシュの排出とブロック KV キャッシュの選択、サンプリング、クラスタリングの方法に分かれています。しかし、これらの二つの方法は、彼らが言うほど良くありません。

主に以下のいくつかの側面の問題があります：局所的スパース、アテンションに適合しない、エンドツーエンドの訓練

局所的スパース：H2O のような方法は、自回帰デコード段階でのみスパースマトリックスを適用しましたが、プレフィルでは計算集約型の前処理が必要です。MInference はプレフィル時にのみスパースアテンションを採用しました。これらの方法はすべての段階でスパースアテンションを実現していないため、プレフィル主導の作業（例えば、書籍の要約、コード補完）やデコード主導の作業（例えば、思考の連鎖）ではパフォーマンスが悪くなります。つまり、下流タスクに対してエンドツーエンドの訓練を実現するための統一されたアーキテクチャがありません。
アテンションに適合しない：ほとんどのスパースマトリックスは MHA のスパースを考慮していますが、MQA や GQA のような構造には適合しない場合があります。例えば、Quest メソッドでは、各アテンションヘッドには独立した kv キャッシュがありますが、MQA や GQA のように共有されることがあります。
エンドツーエンドの訓練：現在ほとんどのスパースマトリックスは推論タスクに対して設計されており、訓練タスクに対するスパースマトリックスが必要ですが、密なマトリックスで訓練されたモデルはスパース推論下でのパフォーマンスが悪く、20％のアテンションが 70％のアテンションスコアをカバーすることしかできません。さらに、ClusterKV や MagicPIG のような作業は不連続な計算グラフを導入し、逆伝播が正常に実行できなくなります。不連続なメモリアクセスは、FlashAttention などの高速アテンション技術への効果的な適応を妨げ、これらの技術は連続したメモリアクセスとブロック計算を利用して高スループットを実現します。

NSA の報告では、主に解決すべき二つの問題が挙げられています：

一つはハードウェアと連携した推論最適化で、プレフィルとデコードの二つの段階で、理論的な最適化を実際の加速に変えるためには、ハードウェアに優しいアルゴリズムが必要です。主にメモリアクセスとハードウェアボトルネックのスケジューリングに関してです。
二つ目は訓練を意識したアルゴリズム設計で、エンドツーエンドの学習スパースパターンをサポートし、従来の「先に訓練し、後に剪定する」性能損失を回避します。

解決策として提案された方法は、主に三つのステップに分かれています：

圧縮粗粒度トークン（cmp）

連続したキー / バリューブロックをブロックレベルの表現に集約し、粗粒度の意味情報をキャッチし、計算負担を軽減します。

言い換えれば、kv の複数の次元を 1 つの次元に統合することです。例えば、1024 次元の kv を 64 次元に変えることです。
選択的に保持された細粒度トークン（slc）

重要なトークンを選択的に保持し、圧縮によって生じる情報損失を補います。

言い換えれば、MIPS のように最も関連性の高いトークンアテンションを探し、他のトークンは注目に値しないということです。
スライディングウィンドウ（win）

局所的コンテキストを専門に処理するスライディングウィンドウブランチを導入し、局所パターンが学習プロセスを支配する可能性のある問題を解決します。

言い換えれば、上記のバンドアテンションのことです（これが本当に役立つとは🤯）。

NSA

デモ#

このプロセスを説明するために、簡単な例を挙げることができます：

現在の入力は $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$ と同じサイズのベクトルブロックに変換します。言い換えれば、多くのブロック $k$ を 1 つの $k$ に変えて、 $k$ が占めるメモリを削減し、計算を加速します。この時、元の $q$ と圧縮された $k'_{:8},v'_{:8}$ を使用してアテンションスコアを計算し、圧縮アテンション $o_{64}^{cmp}$ を得ます。

中間部分は $Selection$ と呼ばれ、圧縮時に得られた圧縮された KV ブロック $k'_{:8},v'_{:8}$ から、最大のいくつかのアテンションスコアを計算します。ここでは $Top-2$ を選択し、仮に $k'_{3,7},v'_{3,7}$ とします。つまり、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$ メモリを節約しました。

背景#

アテンション#

新しいクエリ $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 の計算能力を十分に活用できるようになります

プレフィル段階では、大量の因果自己アテンションが高い算術強度を示し、性能が計算能力に制約されます。自回帰のデコード段階では、各トークンを生成するたびに以前のすべての kv キャッシュにアクセスする必要があるため、メモリ帯域幅に制約されます。このような差異は、プレフィルと訓練の際に計算複雑度を減少させ、デコード段階ではメモリアクセスを減少させるという最適化方向の不一致をもたらします。

方法#

二つの方向：アルゴリズム側の設計とカーネル最適化

全体の概要#

元の $k_{:t},v_{:t}$ をよりコンパクトで情報密度の高い $\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]

圧縮粗粒度トークン（圧縮）#

\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$ より小さい必要があり、情報の断片化を軽減します。

原文では、圧縮表現がより粗い粒度の高次の意味情報をキャッチし、アテンションの計算負担を軽減できると述べられています。言っていることが正しければ、実験結果が良ければそれが王です（🤓）。

選択的に保持された細粒度トークン（選択）#

粗粒度の上記のトークンだけを使用することは必然的に不十分であり、大量の細粒度情報が失われます。モデルがより良く理解できるように、細粒度のブロックも必要です。

ブロックの選択：

ハードウェアに優しい考慮とアテンションスコアの固定分布に基づいています。このステップは、現代の GPU で効率的な計算を実現するために重要です。現代の GPU は連続ブロックアクセスのスループットがランダムインデックスの読み取りよりもはるかに優れており、ブロック計算も GPU のテンソルコアを最大限に活用できます。アテンションスコアは通常、空間的連続性を示し、隣接するキーがしばしば類似の重要性レベルを持つことを示しています。これは 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 ヘッドに対して、彼らの重要なアテンションスコアは同じであり、すべてのアテンションスコアを合計してこの 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}\}]

スライディングウィンドウ#

アテンションメカニズムでは、局所パターンが通常より早く適応し、学習プロセスを支配する可能性があり、これがモデルが前の二つの 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}$ の 6 つの KV 値を取得した後、ゲート制御の方法で結果を取得します。

カーネル設計#

トレーニングとプレフィル段階で FlashAttention レベルの加速を実現するために、Triton を利用してハードウェアに整合したスパースマトリックスを実現しました。

圧縮段階とスライディングウィンドウでは、FlashAttention をうまく利用して最適化できるため、ここで言及されるカーネル最適化は、選択段階で生成される離散アテンションシーケンスの計算に主に焦点を当てています。

GQA と MQA の最適化について、FlashAttention の戦略に従うと、時間的に連続したクエリブロックを SRAM にロードすると、メモリアクセスの効率が低下します。なぜなら、ブロック内のクエリは交差しない KV ブロックを必要とする可能性があるからです。この問題を解決するために、GQA 内のすべての共有同じ kv ブロックのクエリヘッドを SRAM に一緒にロードします。

グループセンターデータのロード：

内ループ内で、 $h$ 個のヘッドを持つクエリ $Q \in \mathbb R^{[h, d_k]}$ をロードし、次に彼が属する圧縮されたインデックス $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

上記の緑色のブロックは、1 つの q と一連の kv 計算を表しています。ここで、t が増加すると、選択した KV ブロックが常に 3 ブロック以下であるため、NSA の加速がより顕著になります。

MoBA（大道至簡、即挿即用）#

背景#

スパース性については、アテンションスコアのスパース性だけでなく、記憶ストレージに関連する脳領域で観察されるスパース接続特性についても言及されています（ACL に投票できるかもしれません🤓）。

従来のスパースアテンションには二つの大きな欠陥があります。一つは、特定のタスクに基づいた事前定義された構造を採用しており、汎用性が非常に低いこと、もう一つは、トークンを動的に選択してスパースアテンションを訓練する方法であり、この方法は一般的に訓練段階には全く役立ちません。

NSA が解決しようとしている推論速度と訓練可能性の二つの問題に似て、MoBA が解決する問題も推論の加速と訓練可能性です。ブロックアテンション混合メカニズム（MoBA、Mixture of Block Attention）は、専門家の混合（MoE）を MLP からアテンションに移行します。

コアメソッドは、各 $q$ に対して関連する過去の $kv$ ブロックを動的に選択する機能です。

方法#

MoBA のコアメソッドは、ブロックパーティショニング（block partitioning）と選択戦略（selection strategy）です（NSA の圧縮ブロックと選択に非常に似ているように聞こえます🤔）。

全体#

MoBA の方法論は非常にシンプルで、 $o=Attn(q_{:N},k_{:N},v_{: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$ 個のブロックを選択します。ここでの $Score$ の計算方法は、 $Score = <q,mean\_pool(k_{I_i:I_{i+1}})>$ であり、外側の尖括弧は内積を示し、mean_pool は平均値を示し、 $i$ ブロックでの平均 $q*k$ 値を計算することに相当します。

MoBA

次に、MoBA は自回帰言語モデルにおいて因果関係を維持することが重要であると述べており、 $q$ が未来の任意の $kv$ ブロックにルーティングできないことを確認する必要があります。特に特異な状況の一つは、「現在のブロック」をクエリトークン自体を含むブロックとして定義することです。現在の $kv$ ブロックへのルーティングは因果関係に違反する可能性があり、なぜなら全体の $kv$ ブロックの平均プールが未来の $kv$ ブロックの情報を意図せずに含む可能性があるからです。この問題を解決するために、各トークンがそれぞれの現在の $kv$ ブロックにルーティングされることを強制し、現在の $kv$ ブロックのアテンション演算中に因果マスクを適用します。

最終的な考察#

MoBA と NSA のコアの違いは以下の通りです：

MoBA は $kv$ を分割し、より小さなブロックを選択して計算しますが、NSA は圧縮後に小さなブロックを選択して計算し、さらにスライディングウィンドウを加えます。この計算のコアロジックは異なり、MoBA の選択は内積の topk によって選択され、勾配が関与する必要はありませんが、NSA の選択は勾配が戻って修正されます。
NSA は KV ブロックの細粒度を取得しますが、MoBA は異なるクエリヘッドが異なる $kv$ ブロックにアクセスできるようにします。重点が異なり、両者は互いにできないことを行います。

著者の疑問：

MoBA の中で小さなブロックに分割した後のアテンションスコアは FlashAttention を使用できますが、なぜ NSA の中で小さなブロックを選択した後にそれを使用できないのでしょうか？