#!/usr/bin/env python3 """ annlab.py — the 150-line version of Project 2, so you can see the shape of the result before you spend seven weeks on the real thing. This is deliberately NOT the deliverable. It is a reference point: a working brute-force baseline, a working navigable-small-world graph, and the recall / latency / efSearch curve they produce. When your own implementation is done, its curve should dominate this one at every point -- if it does not, you have a bug, and knowing that on day 3 instead of week 7 is the entire value of having a baseline. What it demonstrates -------------------- 1. Brute force is not slow because the algorithm is bad. It is slow because it moves n*d*4 bytes per query and does 2*n*d FLOPs of work with an arithmetic intensity of about 2 -- deeply memory-bound. See roofline.py. 2. A greedy graph walk turns an O(n) scan into an O(log n)-ish walk, but it is *approximate*: it can and does get stuck in local minima. The beam width (efSearch) is the knob that trades recall against latency, and the curve is strongly concave -- the first few units of ef buy a lot of recall, the last few buy almost none. 3. Recall must be measured against exact ground truth from the same distance function on the same data. Not a proxy. Not another library's answer. python3 annlab.py # default sweep python3 annlab.py --n 20000 --d 64 --m 16 """ from __future__ import annotations import argparse import heapq import time import numpy as np from metrics import recall_at_k # --------------------------------------------------------------------------- # Distance. One choice, made explicitly, used everywhere. # --------------------------------------------------------------------------- # We L2-normalise all vectors, which makes cosine similarity equal to the dot # product and makes squared Euclidean distance a monotone function of it: # ||a-b||^2 = ||a||^2 + ||b||^2 - 2 a.b = 2 - 2 a.b (for unit vectors) # So ranking by max dot product, max cosine, or min L2 gives the IDENTICAL # ordering. This equivalence only holds under normalisation -- forget it and # your recall silently drops, which is Project 2's normalised-vs-unnormalised # experiment. def normalise(x: np.ndarray) -> np.ndarray: return x / np.linalg.norm(x, axis=-1, keepdims=True) class BruteForce: """The baseline you are allowed to trust. Exact by construction.""" def __init__(self, data: np.ndarray): self.data = data def search(self, q: np.ndarray, k: int) -> list[int]: sims = self.data @ q # (n,) one BLAS call # argpartition is O(n) vs O(n log n) for a full sort: we only need the # top k to be in the first k slots, then we sort just those k. idx = np.argpartition(-sims, k)[:k] return idx[np.argsort(-sims[idx])].tolist() class NSW: """A single-layer navigable small-world graph. Insertion: greedily search the graph built so far, connect the new node to its M nearest found neighbours, and add the reciprocal edges. Because early nodes are inserted when the graph is nearly empty, their edges are long -- and those accidental long-range links are exactly what makes the greedy walk logarithmic instead of linear. That is the small-world property, and it emerges from the insertion order rather than being designed in. HNSW's addition (Project 2, milestone 4) is to stack these graphs in layers with exponentially decaying membership, so the walk starts at a coarse layer and refines -- which removes the dependence on lucky long edges. """ def __init__(self, data: np.ndarray, m: int = 16, ef_construction: int = 100, seed: int = 0): self.data = data self.m = m self.graph: list[list[int]] = [[] for _ in range(len(data))] self.entry = 0 # Distance computations are the currency of ANN search. Wall-clock is # language- and machine-specific; distance count is the algorithmic # quantity that transfers between implementations. Report BOTH, always: # a graph that needs 100x fewer distances but runs slower is telling # you your constant factor is the problem, not your algorithm. self.dist_count = 0 rng = np.random.default_rng(seed) order = rng.permutation(len(data)) for count, node in enumerate(order): if count == 0: self.entry = int(node) continue cands = self._search_internal(data[node], ef_construction, int(node)) neighbours = [i for _, i in heapq.nsmallest(m, cands)] self.graph[node] = neighbours for nb in neighbours: self.graph[nb].append(int(node)) # Degree bound: without it, hub nodes accumulate thousands of # edges, every visit to a hub costs a huge scan, and the search # degenerates toward brute force. Pruning to the m nearest is # the cheap version of HNSW's heuristic neighbour selection. if len(self.graph[nb]) > 2 * m: d = self.data[self.graph[nb]] @ self.data[nb] keep = np.argsort(-d)[: 2 * m] self.graph[nb] = [self.graph[nb][i] for i in keep] def _search_internal(self, q: np.ndarray, ef: int, exclude: int = -1): """Beam search. Returns a list of (distance, node) with distance = 1 - cosine, so smaller is better and heapq gives us a min-heap for free.""" start = self.entry d0 = 1.0 - float(self.data[start] @ q) self.dist_count += 1 visited = {start} candidates = [(d0, start)] # min-heap: closest unexplored first results = [(-d0, start)] # max-heap by negated distance: worst # result on top, so we can evict it while candidates: d, node = heapq.heappop(candidates) # Stopping rule: if the closest unexplored candidate is further # than the worst thing we already have, no unexplored node can # improve the result set -- under the assumption that the graph is # locally metric. That assumption is exactly what makes this # approximate: it is false near a local minimum, and that is where # the missing recall goes. if -results[0][0] < d and len(results) >= ef: break for nb in self.graph[node]: if nb in visited or nb == exclude: continue visited.add(nb) dn = 1.0 - float(self.data[nb] @ q) self.dist_count += 1 if len(results) < ef or dn < -results[0][0]: heapq.heappush(candidates, (dn, nb)) heapq.heappush(results, (-dn, nb)) if len(results) > ef: heapq.heappop(results) return [(-nd, n) for nd, n in results] def search(self, q: np.ndarray, k: int, ef: int = 64) -> list[int]: # ef must be >= k or you cannot possibly return k good answers. got = self._search_internal(q, max(ef, k)) return [i for _, i in heapq.nsmallest(k, got)] def main() -> int: ap = argparse.ArgumentParser() ap.add_argument("--n", type=int, default=10000) ap.add_argument("--d", type=int, default=64) ap.add_argument("--m", type=int, default=16) ap.add_argument("--efc", type=int, default=100) ap.add_argument("--queries", type=int, default=200) ap.add_argument("--k", type=int, default=10) ap.add_argument("--clusters", type=int, default=0, help="0 = uniform on the sphere (adversarial). >0 = that many " "Gaussian clusters, which is what real embeddings look like.") ap.add_argument("--spread", type=float, default=0.05, help="cluster standard deviation. MUST satisfy spread*sqrt(d) << 1 " "or the clusters wash out -- see the warning this prints.") a = ap.parse_args() rng = np.random.default_rng(42) if a.clusters > 0: # Real embedding spaces are not uniform. Text/image encoders produce # points concentrated near a low-dimensional manifold: topics, styles, # entities. We model that crudely with c Gaussian blobs. The *intrinsic* # dimensionality is what governs ANN difficulty, not the ambient d, and # this switch is the cheapest way to see that for yourself. centres = normalise(rng.standard_normal((a.clusters, a.d)).astype(np.float32)) assign = rng.integers(0, a.clusters, size=a.n) data = normalise(centres[assign] + a.spread * rng.standard_normal((a.n, a.d)).astype(np.float32)) qassign = rng.integers(0, a.clusters, size=a.queries) queries = normalise(centres[qassign] + a.spread * rng.standard_normal((a.queries, a.d)).astype(np.float32)) kind = f"{a.clusters} gaussian clusters, spread={a.spread}" # A Gaussian perturbation with per-axis sigma has expected norm # sigma*sqrt(d). The cluster centres are unit vectors. So if # sigma*sqrt(d) >~ 1 the noise is larger than the signal, every point # lands somewhere random on the sphere, and your "clustered" dataset is # uniform data wearing a hat. Measured at d=64: spread=0.25 gives # sigma*sqrt(d)=2.0 and RC=1.39 against uniform's 1.36 -- no difference. # spread=0.05 gives 0.40 and RC=3.37 -- real structure. noise_norm = a.spread * (a.d ** 0.5) if noise_norm > 0.7: print(f"WARNING: spread*sqrt(d) = {noise_norm:.2f}. The perturbation is " f"comparable to the\n unit-norm centres, so these clusters " f"are washed out and this dataset\n is effectively uniform. " f"Use --spread {0.4/(a.d**0.5):.3f} or less at d={a.d}.\n") else: data = normalise(rng.standard_normal((a.n, a.d), dtype=np.float32)) queries = normalise(rng.standard_normal((a.queries, a.d), dtype=np.float32)) kind = "uniform on the unit sphere" bf = BruteForce(data) t0 = time.perf_counter() truth = [bf.search(q, a.k) for q in queries] bf_ms = (time.perf_counter() - t0) / a.queries * 1e3 t0 = time.perf_counter() index = NSW(data, m=a.m, ef_construction=a.efc) build_s = time.perf_counter() - t0 edges = sum(len(g) for g in index.graph) # Index size accounting: the vectors themselves plus the adjacency lists. # Graph ANN indexes are usually 1.2-2x the raw vector bytes, and people who # quote "index size" without saying whether vectors are included are # comparing different things. vec_mb = a.n * a.d * 4 / 1e6 graph_mb = edges * 4 / 1e6 # Relative contrast (He, Kumar & Chang, ICML 2012): the mean distance from # a query to the whole dataset, divided by the distance to its true nearest # neighbour. # # RC = d_mean / d_1 # # RC -> 1 means every point is about as far away as every other point, so # greedy descent has no gradient to follow and the graph walk degenerates # into a random walk. RC, not the ambient dimension d, is what predicts ANN # difficulty -- which is why "works at d=1536" claims are meaningless # without the dataset. Measured here on unit vectors, where # ||a-b|| = sqrt(2 - 2*cos). # # Reference values from this script at n=10k: d=16 -> 2.22, d=64 -> 1.36, # d=128 -> 1.22, d=512 -> 1.10. Real text embeddings sit far above their # ambient dimension's uniform value because they are concentrated on a # low-dimensional manifold. _s = data @ queries.T _dist = np.sqrt(np.maximum(0.0, 2.0 - 2.0 * _s)) _srt = np.sort(_dist, axis=0) rc = float(np.mean(_dist.mean(axis=0) / _srt[0])) dk_d1 = float(np.mean(_srt[a.k - 1] / _srt[0])) print(f"dataset n={a.n} d={a.d} ({kind})") print(f" relative contrast RC = {rc:.3f} d{a.k}/d1 = {dk_d1:.4f} " f"({'HARD' if rc < 1.5 else 'navigable'})") print(f"build {build_s:.2f}s M={a.m} efC={a.efc} " f"{edges/a.n:.1f} edges/node") print(f"size vectors {vec_mb:.2f} MB + graph {graph_mb:.2f} MB " f"= {vec_mb+graph_mb:.2f} MB ({(vec_mb+graph_mb)/vec_mb:.2f}x raw)") print(f"brute {bf_ms:.3f} ms/query -> {1000/bf_ms:.0f} qps (exact)\n") print(f"{'efSearch':>9} {'recall@'+str(a.k):>10} {'ms p50':>9} {'ms p95':>9} " f"{'qps':>9} {'speedup':>9} {'dists/q':>9} {'ns/dist':>9}") print("-" * 82) rows = [] for ef in (10, 16, 24, 32, 48, 64, 96, 128, 192, 256): lat = [] recs = [] index.dist_count = 0 for q, t in zip(queries, truth): s = time.perf_counter() got = index.search(q, a.k, ef=ef) lat.append((time.perf_counter() - s) * 1e3) recs.append(recall_at_k(got, t, a.k)) dpq = index.dist_count / len(queries) lat.sort() p50 = lat[len(lat) // 2] p95 = lat[int(0.95 * len(lat)) - 1] r = sum(recs) / len(recs) rows.append((ef, r, p50, dpq)) print(f"{ef:>9} {r:>10.4f} {p50:>9.3f} {p95:>9.3f} {1000/p50:>9.0f} " f"{bf_ms/p50:>8.2f}x {dpq:>9.0f} {p50*1e6/dpq:>9.1f}") # ---- the model that explains the speedup column -------------------- # Brute force does n distance computations per query inside ONE BLAS call. # The graph does far fewer, but each one is an interpreted Python round # trip. Break-even is where those two effects cancel: # # speedup = (n / dists_per_query) / (ns_per_dist_graph / ns_per_dist_bf) # \_____ algorithmic ____/ \______ constant factor ______/ # # Print both factors so you can see WHICH one you need to fix. ef64 = [r for r in rows if r[0] == 64][0] bf_ns_per_dist = bf_ms * 1e6 / a.n g_ns_per_dist = ef64[2] * 1e6 / ef64[3] algo = a.n / ef64[3] const = g_ns_per_dist / bf_ns_per_dist print(f"\nWhy the speedup column looks like that, at efSearch=64:") print(f" algorithmic win : {a.n:,} / {ef64[3]:.0f} distances = {algo:8.1f}x fewer") print(f" constant-factor : {g_ns_per_dist:.1f} ns/dist (python loop) vs " f"{bf_ns_per_dist:.1f} ns/dist (BLAS) = {const:6.1f}x slower each") print(f" predicted speedup : {algo:.1f} / {const:.1f} = {algo/const:.2f}x") print(f" measured speedup : {bf_ms/ef64[2]:.2f}x") print(f" The algorithm is right and the implementation is wrong. That is a") print(f" DIFFERENT bug from 'the algorithm is wrong', and only the") print(f" distance counter can tell you which one you have.") print("\nRead the curve, not the endpoints:") print(" * recall is concave in efSearch -- doubling ef past the knee buys") print(" fractions of a percent and costs a linear amount of latency.") print(" * p95/p50 ratio grows with ef: longer walks have more variance, so") print(" an index tuned on p50 will miss its p99 SLO.") print(" * the honest operating point is chosen by fixing the recall you") print(" need FIRST and reading the latency off the curve -- never by") print(" picking a latency and reporting whatever recall falls out.") return 0 if __name__ == "__main__": raise SystemExit(main())