BeaconKV：ビーコンクエリで長時間推論モデルのKVキャッシュを圧縮

導入

大規模な推論モデルは、長いChain-of-Thought（CoT）を生成することで複雑な問題に対応します。しかし、自己回帰生成では各トークンに対応するKey-Value（KV）キャッシュが蓄積されるため、推論系列が長くなるほどGPUメモリの消費も増えます。長い思考過程を扱う場面では、キャッシュ容量がコンテキスト長や同時実行数を制限する要因になり得ます。

既存のKVキャッシュ圧縮手法の多くは、直近に生成されたクエリを手がかりに、今後重要になる過去トークンを推定します。短い生成では合理的な近似でも、長期的な推論では不十分な場合があります。モデルは推論の初期に作った計画や条件を、何ステップも後になって再確認することがあるためです。

BeaconKVはこの問題に対し、すべてのクエリ履歴を保持する代わりに、繰り返し現れる問い合わせパターンを代表する「ビーコンクエリ」を保存します。追加学習を必要としない、キャッシュ管理側の最適化です。

主なポイント

• 遠い文脈を再訪するトークンに着目。 論文は、過去の遠いコンテキストへ再び注意を向けるデコード段階をThought Revisiting Tokens（TRT）と呼びます。初期の解法計画や中間結果の検証などが、その例に当たります。
• 直近クエリだけでは見落としがある。 最新のクエリを中心にキャッシュを選ぶと、現在の局所的な文脈は残せても、後で必要になる古いKVペアを削除する可能性があります。
• 再訪に関係するクエリは集約できる。 TRTに対応するクエリは、埋め込み空間で少数の類似グループを形成する傾向があると分析されています。各グループの代表を保存すれば、長い履歴を小さな状態で近似できます。
• 直近情報と大域的な信号を併用。 BeaconKVは、現在の注意パターンを捉える直近クエリと、遠い文脈の再利用を示すビーコンクエリを組み合わせて、保持すべきKVペアを選別します。
• 再学習なしで導入できる。 提供された概要によれば、4つのオープンソース推論モデルと複数の推論ベンチマークで、最大約5.8倍のメモリ削減、4.3倍超のスループット向上が報告されています。完全なキャッシュに近い精度も維持したとされています。

意義と今後

BeaconKVの重要性は、推論モデルのキャッシュ圧縮を単純な「近いものを残し、遠いものを捨てる」問題として扱わない点にあります。通常の文章生成では新しい文脈ほど重要になりやすい一方、長い推論では初期の計画や制約が後の判断を支えることがあります。したがって、キャッシュの重要度を推定するには、局所的な新しさだけでなく、再訪されやすい大域的パターンも必要です。

クエリ履歴を代表ベクトルへ圧縮して利用する設計は、推論モデルを長い系列で動かすサービング基盤に適しています。GPUメモリの圧力を下げれば、より長い推論、同一ハードウェアでの高い同時実行数、あるいは運用コストの削減につながる可能性があります。

ただし、提供資料にはモデル別・ベンチマーク別の詳細な比較や、具体的な圧縮設定までは示されていません。アーキテクチャやタスクが変わった場合にも同じ効果が得られるかは、論文本文とコードによる追加確認が必要です。直近性と長期的な再訪パターンを組み合わせるという発想は、推論向けKV最適化の有力な方向性と言えるでしょう。

出典：Hugging Face Daily Papers