Matryoshka Embedding Truncation

Background

Matryoshka Representation Learning (MRL) trains an embedding model such that any prefix of the output vector is a valid embedding in its own right. At serving time, truncating from $d=1024$ to $d/k$ gives a $k$× smaller index with only a small recall hit — and the embeddings still need to be L2-renormalised after slicing so cosine similarity is meaningful.

Problem statement

Implement truncate_and_renorm(embs, dim) returning (n, dim) L2-normalised embeddings:

1. Slice the first dim columns.
2. Divide each row by its (post-slice) L2 norm.

Input

• embs — np.ndarray shape $(n, d)$, the original embeddings (need not be normalised).
• dim — int in $[1, d]$, the target dimension.

Output

Returns np.ndarray shape $(n, \text{dim})$ with each row's L2 norm equal to $1$.

Examples

Example 1 — shape

Input:  embs shape=(5, 1024), dim=128
Output: shape=(5, 128)

Example 2 — re-normalised

After truncation, ||row||_2 should be 1.0 ± 1e-9 for every row.

Example 3 — dim == d is the full-vector L2-normalised case

Input:  embs=[[3.0, 4.0]], dim=2
Output: [[0.6, 0.8]]

Constraints

• dim must be in [1, embs.shape[1]]. ValueError otherwise.
• Add 1e-12 epsilon to the denominator to avoid division-by-zero on zero rows.
• Original embs is not mutated.

Notes

• Recall hit. Truncating $1024 \to 128$ typically loses ~2pp recall@10 on MRL-trained embeddings; truncating an embedding NOT trained for MRL can lose 10pp+.
• Combine with quantisation. MRL truncation + binary or scalar quantisation gives 30–100× total index reduction — that's the production stack at most vector DB providers.