🤖 machine learning

A JoLT for the KV Cache: Near-Lossless KV Cache Compression via Joint Tucker and JL-Residual Allocation for LLMs

JoLTは、トークン軸と特徴量軸に対して部分的なタッカー分解を適用し、破棄された情報をジョンソン＝リンデンストラウス回転を用いた低ビット残差によって復元することで、大規模言語モデルにおけるKVキャッシュのほぼロスレスな2〜3倍の圧縮を実現しており、これらすべては、パープレキシティおよびダウンストリームタスクにおけるベースライン性能を維持するために、統一されたバイト予算の下で最適化されている。

原著者： Rahul Krishnan, Volker Schulz

公開日 2026-07-15

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Rahul Krishnan, Volker Schulz

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたは、巨大で超スマートなロボットの脳（大規模言語モデル）を動かし、物語を語ろうとしているところだと想像してください。物語を継続するために、ロボットは現在の会話でこれまでに発したすべての単語を記憶しておく必要があります。ロボットはこの記憶を、「KVキャッシュ」と呼ばれる特別な「ノート」に保管しています。

問題は、物語が長くなるにつれて、このノートが巨大化していくことです。ノートが大きくなりすぎると、ロボットのメモリを食いつぶし始め、動作を遅くしてしまいます。それはまるで、マラソンを走っている最中にバックパックの中に図書館を詰め込もうとするようなもので、最終的には動けなくなってしまいます。

科学者たちは以前、このノートを縮小しようと試みてきました。ある者は全体を小さな箱に押し込めようとし（量子化）、またある者は行や列のパターンを見つけることでページを要約しようとしました（低ランク手法）。しかし、この論文の著者である Rahul Krishland と Volker Schulz は、他の手法が見落としていたことに気づきました。それは、ノートは単なる平らな紙の束ではないということです。ノートは、ヘッド（思考の多様な方法）、トークン（単語）、特徴量（詳細）という3つの異なる側面を持つ、3次元のブロックなのです。

彼らは、これら2つの側面には、簡単に押しつぶせる退屈で繰り返しの多い内容が詰まっている一方で、もう一方の側面はユニークであり、押しつぶすとロボットの知能を損なってしまうことを発見しました。

大発見：「JoLT」手法

チームは、JoLT（Joint Tucker and JL-residual allocation）と呼ばれる新しいトリックを考案しました。これは、あなたのバックパックのための超スマートなパッキング・サービスのようです。

スマートな絞り込み（部分的Tucker）: 3次元ブロック全体を押しつぶそうとするのではなく、JoLTはデータを見てこう判断します。「よし、『ヘッド』と『レイヤー』はユニークで貴重だから、そのままにしておこう。でも、『トークン』と『特徴量』は無駄な部分が多いな」。つまり、これら特定の2つの側面だけを圧縮するのです。それは、巨大でふわふわした枕を取り出し、頑丈な端の部分はそのままにして、真ん中の空気だけを抜くようなものです。
セーフティネット（JL-residual）: 枕を押しつぶすと、いくらかの空気が逃げてしまいます。そのままにしておくと、枕は平らで使い物にならなくなります。JoLTはこの「逃げた空気」（失われた情報）を捉え、JL-residualと呼ばれる、非常に効率的な小さなセーフティネットに保存します。このネットはパッキングが非常に得意なので、わずかなビット数で足りない詳細を保持することができます。
完璧なバランス（ラグランジュ・デュアル）: ここに魔法があります。ロボットには、使用できるスペースの厳格な予算（例えば1バイト）があります。JoLTは数学的な「スマート・アロケーター（賢い配分器）」を使用して、枕をどれくらい押しつぶすか、そしてセーフティネットにどれだけのスペースを与えるべきかを正確に決定します。Jo-LTは、メモリの特定のパーツ（「Key」）は圧縮しやすい一方で、別のパーツ（「Value」）は頑固でより多くのセーフティネット・スペースを必要とすることを理解しています。そして、最高の結果を得るために予算をダイナミックに動かします。

彼らが証明したもの（そして証明できなかったもの）

