技術要約: KVEraser

問題提起

長文脈（ロングコンテキスト）の大規模言語モデル（LLM）アプリケーションにおいて、コンテキストは、検索されたドキュメント、ツールの観測結果、または会話のメモリからオンラインで動的に組み立てられることがよくあります。ここで、処理されたコンテキスト内の特定のスパンが、プリフィル（prefill）フェーズの後に「古くなった」「不正確になった」あるいは「有害になった」（例：撤回されたユーザーの好み、誤ったツールの結果、プロンプトインジェクションなど）と特定された場合、重大な課題が生じます。

「事後的なコンテキスト消去（post-hoc context erasing）」の核心的な困難さは、Key-Value（KV）キャッシュの性質にあります。因果的自己注意（causal self-attention）の下では、任意のトークンの表現は、それより前のすべてのトークンに依存します。したがって、あるスパンが編集または削除されると、その後のすべてのトークンのキャッシュされた状態は無効になります。なぜなら、それらは削除された情報を含むプレフィックスに基づいて計算されているからです。

• 完全な消去（Exact Erasing）: 標準的な解決策は、有効なプレフィックスキャッシュを再利用し、残りのサフィックス全体に対してプリフィル・パスを再実行することです。しかし、この計算コストは、消去されたスパンの長さではなく、サフィックスの長さに比例してスケールします。短い誤ったスパンの後に長いコンテキストが続くシナリオでは、このアプローチは法外なレイテンシを招きます。
• 指示のみによる忘却（Instruction-Only Forgetting）: スパンを「忘れる」という指示を単に追加するだけでは、サフィックスのKV状態はすでに汚染されているため、失敗します。モデルはデコーディング中にキャッシュされた証拠を書き換えることができないからです。
• 既存手法の限界: KVキャッシュの再利用（共有プレフィックス用など）や、除去（低アテンションのトークンを削除するなど）に関する先行研究は、あるスパンが最初から存在しなかったかのように振る舞わせるという、反事実的な挙動を誘発するという特定の要件に対処していません。

手法: KVEraser

著者らは、サフィックス全体を再計算することなく、効率的な局所的コンテキスト消去を行うために設計された、学習可能なKVキャッシュ編集手法であるKVEraserを提案しています。

コアメカニズム

KVEraserは、編集されたキャッシュの完全な再構築を、関数的な近似に置き換えます。サフィックスを再計算する代わりに、以下の手順で代替KVキャッシュ ($\widehat{KV}$) を構築します。

1. 有効なプレフィックスキャッシュ ($KV_{1:m-1}$) を保持する。
2. 消去されたスパン ($e$) のKV状態を、学習されたステアリング状態 ($KV_{steer}$) に置き換える。
3. 元の、汚染されたサフィックスキャッシュ ($KV_{n+1:T}$) を変更せずに再利用する。

目標は、ステアリング状態が欠落したスパンとサフィックス内の汚染を補償し、凍結された生成器（generator）に対し、あたかもそのスパンが存在しなかったかのような出力を導くことです。

モデルアーキテクチャ

• 生成器 ($p_\theta$): 凍結されたLLMバックボーン。
• イレイサー（Eraser, $E_\phi$）: 消去された区間のための置き換え用KV状態を予測する、学習可能なモジュール（生成器のバックボーンを、LMヘッドを除いてコピーしたものとして初期化）。
• 条件付け（Conditioning）: イレイサーは、保存されたプレフィックスKVキャッシュおよび消去されたスパン自体に基づいて条件付けられます。これはダウンストリームのクエリやサフィックスには条件付けられません。これにより、ステアリングブロックは異なるクエリに対しても再利用可能になります。

学習パイプライン

コンテキスト消去のための大規模なアノテーション済みデータの欠如を考慮し、著者らは2段階の学習戦略を提案しています。

1. 汎用的スパン・ネイバー事前学習（Generic Span-Neighbor Pre-training）: イレイサーは、Wikipediaのデータを用いた大規模な合成タスクを用いて学習されます。モデルは、挿入された「ディストラクター（妨害）」スパンの影響を抑制し、正しい隣接テキストを復元することを学習します。これにより、キャッシュ編集のための汎用的なメカニズムを教え込みます。
2. タスク固有のファインチューニング: 事前学習されたイレイサーは、以下のダウンストリームタスクに対してファインチューニングされます。
   • Needle-in-a-Haystack (NIAH): 消去された「針（needle）」を無視し、保持された「針」を想起しなければならない、制御された合成ベンチマーク。
   • 事実的ディストラクターの除去（Factual Distractor Removal）: 長いドキュメントの中に挿入された誤解を招く事実的なチャンクを無視しなければならない、QA設定。

