Every Cache Entry Earns Its Place: Global Allocation of Resolution and Coverage for KV Cache Compression

あなたは、42,000ページ目に登場するキャラクターに関するたった一つの質問に答えるために、10万ページの膨大な長編小説を記憶しようとしていると想像してください。あなたの脳はスーパーコンピューターですが、そのデスクは非常に小さく、高価なもので、一度に数ページしか開いておくことができません。次の文章を読むためにページをめくるたびに、あなたはデスクを整理し、新しいページのためのスペースを作るために古いページを床に投げ捨てなければなりません。これは、現代の「大規模言語モデル（LLM）」が長い物語や文書を読み取る際の仕組みそのものです。彼らは、これまで読んできた内容を再計算する手間を省くために、「Key-Valueキャッシュ」（洗練されたメモリ用デスク）と呼ばれるものに情報を保持しています。しかし、物語が進むにつれて、このデスクは混雑し、動作を遅らせ、メモリを使い果たしてしまいます。科学者たちは、この問題を解決するために、「重要度の低い」ページを捨てる（トークン除去）か、似たページ同士を一つの要約シートに貼り合わせる（KV結合）という方法を試みてきました。しかし、これらの古い手法は硬直したルールに基づいています。つまり、あなたがどのような特定の質問をしているかを見ることなく、どのページを残すか、あるいはどのように結合するかを事前に決定してしまうのです。これでは、物語の内容が変わったときに、最も重要な部分にリソースを集中させることは容易ではありません。

この論文は、そのメモリデスクを管理するための、よりスマートで新しい方法であるGraceKVを紹介しています。GraceKVは、メモリを固定されたルールブックに従うのではなく、最も必要とされる場所に自由に使える「柔軟な予算」として扱います。メモリ・トークン（コインのようなもの）の数を決まった数だけ持っていると想像してください。古い手法では、「すべての章の10%を保持しなければならない」とか「10ページごとに結合しなければならない」といったルールを課します。しかし、GraceKVはこう問いかけます。「この特定の質問に対して、最も価値のある情報はどこにあるのか？」と。GraceKVは、物語のあらゆる部分に対して、特別なツリー状のマップを構築します。ツリーの頂点にある単一の「要約コイン」は、巨大なテキストの塊を広くカバーします（広域カバレッジ）。もし物語の特定の箇所で展開が面白くなったり、複雑になったりした場合、GraceKVはその「要約コイン」を分割して、その極めて小さなセクションに対して、より詳細で高解像度なコインを購入することができます（局所解像度）。GraceKVは、物語全体にわたって、広範な要約を維持することと、詳細なスニペットを保持することのどちらに価値があるかを常に比較し、最高の答えを出すための組み合わせに予算を投じます。GraceKVは、メモリを最大128倍まで圧縮しても正確に回答できることを示しており、固定されたルールを使用する他の手法をしばしば凌駕します。それはまるで、単に本のリストに従って保管するものを選ぶだけでなく、必要な一冊が目の前に来るように、リアルタイムで図書館全体を再配置する司書がいるようなものです。たとえ、それによって他のすべてを動かすことになったとしても。

問題点：「長すぎて覚えられない」ジレンマ

大規模言語モデルは、ほとんど何でも読むことができる優秀な学生のようなものですが、短期記憶に問題を抱えています。長い文書を読んで質問に答えるとき、彼らはこれまでに見たすべての言葉の「Key」と「Value」（誰が、何を、どこで、なぜ行ったか）を覚えておく必要があります。このメモリはKVキャッシュと呼ばれ、テキストの長さに比例して線形に増加します。もし10万語の小説を入力すれば、それらのキーと値を保持するために必要なメモリは膨大になり、コンピュータのRAMを占有し、次の単語を生成するプロセスを遅らせてしまいます。

これを修正するために、研究者は主に2つのテクニックを試してきました：

1. トークン除去（Token Eviction）： 「退屈な」言葉を捨て、重要なものだけを残すこと。これは、本の中から関連性がなさそうなページを削除するようなものです。
2. KV結合（KV Merging）： 似た言葉同士を一つの「要約」エントリーに貼り合わせること。これは、物語の10ページ分を取り出し、その要旨を捉えた1つの段落に置き換えるようなものです。

これらの古いテクニックの問題は、それらが硬直していることです。通常、「最後の100語を残す」とか「5語ごとに結合する」といった、あらかじめ設定されたルールに従います。そのため、あなたがしている特定の質問にはうまく適応できません。退屈に見える言葉が答えの鍵であることもあれば、巨大なテキストの塊が無関係であることもあります。古い手法は、カバレッジ（物語全体を覚えていること）と解像度（細かなディテールを覚えていること）のバランスを取るのに苦労します。なぜなら、メモリ予算を自由に動かすことができないからです。

解決策：GraceKVの「グローバル・バジェット」

