TurboQuant: 3-Bit Compression with Zero Overhead
how polarquant and quantized johnson-lindenstrauss combine to achieve 3-bit kv-cache compression with zero memory overhead
Large Language Models (LLMs) suffer from a severe memory bottleneck during inference: the Key-Value (KV) cache. As context windows grow to millions of tokens, storing the keys and values for attention computation dominates GPU memory. Traditional quantization (e.g., INT8 or INT4) introduces memory overhead because it requires storing scale and zero-point parameters for small blocks of channels or tokens to manage the dynamic range of activations.
TurboQuant solves this by achieving 3-bit KV cache compression with virtually zero accuracy loss and zero scaling overhead. It achieves this by unifying two mathematical paradigms: PolarQuant (recursive hyperspherical quantization) and Quantized Johnson-Lindenstrauss (QJL) (1-bit randomized error correction).
Here is the comprehensive mathematical and programmatic breakdown of how this pipeline operates.
Part 1: PolarQuant (The Recursive Geometric Transformation)
The core insight of PolarQuant is that while the Cartesian coordinates ($X, Y, Z, \dots$) of LLM activations have highly volatile ranges (requiring dynamic scaling factors), the angles of these vectors in polar/hyperspherical space do not.
If we map the data to angles, those angles are strictly bounded between $-\pi$ and $\pi$. More importantly, for high-dimensional isotropic data, these angles are near-uniformly distributed. A uniform distribution with absolute, known bounds can be quantized with zero overhead—no scale factors needed.
1. The Mathematics of Recursive Polar Transformation
For a 2D vector $v = [x, y]^T$, the standard Cartesian-to-Polar conversion is: $$r = \sqrt{x^2 + y^2}$$ $$\theta = \text{atan2}(y, x)$$
LLM attention heads operate in $d$-dimensional space (e.g., $d=128$). To handle $d$ dimensions, PolarQuant uses a recursive pairing tree. Assuming $d = 2^k$, the algorithm groups dimensions into pairs, converts them to polar, and recursively repeats the process on the resulting radii.
Let the input vector be $x^{(0)} \in \mathbb{R}^d$.
Step $L=1$: Group $x^{(0)}$ into pairs. For $i \in \{1, \dots, d/2\}$: $$r_i^{(1)} = \sqrt{\left(x_{2i-1}^{(0)}\right)^2 + \left(x_{2i}^{(0)}\right)^2}$$ $$\theta_i^{(1)} = \text{atan2}\left(x_{2i}^{(0)}, x_{2i-1}^{(0)}\right)$$
We now have $d/2$ radii and $d/2$ angles.
Step $L=2$: We recursively apply this to the radii $r^{(1)}$. For $i \in \{1, \dots, d/4\}$: $$r_i^{(2)} = \sqrt{\left(r_{2i-1}^{(1)}\right)^2 + \left(r_{2i}^{(1)}\right)^2}$$ $$\theta_i^{(2)} = \text{atan2}\left(r_{2i}^{(1)}, r_{2i-1}^{(1)}\right)$$
This process repeats for $\log_2(d)$ steps. At the root of the tree, you are left with exactly one global radius $R_{final} = \|x\|_2$ and exactly $d-1$ angles ($\theta$).
2. The Quantization Step
The $d-1$ angles are inherently bounded within $[-\pi, \pi]$ or $[0, \pi/2]$ (since intermediate radii are strictly positive). We can directly quantize them into $2^b$ discrete bins (where $b$ is the bit-width, e.g., 2 bits) using a simple round function, completely eliminating the need to compute or store scale factors.
3. PyTorch Implementation of PolarQuant
import torch
def cartesian_to_polar_pairs(x):
"""Converts pairs of Cartesian coordinates to Polar."""
x_even = x[..., 0::2]
x_odd = x[..., 1::2]
r = torch.sqrt(x_even**2 + x_odd**2)
theta = torch.atan2(x_odd, x_even)
return r, theta
def polarquant_compress(x, bits=2):
"""
Recursively transforms a d-dimensional vector into
1 global radius and d-1 angles, then quantizes the angles.
"""
assert x.shape[-1] % 2 == 0, "Dimension must be divisible by 2"
current_r = x
all_angles = []
# Recursive transformation
while current_r.shape[-1] > 1:
current_r, theta = cartesian_to_polar_pairs(current_r)
# Quantize angles to fixed bins between -pi and pi
# Number of bins = 2^bits (e.g., 4 bins for 2-bit)
bins = 2**bits
# Normalize to [0, 1], scale to bins, round, scale back
theta_norm = (theta + torch.pi) / (2 * torch.pi)
theta_quantized = torch.round(theta_norm * (bins - 1)) / (bins - 1)
theta_dequantized = (theta_quantized * 2 * torch.pi) - torch.pi
all_angles.append(theta_dequantized)
global_radius = current_r
return global_radius, all_angles
# Note: Decompression is just the inverse: x = r * cos(theta), y = r * sin(theta)
Part 2: Quantized Johnson-Lindenstrauss (The 1-Bit Error Checker)
Even with optimal angle binning, PolarQuant introduces a small quantization error. Let the original Key vector be $k$, and the decompressed PolarQuant vector be $\hat{k}$. The residual error is: $$e = k - \hat{k}$$
During attention, we compute the dot product between a Query $q$ and Key $k$: $$q^T k = q^T (\hat{k} + e) = q^T \hat{k} + q^T e$$
We have $q^T \hat{k}$ perfectly, but we are missing $q^T e$. Storing $e$ in full precision destroys our compression gains. TurboQuant solves this using a 1-bit Quantized Johnson-Lindenstrauss (QJL) projection.
1. The Mathematics of QJL and Dot-Product Recovery
The JL Lemma guarantees that projecting high-dimensional data into a lower-dimensional space using a random Gaussian matrix $P \in \mathbb{R}^{m \times d}$ (where $P_{i,j} \sim \mathcal{N}(0, 1/m)$) preserves Euclidean geometry.
QJL takes this further by taking the random projection and aggressively compressing it down to just its sign ($+1$ or $-1$, requiring exactly 1 bit per dimension): $$s_e = \text{sign}(P e)$$
How do we compute the dot product $q^T e$ using only 1-bit signs? We project the Query vector using the same random matrix $P$, but keep it in high precision: $$y_q = P q$$
Because of the statistical properties of Gaussian projections and the arcsin law (Grothendieck’s identity), the dot product of a high-precision projection ($y_q$) and a 1-bit sign projection ($s_e$) yields an unbiased estimator of the original high-dimensional dot product, multiplied by a scaling constant $\alpha$: $$\mathbb{E}[y_q^T s_e] \approx \alpha (q^T e)$$
2. PyTorch Implementation of QJL Error Correction
import math
def qjl_compress_error(error_vector, projection_matrix):
"""
Compresses the residual error to a 1-bit sign vector.
projection_matrix shape: (m, d)
error_vector shape: (..., d)
"""
# Project the error into lower dimensional space
projected_error = torch.matmul(error_vector, projection_matrix.T)
# Extract only the sign (+1 or -1) -> Costs 1 bit per value!
sign_bits = torch.sign(projected_error)
# Handle exact zeros by mapping to 1
sign_bits[sign_bits == 0] = 1
return sign_bits
def qjl_approximate_dot_product(query, sign_bits, projection_matrix):
"""
Approximates the dot product q^T e using only the 1-bit error signs.
"""
# 1. Project the high-precision Query using the SAME matrix
projected_query = torch.matmul(query, projection_matrix.T)
# 2. Compute dot product in the projected space
# (Notice we are doing a highly efficient float-to-bit dot product)
approx_dot = torch.matmul(projected_query, sign_bits.transpose(-2, -1))
# 3. Apply the mathematical scaling constant (derived from Gaussian integrals)
# alpha roughly equals sqrt(2/pi)
alpha = math.sqrt(2 / math.pi)
return approx_dot * alpha
Part 3: TurboQuant Synthesis (Putting it Together)
TurboQuant fuses these two systems in a hardware-aware pipeline. When a new Key vector $k$ is generated during LLM generation, the system does the following:
- Main Signal: $k$ is passed through PolarQuant, yielding a global radius and a set of $d-1$ quantized angles. This utilizes roughly 2 bits per dimension.
- Error Calculation: The system instantly decompresses the angles back to Cartesian to find the reconstructed $\hat{k}$ and computes the residual $e = k - \hat{k}$.
- Residual Compression: The residual $e$ is multiplied by a pre-computed random matrix $P$ and signed. This costs 1 bit per dimension.
- Storage: The KV cache stores the 2-bit angles and the 1-bit signs. Total size: 3 bits per parameter (zero scaling factors).
When computing the Attention Matrix ($QK^T$) during generation:
def turboquant_attention_score(query, compressed_key_angles, global_radius, sign_bits, proj_matrix):
# 1. Reconstruct main Key from angles (Fast math operation)
k_hat = polarquant_decompress(compressed_key_angles, global_radius)
# 2. Compute main attention score q^T \hat{k}
main_score = torch.matmul(query, k_hat.transpose(-2, -1))
# 3. Recover the bias q^T e using 1-bit QJL
error_correction = qjl_approximate_dot_product(query, sign_bits, proj_matrix)
# 4. Final perfectly corrected Attention Logit
final_attention_score = main_score + error_correction
# Apply standard Transformer scaling (1 / sqrt(d_k)) and Softmax
d_k = query.shape[-1]
return torch.softmax(final_attention_score / math.sqrt(d_k), dim=-1)
By separating the magnitude/direction mapping (PolarQuant) from the residual geometric preservation (QJL), TurboQuant breaks the traditional barrier of quantization, pushing massive LLMs into latency and memory regimes previously thought impossible.