Matryoshka: One Vector, Every Budget

Chapter 1: The Nested Doll

A matryoshka doll is not a container. That distinction is the whole idea, so let us be pedantic about it.

A container holds a thing. Open a lunchbox and you find a sandwich; the lunchbox itself is not food. A matryoshka is different: open the biggest doll and inside is a smaller doll that is itself a complete doll — painted, symmetric, standing upright on a shelf without apology. Open that one and there is another. Every level is a finished object. Nothing is packaging.

Now transfer that to a vector. An ordinary compressed embedding is a container: you have a small code, and to use it you decompress it back to the big one. A matryoshka embedding is the doll: the first 64 numbers of the 768-number vector are, by themselves, a finished 64-dimensional embedding. You do not decode them. You do not project them. You use them.

What nesting is not — four things it gets confused with

Because the word “smaller” is doing a lot of work, it is worth ruling out four neighbours explicitly. Each is a real technique that produces a smaller thing, and each differs from nesting in a way that matters operationally.

The early-exit row is the one worth internalizing, because it names the boundary of what this paper can do for you. MRL does not make inference cheaper. You run the full encoder, produce the full 768 numbers, and then choose how many to keep. If your bottleneck is the encoder, this technique is aimed at the wrong resource — and Chapter 9 will point you at the one aimed at the right one.

Why a prefix, and not some other subset

You could imagine asking for a different structure — “any 64 of the 768 coordinates work”, or “here is a learned mask per budget”. Prefixes win for three specific engineering reasons, and it is worth being able to say all three.

One: slicing a prefix is genuinely free. A vector in memory is a base address and a length. The first m coordinates share the base address; you change the length. In NumPy z[:m] returns a view — no copy, no allocation. In C it is the same pointer. On disk you can mmap a 768-dimensional index and read only the first 64 floats of every row if you lay it out row-major, though in practice you store the truncated index separately so the reads are contiguous. An arbitrary subset, by contrast, requires a gather: 64 scattered loads instead of one 256-byte streaming read, which on modern hardware is roughly an order of magnitude worse.

Two: prefixes nest transitively. The first 8 coordinates are a prefix of the first 16, which are a prefix of the first 32. So the granularities form a chain: z_1:8 ⊂ z_1:16 ⊂ … ⊂ z_1:768. That chain is exactly what lets you build the cascade in Chapter 6, where each stage refines the previous stage's shortlist using strictly more information than the stage before. An arbitrary-subset scheme has no chain; the 64-dim mask and the 256-dim mask might not even overlap.

Three: a prefix induces an ordering, and an ordering is exploitable. If the first coordinates matter most, then coordinate index becomes a proxy for importance. You get, for free, an answer to “which parts of this embedding should I keep if I can only keep some?” — and, as Chapter 8 shows, that same ordering tells a quantizer where to spend its bits.

Read the “Nested?” column. Only two rows say yes, and one of them costs a matrix multiply and a corpus-specific artifact. That is the design space, and it is small.

What “free” means at the byte level

“Slicing is free” is the sort of claim that deserves one paragraph of literalism, because the whole deployment story rests on it.

A batch of embeddings in memory is one contiguous block. For B documents at d = 768 float32, the block is B × 3,072 bytes, laid out row-major: document 0's 768 numbers, then document 1's, and so on. Taking a prefix means reading the first 256 float32s of each row and skipping the rest.

what a prefix costs, three ways# 1. In-memory view - no copy at all
Z = np.load("index.npy", mmap_mode="r")   # (12_000_000, 768) float32
P = Z[:, :256]                              # a VIEW. zero bytes allocated.
# but note: P is strided - each row skips 512 floats - so a scan over P
# still touches every cache line of Z. Views are free; SCANS are not.

# 2. Materialised prefix index - what you actually ship
np.ascontiguousarray(Z[:, :256]).tofile("index256.bin")
# 12.3 GB one-off write, then every scan reads 12.3 GB instead of 36.9 GB

# 3. Truncating a single query vector at request time
q256 = q[:256]; q256 /= np.linalg.norm(q256)   # ~1 microsecond

Where does an ordinary embedding put its information?

Everywhere, in roughly equal amounts. And this is not an accident or a flaw — it is the correct answer to the question an ordinary loss asks.

Consider what a standard training objective sees. It gets one representation z ∈ R⁷⁶⁸ and one head W ∈ R^L×768, computes Wz, and pays a loss. Now ask: does that loss have any preference about which coordinates carry the signal? Take any invertible matrix R and replace z with Rz and W with WR⁻¹. The logits are identical:

The loss cannot tell these apart. Every rotation of the representation space is exactly as good, so the training run wanders into whichever one the initialization and the optimizer happen to hand it. In the language of optimization, the objective has a large symmetry group — a set of transformations under which it is invariant — and any property you want that is not invariant under that group must be put into the objective explicitly, because the optimizer will not choose it for you.

This also explains why the failure is so severe rather than merely mild. With information spread uniformly over 768 coordinates, keeping 64 of them keeps 64/768 = 8.3% of the coordinates. Chapter 3 shows that in the presence of noise this translates to keeping only √(64/768) = 28.9% of the signal-to-noise ratio, which for a hard task is the difference between working and not working.

What the nesting constraint costs

Nothing is free. If z_1:8 has to be a good 8-dimensional embedding, then the full z is no longer allowed to be an arbitrary 768-dimensional embedding — you have removed a lot of the rotational freedom described above. So there ought to be a price at d = 768. Why might it be small?

Reason one: the constraint is a chain, not a set of conflicts. The nesting requirements are compatible with each other in a strong sense. Suppose you have found a solution good at m = 8. Now extend to m = 16 with a head whose extra columns are all zero. The logits at 16 are then identical to the logits at 8:

So a solution that is good at every smaller granularity can always be extended to be at least as good at the next one. The family of constraints admits a monotone solution by construction. This is why we should expect quality to be non-decreasing in m rather than trading off between granularities, and it is why the objective in Chapter 2 is a plain sum rather than something adversarial.

Reason two: deep representations are enormously redundant. A ResNet50 penultimate layer has 2048 coordinates, but the set of images it actually sees does not fill 2048 dimensions. Empirically the intrinsic dimension — the smallest number of coordinates needed to parameterize the data manifold without loss — of such representations is far lower than the ambient dimension. If most of the 2048 numbers are describing directions the data never explores, then reorganizing so that the used directions come first costs almost nothing.

Reason three: the price falls where you can afford it. The constraint binds hardest on the small prefixes, because they are the ones being asked to do something they could not otherwise do. At m = 768 the constraint is nearly vacuous — you already have all the coordinates. So the tax lands on the granularity that was previously impossible, which is a strange kind of tax: you are being charged for a capability you did not have.

Set the cut to 8 and compare. The ordinary encoder's cumulative curve is a straight line — every coordinate contributes the same, so keeping 8 of 64 captures 12.5% of the signal. The matryoshka curve is concave and steep at the left: the first 8 coordinates already capture most of what a linear decision needs. The bars are not sorted by a post-processing step. That shape is what the training objective produced.

The extendability argument, on actual numbers

The claim that a good solution at granularity m can always be extended to be at least as good at the next one deserves a concrete instance, because it is the reason to expect the quality curve to be monotone rather than a trade-off frontier.

Suppose a two-class problem with a representation whose first two coordinates already separate perfectly, and a head W whose first two columns are (1.0, 0.0) for class A and (−1.0, 0.0) for class B. On an example with z = [0.8, 0.3, anything, anything], granularity 2 gives logits (0.8, −0.8), a margin of 1.6.

Now extend to granularity 4 by appending two columns of zeros to each row:

So the granularity-4 term can always be made no worse than granularity 2 by zeroing the new columns, and optimization will only leave that solution if the extra columns strictly help. The constraint set therefore contains a monotone solution and the objective has no reason to prefer a non-monotone one.

What the constraint costs, in bits

A rough information-theoretic sanity check, because “the tax is small” should have some quantitative grounding beyond “the experiments say so.”

An unconstrained d-dimensional representation may be rotated freely. The set of rotations of R^d — the orthogonal group — has d(d − 1)/2 degrees of freedom. For d = 2048 that is 2,096,128 free parameters of orientation, every one of which the ordinary loss is blind to.

MRL spends some of that orientation freedom to satisfy the nesting constraints. How much? Roughly, forcing the first m coordinates to span the best m-dimensional task subspace fixes the orientation of that subspace — m(d − m) parameters — while leaving rotations within the prefix and within the tail untouched.

Nine nested constraints together fix a nested flag of subspaces, still a small fraction of the whole. And crucially none of this touches the representation's information content: rotations are invertible, so a rotated representation contains exactly what the original did. The tax is paid in orientation freedom, not in bits — which is a strong hint about why the cost at full width is close to zero, and it is the closest thing to an explanation the literature currently offers.

Choosing M — why exponential spacing

The nesting set M is a design choice. The paper uses exponentially spaced granularities: for a 2048-dimensional ResNet50 representation,

Why exponential? Two independent arguments land on the same answer.

The quality argument. Representation quality is roughly logarithmic in dimension over the useful range: doubling from 8 to 16 buys about as much as doubling from 512 to 1024. If quality is roughly linear in log m, then sampling m on a log grid samples the quality axis uniformly. A linear grid would waste almost all of its granularities in the flat high-dimensional region where they are indistinguishable, and would have nothing to say in the steep low-dimensional region where every doubling matters.

The cost argument. Every granularity is one more classifier head and one more loss term. On a linear grid with stride 8 you would have 2048/8 = 256 granularities: 256 heads, 256 softmaxes, 256 backward passes through heads. On a log grid you have log₂(2048/8) + 1 = 9. Chapter 4 prices this exactly, and the difference between 9 and 256 is the difference between “negligible” and “this doubles training cost.”

Note the 768-dimensional row. 768 is not a power of two, so the chain is built by repeatedly halving from the top: 768, 384, 192, 96, 48, 24, 12. Everything still nests, because nesting only requires that the sizes be a chain of increasing integers — nothing about the arithmetic requires powers of two.

The metaphor, pushed until it breaks

Analogies teach by being right in one place, and they teach twice as much when you find where they go wrong. Four ways a matryoshka doll is not a matryoshka embedding:

What survives the demolition is exactly one claim, and it is the only one we need: every level is a complete, finished object. That is the whole metaphor, and it is worth keeping because the alternative mental models — compression, packaging, downsampling — all carry an implied decode step that does not exist here.

How many dolls? The |M| question

Choosing M has two independent parts: how far down it reaches, and how many granularities it contains. The first is a product decision — the smallest size any consumer will ask for. The second is a modelling one, and it has a clean answer.

The number of granularities is the number of distinguishable importance tiers your coordinates get. With |M| = 9 the multiplicity takes nine values, so the objective can distinguish nine levels of “how broadly useful should this coordinate be” across 2048 coordinates. Within a tier — coordinates 513 through 1024, say — the objective is completely indifferent to ordering, and the model will not produce one.

The last row is the one that shows why exponential spacing is not a stylistic choice. Linear spacing over 2048 dimensions makes the classifier heads ten times larger than the network they sit on, which turns a 0.2% training overhead into a 10× one. Exponential spacing buys nine tiers for the price of two heads, because — as Chapter 2 will compute exactly — the sum of a doubling chain is less than twice its largest term.

What about dimensions that are not in M?

A fair worry: you trained with M = {12, 24, 48, 96, 192, 384, 768}, and now a service asks for 100 dimensions. Is z_1:100 garbage?

It is not, and the reason is structural rather than lucky. The pressure the objective puts on coordinate j is a function of how many granularities include it, and that function is a staircase that decreases smoothly as j grows (Chapter 2 derives the staircase exactly). Coordinate 97 is included in every granularity from 192 upward, coordinate 100 in the same ones. Neighbouring coordinates get near-identical treatment, so the resulting importance ordering is smooth rather than blocky. Truncating between the trained granularities interpolates between their qualities; the paper reports that MRL remains accurate at intermediate dimensions it never explicitly optimized.

Prior art, and the exact gap MRL fills

The idea of an information ordering is old. It is worth seeing the neighbours, because knowing what already existed sharpens what is new.

Read the last column of the slimmable-networks row carefully, because it is the sharpest distinction. Slimmable networks make the model elastic; MRL makes the output elastic. In a retrieval system those are different bills: the model runs once per query while the output is stored once per document, and there are a billion documents and one query. Chapter 9 shows that the two ideas eventually merged.

The property, as a test you could actually write

A definition you can assert in code is a definition you understand. Here is the nesting property expressed as a check, which is also the shape of the acceptance test you should run before believing any claim of matryoshka-ness.

python — the nesting contract, as an assertiondef assert_matryoshka(model, dims, eval_fn, reference, tol=0.02):
    """
    model      : produces (N, d) embeddings
    dims       : the granularities claimed to be supported
    eval_fn    : (embeddings) -> a scalar quality metric
    reference  : {m: quality of a model TRAINED natively at m}
    """
    E = model.encode(eval_set)                  # (N, d) - ONE pass
    prev = -1.0
    for m in sorted(dims):
        q = eval_fn(l2(E[:, :m]))               # slice, THEN normalise

        # (1) each prefix is about as good as a dedicated m-dim model
        assert q >= reference[m] - tol, \
            f"prefix {m} scores {q:.4f}, native model gets {reference[m]:.4f}"

        # (2) quality never decreases with m - Chapter 1's extendability
        assert q >= prev - 1e-3, \
            f"quality fell from {prev:.4f} to {q:.4f} going up to m={m}"
        prev = q

        # (3) the prefix is literally a slice - no transform anywhere
        assert np.shares_memory(E[:, :m], E)

Three assertions, and each one is a different chapter of this lesson. The first is the definition. The second is the monotonicity that the extendability argument guarantees, and its failure means undertraining rather than a fundamental limit. The third is the one people forget to check, and it is the one that separates matryoshka from every neighbour in the table above: if a transform is being applied, you do not have a nested representation, you have a compression scheme wearing its clothes.