著者らは、メモリ圧縮をルールに従うゲームではなく、グローバルなリソース配分問題として扱うシステムであるGraceKVを提案しています。これは、限られた電力予算を管理するスマートな都市計画者のようなものです。すべての近隣地域に同じ量の電力を与えるのではなく、計画者は「今まさに」どこに電力が必要かを判断します。

GraceKVは、主に3つのステップで動作します：

1. ツリー・マップの構築：
   まず、GraceKVは、単なるランダムな切り出しではなく、意味の変化に基づいて長い物語を「スロット」（テキストの塊）に分割します。AIの脳の各レイヤーおよび各アテンション・ヘッドごとに、プロトタイプ・ツリーを構築します。

   • ツリーの**ルート（根）**は、巨大なテキストの塊の単一で粗い要約です。
   • **ブランチ（枝）**は、その塊をより小さく詳細な断片へと分割できます。
   • **リーフ（葉）**は、元の正確な言葉です。
     このツリーにより、システムは、広範な概要から単一の精密な言葉に至るまで、異なるレベルの詳細度で同じテキストを表現できるようになります。

2. 価値のフロー（お宝探し）：
   システムは、現在の質問に対してテキストのどの部分が実際に有用であるかを判断します。単に質問を直接見るだけでなく、情報がテキスト内をどのように流れているか（手がかりの跡を追う探偵のように）を追跡します。もしある言葉が質問の中で言及されていたり、他の重要な言葉と結びついていたりする場合、その言葉には高い「価値スコア」が与えられます。このスコアは、その物語のどの部分にどれだけの「お宝」が隠されているかをシステムに伝えます。

3. 予算のフロー（コインの使い道）：
   ここからが魔法です。GraceKVには、メモリ・スロット（コイン）の固定された予算があります。システムは、物語全体におけるあらゆる可能なアクションを検討します：

   • 追加（Add）： コインを使い、カバーされていない新しいテキストの塊を粗い要約でカバーする（広域カバレッジの拡大）。
   • 分割（Split）： コインを使い、粗い要約をより小さく詳細な断片へと分解する（解像度の向上）。

   あらゆる「追加」または「分割」のアクションは、単一のグローバルなキューの中で競合します。システムは、各アクションの「ユーティリティ（コインあたりの価値）」を計算します。もし特定の非常に細かい言葉が答えにとって極めて重要であれば、その言葉の要約を「分割」することが、非常に高いユーティリティを持つ可能性があります。もし段落全体が退屈であれば、そのための粗い要約を「追加」することが、コインの最善の使い道かもしれません。システムは、予算が尽きるまで、最も価値の高いアクションを貪欲に選択していきます。

   また、**シングルトン・フロア（Singleton Floor）**と呼ばれるセーフティネットもあります。貪欲なアルゴリズムでは、ステップを踏んでいくコストが高すぎるために、超重要な言葉を見逃してしまうことがあります。GraceKVは予算の一部を確保し、いくつかの高価値な言葉をそのままの形で保持することを保証し、重要な詳細が失われないようにしています。

実績

著者らは、長い文書からの質問回答、物語の要約、巨大なデータセットからの特定の事実の検索など、さまざまなタスクでGraceKVをテストしました。彼らは、4倍から128倍までの異なる圧縮レベルにおいて、既存の最高水準の手法（H2O、SnapKV、PyramidKVなど）と比較しました。

• パフォーマンス： GraceKVは、32の異なる設定のうち24においてトップの結果を出しました。メモリ予算が極めて厳しい状況（128倍圧縮）であっても、一貫して1位または2位を維持しました。
• 堅牢性： ある種のタスクでは優れているが別のタスクでは失敗するという他の手法とは異なり、GraceKVはすべてのタスクにおいて強力な性能を維持しました。それは「広域カバレッジ」を必要とするタスク（要約など）と、「精密な検索」を必要とするタスク（特定の名前の発見など）の両方を等しくこなしました。
• 効率性： メモリを圧縮することで、GraceKVは必要なメモリ量を大幅に削減し（フルメモリより最大92%削減）、特に非常に長いコンテキストにおいて、コンピュータのテキスト生成速度を向上させました。
• 再学習不要： GraceKVの素晴らしい点の一つは、再学習を必要としないことです。テキストと質問を処理の最中に分析して動作するため、既存のあらゆるモデルに対してプラグアンドプレイのソリューションとして機能します。

なぜ重要なのか

この論文は、長いコンテキストを持つAIの未来は、何を保持し何を捨てるかという単一の「完璧なルール」を見つけることではないと示唆しています。むしろ、それは柔軟性にあります。メモリを、広範なカバレッジと微細な詳細のバランスを取るために動的に割り当て可能な、共有されたグローバル・リソースとして扱うことで、複雑で長い物語を理解する能力を失うことなく、AIモデルをより効率的にすることができます。GraceKVは、スマートで適応的なメモリ管理のアプローチが、硬直した事前設定のルールを凌駕できることを証明しており、AIが圧倒されることなく図書館全体を読み解くための道を切り拓いています。