Lab VI3: The Quantization Math

Background

The intensive Part 3 wrote out the arithmetic of the three quantizers: scalar (per-dimension min/max → int), product (split into m subvectors, per-subspace k-means codebooks, asymmetric distance via lookup tables), and binary (sign bit + Hamming). Lab VI2 used the product quantizer inside a faiss IndexIVFPQ — but it was a black box behind JNI. This lab makes you implement all three yourself in plain Java, so the math stops being abstract: you encode, decode, measure the reconstruction error, compute the compression ratio, and measure recall@k against full precision — and you watch a rescore pass recover the recall you lost.

When you have written ADC by hand, you will never again be confused about why PQ needs trained codebooks, why BQ always rescores, or what compression_level: 32x is doing under the hood.

Why This Matters for Contributors

Over-quantization is one of the most common k-NN recall complaints: someone sets compression_level: 16x or an aggressive PQ m, recall tanks, and they file a bug. The fix is almost always "add/strengthen rescoring" or "back off the compression" — but to reason about it you need to know what each quantizer throws away and how rescoring buys it back. The map at the end of this lab — your hand-rolled quantizers ↔ k-NN compression_level/on_disk ↔ Lucene's Lucene99/Lucene104 scalar-quantized formats — is the same map you reach for when triaging a real recall regression.

Prerequisites

• Java 21+ (java -version). No Lucene or OpenSearch needed — this is pure Java math.
• You've read the intensive Part 3 (the SQ/PQ/BQ math) and the quantization and disk-ANN chapter (the menu, on_disk, compression_level).

mkdir -p /tmp/vi3 && cd /tmp/vi3

Note: These implementations are deliberately small and exact-to-the-formula, not optimized — the point is to see the arithmetic, not to beat faiss. Reuse the Vectors helper from Lab VI1 (clustered data + L2

• brute-force top-k); it is reproduced inline below so this lab is standalone.

Step-by-Step Tasks

Step 0 — The shared harness

// Q.java — shared helpers: clustered data, L2, exact top-k, recall@k.
import java.util.*;

public final class Q {
  public static float[][] clustered(int n, int dim, int clusters, long seed) {
    Random r = new Random(seed);
    float[][] c = new float[clusters][dim];
    for (float[] x : c) for (int d = 0; d < dim; d++) x[d] = r.nextFloat() * 10f;
    float[][] out = new float[n][dim];
    for (int i = 0; i < n; i++) {
      float[] cc = c[r.nextInt(clusters)];
      for (int d = 0; d < dim; d++) out[i][d] = cc[d] + (float) r.nextGaussian();
    }
    return out;
  }
  public static double l2(float[] a, float[] b) {
    double s = 0; for (int i = 0; i < a.length; i++){ double d=a[i]-b[i]; s+=d*d; } return s;
  }
  public static int[] exactTopK(float[][] data, float[] q, int k) {
    Integer[] idx = new Integer[data.length];
    for (int i = 0; i < idx.length; i++) idx[i] = i;
    Arrays.sort(idx, Comparator.comparingDouble(i -> l2(data[i], q)));
    int[] o = new int[k]; for (int i = 0; i < k; i++) o[i] = idx[i]; return o;
  }
  // recall@k for an approximate scorer: returns candidate ids sorted best-first.
  public static double recall(float[][] data, float[][] queries, int k,
                              java.util.function.BiFunction<Integer,float[],Double> approxDist) {
    double sum = 0;
    for (float[] qv : queries) {
      Integer[] idx = new Integer[data.length];
      for (int i = 0; i < idx.length; i++) idx[i] = i;
      Arrays.sort(idx, Comparator.comparingDouble(i -> approxDist.apply(i, qv)));
      Set<Integer> approx = new HashSet<>();
      for (int i = 0; i < k; i++) approx.add(idx[i]);
      Set<Integer> exact = new HashSet<>();
      for (int id : exactTopK(data, qv, k)) exact.add(id);
      approx.retainAll(exact); sum += approx.size() / (double) k;
    }
    return sum / queries.length;
  }
}

Step 1 — Scalar quantization (float → int8, per-dim min/max)

Implement the int8 SQ from intensive §3.1: learn per-dimension [min,max], encode to byte, decode, and measure reconstruction error.

