技術的概要：RATEQUANT

問題定義

大規模言語モデル（LLM）は、生成中に以前に計算されたすべてのキー・バリュー（KV）ペアをキャッシュする必要があり、これはシーケンス長、バッチサイズ、モデルの深さに比例して線形にスケールするメモリボトルネックを生み出します。KV キャッシュを量子化することでこのフットプリントを削減できますが、既存の手法はすべてのアテンションヘッドに同じビット幅を均一に割り当てます。このアプローチは、ヘッドの重要性におけるよく知られた不均一性を無視しています。

レイヤーレベルまたはチャネルレベルでの混合精度割り当てに関する最近の試みは、特定の量子化器に紐付いたヒューリスティックな規則に依存しており、理論的な保証や転送可能性が欠けています。著者らは、より深く、以前は認識されていなかった障害を特定しました：歪みモデルの不一致です。異なる量子化器は、$D(b) = \alpha \beta^{-b}$ とモデル化される、明確に異なる歪み - レート曲線に従います。減衰率 $\beta$ は設計によって大きく異なります（例：TURBOQUANT では 3.6、QuaRot では 5.3）。ある量子化器の歪みモデルから導き出された混合精度割り当て戦略を、別の量子化器に単純に適用すると、アテンションヘッドの限界利益の順序が逆転し、しばしば均一量子化よりも悪いパフォーマンスをもたらします。

手法

RATEQUANT は、ヘッドごとの KV キャッシュビット割り当てをレート歪み最適化問題として形式化することで、これらの課題に対処します。このフレームワークは 4 つのフェーズから構成されます：

1. 問題定式化と最適割り当て

著者らは、ヘッドごとの量子化平均二乗誤差（MSE）をビット幅の指数関数としてモデル化します：$D(b) = \alpha \beta^{-b}$。目的は、総ビット予算 $B$ の制約のもとで、$N$ 個の量子化グループ（ヘッド）にわたる重み付き歪みの和を最小化することです：
$$\min_{b} \sum_{i=1}^{N} w_i \cdot D(b_i) \quad \text{s.t.} \quad \sum b_i = B$$
ここで、$w_i$ はヘッド $i$ の感度を表します。

ラグランジュの定常性を用いて、著者らは逆水充填法による閉形式解を導出します：
$$b^*_i = \bar{b} + \frac{\ln w_i - \ln \bar{w}}{\ln \beta}$$
ここで、$\bar{b}$ は平均ビット幅、$\bar{w}$ は感度の幾何平均です。この解は、感度の高いヘッドにより多くのビットを割り当てます。理論的には、達成可能な歪み削減は、ヘッド感度の算術平均と幾何平均（AM/GM）の比率によって制限されます。

2. 感度推定

重み $w_i$ を決定するために、RATEQUANT は勾配ベースの感度を採用します。これは、因果言語モデル化の損失に対する KV 投影出力の勾配ノルムの二乗を計算します：
$$w_{l,h} = \mathbb{E}_{x \sim \mathcal{D}} \left[ \frac{1}{T} \sum_{t=1}^{T} \left\| \frac{\partial \mathcal{L}}{\partial K_{l,h,t}} \right\|^2 \right]$$
著者らは、勾配ベースの代理指標が、誤差増幅を測定するのではなく直接的な損失への影響を測定する活性化ベースの指標（チャネルの大きさなど）よりも優れていることを実証しています。勾配ベースの感度は、U 字型の重要性パターン（初期層と後期層で高い感度）を明らかにしますが、活性化ベースの指標は深さとともに単調に増加する傾向を示します。

3. 整数割り当てと拡張

ビット幅は整数でなければならないため、連続的な解は貪欲な限界利益アルゴリズム（アルゴリズム 1）を用いて洗練されます。このアルゴリズムは、予算が尽きるまで、重み付き歪みの削減量が最大となるヘッドにビットを反復的に割り当てます。

2 つの重要な拡張により、RATEQUANT は量子化器に依存しないものになります：