What a matryoshka embedding looks like, numerically

Abstractions are easier to hold once you have seen the numbers they describe. Here is the shape of the coordinate-magnitude profile for two 768-dimensional text embedders — illustrative values of the kind you get from the twenty-minute diagnostic in Chapter 3, not measurements from any specific model.

Both rows of totals are 1.000 because both vectors are unit-norm — the same amount of “stuff”, distributed completely differently. The ordinary column is flat to three decimal places, which is exactly what Chapter 1's rotation-invariance argument predicts: with nothing in the loss expressing a preference, 1/768 = 0.00130 per coordinate is what you get.

The matryoshka column falls by a factor of 26 from the first block to the last. And the first 64 coordinates — 8.3% of the vector — hold 64 × 0.0081 = 0.518, or 52% of the total squared magnitude. Check the second row too: coordinates 65 through 128 add 64 × 0.0034 = 0.218, so the first 128 hold 74%.

Two names for two things, so Chapter 2 reads cleanly

One thing to fix in your head now, because it is the most common misreading of this paper: MRL is a training objective, not an architecture. There is no matryoshka layer. The network is whatever network you already had. What changes is the sum you minimize.

Chapter 2: The Objective, Term by Term

Here is the entire method. One line. Read it once, do not try to understand it yet, and then we will take it apart symbol by symbol until nothing in it is mysterious.

That is Matryoshka Representation Learning. There is nothing else. If you have written a training loop with a cross-entropy loss, you have written 90% of this already; the remaining 10% is a for loop and a slice.

Walking the equation, left to right

min over {W^(m)}_m∈M and θ_F. Two groups of parameters are being learned jointly: the backbone weights θ_F, and one linear head W^(m) for each granularity. Both groups see gradients from the same loss. This matters: the backbone is not frozen and then fitted with nine heads afterwards. The nesting requirement reaches all the way back through the encoder, which is why the encoder learns to produce front-loaded vectors rather than merely being read front-loadedly.

(1/N) ∑_i=1^N. The ordinary average over a minibatch of N examples. Nothing special. Note that unlike a contrastive loss, there is no interaction between examples here — each example's loss is computed independently, so batch size is a throughput knob, not a modelling knob.

∑_m∈M. This is the new sum, and it is the whole paper. Instead of one loss per example there are |M| losses per example, one for each granularity. Nine terms for the ResNet50 setup. Every one of them is a complete, ordinary classification loss.

F(x_i; θ_F). The backbone applied to the input. Concretely for the paper's ImageNet setup: x_i is an image tensor of shape (3, 224, 224); F is a ResNet50 up to and including global average pooling; the output z_i = F(x_i) has shape (2048,). For a text setup: x_i is a sequence of token ids of shape (L,), F is BERT-base with mean or [CLS] pooling, and z_i has shape (768,).

…_1:m. The slice. Take the first m entries of a (2048,) tensor and you have an (m,) tensor. In PyTorch this is z[:, :m] on a batched (N, 2048) tensor, giving (N, m). It allocates nothing and it is differentiable: the backward pass of a slice is “route the incoming gradient to the first m positions and put zeros in the rest.” No parameters. No cost worth measuring.

W^(m) · (…). The granularity head. W^(m) ∈ R^L×m where L is the number of classes — 1000 for ImageNet-1K. Multiplying an (m,) prefix by an (L, m) matrix gives an (L,) vector of logits, the unnormalized scores that a softmax will turn into probabilities.

L( · , y_i ). Multi-class softmax cross-entropy, exactly as you already know it: exponentiate the logits, normalize to a distribution, take the negative log of the probability assigned to the true class y_i.

c_m. A positive scalar setting how much granularity m matters relative to the others. The paper sets c_m = 1 for all m — every granularity weighted equally. Chapter 4 examines whether that is a lazy choice or a considered one.

MRL and MRL-E: independent heads versus a sliced head

There are two variants, and the difference is one line of code with a genuinely different parameter story.

MRL keeps |M| separate matrices: W⁽⁸⁾ ∈ R^1000×8, W⁽¹⁶⁾ ∈ R^1000×16, and so on up to W⁽²⁰⁴⁸⁾ ∈ R^1000×2048. Each granularity gets its own freedom to decide how to read its prefix.

MRL-E (“efficient”) ties them all to slices of one matrix. Keep a single W ∈ R^L×d and define

Now the head is matryoshka too. And this is more elegant than it first looks. Because column j of W multiplies coordinate j of z, tying the heads means every granularity must agree on what coordinate j means. Under untied MRL, granularity 8 could interpret coordinate 3 one way and granularity 2048 another way; under MRL-E they are forced into a single shared reading.

Count the parameters for ImageNet-1K, L = 1000, d = 2048:

Verify the inner sum by hand, because it is a pleasant one: 8+16 = 24, +32 = 56, +64 = 120, +128 = 248, +256 = 504, +512 = 1016, +1024 = 2040, +2048 = 4088. Notice the pattern — each partial sum is one less than double the next term, so the total is just under twice the largest granularity. The full geometric chain costs less than two of its largest element. That single fact is why nine granularities is cheap and why 256 linearly-spaced ones would not be.

The paper reports MRL-E as slightly weaker than MRL, which is the expected direction: tying weights is a constraint, and constraints cost accuracy. The reason MRL-E exists at all is the extreme-classification regime, where L can be in the millions and 4.09M × (L/1000) stops being a rounding error.

Now the gradient — where the magic actually is

Everything so far has been bookkeeping. This section is the mechanism, and it is a genuinely beautiful one because it is not a mechanism at all in the usual sense: nobody wrote a sorting step, an orthogonality penalty, or an importance regularizer. The ordering falls out of a counting argument.

Ask: which loss terms does coordinate j of z appear in?

Coordinate j appears in the prefix z_1:m if and only if m ≥ j. So the total gradient arriving at coordinate j is

Read that sum's index set. For j = 1, every granularity in M satisfies m ≥ 1, so coordinate 1 receives gradient from all nine loss terms. For j = 2000, only m = 2048 qualifies, so coordinate 2000 receives gradient from one.

Define the multiplicity of coordinate j as the number of granularities that see it: μ(j) = |{ m ∈ M : m ≥ j }|. For the ResNet50 nesting set:

Check the widths sum correctly: 8 + 8 + 16 + 32 + 64 + 128 + 256 + 512 + 1024 = 2048. Good.

There is a second, subtler effect hiding in the same sum, and it is worth spelling out because it explains why the small prefixes get good rather than merely getting attention. The nine terms are not nine copies of the same demand. The m = 8 term demands that eight numbers alone separate 1000 classes — a brutal requirement that can only be satisfied by extremely coarse, high-level structure. The m = 2048 term permits fine, specialized distinctions. So the early coordinates are being asked for coarse information by construction, not merely for more information. The representation ends up ordered coarse-to-fine, which is exactly the ordering you want for a cascade.

The per-coordinate gradient, written out

Let us make it fully concrete for softmax cross-entropy, because the shape of the expression explains the previous paragraph.

For a single granularity m, write the logits u^(m) = W^(m)z_1:m ∈ R^L and the softmax probabilities p^(m) = softmax(u^(m)). The classic result for cross-entropy with true class y:

where [k = y] is 1 for the true class and 0 otherwise. Chain that through the head:

And the total, summing over the granularities that contain j:

Now the important observation. The error signal (p^(m) − one-hot) is different for every m, because each granularity has its own predictions and therefore its own residual. Granularity 8, working with almost no capacity, will still be badly wrong about fine-grained classes; its residual is dominated by coarse confusions. Granularity 2048 has nearly solved the coarse problem and its residual is dominated by fine confusions. So coordinate 1 receives a gradient that is mostly the coarse residual repeated nine times, while coordinate 2000 receives only the fine residual once.

That is the coarse-to-fine ordering, derived. Not asserted, not hoped for — it is a direct consequence of summing residuals from classifiers of different capacity.

Turn off the 8 and 16 granularities and watch the staircase's tall left step disappear: with M starting at 32, coordinates 1 through 32 all get identical multiplicity, and the model has no reason to order within that block. That is the design knob — M determines the resolution of the importance ordering, and below the smallest granularity there is no ordering at all.

Why a sum? The combiners that were not chosen

The |M| losses have to be combined into one scalar, and “add them” is such a natural move that it slips past unexamined. It should not, because the alternatives all have advocates and each fails in an instructive way.

The random-sampling row is worth a second look because it clarifies what the sum is doing. Sampling one granularity per step gives the same objective in expectation but a noisier gradient. Summing all of them is variance reduction by exhaustive enumeration — affordable precisely because |M| is small and the heads are cheap. If someone proposed matryoshka over 256 linearly spaced granularities, sampling would suddenly be the right call.

How the gradient actually gets there: the slice's backward pass

One more piece of concreteness, because “autograd handles it” hides the exact thing we care about. A slice is a linear operation, so it has a transpose, and the transpose is what the backward pass applies.

In words: the incoming gradient of length m is written into the first m slots and zeros are written everywhere else. Do that once per granularity and accumulate, and the buffer at z ends up holding exactly the multiplicity sum:

the gradient buffer at z, filled granularity by granularityd = 16, M = [2, 4, 8, 16]

after L2  backward:  [ g g . . . . . . . . . . . . . . ]
after L4  backward:  [ g g g g . . . . . . . . . . . . ]
after L8  backward:  [ g g g g g g g g . . . . . . . . ]
after L16 backward:  [ g g g g g g g g g g g g g g g g ]
                       ─────────────────────── accumulated ───────
number of writes:      4 4 3 3 2 2 2 2 1 1 1 1 1 1 1 1
                       ^                             ^
                    coordinate 1                 coordinate 16
                    written 4 times              written once

That bottom row is the multiplicity staircase again, now visible as literal accumulation counts in a gradient buffer. There is no cleverness anywhere in the implementation — there is a += that happens more often on the left than on the right.

Real shapes, three real models

Rows two and three have no heads at all, which is worth pausing on. In a contrastive setting the “head” that reads the prefix is the dot product itself, and a dot product has no parameters. So the entire matryoshka modification collapses to: slice, normalize, compute the similarity matrix, sum the losses. There is nothing to add to the model — only to the loop.

The implementation, in full

Twenty lines. Every shape annotated. This is not pseudocode; it is the whole thing.

python — pytorch, the entire methodimport torch, torch.nn as nn, torch.nn.functional as F

class MatryoshkaHead(nn.Module):
    """One linear head per granularity. MRL (untied)."""
    def __init__(self, dims, n_classes):
        super().__init__()
        self.dims  = dims                       # e.g. [8,16,32,...,2048]
        self.heads = nn.ModuleList([nn.Linear(m, n_classes, bias=False)
                                    for m in dims])

    def forward(self, z):                    # z: (N, d)
        return [h(z[:, :m]) for h, m in zip(self.heads, self.dims)]
        #      ^ z[:, :m] is a VIEW - no copy, no allocation
        # returns |M| tensors, each (N, n_classes)

class MatryoshkaHeadE(nn.Module):
    """MRL-E: one weight matrix, sliced by columns."""
    def __init__(self, dims, n_classes, d):
        super().__init__()
        self.dims = dims
        self.W    = nn.Parameter(torch.empty(n_classes, d))   # (L, d)
        nn.init.trunc_normal_(self.W, std=0.02)

    def forward(self, z):
        return [z[:, :m] @ self.W[:, :m].T for m in self.dims]
        #       (N,m) @ (m,L) -> (N,L)   the SAME columns, every time

def mrl_loss(logits_list, y, c=None):
    """logits_list: |M| tensors of shape (N, L). y: (N,) int64."""
    if c is None: c = [1.0] * len(logits_list)      # the paper's choice
    return sum(cm * F.cross_entropy(u, y)
               for cm, u in zip(c, logits_list))

# --- the training step is otherwise completely unchanged ---
z      = backbone(x)                 # (N, 2048)  ONE forward pass
logits = head(z)                     # list of 9 tensors, each (N, 1000)
loss   = mrl_loss(logits, y)
loss.backward()                     # gradients merge at z automatically
opt.step()

Two lines deserve a second look. z[:, :m] is a view, so the nine “different inputs” are nine windows onto one tensor — autograd handles the merge at z without you writing anything, and the merge is the multiplicity sum from three sections ago. And sum(...) over the nine cross-entropies builds a single scalar, so there is exactly one backward(): nine loss terms, one backward pass, one optimizer step.

Three sanity checks on the equation itself

Before trusting any objective, poke it at its boundaries. Three degenerate cases, each of which should give an answer you can predict.

What if M = {d} only? The sum has one term and the equation collapses to ordinary training. Every coordinate has multiplicity 1, no ordering pressure exists, and you have exactly the model you had before. Good — the objective contains its own baseline as a special case, which is the first thing any generalization should do.

What if every c_m = 0 except the smallest? Then only z_1:8 is supervised. The remaining 2040 coordinates receive no gradient at all through the heads, so they drift under weight decay toward zero and the model is effectively an 8-dimensional encoder padded with noise. Predictable, and a useful reminder that the large granularity is not a free rider — it is the term keeping the tail meaningful.

What if M = {1, 2, 3, …, d}, every integer? Multiplicity becomes exactly d − j + 1 — a perfectly linear ramp, and the strongest possible ordering pressure. Also 2048 heads totalling L × 2,098,176 parameters, which for L = 1000 is 2.1 billion — eighty times the backbone. The objective does not break; the budget does. This is the cleanest way to see that exponential spacing is a compute decision rather than a modelling one.

Check the last row's arithmetic: the sum of 1 through 2048 is 2048 × 2049 / 2 = 2,098,176, times 1000 classes. The geometric chain summed to 4,088; the arithmetic chain sums to 2,098,176 — five hundred times more. That single comparison is the entire justification for exponential spacing, and it takes one line of Gauss to see.

