技術要約: SEKV – 階層的意味記憶を用いた解像度適応型KVキャッシュ

問題提起

大規模言語モデル（LLM）は、文書解析、マルチターン対話、エージェントによるプランニングなどの長文脈推論を必要とするシナリオでますます導入されています。しかし、デコーディングを加速させるために中間状態を保存するKey-Value（KV）キャッシュは、シーケンス長に対して線形に増加します。128Kトークルを超えるコンテキストに対してフルKVキャッシュを保持することは、GPUメモリの観点から極めて高価であり、圧縮なしでは長文脈の推論を不可能にすることが多々あります。

既存の圧縮戦略は、効率性と忠実度の間で根本的なトレードオフに直面しています：

1. トークン除去（Token Eviction）: アテンション統計量に基づいてトークンを永久に破棄する手法（H2OやSnapKVなど）は、一貫した証拠を断片化したり、生成プロセスの中で後になって重要となる情報を取り返しのつかない形で失ったりするリスクがあります。
2. 静的な意味的グルーピング（Static Semantic Grouping）: トークンをプリフィル時に意味的なチャンクや文章へと圧縮する手法（ChunkKVやSentenceKVなど）は、圧縮解像度を早い段階で固定してしまいます。これらは、圧縮されたスパンがデコーディング中に再び関連性を持った際に、微細なトークンレベルの詳細を復元する能力を欠いています。

核心となる課題は、厳格なメモリ制約の下で、情報の不可逆的な損失を回避しつつ、クエリの関連性に基づいて保存されるコンテキストの解像度を動的に調整できるメモリ組織を設計することです。

手法: SEKV

著者らは、GPUとCPUにまたがる階層的な意味的メモリとしてコンテキストを整理する、解像度適応型の意味的KVキャッシュであるSEKV（Semantic KV Cache）を提案しています。SEKVは、コンテキストを可変解像度で保存し、デコーディング中に適切な詳細レベルへと動的に「ズームイン」することで、不可逆的なトークン除去を回避します。

1. エントロピー誘導型スパン分割

固定サイズのチャンキングではなく、SEKVは**トークンのサプライザル（驚き度）**をゼロコストの境界信号として使用し、入力コンテキストを意味的に一貫したスパンに分割します。

• メカニズム: プリフィルのフォワードパス中に、各トークンのサプライザル $H_t = -\log p(x_t | x_{<t})$ が計算されます。高いサプライザルを示すトークンは、意味的な遷移（トピックの転換、エンティティの導入など）を示すことが多いため、これらをスパンの境界として特定します。
• アンカー・トークン: これらの高サプライザルな境界トークンは、「アンカー・トークン」としてGPU上にフル解像度で保持されます。これらは、後のルーティングと再構成に不可ло欠な、濃縮された意味情報を持っています。
• スパン: 隣接するアンカー間のトークンは、一つの意味的スパン $S$ を形成します。

2. 二重解像度スパン表現とメモリ階ラシー

各意味的スパンは、GPU–CPU階層にわたって、二つの解像度で同時に表現されます。

• GPU（粗いルーティング）: 各スパンは、軽量なサマリー・ベクトル（投影されたキー・ベクトルの、サプライザルで重み付けされた正規化平均）によって表現されます。このサマリーはGPU上に存在し、現在のクエリに対してどのスパンが関連しているかを判断するための、効率的かつ粗い粒度のルーティングに使用されます。
• CPU（微細な再構成）: スパンの完全なトークンレベルの構造は、低ランクSVD基底を用いてCPU上にコンパクトにエンコードされます。軽量な予測器が、各スパンの情報の密度に基づいて、特異値に対して個別のソフトゲートを割り当て、ランク予算 $r_S$ を適応的に保持します。これにより、フル行列をGPUに保存することなく、トークンレベルのKVペアを忠実に、オンデマンドで再構成することが可能になります。

3. 学習されたズームイン・メカニズム

デコーディング中、SEKVはどのスパンを拡張すべきかを動的に選択する学習されたメカニズムを採用します。