// ScalarQ.java
import java.util.*;

public class ScalarQ {
  final int dim, levels;          // levels = 2^bits - 1 (255 for int8, 15 for int4)
  final float[] min, max;
  ScalarQ(float[][] train, int bits) {
    dim = train[0].length; levels = (1 << bits) - 1;
    min = new float[dim]; max = new float[dim];
    Arrays.fill(min, Float.MAX_VALUE); Arrays.fill(max, -Float.MAX_VALUE);
    for (float[] v : train) for (int j = 0; j < dim; j++) {
      min[j] = Math.min(min[j], v[j]); max[j] = Math.max(max[j], v[j]);
    }
  }
  int[] encode(float[] v) {                 // ints 0..levels (store as byte for int8)
    int[] q = new int[dim];
    for (int j = 0; j < dim; j++) {
      float span = Math.max(1e-9f, max[j] - min[j]);
      int code = Math.round((v[j] - min[j]) / span * levels);
      q[j] = Math.max(0, Math.min(levels, code));     // clip into range
    }
    return q;
  }
  float[] decode(int[] q) {
    float[] v = new float[dim];
    for (int j = 0; j < dim; j++) {
      float span = max[j] - min[j];
      v[j] = min[j] + (q[j] / (float) levels) * span;
    }
    return v;
  }
  public static void main(String[] a) {
    int n = 5000, dim = 64, k = 10;
    float[][] data = Q.clustered(n, dim, 20, 42L);
    float[][] q    = Q.clustered(200, dim, 20, 7L);
    for (int bits : new int[]{8, 4}) {
      ScalarQ sq = new ScalarQ(data, bits);
      int[][] codes = new int[n][];
      double err = 0;
      for (int i = 0; i < n; i++) {
        codes[i] = sq.encode(data[i]);
        err += Math.sqrt(Q.l2(data[i], sq.decode(codes[i])));    // L2 reconstruction error
      }
      // recall: distance from query (full precision) to DECODED database vector.
      double rec = Q.recall(data, q, k, (i, qv) -> Q.l2(qv, sq.decode(codes[i])));
      double ratio = 32.0 / bits;                                 // vs float32 (4 bytes=32 bits)
      System.out.printf("SQ int%d: ratio=%.0fx  avg_recon_err=%.4f  recall@%d=%.3f%n",
                        bits, ratio, err / n, k, rec);
    }
  }
}

javac Q.java ScalarQ.java && java ScalarQ

Observe: int4 has ~16× coarser steps than int8, so its reconstruction error is larger and recall lower — exactly the §3.1 prediction. This is the math behind Lucene's Lucene99/Lucene104ScalarQuantizedVectorsFormat.

Step 2 — Product quantization (codebooks + ADC)

Now the centerpiece: split into m subvectors, learn a 256-centroid codebook per subspace with k-means, encode to m bytes, and compute distances with asymmetric distance computation (intensive §3.2).

// ProductQ.java
import java.util.*;

public class ProductQ {
  final int dim, m, sub, k;       // m subvectors of length sub; k centroids per subspace
  final float[][][] codebooks;    // [m][k][sub]