Adapting the objective to a non-classification task

The equation above uses classification cross-entropy because that is what the paper trains on, but nothing about the structure depends on it. Replace L with any per-example loss and you have a matryoshka version of that method.

The contrastive row is the practically important one, because every production text embedder in Chapter 7 is trained contrastively, not with a classification head. Note the crucial detail hidden in it: each prefix must be L2-normalized on its own. If you normalize the full 768-dimensional vector and then truncate, the resulting 64-dimensional vector does not have unit norm — and worse, its norm varies from example to example depending on how much mass sat in the tail. Cosine similarity computed on non-unit vectors silently becomes a length contest. Normalize after slicing, always.

Reading the objective as a constrained problem

One more angle on the same equation, because it clarifies what c_m is rather than merely what it does.

What we actually want is: minimize the full-dimension loss, subject to every prefix also being good. Write that as a constrained program — minimize L_d subject to L_m ≤ τ_m for each m < d — and form the Lagrangian:

The τ terms are constants and vanish from the gradient, so this is exactly the MRL objective with c_d = 1 and c_m = λ_m. The relative importance weights are Lagrange multipliers on the nesting constraints.

Two things follow. First, setting every c_m = 1 is a claim that all the constraints are equally binding — defensible, but a modelling assumption rather than a neutral default. Second, there is a principled procedure available when you need one: run the constrained problem properly, raising λ_m while constraint m is violated and lowering it while it has slack. That is dual ascent, and it is what you would reach for given a hard requirement like “quality at 64 dimensions must be within one point of a native 64-dimensional model.”

The last row matters the moment you apply MRL to a mixed objective — a retrieval loss plus a classification loss, say. Cross-entropies over the same L classes are commensurable and need nothing; a cross-entropy and a mean-squared error are not, and summing them raw quietly lets one term own the gradient.

Chapter 3: Why Plain Embeddings Do Not Truncate

Chapter 1 asserted that cutting an ordinary embedding short destroys it. That was an assertion. This chapter earns it, with a model simple enough to compute in your head and general enough that the conclusion transfers.

We will build the argument in four steps. First a noise model that makes “how much information survives truncation” a number. Then the exact scaling law, derived. Then two worked examples that turn the scaling law into accuracies. Then the post-hoc fixes — PCA, LDA, random projection — and precisely where each one runs out.

A model of what an embedding contains

Every useful embedding is a mixture of two things: signal, the part that varies with what you care about, and nuisance, everything else — sensor noise, phrasing accidents, lighting, the particular way this document happened to be worded.

Take the simplest possible instance. Two classes. An embedding in R^d. Class “+” has mean μ₊ = +s u and class “−” has mean μ₋ = −s u, where u is a unit vector giving the direction that separates the classes and s is how far apart they are along it. Around each mean, add independent Gaussian noise with standard deviation σ in every coordinate.

Two things about this model deserve a defence, because a model you do not believe teaches nothing.

Why isotropic noise? Because that is the neutral assumption — it says the nuisance directions have no preferred orientation, which is what you get when the encoder has no reason to prefer one. If the noise were anisotropic and you knew its shape, you would whiten it away first, and after whitening the noise is isotropic by construction. So this is the post-whitening picture.

Why a single signal direction? For two classes there genuinely is only one direction that matters — the difference of means — and everything orthogonal to it is nuisance by definition. With L classes the signal occupies a subspace of dimension at most L − 1, and the argument below repeats independently along each of its directions. Two classes is not a simplification that hides anything; it is one factor of the general case.

How much signal survives a cut

The best you can do with a full d-dimensional vector is to project it onto u and threshold. That projection gives a scalar: the signal contributes ±s, and the noise contributes a Gaussian of standard deviation σ (projecting isotropic noise onto any unit vector gives variance σ²). Define the standard measure of separability, the discriminability index:

And for two equally likely Gaussian classes with equal covariance, the optimal error rate is

Now truncate to the first m coordinates. What happens to each piece?

The noise is unchanged per coordinate. Each retained coordinate still carries noise of standard deviation σ. You did not remove noise; you removed coordinates, and the ones you kept are as noisy as before.

The signal shrinks by however much of u lived in the prefix. The class means in the truncated space are ±s u_1:m, so the separation is 2s ‖u_1:m‖. Therefore:

Case A: the democratic embedding — what ordinary training gives you

An ordinary encoder has no reason to align u with any particular axis, so the signal direction is effectively a random unit vector — and a random unit vector in high dimensions spreads its mass almost perfectly evenly. Take the exactly-even case:

So for the democratic case:

A square root, not a linear factor. That is worse than the naive intuition in one direction and better in another, so let us be careful: keeping 8.3% of the coordinates keeps 28.9% of the discriminability, not 8.3%. The square root is generous. It is also nowhere near generous enough.

Worked example 1. Take d = 768 and a task where the full embedding is comfortably good: s/σ = 3, so d′ = 6 and the error is Φ(−3) = 0.00135. That is 99.87% accuracy — a solved problem. Now cut.

Check one row by hand to be sure the table is not decoration. At m = 64: √(64/768) = √(1/12). Since 12 = 4 × 3, √(1/12) = 1/(2√3) = 1/3.4641 = 0.28868. Multiply by 6: d′₆₄ = 1.7321. Half of that is 0.8660, and Φ(−0.866) = 1 − 0.8068 = 0.1932. So 80.68% accuracy. ✓

Read the first and last rows together. A task that was 99.87% solved is 62.03% solved — barely better than a coin — when you keep 8 of 768 coordinates. And remember this is two classes. On a 1000-class problem the same √(m/d) shrinkage applies simultaneously to every one of the 999 signal directions, and the accuracy collapse is far steeper because you now need many separations to survive at once.

A sanity check: is the √ law believable?

A derived law you cannot check independently is a law you should not trust, so verify it from a completely different direction — energy.

The signal vector is s u with total squared length s², spread evenly over d coordinates, so each coordinate holds s²/d of signal energy and σ² of noise energy. Keep m coordinates:

Two things fall out, and both are worth holding separately. Energy is linear in m; amplitude is the square root of energy. Discriminability is an amplitude quantity — a distance in units of a standard deviation — so it inherits the square root. Anyone who tells you that keeping 8% of the coordinates keeps 8% of the quality is quoting an energy fraction where an amplitude one belongs.

Second check, from the noise side. The noise in the retained subspace has total energy mσ², so noise amplitude grows as √m while signal amplitude also grows as √m — and their ratio would be constant if the signal were confined to the kept coordinates. It is not: the signal amplitude in the prefix grows as √m only because the signal is spread, while the best detector projects noise down to a single direction regardless. That asymmetry is where the whole effect lives, and it is why concentrating the signal (Case B) beats spreading it.

Case B: the concentrated embedding — what we wish we had

Now suppose the encoder had put the entire signal direction into coordinate 1: u = e₁ = (1, 0, 0, …, 0). Then for every m ≥ 1:

Truncation is free. Cut to one dimension and lose nothing, because everything you were going to use is already there. The other 767 coordinates were carrying nuisance, and throwing away nuisance costs nothing.

Two encoders. Identical information content — you can rotate one into the other and back, losing nothing. Wildly different truncation behaviour. The difference is not what the encoder knows. It is where the encoder put it.

Scaling the argument to many classes — why the collapse is steeper than the table

Two classes gave us one signal direction. With L classes there are up to L − 1 independent directions the representation must keep apart, and the √(m/d) shrinkage applies to each of them at once.

Think about what that does to accuracy. For a single pair of classes to be confused, one separation has to fail. For a 1000-way prediction to be correct, essentially all the separations involving the true class must hold simultaneously. Approximate the classifier as needing all 999 pairwise comparisons against the true class to come out right, each independently succeeding with probability 1 − e(m) where e is the pairwise error at that dimension:

The independence assumption is crude — real confusions are correlated, which makes the true accuracy higher than this — but the qualitative behaviour it predicts is right and dramatic. Take three pairwise error rates and see:

Check the first: ln(0.9987) = −0.0013009; times 999 = −1.2996; e^−1.2996 = 0.2726. (Close to the 0.2597 above; the difference is rounding in the logarithm — either way, the point stands.) A pairwise error rate that looks superb in a two-class problem is already fatal at 1000 classes under this crude model, which is exactly why real classifiers do far better than it predicts: confusions are highly correlated, most class pairs are trivially separable, and only a handful compete.

Then why not just rotate afterwards? Enter PCA

If the problem is orientation, rotate. Principal component analysis takes your embeddings, computes their covariance matrix, and finds the orthogonal directions along which they vary most, largest first. Keep the top m and you have the m-dimensional subspace that best reconstructs your data in a least-squares sense.

PCA genuinely helps. It is cheap, it is standard, and against naive truncation it is a large improvement. But it optimizes variance retained, and variance is not accuracy. Here is the counterexample, small enough to verify by inspection.

Worked example 2 — the two-coordinate trap. An embedding with d = 2. Coordinate 1 encodes a nuisance factor — say document length, or image brightness — that varies wildly and has nothing to do with the label. Coordinate 2 encodes the label cleanly but with small amplitude.

The total variance of coordinate 2 is the within-class variance plus the between-class variance, because the two class means sit at ±0.6 around zero: the mean-square of the means is 0.6² = 0.36. So 0.04 + 0.36 = 0.40.

PCA keeps the direction of largest variance. 9.00 versus 0.40 — a ratio of 22.5 to 1 — so PCA at m = 1 keeps coordinate 1 and discards coordinate 2. Now compute both outcomes:

PCA did not underperform. PCA achieved exactly chance while a one-dimensional representation existed that achieved 99.87%. It made the worst possible choice available to it, and it did so while correctly maximizing the thing it was designed to maximize. The objective was wrong, not the algorithm.

This is not a contrived pathology. Nuisance factors with large amplitude are the normal condition of learned representations: document length, speaker identity, image exposure, JPEG artifacts, template boilerplate. They vary enormously and predict nothing.

Then why not use a supervised rotation? Enter LDA

Fair. Linear discriminant analysis finds the directions maximizing between-class scatter relative to within-class scatter — exactly the right objective for the example above, and it would pick coordinate 2 instantly. So why is post-hoc LDA not the answer?

The deep limit of every post-hoc method. A rotation is invertible: it moves information around without adding or removing any. So if you ask “what is the best m-dimensional linear subspace of this fixed encoder's output for this specific linear task”, a supervised rotation gets you the optimum and MRL cannot beat it. That sounds like a knockout. It is not, and here is why.

The downstream reader is a linear probe, or a cosine similarity, or a nearest-neighbour lookup — all of them linear or near-linear operations. An ordinary encoder, trained with capacity to spare, has no reason to make coarse structure linearly available in few dimensions. It might encode “is this an animal” in a way that requires combining eighty coordinates nonlinearly, because at 768 dimensions with a linear head that is fine — the head has 768 coefficients to play with. No rotation can turn that into a property readable from 8 coordinates, because no rotation changes what is linearly readable in a subspace of that size.

MRL's m = 8 loss term is a linear head on eight coordinates. During training, that term is only satisfiable if the encoder makes coarse structure linearly separable in eight dimensions. So the encoder changes. It produces different vectors than it otherwise would have.

And random projection? The honest bound

One more contender, because it is the theoretically respectable one. The Johnson–Lindenstrauss lemma says a random linear map into m dimensions preserves all pairwise distances among n points to within a factor (1 ± ε), provided

Delightful theorem, and it has the enormous virtue of requiring no assumptions and no fitting. Now put your numbers in. Twelve million documents, 10% distortion tolerance:

Verify the logarithm: ln(1.2 × 10⁷) = ln(1.2) + 7 ln(10) = 0.182 + 16.118 = 16.30. ✓ So the guarantee demands seventeen times more dimensions than the 768 you started with.

This is not a defect of the lemma, it is a lesson about what guarantees cost. JL is distribution-free: it works for adversarially chosen points, so it must pay for the worst case. Real embeddings are not adversarial — they live near a low-dimensional manifold — and a method that exploits that structure will beat a method that refuses to assume it. Random projection is the price of assuming nothing, and it is a high price.

It also reveals a subtlety worth keeping. JL preserves distances. Retrieval only needs the ranking of distances, which is weaker. And a classification task only needs the decision boundary, which is weaker still. Every step down that ladder buys you dimensions, which is exactly why 64 well-chosen coordinates can work when the distance-preservation bound says 13,000.

Three features of that picture are the ones to carry away. The naive-truncation curve falls off a cliff below 128 dimensions, because √(m/d) is brutal down there. The SVD curve tracks the reference well at large m — where there is plenty of room and the top components really do carry the task — and separates below about 64, where the difference between “high variance” and “discriminative” stops being negligible. And the MRL curve hugs the reference across the whole range, which is the paper's central empirical claim.

One more contender: train a small autoencoder on top

A neural compressor. Take the frozen 768-dimensional embeddings, train a small MLP down to 64 dimensions and another back up, minimize reconstruction error. Nonlinear, learned, and it genuinely beats PCA at reconstruction. Four reasons it does not solve this problem:

It is a second model to ship. Encoder weights, decoder weights, versioned and kept in sync with the base model and with the index. When the base model updates, the autoencoder must be refit and every vector re-encoded — the same backfill treadmill as three separate models.

Its output is not a prefix. There is no chain. A 64-dimensional autoencoder code and a 256-dimensional one from a different autoencoder are unrelated, so no cascade is possible — back to Chapter 0's incompatible geometries.

It optimizes reconstruction, which is PCA's mistake with more parameters. A nonlinear reconstruction objective is still a reconstruction objective. It will happily spend capacity representing the brightness nuisance from Worked Example 2, because brightness is a large part of what you have to reconstruct.

