Matryoshka Representation Learning: How a Nesting-Doll Loss Turns One Embedding Model into Many

Every embedding model makes a silent bet about how much storage and compute its users can afford. The model picks a fixed output dimension, say 768 or 1536 floats, and every embedding it produces is exactly that wide, forever. Downstream systems must store, index, and compare those vectors at full width regardless of whether the task demands that precision. A coarse two-class classifier and a fine-grained retrieval system over a billion documents receive the same 1536-dimensional vector even though the classifier would work just as well with 64 dimensions and the retrieval system might benefit from 3072. Aditya Kusupati, Gantavya Bhatt, Aniket Rege, Matthew Wallingford, and their colleagues at the University of Washington, Google, and the Allen Institute identified this mismatch as a fundamental inefficiency in how representation learning pipelines are designed, and in May 2022 they published 'Matryoshka Representation Learning' (arXiv:2205.13147). The paper introduces a training objective that encodes information at multiple granularities simultaneously inside a single embedding vector, so that any prefix of the embedding is itself a valid, useful embedding at its own scale. The name comes from Russian nesting dolls: a single object that contains smaller complete objects inside it.

The fixed-dimension problem and what MRL changes

Standard representation learning trains a model to produce an embedding of size d using a loss function computed over the full d-dimensional output. A ResNet or ViT backbone followed by a linear projection head produces a single vector and optimizes a classification or contrastive loss at that single size. The result is a model that concentrates information across all d dimensions without any guarantee about what partial prefixes of those dimensions contain. If you truncate a standard 2048-dimensional embedding to its first 8 dimensions and try to run nearest-neighbor search, the result is typically worse than an independently trained 8-dimensional model, because the standard model had no reason to front-load useful information into its early dimensions. This is the inefficiency MRL fixes. Its training procedure adds a weighted sum of losses over a set of nested granularities. Rather than optimizing only at dimension d, MRL simultaneously optimizes at dimensions drawn from a set such as 8, 16, 32, 64, 128, 256, 512, 1024, and 2048. At each granularity m, the first m dimensions of the full embedding vector are extracted and passed through a small linear head before the loss is computed. All of these per-granularity losses are summed and backpropagated together. The weight matrix of the backbone is shared; only the linear classification or projection heads differ per granularity, and those add negligible parameters. The final embedding produced during inference is still a single vector of dimension d. What MRL changes is that the model is forced, during training, to make every prefix of that vector coherent and useful on its own.

Why the nesting property matters for retrieval and classification

The nesting property enables several workflows that rigid fixed-dimension models cannot easily support. The most direct application is adaptive storage: a system can choose to store embeddings at 64 or 128 dimensions rather than 2048 and recover most of the retrieval quality at a fraction of the memory cost. Because HNSW and other approximate nearest-neighbor indexes scale in both memory and query latency with embedding dimension, this reduction propagates through the whole retrieval stack. A second application is cascade retrieval: run an initial nearest-neighbor search with short 64-dimensional embeddings to produce a candidate set, then re-rank the candidates using the full 2048-dimensional vectors from the same model without re-encoding any text. The two stages share the same index entries because the short embeddings are exact prefixes of the long ones. A third application is graceful degradation: a system facing memory pressure can truncate stored vectors without replacing them, because any prefix is still a valid MRL embedding. The paper calls these 'coarse-to-fine' representations: earlier dimensions capture the most broadly useful structure, and later dimensions add increasingly fine-grained distinctions.

Benchmark results from the paper

The paper reports results across vision (ResNet-50, ViT-B/16), vision plus language (ALIGN), and language (BERT) on ImageNet-1K, ImageNet-4K, and few-shot classification benchmarks. For ImageNet-1K classification with ResNet-50, an MRL embedding truncated to 1/14th of its full dimension matches the accuracy of an independently trained model at that smaller dimension, demonstrating that MRL representations are at least as informative as separately trained low-dimensional models. The paper reports up to 14x smaller embedding size for ImageNet-1K classification at the same level of accuracy as a full-dimension baseline. For large-scale retrieval on ImageNet-1K and ImageNet-4K, MRL achieves up to 14x real-world speed-ups compared to using full-dimension embeddings while maintaining retrieval quality. On long-tail few-shot classification, MRL embeddings show up to 2% accuracy improvement over full-dimension rigid embeddings of the same model, because the nesting loss encourages the early dimensions to capture broadly transferable features rather than features tuned only for the full-dimension evaluation. The MRL approach is shown to work with no modification to existing training pipelines other than adding the nested losses, and it imposes no overhead during inference or embedding storage: the inference path is identical to a standard embedding model.

• MRL trains a single model to produce embeddings where any prefix of size m from a nested set (8, 16, 32, ..., 2048) is itself a useful embedding at dimension m.
• The training objective sums per-granularity classification or contrastive losses; each uses only the first m dimensions of the shared backbone output.
• Results on ImageNet-1K: up to 14x smaller embedding at equal classification accuracy; up to 14x retrieval speed-up; up to 2% gain on long-tail few-shot classification.
• Compatible with ResNet, ViT, ALIGN (vision plus language), and BERT with no architectural changes and no additional inference cost.
• OpenAI adopted MRL in text-embedding-3-small and text-embedding-3-large (January 2024): users can truncate to any smaller dimension without retraining.
• arXiv:2205.13147, submitted May 26, 2022; authors from University of Washington, Google, Allen Institute for AI; code open-sourced at github.com/RAIVNLab/MRL.

MRL in production: OpenAI text-embedding-3

The clearest confirmation that MRL addresses a real production need came in January 2024, when OpenAI launched text-embedding-3-small and text-embedding-3-large with native support for shortening embeddings. Both models are trained with MRL. text-embedding-3-small produces vectors of up to 1536 dimensions by default, and text-embedding-3-large produces up to 3072 dimensions. Users can pass a dimensions parameter to the API to request any smaller size, and OpenAI guarantees that the truncated vectors are still valid embeddings that maintain quality gracefully as size decreases. On the MTEB English benchmark, text-embedding-3-small scores 62.3% and text-embedding-3-large scores 64.6%, compared to 61.0% for the previous text-embedding-ada-002 model that had no shortening support. On the MIRACL multilingual retrieval benchmark, the improvement is more pronounced: ada-002 scores 31.4%, text-embedding-3-small reaches 44.0%, and text-embedding-3-large reaches 54.9%. The 3-small model is also priced at $0.00002 per 1,000 tokens, a 5x reduction from the $0.0001 price of ada-002. For systems that can operate on shortened embeddings, the effective cost per unit of useful retrieval signal is far lower still. PromptingIndex covers MRL alongside FlashAttention, QLoRA, BitNet, and Grouped Query Attention in its series on efficiency techniques that change what is practical to build and deploy with modern AI models.