著者らは、この手法を2つの有名なロボットの脳、Mistral-7B（「Grouped-Query」スタイルを使用）と LLaMA-2-13B（「Multi-Head」スタイルを使用）でテストしました。

「フリーゾーン」: 彼らは、メモリを2〜3倍（2–3×）縮小しても、ロボットのパフォーマンスが全く低下しないスイートスポットを見つけました。これは「ニアロスレス（ほぼ損失なし）」です。
- GSM8K（数学の問題）や RULER（テキストの干し草の山の中から針を探すタスク）のようなテストにおいて、圧縮されたロボットは、統計的な誤差の範囲内で、圧縮されていないものと全く同じスコアを記録しました。
- メモリを再構成する際の誤差は極めて小さく、Keyについては約0.009、Valueについては約0.006でした。これは、4ビット量子化やクロスレイヤーSVDといった従来の手法よりも、約10倍（1桁）優れた結果です。
「崖（クリフ）」: 彼らは限界も発見しました。メモリを圧縮しすぎると（3倍を超えると）、状況は混乱します。
- Mistralのロボットは、より多く圧縮しても性能が緩やかに低下する、緩やかな劣化を示しました。
- しかし、LLaMAのロボットは、4倍から5倍の圧縮の間で「崖」に突き当たりました。その性能は急落し、パープレキシティ（予測の悪さを示す指標）が5.39から9.07へと大きく跳ね上がりました（これは、予測能力が大幅に悪化したことを意味します）。

高速版：FlashJoLT

完璧な絞り込みを計算するには時間がかかります。これを解決するために、彼らはFlashJoLTを作成しました。毎回重い数学的計算を完璧に行う代わりに、主要なパターンを素早く推測する「ランダム化された」ショートカットを使用します。

結果: 圧縮スピードは、通常の（遅い）完璧なバージョンと比較して5〜13倍速いですが、品質は全く同じに保たれています。

彼らが否定したもの

論文は、この特定の問題において何がうまくいかないのかについても明確に述べています。

すべてを押しつぶすこと: 3つの側面すべて（ヘッド、トークン、特徴量）を圧縮しようとすることは、悪いアイデアです。「ヘッド」と「レイヤー」はあまりにユニークであるため、それらを押しつぶすとロボットの脳を傷つけてしまいます。
固定ビット量子化: すべての数字に対して一律にビット数を下げること（例えば、すべてを4ビットに強制すること）では、この「スイートスポット」である2〜3倍の圧縮には到達できません。圧縮が足りないか、あるいは品質を失いすぎてしまいます。
一律の扱い: 「Key」と「Value」を同じように扱うことはできません。「Value」は圧縮するのがずっと難しいため（2〜3倍難しい）、異なる予算配分が必要です。

結論

著者らはこれを実際のハードウェア（A100 GPU）で測定し、JoNTがロボットの知能を損なうことなく、メモリを2〜3倍に縮小するニアロスレスな方法であることを明らかにしました。

しかし、彼らはこれがあらゆるものに対する魔法の杖ではないことも慎重に述べています。

「フリーゾーン」（2〜3倍）では非常にうまく機能しますが、特定のロボット（LLaMAなど）でさらに強く押し進めると、品質が急激に低下します。
メモリの「保存」自体は小さくなりますが、ロボットが話すたびに、そのメモリを「解凍（展開）」するための計算を行う必要があります。これを実用的なものにするためには、圧縮されたメモリを、展開することなく直接読み取れる特殊なコンピューターチップ（融合カーネル/fused kernels）をエンジニアが構築する必要があると彼らは示唆しています。

要約すると、JoLTは、長い会話のために膨大なスペースを節約する、数学的に高度なパッキング技術ですが、限界があり、フルスピードで動作させるためにはハードウェアからの助けを必要とします。

テクニカル・サマリー：KVキャッシュへのJOLT（劇的な改善）

問題提起

