技術要約: Autonomy-of-Heads (AoH)

問題提起

大規模言語モデル（LLM）における長文脈推論は、現在、アテンション機構の二次的な計算コストと、Key-Value（KV）キャッシュの線形に増加するメモリ要件によってボトルネックとなっています。既存のソリューション（スパース・アテンションやKV圧縮手法など）は、通常、どのトークンやヘッドを保持するかを決定するために、入力依存の信号に依存しています。これらには、実行時のアテンション・スコア（Heavy-Hitterなど）、観測ウィンドウ、キャリブレーション・プロンプト、あるいは学習されたゲーティング機構が含まれます。これらは効果的ではあるものの、実行時の観測、追加のトレーニング、またはキャリブレーションの手順を必要とするため、デプロイのコストが高く、特定の入力挙動に依存してしまいます。

著者らは、凍結された重みのみによって、実行時のスコア、キャリブレーション・プロンプト、または追加のトレーニングに依存することなく、アテンション・ヘッドの機能を決定するための有用な事前知識（prior）を提供できるのではないか、という問いを投げかけています。

手法: Autonomy-of-Heads (AoH)

AoHは、凍結されたQuery-Key射影のスペクトル幾何学から直接、「リトリーバル・ヘッド（retrieval heads）」（グローバルな文脈を必要とするもの）と「ストリーミング・ヘッド（streaming heads）」（ローカル/シンク・トークンで充足するもの）を特定する、データフリーかつトレーニングフリーの手法です。

1. 理論的基礎

本手法は、各ヘッド $h$ に対してカーネル・アテンション演算子を定義します：
$$M_h = W_K^{h\top} W_Q^h \in \mathbb{R}^{d_{model} \times d_{model}}$$
この行列は、ヘッド $h$ に固有の凍結されたquery-keyのマッチング方向を表します。著者らは、$M_h$ のスペクトル幾何学がヘッドの機能的役割をエンコードしていると仮定しています：

• 集中したスペクトル (Concentrated Spectra): 少数の支配的なquery-keyマッチング方向を示しており、特定の文脈をグローバルに探索する必要があるリトリーバル・ヘッドに関連付けられます。
• 拡散したスペクトル (Diffuse Spectra): 支配的なグローバル・マッチング方向の欠如を示しており、シンク・トークンや直近のウィンドウに焦点を当てるストリーミング・ヘッドに関連付けられます。

2. 分類器としての実効ランク (Effective Rank)

AoHは、ウェイト空間におけるヘッド機能の尺度として、$M_h$ の実効ランク (Effective Rank: ER) を使用します。ERは、$M_h$ の正規化された特異値のシャノン・エントロピーの指数として計算されます：
$$\text{eff\_rank}(h) = \exp\left(-\sum_{k=1}^{r_h} \hat{\sigma}_k \log \hat{\sigma}_k\right)$$

• 低いER: 集中したスペクトルを示唆 $\rightarrow$ リトリーバル・ヘッド（フル・グローバル・アテンションを割り当て）。
• 高いER: 拡散したスペクトルを示唆 $\rightarrow$ ストリーミング・ヘッド（限定されたシンクおよび直近ウィンドウのアテンションを割り当て）。

3. 高効率な計算

フルサイズの $d_{model} \times d_{model}$ 行列 $M_h$ を構築することは計算コストが高くなります。著者らは、効率的な $d_{head}$ 次元の計算を導出しました。シルベスターの行列式定理を活用することで、$M_h$ の非ゼロの特異値をより小さなプロキシ行列 $C_h$ から計算できます：
$$C_h = (W_Q^h W_Q^{h\top})(W_K^h W_K^{h\top}) \in \mathbb{R}^{d_{head} \times d_{head}}$$
これにより、計算量は $O(d_{model}^2)$ から $O(d_{head}^2 \cdot d_{model})$ に削減され、ヘッド・ラベルの事前計算が可能になります。

4. デプロイ戦略