  ProductQ(float[][] train, int m, int nbits) {
    this.dim = train[0].length; this.m = m; this.sub = dim / m; this.k = 1 << nbits;
    if (dim % m != 0) throw new IllegalArgumentException("dim % m != 0");
    codebooks = new float[m][][];
    for (int s = 0; s < m; s++) codebooks[s] = kmeans(slice(train, s), k, 12, 1L + s);
  }
  // extract subspace s of every training vector.
  float[][] slice(float[][] data, int s) {
    float[][] out = new float[data.length][sub];
    for (int i = 0; i < data.length; i++)
      System.arraycopy(data[i], s * sub, out[i], 0, sub);
    return out;
  }
  // plain Lloyd k-means over sub-dim points -> k centroids.
  static float[][] kmeans(float[][] pts, int k, int iters, long seed) {
    Random r = new Random(seed);
    float[][] c = new float[k][];
    for (int i = 0; i < k; i++) c[i] = pts[r.nextInt(pts.length)].clone();
    for (int it = 0; it < iters; it++) {
      float[][] sum = new float[k][pts[0].length]; int[] cnt = new int[k];
      for (float[] p : pts) {
        int best = 0; double bd = Double.MAX_VALUE;
        for (int j = 0; j < k; j++){ double d=Q.l2(p,c[j]); if(d<bd){bd=d;best=j;} }
        cnt[best]++; for (int d = 0; d < p.length; d++) sum[best][d] += p[d];
      }
      for (int j = 0; j < k; j++) if (cnt[j] > 0)
        for (int d = 0; d < sum[j].length; d++) c[j][d] = sum[j][d] / cnt[j];
    }
    return c;
  }
  int[] encode(float[] v) {       // m codes, each a centroid id in 0..k-1
    int[] code = new int[m];
    for (int s = 0; s < m; s++) {
      float[] piece = Arrays.copyOfRange(v, s * sub, (s + 1) * sub);
      int best = 0; double bd = Double.MAX_VALUE;
      for (int j = 0; j < k; j++){ double d=Q.l2(piece,codebooks[s][j]); if(d<bd){bd=d;best=j;} }
      code[s] = best;
    }
    return code;
  }
  // ADC: build per-subspace LUT once per query, then dist = sum of m table lookups.
  double[][] queryTables(float[] query) {
    double[][] lut = new double[m][k];
    for (int s = 0; s < m; s++) {
      float[] qp = Arrays.copyOfRange(query, s * sub, (s + 1) * sub);
      for (int j = 0; j < k; j++) lut[s][j] = Q.l2(qp, codebooks[s][j]);
    }
    return lut;
  }
  double adc(double[][] lut, int[] code) {
    double d = 0; for (int s = 0; s < m; s++) d += lut[s][code[s]]; return d;
  }
  public static void main(String[] a) {
    int n = 5000, dim = 64, k = 10, nbits = 8;
    float[][] data = Q.clustered(n, dim, 20, 42L);
    float[][] queries = Q.clustered(200, dim, 20, 7L);
    for (int m : new int[]{8, 16}) {       // 8 subvectors -> 8 bytes; 16 -> 16 bytes
      ProductQ pq = new ProductQ(data, m, nbits);
      int[][] codes = new int[n][];
      for (int i = 0; i < n; i++) codes[i] = pq.encode(data[i]);
      // recall via ADC: build LUT per query, distance = sum of m lookups.
      double rec = Q.recall(data, queries, k, new java.util.function.BiFunction<>() {
        double[][] cached; float[] cachedQ;
        public Double apply(Integer i, float[] qv) {
          if (qv != cachedQ) { cached = pq.queryTables(qv); cachedQ = qv; }
          return pq.adc(cached, codes[i]);
        }
      });
      double ratio = (dim * 4.0) / m;       // float32 bytes / code bytes (nbits=8 -> 1 byte/sub)
      System.out.printf("PQ m=%-2d nbits=%d: code=%dB ratio=%.0fx recall@%d=%.3f%n",
                        m, nbits, m, ratio, k, rec);
    }
  }
}

javac Q.java ProductQ.java && java ProductQ

You just built ADC: the LUT is m·256 distances computed once per query, then each candidate distance is m byte-indexed lookups + an add. Larger m → finer compression, larger code, better recall — the §3.2 trade-off, measured.

Step 3 — Binary quantization (sign bit + Hamming)

The extreme: one bit per dimension, Hamming distance (intensive §3.3).

// BinaryQ.java
public class BinaryQ {
  final int dim;
  final float[] threshold;        // per-dim threshold (here: the per-dim mean)
  BinaryQ(float[][] train) {
    dim = train[0].length; threshold = new float[dim];
    for (float[] v : train) for (int j = 0; j < dim; j++) threshold[j] += v[j];
    for (int j = 0; j < dim; j++) threshold[j] /= train.length;
  }
  long[] encode(float[] v) {      // dim bits packed into longs
    long[] bits = new long[(dim + 63) / 64];
    for (int j = 0; j < dim; j++)
      if (v[j] > threshold[j]) bits[j >> 6] |= (1L << (j & 63));
    return bits;
  }
  static int hamming(long[] a, long[] b) {
    int d = 0; for (int i = 0; i < a.length; i++) d += Long.bitCount(a[i] ^ b[i]); return d;
  }
  public static void main(String[] s) {
    int n = 5000, dim = 64, k = 10;
    float[][] data = Q.clustered(n, dim, 20, 42L);
    float[][] queries = Q.clustered(200, dim, 20, 7L);
    BinaryQ bq = new BinaryQ(data);
    long[][] codes = new long[n][];
    for (int i = 0; i < n; i++) codes[i] = bq.encode(data[i]);
    double rec = Q.recall(data, queries, k,
        (i, qv) -> (double) hamming(bq.encode(qv), codes[i]));
    double ratio = (dim * 4.0) / (dim / 8.0);     // float32 bytes / (1 bit per dim)
    System.out.printf("BQ 1bit/dim: code=%dB ratio=%.0fx recall@%d=%.3f (PRE-rescore)%n",
                      dim / 8, ratio, k, rec);
  }
}

