技術要約: MOMENTKV

問題提起

Transformerベースの言語モデルにおける自己回帰的デコーディングは、以前に計算された射影をデコーディングステップ間で再利用するために、Key-Value (KV) キャッシュに依存している。シーケンス長が数万から数十万トークンへとスケールするにつれ、KVキャッシュのメモリフットプリントは線形に増大し、長文脈推論における主要なボトルネックとなる。

KVキャッシュの退避（eviction）は、固定サイズのKVペアのサブセットを保持し、残りを破棄することでこの問題に対処する。既存の手法は、累積アテンション重みや値認識スコアリング（value-aware scoring）などに基づいたトークン選択戦略を改善してきたが、根本的な限界を共有している。それは、退避後の推論において、保持されたトークンに対してのみアテンションを再正規化し、退避された集合を完全に破棄してしまうという点である。

著者らは、出力の劣化の主な原因は、単に退避されたトークンへの残留アテンション質量（既存の手法がすでに最小化しているもの）にあるのではなく、保持されたトークンセットと退避されたトークンセットの間の**方向性の不一致（directional mismatch）にあることを特定した。実証的な分析により、実際には、退避されたトークンは保持されたトークンに対してほぼ直交（near-orthogonal）**していることが明らかになった。その結果、たとえ退避されたアテンション質量がわずかであっても、標準的な再正規化では回復できない大幅な方向誤差を生じさせる。これにより、「方向性のギャップ（directional gap）」が生じる。つまり、より選択的な保持を行うほど、保持されたセットが特定のクエリに整合した方向に集中するため、退避されたセットとの差がさらに広がり、トークン選択の改善だけでは収穫逓減に陥るのである。

手法: MOMENTKV

この方向性のギャップに対処するため、著者らは、退避されたトークンセットに対してコンパクトで小規模なモーメント統計量を維持するフレームワーク、MOMENTKVを提案する。これらの統計量は、退避の決定を導くことと、推論中のアテンション出力を補正するという二つの目的を果たす。

1. コンパクトなモーメント統計量

MOMENTKVは退避されたトークンを直接保持する代わりに、各アテンションヘッドに対して、各退避時に逐次更新される4つの実行和を保持する：

• カウント ($n_e$): 退避されたトークンの数。
• キー和 ($s_k$): $\sum_{i \in E} k_i$。
• 値和 ($s_v$): $\sum_{i \in E} v_i$。
• 外積和 ($S$): $\sum_{i \in E} v_i k_i^\top$。

これらの統計量により、退避されたトークンを再訪問することなく、キー平均 ($\bar{k}$)、値平均 ($\bar{v}$)、および値-キー共分散 ($\tilde{S}$) を計算することが可能になる。総ストレージオーバーヘッドは、コンテキスト長に依存せず、ヘッドあたり $O(d^2)$ である。

2. 退避された出力の一次近似

推論中、MOMENTKVは、退避されたアテンション出力 $\hat{f}_E(q)$ の閉形式の一次近似を導出する。ロジットを平均キーを中心に再中心化し、ソフトマックス関数に一次テイラー展開を適用することで、著者らは以下を導出した：
$$\hat{f}_E(q) = \bar{v} + \frac{1}{n_e \sqrt{d}} \tilde{S} q$$

• $\bar{v}$ 項は、クエリに依存しないベースラインとして機能し、退避された部分出力の平均的な方向を回復する。
• $\tilde{S} q$ 項は、クエリ適応型の補正として機能し、ベースラインだけでは捉えきれない垂直方向の変動（方向的コンテンツ）を捉える。

この近似により、モデルは失われた退避された寄与を推定して再構築することができ、標準的な退避によって抹消されてしまう方向情報を効果的に復元できる。

3. モーメントに基づく退避（Moment-Informed Eviction）

モーメント統計量は、退避プロセス自体もガイドする。保持される各トークン $j$ について、アルゴリズムは以下のモーメント残差を計算する：
$$r_j = v_j - \bar{v} - \frac{1}{n_e \sqrt{d}} \tilde{S} k_j$$
この残差は、蓄積されたサマリーが既に捉えている内容を超えて、トークン $j$ がどれだけの新しい方向的情報を持っているかを測定する。退避スコアは、アテンションの重みと残差ノルムの積として定義される：$\text{score}(j) = \alpha_j \cdot \|r_j\|$。

