Similarity search

Problem

An embedding model turns documents, images or products into vectors so that related items land near each other. Retrieval then becomes a geometry question: which of the n stored vectors are closest to the query vector? Answering it means picking a measure of "closest", computing it against the stored set fast enough, and knowing when the plain old keyword search it replaced still does better.

Intuition

A vector has a direction and a length. Cosine similarity looks only at direction (the angle between query and document). The dot product looks at direction and length together, so a long vector can win even when it points somewhat away. Euclidean (L2) distance looks at where the tip of the vector lands, so a short vector near the query's tip wins. Three measures, three different rankings, until every vector is scaled to length 1. Then length stops mattering and all three agree.

The diagram plots the query q = (1.0, 0.5) and five documents. Their unit versions sit on the dashed unit circle.

Mechanics

Three measures and when they agree

For a query q and a stored vector x:

dot(q, x) = Σ_j q_j x_j     cos(q, x) = q·x / (‖q‖ ‖x‖)     L2(q, x) = ‖q − x‖

Expanding the squared distance ties them together:

‖q − x‖² = ‖q‖² + ‖x‖² − 2 q·x

• Ranking by L2 is ranking by dot product only when every stored vector has the same length. ‖q‖² is the same for every candidate, so it never changes the order, but ‖x‖² does. A long vector pays a larger ‖x‖² penalty under L2 and gets a larger dot product, so the two measures pull opposite ways.
• Cosine is the dot product of the normalized vectors. On unit vectors ‖q‖ = ‖x‖ = 1, so cos(q, x) = q·x and ‖q − x‖² = 2 − 2 cos(q, x). Smallest distance, largest cosine and largest dot product are then the same item.
• Maximum inner product search (MIPS) is the name for ranking by raw dot product when lengths differ. Some retrievers train for it on purpose and let the length carry a "popularity" or "quality" signal. Then L2 and cosine indexes give the wrong answer.

That is why the usual recipe is to L2-normalize every embedding once at indexing time and the query at search time (PyTorch's F.normalize, scikit-learn's normalize), then rank by dot product. The search becomes one matrix-vector product, and the choice between cosine and L2 disappears. Whether to normalize is a property of the model, not of the search: use the measure the embedding model was trained with.

Exact search by brute force

Stack the n stored vectors as the rows of a matrix X of shape (n, d). Scoring a query is s = Xq, a matrix-vector product costing 2nd FLOPs. The top k of s takes O(n log k) with a heap, or a partial selection (torch.topk, numpy.argpartition). For a batch of m queries, XQ^T is a full matrix multiply. That is exact search: the answer is correct by construction, and it is the ground truth that approximate indexes are graded against (their recall@k).

L2 distance runs on the same matrix multiply. scikit-learn's euclidean_distances computes −2XY^T + ‖x‖² + ‖y‖² and clamps negatives to zero; its docstring warns that this expansion can suffer catastrophic cancellation. torch.cdist uses the same matrix-multiply route for p = 2 once either side has more than 25 rows, unless compute_mode says otherwise.

NearestNeighbors wraps the whole thing. With algorithm="auto" it picks brute force when the data has more than 15 features, so for embeddings it is brute force in practice (the trees from nearest neighbors stop pruning in high dimensions). Its metric="cosine" works only with brute force.

BM25, the lexical baseline

Before embeddings, search engines ranked documents by the query words they contain. BM25 (Okapi BM25) is still the standard version, and still a hard baseline to beat on queries full of names, codes and rare terms. For query terms t, a document D of length |D|, and an average length avgdl over the N documents:

score(D, Q) = Σ_{t ∈ Q} idf(t) · tf(t, D)·(k₁ + 1) / (tf(t, D) + k₁·(1 − b + b·|D|/avgdl))

idf(t) = ln((N − df(t) + 0.5) / (df(t) + 0.5) + 1)

• idf rewards rare terms: a word in every document carries no information. The "+ 1" inside the log is Lucene's variant, which keeps idf positive even for terms in more than half the documents.
• Saturation, k₁ (usually 1.2 to 2). The term-frequency factor rises from 1 at tf = 1 toward k₁ + 1 and never passes it, so the tenth "cat" adds far less than the second.
• Length normalization, b (usually 0.75). A long document contains more words by chance, so its term frequencies are discounted by |D|/avgdl. b = 0 turns this off, b = 1 normalizes fully.