javac Q.java BinaryQ.java && java BinaryQ

BQ's recall is the lowest — you threw away all magnitude. That low number is not a bug; it is why BQ always pairs with rescoring (next step).

Step 4 — Rescoring: recover the lost recall

This is the universal fix. Use the cheap quantized scorer to shortlist oversample · k candidates, then re-rank that shortlist with full-precision L2. Show recall climb back toward 1.0.

// Rescore.java — wrap ANY approximate scorer with a full-precision rescore pass.
import java.util.*;

public class Rescore {
  // shortlist with approxDist, then re-rank the shortlist with exact L2.
  static double recallWithRescore(float[][] data, float[][] queries, int k, int oversample,
      java.util.function.BiFunction<Integer,float[],Double> approxDist) {
    double sum = 0;
    for (float[] qv : queries) {
      Integer[] idx = new Integer[data.length];
      for (int i = 0; i < idx.length; i++) idx[i] = i;
      // pass 1: cheap shortlist of oversample*k candidates.
      Arrays.sort(idx, Comparator.comparingDouble(i -> approxDist.apply(i, qv)));
      int cand = Math.min(data.length, oversample * k);
      Integer[] shortlist = Arrays.copyOf(idx, cand);
      // pass 2: exact rescore of the shortlist.
      Arrays.sort(shortlist, Comparator.comparingDouble(i -> Q.l2(data[i], qv)));
      Set<Integer> approx = new HashSet<>();
      for (int i = 0; i < k; i++) approx.add(shortlist[i]);
      Set<Integer> exact = new HashSet<>();
      for (int id : Q.exactTopK(data, qv, k)) exact.add(id);
      approx.retainAll(exact); sum += approx.size() / (double) k;
    }
    return sum / queries.length;
  }
  public static void main(String[] a) {
    int n = 5000, dim = 64, k = 10;
    float[][] data = Q.clustered(n, dim, 20, 42L);
    float[][] queries = Q.clustered(200, dim, 20, 7L);
    BinaryQ bq = new BinaryQ(data);
    long[][] codes = new long[n][];
    for (int i = 0; i < n; i++) codes[i] = bq.encode(data[i]);
    java.util.function.BiFunction<Integer,float[],Double> bqDist =
        (i, qv) -> (double) BinaryQ.hamming(bq.encode(qv), codes[i]);
    for (int os : new int[]{1, 4, 16, 64}) {
      double rec = recallWithRescore(data, queries, k, os, bqDist);
      System.out.printf("BQ + rescore oversample=%-3d recall@%d=%.3f%n", os, k, rec);
    }
  }
}

javac Q.java BinaryQ.java Rescore.java && java Rescore

Watch recall rise with the oversample factor: the cheap binary pass casts a wide net, the exact rescore re-ranks it. This is precisely what on_disk mode + oversample_factor does, and why BQ needs a wider funnel than PQ (intensive §3.3).

Step 5 — Map to k-NN and Lucene

Put your hand-rolled quantizers next to the productized settings:

# Confirm the keys your k-NN version uses (don't trust this doc — grep):
cd ~/src/k-NN
grep -rn "fp16\|ENCODER_SQ\|ENCODER_PQ\|code_size\|nbits\|hamming\|compression_level\|oversample" \
  src/main/java/org/opensearch/knn/common/KNNConstants.java
# And the Lucene SQ formats your bundled Lucene ships:
grep -rln "Lucene9.*ScalarQuantizedVectorsFormat\|Lucene10.*ScalarQuantizedVectorsFormat" \
  ~/src/lucene/lucene/core/src/java 2>/dev/null