• ルーティング: 現在のクエリは、スパン・サマリーと同じルーティング空間に投影されます。シグモイド関数による関連性ゲートが、そのスパンが関連している確率を計算します。
• 拡張: 学習された、ヘッドごと・レイヤーごとの閾値が、スパンを拡張するかどうかを決定します。選択された場合、システムはCPUから対応する低ランクSVD基底を非同期にフェッチし、トークンレベルのKVペアを再構成して、アテンション計算に統合します。
• 混合解像度アテンション: アテンション・メカニズムは、混合解像度のメモリに対して動作します。アンカー・トークンはフルKVペアを寄与させ、未拡張のスパンは粗いサマリー・ベクトル（サイズ補正付き）を寄与させ、拡張されたスパンは再構成されたトークンレベルのペアを寄与させます。
• 学習: ベースとなるLLMは凍結されたままです。軽量なルーティング投影、ズームイン閾値、およびランク・ゲート予測器のみが、蒸留損失（フルKVの教師モデルとの整合）、ズーム監視、再構成損失、および予算ペナルチを用いて学習されます。

主な貢献

1. SEKVアーキテクチャ: コンテキストを複数の解像度で保存し、クエリの関連性に基づいて動的に詳細度を調整する、解像度適応型の意味的KVキャッシュ。これにより、不可逆的なトークン除去を回避します。
2. エントロピー誘導型セグメンテーション: トークンのサプライザルを利用して一貫した意味的スパンを特定し、信頼性の高い検索のために高サプライザルなアンカー・トークンをフル解像度で保持する手法。
3. 二重表現構造: 粗いルーティングのためのGPU常駐型スパン・サマリーと、CPU常駐型の低ランク基底による近似的なトークンレベルの再構成を組み合わせ、クエリに関連するスパンのみを選択的に拡張する学習されたズームイン・メカニズムを備えた、新しいKV構造。

実験結果

SEKVは、4つのベンチマーク（LongBench, RULER, InfiniteBench, Needle-in-a-Haystack）および5つのバックボーン・モデル（Llama-3のバリアント、Mistral、Qwenを含む）を用いて評価されました。

• パフォーマンス: SEKVは、最強のセマンティック圧縮ベースライン（SentenceKV）を、すべてのベンチマークとモデルにおいて平均で**5.9%**上回りました。20個のベンチマーク・モデル比較のうち、20個すべてで最高の結果を達成しました。
• 検索の忠実度: Needle-in-a-Haystack (NIAH) ベンチマークにおいて、SEKVは、トークン除去法や静的な意味的手法と比較して、様々なコンテキスト長とニードルの深さに対して優れた検索の安定性を示しました。
• メモリ効率: 128Kのコンテキスト長において、SEKVはフルKVキャッシュと比較してGPUメモリ消費量を**53.3%**削減しながら、非圧縮のベースラインに近い性能を維持しました。
• 効率性: SEKVは競争力のあるレイテンシとスループットを実現しており、メモリ帯域幅の圧力を軽減するため、しばしばフルKVキャッシュよりも優れたレイテンシを示しますが、精度の面ではアグレッシブな除去手法を大幅に上回ります。
• パラメータ効率: 本手法は、ベースモデルを凍結したまま、0.05%未満（Llama-3-8Bに対して約4.3Mパラメータ）の学習可能パラメータしか追加しません。

意義と主張

本論文は、効率的な長文脈推論には、静的な圧縮から動的な解像度回復への転換が必要であると主張しています。著者らは、静的な圧縮は粗いコンテキストを保持するものの、生成プロセスの後半で微細な証拠が必要になった際に、それを取得できないことを指摘しています。SEKVは、コンテキストをエントロピー誘導型の意味的スパンとして整理し、オンデマンドで関連する情報を「ズームイン」することを可能にする階層的メモリシステムを利用することで、この問題に対処します。

本研究の意義は、以下のことを示した点にあります：

1. 不可逆的な除去は必ずしも必要ではない: 低ランク表現をCPUにオフロードすることで、高精度な長文脈推論が可能になること。
2. 適応的な解像度が不可欠である: ズームイン・メカニズムを通じて必要な場所にのみ微細な解像度を割り当てることで、静的なポリシーよりも効果的にメモリ制約と検索の忠実度のバランスを取れること。
3. 意味的構造が圧縮を助ける: エントロピー誘導型のセグメンテーションを用いることで、任意のトークンブロックよりも圧縮しやすい単位を作成でき、効率的な低ランク近似が可能になること。

著者らは限界についても認めており、有効性は（高度に断片化された談話に対して苦戦する可能性がある）意味的スパン構築の質に依存すること、および、このアプローチがCPU-GPU間の帯域幅への依存性を導入し、敵対的なクエリ・シナリオにおいてレイテンシに影響を与える可能性があることを述べています。しかしながら、SEKVは、スケーラブルで高忠実度な長文脈LLM推論のための堅牢なフレームワークを提供すると結論付けています。