AIに5000万トークンの再利用可能な記憶を与えるKV状態の永続化

概要

長い文脈を扱うAIでは、履歴を保存できるかだけでなく、同じ履歴を呼び出すたびに再処理しなければならない点も課題になる。Hugging Face Daily Papersで紹介された研究は、galahad-kvというメモリ層を検証した。モデルがテキストを処理して生成したキー・バリュー状態（KV cache）を、約1万6000トークンのブロック単位で暗号化されたローカルNVMeに保存し、必要なときに読み込む仕組みだ。

これは、恒久的な記憶を持つ新しいモデルを学習するのではなく、推論時の中間状態を保存して再利用するアプローチである。

実験で示されたこと

• 5000万トークンを検証。 vLLMを使い、1枚のNVIDIA H100上でGemma 4 12Bと31Bに実際の公開テキストを処理させた。
• 再計算なしで復元。 ストリームの先頭から5000万トークン地点まで、異なる深さから100ブロックを調べた結果、両モデルですべて保存先から読み戻せた。状態はバイト単位で一致した。
• 速度と電力を改善。 ブロックの読み込みは再計算の約2.8～4.3倍速く、GPUエネルギー使用量は約8.8～12.3倍少なかった。履歴が増えてもGPUメモリはほぼ一定だった。
• 回答品質には差がある。 数百万トークン前に埋め込んだ事実への正答率は、12Bが100問中82問、31Bが98問だった。提示された結果では、どちらも回答を捏造しなかった。

何が変わるのか

同じ文書を何度も参照するエージェントや、長期対話、コードベース解析、文書検索では、毎回すべての履歴を再送してプリフィルする処理が大きな負担になる。KV状態を一度作って保存すれば、後続の要求では元テキストの再処理を省ける。GPU上に全履歴を保持する必要もないため、推論サービスのメモリ設計にも利点がある。

5000万トークンの注意とは違う

この方法は、モデルが5000万トークン全体に一度に注意を向けることを意味しない。実験では記憶ブロックを一つずつ読み込む。どのブロックを選ぶか、そしてモデルがその情報を利用できるかによって回答は変わる。したがって、保存層は検索やオーケストレーションを置き換えるものではなく、それらを支える外部状態ストアと考えるべきだ。

制約も明確である。最初の保存には通常の計算が必要で、保存先は数TB規模になる可能性がある。モデル、ブロック構成、質問の種類によっても結果は変わる。評価では単純な文脈長だけでなく、初回書き込み、容量、読み込み遅延、電力、検索精度、回答品質をまとめて見る必要がある。

出典：Hugging Face Daily Papers