大規模言語モデルにおけるアテンションヘッドの線形予測可能性

要約: 大規模言語モデル（LLM）の推論はますます Key-Value（KV）キャッシュによってボトルネックとなっており、注意ヘッドの活性化の細かな構造は依然として十分には理解されていない。事前学習済みの Transformer はヘッド間に広く見られる線形構造を示す。特定のトークンに対して、注意機構のヘッドの Query、Key、Value（QKV）ベクトルは、しばしば同じ層内の限られた数の同僚ヘッドの線形結合として再構成できる。Llama-3.1-8B、Falcon3-10B、OLMo-2-7B、Qwen3-32B にまたがって、わずか 2–5 匹の参照ヘッドが多くのターゲットヘッドを高忠実度で回復する（例：C4 の Keys に対して参照が5つの場合の平均 R^2 は約 0.76、GSM8K では頻繁に R^2 > 0.85）。この予測可能性はアーキテクチャ上の性質ではなく学習によって獲得されるもので、ランダム初期化時にはほとんど欠如しており、OLMo-2 のチェックポイントを追跡する過程で事前学習中に急速に高まる。初期化時の線形予測の平均二乗誤差が大きいことを示す理論的下限によっても支持される。さらに、この出現を Key 投影サブスペースの層内整列の増加と関連付ける。最後に、この冗長性を効率化のために活用し、参照ヘッド KV 状態のみをキャッシュして残りのヘッドを軽量な線形マップでその場で再構成することで、モデル依存の精度トレードオフを伴いながら KV キャッシュを 2 倍削減している（Falcon3-10B および Qwen3-32B で 5 つのベンチマーク全体の平均で 4.5–5.5 ポイントの精度低下、Llama-3.1-8B ではより大きな低下）。また、Keys の再構成は Values の再構成より大幅に害が少ないことがわかった。