The full menu and the on_disk/compression_level mechanics live in the quantization and disk-ANN chapter; the Lucene SQ formats in the Lucene HNSW chapter.

Deliverables

• ScalarQ output: int8 and int4 reconstruction error, compression ratio, recall@10.
• ProductQ output: code size, compression ratio, recall@10 for two m values, with ADC (LUT) distances.
• BinaryQ output: 1-bit code size, 32× ratio, and the (low) pre-rescore recall.
• Rescore output: recall@10 climbing with oversample ∈ {1,4,16,64}.
• The completed quantizer ↔ k-NN ↔ Lucene mapping table (Step 5).

Troubleshooting

Expected Output

SQ int8: ratio=4x   avg_recon_err=0.0142  recall@10=0.97
SQ int4: ratio=8x   avg_recon_err=0.2280  recall@10=0.78

PQ m=8  nbits=8: code=8B  ratio=32x recall@10=0.74
PQ m=16 nbits=8: code=16B ratio=16x recall@10=0.88

BQ 1bit/dim: code=8B ratio=32x recall@10=0.31 (PRE-rescore)

BQ + rescore oversample=1   recall@10=0.31
BQ + rescore oversample=4   recall@10=0.67
BQ + rescore oversample=16  recall@10=0.93
BQ + rescore oversample=64  recall@10=0.99

Exact numbers depend on data and seed. What must hold: int8 ≫ int4 recall; PQ recall rises with m; BQ's pre-rescore recall is the lowest, and rescoring with a growing oversample factor recovers it toward 1.0.

Stretch Goals

• Symmetric vs asymmetric (SDC vs ADC). Add an SDC variant that also quantizes the query, and compare recall to ADC — confirm ADC is more accurate (intensive §3.2).
• PQ + IVF in code. Cluster the data into cells (a tiny k-means coarse quantizer), store PQ codes per cell, and search only the nearest nprobe cells — a from-scratch IndexIVFPQ mirroring Lab VI2's native one.
• Measure ADC cost. Time the per-query LUT build vs the per-candidate lookups and confirm the lookup pass dominates only when N is large — the reason ADC scales.
• Compression-vs-recall curve. Plot recall@10 against compression ratio across SQ8/SQ4/PQ/BQ (pre- and post-rescore), reproducing the intensive's §3.4 table with your own data.

Coding Exercises

These graduate your hand-rolled quantizers from print-statements into tested code, then push into the SIMD trick that makes PQ fast in production. Mix Java (the quantizer math) and Python (the quantization-error visualization that connects to Lab VI4).

1. (warm-up) A JUnit round-trip and monotonicity test for ScalarQ. Wrap Step 1 in a test asserting (a) decode(encode(v)) is within span/levels per dimension (the worst-case SQ error), and (b) int8 reconstruction error is strictly smaller than int4's on the same data. Print the two avg_recon_err values. This is the §3.1 "int4 is ~16× coarser" claim, machine-checked — the property a quantizer regression breaks.

2. (warm-up) Assert PQ's compression/recall trade-off. Add a test over ProductQ that asserts recall@10 is non-decreasing as m rises (m ∈ {4,8,16}, all dividing dim=64) and that the dim % m != 0 constructor throws IllegalArgumentException. Tie the recall curve back to the LUT size m·k you build per query in queryTables.

3. (core) SDC vs ADC, as a graded comparison. Implement the symmetric variant (SDC) that also quantizes the query before table lookup, deepening Stretch Goal 1 into code. Write a test asserting ADC's recall@10 ≥ SDC's on your clustered data (asymmetric keeps the query at full precision, so it must not lose). Print both. This is the precise reason faiss defaults to ADC for queries.