• 低残差のトークンは、モーメントモデルによってよく予測されるため、退避の優先順位が高くなる。
• 高残差のトークン（独自の方向的情報を持つトークン）は、アテンションの重みが低くても保持される。

これにより、相互に強化されるループが形成される。モーメントに基づく退避は、退避されたセットを「幾何学的に規則的」（低いロジット分散 $\sigma$）に保ち、推論補正のための近似境界をタイトにする。逆に、正確な補正は、将来の退避決定に対してより良い残差を提供する。

4. 正規化補正推論（Normalization-Corrected Inference）

最終的なアテンション出力は、保持された出力 ($f_R$) と近似された退避出力 ($\hat{f}_E$) の凸結合として再構成される：
$$\hat{f}(q) = \hat{w}_R f_R(q) + (1 - \hat{w}_R) \hat{f}_E(q)$$
混合重み $\hat{w}_R$ は、イェンセンの不等式を通じて導出された、退避された分割関数の下界を用いて推定される。これにより、モデルは退避された質量の大きさを考慮しつつ、その方向を補正することが保証される。

主な貢献

1. 方向性ギャップの特定: 本論文は、KVキャッシュ退避における支配的なエラー源は、単なる退避されたアテンションの質量の大きさではなく、保持された部分出力と退避された部分出力の間のほぼ直交性であることを確立した。
2. コンパクトなモーメント統計量: 各ヘッドあたり $O(d^2)$ のストレージのみを使用して、退避されたトークンの幾何学的特性（平均と共分散）を追跡するメモリ効率の高い手法を導入した。
3. 閉形式の補正: 平均的な方向とクエリ依存の変動の両方を回復する、退避されたアテンション出力の一次近似を導出した。
4. デュアルロール・メカニズム: モーメント統計が、退避の選択（冗長なトークンの特定）と推論出力の補正（失われた方向情報の復元）の両方を同時に最適化する、統一されたフレームワーク。

実験結果

著者らは、LLaMA-3.1-8B-Instruct および Qwen3-4B-Instruct を用いて、LongBench および RULER ベンチマークで MOMENTKV を評価した。

• パフォーマンス: MOMENTKV は、すべてのキャッシュ予算 ($L \in \{128, 256, 512, 1024\}$) において、すべてのベースライン（H2O, SnapKV, PyramidKV, Ada-KV）を一貫して上回った。
• 積極的な圧縮: 最大の利得は、積極的な圧縮下で観察された。LongBench（$L=128$, LLaMA-3.1-8B）において、MOMENTKV は最強のベースラインである Ada-KV に対して平均 +1.35 のスコア向上を達成し、フルキャッシュ性能の 94.1% を維持した。
• 検索タスク: 精密なトークン検索をテストする RULER において、MOMENTKV は LLaMA および Qwen それぞれに対して Ada-KV よりそれぞれ +3.4 および +3.5 の比例的な利得を示し、検索集中的な設定における有効性を証明した。
• アブレーション研究:
  • 正規化補正推論 (NC) は、モーメントに基づく退避 (MI) 単独（+0.57）よりも大幅に高い貢献（+2.41）をしており、退避された方向を回復することが主要な性能向上要因であることを裏付けている。
  • 両方のメカニズムを組み合わせることで相乗効果（合計利得 +3.17）が得られ、個々の利得の和を超えた。
• 効率性: この手法は、トークンあたりわずか 1.6 ms のレイテンシを追加するだけであり、補助メモリとして 4.1 MB（保持されたキャッシュの 6.4%）を必要とする。一方で、フルキャッシュと比較して、トークンあたりのレイテンシを 65% 以上削減している。

重要性

本論文は、長文脈推論のさらなる向上の道は、単に優れたトークン選択を行うことではなく、現在の退避戦略が抹消してしまう方向的情報を回復することにあると主張している。MOMENTKV は、退避されたセットに対してコンパクトな統計量を維持することで、方向性のギャップを埋めることができ、退避による「不可避なコスト」を、補正可能な近似誤差へと変えられることを示している。このアプローチは、固定予算の KV キャッシュの効率性を損なうことなく、長文脈言語モデルをデプロイするための実用的かつスケーラブルなソリューションを提供する。