Matryoshka Embeddings: Truncating Vectors Without Losing Much

Take a normal embedding, keep the first 256 of its 3072 numbers, throw the rest away, and search with what is left. With an ordinary model this destroys the vector, because no dimension is more important than any other. With a Matryoshka-trained model it mostly works, and that difference is a training objective rather than a trick applied afterwards.

The idea

In a standard contrastive objective, the loss is computed on the full d-dimensional output. Nothing pushes information toward the front of the vector, so the coordinates are interchangeable and a prefix is a random projection of no particular quality.

Matryoshka Representation Learning — Kusupati et al., NeurIPS 2022 — changes the objective so that the same loss is computed at several nested prefix lengths at once, typically a geometric ladder such as 8, 16, 32, 64, ..., 2048. The model is graded simultaneously on how good its first 8 numbers are, its first 16, its first 32, and so on up to the full width. The only way to score well on all of them at once is to put the coarsest, most discriminative structure at the front and use later coordinates for refinement. Hence the nesting-doll name: every prefix is a complete, usable embedding.

The loss that makes prefixes work

# standard contrastive training
loss = infonce(z_query, z_positive, z_negatives)

# matryoshka: same loss, evaluated on every nested prefix
loss = 0
for k in [64, 128, 256, 512, 1024, 2048, 3072]:
    loss += w[k] * infonce(z_query[:k], z_positive[:k], z_negatives[:k])

That is the whole modification. The weights w[k] are commonly left uniform. The cost is a modest amount of extra work per training step; the benefit is that dimension becomes a deployment-time parameter rather than a training-time one. Notice that this is strictly a property of how the model was trained — you cannot make an existing non-Matryoshka model truncatable by wishing, and truncating one is a quiet way to lose a great deal of recall.

What has actually been published

Two sources are worth quoting rather than paraphrasing. The original paper reports, among other results, embeddings up to 14× smaller at the same ImageNet-1K classification accuracy, and correspondingly large speed-ups for large-scale retrieval using an adaptive two-stage scheme. The paper is on vision and language tasks both, and the claim it makes is about the method, not about any particular text model you might use.

For text specifically, OpenAI’s January 2024 announcement of text-embedding-3 stated that the models were trained with this technique and exposed it as a dimensions API parameter, and that text-embedding-3-large shortened to 256 dimensions still scores above unshortened ada-002 at 1536 on MTEB. That is one vendor’s figure for one model pair on one benchmark, which is exactly how much weight it should carry — but it is a published figure and it is a striking one: a 6× smaller vector outperforming the previous generation.

What nobody has published is how far your corpus tolerates truncation, and that is the number that decides your storage bill. The last section is how to get it.

Using it, including the step people miss

If the API takes a dimensions parameter, pass it and the provider does the correct thing. If you are truncating locally — from a full-width vector you already stored, or from an open model — you must re-normalise afterwards:

import numpy as np

def truncate(v, k):
    head = v[:k]
    return head / np.linalg.norm(head)     # <- this line is not optional

Skipping the re-normalisation is the single most common implementation error here. A unit-length 3072-dim vector sliced to 256 dimensions is no longer unit length — it has lost whatever magnitude lived in the discarded tail, and different texts lose different amounts. If you then score with a dot product, you are ranking partly by how much of each document’s energy happened to sit in the first 256 coordinates, which is not a relevance signal. Cosine hides the error, dot product does not, and mixed pipelines where queries are normalised and documents are not produce a ranking that looks plausible and is subtly wrong.

Adaptive retrieval: shortlist small, rescore full

The pattern the paper calls adaptive retrieval is where the technique earns its keep. Keep a small index of truncated vectors in memory, search that to get a generous shortlist, then rescore the shortlist with the full-width vectors fetched from wherever they are cheap to store. Working the numbers for 100 million vectors at 3072 dimensions:

full index, 3072 dims float32:
  100e6 * 3072 * 4 = 1,228.8 GB   -> a sharded cluster

truncated index, 256 dims float32:
  100e6 *  256 * 4 =   102.4 GB   -> one large machine

rescoring 1,000 candidates against full vectors:
  1000 * 3072 * 4  =    12.3 MB of reads per query

One caveat that decides whether this works: the shortlist has to contain the right documents, because rescoring can only reorder what it is given. Recall of the truncated first stage at the shortlist depth is therefore the number to watch, not its recall at 10. A 256-dimension index that finds the correct document somewhere in its top 1,000 for 99% of queries is a perfectly good first stage even if its own top-10 is mediocre, and that is the measurement to run before committing to the architecture.

Twelve megabytes of mostly-random reads per query is comfortable on NVMe and trivial from an object store with a cache in front of it. You have converted a 1.2 TB RAM problem into a 100 GB RAM problem plus a small read amplification, and the final ranking is computed on full-precision vectors so the top of the list is unchanged as long as the shortlist contained the right documents.

Finding your own cut-off

The procedure is cheap because it needs no re-embedding — every prefix is already inside the vectors you have.

• Take your gold set of queries and known-correct chunks. A hundred queries is enough to see the shape of the curve.
• For each k in 64, 128, 256, 512, 1024 and full width: truncate every document vector and every query vector to k, re-normalise, and compute recall@10 by exhaustive search. Exhaustive is fine here — you are measuring the representation, not the index, and mixing in an approximate index would confound the two.
• Plot recall against k. The curve is typically flat then falls off a cliff. Take the smallest k still on the flat part, then step one rung back up as insurance against the queries your gold set does not contain.
• If you plan to quantise as well, run this after choosing the quantisation, not before. Truncation and quantisation both spend from the same accuracy budget and their effects are not additive in any convenient way.