• 歪み較正：普遍的な $\beta$ を仮定する代わりに、RATEQUANT は小さな較正セット（16 系列）を用いて量子化器ごとの歪みモデル（$\alpha, \beta$）をフィットさせます。これにより、「歪みモデルの不一致」という失敗モードを防ぎます。
• K/V 分離：キーとバリューはしばしば異なる量子化スキームを使用すること（例：KIVI におけるキーはチャネルごとの対称量子化、バリューはトークンごとの非対称量子化）を認識し、フレームワークはこれらを異なる歪み曲線を持つ $2N$ 個の独立したコンポーネントとして扱います。これにより、アルゴリズムは非対称なビット予算（例：バリューよりもキーに多くのビットを割り当てる）を発見できます。

4. パイプライン

このプロセスは、以下の 1 回限りのオフラインコストを含みます：

1. 感度較正：小さなデータセットでの順方向・逆方向パス（8B モデルで約 1.6 秒）。
2. 歪みモデリング：5 つのビット幅で $\alpha$ と $\beta$ をフィット（0.1 秒未満）。
3. 割り当て：貪欲アルゴリズムの実行（0.01 秒未満）。
4. 推論：静的なビット幅を適用し、ゼロのランタイムオーバーヘッドを実現。

主要な結果

実験は、WikiText-2 のパープレキシティ（PPL）と下流タスク（ARC-Challenge、HellaSwag など）を使用して、Qwen3 モデル（4B、8B、32B）で行われました。

• 不一致の解決：不一致な歪みモデルを適用する（例：KIVI に対して TURBOQUANT の $\beta$ を使用）と、2.5 ビットの平均ビット幅において KIVI の PPL が 49.3 から 87.0 に悪化しました。較正された RATEQUANT はこれを解決し、PPL を14.9に削減しました（70% の削減）。
• パフォーマンスの向上：
  • KIVI（2.5 ビット）：RATEQUANT は PPL を 49.3 から 14.9 に改善しました。
  • QuaRot（2.5 ビット）：PPL を 6.6 ポイント改善しました。
  • TURBOQUANT（4.0 ビット）：Qwen3-8B において、量子化によって引き起こされた劣化の 66% を回復しました。
• K/V 分離：2.5 ビットの KIVI において、RATEQUANT によって発見された最適な割り当ては、キーで 2.85 ビット、バリューで 2.15 ビットであり、結合割り当てを大幅に上回りました。
• 効率性：この手法はランタイムオーバーヘッドを追加しません。較正は単一の GPU で 2 秒未満で完了します。
• 下流タスクの精度：4.0 ビットにおいて、RATEQUANT は下流タスクで平均 64.2% の精度を達成し、FP16 ベースライン（64.4%）とほぼ同等となり、精度ギャップの 89.8% を回復しました。

意義と主張

本論文は、RATEQUANT が、キーとバリューの構造的な違いに対処しながら、ヘッドごとの粒度で KV キャッシュ量子化にレート歪み理論を適用する最初のフレームワークであると主張しています。

中心的な発見は、歪みモデルの不一致が混合精度割り当てにおける重要な失敗モードであるということです。特定のベース量子化器に合わせて歪み曲線（$\beta$）を較正しなければ、混合精度戦略は有害になる可能性があります。以下の 3 つを組み合わせることで：

1. 勾配ベースの感度（活性化ベースの代理指標よりも優れている）、
2. 実証的な歪み較正（量子化器に一致させるため）、
3. K/V 分離（非対称な誤差特性を処理するため）、

RATEQUANT は、混合精度 KV 量子化を、汎用的で量子化器に依存しない割り当てレイヤーへと変換します。著者らは、ヘッド感度の AM/GM 比率が、混合精度割り当てが有意な利益をもたらすかどうかの実用的な予測指標として機能すると仮説を立てています。この研究は、レート歪みの視点が、KV キャッシュを超えた他の不均質なニューラルネットワークコンポーネントにおけるリソース割り当てにも一般化できることを示唆しています。