IVF-PQ

Deep dive

How it works

Parts:

• coarse quantizer — nlist k-means centroids that partition the space
• product-quantizer codebooks — m subspaces × 256 sub-centroids, one byte per code
• inverted lists — per centroid, contiguous (id, m-byte code) pairs
• lookup table (LUT) — per query and probed list, m × 256 partial distances
• re-ranker (optional) — exact distances on the best k′ candidates from full vectors

Operations: train, add (below) and search with nprobe lists.

Train (on a sample, e.g. 30-256 points per centroid):

1. Coarse k-means: nlist centroids C.
2. Residuals r = x − C[assign(x)].
3. Split each residual into m sub-vectors of dim/m; per subspace, k-means with 256 centroids (one byte per code).

Add x: list l = nearest centroid; code[s] = nearest codeword of subspace s to the residual's s-th part; append (id, code) to list l.

Search q:

1. Distances to all nlist centroids; take the nprobe nearest.
2. Per probed list: residual q − C[l]; lookup table LUT[s][j] = ‖(q − C[l])_s − book_s[j]‖² (m × 256 floats).
3. For each code in the list: d ≈ Σ_s LUT[s][code[s]] (asymmetric distance: exact query, quantized database). Keep the top k with a heap (top-k selection).
4. Optional: re-rank the best k′ (100-1000) with full vectors from RAM or SSD.

x (32 floats) - centroid -> residual   [r0 r1 r2 r3 | r4 .. r7 | ... | r28 .. r31]   m = 8 subspaces
                                         ↓ nearest of 256        ↓                      ↓
code (8 bytes)                          [  17       |   203     | ... |   88     ]
distance = LUT[0][17] + LUT[1][203] + ... + LUT[7][88]        (8 lookups, LUT = 8 KB, in L1)

Code

use std::collections::BinaryHeap;

fn l2(a: &[f32], b: &[f32]) -> f32 { a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum() }

fn nearest(cents: &[f32], dim: usize, v: &[f32]) -> usize {
    (0..cents.len() / dim).min_by(|&i, &j| {
        l2(&cents[i * dim..][..dim], v).total_cmp(&l2(&cents[j * dim..][..dim], v))
    }).unwrap()
}

/// Lloyd's k-means; returns k × dim centroids. Init: k evenly spaced samples (use k-means++ in practice).
pub fn kmeans(data: &[f32], dim: usize, k: usize, iters: usize) -> Vec<f32> {
    let n = data.len() / dim;
    assert!(n >= k);
    let mut c: Vec<f32> = (0..k).flat_map(|j| data[j * n / k * dim..][..dim].to_vec()).collect();
    for _ in 0..iters {
        let (mut sum, mut cnt) = (vec![0f32; k * dim], vec![0u32; k]);
        for x in data.chunks_exact(dim) {
            let a = nearest(&c, dim, x);
            cnt[a] += 1;
            for (s, v) in sum[a * dim..][..dim].iter_mut().zip(x) { *s += v; }
        }
        for j in (0..k).filter(|&j| cnt[j] > 0) { // empty cluster keeps its old centroid
            for d in 0..dim { c[j * dim + d] = sum[j * dim + d] / cnt[j] as f32; }
        }
    }
    c
}

pub struct IvfPq {
    dim: usize, m: usize, dsub: usize,
    coarse: Vec<f32>,             // nlist × dim
    books: Vec<f32>,              // m codebooks × 256 centroids × dsub
    lists: Vec<(Vec<u32>, Vec<u8>)>, // per list: ids and m-byte codes
}

impl IvfPq {
    pub fn train(sample: &[f32], dim: usize, nlist: usize, m: usize) -> Self {
        assert!(dim % m == 0);
        let dsub = dim / m;
        let coarse = kmeans(sample, dim, nlist, 10);
        let resid: Vec<f32> = sample.chunks_exact(dim).flat_map(|x| {
            let c = &coarse[nearest(&coarse, dim, x) * dim..][..dim];
            x.iter().zip(c).map(|(a, b)| a - b).collect::<Vec<_>>()
        }).collect();
        let mut books = Vec::with_capacity(m * 256 * dsub);
        for s in 0..m { // one 256-centroid codebook per subspace, trained on residuals
            let sub: Vec<f32> = resid.chunks_exact(dim).flat_map(|r| r[s * dsub..][..dsub].to_vec()).collect();
            books.extend(kmeans(&sub, dsub, 256, 10));
        }
        IvfPq { dim, m, dsub, coarse, books, lists: vec![(Vec::new(), Vec::new()); nlist] }
    }

