GoogleのKV cache圧縮技術 TurboQuant について

はじめに

3/24にGoogleが発表した新技術TurboQuantに関する記事「TurboQuant: Redefining AI efficiency with extreme compression」が非常に面白いと感じたので、その内容についてClaudeと色々議論して勉強した内容をまとめる。

概要

• GoogleがLLMのKV cacheを圧縮するアルゴリズム「TurboQuant」を発表
• 本技術によってKV cacheメモリを平均6倍削減し、NVIDIA H100 GPU上でのAttention計算を最大8倍高速化する
• TurboQuantはベクトルを量子化するための技術であるため、適用範囲はKV cacheに留まらず、RAG、異常検知、顔認識などにも応用可能

TurboQuantの詳細

TurboQuantは大きく分解するとPolarQuantとQJLという2つのアルゴリズムから構成される。

PolarQuant

PolarQuantは、ベクトルを通常のデカルト座標系から極座標に変換したのち、量子化によってデータ量を圧縮するプロセスを担う。

2次元における極座標への変換は、すなわち $(x, y) \to (r, \theta)$ である。3次元では球面座標で $(x, y, z) \to (r, \theta, \phi)$ となり、角度が2つに増える。一般に $d$ 次元では、1つの半径 + $(d-1)$ 個の角度で表現され、超球面座標と呼ばれる。

これって $d$ 次元から $d$ 次元への変換だから情報量変わってなくない？

と思えるが、この極座標への変換は前処理のイメージに近い。本質としては、 高次元空間における極座標表現はデカルト座標表現に比べて値の分布が小さく、予測しやすい形で収束しやすい という傾向にあることが重要である。すなわち、ベクトルの半径はほぼ一定値に収束しやすく、角度成分も特定の分布に従って集中しやすい。したがって、この予測可能な分布に対して最適な量子化コードブック（どの範囲の値をどの離散値に割り当てるか）を事前に一度計算しておけばよい。

これが従来手法との最も大きな差別化点である。デカルト座標の場合は、各次元の値の分布はデータに依存して大きく変わりうるため、データごとに正規化定数を保存する必要があった（= 1 ~ 2ビット分の量子化オーバーヘッド）。PolarQuantでは分布が事前にわかっているため、この正規化定数の保存が不要になり、その分のビットを圧縮に回すことができる。

QJL

QJLのプロセスでは、量子化した際の残差が持つ系統的なバイアスを除去する。QJLはJohnson-Lindenstrauss変換という数学的手法を使って、複雑な高次元データをデータ点間の本質的な距離や関係性を保持したまま圧縮する。

PolarQuantで量子化されたベクトル $\hat{x}$ には、元ベクトル $x$ との間に残差 $e = x - \hat{x}$ が生じる。この残差のせいで、別のベクトル $y$ との内積を計算すると $<y, \hat{x}> \neq <y, x>$ となり、系統的なバイアス（常に少し大きい、または小さい方向に値がズレる）が生じる。LLMにおけるAttention機構は内積でスコアを計算するため、系統的なバイアスはsoftmaxの分布を歪める。これは注目すべきトークンを見逃すような劣化を引き起こしうる。

そこでQJLは、この残差 $e$ に対してランダム射影 + 1ビット符号量子化を行い、内積推定のバイアスをゼロにする。すなわち、期待値として内積が完全に一致する。個々の推定には分散が生じるが、この分散にはバイアスが無いため、系統的な歪みは生じない。さらに、高次元ではこの分散も小さくなる。

QJLの1ビット符号量子化には、PolarQuantで節約した1 ~ 2ビット分の量子化オーバーヘッドに当てられていたリソースが使用される。言い換えると、従来は1 ~ 2ビット分を値の精度を上げるために使用していたが、TurboQuantではそれを残差のバイアス除去に使用する。前者は値の精度を多少上げることができても、系統的なバイアスが残る以上、Attentionの精度は保証されない。一方で後者は系統的なバイアスをゼロにすることにリソースを割くことで、量子化前後の精度低下をゼロに抑えることができる。

個人的に興味深い点

本技術はKV cacheだけのものではない

この技術はあくまで「ベクトルを量子化するための技術」であるため、応用範囲はLLMにおけるKV cacheに留まらない点が非常に面白いと感じた。例えばRAGにおけるベクトル検索も高速化対象である。

具体的には、クエリベクトルとDB内の大量のベクトルの内積（またはコサイン類似度）を計算する過程において、DBから大量のベクトルを読み出す際のデータ量を大幅に削減することができる。さらに、クエリベクトルとDB内のベクトルとの内積計算の精度はQJLによって維持される。

この技術によって、RAGにおける検索速度の大幅な向上が期待される。最近はユーザーの入力に対して一度のみ検索実行するのではなく、ユーザーの目的を達成できるまで何度もエージェンティックに検索を繰り返すような構成が基本になってきているため、検索の効率化は見た目以上に効果が大きいと期待される。

LLM推論専用チップへの応用可能性

もう一点、過去の記事で紹介したTaalas社開発のLLM推論専用チップとの親和性の高さにも注目したい。

このチップは、ストレージ（ROM）に対応する部位にLLMのモデルの重み情報を直接焼き付けることで、既存技術をはるかに凌ぐ速度での推論を可能とするものである。

この技術では、モデルの重みを乗せるためのメモリ領域を必要としない一方、KV cacheは通常通りメモリに乗せる必要があった。したがって、モデルサイズが大きくなれば、KV cacheのためにより大きなメモリ領域を必要とする。これがモデル拡張のボトルネックの一つであると考えられた。

そこに今回のTurboQuantを適用することで、KV cacheに必要なメモリ領域を大幅に削減することができるようになり、モデルの重みを焼き付けたROM部分により多くの領域を割くことができるようになることが期待される。

さいごに

GoogleはTurboQuantのアルゴリズムを完全にオープンにしているわけではなく、今回は論文での報告に留めたようである。しかしながらコミュニティへのインパクトは大きかったようで、早速オープンソースでTurboQuantを再現する動きや再現実装を用いてLLMを動かす試みが活発化している。

応用範囲が広くて大変面白い技術だと感じるため、引き続き注目していきたい。