Token Sparse Attention: Efficient Long-Context Inference with Interleaved Token Selection

論文「Token Sparse Attention: Efficient Long-Context Inference with Interleaved Token Selection」の詳細な技術的サマリーを以下に示す。

1. 問題提起

大規模言語モデル（LLM）は、プレフィル段階におけるアテンションメカニズムの**二次的な複雑さ（$O(L^2)$）**により、長コンテキスト推論において根本的なボトルネックに直面している。FlashAttention などのハードウェア最適化はメモリアクセスのオーバーヘッドを削減するが、アルゴリズム的な複雑さを変えるものではない。

既存の高速化手法は、主に 2 つの限界を抱えている。

• 構造化スパース性: Minference や FlexPrefill などの手法はブロックレベルのスパース性を適用する。これは、重要なトークンと同じブロックに無関係なトークンがグループ化されている場合、それらの無関係なトークンも保持されてしまい、結果としてスパース性が最適化されないことが多い。
• トークンエヴィクション: FastKV や GemFilter などの手法は、初期レイヤーでトークンを恒久的に破棄する。これはトークンの重要性が静的であると仮定しており、異なるレイヤーやアテンションヘッド間でのトークンの関連性の動的な性質を無視している。一度エヴィクトされたトークンは再検討できず、精度の低下を招く。

2. 手法：Token Sparse Attention

著者らは、軽量で動的かつ可逆的なトークンレベルのスパシフィケーションメカニズムであるToken Sparse Attentionを提案する。エヴィクションベースの手法とは異なり、トークンを恒久的に削除するのではなく、計算空間を圧縮してからシーケンスを復元する。

コアアーキテクチャ

この手法は、各アテンションヘッドごとに 2 つの段階で動作する。

1. 圧縮（QKV 選択）:
   • 各アテンションヘッド $h$ に対して、重要度スコアに基づいてトークンインデックスのサブセット $S_h$ が選択される。
   • Query（$Q$）、Key（$K$）、Value（$V$）テンソルを集約し、シーケンス長を $L' \ll L$ に削減した圧縮テンソル $\hat{Q}, \hat{K}, \hat{V}$ を形成する。
   • 主要な特徴: 選択はヘッドごとに独立して行われるため、異なるヘッドが異なるトークンのサブセットに焦点を当てることができ、ヘッドごとの不均一性に対処する。
   • 圧縮されたテンソルは密で連続したままであるため、FlashAttention などの最適化されたカーネルを使用できる。
2. 解凍（出力復元）:
   • アテンション出力 $\hat{O}$（形状 $L' \times d$）は、ヘッド固有のインデックスを使用して、元のシーケンス長 $L \times d$ のゼロ初期化テンソルに散らばる（スキャッターされる）。
   • 主要な特徴: 選択されなかった位置はゼロにされるが、前のレイヤーからの残差接続がスキップされたトークンの情報を保持する。これにより、モデルは後続のレイヤーでトークンの重要性を「再訪問」し、再評価できる。

動的トークンカバレッジ

どのトークンを保持し、どれだけの数を保持するかを決定するために、論文ではDynamic Token Coverageを導入している。

• スコアリング: 最近のクエリを用いてすべてのキーに対する軽量なプロキシアテンションマップを計算する。トークンの重要度スコアは、アテンション重みを垂直方向に合計することで導き出される。
• 予算配分: 上位 $k$ 個のトークンを保持するのではなく、トークンを昇順の重要度でソートし、「尾部」（最も重要度の低いトークン）を累積質量がカバレッジ閾値 $\tau$ を超えるまで剪定する。これは、アテンションノイズに対する構造的な正則化として剪定を扱うものである。
• ヘッド別選択: グローバルな予算を決定した後、各ヘッドは自身の固有のスコアに基づいて、それぞれ独立して上位 $k$ 個のトークンを選択する。

スパースレイヤー選択

著者らは、すべてのレイヤーにスパース性を適用すると性能が低下することを観察した。そこで、レイヤー間のトークン表現の安定性を測定するInter-Layer Representation Drift（$R_\ell$）を導入する。

• ドリフトが低い（表現が安定している）レイヤーがスパシフィケーションの対象として選択される。
• ドリフトが高い（表現が著しく変化する）レイヤーは、モデルの挙動を維持するために密なまま保持される。
• この選択は、モデルごとに一度だけ行われる前処理ステップである。

3. 主な貢献

1. 可逆的なトークンレベルのスパシフィケーション: トークンを恒久的にエヴィクトすることなく計算コストを削減し、レイヤー間でのトークン重要度の動的な再評価を可能にする、新しい「圧縮と解凍」の設計。
2. ヘッドごとの不均一性: 各アテンションヘッドに対して独立したトークン選択をサポートし、異なるヘッドが異なる意味的特徴を優先するという事実を捉える。
3. シームレスな互換性: このアプローチは、既存の密（FlashAttention）およびスパース（Minference、FlexPrefill）カーネルと完全に互換性があり、代替手段ではなく相補的なアクセラレータとして機能する。
4. 動的予算配分: 固定比率ではなく、トークン重要度の分布とレイヤーごとの安定性に基づいてスパース性レベルを決定するデータ駆動型のアプローチ。

4. 実験結果

実験は、RULER、InfiniteBench、Needle-in-a-Haystackなどのベンチマークを使用して、LLaMA-3.1-8B-InstructおよびMistral-Nemo-12B-Instructで行われた。

• 精度とレイテンシのトレードオフ:
  • Token Sparse Attention は、一貫して速度と精度のトレードオフを改善する。
  • 128K コンテキストにおいて、FlashAttention と比較して最大 $\times 3.23$ のアテンション高速化を達成し、1% 未満の精度低下に留まっている。
  • FlexPrefillと組み合わせると、平均精度（87.27%）を維持したまま、高速化を $\times 2.44$ から $\times 2.76$ に向上させる。
• エヴィクション手法との比較:
  • トークンエヴィクション手法（PyramidInfer、FastKV、GemFilter）を大幅に上回る。例えば、128K において、Token Sparse Attention は 86.84% の精度を達成するのに対し、PyramidInfer は同程度の高速化で 78.49% にとどまる。
• スケーラビリティ:
  • 高速化の恩恵はコンテキスト長とともに増大する（例：4K では最小限、128K では顕著）。これは、シーケンス長が増加するにつれてアテンションスパース性も増加するためである。
• オーバーヘッド:
  • スコアリングとインデックス管理のオーバーヘッドは最小限であり、高いスパース性レベルであっても、総アテンションレイテンシの11% 未満を占める。

5. 意義

この論文は、不可逆的なトークンエヴィクションから動的でインターリーブされたスパシフィケーションへと移行することで、長コンテキスト LLM 推論における重要なギャップを埋める。

• 理論的洞察: トークンの重要性はレイヤーやヘッドを超えて非常に動的であることを示しており、静的な剪定ではなく柔軟な選択メカニズムが必要であることを実証する。
• 実用的な影響: 既存の最先端アテンションカーネル（FlashAttention、Minference など）の上にレイヤーとして適用可能な「プラグアンドプレイ」の最適化を提供し、モデルの再学習や大幅な精度の犠牲なしに、長コンテキストシナリオで莫大な高速化を達成する。
• 将来の方向性: 永久的なトークン削除と比較して、可逆的な圧縮が $O(L^2)$ のボトルネックを処理するための優れた戦略であることを示唆し、スケーラブルな推論のための新しい設計点を確立する。