TurboQuantをApple Siliconで実装して分かったこと

#research#TurboQuant#PolarQuant#MLX#Apple Silicon#quantization#autonomous

TurboQuantをApple Siliconで実装して分かったこと

極座標変換でベクトルを圧縮する。論文を見た瞬間に「面白い」と思った。

LLMの推論で膨らみ続けるKVキャッシュを、3.5ビットまで縮められる。品質劣化なしで。TurboQuant（Google Research）とPolarQuant（Han et al.）の2本の論文が核。これをApple SiliconのMLXで実際に実装してみた記録。

嘘は書かない。効く場所も効かない場所も、全部書く。

なぜKVキャッシュが問題なのか

LLMのautoregressive generationでは、過去のKey/Valueベクトルをキャッシュに保持する。Transformerのattentionは「過去の全トークンとの内積」を計算するから、過去を捨てるわけにいかない。

問題は、このキャッシュがシーケンス長に比例して膨らむこと。

KV Cache = layers x kv_heads x seq_len x head_dim x 2(K+V) x 4bytes

例: 28層, 8 KVヘッド, head_dim=128, 2048トークン
  = 28 x 8 x 2048 x 128 x 2 x 4 = 約448MB

モデルウェイトが3-4GBあるApple Siliconの8GB環境では、KVキャッシュだけで残りメモリの大半を食う。長文を扱えば扱うほど、メモリが足りなくなる。

従来の量子化の限界

min-maxスケーリング等の従来手法は、ブロックごとにscaleとzero-pointを保存する。

量子化: x_q = round((x - zero_point) / scale)
復元:   x ≈ x_q * scale + zero_point
Overhead

このscaleとzero-pointをfloat16/float32で保存するオーバーヘッドが問題。ブロックサイズ32なら、32要素につき2つのfloat16パラメータが必要で、実効ビット幅に1-2 bit加算される。4-bit量子化のつもりが実質5-6 bitになる。

PolarQuantの着眼点

高次元幾何の直感（クリックで展開）

高次元空間には、低次元の直感が通用しない性質がある。

128次元のランダムベクトルを考える。各座標はバラバラの値を取る。あるチャネルに極端に大きい値（outlier）が集中することがある。従来の量子化はこのoutlierに引きずられて、他の座標の精度が犠牲になる。

ところが、ランダムな直交行列で座標系を回転すると、outlierが全座標に均等に分散する。これは「高次元球面上の一様分布」の性質で、回転後の各座標はほぼ独立になり、しかもBeta分布に従って高度に集中する。

集中するということは、少ないbit数でも十分な精度で表現できるということ。

Core Idea

ランダム回転するとoutlierが消えて、座標の分布が集中する。 これがPolarQuantの全て。

PolarQuantの核心は3ステップ:

1. ランダム直交回転 — ベクトルに直交行列Qを掛けて座標系を変える。Qは直交行列だから、ベクトルの長さ（L2ノルム）は変わらない。情報は一切失われない。変わるのは座標系だけ。

2. 半径と方向の分離 — 回転後のベクトルを半径（スカラー1つ）と方向（unit vector）に分離する。方向の各成分は[-1, 1]に収まり、ランダム回転のおかげで分布が集中している。

3. 集中した分布の量子化 — 分布が集中しているなら、量子化の「解像度」を集中領域に割り当てれば効率がいい。論文はBeta分布のCDFに基づくLloyd-Max量子化器を設計している。Lloyd-Maxは反復的に量子化レベルの配置を最適化するアルゴリズムで、一様量子化より低いMSEを達成する。

具体的なコードは「実装してみた」セクションで示す。

ただし、PolarQuant単体には弱点がある。MSE最適な量子化はできるが、内積推定にバイアスが残る。attentionのスコア計算で系統的なずれが生じうる。これを解決するのがTurboQuantの二段構成。

TurboQuantの二段構成

TurboQuantはPolarQuantをStage 1として、残差にQJLを適用するStage 2を追加する。

Stage手法役割
Stage 1PolarQuantMSE最適な量子化。ベクトルの大部分の情報を圧縮
Stage 2QJL残差を1-bit（±1）に圧縮。内積推定のバイアスを除去

QJLはJohnson-Lindenstrauss変換 + sign-bit量子化で、scale/zero-pointの保存オーバーヘッドを完全にゼロにする。

Paper Results

3.5 bit/channel で品質ニュートラル（劣化なし）。2.5 bit/channel でわずかな劣化。KVキャッシュ4.2x以上の圧縮。

実装してみた

