Matryoshka Embeddings for Flexible Dims

By Michael Samuel Naeem · December 13, 2025 · 5 min read

You shipped vector search at 3072 dimensions and the index fits — until the catalog triples. Re-embedding with a smaller model means dual indexes, migration downtime, and inconsistent scores across old and new vectors. Matryoshka embeddings fix the dimensionality trade-off at training time: the model learns that the first 256 components already encode most semantic signal, with each additional prefix dimension refining distance estimates. One trained model, multiple deployment profiles — 256-dim for mobile edge cache, 1536-dim for server rerank — without maintaining separate encoders.

How MRL training works

Standard contrastive loss on full-dimension vectors plus auxiliary losses on truncated prefixes:

[ \mathcal{L} = \sum_{d \in {64,128,256,...,D}} w_d \cdot \mathcal{L}{\text contrastive}(x{1:d}, y_{1:d}) ]

Each prefix must be independently useful for matching — not just padding that only makes sense at full D.

OpenAI exposes this via API:

from openai import OpenAI
client = OpenAI()

response = client.embeddings.create(
    model="text-embedding-3-large",
    input="reset password for SSO tenant",
    dimensions=256,  # truncate Matryoshka-trained vector
)
vec = response.data[0].embedding  # length 256

Open-source with Sentence Transformers — use checkpoints trained with MRL loss (nomic, some BGE variants).

Storage and latency math

For N vectors, float32 storage ≈ N × d × 4 bytes.

HNSW search latency scales sublinearly with d but constant factors matter at scale. Truncation reduces RAM and improves cache locality.

Tiered retrieval pattern

Query ──► embed @ 256d ──► coarse HNSW (full catalog)
                │
                ▼ top 500 ids
         re-fetch stored 1536d vectors (or re-embed docs)
                │
                ▼
         cosine rerank ──► top 20

Store full-dimension vectors on disk/object storage; keep short prefixes in RAM index. Alternatively store both prefixes in one record with offset layout.

Do not truncate non-MRL models and expect this to work — arbitrary prefix slicing on legacy embeddings scrambles ranking.

Evaluating truncation on your data

import numpy as np

def recall_at_k(full_embs, trunc_dim, queries, relevant, k=10):
    q = queries[:, :trunc_dim]
    d = full_embs[:, :trunc_dim]
    q = q / np.linalg.norm(q, axis=1, keepdims=True)
    d = d / np.linalg.norm(d, axis=1, keepdims=True)
    sims = q @ d.T
    # argsort, compute recall...

Plot recall@10 vs dimension {64, 128, 256, 512, full}. Pick the knee where marginal recall per megabyte drops.

Training your own MRL head

Fine-tuning existing encoders with nested losses (simplified sketch):

# pseudo: sum losses over truncations
for d in [64, 128, 256, 512]:
    z_q = F.normalize(q_emb[:, :d], dim=-1)
    z_d = F.normalize(d_emb[:, :d], dim=-1)
    loss += weighted_contrastive(z_q, z_d)

Weights often increase with d — full dimension dominates but small prefixes must still pull their weight.

Operational notes

• Index versioning — document which truncation level each index uses.
• Score calibration — cosine scores at 256d are not comparable to 1536d; do not merge ranked lists naively.
• Consistency — query and document must use same truncation at each stage.

Matryoshka embeddings turn dimensionality from a one-way door into a runtime knob — valuable when catalog growth and infra cost share the roadmap with search quality.

Open-source MRL models

Several open models support Matryoshka truncation without retraining:

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("nomic-ai/nomic-embed-text-v1.5")
full_emb = model.encode("query text")           # 768d
fast_emb = model.encode("query text", truncate_dim=128)  # 128d

Use truncate_dim at query time — same model, different index granularity per use case.

Two-stage retrieval with MRL

Stage 1: fast coarse search at low dimension. Stage 2: rerank top-K at full dimension:

def search(query, catalog, k=10):
    # Stage 1: 128d index — fast, approximate
    coarse_emb = model.encode(query, truncate_dim=128)
    candidates = index_128d.search(coarse_emb, k=200)

    # Stage 2: full 768d rerank — precise
    full_emb = model.encode(query)  # 768d
    candidate_embs = [catalog[id].full_emb for id in candidates]
    scores = [cosine(full_emb, c) for c in candidate_embs]
    return sorted(zip(candidates, scores), key=lambda x: -x[1])[:k]

128d index is 6× smaller than 768d — fits in L3 cache. Rerank 200 candidates at full dimension adds ~2ms vs 50ms for full-dimension search over entire catalog.

Index sizing with MRL

Catalog: 10M documents

768d float32 index: 10M × 768 × 4 bytes = 30 GB
128d float32 index: 10M × 128 × 4 bytes =  5 GB
128d int8 index:    10M × 128 × 1 byte  =  1.3 GB

MRL lets you choose the memory/latency tradeoff at deployment time — not at model selection time. Start with 128d index; upgrade to 256d if recall insufficient.

Failure modes

• Query at 128d, index at 768d — dimension mismatch; silent garbage results
• Score comparison across dimensions — cosine at 128d ≠ cosine at 768d; don't merge ranked lists
• Truncation without MRL-trained model — arbitrary truncation destroys semantic quality
• No rerank stage — low-dimension recall loss unacceptable for precision-critical apps
• Index version mismatch — query uses v2 model, index built with v1

Production checklist

• Query and index use same truncation dimension per stage
• MRL-trained model used (not arbitrary truncation)
• Two-stage retrieval: coarse (low-d) + rerank (full-d)
• Index dimension documented in index metadata
• Recall@K evaluated at each truncation level before deployment
• Model version tagged on all index entries

Resources

• Matryoshka Representation Learning paper
• OpenAI embedding models guide
• Nomic embed — open MRL model
• Sentence Transformers documentation
• FAISS index types and memory trade-offs