学習目的は、教師強制（teacher-forced）による負の対数尤度損失であり、サロゲートキャッシュからの生成器の出力が、クリーンで編集されたプロンプトからの出力と一致するようにイレイサーを最適化します。

主な貢献

1. 局所的KVステアリングの定式化: 本論文は、コンテキスト消去を、グローバルなキャッシュ汚染を補償するローカルなステアリングブロックを学習する問題として定式化し、計算の依存関係をサフィックス長から消去されたスパン長へとシフトさせました。
2. KVEraserフレームワーク: 小さなKVブロックの置き換えのみを必要とし、既存のLLMサービングシステムとシームレスに統合できる、実用的な学習可能アーキテクチャを提供します。
3. 2段階の学習戦略: 汎用的な事前学習を利用して転移可能な消去能力を学習し、続いて特定のダウンストリームシナリオに適合させるという、新しいパイプラインを活用しています。
4. 包括的な評価: コンテキスト長1Kから32Kトークンにわたる、合成および自然な長文ドキュメントQAタスクにおける評価を実施しました。

実験結果

制御されたベンチマーク（NIAH）における性能

• 精度: KVEraserは、すべてのコンテキストサイズ（1K–32K）において、完全な再計算の性能に匹敵する**ほぼ完璧な完全一致（exact match）**の性能を達成しています。
• ベースライン: 近似的なベースライン（Delete-and-shift、Instruction-only forgetting、Local suffix repair）は、コンテキスト長が増加するにつれて精度が低下するか、低い精度から始まります。
• 効率性: 完全な再計算のレイテンシは、コンテキストが1Kから32Kに移動する際に17.6倍増加しますが、KVEraserのレイテンシの増加はわずか**24%**です。これは、そのコストがサフィックスの長さ ($|s|$) ではなく、消去されたスパンの長さ ($|e|$) に比例するためです。

自然な長文ドキュメントQAにおける性能

• 品質と効率のトレードオフ: 有害な事実的ディストラクターを含む未知のQAデータセット（2WikiMultiHopQA, MuSiQue, IIRC）において、KVEraserはすべての近似ベースラインの中で最高の性能を達成しました。
• 完全な再計算との比較: 完全な再計算は最も高い絶対精度を示しますが、KVEraserよりも3〜4倍遅いです。KVEraserは、品質と効率の最適なバランスを提供するパレート・フロントに位置しています。

アブレーション研究

• 条件付けのソース: 著者らは、イレイサーを保存されたプレフィックスKVおよび消去されたスパンに条件付けることが最適であることを見出しました。クエリやサフィックスの情報を追加しても一貫した改善は見られず、クエリに依存しない設計の妥当性が検証されました。
• 失敗モード: 近似的なベースラインは、消去された値と保持された値の両方を出力したり（汚染の兆候）、無関係な値を生成したりすることで失敗することがよくあります。KVEraserの失敗は稀であり、限定的です。

意義と主張

本論文は、学習可能な局所的KVステアリングが、効率的なコンテキスト消去のための実行可能かつ有望なメカニズムであることを主張しています。

• スケーラビリティ: KVEraserは、コストが残りのコンテキストの長さではなく、エラーのサイズによって決定されるという、完全な消去におけるスケーラビリティのボトルネックを解決します。
• 信頼性: 指示ベースの忘却や単純なキャッシュ削除とは異なり、KVEraserは消去されたスパンの影響を効果的に抑制し、モデルが古くなった、あるいは有害な情報に依存することを防ぎます。
• 実用性: 本手法は、完全な再計算と比較して（3〜4倍の）大幅な高速化を実現しつつ、完全な手法に匹敵する性能を維持しており、コンテキストの更新が頻繁に発生するRAGシステム、ツール拡張型エージェント、会話型AIなどの実世界のアプリケーションに適しています。

著者らは、自然なQAタスクにおけるKVEraserと完全な再計算との間の残された差は、学習データの規模と多様性に起因する可能性があると認めており、さらなるスケーリングによってこの差を埋められる可能性を示唆しています。