Vector Databases Approximate Nearest Neighbor Search
# Vector Databases & Approximate Nearest Neighbor Search
## 1. Introduction & Motivation
Modern machine learning systems increasingly represent objects — documents, images, products, user profiles — as dense embedding vectors, and a huge fraction of practical ML applications reduce, at some point, to a nearest-neighbor search: given a query vector, find the most similar vectors in a large collection. Retrieval-augmented generation looks up relevant passages this way, recommendation systems find similar items this way, image search finds visually similar photos this way, and deduplication pipelines find near-duplicate records this way. As embedding collections have grown from thousands to billions of vectors, exact nearest-neighbor search — computing the distance from a query to every stored vector — has become computationally infeasible for latency-sensitive applications, giving rise to an entire subfield built around approximate nearest neighbor (ANN) search and the specialized database systems, called vector databases, built to serve it at scale.
The central insight enabling ANN search is that most applications do not actually require the mathematically exact nearest neighbor; a result that is highly likely to be among the true top-k neighbors, found in a small fraction of the time and memory that exact search would require, is a far better engineering trade-off. This has produced a rich family of indexing algorithms — inverted-file indexes, hierarchical graph-based indexes like HNSW, locality-sensitive hashing, and product quantization — each making a different trade-off between search speed, memory footprint, index build time, and recall (the fraction of true nearest neighbors actually returned).
Vector databases (Pinecone, Weaviate, Milvus, Qdrant, and vector-search extensions to traditional databases like pgvector) package these indexing algorithms into production-ready systems with the operational features expected of any database: durability, horizontal scaling, metadata filtering, and incremental updates — turning what was originally a purely algorithmic research area into critical infrastructure underpinning retrieval-augmented generation, semantic search, and recommendation at scale.
## 2. Core Concepts & Theory
The nearest-neighbor search problem, formally, is: given a query vector $q$ and a collection of $n$ vectors $\{v_1, \ldots, v_n\}$ in $\mathbb{R}^d$, find the $k$ vectors minimizing (or maximizing, for similarity measures) a distance function, most commonly Euclidean distance, cosine similarity, or dot product. Brute-force (exact, "flat") search computes this distance for every stored vector, giving $O(nd)$ query time — exact, but linear in the size of the collection, which becomes prohibitive once $n$ reaches millions or billions.
ANN algorithms accelerate this by avoiding a full linear scan, using one of several core strategies. Partitioning-based methods (inverted file indexes, IVF) cluster the vector space (typically via k-means) into a modest number of partitions ("Voronoi cells"), and at query time only search the partitions nearest to the query, examining a small fraction of the full collection rather than all of it. Graph-based methods (HNSW, Hierarchical Navigable Small World graphs) build a multi-layer proximity graph over the vectors and perform greedy best-first search through the graph, exploiting the "small world" property that any two nodes can be reached via relatively few hops through well-chosen edges. Hashing-based methods (locality-sensitive hashing, LSH) use hash functions specifically designed so that similar vectors are more likely to collide into the same hash bucket than dissimilar ones, converting similarity search into a hash-table lookup. Quantization-based methods (product quantization, PQ) compress vectors into compact codes to shrink memory footprint and accelerate distance computation, often used in combination with a partitioning index rather than as a standalone search structure.
## 3. Mathematical Formulation
For the inverted-file (IVF) approach, the space is first partitioned via k-means clustering into $n_{ ext{list}}$ clusters with centroids $\{c_1, \ldots, c_{n_{ ext{list}}}\}$; each database vector is assigned to its nearest centroid, forming inverted lists. At query time, the query's distance to all centroids is computed, and only the vectors belonging to the $n_{ ext{probe}}$ nearest clusters are exhaustively compared:
$
ext{candidates}(q) = \bigcup_{c \in ext{top-}n_{ ext{probe}}(\{d(q, c_i)\})} ext{list}(c)
$
This reduces the expected number of full distance computations from $n$ to roughly $n \cdot n_{ ext{probe}} / n_{ ext{list}}$, at the cost of some recall loss whenever a true nearest neighbor happens to lie in a cluster that was not probed (an edge effect near cluster boundaries).
For locality-sensitive hashing based on random hyperplanes (SimHash), a hash bit is generated per random hyperplane $w \sim \mathcal{N}(0, I)$ as $h_w(v) = \mathbb{1}[w \cdot v > 0]$. A key theoretical property of this scheme is that the probability two vectors collide on a given hash bit is directly related to the angle between them:
$
P\big[h_w(u) = h_w(v)\big] = 1 - \frac{ heta(u, v)}{\pi}, \qquad heta(u, v) = \arccos\left(\frac{u \cdot v}{\|u\|\|v\|}
ight)
$
This is precisely what makes SimHash a valid locality-sensitive hash for cosine similarity: vectors with high cosine similarity (small angle) collide with high probability, and dissimilar vectors collide close to 50% of the time (pure chance), letting hash bucket membership serve as a fast approximate similarity filter.
For product quantization, a $d$-dimensional vector is split into $m$ sub-vectors of dimension $d/m$, and each sub-vector is independently quantized against its own learned codebook of $K$ codewords (via k-means on that sub-space), so that a full vector is represented as $m$ small integer codes rather than $d$ floating-point numbers, reducing memory from $O(d)$ floats to $O(m \log_2 K)$ bits per vector while allowing approximate distances to be computed via fast codebook lookup tables.
## 4. Advanced Theory & Extensions
HNSW's theoretical grounding traces to navigable small-world graph theory: by constructing a hierarchy of proximity graphs, with sparser graphs at higher layers (used for coarse, long-range navigation) and denser graphs at the base layer (used for fine-grained local search), greedy search can navigate from an arbitrary entry point to the query's true neighborhood in roughly logarithmic expected number of hops, giving HNSW its strong empirical query-time scaling of approximately $O(\log n)$ despite the underlying problem being combinatorially hard in the worst case.
IVF-PQ, combining inverted-file partitioning with product quantization, is the workhorse of large-scale ANN systems (as used in Facebook's FAISS library): IVF narrows the candidate set to a manageable fraction of the collection, and PQ compresses those candidates (and, more importantly, the full database) into a memory footprint that fits in RAM even for billion-scale collections, with an optional final "re-ranking" pass that recomputes exact distances on the top approximate candidates using the original, uncompressed vectors to recover accuracy lost to quantization.
Filtered ANN search — finding nearest neighbors subject to metadata constraints (e.g., "similar products, but only in stock") — is a genuinely difficult extension of the base problem, since a naive pre-filter can shrink the effective candidate pool so much that graph- or partition-based indexes lose their efficiency advantages, and a naive post-filter can return far fewer than $k$ results if too many top candidates fail the filter; production vector databases use specialized techniques (filtered graph traversal that respects predicates during the graph walk, or hybrid inverted-index-plus-vector-index architectures) to address this efficiently.
## 5. Computational Considerations
The core trade-off in ANN system design is a three-way balance between query latency, memory footprint, and recall, and no single algorithm dominates on all three simultaneously: HNSW typically offers the best latency-recall trade-off but has a larger memory footprint (storing graph edges alongside vectors) and slower, more memory-intensive index construction; IVF-PQ offers the smallest memory footprint (crucial at billion-vector scale, where storing full-precision vectors may not fit in RAM at all) at some cost to recall and latency relative to HNSW; and pure brute-force search, while offering perfect recall, only remains practical up to roughly hundreds of thousands to low millions of vectors on modern hardware, or when heavily parallelized on GPUs.
Index build time itself is a real operational cost distinct from query time: HNSW graph construction is $O(n \log n)$ but with a large constant factor, and rebuilding an index from scratch for very large collections can take hours, motivating incremental index update support (inserting new vectors without a full rebuild) as a key differentiator among production vector database systems. Distance computation itself is often the dominant per-query cost at scale, which is why quantization-based memory compression doubles as a computational accelerant — smaller codes mean more of the working set fits in cache and more distance computations can be done via fast table lookups rather than full floating-point arithmetic.
## 6. Practical Implementation Strategies
A practical vector search pipeline typically combines a partitioning or graph-based coarse index with a re-ranking step: retrieve a modestly over-sized candidate set (e.g., $10k$ candidates for a $k$-nearest-neighbor query) via the fast approximate index, then compute exact distances on just that candidate set to produce the final, more accurate top-$k$ ranking — this two-stage retrieve-then-rerank pattern recovers much of the accuracy lost to approximation while still avoiding a full linear scan. Choosing between HNSW and IVF-PQ in practice usually comes down to scale and update pattern: HNSW is preferred when memory is not the binding constraint and highest recall at low latency is the priority, while IVF-PQ (or hybrid IVF-HNSW variants) is preferred once the collection is too large to fit as full-precision vectors in available memory.
Embedding normalization matters more for vector search than it might initially appear: cosine similarity search is typically implemented as a dot-product search on unit-normalized vectors (since $\cos( heta) = \frac{u \cdot v}{\|u\|\|v\|}$ reduces to a plain dot product once vectors are pre-normalized to unit length), which is faster than repeatedly computing norms at query time and is the standard convention in most production embedding pipelines. Metadata filtering should be planned for from the start of index design rather than bolted on afterward, since retrofitting predicate support onto a pure vector index (without index-aware filtering support) commonly degrades either recall or latency substantially under selective filters.
## 7. Benchmark Datasets & Evaluation
ANN-Benchmarks and the closely related BigANN Benchmark (used in the NeurIPS Billion-Scale ANN competition) are the standard evaluation suites for comparing ANN algorithms, using datasets ranging from SIFT1M and GIST1M (classical, hand-crafted image descriptor benchmarks) to billion-scale deep-learning embedding sets (DEEP1B) specifically constructed to stress-test index scalability. The standard evaluation methodology reports recall@k (the fraction of true top-k nearest neighbors returned by the approximate index, computed against a brute-force ground truth) against queries-per-second at a given recall target, typically visualized as a Pareto frontier — since virtually every ANN algorithm has a tunable parameter (nprobe for IVF, ef_search for HNSW) trading recall against latency, comparing algorithms at a single fixed operating point can be misleading.
For retrieval-augmented generation and semantic search applications specifically, end-to-end task metrics (downstream answer accuracy, or human-judged relevance of retrieved passages) are increasingly reported alongside pure recall@k, since a vector index with excellent recall on paper can still underperform if the underlying embedding model itself poorly captures the notion of relevance the downstream application actually needs.
## 8. Key Challenges & Limitations
The curse of dimensionality fundamentally limits ANN algorithm effectiveness: as embedding dimensionality grows very large, distances between points tend to concentrate (the ratio of the farthest to nearest point distance approaches 1), which erodes the discriminative signal that both partitioning and graph-based methods rely on, and is part of why dimensionality reduction is sometimes applied as a pre-processing step before indexing extremely high-dimensional embeddings. Recall-latency-memory trade-offs are fundamentally irreducible, not merely an engineering inconvenience to be optimized away — any ANN system operating below the brute-force memory and compute budget will sacrifice some recall, and system designers must explicitly choose an acceptable point on this trade-off curve for their application rather than expecting a free lunch.
Index staleness under high-throughput updates is a genuine operational challenge: graph-based indexes like HNSW are not naturally suited to high-volume deletion (deleted vectors are typically only "tombstoned" and periodically purged via a full or partial rebuild, since removing a node from a navigable small-world graph without breaking connectivity is non-trivial), which complicates vector database deployments with rapidly changing data, such as real-time personalization systems.
## 9. Hyperparameter Tuning
For IVF, the number of clusters $n_{ ext{list}}$ trades index granularity against per-query overhead — too few clusters means each cluster (and hence each probe) contains too many vectors, eroding the speedup; too many clusters means the initial centroid-comparison step itself becomes expensive, and a common rule of thumb is $n_{ ext{list}} \approx \sqrt{n}$ for a collection of $n$ vectors. The number of probed clusters $n_{ ext{probe}}$ is the primary recall/latency dial at query time, and should be tuned against a validation recall target rather than fixed arbitrarily, since the correct value is highly dependent on data distribution and cluster balance.
For HNSW, the graph degree $M$ (number of edges per node) trades memory and build time against search quality, and the search-time parameter $ ext{ef\_search}$ (candidate list size during the greedy graph walk) is the primary recall/latency dial, analogous to $n_{ ext{probe}}$ for IVF — larger $ ext{ef\_search}$ explores more of the graph per query, improving recall at the cost of latency. For product quantization, the number of sub-vectors $m$ and per-subspace codebook size $K$ jointly determine both compression ratio and reconstruction (hence search) accuracy, with the sub-vector count in particular needing to divide the original dimensionality evenly and to be chosen so that each sub-space retains enough structure for its codebook to be meaningfully discriminative.
## 10. Real-World Applications & Case Studies
Retrieval-augmented generation systems depend critically on fast, accurate vector search to retrieve relevant context passages at generation time, and the latency budget for this retrieval step directly constrains what index configuration is viable in a production chat application where total response latency is user-facing and tightly bounded. E-commerce recommendation and visual search systems (finding visually or semantically similar products) at companies like Amazon, Alibaba, and Pinterest operate some of the largest known production ANN deployments, indexing billions of product or image embeddings and serving nearest-neighbor queries at very high throughput with tight latency SLAs.
Facebook's FAISS library, originally developed for large-scale image similarity search, became a foundational open-source building block underlying many production vector databases and RAG pipelines, popularizing IVF-PQ and GPU-accelerated brute-force search as practical, well-engineered baseline choices. Deduplication and fraud-detection pipelines use ANN search to identify near-duplicate records or transactions at scale — a task where high recall on true near-duplicates matters more than perfect precision, since a human review or downstream verification step typically follows automated candidate identification.
## 11. Integration with Other Methods
Vector search integrates directly with retrieval-augmented generation, chain-of-thought-style iterative retrieval (where an LLM issues multiple retrieval queries during a single reasoning chain), and hybrid search architectures that combine dense vector similarity with traditional sparse keyword search (BM25) via a re-ranking or score-fusion step, since dense and sparse retrieval methods tend to have complementary strengths — dense embeddings capture semantic similarity while sparse keyword search reliably catches exact term matches that embeddings can sometimes miss. It also connects naturally to model compression and quantization research, since product quantization for ANN search and weight quantization for model compression share the same underlying mathematical machinery of vector-space compression via learned codebooks.
Embedding model training itself (contrastive learning, sentence-embedding fine-tuning) is tightly coupled to downstream ANN search quality, since an embedding space that does not cleanly separate semantically similar from dissimilar items no amount of index sophistication can compensate for — this makes embedding model selection and fine-tuning, not just index algorithm choice, a first-order lever for retrieval quality in practice.
## 12. Future Research Directions
Learned indexes — using small neural networks to directly predict a vector's approximate location or partition, rather than relying purely on hand-designed clustering or graph-construction heuristics — represent an active research direction seeking to further improve the recall-latency-memory trade-off curve beyond what classical algorithms achieve. GPU- and hardware-accelerator-native ANN index designs, exploiting massive parallelism for both index construction and query serving, are increasingly important as embedding collection sizes and query throughput requirements continue to grow, particularly for latency-sensitive, high-QPS production RAG deployments.
Better support for hybrid structured-and-unstructured queries (combining vector similarity with complex relational predicates, joins, and aggregations) remains an open systems challenge, as does developing standardized, more rigorous benchmarking methodology that accounts for realistic update patterns (streaming insertions and deletions) rather than the largely static, build-once-query-many benchmarks that dominate current ANN algorithm comparisons.
## 13. Summary & Key Takeaways
Approximate nearest neighbor search addresses the fundamental scalability limitation of exact nearest-neighbor search — its linear dependence on collection size — through a family of complementary algorithmic strategies: partitioning-based indexes that narrow the search to promising clusters, graph-based indexes that enable efficient greedy navigation through a proximity structure, hashing-based methods that convert similarity into fast bucket lookups, and quantization-based compression that shrinks memory footprint and accelerates distance computation. Every practical ANN system embodies an explicit, tunable trade-off between recall, query latency, memory footprint, and index build/update cost, and production vector databases (built atop libraries like FAISS, or as purpose-built systems like Pinecone, Milvus, and Qdrant) exist precisely to package these algorithmic trade-offs into operationally robust infrastructure. As embedding-based retrieval becomes foundational to retrieval-augmented generation, recommendation, and semantic search at ever-larger scale, the algorithms and systems covered in this article represent critical, load-bearing infrastructure rather than a narrow academic subfield.
Keywords: vector database, approximate nearest neighbor search, ANN, HNSW, inverted file index, IVF, product quantization, PQ, locality-sensitive hashing, LSH, FAISS, recall at k, cosine similarity, vector index, embedding search, semantic search, Pinecone, Milvus, Qdrant, filtered vector search
---
## Appendix: Practical Labs
### Lab 1: IVF Index Recall Improves With More Probed Clusters
import numpy as np
def kmeans(vectors, n_clusters, n_iters, rng):
"""A minimal Lloyd's-algorithm k-means implementation used to build the
coarse quantizer (cluster centroids) for an inverted-file (IVF) index."""
n = vectors.shape[0]
idx = rng.choice(n, n_clusters, replace=False)
centroids = vectors[idx].copy()
for _ in range(n_iters):
dists = np.linalg.norm(vectors[:, None, :] - centroids[None, :, :], axis=2)
assignments = np.argmin(dists, axis=1)
for c in range(n_clusters):
members = vectors[assignments == c]
if len(members) > 0:
centroids[c] = members.mean(axis=0)
return centroids, assignments
def build_ivf_index(vectors, n_clusters, rng, n_iters=10):
centroids, assignments = kmeans(vectors, n_clusters, n_iters, rng)
inverted_lists = {c: np.where(assignments == c)[0] for c in range(n_clusters)}
return centroids, inverted_lists
def ivf_search(query, vectors, centroids, inverted_lists, nprobe, k):
centroid_dists = np.linalg.norm(centroids - query, axis=1)
probe_clusters = np.argsort(centroid_dists)[:nprobe]
candidate_idx = np.concatenate([inverted_lists[c] for c in probe_clusters]) if len(probe_clusters) else np.array([], dtype=int)
if len(candidate_idx) == 0:
return np.array([], dtype=int)
dists = np.linalg.norm(vectors[candidate_idx] - query, axis=1)
order = np.argsort(dists)[:k]
return candidate_idx[order]
def brute_force_search(query, vectors, k):
dists = np.linalg.norm(vectors - query, axis=1)
return np.argsort(dists)[:k]
def recall_at_k(approx_idx, true_idx):
return len(set(approx_idx) & set(true_idx)) / len(true_idx)
def test_ivf_recall_increases_with_nprobe():
rng = np.random.RandomState(0)
n, dim = 2000, 16
vectors = rng.normal(0, 1, size=(n, dim))
n_clusters = 20
centroids, inverted_lists = build_ivf_index(vectors, n_clusters, rng)
n_queries, k = 30, 10
queries = rng.normal(0, 1, size=(n_queries, dim))
nprobes = [1, 3, 8, 20]
mean_recalls = []
for nprobe in nprobes:
recalls = []
for q in queries:
true_idx = brute_force_search(q, vectors, k)
approx_idx = ivf_search(q, vectors, centroids, inverted_lists, nprobe, k)
recalls.append(recall_at_k(approx_idx, true_idx))
mean_recalls.append(np.mean(recalls))
print(f"{'nprobe':>7} | {'mean recall@10':>14}")
for np_, r in zip(nprobes, mean_recalls):
print(f"{np_:7d} | {r:14.4f}")
assert mean_recalls[0] < mean_recalls[1] < mean_recalls[2] < mean_recalls[3], \
"Recall should increase monotonically as more clusters are probed"
# Probing every cluster is mathematically equivalent to brute-force search
assert mean_recalls[-1] == 1.0, "Probing all clusters should recover exact brute-force recall"
print("IVF recall-vs-nprobe test passed.")
if __name__ == "__main__":
test_ivf_recall_increases_with_nprobe()### Lab 2: Locality-Sensitive Hashing Collision Probability Matches Theory
import numpy as np
def make_pair_with_target_cosine(target_cos, dim, rng):
"""Constructs two unit vectors with a specified cosine similarity, by
combining a base vector with an orthogonal noise component."""
a = rng.normal(0, 1, dim)
a /= np.linalg.norm(a)
noise = rng.normal(0, 1, dim)
noise_orth = noise - np.dot(noise, a) * a
noise_orth /= np.linalg.norm(noise_orth)
b = target_cos * a + np.sqrt(max(0.0, 1 - target_cos ** 2)) * noise_orth
b /= np.linalg.norm(b)
return a, b
def simhash_collision_fraction(a, b, planes):
"""Fraction of random-hyperplane hash bits on which two vectors agree,
i.e., the empirical single-bit collision rate for SimHash."""
ha = (planes @ a > 0).astype(int)
hb = (planes @ b > 0).astype(int)
return np.mean(ha == hb)
def theoretical_collision_probability(cos_sim):
"""Closed-form SimHash collision probability: 1 - theta/pi, where theta
is the angle between the two vectors."""
theta = np.arccos(np.clip(cos_sim, -1.0, 1.0))
return 1 - theta / np.pi
def test_lsh_collision_probability_matches_theory_and_increases_with_similarity():
rng = np.random.RandomState(1)
dim, n_bits, n_trials = 32, 128, 300
planes = rng.normal(0, 1, size=(n_bits, dim))
target_coss = [0.0, 0.3, 0.6, 0.9]
empirical = []
for tc in target_coss:
agreements = []
for _ in range(n_trials):
a, b = make_pair_with_target_cosine(tc, dim, rng)
agreements.append(simhash_collision_fraction(a, b, planes))
empirical.append(np.mean(agreements))
theoretical = [theoretical_collision_probability(tc) for tc in target_coss]
print(f"{'cos sim':>8} | {'empirical':>10} | {'theoretical':>12}")
for tc, emp, th in zip(target_coss, empirical, theoretical):
print(f"{tc:8.2f} | {emp:10.4f} | {th:12.4f}")
assert abs(emp - th) < 0.03, f"Empirical collision rate should closely match theory at cos_sim={tc}"
# Higher cosine similarity should yield strictly higher collision probability
assert empirical[0] < empirical[1] < empirical[2] < empirical[3], \
"LSH collision probability should increase monotonically with cosine similarity"
print("LSH collision probability test passed.")
if __name__ == "__main__":
test_lsh_collision_probability_matches_theory_and_increases_with_similarity()### Lab 3: Product Quantization Reconstruction Error Shrinks With Codebook Size
import numpy as np
def kmeans_subspace(subvectors, n_codewords, n_iters, rng):
n = subvectors.shape[0]
idx = rng.choice(n, n_codewords, replace=False)
codebook = subvectors[idx].copy()
for _ in range(n_iters):
dists = np.linalg.norm(subvectors[:, None, :] - codebook[None, :, :], axis=2)
assignments = np.argmin(dists, axis=1)
for c in range(n_codewords):
members = subvectors[assignments == c]
if len(members) > 0:
codebook[c] = members.mean(axis=0)
return codebook, assignments
def product_quantize(vectors, n_subvectors, n_codewords, rng, n_iters=15):
"""Splits each vector into n_subvectors chunks and independently
vector-quantizes each chunk against its own learned codebook, the core
compression scheme behind product-quantized ANN indexes."""
dim = vectors.shape[1]
sub_dim = dim // n_subvectors
codebooks = []
codes = np.zeros((vectors.shape[0], n_subvectors), dtype=int)
for m in range(n_subvectors):
sub = vectors[:, m * sub_dim:(m + 1) * sub_dim]
codebook, assignments = kmeans_subspace(sub, n_codewords, n_iters, rng)
codebooks.append(codebook)
codes[:, m] = assignments
return codebooks, codes, sub_dim
def reconstruct(codebooks, codes, sub_dim):
n, n_subvectors = codes.shape
dim = n_subvectors * sub_dim
recon = np.zeros((n, dim))
for m in range(n_subvectors):
recon[:, m * sub_dim:(m + 1) * sub_dim] = codebooks[m][codes[:, m]]
return recon
def test_pq_reconstruction_error_decreases_with_codebook_size():
rng = np.random.RandomState(2)
n, dim = 500, 16
vectors = rng.normal(0, 1, size=(n, dim))
codeword_counts = [2, 4, 16, 64]
mses = []
for n_codewords in codeword_counts:
codebooks, codes, sub_dim = product_quantize(vectors, n_subvectors=4, n_codewords=n_codewords, rng=rng)
recon = reconstruct(codebooks, codes, sub_dim)
mse = np.mean(np.sum((vectors - recon) ** 2, axis=1))
mses.append(mse)
print(f"{'codewords':>10} | {'reconstruction MSE':>20}")
for cw, mse in zip(codeword_counts, mses):
print(f"{cw:10d} | {mse:20.4f}")
assert mses[0] > mses[1] > mses[2] > mses[3], \
"Reconstruction error should decrease monotonically as codebook size grows"
assert mses[0] / mses[-1] > 3, \
"Largest codebook should reconstruct substantially more accurately than the smallest"
print("Product quantization compression test passed.")
if __name__ == "__main__":
test_pq_reconstruction_error_decreases_with_codebook_size()### Lab 4: IVF Search Examines Far Fewer Candidates Than Brute Force at Scale
import numpy as np
def kmeans(vectors, n_clusters, n_iters, rng):
n = vectors.shape[0]
idx = rng.choice(n, n_clusters, replace=False)
centroids = vectors[idx].copy()
for _ in range(n_iters):
dists = np.linalg.norm(vectors[:, None, :] - centroids[None, :, :], axis=2)
assignments = np.argmin(dists, axis=1)
for c in range(n_clusters):
members = vectors[assignments == c]
if len(members) > 0:
centroids[c] = members.mean(axis=0)
return centroids, assignments
def build_ivf_index(vectors, n_clusters, rng, n_iters=8):
centroids, assignments = kmeans(vectors, n_clusters, n_iters, rng)
inverted_lists = {c: np.where(assignments == c)[0] for c in range(n_clusters)}
return centroids, inverted_lists
def ivf_search_with_count(query, vectors, centroids, inverted_lists, nprobe, k):
centroid_dists = np.linalg.norm(centroids - query, axis=1)
probe_clusters = np.argsort(centroid_dists)[:nprobe]
candidate_idx = np.concatenate([inverted_lists[c] for c in probe_clusters]) if len(probe_clusters) else np.array([], dtype=int)
n_candidates = len(candidate_idx)
if n_candidates == 0:
return np.array([], dtype=int), 0
dists = np.linalg.norm(vectors[candidate_idx] - query, axis=1)
order = np.argsort(dists)[:k]
return candidate_idx[order], n_candidates
def brute_force_search(query, vectors, k):
dists = np.linalg.norm(vectors - query, axis=1)
return np.argsort(dists)[:k]
def recall_at_k(approx_idx, true_idx):
return len(set(approx_idx) & set(true_idx)) / len(true_idx)
def test_ivf_reduces_candidates_examined_while_maintaining_reasonable_recall():
rng = np.random.RandomState(3)
dim, k, n_queries = 16, 10, 20
collection_sizes = [500, 2000, 8000]
fractions_examined = []
recalls_by_size = []
for n in collection_sizes:
vectors = rng.normal(0, 1, size=(n, dim))
n_clusters = int(np.sqrt(n))
nprobe = max(1, n_clusters // 4)
centroids, inverted_lists = build_ivf_index(vectors, n_clusters, rng)
queries = rng.normal(0, 1, size=(n_queries, dim))
recalls, candidate_counts = [], []
for q in queries:
true_idx = brute_force_search(q, vectors, k)
approx_idx, n_cand = ivf_search_with_count(q, vectors, centroids, inverted_lists, nprobe, k)
recalls.append(recall_at_k(approx_idx, true_idx))
candidate_counts.append(n_cand)
mean_recall = np.mean(recalls)
frac = np.mean(candidate_counts) / n
fractions_examined.append(frac)
recalls_by_size.append(mean_recall)
print(f"N={n:6d} clusters={n_clusters:3d} nprobe={nprobe:3d} "
f"mean_recall={mean_recall:.3f} fraction_examined={frac:.3f}")
# IVF should always examine well under half the full collection per query
assert all(f < 0.35 for f in fractions_examined), \
"IVF should examine only a small fraction of the collection compared to brute force (fraction=1.0)"
# Recall should stay reasonably high across all tested collection sizes
assert all(r > 0.65 for r in recalls_by_size), \
"IVF search should maintain reasonably high recall even while examining a small fraction of the data"
print("IVF candidate-reduction efficiency test passed.")
if __name__ == "__main__":
test_ivf_reduces_candidates_examined_while_maintaining_reasonable_recall()