OCTOPUS: 最適二乗誤差量子化下での八面体パラメータ化によるTransformer向け最適化KVキャッシュ

要旨

キー値（KV）キャッシュは、長いコンテキストの自己回帰推論においてメモリ帯域幅とフットプリントを支配する要因である。最近の回転前処理を施した符号化方式（TurboQuant、PolarQuant）は、構造化されたランダム回転の後に、解析的に扱いやすい周辺分布に適合した座標ごとのスカラ量子化器を適用することが、KV圧縮のほぼ最適な方法であることを示している。OCTOPUSは、回転された座標の三つ組（トリプレット）を共同で量子化することにより、このパラダイムを前進させる。各トリプレットの方向は八面体パラメトリゼーションを介して正方形にマッピングされ、得られた2つの座標とトリプレットのノルムは、実装に合わせた周辺分布に対してロイド・マックス量子化される。トリプレットごとの二乗誤差を最適化することで、キーの総次元数のみに依存する厳密に非一様なビット割り当てが得られる。我々は、有限次元での品質最適点が、テストしたすべての実際のデコーダで一定であることをスイープにより発見した。この符号化方式はデータ非依存、オンライン、そしてシードが与えられれば決定論的である。テキスト、ビデオ、オーディオにわたって、OCTOPUSは報告されたすべてのビット幅とメトリックにおいて、過去のすべての回転符号化方式に匹敵するか凌駕しており、そのリードは極端な圧縮のためにビットが減少するにつれて大きくなる。さらに、融合されたTriton実装により、非圧縮キーを実体化することなくオンザフライでキーを再構築するため、この符号化方式は既存の逆量子化に比べてデコード時の帯域幅やレイテンシを追加しない。プロジェクトページ：https://octopus-quant.github.io/

English

The key-value (KV) cache dominates memory bandwidth and footprint in long-context autoregressive inference. Recent rotation-preconditioned codecs (TurboQuant, PolarQuant) show that a structured random rotation followed by a per-coordinate scalar quantizer matched to an analytically tractable marginal is a near-optimal recipe for KV compression. OCTOPUS advances this paradigm through joint quantization of rotated coordinate triplets. Each triplet's direction is mapped to a square via an octahedral parameterization, and the two resulting coordinates and the triplet norm are Lloyd-Max quantized against implementation-matched marginals. Optimizing the per-triplet squared error gives a strictly non-uniform bit allocation depending only on the total dimensionality of the keys. We find the finite-dimensional quality optimum with sweeps to be constant on every real decoder we test. The codec is data-oblivious, online, and deterministic given a seed. Across text, video, and audio, OCTOPUS matches or beats every prior rotation codec at every reported bit width and metric, with a lead that grows as bits drop for extreme compression. Furthermore, a fused Triton implementation reconstructs keys on the fly without materializing the uncompressed key, so the codec adds no decode-time bandwidth or latency over the existing dequantization. Project Page: https://octopus-quant.github.io/