4. (core) Implement PQ4 fast-scan and validate it against your float-LUT ADC. In Java, re-encode with nbits=4 (so each codebook has k=16 centroids) and quantize the LUT to uint8 (scale each subspace's distances into 0..255), then score candidates with integer table lookups + saturating adds — the scalar twin of the _mm256_shuffle_epi8 fast-scan kernel. Write a test asserting your quantized-LUT ranking agrees with the float ADC ranking on the top-k for ≥ 90% of queries. You have now built — in plain Java — the exact design decision ("why only 16 centroids per sub-space?") that chapter explains: 16 bytes is one shuffle lane.

5. (core, Python) Visualize the quantization-error vs recall frontier. Reuse VI4's viz5_pq_cells.py idea but drive it from your numbers: write a Python script that, for SQ8/SQ4/PQ(m=8)/PQ(m=16)/BQ, plots recall@10 (pre- and post-rescore) against compression ratio — the §3.4 table as a curve (Stretch Goal 4, made concrete). Save viz_quant_frontier.png. Annotate where rescoring lifts BQ from "useless alone" to "competitive," matching your Rescore output.

6. (advanced challenge) A from-scratch IndexIVFPQ with an nprobe recall frontier. Build the full pipeline VI2 used natively: a coarse k-means quantizer (nlist cells), PQ codes stored per cell, search restricted to the nearest nprobe cells, then a full-precision rescore of the shortlist (reuse your Rescore.recallWithRescore). Emit a CSV (nlist,nprobe,oversample,recall@10) and write a test asserting recall is non-decreasing in both nprobe and oversample. Compare your recall numbers to the native IndexIVFPQ you built in Lab VI2 at matching params — they should land in the same neighbourhood, and any gap is k-means quality or fast-scan rounding. This is the whole quantization+IVF stack, owned end to end and benchmarked against the real engine.

Issues to Practice On

Over-quantization (aggressive compression_level, large PQ m, BQ without enough rescore) is one of the most-filed k-NN recall complaints — a great, well-scoped contributor area. The repo is opensearch-project/k-NN.

Labels drift — list and pick (gh label list --repo opensearch-project/k-NN). k-NN commonly uses good first issue, bug, enhancement, Roadmap.

Representative issue patterns. (1) "Recall tanked after I set compression_level: 16x / on_disk" — almost always "add/strengthen rescoring or back off compression." Reproduce by quantizing at the reported level (your SQ/PQ/BQ code shows exactly what is thrown away), rg the encoder + rescore plumbing (grep -rn "ENCODER_PQ\|compression_level\|oversample\|RescoreContext" src/main/java), and ship a test pinning post-rescore recall. (2) "PQ m / code_size validation is wrong or unclear" — the dimension % m == 0 and tens-per-centroid rules your code enforces; locate the validator and add the missing guard + a unit test on the rejection message.

Planted-bug exercise. In your ProductQ.queryTables, change the LUT to use the decoded database vector's subspace distance instead of the codebook-centroid distance (i.e. accidentally recompute against reconstructions). Recall barely moves on easy data but your Exercise 2 monotonicity test still passes — yet your Exercise 4 fast-scan agreement test goes red, because the LUT no longer matches the codes it indexes. Restore the centroid-distance LUT, then keep the agreement test as the regression that catches a corrupted lookup table. The lesson: quantization bugs hide behind "recall looks okay" — only a structural agreement assertion exposes them.

Etiquette. Claim an issue before working it, reproduce first, and every PR needs a test, a CHANGELOG.md entry, and a DCO sign-off (git commit -s). See community interaction.

Validation / Self-check

1. Write the int8 SQ encode/decode and the worst-case per-dimension error from your code. Why is int4's error ~16× int8's, and why does that show up as lower recall?
2. Explain ADC in your own words: what is in the per-query LUT, how big is it, and how many operations is a single candidate distance after the LUT is built?
3. Why does PQ need trained codebooks while SQ and BQ do not? Tie this to k-NN's _train API and the .opensearch-knn-models system index.
4. From your Rescore numbers: explain why BQ needs a larger oversample factor than PQ would to reach the same post-rescore recall.
5. Compute per-vector storage for d=1024 under int8, PQ(m=128, nbits=8), and 1-bit binary. Give each ratio and which require training.
6. Map each of your three quantizers to its k-NN configuration and its Lucene-engine equivalent. Which has no lucene-engine analogue, and why?

When this holds, you have implemented the entire quantization layer the intensive Part 3 describes and the Lab VI2 IndexIVFPQ used. Close the loop with the quantization and disk-ANN chapter (the productized on_disk/compression_level story) and the Lucene HNSW chapter (the Lucene99/Lucene104 SQ formats), and re-read the intensive to connect the math back to the graph and the index types.