LowRAM、KeyとValueの予測関係を使ってKV Cacheを圧縮

概要

自己回帰生成では、過去のtokenに対応するKeyとValueをKV Cacheに保持し、各ステップで再利用します。コンテキストが長くなるほどキャッシュは線形に膨らむため、必要なメモリが増え、バッチサイズやデコードスループットの制約になります。既存の圧縮手法は、token間や層間の冗長性、特徴次元の冗長性、あるいは低い数値精度を利用するものが中心でした。

LowRAMは、KeyとValueを同じ方法で扱う必要はないという見方を示します。論文では、同一のアテンションヘッド内で、あるtokenのValueが対応するKeyから比較的よく予測できると観察しました。この関係を利用し、テスト時に適応できる非対称な圧縮パイプラインを構成しています。

主な仕組み

• Keyの低ランク化：層をまたぐ相関を利用した低ランク因子分解によって、保存するKeyの量を減らします。
• Valueの予測保存：ヘッドごとに軽量な予測器を学習し、Keyを入力としてValueを再構成します。予測しやすいtokenについては、完全なValueを保存する必要がありません。
• 残差の選択保存：再構成誤差が大きいtokenにだけ補正用の残差を保持します。これにより、予測が難しいデータへ保存容量を集中できます。報告されたValue圧縮率は、ほぼ無損失の品質で4倍です。
• テスト時の適応：予測器はメタラーニングで初期化でき、prefill時の適応コストを抑えられます。一度の追加コストを許容できるマルチクエリ推論では、ゼロから適合させることも可能です。
• 選択的な再構成：デコード時には全KVペアを復元せず、選択された部分だけを再構成します。4倍のKV圧縮、262Kのコンテキストで、ピークデコードスループットが最大5倍向上したと報告されています。

意義と注意点

LowRAMの重要性は、単に圧縮率を高めた点だけではありません。KeyとValueの構造的な違いを利用するため、量子化やキャッシュ削除とは異なる圧縮軸を追加できます。論文の組み合わせ結果では、圧縮率26.9倍でもベースライン性能の94.5%を維持しました。

長文書QAや、同じコンテキストに対して複数の問い合わせを処理するサービスでは、適合済みの予測器を再利用できる可能性があります。一方で、予測器の適合、残差の選択、再構成処理には追加コストがあります。実際の効果は、ハードウェア、コンテキスト長、負荷の特性に左右されるでしょう。それでもLowRAMは、KV Cacheを低精度化したり単純に削除したりするだけでなく、内部表現間の予測関係から圧縮余地を探る方向を示しています。

出典：Hugging Face Daily Papers