It costs a forward pass per query and per document. Small, but nonzero, and it sits directly on the latency path that motivated all of this.

The one way truncation can genuinely help

This chapter has been an argument that truncation destroys information, so intellectual honesty requires the counter-current: there is a regime where fewer dimensions makes retrieval better, and it is worth understanding because it explains occasional surprising measurements.

In high dimensions, distances concentrate. For points drawn from a distribution in R^d, the spread of pairwise distances shrinks relative to their mean as d grows — the classic statement is that

which means the very notion of a “nearest” neighbour loses contrast: everything is roughly equally far from everything. Real embeddings are not drawn from a full-dimensional distribution — they live near a low-dimensional manifold, which is why nearest-neighbour search works at all — but a residue of the effect survives. Coordinates that carry no task signal contribute noise to every distance, flattening the gap between a genuine match and a near-miss.

So cutting a representation back toward its intrinsic dimension does two opposing things at once: it removes signal (the √(m/d) law) and it removes nuisance that was diluting contrast. For an ordinary embedding the first dominates overwhelmingly. For a matryoshka embedding, where the removed coordinates were specifically the least broadly useful ones, the two can be closer — and the practical consequence is that you will sometimes measure a truncated index retrieving slightly better than the full one on some query sets.

The four baselines, and what each one is actually testing

Notice which two rows have zero query cost and no shipped artifact: naive truncation and MRL. That column, not the accuracy column, is why MRL got adopted so fast. A technique whose deployment story is “use fewer bytes of the vector you already have” needs no infrastructure at all.

How to find out whether a model you did not train front-loads

You are handed an embedding API. The docs say nothing about matryoshka. Twenty minutes and a thousand documents will tell you whether truncation is safe, and the procedure is worth having in your fingers.

python — is this model truncatable? three testsimport numpy as np
E = embed(docs)                       # (10000, d) float32, unit norm
d = E.shape[1]

