TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate
Amir Zandieh, Majid Daliri, Majid Hadian, Vahab Mirrokni
TL;DR
TurboQuant tackles online vector quantization by delivering near-optimal distortion for both $D_{\text{mse}}$ and $D_{\text{prod}}$ in high-dimensional spaces. It combines a random rotation to induce per-coordinate Beta distributions with per-coordinate Lloyd-Max scalar quantization, plus a residual stage using Quantized Johnson-Lindenstrauss to achieve an unbiased inner-product estimator $\mathbb{E}[\langle y, Q^{-1}(Q(x))\rangle] = \langle y, x\rangle$ and distortion bounds $D_{\text{mse}} \le (\sqrt{3}\pi/2) 4^{-b}$ and $D_{\text{prod}} \le (\sqrt{3}\pi^2 \|y\|_2^2 / d) 4^{-b}$. The authors establish information-theoretic lower bounds via the Shannon lower bound and Yao's minimax principle, showing TurboQuant is within a small constant factor (about $2.7$ for MSE) of the optimum. Empirically, TurboQuant delivers strong performance on KV-cache quantization and nearest-neighbor search, achieving near-perfect long-context retrieval at modest bit-rates and outperforming product-quantization baselines in recall with near-zero indexing time. These results demonstrate online, accelerator-friendly quantization that preserves essential geometry for inner products and distances in large-scale AI and retrieval systems.
Abstract
Vector quantization, a problem rooted in Shannon's source coding theory, aims to quantize high-dimensional Euclidean vectors while minimizing distortion in their geometric structure. We propose TurboQuant to address both mean-squared error (MSE) and inner product distortion, overcoming limitations of existing methods that fail to achieve optimal distortion rates. Our data-oblivious algorithms, suitable for online applications, achieve near-optimal distortion rates (within a small constant factor) across all bit-widths and dimensions. TurboQuant achieves this by randomly rotating input vectors, inducing a concentrated Beta distribution on coordinates, and leveraging the near-independence property of distinct coordinates in high dimensions to simply apply optimal scalar quantizers per each coordinate. Recognizing that MSE-optimal quantizers introduce bias in inner product estimation, we propose a two-stage approach: applying an MSE quantizer followed by a 1-bit Quantized JL (QJL) transform on the residual, resulting in an unbiased inner product quantizer. We also provide a formal proof of the information-theoretic lower bounds on best achievable distortion rate by any vector quantizer, demonstrating that TurboQuant closely matches these bounds, differing only by a small constant ($\approx 2.7$) factor. Experimental results validate our theoretical findings, showing that for KV cache quantization, we achieve absolute quality neutrality with 3.5 bits per channel and marginal quality degradation with 2.5 bits per channel. Furthermore, in nearest neighbor search tasks, our method outperforms existing product quantization techniques in recall while reducing indexing time to virtually zero.
