Context Memorization for Efficient Long Context Generation

技術的概要：効率的な長文脈生成のためのアテンション状態メモリ

問題定義
現代の大規模言語モデル（LLM）アプリケーションは、推論を導くために、長文脈の条件付けプレフィックス（例えば、インコンテキスト例、検索されたドキュメント、またはエージェント的指示）への依存度を高めています。しかし、プレフィックス付加による推論は、以下の 2 つの主要な構造的限界に直面しています：

1. プレフィックスの減衰：生成が進むにつれて、モデルのアテンションは新たに生成されたトークンに分散し、特に長文脈シナリオにおいてプレフィックスの影響が薄れていきます。
2. 推論の非効率性：プレフィックスに対するアテンション計算は、その長さに比例して線形にスケーリングします（$O(L)$）。これにより、プリフィルおよび各デコードステップの両方で、大幅なレイテンシとメモリオーバーヘッドが生じます。プレフィックスキャッシングはプリフィルコストを償却しますが、依然として多大なメモリ消費を伴います。

既存の解決策は、いずれも欠点を伴う 2 つのカテゴリーに分類されます：

• プレフィックスの内部化：モデルパラメータを微調整するか、ハイパーネットワークを使用してプレフィックスの振る舞いを内部化する方法です。これらは勾配ベースの学習を必要とするため、リソース集約的で更新が遅く、動的なプレフィックスの変化には不適切です。
• プレフィックスの圧縮：プレフィックスをアテンション機構内に保持しつつ圧縮する方法（例えば、トークン剪定または KV キャッシュの排除）です。これらは依然として圧縮されたプレフィックスへのアテンションを推論時に必要とするため、レイテンシは圧縮された長さに比例してスケーリングし、プレフィックスの影響は減衰の対象となります。

手法：アテンション状態メモリ（ASM）
著者らは、トレーニング不要なアプローチである**アテンション状態メモリ（Attention-State Memory）**を提案します。これは、プレフィックスを、軽量でルックアップベースの、事前計算されたアテンション状態のメモリへと外部化する手法です。この手法は、連結されたキーブロックに対するアテンションを、部分ブロックに対するアテンションへと分解するために、オンラインソフトマックス恒等式を利用します。

核心的な洞察は、特定のクエリに対して、アテンション出力と正規化項からなるアテンション状態 $(a, Z)$ が、元のキーとバリューへの再アテンションなしに、プレフィックスのアテンションへの寄与を再構成するのに十分であるという点です。

このフレームワークは 2 つのフェーズで動作します：

1. オフライン較正フェーズ（メモリ構築）：

   • 収集：モデルは、プレフィックスと連結された代表的なクエリ - 応答トレースのセットに対してフォワードパスを実行します。応答内のすべてのクエリトークンに対して、システムはクエリベクトルとプレフィックスに対する対応するアテンション状態 $(a, Z)$ を記録します。
   • クラスタリング：収集されたクエリベクトルは（K 平均法などを通じて）$K$ 個の重心にクラスタリングされます。
   • 集約：各クラスタに対して、単一の代表エントリが作成されます。重心クエリは、そのクラスタ内のクエリの平均です。アテンション状態 $(\bar{a}, \bar{Z})$ は、アテンションを考慮した集約を通じて計算され、オンラインソフトマックス更新規則を使用してクラスタ内の全メンバーの状態をマージします。これにより、重心がクラスタのプレフィックスアテンションの損失なしの平均を代表することが保証されます。
   • 効率性：GPU メモリを超えずに非常に長いプレフィックス（例えば 16K 以上のトークン）を処理するために、プレフィックスをチャンクに分割し、独立してエンコードし、オンラインソフトマックス恒等式の合成性を利用してメモリ内でマージすることができます。

2. オンライン推論フェーズ：

   • 検索：入力されたクエリに対して、システムはコサイン類似度（オプションでホワイトニング変換と RoPE 処理を含む）を使用して、最も近いメモリ重心を特定します。
   • マージ：検索された事前計算されたアテンション状態 $(\bar{a}, \bar{Z})$ は、オンラインソフトマックス恒等式で定義されたマージ演算子を使用して、クエリの自己アテンション出力（プレフィックス以外のトークン上で計算されたもの）とマージされます。これにより、プレフィックスが存在した場合に生成されたはずの正確なアテンション出力が、プレフィックストークンへのアテンションなしに回復されます。
   • 計算量：階層的インデックス化を通じて、検索コストはメモリサイズに対して対数的に（$O(\log K)$）スケーリングし、元のプレフィックスの長さには依存しません。

主要な貢献

1. トレーニング不要な外部化：本論文は、勾配ベースの学習や補助モデルを必要とせずに、長文脈プレフィックスをコンパクトなメモリ辞書へ外部化する手法を提案します。
2. オンラインソフトマックスの拡張：著者らは、効率的なアテンション計算から、クエリ間プレフィックスの再利用のドメインへとオンラインソフトマックス恒等式を拡張し、検索によるプレフィックスアテンションの損失なしの再構成を可能にしました。
3. 性能と効率性：実験により、ASM がフルアテンションの性能と同等かそれ以上でありながら、計算コストを大幅に削減することが示されました。

実験結果
この手法は、ManyICLBench（インコンテキスト学習）およびRuleArena/NBA（検索拡張生成）において、LLaMA-3.1-8Bを使用して評価されました。

• インコンテキスト学習（ICL）：ManyICLBench において、ASM は 1K から 8K のエントリまでのメモリ予算全体にわたり、標準的なインコンテキスト学習よりも精度を向上させました。8K のメモリ予算において、フルアテンションと比較してアテンションレイテンシを1.36 倍削減しました。ASM は、ラベルトークンのような重要な情報をしばしば失う KV キャッシュ圧縮ベースライン（KVzip）を上回りました。
• 検索拡張生成（RAG）：NBA ベンチマーク（約 20K トークンのルールを含む）において、ASM はメモリフットプリントのわずか**20%**を使用して、フルアテンション RAG の性能を上回りました。最適な性能は中間的なコードブックサイズ（4K エントリ）で観察され、エントリ数は単に最大化すべきものではなく、タスク固有のハイパーパラメータであることを示唆しています。
• レイテンシ：階層的ルックアップは非常に小さなメモリ予算ではより高いオーバーヘッドを伴いますが、レイテンシの増加は非線形です。ASM は約 4K のメモリエントリでフルアテンションよりも高速になり、16K エントリでは1.8 倍の高速化を達成しました。

意義と主張
本論文は、アテンション状態メモリが、プレフィックス知識をクエリごとのアテンション計算から切り離すことで、長文脈生成の構造的コストに対処すると主張しています。プレフィックスを事前計算されたルックアップベースのメモリへ外部化することで、この手法はプレフィックスの長さに対する推論レイテンシの線形スケーリングを排除し、標準的なアテンション機構に関連するプレフィックス減衰の問題を解消します。

著者らは、この研究を、テキストプロンプトやモデルパラメータを超えて存在する、コンパクトで再利用可能な表現へと LLM の知識を外部化する、より広範な研究方向への一歩として位置づけています。彼らは、このアプローチが精度と効率性の間で有利なトレードオフを提供し、パラメータの内部化に対するトレーニング不要な代替手段となり、プレフィックス圧縮に対するより効率的な代替手段を提供することを強調しています。