# TEST 1 - the magnitude profile. 30 seconds, no ground truth needed.
mag = (E ** 2).mean(0)                # (d,) mean squared magnitude per coordinate
print("front/back ratio:", mag[:d//8].mean() / mag[d//2:].mean())
# ~1.0  -> democratic, do not truncate
# >5    -> something front-loaded this deliberately

# TEST 2 - the operational one. Needs queries; ground truth is the FULL model.
def l2(X): return X / np.linalg.norm(X, axis=-1, keepdims=True)
truth = np.argsort(-(Q @ E.T), axis=1)[:, :10]      # (n_q, 10)
for m in [64, 128, 256, 512, d]:
    approx = np.argsort(-(l2(Q[:, :m]) @ l2(E[:, :m]).T), axis=1)[:, :10]
    rec = np.mean([len(set(a) & set(t)) / 10 for a, t in zip(approx, truth)])
    print(m, "recall@10 vs full:", round(rec, 3))

# TEST 3 - the control. Is truncation beating a post-hoc rotation?
U, S, Vt = np.linalg.svd(E - E.mean(0), full_matrices=False)
for m in [64, 256]:
    Esvd = l2((E - E.mean(0)) @ Vt[:m].T)
    # ... same recall computation. If SVD wins, the model is NOT matryoshka
    # and you should ship the projection matrix rather than a slice.

Test 1 is a screen, not a proof — magnitude is not the same as usefulness, and a model could in principle front-load information without front-loading magnitude. Test 2 is the real answer, and it is the number to put in a decision document. Test 3 is the control that stops you from congratulating yourself: if a plain SVD at 64 dimensions beats the prefix at 64 dimensions, the model was never matryoshka and you have learned something valuable in twenty minutes.

Chapter 4: Training Mechanics

The objective is one line and the code is twenty. What remains is the set of decisions you have to make to actually run it, and the arithmetic that justifies the paper's claim that the extra cost is negligible. “Negligible” is a word people use when they have not measured. Let us measure.

What the nine granularities actually cost

Take the paper's headline setting: ResNet50 backbone, ImageNet-1K, d = 2048, L = 1000, M = {8, 16, …, 2048}.

Parameters. A standard ResNet50 has 25,557,032 parameters, of which the final fully-connected layer contributes 2048 × 1000 + 1000 = 2,049,000 — about 8% of the model. Replace that single head with nine bias-free heads:

Under MRL-E the single sliced head is 1000 × 2048 = 2,048,000, so the total is 25,556,032 — within a thousand parameters of the baseline. Zero, for practical purposes.

Compute. This is where the real answer lives, because parameters sit in memory while FLOPs cost time. A ResNet50 forward pass on a 224×224 image is about 4.1 GFLOP. The heads:

Two-tenths of one percent. And that is the total for all nine heads — the single baseline head was already 0.10% of the same figure, so the marginal cost is 0.10%. The backward pass roughly triples the forward cost for both, so the ratio holds.

Activation memory. The one thing that genuinely grows is the stored logits, since autograd must keep each granularity's output for the backward pass. At batch size 256:

ResNet50 activations at batch 256 run to several gigabytes. Nine megabytes is a rounding error, and it is the only line item that scales with |M|. Even a hundred granularities would add only 102 MB — the reason to keep |M| small is the head parameters and the loop overhead, not activation memory.

The same arithmetic for a text embedder, where it is even cheaper. BERT-base has about 110 million parameters and a forward pass on a 128-token sequence costs roughly 22 GFLOP. Suppose you train it contrastively with M = {12, 24, 48, 96, 192, 384, 768}. There are no heads at all — the similarity is a dot product. The only extra work is computing six additional similarity matrices:

Under two thousandths of one percent. In the contrastive setting matryoshka training is, to measurement precision, free — which is exactly why every production text embedder now does it.

What does grow: the similarity matrices in memory. Each is (N, N) float32. At N = 256 that is 256 KB each, so seven of them is 1.8 MB. At the batch sizes contrastive training actually wants — 8,192 or more — each matrix is 268 MB and seven is 1.9 GB, which is a real number on a real GPU. The mitigation is the same one large-batch contrastive training already uses: compute the loss in chunks, or free each granularity's matrix before building the next.

The relative importance weights c_m

The objective has a free scalar per granularity, and the paper sets them all to 1. Is that principled or lazy? Three arguments say it is principled, and one caveat says know what you are getting.

Argument one: the terms are already commensurable. All nine terms are softmax cross-entropies over the same 1000 classes. At initialization each sits near ln(1000) = 6.908, and each is bounded below by 0. There is no unit mismatch to correct — unlike, say, mixing a cross-entropy with a mean-squared reconstruction term, where the relative scale is arbitrary and must be tuned.

Argument two: there is already an implicit weighting, and it points the right way. This is the subtle one. Uniform c_m does not mean uniform emphasis on the coordinates, because of the multiplicity staircase from Chapter 2. Coordinate 1 accumulates nine gradient contributions; coordinate 2000 accumulates one. So even with equal term weights, the early coordinates are under 9× the pressure. Uniform c_m is not neutral — it is already strongly biased toward making small prefixes work, which is exactly the bias you want.

Argument three: it is one fewer hyperparameter, and hyperparameters are the enemy of adoption. A method whose recipe is “change your loss to a sum, keep everything else” will be tried by a hundred times more people than one requiring a nine-dimensional weight sweep.

The caveat. If you know at training time which granularity you will actually serve at, weighting it up is free and sensible. Concretely, in Chapter 5's toy the total gradient pressure on coordinate 1 is −0.901 with uniform weights; raising c₂ from 1 to 2 takes it to −1.222, a 36% increase in how fast that coordinate moves. Whether that improves the served metric is an empirical question, but the knob is real and it is one line of code.

Push the tilt all the way toward small m and watch the cumulative curve become almost vertical at the left: nearly all the ordering pressure lands in the first sixteen coordinates and the tail is essentially unsupervised. Push it the other way and the staircase flattens toward the ordinary, order-blind objective. Uniform sits in a sensible middle, which is a reasonable defence of the paper's choice.

Five decisions that are not in the equation

1. Where in the network do you apply it? To the final representation — the vector you would actually ship. For a CNN that is the post-global-average-pool feature; for a Transformer it is the pooled sequence output; for a model with a projection head it is whatever the serving system stores. Applying MRL to an internal projection you throw away at inference accomplishes nothing, and this is the single most common way to waste a training run.

2. What kind of head? Linear. This is not a stylistic preference. The head is a stand-in for whatever will read the embedding downstream, and downstream readers are linear probes and cosine similarities. If you use a two-layer MLP head, the prefix becomes good only when read through that MLP — which nobody at serving time has. The MLP absorbs the ordering into its own weights and the shipped vector learns nothing.

3. Normalization. If your pipeline L2-normalizes embeddings — and every retrieval pipeline does — normalize each prefix separately, inside the loss:

python — the contrastive variant, correctlydef mrl_contrastive(z_a, z_b, dims, temp=0.05, c=None):
    """z_a, z_b: (N, d) paired encodings. Returns a scalar."""
    if c is None: c = [1.0] * len(dims)
    total = 0.0
    labels = torch.arange(z_a.shape[0], device=z_a.device)
    for cm, m in zip(c, dims):
        a = F.normalize(z_a[:, :m], dim=-1)     # slice FIRST
        b = F.normalize(z_b[:, :m], dim=-1)     # then normalize
        sim = a @ b.T / temp                    # (N, N)
        total = total + cm * 0.5 * (F.cross_entropy(sim, labels) +
                                       F.cross_entropy(sim.T, labels))
    return total

# WRONG, and it will look like it works for a while:
#   a = F.normalize(z_a, dim=-1)[:, :m]
# now |a| varies per example and long-tailed docs win every comparison

4. The nesting set. Exponential spacing down to the smallest size any consumer will plausibly ask for. Include sizes your infrastructure actually wants — if your ANN library is happiest at 128 or your cache line holds 64 float32s, put those numbers in M. Every extra granularity costs one head, which Chapter 4's arithmetic says is nothing.

5. Everything else stays. Same epochs, same learning-rate schedule, same augmentation, same optimizer, same weight decay. That is a genuine claim and a large part of the appeal: MRL is not a training recipe, it is a modification to one line of the loss.

Numerics, distribution, and the things that bite at scale

Four practical concerns that do not appear in the equation but appear in every real run.

Mixed precision. The nine cross-entropies are summed before the backward pass, so the loss scalar is roughly |M| times larger than a baseline loss — about 62 at initialization instead of 6.9. With float16 gradient scaling that is harmless (the scaler adapts), but if you have a hand-tuned static loss scale, divide it by |M| or you will start clipping. The safer habit is to log the mean of the per-granularity losses for monitoring while backpropagating the sum, so your dashboards remain comparable with the baseline run.

Gradient norm and clipping. The gradient at z is up to |M| times larger on the early coordinates, so the global gradient norm rises. If your recipe clips at a fixed global norm, that clip now binds more often and it binds asymmetrically — it scales down the whole gradient, which disproportionately throttles the tail coordinates that were already receiving the least. Either raise the clip threshold or clip per parameter group.

Weight decay on the heads. Under untied MRL you now have nine head matrices instead of one, so the total weight-decay penalty on the head grows about twofold. This is usually invisible against a 25-million-parameter backbone, but in a small model where the head is a large fraction of the parameters it will silently shrink your logit scale. Exclude the heads from decay, or scale their decay by 1/|M|.

Data-parallel training. Nothing special: the loss is a sum of per-example terms, so gradients all-reduce exactly as before. The one wrinkle is contrastive training, where negatives come from the batch — if you gather across ranks to build a large effective batch, you must gather once and reuse the gathered tensors for all |M| similarity matrices, not gather |M| times. Gathering per granularity multiplies your communication volume by |M| and turns a 0.002% compute overhead into a real slowdown.

Retrofitting an existing model

You already have a trained 768-dimensional encoder with an index built from it, and you would rather not retrain from scratch. Two options, in increasing order of cost and quality.

The first row is worth dwelling on, because it is the intuitive thing to try and it does not work. Attaching a nine-headed classifier to a frozen encoder trains nine linear probes. Linear probes on a frozen representation are precisely the post-hoc rotation family from Chapter 3, subject to exactly the same ceiling. The encoder has to be allowed to move, because moving the encoder is the mechanism.

The second row is what most teams actually do, and there is a growing literature on doing it well — adaptor-style approaches that learn a small transformation plus a short fine-tune, trading a little quality for a lot of compute.

The training loop, with the logging that saves you a week

python — the whole loop, instrumentedDIMS = [8, 16, 32, 64, 128, 256, 512, 1024, 2048]
head  = MatryoshkaHead(DIMS, n_classes=1000).to(dev)
opt   = torch.optim.SGD(list(backbone.parameters()) + list(head.parameters()),
                        lr=0.1, momentum=0.9, weight_decay=1e-4)   # UNCHANGED

for x, y in loader:
    z          = backbone(x)                      # (N, 2048)  ONE pass
    logits     = head(z)                          # list of 9 x (N, 1000)
    per_m      = [F.cross_entropy(u, y) for u in logits]
    loss       = sum(per_m)                       # c_m = 1
    opt.zero_grad(); loss.backward(); opt.step()

    # --- the logging that matters ---
    for m, l in zip(DIMS, per_m):
        log(f"loss/m={m}", l.item())          # NINE scalars, never one
    # the front-loading smoke test, cheap enough to run every 500 steps
    with torch.no_grad():
        mag = (z ** 2).mean(0)                  # (2048,) mean squared magnitude
        log("mag/first8",   mag[:8].mean().item())
        log("mag/last1024", mag[1024:].mean().item())
        # the RATIO is the signal. Ordinary training holds it near 1.0;
        # a working MRL run drives it into the tens within a few epochs.

The last three lines are the highest-value diagnostic in the whole recipe, because they distinguish “my loss is going down” from “my representation is reorganizing.” Those are different claims and only the second one is what you are paying for. A run whose nine losses all fall beautifully while the magnitude ratio sits at 1.0 has a bug — almost always a slice applied to the wrong axis, or MRL applied to a projection that is discarded before z is stored.

Three recipe cards

That last cell is a real failure mode. In a dual encoder the two towers must be nested at the same granularities, because a 64-dimensional image prefix will be compared against a 64-dimensional text prefix. Nest only the image side and the text side has no reason to put anything useful in its first 64 coordinates, so the comparison at m = 64 is a good vector against a random one.

When full-width quality actually drops: a decision procedure

The most alarming outcome is a run where the nine losses all fall but your production metric at m = d is half a point below the old model. The nesting constraint is binding. Work through this in order, because the causes are ranked by how often they turn out to be the real one.

Step 3 deserves emphasis because it is counter-intuitive. More constraints does not mean a worse optimum — Chapter 1's extendability argument says a monotone solution exists — but it can mean a harder optimization landscape, and “harder to reach” looks identical to “worse” if you stop at a fixed step count.

A natural extension: distilling across granularities

One modification appears repeatedly in follow-up work and is worth knowing about, because it is the first thing you would invent yourself. Right now the nine granularities are trained independently against the ground-truth label; they never talk to each other. But granularity 2048 is a better model than granularity 8 — so let the small one learn from the big one:

The second term is ordinary knowledge distillation, with the full-width granularity as the teacher and each prefix as a student, and the teacher is free because it is already being computed. The argument for it: a hard label says only “cat”, while the teacher's distribution says “cat, but somewhat lynx, definitely not trombone” — and that similarity structure is exactly the coarse information a small prefix has room to represent. The argument against: you must stop the gradient through the teacher, or the big granularity will learn to make itself easy to imitate rather than accurate.

Diagnostics: what to log, and what each signal means

Log every granularity's loss separately. One scalar is nine scalars in disguise and the disguise hides every interesting failure.

What this does to your evaluation cost

One overhead nobody warns you about, and it is entirely self-inflicted. You now have nine granularities to evaluate instead of one. Done naively — re-embed the eval set once per granularity — your evaluation cost goes up ninefold, and evaluation is often already the slowest part of a research loop.

Done correctly it costs about twice one evaluation, and the reason is the same geometric-sum fact that made the heads cheap.

python — evaluate every granularity for the price of twoE = encode(eval_set)             # (N, 2048)  ONCE. This is the expensive part.
Q = encode(queries)              # (n_q, 2048) ONCE.

results = {}
for m in DIMS:
    e = l2(E[:, :m])              # a view + one normalise. No encoder.
    q = l2(Q[:, :m])
    results[m] = ndcg_at_10(q @ e.T, qrels)

# The mistake that costs 9x:
#   for m in DIMS: E = encode(eval_set, dimensions=m)   # re-runs the model
# The prefix is a slice. There is nothing to re-encode. Ever.

The same reasoning applies to your production monitoring: instrument one embedding pass and derive every granularity's metric from it. If a dashboard is re-encoding per dimension, someone has misunderstood what a prefix is, and the fix is a two-line change that reclaims 89% of the job's runtime.

Reading the training curves

Nine loss curves on one chart is a lot of ink. Here is what the shapes mean, so you can glance rather than squint.

The last row is the one to watch for, because it is the whole bet. If L₂₀₄₈ overlays your baseline curve step for step, you are getting eight extra operating points for free, and you can stop worrying about whether the constraint is costing you something.

What you actually ship

A final piece of concreteness: the artifacts that leave the training job, because the short list is the argument for the whole method.

Chapter 5: The 8-Dimensional Toy, By Hand

Every equation in Chapter 2 now gets computed with a pencil. Eight dimensions, three classes, three granularities. By the end of this chapter you will have produced a single scalar loss the way the optimizer does, and computed the gradient arriving at two different coordinates so you can see the multiplicity effect as a number rather than a claim.

The setup

Three classes: cat, dog, bird. Representation dimension d = 8. Nesting set M = {2, 4, 8}. We use MRL-E, so there is one head matrix W ∈ R^3×8 and granularity m uses its first m columns.

One training example, an image of a cat, produces the representation:

Front-loaded, as an MRL-trained encoder should be. Coordinate 1 has magnitude 0.90; coordinate 8 has 0.05, eighteen times smaller.

The head, written as a table because you are about to read its columns as well as its rows:

True label y = cat. Relative importance c_m = 1 for all three granularities, as in the paper.

Granularity 1: m = 2

The prefix is z_1:2 = [0.90, 0.40], and the head slice is the first two columns. Three dot products:

Exponentiate:

Normalize:

The loss is the negative log of the true class's probability:

Sanity check the magnitude before moving on. Guessing uniformly over three classes gives −ln(1/3) = ln 3 = 1.0986. We are at 0.5458, comfortably below the guessing baseline, so two coordinates already carry real information. Good — that is the whole promise of the small granularity, verified on the first computation.

Granularity 2: m = 4

Two more coordinates, z₃ = 0.30 and z₄ = 0.20, and two more columns. Because the head is sliced rather than retrained, the first two terms are identical to before — we can just add:

That incremental structure is worth pausing on: under MRL-E, going from granularity m to the next one adds terms to the logits and never rewrites them. That is the mechanical reason the “can always be extended” argument from Chapter 1 works — and it is why a well-trained MRL-E model has quality monotone in m.

Granularity 3: m = 8

The last four coordinates are [0.15, 0.10, 0.05, 0.05] — small, by construction. Add their contributions:

Notice how little the last four coordinates bought: 0.5204 → 0.5060, a gain of 0.0144, versus the 0.0254 that coordinates 3 and 4 bought and the enormous jump from the 1.0986 guessing baseline that coordinates 1 and 2 bought on their own. That decreasing-returns shape is what a front-loaded representation looks like from the inside.

An aside on how you would actually compute this

We exponentiated raw logits above because the numbers were small and it kept the arithmetic legible. Real code does not do that, and the reason is worth one paragraph because it is a bug you will otherwise write.

Logits in a trained network routinely reach 20 or 30. Then e³⁰ is about 10¹³, and with a few more it overflows float32 entirely, producing inf, then inf/inf = NaN, then a silently poisoned training run. The fix is to subtract the maximum first, which changes nothing because softmax is shift-invariant:

Verify on our own m = 8 numbers. The logits were (1.290, 0.515, −0.330) with maximum 1.290. Shift them to (0, −0.775, −1.620):

And the loss falls out even more directly, without ever forming a probability:

Matching the 0.5060 from the long way round. The largest exponent is always exactly 1, so nothing can overflow, and F.cross_entropy in PyTorch does precisely this — which is why you should hand it logits rather than probabilities. Passing it a softmax output is the second most common bug in this territory, after normalizing before slicing.

The total objective

That single number is what loss.backward() is called on. One scalar, three classifiers of three different capacities, one shared representation.

And immediately: what would the weighted variants give? With c = (0.5, 1, 1) — de-emphasizing the tiny granularity:

With c = (2, 1, 1) — leaning hard on the smallest, because that is what you plan to serve:

The absolute value is meaningless — scaling a loss scales nothing that matters. What matters is the ratio of the terms, because that sets which granularity's residual dominates the gradient. Which is where we go next.

The gradient, and the multiplicity effect as a number

Recall the cross-entropy gradient with respect to logits: ∂L/∂u_k = p_k − [k = y]. With y = cat (index 1), the three residual vectors are:

Every residual has the same structure — push the true class's logit up, push the others down — and its magnitude shrinks as the classifier gets more capacity and gets more of the answer right. Now chain into a coordinate. For coordinate j:

Coordinate 1, whose column of W is (1.0, 0.6, −0.5). It appears in all three granularities, so all three residuals reach it:

Coordinate 5, whose column is (0.1, −0.2, 0.3). It appears only in the m = 8 prefix, so only one residual reaches it:

Now the comparison that makes the mechanism visible. Coordinate 1 receives −0.9011. If MRL had not been used — if there were only the ordinary m = 8 loss — coordinate 1 would receive only its m = 8 contribution, −0.2901. So:

Three granularities, roughly 3× the update speed on the first coordinates. Under SGD with learning rate η, coordinate 1 moves 3.11 times as far per step as it would under the ordinary objective, while coordinate 5 moves exactly as far as it always would. Repeat that asymmetry for tens of thousands of steps and the representation reorganizes: whatever is most broadly useful migrates to where the pressure is highest.

One SGD step, so the reorganization is not hypothetical

Treat z as if it were directly updatable — it stands in for whatever the backbone would do — and take one step with learning rate η = 0.1. We have the two gradients we need plus coordinate 2's, computed the same way from column (0.5, −0.4, 0.3):

Coordinates 1 and 2 each moved by about 0.09; coordinate 5 moved by 0.006 — fifteen times less. Now recompute the small granularity with the updated prefix:

Down from 0.5458 in a single step — a 9.3% improvement, driven overwhelmingly by the two coordinates that were under triple pressure. Run that ten thousand times and the front-loading is not a tendency, it is the shape of the representation.

The head is matryoshka too

One more gradient, because it reveals that the ordering pressure applies to both sides of the product. For MRL-E the head is a single W, and column j of W is used by every granularity m ≥ j — the same index set. So:

Evaluate for the cat row at coordinates 1 and 5:

An eighteen-fold difference — and it decomposes cleanly into two independent causes. Three residuals instead of one gives a factor of about 3.08 (1.2234/0.3971), and the front-loaded representation itself gives a factor of 6 (0.90/0.15). The multiplicity and the magnitude compound, so the head's early columns learn far faster than its late ones. The nesting propagates to every tensor the objective touches.

When the small granularity is simply wrong

The example so far was easy: two coordinates already got the answer right. Real training is full of examples where they do not, and those are where the interesting gradient lives. Take a fine-grained cat whose coarse features look like a dog:

At m = 2 the logits are 1.0(0.30) + 0.5(−0.50) = 0.05 for cat, 0.6(0.30) − 0.4(−0.50) = 0.38 for dog, and −0.5(0.30) + 0.3(−0.50) = −0.30 for bird.

Two things at once: the prediction is dog, which is wrong, and the loss 1.1300 is above the guessing baseline ln 3 = 1.0986. The two-dimensional classifier is worse than a coin flip on this example. At m = 4 it is still wrong (L₄ = 0.9103, predicting dog at 0.41 against cat at 0.35). Only at m = 8 does the tail pull it back:

The counterfactual: what a democratic representation costs

Take the same head, the same class, and a representation with the same total magnitude spread evenly. The original z has norm

Spread that evenly over eight coordinates: each entry is 1.0665/√8 = 0.377.

At m = 2 the logits become

Against the front-loaded representation's 0.5458. At full width the democratic vector reaches L₈^dem = 0.6687 (the row sums of W are 2.4, 0.8, and 0.0, so the logits are 0.9048, 0.3016, 0.0). Line the two up:

The democratic vector degrades 2.3 times as much when you cut it to a quarter of its length. The effect is real and it points the right way — but be honest about its size: this is a mild effect, and it is mild for a reason worth understanding.

Changing M, and watching the pressure move

One last variation, because M is the only genuinely new hyperparameter and its effect should be something you can compute rather than something you take on faith. Keep everything identical and change the nesting set from {2, 4, 8} to {4, 8}. What happens to coordinate 1?

With M = {4, 8} the multiplicity of coordinates 1 through 4 is 2, and of coordinates 5 through 8 is 1. Coordinate 1's total gradient loses its m = 2 contribution:

And note what else changes: with M = {4, 8}, coordinates 1 through 4 all have multiplicity 2, so they are indistinguishable to the objective. The model has no reason to order coordinate 1 above coordinate 3. Truncating to 2 dimensions would land in unsupported territory.

Seeing the geometry directly

The clearest way to feel why prefixes work is to look at what taking a prefix does geometrically. Taking the first m coordinates of a vector is an orthogonal projection onto the coordinate subspace spanned by e₁, …, e_m. It casts a shadow. Whether the classes remain distinguishable in the shadow depends entirely on whether the direction separating them was aligned with the wall.

With the ordinary encoder the three clusters are perfectly separated in 3D — the representation is excellent — and their shadows on the e₁–e₂ plane overlap into mush, because the separating direction points mostly along e₃. With the matryoshka encoder the 3D separation is essentially the same and the shadows are clean. Same information, same quality at full width, different orientation, completely different behaviour under projection. That single picture is the paper.

Do one yourself

Two short exercises using the same W and the same conventions. Work them before reading the answers; the arithmetic is small and doing it is the difference between recognizing the mechanism and owning it.

Every number in this chapter, checked

Hand arithmetic that nobody verifies is hand arithmetic that is wrong. Here is the audit, so you can confirm the chapter rather than trusting it.

The fourth row is the most useful invariant to remember. A softmax cross-entropy residual p − onehot always sums to zero, because p sums to 1 and the one-hot sums to 1. If any residual you compute does not, you have a bug — and it is a one-line check you can assert in code.

The sixth row is the empirical confirmation of Chapter 1's extendability argument on real numbers: quality improved at every step of m, and the mechanism (each granularity adds terms without rewriting the previous ones) explains why it had to.

Chapter 6: Adaptive Retrieval and the Funnel

Everything so far has been about making a small embedding good. This chapter is about the thing that small embedding unlocks, and it is a bigger deal than the memory saving: the cheap representation can filter work for the expensive one, because they live in the same space. Chapter 0 named that as the property three independent models cannot have. Here we cash it in.

The retrieval problem, priced from zero

You have N documents, each embedded as a d-dimensional unit vector, and a query embedded the same way. Find the k documents with the largest cosine similarity to the query.

Because everything is unit-norm, cosine similarity is the dot product, so one document costs d multiply-accumulate operations. Exhaustive search costs:

Use the paper's setting to make it concrete: the ImageNet-1K training set as a database, N = 1,281,167 images, d = 2048.

Two stages, and the cost model that falls out

Now the idea. Do not score every document at full width. Score every document at a shortlist dimension D_s, keep the best K of them, and re-score only those K at the full rerank dimension D_r. Return the top k of the re-scored set.

Read the two terms. The first is huge in N and tiny in dimension; the second is tiny in count and huge in dimension. The whole design is arranging for neither to dominate.

Worked example. D_s = 16, K = 200, D_r = 2048:

Verify: 20,908,272 × 125 = 2,613,534,000; the remainder 10,296,016 divided by 20,908,272 is 0.49. So 125.5. That is where the paper's “roughly 128× theoretical” figure comes from — and notice it is almost exactly D_r/D_s = 2048/16 = 128, because stage 2 contributes only 2% of the total. When the rerank stage is negligible, the theoretical speed-up is just the dimension ratio.

How big must K be? The recall question

The cost model above assumed the shortlist contains the right answers. It might not. Stage 2 can only rank what stage 1 handed it, so the pipeline's quality ceiling is:

If a document that would have been in the true top-10 fails to make the 16-dimensional top-200, it is gone forever and no amount of reranking recovers it. So K is the knob trading cost against recall, and you want the smallest K with recall ≈ 1.

Reason about it through Chapter 3's noise model. The low-dimensional score is the full-dimensional score plus an error — call its standard deviation δ. A document truly at rank 10 falls out of the top K only if the noise pushes it below K − 10 documents that were genuinely worse. So the question is whether the true score gap between rank 10 and rank K exceeds a few multiples of δ.

Worked example. Suppose the 16-dimensional score error has δ = 0.02 in cosine units, and on your corpus the true cosine at rank 10 is 0.74 while at rank 200 it is 0.66. The gap is 0.08 = 4δ. A rank-10 document needs a 4-sigma downward error to be displaced past rank 200, which happens for roughly 1 document in 30,000. Recall@200 is effectively 1. Now try K = 20, where the gap might be only 0.01 = 0.5δ: displacements are routine and recall collapses.

The shape of that table is the practical guidance: K in the range of 10× to 20× the number of results you actually return costs about 2% of the pipeline and buys essentially all the recall. Beyond that you are spending real money to protect against errors that are not happening.

Where the optimum sits — solving for the best D_s

The cost model is simple enough to differentiate, which turns “pick a shortlist dimension” from a sweep into an equation. Treat K as depending on D_s, since a coarser shortlist needs more candidates to keep recall — a reasonable model is K = K₀ · D₀/D_s, so halving the shortlist dimension doubles the candidates you must keep. Then:

Put the paper's numbers in. Anchor the recall model at K₀ = 200 when D₀ = 16, so K₀D₀ = 3,200:

The optimum is below 4, which says the cost model alone would push you all the way to the smallest dimension available. So why does the paper use 16? Because the cost model is not the binding constraint — recall is. Below about 16 dimensions the shortlist stops being a reliable filter no matter how many candidates you keep, because the score noise swamps the true gaps. The optimum of the cost function sits in a region the quality constraint has already ruled out.

Funnel retrieval — the cascade, and its lovely accident

Two stages is the simple version. The paper's funnel retrieval generalizes it: instead of one jump from 16 dimensions to 2048, walk up the chain, halving the candidate set at each rung and doubling the dimension.

And that invariance is only possible because the granularities are nested. The rung at 128 dimensions is refining a shortlist produced at 64 dimensions using strictly more of the same vector. Under three independently trained models there is no such chain: each stage would be re-scoring with a different, incomparable notion of similarity, and the shortlist from stage 1 would carry no guarantee about stage 2 at all.

Drag D_s down to 8 and then to 4 and watch the total stop improving: once stage 1 is cheap enough, stage 2 and the fixed overheads dominate and further shrinking is wasted. Drag K up and watch stage 2 take over. The sweet spot is a broad basin, not a knife edge, which is why this technique is robust in practice.

Setting the three parameters, in order

Adaptive retrieval has exactly three knobs — D_s, K, and D_r — and they should be set in a specific order, because two of them are determined by constraints and only one is a genuine trade-off.

People usually do this backwards — they tune D_s for quality and treat K as fixed — and then spend a week discovering that D_s barely matters over a broad range while K matters enormously below a threshold. The reason is in the cost model: quality is governed by whether the true answers survive the shortlist, and that is far more sensitive to how many you keep than to how precisely you scored them.

Why 128× in theory becomes 14× on a clock

The paper reports about 128× in FLOPs and about 14× in wall-clock time. That is a factor-of-nine discrepancy, and understanding it is more useful than either number, because the same reasoning applies to every optimization you will ever ship.

Three costs do not shrink when you shrink the dimension:

1. Top-K selection. After scoring, you must find the best K of N scores. That is N comparisons against the current threshold no matter what d was — 1.28 million branches, independent of dimension. At roughly 1.5 ns each that is about 2 ms, and it is a hard floor.

2. Approximate-nearest-neighbour traversal. Nobody runs an exhaustive scan in production; they run an HNSW graph or an IVF index. Graph traversal is dominated by pointer chasing — each hop is a cache miss costing about 100 ns of pure latency — and a cache miss costs the same whether the vector at the other end is 16 or 2048 numbers wide. Reducing d shrinks the distance computation and the bytes read, but not the latency of getting there.

3. The baseline is already optimized. The 128× is against a naive full scan. Against a well-tuned ANN index — the thing you would actually deploy — much of the easy win is already taken.

Put a number on it with Amdahl's law. Let f be the fraction of baseline time that scales with dimension, and (1 − f) the fraction that does not. Shrinking the dimension by 128× gives:

Solve for the observed 14×:

So about 94% of the baseline's time was dimension-proportional and 6% was not. Check it: 1/(0.064 + 0.936/128) = 1/(0.064 + 0.00731) = 1/0.07131 = 14.0. ✓

What this does to an ANN index

Chapter 6 has so far assumed brute-force scanning, which is the honest way to derive a cost model but not what you deploy. Fold in an approximate index and three separate things improve, only one of which is the one people expect.

The second and fourth rows are the ones that explain the Amdahl fraction. An HNSW graph over 1.28 million nodes costs the same 164 MB regardless of dimension, and each hop is a cache-missing pointer chase whose ~100 ns of latency does not shrink. Reducing the dimension by 128× turns an 8 KB read per hop into a 64-byte read, but does not remove the miss.

The two-stage pipeline, in code

python — adaptive retrieval, end to endimport numpy as np

# Built once, offline, from ONE matryoshka index.
#   Z: (N, 2048) float32 full archive, memory-mapped or on SSD
#   S: (N,   16) float32 contiguous shortlist copy - 82 MB, stays in RAM
Z = np.load("index_2048.npy", mmap_mode="r")
S = np.load("index_16.npy")              # np.ascontiguousarray(Z[:, :16]) renormalised

def search(q, k=10, K=200):
    # stage 1 - sweep every document at 16 dims
    q16 = q[:16]; q16 = q16 / np.linalg.norm(q16)   # slice THEN normalise
    scores = S @ q16                              # (N,)  20.5 MMAC, 82 MB read
    cand   = np.argpartition(-scores, K)[:K]     # (K,)  O(N), dimension-free

    # stage 2 - rerank only those K at full width
    full   = np.asarray(Z[cand])                  # (K, 2048) - 1.6 MB of SSD reads
    exact  = full @ q                             # (K,)  0.41 MMAC
    order  = np.argsort(-exact)[:k]
    return cand[order], exact[order]

# Measuring the ONLY thing that matters for correctness:
def recall_at_k(qs, k=10, K=200):
    hits = 0
    for q in qs:
        approx = set(search(q, k, K)[0].tolist())
        truth  = set(np.argsort(-(Z @ q))[:k].tolist())   # exhaustive ground truth
        hits  += len(approx & truth)
    return hits / (len(qs) * k)
# Sweep K until this plateaus, then stop. That is the whole tuning procedure.

Two lines carry the argument. q[:16] then normalize — the same rule as everywhere else. And np.argpartition, which is the dimension-independent O(N) top-K selection that sets the Amdahl floor: it does not get faster when the vectors get shorter, because it never looks at a vector.

The 20 ms budget, revisited

Return to Chapter 0's autocomplete path and re-itemize it with a matryoshka index. The point is not the total; it is which line moved.

The encoder line did not move, and it never will — MRL shrinks the output, not the network. If your budget is dominated by encoding rather than searching, this entire chapter buys you nothing and you want Chapter 9's depth-nesting instead. Knowing which line of the budget you are actually fighting is the difference between a technique that works and one that merely sounds good.

The gotcha nobody warns you about: scores are not comparable across m

Here is a bug that ships to production regularly. You have a relevance threshold: “show the result if cosine > 0.75.” You switch a service from 768 dimensions to 256 to save memory. The threshold no longer means what it meant.

Why: cosine similarity at m dimensions is computed over a different number of terms, and the baseline similarity between two unrelated vectors depends on m. For two random unit vectors in R^m, the expected cosine is 0 but its standard deviation is 1/√m:

At 64 dimensions two completely unrelated documents routinely score 0.25 — two standard deviations of pure chance — whereas at 768 that same score is nearly seven standard deviations and would be a strong match. The noise floor rises as the dimension falls, so every threshold, every cutoff, and every “confidence” heuristic must be recalibrated per granularity.

Sharding, replication, and the tail latency nobody budgets for

One consequence of index size is so large and so frequently missed that it deserves its own arithmetic: how many machines you need, and what that does to your p99.

Take the 12-million-document, 768-dimensional index at 36.9 GB. Assume 64 GB machines with about 32 GB usable for the index after the OS, the process, and headroom. You need two shards, and with a replication factor of three for availability that is six machines.

Now the part people do not budget for. A sharded query fans out and waits for the slowest shard, so the overall latency is the maximum over shards. If shard latencies are independent and each shard is under t with probability p, then:

Verify the third row: 0.99^1/8 = exp(ln(0.99)/8) = exp(−0.010050/8) = exp(−0.0012563) = 0.998744, which is the 99.87th percentile. Since latency distributions have long tails, the p99.87 of a shard is typically two to three times its p99 — so fanning out to eight shards can double or triple your tail latency even though every shard is doing less work.

This is the largest practical benefit in the whole lesson, and it never appears as a number in the paper because it is not a property of the representation. It is a property of thresholds in your infrastructure: 32 GB per machine, 64 GB per machine, whatever your fleet has. Crossing back under one of those thresholds is discontinuous in a way that a smooth accuracy curve cannot express.

Keeping two indexes in sync — and why that is easier than it sounds

Adaptive retrieval means you now maintain two artifacts: the full archive and the derived shortlist index. Two artifacts is normally two chances to drift, so the operational question is fair: what happens when documents arrive, change, and are deleted?

The stale-update row is the subtle one, because it fails in the direction that looks like a quality problem rather than a bug. But here is the property that makes all of this manageable, and it is the operational payoff of “a prefix is a slice”:

Compare the alternatives. If your small index came from a PCA projection, rebuilding requires the projection matrix from whenever it was fitted — and if you refit it on the current corpus you get a different projection, so every stored vector must be recomputed and every threshold recalibrated. If it came from a separately trained small model, rebuilding means re-encoding twelve million documents on GPUs. The matryoshka version is the only one where the derived artifact is genuinely a function of the primary artifact and nothing else.

Chapter 7: Results, Honestly — and Production

The method is one line and the mechanism is a counting argument. The scientific content is entirely in the measurements: is the nesting constraint cheap, and does the recipe survive contact with modalities and datasets it was not designed for?

The four claims in the abstract

The paper makes four quantitative claims, and they are worth separating because they measure different things.

The third row deserves a moment. Why would forcing a representation to be usable at eight dimensions make it better on rare classes? A plausible account: the small-granularity terms cannot afford to memorize fine distinctions, so they push the encoder toward coarse, shared, reusable structure. Rare classes have few examples and therefore benefit most from structure shared with common classes rather than from class-specific memorization. The nesting constraint acts like a regularizer that happens to encourage exactly the kind of representation that transfers.

The four claims, re-derived from what you already know

You now have every mechanism needed to predict each of those results before reading them, which is the best test that the previous chapters landed. Try it.

Why should a 14× reduction be possible at matched accuracy? Because Chapter 1's redundancy argument says a 2048-dimensional representation of ImageNet does not use 2048 independent directions, and Chapter 3's PCA discussion says the useful directions exist but are not axis-aligned. MRL rotates them into alignment and orders them. The size of the reduction is a measurement of how much redundancy was there, and 14× says: a lot.

Why should the retrieval speed-up match the compression? Because Chapter 6's cost model is dominated by N · D_s, which is linear in the shortlist dimension. It is not a coincidence that both numbers are 14 — they are the same fact, measured on the memory axis and on the clock.

Why should few-shot on rare classes improve? Because Chapter 5's fine-grained example showed that the small-granularity terms generate residuals dominated by coarse, broadly shared structure — and coarse shared structure is exactly what a class with eight examples has to lean on, since it cannot afford class-specific memorization. The nesting objective is, accidentally, a regularizer that favours transferable features.

Why should it transfer across modalities? Because nothing in Chapter 2's derivation mentions images. The mechanism is an index-set inequality over a sum of losses; it is agnostic to what produced z. If it had failed on text that would have been the surprising result, and it would have implied something specific about text representations that nobody has a theory for.

The shape of the curve, and what to expect from it

The consistent finding across the paper's settings is easiest to hold as a shape rather than a table:

And the honest caveats, stated plainly because a lesson that only reports wins is advertising:

Which metric to read: linear probe or 1-NN?

The paper reports both, and if you are building retrieval you should weight them very differently. The distinction is worth twenty seconds of thought because it changes which number predicts your outcome.

The critical difference is the second row. A linear probe gets to refit for every dimension, so it can partly compensate for a badly oriented representation by choosing different coefficients. Retrieval cannot: cosine similarity has no parameters, so it is stuck with whatever geometry the encoder produced. That makes 1-NN the honest proxy, and it makes probe accuracy systematically optimistic about how well truncation will go.

Robustness: does the ordering survive distribution shift?

A reasonable worry. The coordinate ordering was learned on one training distribution. If you evaluate on shifted data — ImageNetV2 (a fresh test set collected the same way), ImageNet-R (renditions: art, sketches, cartoons), ImageNet-A (natural adversarial examples), ObjectNet (unusual viewpoints and backgrounds) — does the first coordinate still carry the most broadly useful signal, or was the ordering fitted to a quirk of the training set?

The paper evaluates exactly this, and the finding is that MRL's robustness tracks the baseline's: the ordering is a property of the task structure, not of the specific test distribution. That is what the mechanism predicts. Coordinate 1 became broadly useful because nine classifiers of differing capacity all demanded broad usefulness from it, and “broadly useful” is close to the definition of “transfers.” The pressure that creates the ordering is the same pressure that makes representations robust.

MRL in production — the fastest adoption of any 2022 representation paper

Something unusual happened to this paper. Between 2022 and 2024 the technique went from a NeurIPS paper to a parameter in the largest commercial embedding APIs. Very few representation-learning papers make that journey, and the reason is structural rather than about quality.

OpenAI's announcement makes the value proposition unusually legible, because they published the comparison. Their older text-embedding-ada-002 scored 61.0% average on MTEB — the Massive Text Embedding Benchmark, a suite of 50-odd retrieval, clustering, and classification tasks — at 1536 dimensions. The new text-embedding-3-large scores 64.6% at 3072 dimensions. And the sentence that made everyone pay attention: shortened to 256 dimensions, text-embedding-3-large still outperforms the full 1536-dimensional ada-002.

Run that as a bill. One hundred million chunks, which is a mid-sized enterprise RAG corpus:

Check one row: 100,000,000 × 1,024 bytes = 102,400,000,000 = 102.4 GB. ✓ The third row is a better model than the first, in a sixth of the space. The fourth row — which is Chapter 8 — fits an entire hundred-million-chunk corpus into the RAM of a single large machine.

What MTEB is measuring, so the numbers mean something

Every production claim in this chapter is denominated in MTEB points, so it is worth two paragraphs on what a point is. The Massive Text Embedding Benchmark aggregates dozens of datasets across seven task families, and an “average” is a mean over those families, not over a single metric.

The right-hand column is the practical warning. An average MTEB score that barely moves under truncation can hide a retrieval score that moved a lot, because retrieval is one family among seven and it is the one most exposed. If retrieval is your product, read the retrieval number, not the average. The published headline figures are averages; your decision should not be.

The three-year bill

Money makes the trade-off legible in a way percentages do not. Take a 100-million-chunk corpus on a 3072-dimensional embedder, and use a round illustrative figure of $5 per gigabyte-month for hot vector-index memory in a managed service. Substitute your own unit cost; the ratios are what carry.

Check one line: 100,000,000 × 3072 × 4 = 1.229 × 10¹² bytes = 1,229 GB; at $5 that is $6,145 per month. The last row is a better model than the ada-002 baseline, at 2% of the first row's cost, reachable by changing one API parameter and adding a per-dimension quantizer.

Quality per byte — the frontier nobody plots

Model comparisons are almost always “quality at the native dimension,” which silently prices bytes at zero. Once dimension is a runtime parameter, the honest comparison is a frontier: quality against bytes stored. Plotting it changes which model you would pick.

The “points per KB” column is deliberately a crude ratio — MTEB points are not linear and zero is not a meaningful origin — but the ordering it produces is the useful part, and it inverts the ordering you get from quality alone. On raw quality, 3-large at 3072 wins. On quality per byte it is the worst row in the table, twenty-four times worse than the same model truncated and quantized.

Triaging a quality regression after you turn truncation on

The first two rows account for most real incidents, and both are the same one-line mistake. Put an assertion in the serving path: every stored and every queried vector has unit norm to within 10⁻⁵. It costs nothing and it catches the entire failure class before it reaches a dashboard.

The API contract, and the one line every provider repeats

python — the dimensions parameter, and the thing you must not forgetimport numpy as np
from openai import OpenAI
client = OpenAI()

# Server-side truncation: ask for the prefix you want.
r = client.embeddings.create(model="text-embedding-3-large",
                              input="a query about kalman filters",
                              dimensions=256)
v = np.array(r.data[0].embedding)        # (256,) already re-normalized

# Client-side truncation from a full vector - YOU must re-normalize.
full = np.array(client.embeddings.create(
    model="text-embedding-3-large",
    input="a query about kalman filters").data[0].embedding)   # (3072,)

cut  = full[:256]                          # (256,)  NOT unit norm
print(np.linalg.norm(cut))                  # e.g. 0.83, and it VARIES per input
cut  = cut / np.linalg.norm(cut)            # now it is comparable

# Why it matters: without this line, a document whose signal happens to
# concentrate in the first 256 dims has a LONGER truncated vector, so its
# raw dot product with every query is inflated. Your ranking becomes a
# popularity contest between documents that front-load well.

That re-normalization line appears in the documentation of every provider that ships this feature, and it is Chapter 2's gotcha in production form. It is worth internalizing why it is not a formality: MRL deliberately front-loads magnitude, so the truncated norm is systematically large and systematically variable across documents. It is the one failure mode of matryoshka embeddings that produces plausible-looking wrong answers rather than obviously broken ones.

The evaluation protocol you should actually run

A number from a paper is a claim about their setup. Here is the protocol that turns it into a claim about yours, and it is short enough that there is no excuse for skipping it.

Step 3 is the one teams skip and then regret, because it is the only one that measures the thing adaptive retrieval actually depends on. Full-width parity and a pretty truncation curve tell you the representation is good; recall of the shortlist tells you the pipeline is correct. They are different questions and only the second can silently lose documents in production.

Migrating an index that already exists

You have a live 768-dimensional index and a new matryoshka model. Both cannot be true at once, and a flag day is unacceptable. The sequencing that works:

Step 3 exists purely to keep two variables from moving together, and it is the step schedules always try to delete. Do not let them: if you change the model and turn on truncation in the same deploy and quality moves, you will not know which one did it, and you will burn a week finding out.

A practitioner's FAQ

“Can I truncate an OpenAI embedding to 137 dimensions?” Mechanically yes, and it will probably work — the pressure staircase is smooth, so nearby coordinates are treated nearly identically. But you are outside the documented sizes, so measure it. And re-normalize.

“Do I have to re-embed my corpus to change dimension?” No, and this is the single most valuable operational property. Store the full-width vectors once. Any smaller index is a strided copy away, produced in minutes with no GPU. Changing your served dimension is a rebuild of a derived artifact, not a migration.

“Can I mix dimensions in one index?” No. A dot product between a 256-dimensional vector and a 768-dimensional one is not defined, and truncating one side at query time to match gives you a comparison at the smaller dimension — which is fine, as long as you know that is what happened. Pick a dimension per index.

“Does the query have to use the same dimension as the documents?” Yes, for the obvious reason, and with the same normalization rule on both sides. Truncate and normalize the query exactly as you truncated and normalized the documents. Asymmetry here is a silent ranking bug.

“Is 3-large at 256 dimensions really better than ada-002 at 1536?” On MTEB average, that is the published claim. On your corpus, run the twenty-minute test from Chapter 3. Benchmark averages are a prior, not a measurement of your system.

“What if my provider does not document matryoshka support?” Then it probably does not have it, and truncating will fail in the way Chapter 3 derives. Test before assuming: the magnitude profile takes thirty seconds and the recall test takes twenty minutes.

“Does this interact with hybrid search?” Only through the score-fusion weights. If you blend a dense score with BM25, the dense score's scale changes when you change dimension — the 1/√m floor from Chapter 6 — so the blend weight has to be retuned. Reciprocal-rank fusion, which uses only ranks, is immune, and that is a good reason to prefer it.

“Should I truncate or quantize first?” Quantize first if you cannot retrain, because it is post-hoc and nearly free. Truncate first if you can, because it is free at query time and needs no artifact. Then do both — Chapter 8 is about how they multiply.

Where the technique does not help

What the paper does not tell you, and where to look instead

Reading a paper well includes noticing its silences. None of these are criticisms — a conference paper has a page limit — but each is a question you will need answered and will not find there.

The third row generalizes into a habit worth keeping. Wall-clock speed-ups are measurements of a system, not properties of a method, and they do not survive a change of hardware, index library, batch size, or baseline. What transfers is the structure of the result — that the gain is dimension-proportional up to an Amdahl fraction you can measure yourself in an afternoon. Take the mechanism from a paper and the numbers from your own machine.

Chapter 8: Matryoshka × Quantization

An embedding's storage cost is a product of two independent numbers: how many coordinates and how many bits each coordinate takes. Matryoshka attacks the first. Quantization attacks the second. They are orthogonal, so they multiply — and the modern serving stack uses both.

Scalar quantization, derived

Quantization means replacing a continuous value with one of a finite set of levels. The simplest useful version, affine scalar quantization, maps a real range onto the 256 values an unsigned byte can hold.

Pick a range [a, b] covering the values a coordinate takes. Define the step size and the maps:

The rounding error is uniform on [−Δ/2, +Δ/2], so its root-mean-square is Δ/√12 = Δ/3.464. (A uniform distribution on an interval of width w has variance w²/12 — that is the standard result and the only fact you need.)

Worked example. A 1024-dimensional unit-norm embedding. Each coordinate has typical magnitude around 1/√1024 = 0.031, so a range of roughly ±4 typical magnitudes covers essentially everything: [a, b] = [−0.15, +0.15].

Relative to a typical coordinate magnitude of 0.031 that is 1.1% error per number, which sounds alarming. It is not, and here is the derivation that shows why. The quantity you actually care about is the dot product, and the errors in it partially cancel. Writing σ_e for the per-coordinate error standard deviation and treating errors as independent:

Five ten-thousandths of a cosine unit. Score gaps that decide rankings are typically in the hundredths. int8 quantization is, for retrieval purposes, free — four times less memory for an error two orders of magnitude below what matters. Notice also that the error does not grow with dimension: the unit-norm constraint means adding coordinates shrinks each one, and the two effects cancel exactly.

Binary quantization, derived

Push harder. Keep one bit per coordinate: is it positive?

Comparison becomes Hamming distance: XOR the two bit strings and count the set bits. A modern CPU does 64 bits per popcount instruction, so a 1024-dimensional comparison is 16 instructions instead of 1024 multiply-adds — and reads 128 bytes instead of 4,096.

Why does counting sign agreements estimate an angle? The classic result behind SimHash: for two vectors separated by angle θ, a random hyperplane separates them with probability θ/π. If the embedding's coordinates behave isotropically, each coordinate's sign acts like such a hyperplane test, so:

Each bit is a Bernoulli trial, so with m bits the estimate has a binomial standard error. Worked example at cosine 0.8, i.e. θ = 0.6435 rad = 36.87°, so p = θ/π = 0.2048:

Fifty times worse than int8's 0.00048, but still small next to the score gaps that decide a top-10 list — which is why binary quantization with a rescoring pass works well in practice. Now do it at 256 dimensions:

The memory ladder

One million documents, native dimension 1024. Every row is a real configuration someone ships.

Verify the headline row: 1,000,000 × 32 bytes = 32,000,000 bytes = 32 MB, against 4.10 GB for the naive configuration. A corpus that needed a dedicated machine now fits in a CPU's last-level cache neighbourhood. That 128× is the product of 4× from matryoshka and 32× from binarization, and neither knew the other existed.

Walk the grid diagonally from the top-left and notice that the two axes are not symmetric. Going from float32 to int8 costs almost nothing at any dimension. Going from 1024 to 256 costs a little. Going to binary costs a lot at low dimension and little at high dimension. That asymmetry is the whole guidance: quantize aggressively at high dimension, truncate aggressively at high precision, and be careful about doing both at once.

The calibration trap that is specific to matryoshka

Here is a bug that only exists because MRL front-loads magnitude, which makes it exactly the kind of thing a team discovers three weeks after shipping.

Suppose you quantize with a global range shared by all coordinates — a single (a, b) pair for the whole vector. On an ordinary embedding that is fine, because all coordinates have similar scale. On an MRL embedding they emphatically do not. Say coordinate 1 ranges over ±0.40 while coordinate 900 ranges over ±0.02.

The tail coordinates get about 3.7 bits of the 8 you paid for, and they are the coordinates carrying the fine distinctions that motivated keeping 1024 dimensions in the first place. You bought a full-precision tail and quantized it down to a coarse one.

The third axis: product quantization, and how it relates

Scalar quantization treats each coordinate independently. Product quantization (PQ) does something cleverer: it chops the vector into chunks and replaces each chunk with the id of its nearest entry in a learned codebook.

Concretely for d = 1024: split into 64 subvectors of 16 dimensions each. For each subspace, run k-means over a sample of the corpus to get 256 centroids. Now a document is 64 bytes — one centroid id per subspace.

And the query trick is elegant: precompute a 64 × 256 table of distances from the query's subvectors to every centroid, then each document's distance is 64 table lookups and 63 adds — no multiplications at all.

PQ is more powerful per byte and MRL is more powerful per unit of engineering effort. They are complements, not competitors: PQ is a post-hoc, fitted, corpus-specific artifact that squeezes precision; MRL is a training-time property that removes dimensions and needs nothing at serving time. The reason MRL spread faster is the “fitted artifact” row.

The quantizers, in code

python — per-dimension int8, and binary with rescoringimport numpy as np

def fit_int8(Z):
    """Z: (n_sample, m) of the TRUNCATED prefix you will actually serve."""
    lo = Z.min(0); hi = Z.max(0)          # (m,) each - PER DIMENSION
    return lo, (hi - lo) / 255.0              # 2m floats total, 2 KB at m=256

def to_int8(Z, lo, step):
    return np.clip(np.round((Z - lo) / step), 0, 255).astype(np.uint8)

def from_int8(Q, lo, step):
    return lo + Q.astype(np.float32) * step

# A GLOBAL range would look like this - and it is the bug:
#   step = (Z.max() - Z.min()) / 255.0     # ONE scalar
# On a matryoshka embedding coordinate 1 spans +-0.40 and coordinate 900
# spans +-0.02, so the tail gets ~13 of the 256 levels. Per-dimension fixes it.

def to_binary(Z):
    return np.packbits(Z > 0, axis=-1)         # (n, m) -> (n, m//8) uint8

def search_binary_then_rescore(qb, qf, B, F32, K=1000, k=10):
    """B: (N, m//8) packed bits. F32: (N, m) full precision, on SSD."""
    ham  = np.unpackbits(B ^ qb, axis=-1).sum(-1)   # popcount per row
    cand = np.argpartition(ham, K)[:K]
    return cand[np.argsort(-(F32[cand] @ qf))[:k]]

# Hamming -> cosine, when you need a comparable score from stage 1:
#   theta_hat = pi * hamming / m ;  cos_hat = np.cos(theta_hat)
# Its std is about 0.76/sqrt(m) - fine for shortlisting, useless as a threshold.

The last comment is Chapter 6's threshold warning wearing a different hat. A Hamming-derived cosine is an estimator with a known, dimension-dependent standard error, so it is perfectly good for “are these among the closest thousand” and hopeless for “is this above 0.75.” Use each score for the question its noise level can answer.

A gallery of ways this goes wrong

Every failure below has been shipped by somebody. They are collected here because each one produces plausible-looking output rather than an error, which is the worst property a bug can have.

Six of the seven share a root cause worth naming: a step in the pipeline assumed something about the representation that matryoshka training deliberately made false — that coordinates are interchangeable, that norms are comparable, that scores mean the same thing at every width. When you adopt an anisotropic representation, audit every downstream component for an isotropy assumption. Most of them have one, and none of them will tell you.

Which axis do you shrink first?

A full int8 round trip, on real numbers

Abstract error bars are less convincing than one number going in and coming out. Take a 256-dimensional matryoshka prefix and quantize two of its coordinates — an early one and a late one — both ways.

Calibrated on a corpus sample, the two coordinates have very different ranges, which is precisely the matryoshka signature:

Per-dimension calibration. Each coordinate gets its own step:

Global calibration — the bug. One range for everything, [−0.42, +0.44], step 0.0033725. Coordinate 1 is unaffected. Coordinate 200:

Twelve times the error, on the coordinates that were the whole reason you kept 256 dimensions instead of 64. And the structural version of the same fact: coordinate 200's entire range spans

Every tail coordinate is being stored at half the precision you are being charged for. The fix is 2d extra floats — 2 KB at m = 256, for the entire index — and it is not optional on a matryoshka embedding the way it nearly is on an ordinary one.

The real decision is which memory tier you land in

The tables above rank configurations by bytes, which makes the trade-off look continuous. It is not. Memory is a hierarchy with hard steps between tiers, and the only question that changes your latency by an order of magnitude is which tier your index falls into.

Now place the configurations for one million documents:

The last row is not sixteen times faster than the second, it is roughly sixty times faster, because it crossed a tier boundary rather than moving along a line. Fitting an entire retrieval index in last-level cache changes what kind of system you are operating.

Rescoring: the same cascade, one axis over

Chapter 6 built a cascade over dimensions. Precision gives you a second axis for exactly the same trick, and the two compose into the stack most modern retrieval systems now run.

Each stage is cheap because the previous stage made its input small, and each is accurate enough because it only has to preserve the true answers rather than rank everything correctly. That is the single principle running through this entire lesson: an approximation does not need to be right, it needs to be right about what it forwards.

Total cost is dominated by stage 1, which touches 128 bytes per document instead of 4,096 — and 128 MB of hot index instead of 4.1 GB. Stages 2 and 3 together touch 1,000 × 256 + 100 × 4,096 = 256,000 + 409,600 = 665,600 bytes, less than a megabyte, which is nothing.

The other precisions: fp16, fp8, int4

Chapter 8 has treated precision as three points — float32, int8, binary — because those are the three that dominate. The full ladder is worth having, along with the one non-obvious ordering result on it.

Rerun the scalar-quantization derivation for four bits. Sixteen levels over a per-dimension range of 0.30:

Seventeen times worse than int8 — exactly as expected, since each dropped bit doubles the step and doubles the error, and 8 to 4 bits is a factor of 2⁴ = 16 in step size.

Chapter 9: Two Dimensions, Open Questions, and the Cheat Sheet

MRL nests along one axis: the width of the representation. Once you see that a nesting constraint is cheap, the obvious question is which other axes admit the same trick. The answer, as of 2026, is: most of them, and the follow-up literature is essentially a catalogue of axes.

Axis two: depth

A Transformer encoder with 12 layers computes 12 successive representations of its input, and you throw away the first 11. But a layer is a budget just as a dimension is: exiting at layer 6 costs half the compute. So nest over layers as well as over widths.

That is the idea behind 2D Matryoshka approaches (published as 2D Matryoshka Sentence Embeddings, and related work under the name Espresso Sentence Embeddings). The objective becomes a double sum:

Where h^(ℓ) is the pooled output of layer ℓ. If you nest over 4 exit layers and 7 widths you get 28 operating points from one training run — a two-dimensional menu of (latency, memory) trade-offs, all sharing one geometry.

MatFormer is the most striking of these because it moves the nesting inside the network rather than onto its output. Each Transformer block's feed-forward layer is trained so that its first k hidden units form a working FFN on their own, for several values of k. Because the choice can be made independently per layer, you can then “mix and match” a submodel per deployment target — hundreds of extractable models from one training run, all with consistent behaviour because they share weights. This line of work feeds directly into shipped on-device model families that expose a smaller nested variant of a larger model.

The four-year arc, briefly

Read the right-hand column as one sentence. In four years the field moved from “here is a clever training objective” to “a model that cannot be run at several sizes is a model with a missing feature.” That is a fast transition for an idea whose entire content is a sum over a loss, and it happened because the transition cost was near zero at every step.

MatFormer, mechanically

Because it is the most consequential descendant, it is worth seeing how the trick moves inside the network. A Transformer block's feed-forward layer is

with d_ff typically 4d. The hidden units are exchangeable — permuting them and permuting W₂'s columns to match gives an identical function — which is precisely the rotation-invariance situation of Chapter 1, one level in. So nest them: for each k in a set of widths, define a sub-FFN using only the first k hidden units,

and add a loss term for the whole network run with every block using width k. Same objective shape, same multiplicity argument, same near-free constraint. The payoff is different, though, and bigger in one specific way: because each block's width can be chosen independently at extraction time, a network with B blocks and W widths per block yields W^B possible submodels, all sharing weights and all behaving consistently because they were trained jointly. In practice you pick a handful along the accuracy–latency frontier.

What is genuinely still open

A good paper leaves better questions than it answers. These are the ones this one leaves.

What would make this obsolete

Every technique has a world in which it stops mattering, and naming that world is a good test of whether you understand what the technique is for. Four candidates, with an assessment of each.

The third row is worth sitting with because it inverts the usual intuition. In late-interaction schemes a document is represented by one vector per token rather than one per document, which multiplies the index by a factor of a hundred or more. That makes per-vector dimension the dominant cost term, and a nested representation the obvious response. Techniques do not only get obsoleted by the future; sometimes the future makes them load-bearing.

A glossary of everything this lesson bolded

Five experiments you could actually run this week

The open questions above are large. These are small, and each one has a chance of producing a real answer on a single GPU.

1. Does the intrinsic dimension predict the tax? Train MRL on several datasets whose intrinsic dimension you can estimate (a nearest-neighbour or maximum-likelihood estimator on the representations will do), and plot the full-width quality gap against that estimate. The hypothesis from Chapter 1 is that the gap shrinks as redundancy grows. Nobody has published the scatter plot.

2. Where exactly does the ordering stop? Train with M = {64, 128, …} and evaluate at m = 8, 16, 32 — strictly below the smallest granularity. Chapter 1 argues there should be no ordering inside the first block. Does quality fall off a cliff at 64, or does it decay smoothly because the pressure staircase leaks? The answer tells you how conservative to be when choosing M.

3. Does a curriculum beat a sum? Train phase one with only m = d, then phase two with the full nested objective, and compare the frontier against joint training at matched total compute. If the two-phase version wins, MRL becomes a cheap fine-tune of any existing model and the retrofitting story in Chapter 4 changes completely.

4. Is the ordering robust under shift? Take a matryoshka model, measure the truncation curve on in-distribution data and on a shifted set, and compare the gaps rather than the absolute numbers. Chapter 7 reports that robustness tracks the baseline; nobody has asked whether truncatability specifically degrades faster than quality.

5. Build the index at which dimension? Construct an HNSW graph at 16 dimensions and another at 2048, then use each to shortlist and rerank at full width. The 16-dimensional graph is far cheaper to build and traverse, but it encodes a coarser notion of neighbourhood. Does the cheap graph lose recall, and by how much? This is a two-day experiment with an immediately actionable answer.

The cheat sheet — every symbol

Every equation, in order

The numbers worth remembering

Build it yourself — the checklist

Should your model be matryoshka? A decision procedure

The honest summary of ten chapters, arranged so you can act on it.

And the summary of the summary: MRL is a cheap way to buy optionality on the one axis that governs serving cost, paid for with a constraint that turns out to be nearly free because learned representations are redundant. It is not an accuracy technique. It is a technique for making one artifact serve requirements that used to need several, and for letting the cheap version legitimately hand work to the expensive one.

Where to go from here

The whole thing, in one paragraph

If you have to explain this to a colleague in an elevator: an ordinary embedding spreads its information evenly across all its coordinates, because nothing in the training objective ever expressed a preference — the loss is invariant to rotations of the representation space, and truncatability is not. So cutting a 768-dimensional vector to 64 numbers keeps only √(64/768) = 29% of the discriminability and the retrieval collapses. Matryoshka Representation Learning fixes that by adding one term to the loss for each of a handful of prefix sizes. Because coordinate j appears in every prefix of length m ≥ j, the first coordinates accumulate gradient from every term while the last accumulate from one, and the optimizer responds by loading the most broadly useful signal into the front. Nothing sorts anything; the ordering falls out of an index-set inequality. The result is that the first m numbers of the vector are a complete m-dimensional embedding for every m you trained, slicing is a pointer and a length, and every prefix lives in the same geometry — so a 16-dimensional sweep can legitimately shortlist for a 2048-dimensional rerank. That is 128× fewer FLOPs and about 14× on a clock, and it costs about 0.2% of a training step.

What to remember in five years

Reading paths through this lesson

Ten chapters is a lot to re-read. Depending on why you came back, here is the short route.

And if you have five minutes and a whiteboard, the two things to draw are the multiplicity staircase — nine, eight, seven, down to one, over exponentially growing blocks — and the two-stage cost model, N D_s + K D_r. Those two pictures generate everything else in the lesson.

References

1. Kusupati, A., Bhatt, G., Rege, A., Wallingford, M., Sinha, A., Ramanujan, V., Howard-Snyder, W., Chen, K., Kakade, S., Jain, P., Farhadi, A. “Matryoshka Representation Learning,” NeurIPS 2022 — arXiv:2205.13147. The paper this lesson is built on.
2. Kudugunta, S. et al. “MatFormer: Nested Transformer for Elastic Inference,” 2023 — arXiv:2310.07707. Matryoshka nesting moved inside the FFN, giving hundreds of extractable submodels.
3. Li, X. et al. “2D Matryoshka Sentence Embeddings,” 2024 — arXiv:2402.14776. Nesting over encoder depth as well as embedding width.
4. Nussbaum, Z. et al. “Nomic Embed: Training a Reproducible Long Context Text Embedder,” 2024 — arXiv:2402.01613. An open MRL-trained text embedder with documented truncation points.
5. Muennighoff, N., Tazi, N., Magne, L., Reimers, N. “MTEB: Massive Text Embedding Benchmark,” 2022 — arXiv:2210.07316. The yardstick every production claim in Chapter 7 is measured against.
6. Malkov, Y. A. & Yashunin, D. A. “Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphs,” 2016 — arXiv:1603.09320. The index whose latency-bound traversal explains Chapter 6's Amdahl ceiling.
7. Johnson, J., Douze, M., Jégou, H. “Billion-scale similarity search with GPUs” (FAISS), 2017 — arXiv:1702.08734. The reference implementation of everything in Chapter 6.
8. Guo, R. et al. “Accelerating Large-Scale Inference with Anisotropic Vector Quantization” (ScaNN), 2019 — arXiv:1908.10396. Quantization designed for the inner product rather than for reconstruction — the same “optimize the right objective” lesson as Chapter 3's PCA critique.
9. Charikar, M. “Similarity Estimation Techniques from Rounding Algorithms” (SimHash), STOC 2002. The θ/π result behind Chapter 8's binary-quantization error analysis.
10. Devlin, J. et al. “BERT,” 2018 — arXiv:1810.04805; Jia, C. et al. “ALIGN,” 2021 — arXiv:2102.05918. Two of the backbones MRL was shown to transfer to.
11. Radford, A. et al. “Learning Transferable Visual Models From Natural Language Supervision” (CLIP), 2021 — arXiv:2103.00020. The contrastive setting in which most production matryoshka embedders are actually trained.
12. OpenAI. “New embedding models and API updates,” January 2024. The dimensions parameter, the MTEB comparison, and the instruction to re-normalize after truncating.