Product Quantization (PQ)

Product Quantization (PQ) tackles the curse of dimensionality by compressing high-dimensional vectors into compact codes while preserving approximate distances. Let’s break it down using our toy dataset and formal math.

1. The Problem: Why PQ?

Storing and searching 1536-dimensional vectors (common in NLP/vision) is expensive. For 1B vectors:

• Raw storage: $1B\times 1536\times 4\text{ bytes}=6.144\text{ TB}$.
• Brute-force search: $O(nd)$ complexity is infeasible.

Solution: PQ compresses vectors into 8–16 bytes each and accelerates distance computations.

2. PQ Workflow

Step 1: Subspace Decomposition

Split a $d$-dimensional vector into $m$ disjoint subspaces:

Example: For our 2D toy vectors split into $m=2$ subspaces:

• $\mathbf{x}=[\underset{\text{Subspace 1}}{\underbrace{x_1}},\underset{\text{Subspace 2}}{\underbrace{x_2}}]$

Step 2: Codebook Training (Mathematical Deep Dive)

Let’s formalize the codebook training process with equations, using the example subspaces from your dataset.

1. K-means in Subspaces

For each subspace $j$, we solve:

where:

• $X^{(j)}$: Subvectors in subspace $j$.
• ${\mathcal{C}}_j=\{{\mathbf{c}}_{j,1},...,{\mathbf{c}}_{j,h}\}$: Centroids for subspace $j$.

2. Example: Subspace 1 (x₁)

Data: $X^{(1)}=\{1.1,1.0,5.2,5.1,5.3,9.0,9.5,8.9\}$ (1D subvectors).
Goal: Train $h=3$ centroids ${\mathcal{C}}_1=\{c_{1,1},c_{1,2},c_{1,3}\}$.

Phase 1: Initialization (k-means++)

1. Step 1: Randomly pick $c_{1,1}^{(0)}=1.1$.
2. Step 2: Compute distances $D(\mathbf{x})$ from all points to $c_{1,1}^{(0)}$:$$D=\{(1.1-1.1{)}^2,(1.0-1.1{)}^2,(5.2-1.1{)}^2,...,(8.9-1.1{)}^2\}=[0,0.01,16.81,...,60.84]$$
3. Step 3: Sample next centroid $c_{1,2}^{(0)}=9.5$ (far from $c_{1,1}^{(0)}$).
4. Step 4: Update $D(\mathbf{x})=min(\parallel \mathbf{x}-c_{1,1}^{(0)}{\parallel }^2,\parallel \mathbf{x}-c_{1,2}^{(0)}{\parallel }^2)$.
5. Step 5: Sample $c_{1,3}^{(0)}=5.2$ (midpoint between 1.1 and 9.5).

Initial Centroids: ${\mathcal{C}}_1^{(0)}=\{1.1,9.5,5.2\}$.

Phase 2: Lloyd’s Algorithm

Iteration 1:

1. Assignment:
   Assign each ${\mathbf{x}}^{(1)}$ to the nearest centroid:

2. Update:
   Recompute centroids as cluster means:$$c_{1,1}^{(1)}=\frac{1.1+1.0}{2}=1.05$$$$c_{1,2}^{(1)}=\frac{5.2+5.1+5.3}{3}=5.2$$$$c_{1,3}^{(1)}=\frac{9.0+9.5+8.9}{3}=9.13$$

Iteration 2:
Reassign subvectors to new centroids $\{1.05,5.2,9.13\}$:

• Assignments remain unchanged → convergence.

Final Centroids: ${\mathcal{C}}_1=\{1.05,5.2,9.13\}$.

3. Example: Subspace 2 (x₂)

Data: $X^{(2)}=\{0.9,1.2,5.0,4.8,5.1,1.0,0.8,1.2\}$.
Goal: Train $h=3$ centroids ${\mathcal{C}}_2=\{c_{2,1},c_{2,2},c_{2,3}\}$.

Phase 1: Initialization (k-means++)

1. Step 1: Randomly pick $c_{2,1}^{(0)}=1.2$.
2. Step 2: Compute distances $D(\mathbf{x})$ from all points to $c_{2,1}^{(0)}$.
3. Steps 3–5: Sample $c_{2,2}^{(0)}=5.1$ and $c_{2,3}^{(0)}=0.8$.

