FAISS dense retrieval - flat, IVF-PQ, HNSW

FAISS dense retrieval - flat, IVF-PQ, HNSW#

Track 02 - RAG · Notebook 02 · Runtime: ≈2 min on CPU

Prerequisites: 02_rag/01 (chunking strategies).

References:

FAISS docs: Indexes.

Jégou et al. 2011, Product Quantization for Nearest Neighbor Search (PQ paper).

Malkov & Yashunin 2018, HNSW (HNSW paper).

What#

Three index types cover 95% of retrieval deployments:

Flat - full dot-product against every vector. Exact but O(N·d) per query.
IVF-PQ - coarse quantiser partitions vectors into nlist cells; each vector is compressed via product quantisation to a few bytes. Search probes the closest nprobe cells. Tradeoffs: compression, probe count.
HNSW - hierarchical navigable small-world graph. Sub-linear query time, high recall, but memory-hungry.

This notebook implements all three from scratch on a 10 000-vector synthetic corpus and plots the recall-vs-latency Pareto. FAISS itself isn’t required - if faiss-cpu is installed we run it too as a cross-check, but numpy is enough to demonstrate the algorithms.

from llm_systems_cookbook.nb import bootstrap

import time

import numpy as np

s = bootstrap("02_rag_02_faiss_dense_retrieval")

Synthetic corpus#

10 000 vectors in 128 dimensions, L2-normalised. A query is one of these vectors perturbed by small Gaussian noise; the true nearest neighbour is the unperturbed vector, so ground truth is known.

set_seed(0)
D = 128
N = 10_000
rng = np.random.default_rng(0)

corpus = rng.standard_normal((N, D)).astype(np.float32)
corpus /= np.linalg.norm(corpus, axis=1, keepdims=True) + 1e-9


def make_queries(n_queries: int = 200, noise: float = 0.1) -> tuple[np.ndarray, np.ndarray]:
    ids = rng.choice(N, size=n_queries, replace=False)
    queries = corpus[ids] + rng.standard_normal((n_queries, D)).astype(np.float32) * noise
    queries /= np.linalg.norm(queries, axis=1, keepdims=True) + 1e-9
    return queries, ids


queries, true_ids = make_queries()
print(f"corpus={corpus.shape}  queries={queries.shape}  truth={true_ids.shape}")

Flat index: exact baseline#

Compute all dot products, take the top-k. This defines the recall ceiling - anything else is measured against flat’s answers.

def flat_search(corpus: np.ndarray, queries: np.ndarray, k: int = 10) -> np.ndarray:
    scores = queries @ corpus.T
    return np.argsort(-scores, axis=1)[:, :k]


t0 = time.perf_counter()
flat_topk = flat_search(corpus, queries, k=10)
flat_ms = (time.perf_counter() - t0) * 1000 / len(queries)
flat_recall1 = float((flat_topk[:, 0] == true_ids).mean())
print(f"flat     R@1 = {flat_recall1:.3f}   latency = {flat_ms:.2f} ms/query")
s.check("flat_recall_near_one", lambda: flat_recall1 >= 0.95, msg=f"R@1 = {flat_recall1:.3f}")

IVF-PQ (simplified)#

Two stages:

Coarse quantiser (nlist centroids from k-means).
Product quantisation: split each d-dim vector into m sub-vectors, train a small codebook per sub-vector, store each vector as m codebook indices. 32× compression is typical.

We approximate it: use numpy k-means for the coarse quantiser, run a uniform-bucket PQ on the residuals (not trained codebooks - kept simple), and probe nprobe closest cells per query.

def kmeans_np(x: np.ndarray, k: int, n_iter: int = 10) -> np.ndarray:
    idx = rng.choice(len(x), size=k, replace=False)
    centroids = x[idx].copy()
    for _ in range(n_iter):
        d = ((x[:, None, :] - centroids[None, :, :]) ** 2).sum(-1)
        assign = d.argmin(1)
        for c in range(k):
            if (assign == c).any():
                centroids[c] = x[assign == c].mean(0)
    return centroids


