One Embedding Model, Any Dimension: Cut Vector Storage 8x With Matryoshka Representation Learning

TL;DR: MatryoshkaLoss trains a single embedding model to remain accurate when truncated — 1024d for peak quality, 128d for 8x storage savings — from the exact same weights, with minimal quality loss.

What you'll be able to do after this:

• Train an embedding model with MatryoshkaLoss so vectors degrade gracefully at any truncation point (768 → 512 → 256 → 128 → 64 dims)
• Switch between embedding sizes at inference time with a single slice — no retraining, no separate models to maintain
• Build a two-stage retrieval pipeline: pre-screen millions of chunks at 128 dims with fast ANN search, then re-rank the top-k at full 768 dims for final quality

Why this matters

Indexing a million chunks at 1536 dims (OpenAI's default) consumes ~6 GB of vector storage. Matryoshka Representation Learning (MRL) solves this without retraining separate models: it trains the model to pack the most important signal into the first N dimensions, so you can truncate any time.

The trick: during training, the contrastive loss is computed not just on the full embedding, but also on truncated prefixes — at 512, 256, 128, and 64 dims simultaneously. The model learns that its first 64 dimensions must carry meaningful signal, with later dimensions adding progressively more. Result: 97%+ of full-quality retrieval at 256 dims on standard benchmarks.

Walk-through

pip install sentence-transformers datasets

Wrap your base loss in MatryoshkaLoss:

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer
from sentence_transformers.losses import MultipleNegativesRankingLoss, MatryoshkaLoss
from datasets import load_dataset

model = SentenceTransformer("BAAI/bge-base-en-v1.5")

base_loss = MultipleNegativesRankingLoss(model)
matryoshka_loss = MatryoshkaLoss(
    model,
    base_loss,
    matryoshka_dims=[768, 512, 256, 128, 64],  # all dims get supervised
)

dataset = load_dataset(
    "sentence-transformers/msmarco-co-condenser-margin-mse-sym-mnrl-mean-v1",
    split="train[:20000]"
)
trainer = SentenceTransformerTrainer(model=model, train_dataset=dataset, loss=matryoshka_loss)
trainer.train()
model.save_pretrained("my-matryoshka-bge")

Truncate at inference time — it's just a slice:

import numpy as np

query_full = model.encode("How does transformer attention work?")  # shape (768,)

# Pre-screening at 128 dims: 6x less memory, nearly as fast
query_small = query_full[:128]
query_small = query_small / np.linalg.norm(query_small)  # re-normalize after truncation

The two-stage retrieval pattern

For large-scale retrieval: build your ANN index at 128 dims (fast, cheap), pre-screen to the top 50 candidates, then re-rank those 50 with the full 768-dim vectors for final precision. This cuts ANN index memory 6x while keeping full-quality results for the user.

Many off-the-shelf models already include MRL — check matryoshka_dimensions in their config_sentence_transformers.json. If it's there, truncation is free and no training is needed. Models with MRL support: Nomic Embed, e5-mistral-7b-instruct, Cohere's embed-v3 family.

Sources: Introduction to Matryoshka Embedding Models — HuggingFace Blog · Matryoshka Embeddings — Sentence Transformers Docs · Train Embedding Models with Sentence Transformers — HuggingFace Blog