• ヘッドの分類: 推論前に、各ヘッドのERによってランク付けを行います。低ERのトップ $k$ 個のヘッドをリトリーバル・ヘッドとし、残りをストリーミング・ヘッドと指定します。
• KVキャッシュ管理:
  • リトリーバル・ヘッド: フル・グローバルKV状態を保持。
  • ストリーミング・ヘッド: 固定サイズの「シンク」キャッシュ（例：最初の128トークン）と「直近」ウィンドウ（例：最後の256トークン）のみを保持。
• 実装: 本手法は、ヘッドレベルのスコアをグループレベルに集計（平均集計を使用）することで、Grouped Query Attention (GQA) と互換性があります。また、専用のカーネルを必要とせず、標準的なチャンク・プリフィル（chunked prefill）やFlashAttention-2と統合可能です。

主な貢献

1. データフリーのヘッド選択: 凍れるquery-keyの幾何学から直接、リトリーバルおよびストリーミング・ヘッドを特定する手法であるAoHを導入しました。これはキャリブレーション・データや実行時の観測を必要としません。
2. 実効ランク分類器: $M_h$ の実効ランクに基づく分類器と、プロキシ行列を介してそのスペクトルを計算する効率的なアルゴリズムを開発しました。
3. スパース・アテンション・ポリシー: AoHを、精度を維持しながらレイテンシとメモリを大幅に削減する、シンプルでトレーニングフリーなスパース・アテンション・ポリシーとして具現化しました。

実験結果

著者らは、複数のモデル（Qwen2.5-7B, Qwen3-8B, Llama3.1-8B-Instruct）を用いてLongBenchでAoHを評価しました。

• 精度: 50%のスパース性（50%のヘッドをリトリーバル、50%をストリーミングとする）において、AoHはフル・アテンションの性能の平均**96.5%**を維持しました。
  • Qwen3-8Bにおいて、AoHはフル・アテンションに対しわずか0.32ポイントの差しかありません。
  • Llama3.1-8B-Instructにおいて、その差は1.10ポイントです。
  • AoHは、H2O、SnapKV、Quest、DuoAttentionといったベースライン、およびランダムおよび逆転させたヘッド選択のアブレーション（比較実験）を一貫して上回りました。
• 効率性:
  • レイテンシ: 256Kのコンテキスト長において、75%のスパース性を用いた場合、AoHはLlama3.1-8Bにおいてフル・アテンションと比較して、プリフィル速度で最大3.24倍、デコード速度で最大9.14倍の高速化を達成しました。
  • メモリ: 75%のスパース性においてKVキャッシュのメモリ使用量を最大3.98倍削減し、256Kトークンにおいて50%のスパース性で**50.0%**削減しました。

意義と主張

本論文は、**「ヘッドは自らが何を知っているかを知っている（heads "know what they know"）」**と主張しています。すなわち、その機能的役割（リトリーバルかストリーミングか）は、凍結された重みのスペクトル幾何学にエンコードされているのです。AoHは、スパース・アテンションのための、シンプルで拡張可能、かつトレーニングフリーな事前知識を提供します。

• 独立性: 従来の研究とは異なり、AoHは入力依存の信号、学習されたゲート、またはキャリブレーション・プロンプトに依存しません。これにより、実行時の観測を待たずに、最初のトークン（プリフィル）からスパース・アテンション戦略を適用することが可能になります。
• 互換性: 本手法は既存のアーキテクチャ（GQAを含む）および標準的な推論エンジン（FlashAttention）と互換性があり、実用的なデプロイが可能です。
• 堅牢性: 著者らは、実効ランクによる順序付けが異なるコンテキスト長に対しても安定しており、RoPE（回転位置エンコーディング）の回転を考慮した場合でも非常に一貫していることを示しています。

結論として、AoHは、追加のトレーニングや複雑な実行時ヒューリスティックのオーバーヘッドなしに、長文脈推論の効率を向上させるための実行可能な道筋を提供します。