技術的概要：TurboQuant に着想を得た KV キャッシュ量子化の統計的推論と品質指標

問題定義

KV（キー・バリュー）キャッシュは、トランスフォーマ推論における主要なメモリボトルネックであり、コンテキスト長に比例して線形に増加する。量子化による圧縮は不可欠であるが、標準的なスカラー量子化（MSE）は、アテンション機構に必要な幾何学的性質を保持できないことが多い。TurboQuant は、ランダム回転＋スカラー量子化＋バイアス補正のための 1 ビット QJL スケッチという 2 段階のパイプラインを提案しているが、キー（K）キャッシュとバリュー（V）キャッシュ間のビット配分の最適解は未だ不明である。具体的には、（バリューに加えて）キーにも QJL 残差補正を適用することが、異なるビット予算やデータ分布（例：低ランク、重尾）の下で、アテンション品質を向上させるのか、それとも劣化させるのかが不明である。

既存の評価指標である平均二乗誤差（MSE）や KL 発散さえも、スケール誤差と方向誤差を混同したり、単一の誤量子化された支配的なキーがアテンション分布を崩壊させる壊滅的な「argmax シフト」失敗を捉えきれないことが多い。

手法

理論的枠組み

著者らは、その分析を高次元球面の幾何学に基づいている。超球面 $S^{d-1}$ 上で一様に分布するベクトルの場合、座標の周辺分布はベータ分布 $\text{Beta}(\frac{d-1}{2}, \frac{d-1}{2})$ に従うことを確立している。次元 $d \to \infty$ において、これはガウス分布に収束する。

• イェンセンの不等式とソフトマックス： 論文は、ソフトマックス関数がキーの内積における分散を増幅するという重要な非線形性を指摘している。量子化誤差が平均ゼロであっても、指数関数 ($e^x$) の凸性によりバイアス（イェンセン・バイアス）が生じ、高い量子化分散を持つキーのアテンション重みが膨張する。
• K-V 非対称性： キーテンソルは非線形な内積を通じてアテンション機構に入り、バリューテンソルは線形な重み付き和を通じて入る。したがって、キーにおける分散の増大は壊滅的であるのに対し、バリューにおける不偏なバイアス補正は有益である。

量子化方式

厳密な等しいビット予算制約（ベクトルあたり $n$ ビット）の下で、3 つの方式を比較する。

1. KV： K と V の両方にスカラー MSE 量子化（ロイド・マックス）を適用。回転もスケッチもない。
2. KQV： 両方にウォルシュ・アダマール変換（WHT）を適用。K はスカラー MSE（$n$ ビット）で量子化。V はスカラー MSE（$n-1$ ビット）＋1 ビット QJL スケッチで量子化。
3. QKQV： 両方に WHT を適用。K と V の両方が（$n-1$ ビット）＋1 ビット QJL スケッチを使用。

評価指標

スカラー指標を超えて、論文はキーキャッシュ、バリューキャッシュ、出力（$T$）の 3 つの場所における誤差をスケール（SNR）と方向（$\phi$）に分解する6 次元誤差フレームワークを導入する。

• シャノン-$\phi$ 平面： スケール誤差対方向誤差のプロット。
• 統計的ツール： マン・ホイットニーの $U$ 検定（ランク二系列相関）、コルモゴロフ・スミルノフ検定（CDF 分離）、エネルギー距離（6 次元指標）。
• 機構的代理指標： 参照分布と量子化されたソフトマックス分布間の KL 発散、およびトップ 5 ルーティングリコール。

主要な貢献

1. キーに対する QJL の有害性に関する理論的証明

論文は、キーにおけるスカラー量子化を 1 ビット QJL スケッチに置き換えることが、スカラー残差に対して内積分散を少なくとも $\pi/2$ 倍（約 1.57 倍）、完全な $n$ ビットスカラー方式と（$n-1$ ビットスカラー＋1 ビット QJL）方式を比較する場合は $2\pi$ 倍に増大させることを証明している。

• メカニズム： QJL は残差の不偏推定量を提供するが、符号ビットに置き換えることで大きさの情報を破棄する。この推定量の分散は $\frac{\pi}{2} \frac{\|r\|^2}{d}$ である。
• 帰結： この増大した分散は、ソフトマックスの指数関数によって非線形的に増幅され（イェンセンの不等式を通じて）、アテンション重みの体系的な歪みをもたらす。したがって、公平な予算の下では、キーへの QJL 適用はスカラー量子化よりも構造的に劣る。

2. K-V 非対称性と最適方式（KQV）

