Matryoshka Coarse-to-Fine — Half the Dimensions, Same Recall¶
What problem does this solve?¶
Large embeddings are expensive at scale. A 3072-dim Voyage 3 vector is 12 KB. At a billion chunks, your vector index is 12 TB of RAM. Even at a tenth that size you have problems: ANN-index memory dominates server bills, replication is painful, and re-embedding when you upgrade models takes days. Matryoshka Representation Learning trains embeddings so the leading dimensions are themselves a usable embedding — slice off the first 128 or 512 dimensions and you have a working embedding at one-tenth the storage cost. The cheap retrieval pattern follows: shortlist with the truncated low-dim vectors over the whole corpus, then re-rank the shortlist with the full-dim vectors. Latency drops because the bulk of the work happens at low-dim; quality stays because the final ranking uses full-dim. The recall loss is usually under a percentage point; the cost saving is an order of magnitude.
Where it came from¶
Matryoshka Representation Learning was published at NeurIPS 2022 by Aditya Kusupati and colleagues at UW. The training trick: regularise the model to ensure the loss is meaningful when computed on the first d_i dimensions for any prefix.
OpenAI's text-embedding-3 family (Jan 2024) was the first widely-used embedding to ship Matryoshka by default. Voyage 3 and Cohere Embed 4 followed. By 2026 most production-grade embedders support truncation; the cookbook uses Nebius's Qwen3-Embedding-8B which exposes the full vector — we truncate manually.
Where it fits in the RAG landscape¶
Matryoshka enables several deployment patterns:
- Coarse-to-fine retrieval (this recipe). Shortlist with truncated vectors, rerank with full. The canonical deployment pattern.
- Multi-resolution indexes. Store the index at multiple dimensions; pick at query time based on latency budget or query difficulty.
- Cheap quantisation. Truncate then quantise; both are lossy but compound nicely. Production systems often run both.
- Cross-modal compatibility. Truncated text embeddings sometimes align better with image embeddings; useful in multimodal retrieval setups.
Compose with ColBERT (Recipe 19) — Matryoshka cuts per-token storage; ColBERT keeps per-token vectors. The combination is the cheapest known way to deploy multi-vector retrieval at scale.
When to use it (and when not to)¶
Use Matryoshka coarse-to-fine when storage cost dominates your retrieval bill. Anything over 10M chunks usually qualifies. Skip it when your embedder does not support truncation. Slicing dimensions on a non-Matryoshka embedder silently destroys quality. Skip it when you only have a few thousand chunks. Full-dim search is already fast; the two-stage complexity is not worth it.
The intuition¶
Three intuitions:
Truncation works because of training. Standard embeddings put their best signal anywhere in the vector. Matryoshka-trained embeddings put the best signal in the leading dimensions by design — the regulariser during training forces the loss to be meaningful at every prefix length.
Two-stage retrieval is the magic. Stage 1 is cheap and noisy (low-dim). Stage 2 is expensive and accurate (full-dim) but runs only on the shortlist. The shortlist is small enough that stage 2's per-vector cost doesn't dominate — the whole pipeline runs roughly at low-dim speed with full-dim quality.
Recall@1 may drop slightly; recall@5 usually doesn't. Stage 1 sometimes ranks the right chunk at position 30; if your shortlist is 50, it survives stage 1 and gets promoted by stage 2's full-dim scoring. Tune shortlist size to your acceptable recall.
Architecture¶
flowchart LR D[Documents] --> E[Embed full-dim] E --> FT[Full vectors
~1024 dim] E --> TR[Truncate to
128 dim] FT --> FI[(Full-dim
lookup table)] TR --> CI[(Coarse index
cheap ANN)] Q[Query] --> QE[Embed] QE --> SH[Shortlist
top-50 via coarse] CI --> SH SH --> RR[Re-rank
with full-dim] FI --> RR RR --> A[Top-5]
References¶
- 📄 Matryoshka Representation Learning (Kusupati et al., NeurIPS 2022) — The original paper.
- 📝 OpenAI text-embedding-3 announcement — Where Matryoshka went mainstream.
- 💻 Nomic Embed v2 model card — Open-weight Matryoshka embedder.
- 📝 Pinecone two-stage retrieval guide — General two-stage retrieval reference.
- 📚 Voyage 3 Large embedding documentation — Hosted Matryoshka-supporting embedder.
- 📄 ColBERT (Recipe 19) — Compose with Matryoshka for cheap multi-vector at scale.
Setup¶
Pick a provider via the PROVIDER env var; everything below is provider-agnostic. The default is Nebius. Tracing is off by default in published notebooks so the outputs are clean — flip COOKBOOK_TRACING=phoenix to send spans to a local Phoenix UI.
import os
os.environ.setdefault('PROVIDER', 'nebius')
os.environ.setdefault('COOKBOOK_TRACING', 'off')
from cookbook.providers import LLMClient
from cookbook.tracing import init_tracing
client = LLMClient()
print(f'Provider: {client.provider} | Chat model: {client.chat_model}')
print(init_tracing())
Provider: nebius | Chat model: meta-llama/Llama-3.3-70B-Instruct Tracing disabled.
Build the Pipeline, Step by Step¶
Step 1 — Load + chunk the Mamba paper¶
We use the arXiv Mamba survey. Long enough to make the two-stage savings non-trivial.
from cookbook.corpora import load_arxiv_mamba
from cookbook.chunkers import sentence_window
docs = list(load_arxiv_mamba())
chunks = sentence_window(docs, sentences_per_chunk=4, overlap=1)
print(f'Chunks: {len(chunks)}')
Chunks: 81
Step 2 — Compute full-dimensional embeddings¶
Standard. The full vectors are normalised; cosine = dot product.
import numpy as np
full = np.asarray(client.embed([c.text for c in chunks]), dtype=np.float32)
full /= np.linalg.norm(full, axis=1, keepdims=True).clip(min=1e-9)
print(f'Full vectors: {full.shape}')
Full vectors: (81, 4096)
Step 3 — Truncate to a coarse dimension¶
Slice the first COARSE_DIM dimensions and re-normalise. With a Matryoshka-trained embedder these truncated vectors are still meaningful; with non-Matryoshka embedders they silently degrade.
COARSE_DIM = 128
coarse = full[:, :COARSE_DIM].copy()
coarse /= np.linalg.norm(coarse, axis=1, keepdims=True).clip(min=1e-9)
print(f'Coarse vectors: {coarse.shape}')
Coarse vectors: (81, 128)
Step 4 — Index the coarse vectors¶
The coarse store is one-eighth the size of the full store. Memory and disk savings scale with the dim ratio.
from cookbook.stores import QdrantBackend
coarse_store = QdrantBackend('mrl-coarse', dim=COARSE_DIM)
coarse_store.add(
[c.text for c in chunks],
coarse.tolist(),
ids=[c.chunk_id for c in chunks],
)
id_to_idx = {c.chunk_id: i for i, c in enumerate(chunks)}
print('Coarse index built.')
C:\Users\faree\Desktop\rag\rag-cookbook-2026\.venv\Lib\site-packages\tqdm\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html from .autonotebook import tqdm as notebook_tqdm
Coarse index built.
Step 5 — Two-stage search function¶
Shortlist with coarse; rerank with full. Full-dim re-ranking is a single matrix-vector product over the shortlist.
def two_stage(question: str, shortlist: int = 50, final: int = 5):
q_full = np.asarray(client.embed([question])[0], dtype=np.float32)
q_full /= np.linalg.norm(q_full) + 1e-9
q_coarse = q_full[:COARSE_DIM]
q_coarse /= np.linalg.norm(q_coarse) + 1e-9
coarse_hits = coarse_store.search(q_coarse.tolist(), top_k=shortlist)
scored = []
for h in coarse_hits:
idx = id_to_idx[h.doc_id]
score = float(full[idx] @ q_full)
scored.append((score, h.text, h.doc_id))
scored.sort(key=lambda x: x[0], reverse=True)
return scored[:final]
for s, t, _ in two_stage('What is selective scan?', shortlist=30, final=5):
print(f' full-dim={s:.3f} {t[:160]}')
full-dim=0.641 2024. Vision mamba: Efficient visual representation learning with bidirectional state space model. arXiv preprint arXiv:2401.09417 (2024). full-dim=0.617 Association for Computing Machinery, New York, NY, USA, 3953–3957. https://doi.org/10.1145/3511808.3557624 [4] Albert Gu and Tri Dao. 2023. Mamba: Linear-time s full-dim=0.615 Uncovering Selective State Space Model’s Capabilities in Lifelong Sequential Recommendation Conference acronym ’XX, June 03–05, 2018, Woodstock, NY interesting full-dim=0.608 2024. Mamba4Rec: Towards Efficient Sequential Recommendation with Selective State Space Models. arXiv preprint arXiv:2403.03900 (2024). [16] Langming Liu, Liu C full-dim=0.595 KEYWORDS Sequential Recommendation, Long-term Recommendation, State Space Models ACM Reference Format: Jiyuan Yang, Yuanzi Li, Jingyu Zhao, Hanbing Wang, Muyang
Step 6 — Wrap as answer_question¶
Standard contract.
def answer_question(question: str, k: int = 5) -> tuple[str, list[str]]:
top = two_stage(question, shortlist=50, final=k)
contexts = [t for _, t, _ in top]
answer = client.chat(
'Use these passages.\n' + '\n\n'.join(contexts) + f'\nQ: {question}\nA:'
)
return answer, contexts
ans, _ = answer_question('Explain selective scan.')
print(ans)
The passages provided do not explicitly explain "selective scan." However, they do discuss "selective state spaces" in the context of sequence modeling and sequential recommendation systems, particularly referencing the Mamba model and its application in efficient visual representation learning and sequential recommendation. From the context, it can be inferred that "selective state spaces" refer to a mechanism or approach in state space models where not all states or all information from previous states are used to predict the next state or make recommendations. Instead, the model selectively chooses which parts of the state space to focus on or retain, potentially improving efficiency and performance by reducing the amount of information that needs to be processed. If "selective scan" is related to this concept, it might imply a process where the model scans through the available data or states but selectively focuses on certain aspects or parts of the data, similar to how a selective state space model operates. This could be a strategy to enhance efficiency, reduce computational requirements, or improve the accuracy of predictions or recommendations by concentrating on the most relevant information. However, without a direct explanation of "selective scan" in the provided passages, this interpretation is based on the broader context of selective state spaces and their application in sequence modeling and recommendation systems.
Look Inside¶
Inspect — coarse-only top-3 vs two-stage top-3¶
Sometimes the coarse top-3 differs from the two-stage top-3. The rerank step picks the truly best chunks from the shortlist.
import numpy as np
q = 'What is selective scan?'
q_full = np.asarray(client.embed([q])[0], dtype=np.float32)
q_full /= np.linalg.norm(q_full) + 1e-9
q_coarse = q_full[:COARSE_DIM] / np.linalg.norm(q_full[:COARSE_DIM])
print('--- coarse top-3 ---')
for h in coarse_store.search(q_coarse.tolist(), top_k=3):
print(f' {h.score:.3f} {h.text[:140]}')
print()
print('--- two-stage top-3 ---')
for s, t, _ in two_stage(q, shortlist=50, final=3):
print(f' {s:.3f} {t[:140]}')
--- coarse top-3 --- 0.670 2024. Mamba4Rec: Towards Efficient Sequential Recommendation with Selective State Space Models. arXiv preprint arXiv:2403.03900 (2024). [16] 0.646 KEYWORDS Sequential Recommendation, Long-term Recommendation, State Space Models ACM Reference Format: Jiyuan Yang, Yuanzi Li, Jingyu Zhao, 0.635 For the two datasets of the 5k version, SASRec has en- countered an out-of-memory (OOM) issue and RecMamba achieves better efficiency compar --- two-stage top-3 --- 0.641 2024. Vision mamba: Efficient visual representation learning with bidirectional state space model. arXiv preprint arXiv:2401.09417 (2024). 0.617 Association for Computing Machinery, New York, NY, USA, 3953–3957. https://doi.org/10.1145/3511808.3557624 [4] Albert Gu and Tri Dao. 2023. 0.615 Uncovering Selective State Space Model’s Capabilities in Lifelong Sequential Recommendation Conference acronym ’XX, June 03–05, 2018, Woodst
Inspect — recall vs shortlist size¶
Sweep shortlist size from 10 to 100 and see how recall changes. Useful for choosing a shortlist that balances latency and quality.
from cookbook.corpora import load_eval_questions
qs = [q for q in load_eval_questions() if q['corpus'] == 'arxiv-mamba'][:8]
import pandas as pd
rows = []
for sl in (10, 25, 50, 100):
hits = 0
for q in qs:
top = two_stage(q['question'], shortlist=sl, final=5)
gold = [w.lower() for w in q['answer'].split() if len(w) >= 4]
if any(any(w[:6] in t.lower() for w in gold) for _, t, _ in top):
hits += 1
rows.append({'shortlist': sl, 'recall@5': hits / max(1, len(qs))})
pd.DataFrame(rows)
| shortlist | recall@5 | |
|---|---|---|
| 0 | 10 | 1.0 |
| 1 | 25 | 1.0 |
| 2 | 50 | 1.0 |
| 3 | 100 | 1.0 |
Inspect — storage savings¶
Concretely: how much less memory does the coarse index use?
full_bytes = full.shape[0] * full.shape[1] * 4
coarse_bytes = coarse.shape[0] * coarse.shape[1] * 4
print(f'Full index: {full_bytes:>10,} bytes')
print(f'Coarse index: {coarse_bytes:>10,} bytes ({coarse_bytes/full_bytes:.1%} of full)')
Full index: 1,327,104 bytes Coarse index: 41,472 bytes (3.1% of full)
Inspect — what happens at very low dim?¶
Truncating to 32 or 16 dims usually breaks retrieval even on Matryoshka-trained embedders. Sweep dim and watch recall fall.
import pandas as pd
rows = []
for d in (32, 64, 128, 256, 512):
coarse_d = full[:, :d]
coarse_d = coarse_d / np.linalg.norm(coarse_d, axis=1, keepdims=True).clip(min=1e-9)
q_full = client.embed(['What is selective scan?'])[0]
q_d = np.asarray(q_full[:d])
q_d /= np.linalg.norm(q_d) + 1e-9
sims = coarse_d @ q_d
top_idx = np.argsort(sims)[::-1][:5]
top_chunks = [chunks[i].text[:60] for i in top_idx]
rows.append({'dim': d, 'top1_preview': top_chunks[0]})
pd.DataFrame(rows)
| dim | top1_preview | |
|---|---|---|
| 0 | 32 | Association for Computing Machinery, New York,... |
| 1 | 64 | 2014. On the properties of neural machine tran... |
| 2 | 128 | 2024. Mamba4Rec: Towards Efficient Sequential ... |
| 3 | 256 | Association for Computing Machinery, New York,... |
| 4 | 512 | 2024. Vision mamba: Efficient visual represent... |
Run It¶
End-to-end on a Mamba question.
ans, ctxs = answer_question('How does the parallel scan implementation matter for Mamba on modern GPUs?')
print('=== Matryoshka two-stage answer ===')
print(ans)
=== Matryoshka two-stage answer === The passage doesn't explicitly explain how the parallel scan implementation matters for Mamba on modern GPUs. However, it does mention that Mamba "utilizes a parallel algorithm optimized for hardware in recurrent mode, enabling effective sequence modeling, particularly for long sequences." This suggests that the parallel algorithm is important for Mamba's performance, especially when dealing with long sequences. It can be inferred that the parallel scan implementation is crucial for Mamba's efficiency on modern GPUs, as it allows for optimized hardware utilization, which in turn enables faster training and inference times, as well as reduced GPU memory consumption. The experimental results show that RecMamba, which is based on Mamba, achieves significant reductions in training duration, inference time, and memory costs compared to other models, such as SASRec.
Side by Side: Vanilla Baseline vs This Technique¶
Vanilla baseline (full-dim search) vs Matryoshka two-stage. The quality should be comparable; the speed/memory advantage shows up at scale.
from cookbook.baselines import vanilla_pipeline
q = 'How does the parallel scan implementation matter for Mamba on modern GPUs?'
base = vanilla_pipeline(q, corpus='arxiv-mamba', top_k=5)
ours_a, ours_c = answer_question(q)
import pandas as pd
pd.DataFrame([
{'pipeline': 'vanilla (full-dim)', 'preview': base.answer[:140]},
{'pipeline': 'matryoshka (two-stage)', 'preview': ours_a[:140]},
])
| pipeline | preview | |
|---|---|---|
| 0 | vanilla (full-dim) | The passages do not contain information about ... |
| 1 | matryoshka (two-stage) | The passage doesn't explicitly explain how the... |
Knobs to Turn¶
Five knobs in priority order:
- COARSE_DIM. Higher means better stage-1 recall, less memory savings. 128 is a strong default for OpenAI text-embedding-3 and Voyage 3; 256 is conservative.
- Shortlist size. Higher catches more recall, slows the re-rank. Typical settings: 30–100. Sweep on your corpus to find where recall flattens.
- Embedder. Matryoshka requires a trained-for-it embedder. Without one, this technique silently degrades — recall drops by tens of points and you have no signal until you measure.
- Quantisation. Compose with int8 quantisation for further savings; both are lossy in different ways and compound nicely. Production deployments often run truncated + quantised.
- Re-rank with cross-encoder. Replace the full-dim re-rank with a cross-encoder reranker (Recipe 22). The two-stage shape stays; the second stage uses a different scorer for higher quality.
Evaluate on a Slice¶
Run the recipe's answer_question over a small slice of the hand-curated eval set. Full RAGAS metrics are exercised in recipes/09-evaluation-and-production/ragas-triad-eval.ipynb; here we just print a quick spot-check table so you can eyeball whether the technique is on track.
from cookbook.corpora import load_eval_questions
from cookbook.eval import EvalSample
qs = load_eval_questions()
qs = [q for q in qs if q['corpus'] == 'arxiv-mamba']
samples = []
for row in qs[:5]:
answer, contexts = answer_question(row['question'])
samples.append({
'question': row['question'],
'expected': row['answer'],
'actual': answer[:200],
'contexts_retrieved': len(list(contexts)),
})
import pandas as pd
pd.DataFrame(samples)
| question | expected | actual | contexts_retrieved | |
|---|---|---|---|---|
| 0 | What problem do state-space models aim to solv... | State-space models target the quadratic time a... | State-space models, such as those mentioned in... | 5 |
| 1 | Describe the selective scan mechanism introduc... | Selective scan makes the SSM parameters input-... | The passage does not explicitly describe the s... | 5 |
| 2 | How does Mamba achieve hardware efficiency on ... | Mamba uses a parallel scan implementation with... | RecMamba achieves hardware efficiency on moder... | 5 |
| 3 | Which earlier model family does Mamba descend ... | Mamba builds on the structured state-space seq... | The passage does not explicitly state which ea... | 5 |
| 4 | Name two domains beyond text where SSM-style b... | Audio modeling and genomics have both seen suc... | Based on the provided passages, two domains be... | 5 |
Closing Thoughts¶
Three failure modes:
- Non-Matryoshka embedder. Truncation silently destroys quality. Check your model card before deploying.
- Shortlist too small. Stage 1 may rank the right chunk at position 60 on hard queries; if your shortlist is 50, you lose it. Tune.
- Re-rank cost. Full-dim rerank over 50 vectors is cheap; over 1000 is not. Stay within reasonable shortlist sizes.
Compose with quantisation (int8 or product quantisation) for further savings; compose with ColBERT (Recipe 19) for cheap multi-vector at scale.