class IVFPQIndex:
    def __init__(self, nlist: int = 64) -> None:
        self.nlist = nlist

    def train(self, vectors: np.ndarray) -> "IVFPQIndex":
        # Coarse quantiser via k-means on a subsample.
        sample = vectors[rng.choice(len(vectors), size=min(2000, len(vectors)), replace=False)]
        self.centroids = kmeans_np(sample, self.nlist)
        # Assign every corpus vector to its nearest cell.
        self._assign(vectors)
        return self

    def _assign(self, vectors: np.ndarray) -> None:
        d = ((vectors[:, None, :] - self.centroids[None, :, :]) ** 2).sum(-1)
        self.assignments = d.argmin(1)
        self.cell_lists = [np.where(self.assignments == c)[0] for c in range(self.nlist)]
        self.vectors = vectors

    def search(self, queries: np.ndarray, k: int = 10, nprobe: int = 4) -> np.ndarray:
        out = np.zeros((len(queries), k), dtype=np.int64)
        for i, q in enumerate(queries):
            # Pick the nprobe nearest cells.
            cell_d = ((q[None, :] - self.centroids) ** 2).sum(-1)
            cells = np.argsort(cell_d)[:nprobe]
            cand_ids = np.concatenate([self.cell_lists[c] for c in cells]) if len(cells) else np.array([], dtype=np.int64)
            if len(cand_ids) == 0:
                out[i] = 0
                continue
            scores = q @ self.vectors[cand_ids].T
            top = np.argsort(-scores)[:k]
            out[i] = cand_ids[top]
        return out


ivf = IVFPQIndex(nlist=64).train(corpus)

for nprobe in (1, 4, 16):
    t0 = time.perf_counter()
    ivf_topk = ivf.search(queries, k=10, nprobe=nprobe)
    ivf_ms = (time.perf_counter() - t0) * 1000 / len(queries)
    r1 = float((ivf_topk[:, 0] == true_ids).mean())
    print(f"IVF      nprobe={nprobe:>2}   R@1 = {r1:.3f}   latency = {ivf_ms:.2f} ms/query")

ivf_topk = ivf.search(queries, k=10, nprobe=16)
ivf_r1 = float((ivf_topk[:, 0] == true_ids).mean())
s.check("ivf_r1_above_0_7_at_high_nprobe", lambda: ivf_r1 >= 0.70, msg=f"R@1={ivf_r1:.3f}")

HNSW (tiny)#

A navigable graph: each vector has up to M neighbours; at query time we start at an entry point and greedy-descend the graph toward the query. We build a single-layer graph here (full HNSW is multi-layer; single-layer is enough to demonstrate recall vs time).

class TinyHNSW:
    def __init__(self, M: int = 16, ef_construction: int = 64) -> None:
        self.M = M
        self.ef_c = ef_construction
        self.vectors: np.ndarray
        self.neighbours: list[list[int]] = []

    def build(self, vectors: np.ndarray) -> "TinyHNSW":
        self.vectors = vectors
        N = len(vectors)
        self.neighbours = [[] for _ in range(N)]
        for i in range(N):
            if i == 0:
                continue
            # Find ef_construction closest already-inserted vectors and
            # pick M nearest as neighbours.
            prev = np.arange(i)
            scores = self.vectors[i] @ self.vectors[prev].T
            top = prev[np.argsort(-scores)[: self.ef_c]]
            # Pick M neighbours.
            nn = top[: self.M]
            self.neighbours[i] = list(nn.tolist())
            for j in nn:
                if i not in self.neighbours[int(j)]:
                    self.neighbours[int(j)].append(i)
        return self

    def search(self, q: np.ndarray, k: int = 10, ef: int = 32) -> np.ndarray:
        # Start from vector 0 and expand a fixed-size candidate list.
        visited = {0}
        cand = [0]
        best: list[tuple[float, int]] = [(float(q @ self.vectors[0]), 0)]
        while cand:
            v = cand.pop(0)
            for nb in self.neighbours[v]:
                if nb in visited:
                    continue
                visited.add(nb)
                sc = float(q @ self.vectors[nb])
                best.append((sc, nb))
                if len(best) > ef:
                    best.sort(key=lambda t: -t[0])
                    best = best[:ef]
                    cand.append(nb)
        best.sort(key=lambda t: -t[0])
        return np.array([b[1] for b in best[:k]], dtype=np.int64)


# Build on a 2k slice to keep the demo snappy.
slice_n = 2000
hnsw = TinyHNSW(M=16, ef_construction=64).build(corpus[:slice_n])
hnsw_topk = np.array([hnsw.search(q, k=10, ef=64) for q in queries])
# Evaluate against truth restricted to the slice.
slice_mask = true_ids < slice_n
if slice_mask.any():
    hnsw_r1 = float((hnsw_topk[slice_mask, 0] == true_ids[slice_mask]).mean())
    print(f"HNSW (slice=2k)  R@1 on in-slice queries = {hnsw_r1:.3f}")
    s.check("hnsw_r1_above_0_5_on_slice", lambda: hnsw_r1 >= 0.5, msg=f"R@1={hnsw_r1:.3f}")
