AKV、履歴を書き換えずにエージェントのKVキャッシュを管理

検索、ツール呼び出し、観測、行動を繰り返すLLMエージェントでは、処理の軌跡が急速に長くなります。軌跡がモデルのコンテキスト上限を超えると、古い履歴をどう扱うかが問題になります。しかし負担は入力トークン数だけではありません。軌跡に対応するKVキャッシュを保存し、GPUやメモリの間で読み書きし、後続リクエストで復元するコストも増大します。長時間稼働するエージェントでは、KVキャッシュのI/Oと保存が主要なサービングコストになり得ます。

履歴の要約が抱える再計算コスト

従来のコンテキストエンジニアリングでは、古い履歴を削除したり、要約に置き換えたりして入力を短くします。ところが、Transformerの後続状態は前方の計算に依存します。入力のプレフィックスを変更すると、残したテキストが同じでも、その後に計算されたキャッシュが無効になる可能性があります。その結果、保持したコンテキストを再びプリフィルしなければなりません。

再計算を抑えるために履歴を多く残せば、メモリとI/Oが増えます。逆に、履歴を削りすぎればタスクの正確さに影響する可能性があります。論文は、このトレードオフを「テキストの書き換え」ではなく「キャッシュの管理」によって解決しようとします。

dropとreposによるキャッシュエンジニアリング

提案手法では、エージェントの元の軌跡を変更せず、後続の注意計算に参加させるキャッシュだけを制御します。主な操作は次の通りです。

• drop：選択したキャッシュを後続の注意機構から除外する。
• repos：保持するキーの位置エンコーディングを調整し、管理後のキャッシュを利用しやすくする。
• 継続的なポリシー運用：コンテキストが閾値に達してから一括要約するのではなく、ツール応答などの区切りに応じて段階的にキャッシュを整理する。

つまり、完全なテキスト履歴は記録として残しつつ、モデルが現在参照する作業集合だけを小さく保つという発想です。

キャッシュ状態をどう正しく再利用するか

この方法は新たなサービング上の課題も生みます。同じトークン列であっても、過去にどのdropやrepos操作を受けたかによって、最終的なキャッシュ状態が異なる可能性があります。テキストのプレフィックスだけを照合するキャッシュシステムでは、互換性のない状態を誤って再利用する危険があります。

そこで研究チームはAKVサービングエンジンを設計しました。外部には既存のサービングAPIと互換性のあるステートレスなインターフェースを提供しながら、内部ではキャッシュの再利用、追い出し、復元を状態付きで管理します。呼び出し側のインターフェースを大きく変えずに、操作履歴を考慮したキャッシュ管理を実現する設計です。

実験結果と意味

実験では、ツール応答を基準にしたrolling-dropポリシーが、閾値到達後に要約する方式と比べて、平均ベンチマークスコアを1.7ポイント向上させ、平均コストを27.7%削減しました。AKVを組み合わせると、出力スループットは最大1.2倍、P95リクエストレイテンシーは最大1.5倍の改善が報告されています。さらに、再帰的な自己改善でキャッシュポリシーを探索すると、専門家が設計したポリシーよりコストを7.1%追加削減できました。

この研究の示唆は、エージェントの「記憶」を完全なテキスト履歴と同一視する必要はないという点です。KVキャッシュを実行時リソースとして扱い、タスクやツール呼び出しの進行に応じて管理できます。一方で、削除しすぎると後で必要な情報を失う可能性があり、ポリシーはモデルとタスクに合わせて設計する必要があります。AKVは、長時間エージェントの効率化に向けた有望なサービング設計を示すものです。

出典：Hugging Face Daily Papers