技術要約: MemDecay

問題提起

大規模言語モデル（LLM）は、マルチターン対話、ツール利用、および反復的なプランニングを含む長期間のワークフローを実行するエージェントとして、ますます多く展開されています。これらのエージェントのトレースは、静的な長文脈タスクとは異なり、システム指示、安全性制約、進化する計画、ユーザーのターン、検索されたドキュメント、ツールの出力、および中間推論状態といった、不均一なコンポーネントを蓄積します。

自己回帰的デコーディングに必要なKey-Value（KV）キャッシュは、このトレースに伴って線形に増加し、重大なメモリオーバーヘッド、帯域幅への要求、およびレイテンシを引き起こします。既存のKVキャッシュ退避（eviction）ポリシーは、通常、モデル中心の信号（累積アテンション、最新性ウィンドウ、またはアテンションシンクなど）に依存しており、シーケンス全体に対して一律の保持ルールを適用します。このアプローチは、オーケストレーションされたプロンプトの固有のセマンティック構造を活用できていません。エージェントのワークロードにおいて、異なるコンポーネントは極めて異なる機能的役割と期待される寿命を持っています。例えば、システム指示はエピソードを通じて存続しなければなりませんが、スクラッチパッドの推論や古いツールの出力は、短期間のみ関連する場合があるかもしれません。すべてのトークンを同じ期待寿命として扱うことは、この構造を無視することになり、最適ではないメモリ管理を招きます。

手法: MemDecx

著者らは、長期間のLLMエージェント向けに設計された、トレーニング不要でリージョン認識型のKVキャッシュ退避フレームワークであるMemDecayを提案しています。この手法は、固定されたキャッシュ予算の下で動作し、オーケストレーション層からの構造的事前知識と、アテンション由来の重要性を統合します。

コアコンポーネント

1. リージョンラベル: オーケストレーション層が、連続するプロンプトのスパンに対するリージョンラベル（例：system, plan, user, tool_in, tool_out, retrieval, scratchpad）を提供します。
2. 保持スコア: 各トークン $i$ には、時刻 $t$ における保持スコア $s_i(t)$ が割り当てられ、次のように計算されます：
   $$s_i(t) = b_{r(i)} \exp(-\lambda_{r(i)} \Delta t_i(t)) + \alpha a_i(t)$$
   • 構造的事前知識: $b_{r(i)}$ は、トークンのリージョンに特有のベースとなる優先度です。
   • 時間的減衰: $\lambda_{r(i)}$ は、リージョン固有の減衰率です。項 $\Delta t_i(t)$ は、トークンが最後に「リフレッシュ」されてからの経過時間を表します。
   • アテンションの重要度: $a_i(t)$ は、トークンが受け取ったアテンション・マス（attention mass）の指数移動平均（EWMA）です。
   • リフレッシュ・メカニズム: トークンへのアテンションが閾値 $\tau$ を超えた場合、その減衰クロックがリセットされ（$\Delta t_i$ が0になる）、継続的な使用に基づいて実質的に寿命が延びます。
3. ピン留め（Pinning）: システム領域（およびオプションでプランのセグメント）は「ピン留め」することができ、これはスコアに関わらず決して退避されないことを意味します。
4. 退避ポリシー: キャッシュはページ粒度（例：16トークンあたりの1ページ）で管理されます。予算を超過した場合、ページ内の構成トークンの平均保持スコアによってページがスコアリングされます。最も低いスコアを持つページが最初に退避されます。

キャリブレーション

MemDecayはモデルの重みに関してはトレーニング不要ですが、各リージョンの減衰率（$\lambda_r$）を適合させるための軽量な統計的キャリブレーションを採用しています。このプロセスは、エージェントのトレースからリージョン条件付きのアテンション寿命を推定し、LLMを再学習させることなく、特定のワークロード特性に適応することを可能にします。

主な貢献