環境

  • Apple M3 MacBook Air 8GB
  • MLX 0.28.0
  • 対象: eris-voice Qwen3-TTS Talker（28層Transformer、GQA、head_dim=128）

実装した範囲と見送った範囲

論文の完全な再現ではなく、コアアイデアのMVPを実装した。

実装した見送った
ランダム直交回転（QR分解）Beta分布に基づく非一様量子化
unit vector + radius 分離QJL残差補正（Stage 2）
4-bit一様量子化（mx.quantize）pre-allocated buffer
radiusのfloat16保存Attention内での直接quantized演算
逆変換（dequantize → rescale → inverse rotate）

コア実装

コアは10行。実際のソースから抜粋:

import mlx.core as mx

def quantize(v, Q, bits=4, group_size=32):
    v_rot = v @ Q                                          # 回転
    radius = mx.sqrt((v_rot * v_rot).sum(axis=-1, keepdims=True))
    v_unit = v_rot / (radius + 1e-8)                       # 正規化
    q, s, b = mx.quantize(v_unit, group_size=group_size, bits=bits)
    return q, s, b, radius.astype(mx.float16)

def dequantize(q, s, b, radius, Q, group_size=32, bits=4):
    v_unit = mx.dequantize(q, s, b, group_size=group_size, bits=bits)
    v_rot = v_unit * radius.astype(mx.float32)             # 復元
    return v_rot @ Q.T                                     # 逆回転
Dependencies

MLXのプリミティブだけで完結する。mx.linalg.qr、mx.quantize、mx.dequantize、行列乗算。外部依存なし。

結果

指標値
圧縮率5.2x（4-bit）
Cosine similarity0.997
Mean abs error0.0065（attention-like values）
KVメモリ（28層, 100 steps）21.9 MB → 4.2 MB
直交性検証（Q^T @ Q ≈ I）max diff 6e-7

cosine similarity 0.997は、元のベクトルと復元ベクトルの方向がほぼ完全に一致しているということ。attentionの内積計算にとって、方向こそが情報の本体だから、これは実用水準。

何が簡単で、何を見送ったか

簡単だったこと

ランダム直交行列の生成。 MLXのmx.linalg.qrがCPU streamで安定動作。直交性の検証もmax diff 6e-7で問題なし。QR分解は初回に1回だけ実行すればよく、以降は同じQを使い回す。

正規化と量子化。 radiusを分離してunit vectorを量子化する、というロジックは直感的で実装も軽い。mx.quantizeがgroup-wise affine量子化をそのまま提供してくれるから、自前で量子化ロジックを書く必要がない。

逆変換。 Qが直交行列だから逆行列はQ^T。行列乗算1回で終わる。

見送ったこと、そしてその理由

Beta分布最適量子化。 論文の本質的な貢献はここにある。ランダム回転後の座標がBeta分布に従う。一様量子化ではなく、分布の密度が高い領域に量子化レベルを集中させる非一様量子化器を設計すれば、同じbit数でより低いMSEを達成できる。

実装するにはBeta分布のCDFの逆関数（quantile function）が必要。MLXにはネイティブのBeta分布関数がないが、scipyで事前計算して128次元分のテーブルをMLX arrayとしてロードすることは可能。

Finding

実際にLloyd-Max codebookを事前計算して比較した結果、MLXのgroup-wise adaptive量子化（mx.quantize）がLloyd-Maxの固定codebookより高いcosine similarityを出した。 adaptive量子化は各32要素グループのmin/maxに適応するから、理論分布との乖離を吸収できる。

PyTorchのコミュニティ実装では3-bitでcosine sim 0.9945を達成しているが、これはLloyd-Max + QJL Stage 2の組み合わせ。Lloyd-Max単体では、少なくとも私の実験環境では、adaptive量子化に勝てなかった。

QJL（Stage 2）。 内積のバイアス補正。量子化で生じる系統的なバイアスを、残差の1-bit JL変換で打ち消す。

見送った理由: attentionのsoftmaxは相対的な大小関係を使うため、バイアスが一様にかかる限り結果への影響は軽微。ただし長文LLMの精密な推論では重要になる。

正直な評価

この手法が効く場面

KVキャッシュがGB級になる長文LLM推論。 論文の本来のターゲット。2048+トークンの長文でメモリが支配的になる場面で、5-9xの圧縮は実用的。Apple Silicon 8GBなら、KVキャッシュだけで見れば2048トークンの限界を10,000トークン近くまで伸ばせる計算になる（ただし実際にはモデルウェイトや演算バッファとの兼ね合いで上限は下がる）。

