Grouped Value Attention、キーのオンデマンド再構成でKV Cacheを圧縮

導入

長いコンテキストを扱う生成モデルでは、モデル重みよりもKV Cacheが推論の制約になることがあります。トークンを1つ生成するたびにKeyとValueがキャッシュへ追加され、その後のステップでは長くなった履歴を読み続けなければなりません。Grouped-query attention（GQA）は複数のQueryヘッドでKey-Valueヘッドを共有し、この負担を抑えました。しかし、各位置についてKeyとValueの両方を保存する構造は残っています。

論文「Grouped Value Attention: Efficient KV Caching via On-Demand Key Reconstruction」は、ヘッド数だけでなく、キャッシュする表現そのものを見直します。

主なポイント

• グループ化したValueを保存。 GVAはValueをグループ単位で保持し、通常の構成のように各位置の完全なContent Keyを永続化しません。
• Content Keyを必要時に再構成。 Value表現とContent Keyを結ぶ学習済み線形写像を用います。推論時にはこの写像をQuery側の変換へ吸収できるため、想定されたデコード経路ではContent Keyを明示的に生成する必要がありません。
• 位置情報は別経路で維持。 内容表現を圧縮すると位置情報が弱くなる可能性があるため、小さく共有された分離型RoPEチャネルを導入し、別のPositional Keyをキャッシュします。実験では16次元の位置チャネル構成も評価されています。
• キャッシュのスカラー数を削減。 検証された構成では、同等条件のGQAと比べて永続キャッシュのスカラー数を約45〜47%減らしました。これは保存量の比較であり、そのまま遅延やスループットの改善を意味するものではありません。
• GQAに近いベンチマーク精度。 350Mパラメータ規模で30BのFineWeb-Eduトークンを用いた実験では、16次元位置構成の5タスク平均精度が44.18でした。GQAは44.36、MLAは43.88でした。

意義と残された課題

GVAのポイントは、KV Cache内の内容情報と位置情報を分離し、キーの一部の計算をQuery側へ移すことです。長い履歴を扱うサービスや多数のリクエストを同時処理する環境では、キャッシュ容量と読み出しトラフィックを減らし、より大きなバッチや長いコンテキストに使える余地を増やす可能性があります。

一方、キャッシュが小さくなっても、自動的に生成が速くなるわけではありません。再構成に必要な計算、メモリアクセス、カーネル融合、バッチサイズ、ハードウェアの利用効率が実際の性能を左右します。著者らは専用のデコードカーネルを開発し、エンドツーエンド推論性能を評価中としています。オープンソース版も予定されていますが、現時点の素材からは具体的な高速化幅を確認できません。

今後は、より大きなモデルや長いコンテキストで精度と容量削減が維持されるか、カーネル実装が低いキャッシュトラフィックを実測上の低遅延へ変換できるかが重要になります。GVAは有望なKV Cache表現ですが、汎用的な高速化手法として結論づけるには、追加のシステム評価が必要です。

出典：Hugging Face Daily Papers