How Lloyd-Max Scalar Quantization Achieves Near-Optimal Distortion in TurboVec

In TurboVec, Lloyd-Max scalar quantization achieves near-optimal distortion by pre-computing optimal boundaries and centroids for a known Beta marginal distribution via high-precision adaptive Simpson integration, then aligning real-world embeddings with that theoretical distribution using TQ+ calibration.

Lloyd-Max scalar quantization is the core compression mechanism behind TurboVec (RyanCodrai/turbovec), a Rust implementation of Google’s TurboQuant algorithm. Because the quantizer is pre-computed for a distribution known a priori—rather than learned per dataset—it minimizes mean-squared error across all vectors without online training. The following sections trace the exact source code paths that turn this theoretical guarantee into the low-distortion results reported in the repository benchmarks.

Why a Known Distribution Matters: The Beta Marginal

After normalizing each vector to unit length, TurboVec applies a single random orthogonal rotation to the entire dataset. This rotation makes every coordinate independent and forces its marginal distribution to converge, in high dimensions, to a Beta((d-1)/2, (d-1)/2) distribution on [-1, 1]. Crucially, this convergence holds independently of the original data.

Because the distribution is known in advance, TurboVec computes an optimal scalar quantizer once and reuses it for all vectors, eliminating the need for per-dataset k-means training. The formal statement of this Beta marginal appears at lines 3-5 of turbovec/src/codebook.rs, while the README covers the rotation-to-Beta mapping at lines 90-96.

Lloyd-Max Quantizer Implementation in codebook.rs

The Lloyd-Max algorithm finds the boundaries and centroids that minimize mean-squared quantization error for a given probability density. In TurboVec, this algorithm is implemented entirely in turbovec/src/codebook.rs.

Entry Point and Beta Parameters

The function pub(crate) fn codebook(bits, dim) -> (Vec<f32>, Vec<f32>) at lines 17-19 serves as the entry point called from the index constructor (TurboQuantIndex::new). It first derives the Beta shape parameter:

let a = (dim as f64 - 1.0) / 2.0;

This a defines the symmetric Beta distribution that each rotated coordinate follows. Lines 28-33 initialize centroids uniformly inside ±3 standard deviations of the Beta, providing a high-quality starting point for iterative refinement.

Iterative Refinement via Conditional Means

The core optimization loop runs at lines 35-87:

for _ in 0..max_iter {
    // recompute boundaries as midpoints, then update centroids
}

Inside the loop, TurboVec updates each centroid to the conditional mean of the Beta over its interval. At lines 60-73, the code evaluates:

adaptive_simpson(...) / prob

The adaptive Simpson integration routine (lines 97-133 in the same file) computes the exact integral of x * p(x) over each bucket, divided by the bucket probability. This guarantees that every centroid is the true conditional mean—the optimal reconstruction value for minimizing MSE—even for the tiny intervals that appear at 2-bit or 4-bit quantization.

Convergence Guarantee

The iteration halts when centroids move less than 1e-12:

if max_change < tol { break; }

Line 84 enforces this strict tolerance. Upon convergence, the function returns boundaries and centroids_f32 (lines 89-94), which are stored once per index and applied to every vector.

TQ+ Calibration in encode.rs

In real-world embeddings, the empirical Beta marginal can drift—especially at low bit-widths. TurboVec closes this gap with TQ+ calibration in turbovec/src/encode.rs (lines 10-25).

The routine compute_tqplus_calibration (lines 60-84) fits an affine transform (x + shift) * scale per coordinate by mapping the empirical 5% and 95% quantiles onto the canonical Beta quantiles. This calibration runs once on the first batch of vectors and then freezes for future additions, keeping the index online. After calibration, vectors are quantized with the same pre-computed Lloyd-Max codebook, so the quantizer remains near-optimal while adapting to dataset-specific skew.

Four Design Choices That Keep Distortion Near-Optimal

Several tightly coupled mechanisms in RyanCodrai/turbovec work together to bound quantization distortion:

