Product Quantization (PQ) Encode + ADC

Background

Product Quantization (PQ) splits each vector into $M$ subspaces, learns $K$ centroids per subspace (typically $K = 256$ → 8-bit codes), and encodes each db vector as $M$ small codes. Asymmetric Distance Computation (ADC) then approximates the distance from query to db vector via $M$ precomputed lookup-table sums. Foundation of IVF-PQ and FAISS.

Problem statement

Implement two functions:

• pq_encode(vecs, codebooks) — for each row, split into $M$ subspaces; encode each subspace as the index of its nearest centroid. Return shape (n, M) of ints.
• pq_distance(query, codes, codebooks) — build the lookup table $\text{LUT}[m, k] = \|q_m - c_{m, k}\|^2$. Return per-db-row distance $\sum_m \text{LUT}[m, \text{codes}[i, m]]$.

codebooks has shape $(M, K, d/M)$.

Input

• vecs — np.ndarray shape $(n, d)$.
• codebooks — np.ndarray shape $(M, K, d/M)$.
• query — 1-D np.ndarray of length $d$.
• codes — np.ndarray shape $(n, M)`, integer codes.

Output

• pq_encode returns np.ndarray shape $(n, M)ofint`.
• pq_distance returns np.ndarray of length $n$ of float.

Examples

Example 1 — code shape

Input:  vecs (10, 8); codebooks (4, 16, 2)   # M=4, K=16, d/M=2
Output: codes shape (10, 4)

Example 2 — ADC = 0 when query exactly matches a centroid-aligned db row

Input:  db row built by concatenating centroid 0 from each subspace
Output: pq_distance(that row's vec, codes, codebooks)[0] ≈ 0

Example 3 — ADC is non-negative

Input:  any random query and codes
Output: all distances ≥ 0

Constraints

• $d$ must be divisible by $M$.
• Distances are squared L2 (not Euclidean) — same ordering, cheaper to compute.
• Returns np.ndarray.

Notes

• Why asymmetric. "Asymmetric" because the query is not quantised; only the db is. Symmetric distance computation quantises both, which is faster but loses recall.
• IVF-PQ. PQ is typically combined with IVF: IVF picks candidate cells, PQ scores candidates within those cells. Billion-scale search on a single machine.