Matryoshka Representation Learning

A vector index is sized on the day you build it. Ten million documents at 3072 float32 dimensions is about 123 GB before any index overhead; the same corpus at 256 dimensions is 10 GB, and a brute-force scan touches twelve times less memory bandwidth. Pick the large vector and you pay for it on every query forever. Pick the small one and you have capped the ceiling on recall. The usual way to change your mind is to re-embed everything with a different model, which costs an inference run over the whole corpus and invalidates every stored vector. Matryoshka Representation Learning (MRL) takes that decision out of build time: one model emits a vector whose leading coordinates are themselves a working embedding, so the same stored bytes can be read at 3072, 512, or 64 dimensions depending on what a given query is worth (Kusupati et al., 2022, Matryoshka Representation Learning, arXiv:2205.13147).

The loss does all the work

The architecture does not change. No extra heads, no bottleneck layers, no distillation. What changes is that the training objective is evaluated at several prefix lengths at once. Fix a nesting set \(M = \{64, 128, 256, 512, \dots, d\}\) and let \(z_{1:m}\) denote the first \(m\) coordinates of the output vector \(z\). Then

where \(\mathcal{L}\) is whatever loss you were already using (contrastive, InfoNCE, softmax classification) and \(c_m\) weights the granularities, often uniformly. Sentence Transformers ships this as MatryoshkaLoss, typically wrapping a base loss at 768, 512, 256, 128 and 64 dimensions (Sentence Transformers docs).

The consequence is a gradient asymmetry. Coordinate 5 receives gradient from every term in the sum; coordinate 3000 receives gradient only from the largest granularity. The model therefore has an incentive to pack the most broadly useful, coarsest structure into the early coordinates and reserve the tail for fine distinctions that only the full vector needs. Information ends up ordered by importance rather than spread evenly, which is exactly the property that makes truncation safe.

Truncation is not post-hoc compression

Slicing an MRL vector and running PCA on an ordinary one both produce a smaller vector, and they are not the same operation.

PCA, random projection, and autoencoder compression all produce a second artifact: a projection matrix that has to be fitted on a corpus, versioned alongside the index, and applied identically to queries and documents forever. MRL truncation is z[:256]. There is no state, no fitting step, and no risk of a query being projected with a stale matrix. Renormalise after slicing if your similarity assumes unit norm, since a prefix of a unit vector is not itself unit-norm.

The alternative of training separate 64-, 256- and 768-dimensional models is worse in a subtler way. Those models occupy unrelated vector spaces, so a 64-dimensional document vector cannot be compared against a 768-dimensional query vector at all. Under MRL the coarse view lives inside the same space as the fine one, which is what makes the mixed-resolution pipeline below possible.

Truncating a model that was not trained with MRL usually destroys it, because nothing ever asked the early coordinates to stand alone. The Sentence Transformers writeup reports a Matryoshka-trained model retaining 98.37% of its performance at 8.3% of the embedding size, far above the same architecture trained conventionally.

Adaptive retrieval is the payoff

The deployment pattern that MRL enables is a funnel. Score the entire corpus with the truncated prefix, which is cheap in both memory and bandwidth, then rescore a shortlist of a few hundred to a few thousand candidates with the full vector. Both stages read the same stored embedding at different offsets, so there is one index, not two. The original paper reports up to 14x real-world speedups for large-scale retrieval and up to 14x smaller embeddings at matched ImageNet-1K accuracy.

Commercial embedding APIs adopted it quickly. OpenAI's text-embedding-3-large produces up to 3072 dimensions and exposes a dimensions parameter that shortens the vector by dropping coordinates from the end; the company reports that the shortened 256-dimensional version still outscores the older 1536-dimensional text-embedding-ada-002 on MTEB (OpenAI, 2024). nomic-embed-text-v1.5 supports any dimension between 64 and 768 (Nomic, 2024). If an API lets you ask for fewer dimensions and does not require re-encoding, MRL or something close to it is the reason.

When it breaks

Nothing here makes the model cheaper. Training gains extra loss terms, and inference produces the same full-width vector it always did; only the storage and the search get smaller. A team hoping MRL will speed up their encoder has misread it.

Quality does not fall off linearly either. The typical curve is flat for a long stretch and then drops sharply below some corpus-specific \(m\), and where that cliff sits depends on how fine-grained the distinctions in your data are. Domains full of near-duplicates, where the useful signal is precisely the tail dimensions, degrade earliest. Measure the cliff on your own recall@k rather than assuming the vendor's MTEB curve transfers.

Two operational hazards follow. Vectors truncated to different lengths are not comparable, so the truncation length is index metadata and every writer must agree on it. And truncation composes with int8 or binary quantisation, but the errors compound in ways neither technique's published numbers predict, so evaluate the combination end to end instead of assuming the losses add.