自己回帰型トランスフォーマーの推論において、Key-Value (KV) キャッシュは支配的なメモリコストを占め、その規模は $O(\text{batch} \times T \times L \times n_h \times d_h)$ でスケールする。長いコンテキストや大きなバッチサイズにおいては、モデルの重みではなく、KVキャッシュが主要なメモリのボトルネックとなり、スループットを制限する。既存の圧縮手法には2つの系統があるが、いずれもキャッシュの構造的冗長性を十分に活用できていない：

低ランク手法（Low-rank methods）: これらは2次元のスライス（例：ヘッドごとの行列やクロスレイヤー・ブロック）を因子分解するが、単一の軸の冗長性にのみ着目し、他の軸を無視している。
量子化手法（Quantization methods）: すべてのエントリのビット幅を削減するが、固定されたビットレート（例：4ビット）に制限されるため、品質が維持されやすい中間的な圧縮率（2〜3倍）を実現できない「圧縮の底（compression floor）」が生じている。

本研究の動機となった核心的な観察は、あるレイヤーにおけるKVキャッシュは第3次テンソル（ $n_h \times T \times d_h$ ）であり、各軸が極めて異なる量の冗長性を保持しているということである。具体的には、ヘッドおよびレイヤーの軸はほとんど圧縮不可能であるが、トークンおよび特徴量の軸は大部分の圧縮可能な構造を保持している。さらに、経験的に、値（Values）はキー（Keys）と比較してスペクトル特性が異なるため、圧縮が2〜3倍難しい。

手法：JoLT

提案手法である JoLT（Joint Tucker and JL-residual allocation）は、キャッシュを第3次テンソルとして扱い、ターゲットを絞った圧縮戦略を適用することで、これらの課題に対処する。

1. 部分的タッカー分解（Partial Tucker Decomposition）

すべての軸を圧縮する完全なテンソル分解とは異なり、JoLTは部分的タッカー分解を適用する。冗長性を持つトークンおよび特徴量の軸のみを切り詰め、圧縮不可能なヘッドおよびレイヤーの軸はそのまま保持する。これは、逐次的高次特異値分解（ST-HOSVD）を用いて実現される。

動機: ヘッドやレイヤーの軸を圧縮すると大きな誤差が生じるが、トークンと特徴量の軸を切り詰めることで、冗長性の大部分を捉えることができる。

2. JL回転残差（JL-Rotated Residual）

純粋な低ランク・バックボーンだけでは、特にスペクトルが平坦で切り捨てによってエネルギーが失われやすい「値（Values）」において、ニアロスレス（ほぼ無損失）の忠実度を得るには不十分である。JoLTは、ジョンソン＝リンデンシュトラウス（JL）回転を用いた低ビット残差を用いて、この失われたエネルギーを復元する：

切り捨てられた残差 $R = X - \tilde{X}$ を計算する。
残差に対してランダムなJL回転を適用する。
回転された残差を $b \in \{0, 2, 4, 8\}$ ビットで一様量子化する。
デコード時には回転を反転させる。この残差は、損失のあるバックボーンをニアロスレスなものへと変換するレバーとして機能し、特に「値」において価値を発揮する。

3. 結合ラグランジュ配分（Joint Lagrangian Allocation）

本手法は、単一のラグランジュ双対を用いて、2つのステージ（タッカーのランクと残差のビット幅）および異なるターゲット（Keys vs. Values、レイヤーグループ）の間でリソースを共同で配分する。

目的: グローバルなバイト予算の制約下で、再構成誤差を最小化する。
メカニズム: 双対変数は、ランクと残差の間、およびKeyとValueの間でビットを動的に移動させ、圧縮性の非対称性（例：Valueの残差に多くのビットを割り当て、Keyの残差ビットを減らすなど）を最大限に活用することを可能にする。
粒度: 配分はレイヤーグループごとに行われ、KeyとValueについても個別に実行される。

4. FlashJoLT（高速版）

長いコンテキストにおける正確なSVDの計算コストに対処するため、FlashJoLTは正確なトークンモードSVDをランダム化低ランクSVDに置き換える。