1. 定式化: 本論文は、リージョン認識型の選択的忘却を、オーケストレーションされたプロンプトのセマンティック構造に根ざしたLLMエージェントのメモリ管理問題として定式化しました。
2. アルゴリズム: リージョン固有のベース優先度、時間的減衰、およびアテンションによるリフレッシュを単一の解釈可能な保持スコアに組み合わせたポリシーであるMemDecayを導入します。これは、既存のサービングスタック（例：PagedAttention）と互換性のあるページ粒度のレイアウトで動作します。
3. 測定プロトコル: 著者らは、リージョン条件付きのアテンション寿命を推定し、減衰率をキャリブレーションするためのプロトコルを規定しています。これらを2つのコンテキスト長（約450および約1,700トークン）と2つのモデルサイズ（Qwen2.5-1.5Bおよび3B）にわたって実行し、合計4,320個のスコア付き退避プローブを実施しました。

実験結果

本研究では、ランダム退避、最新性ベース（ストリーミング）保持、および累積アテンション保持（H2Oスタイル）を含むベースラインと比較して、MemDecayを評価しています。

主な知見

• リージョンの寿命階層: 測定されたアテンション寿命は、リージョン間で桁違いに異なります。すべての設定において、systemトークンの半減期は148〜189デコードステップである一方、scratchpadトークンは急速に減衰します（14〜16ステップ）。Retrievalトークンは、最も長く生存する可動リージョンであることが判明し、これらが急速に減衰するというデフォルトの仮定を覆しました。
• ピン留めの有効性: システム領域をピン留めすることで、指示の生存が保証されます。MemDec明は、フルキャッシュの天井において、全設定を通じて24個中21〜24個のシステムプローブを保持しましたが、最新性ベースのベースラインは0〜2個しか保持できませんでした。
• 最新性 vs 構造: 最新性ベースの保持（ストリーミング）は、コンテキスト長が増加するにつれて崩壊し、長いティアではほぼゼロの性能となります。対照的に、MemDecayの構造的事前知識は効果を維持します。
• 累積アテンション: ピン留めされていないコンテンツにおいて、累積アテンション（H2Oスタイル）はモデルスケールとともに強化され、3Bモデルの設定においてMemDecayよりもアンピン（unpinned）な想起において優れた性能を示しました。
• MemDecayの限界: MemDecayは、古い事実（例：最初のターンのユーザー領域の事実）に関するアンピンな想起において、アテンションベースのベースラインと比較して特定の損失を示しました。アブレーション研究により、アテンション項（$\alpha a_i$）の大きさが、これらのトークンの構造的減衰に対抗するには小さすぎることが明らかになりました。著者らは、アテンション項が効果的であるためには、単なる重みのチューニングではなく、大きさの正規化（例：一様アテンションレベルに対する相対化）が必要であると指摘しています。

意義と主張

本論文は、MemDecayを、エージェントのワークフローのセマンティック構造にキャッシュ保持を一致させる、実用的かつトレーニング不要なソリューションとして位置付けています。その意義は以下の通りです：

• 解釈可能性: ポリシーは明示的なリージョンラベルと減衰率によって駆動されており、保持の決定を透明にします。
• 互換性: モデルの再学習や補助モジュールを必要とせず、ページ粒度のキャッシュで動作するため、既存のサービングスタックに適合します。
• 安定性の保証: 重要なリージョンをピン留めすることで、メモリプレッシャーによるシステム指示や安全性制約の消失を防ぎます。これは、アテンションや最新性ベースの手法では提供できない保証です。
• 適応性: キャリブレーションプロトコルは、リージョンの寿命がワークロードに依存しており、減衰率を最適化するために測定可能であることを示しており、一律の静的な退避ポリシーの使用に疑問を投げかけています。

著者らは、MemDecayが構造的事前知識とピン留めされた指示の保持には成功しているものの、今後の課題として、アンピンな想起を改善するためのアテンション項の大きさの正規化や、不可逆的な退避を扱うためのティア拡張（例：ホストメモリへのオフロード）を検討する必要があると結論付けています。