実験結果は、テストされたすべての分布と予算において、KQV（V のみ QJL 適用）がQKQV（両方に QJL 適用）より優れていることを確認している。

• KL 発散： 実用的に支配的な 4 ビット予算（$n=4$）において、KQV は QKQV よりも KL 発散が 2.25 倍低く（$p < 10^{-23}$）。
• ルーティング： KQV における低い KL 発散は、高いトップ 5 ルーティングリコールと低い出力方向誤差と直接相関する。
• 結論： 1 ビット QJL スケッチは、分散ペナルティなしにバイアスを補正できるバリューキャッシュに留保すべきであり、致命的な分散を導入するキーキャッシュには決して適用すべきではない。

3. 低ランクの病理

論文は、スカラー指標が見逃す低ランクキー分布における重要な失敗モードを特定している。

• 周辺分布対結合分布： WHT は周辺座標をベータ分布（コードブックに一致）にするが、結合低ランク構造を破壊するものではない。
• 壊滅的な KL 発散： 高次元（$d=1024$）において、周辺ヒストグラムが同一であっても、低ランクキーは重尾キーよりも桁違いに悪い KL 発散を生み出す。ソフトマックスは部分空間内の支配的なキーにおける小さな誤差を増幅し、「argmax シフト」失敗を引き起こす。
• 含意： 標準的な回転ベースの量子化は、低ランクアテンション行列には不十分であり、結合構造（特に低ランク）に対処する必要がある。

4. 予算依存の交差

幾何学的な K 再構成品質における交差という直感に反する発見がある。

• 低予算（$n \in \{2, 3\}$）では、QKQV（K への QJL 適用）が、KQV よりも幾何学的な K 再構成（低い方向誤差）を達成することがある。
• しかし、この幾何学的な改善は、分散増大が依然としてソフトマックスを害するため、アテンション品質（KL またはルーティング）の向上にはつながらない。
• $n=4$ において、KQV はすべての指標で勝利する。この交差はランクや尾部の重みに依存せず、スカラー精度と JL 残差補正間の根本的なレート歪みトレードオフを示唆している。

結果

• KQV の優位性： $n=4$ において、KQV はすべての分布（ガウス、重尾、低ランク、集中）において、KL 発散、幾何学的 K 誤差、6 次元距離の面で QKQV と KV を凌駕する。
• キーへの QJL は有害： キーにおけるスカラービットを QJL に置き換えることは、分散増大により KL 発散を常に増加させる。これは、最終出力に利益をもたらさない特定の低予算幾何学的再構成シナリオを除く。
• 肥満尾部への頑健性： WHT は、周辺集中（肥満尾部）に対する無条件の保護を提供する。KQV は尾部の重みが増加（$\nu \to 1.05$）しても安定した方向誤差（$\epsilon_{dir}^K \approx 0.027$）を維持するのに対し、KV は壊滅的に失敗する。
• 低ランク失敗： 低ランク分布は、周辺正規化にもかかわらず結合部分空間構造が保持されるため、肥満尾部の問題を遥かに上回るアテンション品質の深刻な劣化を引き起こす。

意義と主張

この論文は、TurboQuant のオリジナルの QJL 適用（対称的に拡張された場合）が有害となる理由について、厳密な統計的および幾何学的説明を提供すると主張している。

1. 指標の洗練： 平均 KL は楽観的な指標であると主張し、KV キャッシュ量子化の信頼性の真の尺度は、最悪ケースの KL とルーティングリコールであると述べている。
2. 設計原則： ソフトマックスの非線形性により、K と V キャッシュは根本的に異なる量子化戦略を必要とすることを確立している。QJL は線形出力（V）に適したバイアス補正ツールであるが、非線形出力（K）には有害な分散注入ツールである。
3. 周辺保証の限界： WHT による周辺分布がコードブックと一致することの確保は必要だが不十分であることを示している。結合構造（特に低ランク）が決定的な失敗モードを決定する。
4. 未解決問題： 幾何学的再構成における予算依存の交差を、未解決のレート歪み問題として特定している。非常に低いビット数（$n=2,3$）では QKQV が幾何学的に勝利するが、分散のイェンセン増幅により機能的には敗れることを指摘している。

この研究は、堅牢な KV キャッシュ量子化には、SVD や部分空間意識型コードブックなどによる低ランク構造の明示的な処理が必要であり、現在のビット予算において KQV 方式（K への WHT＋スカラー、V への WHT＋スカラー＋QJL）が最適な構成であると結論付けている。