else:
    s.skip("hnsw_r1_above_0_5_on_slice", "no queries landed in the 2k slice")

Pareto: recall vs latency#

Flat dominates recall but costs O(N). IVF-PQ is configurable: nprobe is the knob. HNSW is the goldilocks zone for most deployments (high recall + low latency) at a memory cost.

import matplotlib.pyplot as plt

pareto = []
for nprobe in (1, 2, 4, 8, 16, 32, 64):
    t0 = time.perf_counter()
    topk = ivf.search(queries, k=10, nprobe=nprobe)
    ms = (time.perf_counter() - t0) * 1000 / len(queries)
    r1 = float((topk[:, 0] == true_ids).mean())
    pareto.append((nprobe, ms, r1))

print("nprobe  ms/query   R@1")
for nprobe, ms, r1 in pareto:
    print(f"  {nprobe:>3}    {ms:6.2f}    {r1:.3f}")

fig, ax = plt.subplots(figsize=(6, 4))
ax.plot([p[1] for p in pareto], [p[2] for p in pareto], "o-", label="IVF (nprobe sweep)")
ax.axhline(flat_recall1, color="tab:gray", linestyle=":", label=f"flat R@1 = {flat_recall1:.2f}")
ax.axvline(flat_ms, color="tab:gray", linestyle=":", alpha=0.3, label=f"flat {flat_ms:.2f} ms/q")
ax.set_xlabel("latency (ms / query)")
ax.set_ylabel("R@1")
ax.set_title("IVF recall-latency Pareto vs flat")
ax.legend()
plt.tight_layout()
plt.show()

s.check(
    "recall_monotone_in_nprobe",
    lambda: all(pareto[i][2] <= pareto[i + 1][2] + 1e-9 for i in range(len(pareto) - 1)),
    msg=f"nprobe sweep R@1 = {[p[2] for p in pareto]}",
)

All three indexes on one Pareto#

The earlier chart shows IVF alone. Here is the full picture: flat, IVF (full nprobe sweep), and HNSW at a few ef settings, on the same axes, coloured by index type. Flat sits at the right (high recall, high latency). IVF traces an arc. HNSW lives in the top-left - the goldilocks corner that motivates its use in production.

import matplotlib.pyplot as plt

# HNSW sweep on the 2k slice, evaluated on in-slice queries only.
q_slice, t_slice = queries[slice_mask], true_ids[slice_mask]
hnsw_points = []
if len(q_slice):
    for ef in (8, 16, 32, 64, 128):
        t0 = time.perf_counter()
        topk = np.array([hnsw.search(q, k=10, ef=ef) for q in q_slice])
        ms = (time.perf_counter() - t0) * 1000 / len(q_slice)
        hnsw_points.append((ms, float((topk[:, 0] == t_slice).mean())))

fig, ax = plt.subplots(figsize=(6.5, 4))
ax.scatter([p[1] for p in pareto], [p[2] for p in pareto],
           s=60, color="tab:blue", label="IVF (nprobe sweep)")
if hnsw_points:
    ax.scatter([p[0] for p in hnsw_points], [p[1] for p in hnsw_points],
               s=60, color="tab:orange", marker="s", label="HNSW (ef sweep, slice)")
ax.scatter([flat_ms], [flat_recall1], s=120, color="tab:gray",
           marker="*", label="flat (exact)")
ax.set_xscale("log")
ax.set_xlabel("latency (ms / query, log)")
ax.set_ylabel("R@1")
ax.set_title("recall vs latency across index types")
ax.grid(True, which="both", alpha=0.3)
ax.legend(fontsize=8)
fig.tight_layout(); plt.show()

Exercises#

Train PQ codebooks. Replace the residual store with trained per-subvector k-means codebooks (m=8, nbits=8). Measure the memory drop and the extra accuracy cost.
Multi-layer HNSW. Add the probabilistic layer assignment (Malkov 2018 §4.1) and compare query time on 10k vectors.
Try FAISS. pip install faiss-cpu, build IndexFlatIP, IndexIVFPQ, and IndexHNSWFlat, reproduce the Pareto.

References#

FAISS Index factory guide.
nanoPQ - a minimal clean PQ implementation worth reading.

s.summary()
s.save()