BM25 needs no training and no GPU, runs from an inverted index (term → documents containing it), and is exact about words. It is blind to meaning: "cats" does not match "cat" without stemming, and "kitten" never does. Dense embeddings have the opposite profile.

Hybrid search

Hybrid search runs both retrievers and merges the results. Their scores live on different scales (BM25 is unbounded, cosine sits in [−1, 1]), so the simplest robust merge ignores scores and uses only ranks. Reciprocal rank fusion (RRF) gives each document

RRF(d) = Σ_{r ∈ rankers} 1 / (k + rank_r(d)),   k = 60

where a document a ranker did not return adds nothing for that ranker. The constant k damps the advantage of the very top ranks, so a document both retrievers like beats one that only one retriever puts first. The alternative is a weighted sum of normalized scores, which needs tuning, or a learned reranker on top.

The dense list is an assumed ranking for illustration (no embedding model runs on this site); the BM25 scores and the fused scores are computed from the formulas above.

Worked example

Implementation

import torch
import torch.nn.functional as F

# Index: n embeddings of dimension d, normalized once.
X = F.normalize(embeddings, dim=1)          # (n, d), every row has norm 1

def search(queries, k=10):
    Q = F.normalize(queries, dim=1)         # (m, d)
    scores = Q @ X.T                        # (m, n), cosine = dot on unit vectors
    return scores.topk(k, dim=1)            # values and indices, each (m, k)

# The same thing in scikit-learn.
from sklearn.neighbors import NearestNeighbors
nn = NearestNeighbors(n_neighbors=10, metric="cosine", algorithm="brute").fit(E)
dist, ind = nn.kneighbors(E_query)          # cosine distance = 1 - cos

# Reciprocal rank fusion of any number of ranked id lists.
def rrf(rankings, k=60):
    scores = {}
    for ranking in rankings:
        for rank, doc in enumerate(ranking, start=1):
            scores[doc] = scores.get(doc, 0.0) + 1 / (k + rank)
    return sorted(scores, key=scores.get, reverse=True)

For indexes larger than one device's memory, split X into shards, take each shard's top k, and merge: the global top k is always among the shards' top k lists. Storing X in fp16 halves memory and bandwidth with negligible effect on the ranking. BM25 comes from a search library (Lucene, Elasticsearch, OpenSearch, or a small Python package), not from PyTorch or scikit-learn; scikit-learn's TfidfVectorizer is the nearest built-in lexical scorer.

Tradeoffs

• Exact versus approximate. Brute force has perfect recall, no build step and trivial updates, but cost grows linearly with n. Approximate indexes cut the cost to roughly logarithmic or sublinear at the price of some recall and a build step.
• Measure. Use what the embedding model was trained with. Cosine (normalize, then dot) is the safe default; raw dot product keeps length as a signal; L2 on unnormalized vectors mixes length and direction in a way few models intend.
• Lexical versus dense. BM25 is exact about rare words, IDs and new vocabulary, needs no model, and explains itself. Dense retrieval handles paraphrase and synonyms but can miss exact matches and fails quietly out of domain. Hybrid search usually beats either alone.
• Fusion. RRF needs no score calibration and no tuning beyond k, but throws away score gaps. Weighted score fusion keeps them but needs per-collection normalization.
• Precision of the index. fp16 or int8 storage cuts memory and doubles or quadruples throughput on a memory-bound scan; near-ties can swap order.

Connections

References

• Robertson, Walker, Jones, Hancock-Beaulieu and Gatford. Okapi at TREC-3. TREC, 1994 (BM25).
• Robertson and Zaragoza. The Probabilistic Relevance Framework: BM25 and Beyond. Foundations and Trends in Information Retrieval, 2009.
• Cormack, Clarke and Büttcher. Reciprocal Rank Fusion Outperforms Condorcet and Individual Rank Learning Methods. SIGIR, 2009.
• Karpukhin et al. Dense Passage Retrieval for Open-Domain Question Answering. EMNLP, 2020. arXiv:2004.04906.
• Johnson, Douze and Jégou. Billion-scale similarity search with GPUs. 2017. arXiv:1702.08734.
• Manning, Raghavan and Schütze. Introduction to Information Retrieval. Cambridge University Press, 2008: chapters 6 and 11.
• PyTorch v2.15.0: distance.py, F.normalize, torch.cdist.
• scikit-learn 1.9.1: NearestNeighbors, cosine_similarity, euclidean_distances.