Phase 2: Lloyd’s Algorithm

After iterations, centroids converge to ${\mathcal{C}}_2=\{1.05,4.97,1.0\}$.

4. PQ Encoding (Mathematical Formulation)

For a vector $\mathbf{x}=[x_1,x_2]$:

1. Split: ${\mathbf{x}}^{(1)}=x_1$, ${\mathbf{x}}^{(2)}=x_2$.
2. Quantize:
   • Subspace 1: $k_1=\arg \underset{k}{min}\parallel x_1-c_{1,k}{\parallel }^2$
   • Subspace 2: $k_2=\arg \underset{k}{min}\parallel x_2-c_{2,k}{\parallel }^2$
3. PQ Code: $(k_1,k_2)$.

Example:
For $\mathbf{x}=[5.2,5.0]$:

• Subspace 1: $\parallel 5.2-1.05{\parallel }^2=17.22$, $\parallel 5.2-5.2{\parallel }^2=0$, $\parallel 5.2-9.13{\parallel }^2=15.40$ → $k_1=2$.
• Subspace 2: $\parallel 5.0-1.05{\parallel }^2=15.60$, $\parallel 5.0-4.97{\parallel }^2=0.0009$, $\parallel 5.0-1.0{\parallel }^2=16.00$ → $k_2=2$.
• PQ Code: $(2,2)$.

5. Lookup Table (LUT) Construction

For a query $\mathbf{q}=[q_1,q_2]$:

1. Subspace 1:$${\text{LUT}}_1[k]=\parallel q_1-c_{1,k}{\parallel }^2\text{for }k=1,...,h$$
2. Subspace 2:$${\text{LUT}}_2[k]=\parallel q_2-c_{2,k}{\parallel }^2\text{for }k=1,...,h$$

Example: For $\mathbf{q}=[5.4,5.2]$:

• Subspace 1 LUT:$${\text{LUT}}_1=[\parallel 5.4-1.05{\parallel }^2,\parallel 5.4-5.2{\parallel }^2,\parallel 5.4-9.13{\parallel }^2]=[18.92,0.04,13.58]$$
• Subspace 2 LUT:$${\text{LUT}}_2=[\parallel 5.2-1.05{\parallel }^2,\parallel 5.2-4.97{\parallel }^2,\parallel 5.2-1.0{\parallel }^2]=[17.22,0.05,17.64]$$

6. Distance Approximation

For a PQ code $(k_1,k_2)$:

Example:
For PQ code $(2,2)$:

7. Why This Works

• Training: Codebooks adapt to data distribution, minimizing reconstruction error.
• Encoding: Subspace independence allows parallel computation.
• Querying: LUTs reduce distance computation to $O(m)$ additions vs. $O(d)$ operations.

This mathematically rigorous process ensures IVF-PQ’s efficiency-accuracy balance!

3. Key Parameters & Trade-offs

Mathematical Trade-off

The quantization error for PQ is:

• Lower ${ϵ}_{\text{PQ}}$ means ADC distances better approximate true distances.

4. Example Walkthrough

1. Training:

   • Split toy dataset into subspaces.
   • Train codebooks with $h=3$ centroids per subspace.

2. Encoding:

   • Compress vectors into PQ codes (e.g., $(2,2)$ for $[5.2,5.0]$).

3. Querying:

   • Precompute LUTs for $\mathbf{q}=[5.4,5.2]$.
   • Compute approximate distances in 1 operation per PQ code vs. 2 operations for brute-force.

5. Why This Matters for Real-World Data

For 1B vectors in ${\mathbb{R}}^{1536}$:

• Storage: PQ reduces memory from 6.144 TB to 16 GB ($16\text{ bytes/vector}\times 1B$).
• Speed: Distance computations drop from 1536 operations to 16 LUT additions.

6. Summary

Putting It All Together: IVF-PQ

1. IVF narrows the search to relevant clusters.
2. PQ approximates distances within those clusters.
3. Result: Billion-scale search in milliseconds!

Next, let’s explore IVF-PQ Integration and parameter tuning!