Approximate Indexing (IVF)
24 min
Embeddings' use-cases module did brute-force search: compare the query
against EVERY candidate. That's exact, but scales linearly with the
number of vectors — too slow for millions of them. Real vector
databases use approximate indexes instead. IVF (Inverted File
Index) is one of the simplest: pre-cluster all vectors around a fixed
set of centroids, then at search time only compare the query
against vectors in its OWN nearest cluster — a small fraction of the
total.
def cosine_similarity(a, b):
dot = sum(x * y for x, y in zip(a, b))
mag_a = sum(x ** 2 for x in a) ** 0.5
mag_b = sum(x ** 2 for x in b) ** 0.5
return dot / (mag_a * mag_b)
def nearest_centroid(vector, centroids):
best_idx = 0
best_sim = -2
for i, c in enumerate(centroids):
sim = cosine_similarity(vector, c)
if sim > best_sim:
best_sim = sim
best_idx = i
return best_idx
def build_ivf_index(vectors, centroids):
index = {i: [] for i in range(len(centroids))}
for vec_id, vec in vectors.items():
cluster = nearest_centroid(vec, centroids)
index[cluster].append(vec_id)
return index
In a real system, the centroids themselves are learned via k-means clustering over the actual vector data (not chosen by hand) -- but the assignment step (nearest_centroid) works identically either way.
Once vectors are clustered, a search only needs to: (1) find which cluster the QUERY itself is nearest to, then (2) rank only that cluster's members by similarity. This is dramatically faster than scanning everything — the tradeoff is that a true nearest neighbor sitting in a neighboring cluster, just across the boundary, gets missed entirely. That's the "approximate" in Approximate Nearest Neighbor.
def ivf_search(query, vectors, centroids, index, k):
cluster = nearest_centroid(query, centroids)
candidates = index[cluster]
scored = [(vec_id, cosine_similarity(query, vectors[vec_id])) for vec_id in candidates]
scored.sort(key=lambda pair: -pair[1])
return [vec_id for vec_id, _ in scored[:k]]
Using the provided `nearest_centroid`, write `build_ivf_index(vectors, centroids)`: `vectors` is a dict of `id -> vector`. Assign every vector to its nearest centroid's cluster, and return a dict mapping cluster index to the list of vector IDs assigned to it (one empty list per centroid, even if unused).
Using `nearest_centroid`, write `ivf_search(query, vectors, centroids, index, k)`: find the query's nearest cluster, then rank only the vectors already assigned to THAT cluster by similarity to `query`, returning the top `k` vector IDs.
Why is an IVF search only APPROXIMATE, not guaranteed to find the true nearest neighbors, unlike Embeddings' brute-force top_k_similar?
You can build a cluster-based approximate index and search within just the query's nearest cluster, and understand why that trades some recall for a large speedup over brute-force search.