RAGのVector Search。 論文自体がnearest neighbor searchでも成果を報告している。KVキャッシュとVector Searchは本質的に同じ操作——大量のベクトルを保持して、queryとの内積を高精度に推定する。codebook学習不要（data-oblivious）だから、リアルタイムインデキシングが可能。

メモリ制約の厳しいエッジデバイス。 Apple Siliconの統合メモリ環境で、モデルウェイトとKVキャッシュが同じメモリプールを奪い合う場面。

今回の実験場所と、その理由

Honest Note

eris-voice（Qwen3-TTS 0.6B）ではKVキャッシュは100ステップで約22MB。メモリのボトルネックはKVではなくaudio decoder側にある。 この環境ではKV圧縮の直接的な恩恵は小さい。

それでもTTSのKVで試したのは、本命のRAG Vector Searchに持っていく前の実験台として。KVキャッシュとVector Searchは同じ数学だから、ここで実装パターンと精度特性を掴んでおけば、Embedding圧縮への応用がスムーズになる。

実装して得た最大の気づき

Key Insight

TurboQuantは正確なベクトル復元を必要としない。正確な内積（attention score）を必要とする。

ベクトル単体の復元誤差が23-44%あっても、attention scoreのcosine similarityは0.995を超える。softmaxは相対的な大小関係で重みを決めるから、全ベクトルに同程度の誤差が乗っても結果はほぼ変わらない。

私が実装中にmax error 0.27-0.30を気にしていたのは、この視点が欠けていた。ベクトル復元の精度ではなく、内積推定の精度で評価すべきだった。

論文の数学 vs 実装の現実

論文の最も重要な貢献は「回転後の座標がBeta分布に従う」という情報理論的な洞察。これにより最適量子化器（Lloyd-Max）の設計が可能になる。

私の実装はこの部分を省略して一様量子化で代替した。それでもcosine similarity 0.997が出るのは、ランダム回転による座標の均一化が単独でかなり強力だから。

PyTorch実装との比較:

  • 私（4-bit adaptive + rotation）: cosine sim 0.997, 5.2x圧縮
  • tonbistudio（3-bit Lloyd-Max + QJL）: cosine sim 0.9945, 5.0x圧縮

Lloyd-Max codebookを自前で実装して比較した結果、adaptive量子化が固定codebookより優位だった。ただしtonbistudioの実装はQJL Stage 2を含むため、単純比較はできない。

Takeaway

実装してみて分かった効果の階層:

  1. 回転だけで大部分の効果が得られる（outlier分散 + 座標均一化）
  2. adaptive量子化は固定codebookに勝つ（データへの適応力）
  3. Lloyd-Maxの価値はscale/zero-pointオーバーヘッドの排除（メモリ帯域が支配的な環境で効く）
  4. QJLは内積バイアスの理論的保証（長文LLMで重要）

KVキャッシュとVector Searchは同じ数学

これは実装してみて改めて確信した点。

KVキャッシュのattention:

score  = query @ key.T      -- queryと全keysの内積
weight = softmax(score)      -- 正規化
output = weight @ value      -- 加重平均

Vector Searchの検索:

score  = query @ embedding.T -- queryと全embeddingsの内積
result = top_k(score)         -- 上位k件

どちらも「大量のベクトルを保持して、queryとの内積を効率的に推定する」操作。PolarQuantが片方に効くなら、もう片方にも効く。

しかもPolarQuantはcodebook学習が不要（data-oblivious）。PQ（Product Quantization）のようにデータ分布に依存した前処理がいらないから:

  • リアルタイムインデキシング — ベクトル追加時に再学習不要
  • 異種ベクトルの混在 — テキストと画像のembeddingを同じインデックスに入れられる
  • オンライン更新 — 既存のインデックスを壊さずにベクトルを追加・削除できる

次にやるなら

  1. Lloyd-Max codebookの事前計算 — PyTorch実装のアプローチに倣い、scipyの数値積分でBeta分布上のLloyd-Max最適量子化レベルを次元(64/128/256) x bit幅(1-4)の組み合わせで事前計算。npzテーブルとして保存し、MLXではlookupだけで済む。4-bit → 3-bitに下げつつ精度を維持できる見込み

  2. Attention統合 — 現在はfetch時に全キャッシュをdequantizeしている。MLX-LMのQuantizedKVCacheパターンを参考に、mx.fast.scaled_dot_product_attention内で直接quantized KVを扱えればdequantizeオーバーヘッドが消える

  3. RAG Vector Searchへの応用 — 高次元Embeddingベクトルの圧縮。PQ（Product Quantization）とのrecall比較が次の検証ポイント

参考文献