In short: three metrics (inner product, cosine, negated squared L2) are all scored so that higher means more
similar, and cosine is turned into an inner product by storing every vector at unit length. A size-k min-heap,
internal::TopK, keeps the best results with ties broken towards the lower id. FlatIndex
scores every stored vector, O(n·d) per query, and is the exact baseline. IvfIndex trains
num_lists centroids with rml/classical's k-means, files each added vector
under its nearest centroid, and at search time scans only the num_probes lists whose centroids score
best; probing every list gives back the exact answer. Both are plain single-threaded C++ over
std::vector<float>, and errors come back as Status.
The concepts are on /ml-models: embeddings,
similarity search, vector indexes,
clustering, two-tower models and
retrieval-augmented generation. The module is two libraries,
similarity and ivf, each with its own test
(BUILD.bazel).
Metrics and scores
Every search returns Neighbor{id, score} with a higher score meaning more similar, so one heap and
one ordering serve all three metrics (similarity.h L7–12).
That is why L2 is negated: the nearest vector has the largest −|q − x|².
| Item | Type/Signature | Meaning | Source |
|---|---|---|---|
Metric | enum class { kInnerProduct, kCosine, kL2 } | Inner product ⟨q, x⟩, what a model trained with dot-product scores wants; cosine ⟨q, x⟩ / (|q| |x|), where only direction counts; L2 as −|q − x|². | similarity.h L26 |
Neighbor | struct { int64_t id; float score; } | One result; defaulted == so tests compare whole result lists. | L28–33 |
Similarity | float(Metric, span<const float> q, span<const float> x) | The score of x for q. Different lengths are a programmer error (RML_CHECK). Cosine with a zero vector scores 0 rather than dividing by zero. Sums accumulate in float, one element at a time. | .h L35–36, .cc L39–55 |
SearchResults | vector<vector<Neighbor>> | The k best for each query row, best first. | .h L38–39 |
A worked example
The metric test scores q = (1, 2, 2) against x = (2, 0, 1) (similarity_test.cc L20–27):
| Metric | Computation | Score |
|---|---|---|
kInnerProduct | 1·2 + 2·0 + 2·1 | 4 |
kCosine | 4 / (|q| |x|) = 4 / (3 · √5) | ≈ 0.596 |
kL2 | −((1−2)² + (2−0)² + (2−1)²) | −6 |
kCosine against (0, 0, 0) | norms are 0 | 0 |
Cosine and inner product can disagree on the winner. With stored vectors (1, 0) and (10, 10) and the query (1, 0.1), cosine picks id 0 (almost the same direction) while inner product picks id 1 (much longer) (similarity_test.cc L86–95).
Keeping the k best
internal::TopK is a min-heap of at most k results ordered by Better: higher score
first, then lower id (similarity.cc L20–23). Because the
heap's comparator is "better", the front of the heap is the worst result kept, so each push is one
comparison against it and, if the candidate wins, a pop and a push: O(log k) per scored vector and O(k) memory
however many vectors are scored.
| Item | Type/Signature | Meaning | Source |
|---|---|---|---|
TopK | explicit TopK(int64_t k) | k must be positive (RML_CHECK). | .h L63–74, .cc L59–61 |
Push | void(int64_t id, float score) | Fills the heap up to k, then replaces the front only when the candidate is strictly better, so an equal score with a higher id never displaces a lower one. | .cc L63–73 |
Take | vector<Neighbor>() | Sorts best first and moves the heap out, leaving it empty. | .cc L75–78 |
PrepareRows | Expected<vector<float>>(const Tensor&, int64_t dim, Metric, const char* what) | The shared input check: InvalidArgument unless the tensor is f32 of rank 2 with dim columns, naming what and the shape it got. Copies the rows out contiguously (views are materialized with Contiguous()) and, under cosine, normalizes each row; a zero row stays zero. | .h L76–79, .cc L80–94 |
The test pushes scores 0.5, 2, 1, 2, −1, 1 for ids 0–5 into a TopK of 3 and gets {(1, 2), (3, 2), (2, 1)}: both 2s, and of the two 1s the lower id (similarity_test.cc L34–39).
FlatIndex: exact search
FlatIndex stores every vector in one row-major std::vector<float>, unit length
under cosine, and compares each query with all of them (similarity.h
L13–15). It is always exact, and it is what the IVF tests measure recall against.
| Item | Type/Signature | Meaning | Source |
|---|---|---|---|
| constructor | FlatIndex(int64_t dim, Metric metric) | dim must be positive (RML_CHECK). | .h L41–59, .cc L98–100 |
Add | Status(const Tensor& vectors) | Appends the rows of f32 [n, dim]; ids continue from size(), so adding in two parts gives the same ids as one. | .cc L102–107 |
Search | Expected<SearchResults>(const Tensor& queries, int64_t k) const | For each row of f32 [m, dim], the min(k, size()) best stored vectors, ties towards the lower id. k ≤ 0 is InvalidArgument. Under cosine both sides are already unit length, so it scores with plain inner product. An empty index returns empty lists. | .h L47–49, .cc L109–128 |
| accessors | dim(), metric(), size() | size() is the stored floats divided by dim. | .h L51–53 |
IvfIndex: the inverted file
An inverted-file index splits the stored vectors into num_lists clusters and, for each query,
scores only the vectors in the num_probes clusters whose centroids look best, about
num_probes / num_lists of the data. A true neighbour just across a cluster boundary is missed, so recall trades
against speed through num_probes (ivf.h L4–16).
NearestLists is a TopK over
centroid ids (ivf.cc L56–68).| Item | Type/Signature | Meaning | Source |
|---|---|---|---|
Options | num_lists = 16, num_probes = 1, seed = 0 | Cluster count, lists scanned per query, and the k-means seed. | ivf.h L29–33 |
| constructor | IvfIndex(int64_t dim, Metric, const Options&) | dim and num_lists must be positive, and the probe count goes through set_num_probes. | ivf.cc L26–31 |
Train | Status(const Tensor& sample) | Needs at least num_lists rows, else InvalidArgument. Runs classical::KMeans with k = num_lists, num_init = 1, max_iterations = 50 and the options' seed (k-means' own defaults are 4 inits and 300 iterations), copies the centers into centroids_, and resets every list and size_, so training again drops what was added. | .cc L39–54, kmeans.h L30–38 |
Add | Status(const Tensor& vectors) | InvalidArgument before Train. Each row goes to its single nearest list with the next id. | .cc L70–82 |
Search | Expected<SearchResults>(const Tensor& queries, int64_t k, int64_t* scanned = nullptr) const | The min(k, scanned) best vectors in the probed lists, best first. scanned, when given, receives the total vectors scored over all queries. Fails before training or with k ≤ 0. | .h L43–47, .cc L84–108 |
set_num_probes | void(int64_t) | Must be in [1, num_lists] (RML_CHECK); changes recall without retraining. | .cc L33–37 |
| accessors | trained(), size(), ListSizes() | Trained once centroids exist; ListSizes gives how many vectors each list holds. | .h L49–53, .cc L110–114 |
| storage | struct List { ids; vectors; } | One row-major vector block per list with the ids beside it; centroids are row-major [num_lists, dim]. | .h L56–70 |
How the metric reaches the clusters
Under cosine every vector is unit length before k-means sees it, so the clustering, the assignment and the
scan all work on directions (ivf.h L15–16). The centroids are means
of unit vectors and are not normalized again; ranking them by inner product still works because the query is
unit length (ivf.cc L60–62). ScoringMetric maps
cosine to inner product everywhere else (L19–22).
The centroids always come from k-means, which clusters by Euclidean distance, but lists are assigned and
probed with the index's own metric. Under kL2 the two agree. Under kInnerProduct a
vector goes to the centroid with the largest inner product, which favours long centroids, so the lists need not
match k-means' own labels; the parameterized tests show recall still rises with probes for all three metrics.
Recall against probes
ivf_test builds 4,096 points around 32 shared, overlapping centers in 16 dimensions, plus 100
queries from the same distribution, and indexes them with 32 lists
(ivf_test.cc L20–38,
L52–68). Recall is the share of the exact top 10 that IVF
returns. The asserted bounds, for every metric (L92–115):
| num_probes | Share of data scanned (about) | Asserted recall@10 |
|---|---|---|
| 1 | 1/32 | below 0.85: overlapping clusters put many neighbours in other lists |
| 2 | 1/16 | at least the 1-probe recall |
| 4 | 1/8 (asserted under 1/6) | at least 0.9 |
| 8 | 1/4 | at least 0.95 |
| 32 | all 409,600 scores | exact: same ids as FlatIndex, scores within 10−4 |
The scanned shares in the middle column are num_probes / num_lists, which is what the lists average; only the 1/6 bound at 4 probes and the full count at 32 are checked.
Who uses it
| Caller | Index | How | Source |
|---|---|---|---|
examples/rag | FlatIndex, cosine | The knowledge base adds one embedding per entity at construction, and Retrieve returns the top-k entity ids per question for the reader to condition on. | rag.cc L43–53, L59–69 |
examples/two_tower | IvfIndex, inner product | Serves the trained item embeddings with 32 lists and 8 probes by default, and compares the top 10 for each test user against the exact top 10 of the same embeddings, recording recall and the fraction of items scored. | two_tower.h L112–119, two_tower.cc L258–293 |
The two-tower header records that the IVF index finds 79% to 95% of the exact top 10 over seeds 0 to 3 while scoring 10% to 24% of the items (two_tower.h L37–40). Nothing in the inference stack or the LLM runtime uses this module. The examples are described on Examples.
Tests
| File | Tests | What they pin down |
|---|---|---|
similarity_test.cc | 7 (one parameterized over the 3 metrics) | Each metric's score, and cosine with a zero vector; vectors of different lengths die; TopK keeps the best and breaks ties by id; FlatIndex matches a stable-sorted brute force for every metric, with ids continuing across two Adds; cosine ignores length where inner product does not; at most size() results, including from an empty index; wrong widths, ranks and k ≤ 0 are InvalidArgument. |
ivf_test.cc | 5 (three parameterized over the 3 metrics) | Every vector filed exactly once; probing all 32 lists is exact; recall never falls as probes grow and meets the bounds above; Add and Search before Train, too few sample rows and the wrong width fail; probe counts of 0 or above num_lists die. |