• Known theoretical distribution: Optimal bucket boundaries are exact for the Beta density rather than approximate cluster centers learned from data.
• Adaptive Simpson integration: Centroids are true conditional means of the Beta, not sample averages, ensuring the MSE for each bucket is minimized.
• TQ+ calibration: Aligns empirical data with the target Beta, eliminating bias that would otherwise inflate distortion at aggressive bit-rates.
• Length renormalization: The scale_from_inner routine in encode.rs (lines 98-104) corrects the systematic inner-product bias introduced by scalar quantization, removing the final error term that would degrade nearest-neighbor recall.

Together, these design points keep TurboVec within approximately 2.7× the Shannon distortion-rate limit, yielding roughly 0.2% recall loss versus FAISS PQ at 4-bit compression.

Code Examples

Python (High-Level)

from turbovec import TurboQuantIndex

# Build a 1536-dim, 4-bit TurboQuant index

index = TurboQuantIndex(dim=1536, bit_width=4)

# Add vectors; first add triggers TQ+ calibration and Lloyd-Max codebook generation

index.add(vectors)

# Search

scores, ids = index.search(query, k=10)

The Python bindings call directly into the Rust library where codebook.rs generates the Lloyd-Max codebook.

Rust (Low-Level)

use turbovec::TurboQuantIndex;

let mut idx = TurboQuantIndex::new(1536, 4).unwrap();   // creates codebook in codebook.rs
idx.add(&vectors).unwrap();
let (scores, ids) = idx.search(&queries, 10);

The TurboQuantIndex::new constructor invokes codebook(bits, dim), which internally runs the Lloyd-Max iteration described at lines 17-95 of turbovec/src/codebook.rs.

Inspecting the Generated Codebook

let (boundaries, centroids) = turbovec::codebook::codebook(4, 1536);
println!("Boundaries: {:?}", boundaries);
println!("Centroids:  {:?}", centroids);

This function lives in turbovec/src/codebook.rs and exports the pre-computed quantization tables used by the encoder.

Summary

• Lloyd-Max scalar quantization in TurboVec is pre-computed for a Beta((d-1)/2, (d-1)/2) marginal that every coordinate follows after random orthogonal rotation.
• The codebook.rs module iteratively refines boundaries and centroids using adaptive Simpson integration to evaluate exact conditional means, stopping at a tolerance of 1e-12.
• TQ+ calibration in encode.rs aligns real-world embeddings with the theoretical Beta distribution using a frozen per-coordinate affine transform.
• Length renormalization corrects quantization-induced inner-product bias, preserving recall.
• Together these techniques achieve near-optimal distortion without per-dataset training.

Frequently Asked Questions

What is the role of the Beta distribution in TurboVec's quantization?

After a random orthogonal rotation, every normalized coordinate converges to a Beta((d-1)/2, (d-1)/2) marginal. Because this distribution is known in advance, TurboVec can compute a globally optimal Lloyd-Max codebook once and apply it to all vectors, eliminating the need for dataset-specific training.

How does the Lloyd-Max algorithm know when to stop iterating?

The implementation in turbovec/src/codebook.rs stops when the maximum centroid movement across all buckets falls below 1e-12. This strict tolerance ensures that the returned boundaries and centroids are numerically indistinguishable from the true fixed point of the Lloyd-Max equations.

Why does TurboVec need TQ+ calibration if the distribution is already known?

Real-world embedding dimensions are finite, so empirical coordinate distributions can deviate slightly from the theoretical Beta. TQ+ calibration in encode.rs fits a single affine shift-and-scale per coordinate to map empirical quantiles onto canonical Beta quantiles, removing drift without retraining the codebook.

How does length renormalization improve search recall?

Scalar quantization biases inner-product estimates because quantized coordinates have different expected magnitudes than the originals. The scale_from_inner function in encode.rs corrects this systematic bias, restoring accurate distance estimates and preventing the recall drops that would otherwise appear at aggressive bit-widths.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}