    pub fn add(&mut self, id: u32, x: &[f32]) {
        let (d, ds) = (self.dim, self.dsub);
        let l = nearest(&self.coarse, d, x);
        let c = &self.coarse[l * d..][..d];
        for s in 0..self.m {
            let r: Vec<f32> = (0..ds).map(|i| x[s * ds + i] - c[s * ds + i]).collect();
            let code = nearest(&self.books[s * 256 * ds..][..256 * ds], ds, &r) as u8;
            self.lists[l].1.push(code);
        }
        self.lists[l].0.push(id);
    }

    /// Approximate k nearest by asymmetric distance (exact query, quantized database).
    pub fn search(&self, q: &[f32], k: usize, nprobe: usize) -> Vec<(f32, u32)> {
        let (d, ds, m) = (self.dim, self.dsub, self.m);
        let mut order: Vec<(f32, usize)> = self.coarse.chunks_exact(d).map(|c| l2(q, c)).zip(0..).collect();
        order.sort_unstable_by(|a, b| a.0.total_cmp(&b.0));
        let mut lut = vec![0f32; m * 256];
        let mut heap: BinaryHeap<(u32, u32)> = BinaryHeap::new(); // max-heap of (dist bits, id), size k
        for &(_, l) in order.iter().take(nprobe) {
            let c = &self.coarse[l * d..][..d];
            for s in 0..m { // LUT: distance from the query residual to every codeword of subspace s
                let r: Vec<f32> = (0..ds).map(|i| q[s * ds + i] - c[s * ds + i]).collect();
                for j in 0..256 { lut[s * 256 + j] = l2(&r, &self.books[(s * 256 + j) * ds..][..ds]); }
            }
            let (ids, codes) = &self.lists[l];
            for (&id, code) in ids.iter().zip(codes.chunks_exact(m)) {
                let dist: f32 = code.iter().enumerate().map(|(s, &c)| lut[s * 256 + c as usize]).sum();
                let key = dist.to_bits(); // non-negative f32: bit order = numeric order
                if heap.len() < k { heap.push((key, id)); }
                else if key < heap.peek().unwrap().0 { *heap.peek_mut().unwrap() = (key, id); }
            }
        }
        heap.into_sorted_vec().into_iter().map(|(b, id)| (f32::from_bits(b), id)).collect()
    }
}

Complexity and constants

Measured (code above; 100k clustered 32-dim vectors; nlist = 256; trained on 20k; 200 held-out queries; one Apple M-series core). Brute force: 513 µs per query.

• Scan cost: m table lookups + adds per vector, order of a few ns per vector in scalar code like the above. FAISS "fast scan" (4-bit codes, LUTs held in SIMD registers and looked up with pshufb) is order of 4-6× faster (André et al.).
• LUT build: 256 × dim multiply-adds per probed list (8 KB of floats for m = 8). With large nprobe and short lists it can dominate; for L2 it can be decomposed into precomputed per-list terms.
• Training is the slow part: k-means is O(n_sample × k × dim × iters). Here 5 s for m = 8 and 14 s for m = 16; GPU or multithreaded k-means is standard.

Efficiency