上位 $q_{cap}$ 個のトークンモード方向のみを計算する。
**テール質量計上（tail-mass accounting）**を採用し、省略されたスペクトルの裾（テール）を合成的な残差特異値として加算することで、特定のランクによって捉えられたエネルギーが過小評価されるのを防ぐ。
このバリアントは、品質の低下をほとんど伴わずに、圧縮時間を5〜13倍高速化する。

主な貢献

スペクトル分析: ヘッドとレイヤーの軸は圧縮不可能であり、トークンと特徴量の軸は冗長性を保持していること、また値はキーよりも大幅に圧縮が難しいということを特定した経験的研究。
JoLTアルゴリズム: 部分的タッカー分解とJL回転残差を組み合わせ、単一のラグランジュ双対によってランクとビットを共同配分する新しいコンプレッサ。
FlashJoLT: 圧縮時間を劇的に短縮しながら品質を維持する、ランダム化SVDバリアント。
実証的検証: GQA（Mistral-7B）およびMHA（LLaMA-2-13B）モデルを用いた包括的な調査により、「フリーゾーン（free zone）」におけるニアロスレスな圧縮を確立。

結果

「フリーゾーン」（2〜3倍の圧縮）

本論文は、モデルの品質が未圧縮のベースラインの統計的ノイズ範囲内に収まる2〜3倍の圧縮率を「フリーゾーン」として特定している：

パープレキシティ: Mistral-7B (GQA) と LLaMA-2-13B (MHA) の両方において、コンテキスト長1024までの2〜3倍の圧縮におけるパープレキシティのドリフトは無視できる程度である。
再構成誤差: 2倍の圧縮において、JoLTは相対フロベニウス誤差 0.009 (Keys) および 0.006 (Values) を達成した。これは、4ビット量子化やクロスレイヤーSVDのベースラインよりも約1桁低い。
ダウンストリームタスク:
- GSM8K: 圧縮されたモデル（2〜3倍）は、ベースラインの95%信頼区間内に留まっている。
- RULER (Needle-in-a-Haystack): Mistralは2〜3倍の圧縮でも100%の検索精度を維持し、LLaMAはベースラインの天井値に一致している。

アーキテクチャによる乖離

フリーゾーンを超えると、2つのアーキテクチャは乖離する：

Mistral (GQA): 緩やかに劣化し、4倍を超える整数比ステップごとに約4%のパープレキシティ低下が発生する。
LLaMA (MHA): 4倍から5倍の間で急激に劣化し、4倍でパープレキシティが約27%、5倍で約68%増加する。LLaMAの推論タスク（GSM8K）は、高い圧縮比で著しく崩壊する。

効率性

メモリ: 永続的なKVフットプリントは圧縮率に応じて線形に縮小する（例：Mistralのctx-2048において、2倍圧縮で268.4 MBから134.6 MBへ）。ピークプロセスメモリは重みやアクティベーションによって制限される。
速度: FlashJoLTは、正確なバックボーンと比較して、圧縮時間を5〜13倍高速化する。

重要性と主張

本論文は、JoNTがKVキャッシュのニアロスレスな2〜3倍の圧縮を提供し、モデルの品質を損なうことなく、長コンテキスト推論におけるメモリのボトルネックを事実上取り除くことができると主張している。キャッシュをテンソルとして扱い、KeyとValueの間の特定のスペクトル非対称性を活用することで、JoLTは既存の低ランク手法や量子化手法を凌駕する。

著者らは、本手法を長コンテキスト・サービングのコストを削減するための直接的な手段として位置づけている。また、以下の2つの主要な境界についても言及している：

アーキテクチャの限界: 本手法はGQAモデルには堅牢であるが、MHAモデルでは4倍を超えると急激な劣化が見られる。これは、2〜3倍の帯域が安全かつモデルに依存しない推奨値であることを示唆している。
デプロイメントコスト: ストレージは削減されるが、現在のデコーディングには各ステップで因子からキャッシュを再構成する必要がある。著者らは、推論のパリティ（同等性）を実現し、本手法のポテンシャルを完全に引き出すための次の重要なステップとして、**融合された再構成・アテンション・カーネル（fused reconstruction-and-attention kernel）**の開発を挙げている。

大発見： 「JoLT」手法