• Cost: m bytes plus an id per vector (768-dim float: 3 KB → 64 B). Scanning is m table lookups and adds per code, order of a few ns per vector in scalar code. Measured with the code above (100k clustered 32-dim vectors, nlist = 256, one Apple M-series core): brute force 513 µs per query; m = 8, nprobe = 1: 10 µs at recall@10 0.41, 0.72 with re-rank of the top 100; m = 16, nprobe = 4: 37 µs, 0.69 and 0.97 with re-rank. Bound: memory bandwidth over the codes, LUT build for large nprobe.
• Levers:
  • re-rank the top 100-1000 with full vectors — recall 0.69 → 0.97 at m = 16, nprobe 4 (measured above) — whenever codes alone are too lossy
  • 4-bit codes with in-register LUTs (fast scan) — order of 4-6× faster scan (André et al. 2015) — CPU search with SIMD
  • residual encoding — better accuracy than encoding raw vectors, one LUT per probed list — default
  • OPQ rotation — a few points of recall at no query cost (Ge et al. 2014) — correlated dimensions
  • GPU IVF-PQ (FAISS, cuVS) — billion-scale search on GPUs (Johnson et al.) — batch queries, data fits GPU memory or shards
• Naive → standard → state of the art:
  • brute-force scan of float vectors · 513 µs per query on 100k × 32-dim (measured above)
  • IVF-PQ, 8-bit codes, re-rank · 10-37 µs at recall 0.72-0.97 (measured above)
  • fast-scan 4-bit codes or GPU IVF-PQ with re-rank · order of 4-6× faster scan (André et al.), billion scale (Johnson et al.)

Implementation notes

• Residual encoding (as above) gives much better accuracy than encoding raw vectors, at the cost of one LUT per probed list.
• OPQ: learn a rotation before PQ so that variance is balanced across subspaces; a few points of recall for free at query time.
• 4-bit PQ (fast scan): 16 centroids per subspace, codes packed in blocks of 32 vectors, and the LUT quantized to u8 so pshufb does 32 lookups per instruction.
• Lists as contiguous arrays of codes and ids (structure of arrays): the scan is a sequential stream, which is prefetcher-friendly.
• Parameter defaults: m such that dim/m is 2-8 (e.g. m = 64 for 768-dim), 8 bits per code, nprobe 8-64 for nlist in the thousands; tune nprobe for recall.
• Re-ranking: store full vectors (or int8 vectors) separately; re-rank k′ = 10·k candidates.
• Training sample: FAISS warns below 39 points per centroid; use at least about 30-50× nlist samples and 256 × 39 ≈ 10k for the codebooks.
• Inner product / cosine: normalize vectors, or use IP-specific LUTs; IVF assignment should use the same metric.

Pitfalls

• Reporting recall of PQ distances against PQ ground truth: always compare with exact brute-force neighbours.
• Too few training points or training on a non-representative sample: empty or skewed clusters, and some lists get huge, so queries that hit them are slow.
• Unbalanced lists (popular regions) make latency unpredictable; cap list sizes or split big clusters.
• k-means init by taking the first k points (as in this compact code) is fragile on sorted data; use k-means++ or random sampling.
• Forgetting that m must divide dim; pad with zeros otherwise.
• After heavy deletes and inserts the codebooks no longer fit the data; schedule retraining and re-encoding.

Variants and alternatives

In production

ML retrieval: FAISS (IndexIVFPQ, IndexIVFPQFastScan, GPU IVF-PQ), NVIDIA cuVS (formerly RAFT) IVF-PQ, Databases and search services: Milvus (IVF_PQ index type), LanceDB (IVF-PQ as its main vector index). Related: Google's ScaNN, which uses partitioning plus quantization in the same spirit.

Implementations

Sources

• H. Jégou, M. Douze, C. Schmid, Product Quantization for Nearest Neighbor Search, IEEE TPAMI, 2011. paper
• T. Ge, K. He, Q. Ke, J. Sun, Optimized Product Quantization, IEEE TPAMI, 2014. paper
• J. Johnson, M. Douze, H. Jégou, Billion-scale Similarity Search with GPUs, IEEE Transactions on Big Data, 2019. paper · arXiv · code
• F. André, A.-M. Kermarrec, N. Le Scouarnec, Cache Locality is not Enough: High-Performance Nearest Neighbor Search with Product Quantization Fast Scan, VLDB, 2015. paper
• R. Guo, P. Sun, E. Lindgren, Q. Geng, D. Simcha, F. Chern, S. Kumar, Accelerating Large-Scale Inference with Anisotropic Vector Quantization, ICML, 2020. arXiv · code