Matryoshka training makes every prefix a usable representation on its own

Shrinking an embedding index is arithmetic about bytes before it is anything else. A retrieval system stores one vector for every chunk it has indexed — a chunk being one passage-sized piece of a document, held alongside the embedding computed from it — and once there are fifty million of them, those vectors are the largest object the system owns. A 1024-dimension vector holding four bytes a number occupies 4,096 bytes, so fifty million of them occupy 200 GB, and that is 200 GB of memory rather than of disk, because the vectors have to sit somewhere fast enough to be searched thousands of times a second. Two levers make that number smaller: store fewer numbers per vector, or store each number in fewer bits.

Dimension truncation is the first lever in its crudest possible form: keep the leading part of each embedding and discard the rest. Cutting 1024 numbers down to their first 256 turns a 4,096-byte vector into a 1,024-byte one, a 200 GB index into a 50 GB one, and every distance computed during a search into a quarter of the arithmetic it was. That saving is the motivation, and at fifty million chunks it is a large one — large enough to change which machines the system needs, which is why the idea gets tried before anybody checks whether it is allowed.

Truncation is not an operation you can apply to any vector, though; it is a property of the model that produced it, and most models lack it. Matryoshka embeddings, named after the Russian nesting dolls in which every doll contains a smaller complete doll, come from encoders trained under Matryoshka Representation Learning (MRL), a method published by Kusupati and colleagues in 2022 that applies the training objective to the leading slices of the output as well as to the whole of it. The first 256 coordinates of such a model are a working 256-dimension embedding in their own right and the coordinates after them add finer detail, whereas an ordinary model puts nothing in particular in its first 256 slots, so cutting there leaves numbers that were never asked to describe a document together. Which kind of model you have is stated in documentation and nowhere else — a dimensions parameter on an embedding API, a ladder of supported widths on a model card — since the vectors themselves look identical either way.

An ordinary embedding has no privileged axes#

Retrieval encoders are trained with a contrastive objective — one that pulls a query and its correct document towards each other in the space while pushing unrelated pairs apart — expressed over cosine or dot-product similarity. That objective is invariant under rotation of the embedding space. The statement sounds like a technicality about linear algebra, and it is the reason the rest of this chapter has anything to say, so it is worth walking through rather than accepting.

An orthogonal matrix R is a square matrix whose columns are unit-length and mutually perpendicular, and multiplying a vector by one rotates that vector — possibly with a reflection — without stretching it. Take any two of a model's output vectors and multiply both by the same R. Their dot product is unchanged, because rotating both operands moves the whole configuration rigidly: neither vector's length changes, the angle between them does not change either, and the dot product is exactly the length of the first times the length of the second times the cosine of that angle. Cosine similarity is the same quantity with the two lengths divided back out, so it is unchanged for the same reason. Every pairwise similarity in the space comes out identical after the rotation.

A contrastive loss reads nothing but those similarities. Feed it the rotated vectors and it computes the numbers it computed before, so the loss on the rotated model equals the loss on the original — at initialisation, at every gradient step, and at convergence. Training cannot tell the two apart, which means nothing in training ever expressed a preference for one set of axes over another. "Coordinate 3 and coordinate 900 are interchangeable" is not a loose way of gesturing at that; it is what the invariance says, because there is a rotation that exchanges those two axes and the objective cannot see it happen.

Whatever ended up in the first 256 slots of a bge-large vector is therefore an accident of initialisation and optimiser path. Slicing there keeps a projection onto a subspace chosen by the order the coordinates happen to be written down in, and that order is not a fact about meaning. The subspace carries no particular claim to being the important one.

It is also not worthless, which is the part that makes the failure hard to see. Keeping an arbitrary coordinate subspace behaves roughly like a random projection, and the Johnson–Lindenstrauss lemma says random projections preserve pairwise distances up to a distortion that shrinks as the target dimension grows. The guarantee is weak at the sizes involved here, and how weak is something you can compute rather than something you have to take on trust.

The standard bound is 8 ln(n) / ε²: that many dimensions suffice to hold the distortion of every pairwise distance below ε, where n is how many points are being embedded, ln is the natural logarithm, and ε is the fraction by which a distance is permitted to move. The usual reading runs left to right — pick the distortion you can tolerate, get the dimension you need. Here the dimension is not free, because 256 is what you have decided to keep, so set the bound equal to the width you kept and read it the other way, as 256 = 8 ln(n) / ε².

Rearranging that is allowed because ε is the only unknown in it. The count n is fifty million, the 8 and the 256 are fixed, and isolating a symbol that appears nowhere else cannot smuggle an assumption in behind it. Multiply both sides by ε² to clear it out of the denominator, which leaves 256 ε² = 8 ln(n), then divide both sides by 256, which leaves ε² = 8 ln(n) / 256.

Now put fifty million in for n and read the result off in stages. The natural logarithm of fifty million is a bit under eighteen. Eight times a number a bit under eighteen is about a hundred and forty. Dividing about a hundred and forty by 256 leaves a shade over a half, so ε² is a shade over a half, and ε is the square root of that — the ε ≈ 0.74 this case works out to.

What that number licenses is nothing. A bound permitting distances to move by seventy-four percent permits a distance to move by roughly three quarters of itself in either direction, and an ordering of documents whose true similarities differ by two percent is settled by gaps hundreds of times narrower than the permitted error. The bound cannot adjudicate that ordering. It does not predict that the ordering will break; it says it has nothing whatever to offer on the question. Real distortion is much smaller than the worst case, which is why coarse topical structure survives and a query about flood exclusions still returns flood documents. Fine ordering does not survive, and fine ordering is what recall past the first few results is made of.

Real embedding spaces also violate the isotropy that makes random projection well behaved. Isotropy is the property of a cloud of points looking statistically much the same in every direction, and it is the condition under which one arbitrary set of axes really is as good as another. Transformer embeddings are strongly anisotropic and carry a handful of outlier dimensions with magnitudes an order of magnitude above the rest, artefacts of the attention machinery rather than of meaning. Whether one of those rogue dimensions lands inside your slice is luck, and if it does, it dominates the dot product for every document and flattens the distinctions you were relying on.

One insurance search index built on bge-large-en-v1.5 had two such coordinates, index 588 and index 941, whose absolute values averaged 3.4 and 2.9 against a per-coordinate mean of 0.021 across the rest. Those three figures decide how much of a vector a slice throws away, once you remember what a Euclidean length does with a coordinate: it squares the coordinate before adding it to a running total, and takes the square root only at the very end. Squaring is not neutral bookkeeping here. It widens the distance between a large coordinate and a small one before anything is summed, so the largest few coordinates end up owning the total.

Work it through with those averages. Both rogue coordinates sit near three in absolute value, and a number near three squares to about nine. An ordinary coordinate sits near two hundredths, and two hundredths squared is about four ten-thousandths. Each outlier therefore puts about nine into the sum of squares while each ordinary coordinate puts in four ten-thousandths, and a thousand ordinary coordinates together still come to well under one. Two coordinates hold the overwhelming majority of the vector's squared length, and neither of them is inside the first 256 slots. The slice does two damaging things in one move: it keeps a subspace that was never asked to separate anything, and it throws away most of the length of every vector it touches.

It is worth being precise about what the arbitrariness does and does not mean. The claim is not that the first 256 coordinates are noise; they carry real signal, in the same sense that any 256-dimensional random view of a 1024-dimensional cloud carries real signal. The claim is that they carry no more of it than coordinates 400 to 655 would, and that the amount they carry was never optimised for. When you truncate an ordinary model you are choosing a compression ratio without choosing what to keep, and the model had no opportunity to arrange its output so that the answer to that question would be a good one.

The objective that makes a prefix mean something#

The nested objective removes that arbitrariness by putting the loss on the prefixes directly. Instead of computing the contrastive objective once on the full output, you compute it at a nested ladder of dimensions — 64, 128, 256, 512 and so on up to the model's native width — and sum the results.

import torch
import torch.nn.functional as F

NESTED_DIMS = (64, 128, 256, 512, 1024, 1536, 3072)


def matryoshka_infonce(q: torch.Tensor, d: torch.Tensor, temperature: float = 0.05):
    """q, d: (batch, 3072). d[i] is the positive for q[i]; the rest are in-batch negatives."""
    labels = torch.arange(q.size(0), device=q.device)
    total = q.new_zeros(())
    for m in NESTED_DIMS:
        qm = F.normalize(q[:, :m], dim=-1)
        dm = F.normalize(d[:, :m], dim=-1)
        logits = (qm @ dm.T) / temperature
        total = total + F.cross_entropy(logits, labels)
    return total / len(NESTED_DIMS)

Follow what that sum does to the gradient, because the arrangement in the output is a consequence of it rather than something the code asks for anywhere. The loss is a sum of terms, one per rung of the ladder. The 64-dimension term is computed from q[:, :64] and d[:, :64] and nothing else, so when the optimiser differentiates it, the derivative with respect to every coordinate past the sixty-fourth is zero — those coordinates do not appear in that term at all. The same holds one rung up: the 128-dimension term differentiates only with respect to the first 128 coordinates. Adding terms together adds their gradients, so the first 64 coordinates receive gradient from every term in the sum, coordinates 65 to 128 from all but one, and the widest block from exactly one. An optimiser reducing a total is served best by improving whatever appears in the most terms, so the coarse separation that every rung needs migrates into the coordinates that every rung reads, and the distinctions only the widest rung can afford get pushed outward into the tail.

One link in that chain is a requirement and the rest are decisions. That a term carries no gradient into coordinates it never reads is arithmetic and could not be otherwise. Which rungs exist is a choice — NESTED_DIMS in the loop above is a tuple somebody wrote down — and so is how heavily each rung counts, which here is equally, since the loop divides the total by the number of rungs. Weight the narrow rungs harder and the prefixes improve at the full width's expense; drop a rung and nothing in particular protects that width.

The whole method is that loop. There is no special architecture, no extra projection head and no second training stage; the model is the same encoder it would otherwise have been, trained against a sum of objectives rather than one. The cost is a few percent of training throughput and some care in weighting the terms, since giving the 64-dimension term equal weight with the 3072-dimension one trades away a little full-width quality to buy prefix quality. Published models make that trade differently, and it is why text-embedding-3-large at its full 3072 dimensions scores slightly below what a non-nested model of the same size and data would, while remaining usable at a sixteenth of its width.

Running the same slice-and-measure comparison over both kinds of model makes the difference legible. The table below reports recall@10 — the share of a query's genuinely relevant documents that the top ten results contain — as a fraction of each model's own full-width recall@10, so the two are compared on how much they lose rather than on where they start; the absolute baselines were 0.79 for bge-large-en-v1.5 at 1024 dimensions and 0.83 for text-embedding-3-large at 3072.

The rows are labelled by fraction of native width, and that is not the same thing as dimension count. On the 25% row, the left-hand model is being cut to 256 coordinates because it is 1024 wide, while the right-hand model is being cut to 768 because it is 3072 wide. Comparing at equal fractions is what makes the comparison fair: each model is asked how much of its own retrieval quality it keeps, not how good it is in absolute terms, and the two columns start from different baselines precisely so that question can be asked. It is also what makes the table easy to misread. The right-hand column is not a claim about equal byte counts — on that row it is storing three times as many numbers per vector as the left — and reading it as one turns a statement about graceful degradation into a shopping recommendation it does not support. What the row does say is that cutting a nested model to a quarter of its own width costs about a point of its own recall, while cutting an ordinary model to a quarter of its own width costs nearly a fifth.

The models that offer this are a short and identifiable list rather than a general capability. OpenAI's text-embedding-3-small and text-embedding-3-large accept a dimensions parameter that returns a truncated and re-scaled vector directly from the API. Nomic's nomic-embed-text-v1.5 documents a nested ladder from 768 down to 64. Several of the mxbai and Snowflake arctic-embed releases advertise the same, and the sentence-transformers library ships a MatryoshkaLoss wrapper that turns any existing training script into the loop above, which means a team fine-tuning their own encoder can add the property for the cost of a retrain. What no team can do is add it to a checkpoint after the fact by slicing, because the arrangement the property depends on has to be created by the gradient.

Nothing about a 1024-dimension float array tells you whether its first quarter was trained to stand alone, and no runtime check will raise if it was not. A slice of a nested model and a slice of an ordinary one are the same kind of object — an array of numbers of the expected length — and the index, the distance function and the client library will accept either without comment.

The spot check that cannot fail#

The defence that failed at Halvorsen is worth naming in its strongest form, because in that form it is obviously useless and in its ordinary form it looks like diligence. Three queries returning the same top result before and after a change is a test whose answer does not depend on whether the change was safe. It returns "pass" on the intact index and "pass" on the damaged one, so running it tells you which world you are in exactly as much as not running it does. A defence that stays constant while the thing it is meant to discriminate varies is not a weak signal but an absence of one, and no number of repetitions makes it stronger.

The reason it is tempting is that the property making a query easy to check by hand is the same property that makes it insensitive. To verify a result without a labelled set you need a query with an obvious right answer, and an obvious right answer is one that beats its rivals by a wide margin. A query whose best answer sits at cosine 0.71 while the next-best candidate sits at 0.58 has thirteen points of margin, and a projection that perturbs similarities by two or three points cannot move a document across a gap that size. So the queries an engineer keeps in a notebook are precisely the ones selected for being unable to detect this class of damage.

What the check cannot see is where the ranking is actually decided. Three queries returning the same top result is compatible with a fifteen-point recall drop, because those two things measure different parts of the ranking. The damage lands on the queries where the top twenty candidates are separated by half a point of cosine each — near-identical policy wordings from different underwriting years, two versions of the same handbook section, an endorsement and the clause it endorses. Those are the queries nobody keeps in a notebook, for the same reason: you cannot tell by looking whether their answer is right.

The habit worth building is to measure the change on the ranking rather than on the result. Slicing and re-ranking two million chunks against a two-hundred-query set takes a few minutes of brute-force numpy on one machine, and it produces recall at several cutoffs instead of a yes or no on three examples. Halvorsen's eventual harness did exactly that, and the first thing it reported was that their 256-dimension slice of bge-large had a mean rank displacement of 41 positions for documents that had been in the true top 100.

Rank displacement is a more sensitive instrument than recall at any single cutoff, because it registers movement that has not yet crossed a threshold. A document that falls from rank 6 to rank 9 changes no recall@10 number and is a warning that the next increment of compression will push it out. Measuring it costs nothing extra once you are already computing exact neighbours for the baseline, since it is the difference between two ranked lists you already have in memory, and it gives you a distribution rather than a scalar — the mean of 41 came with a ninetieth percentile of 137, which said the damage was concentrated rather than uniform.

Where truncation is the wrong tool#

Fitting a PCA projection — principal component analysis, which learns from your data which combinations of coordinates carry the most variance and keeps those — on your own corpus is the honest alternative for a model that was not trained for nesting, and it is genuinely better than slicing: on the same pilot, PCA to 256 dimensions fit over two million sampled bge-large vectors retained 0.91 of full-width recall@10, against 0.81 for the arbitrary slice. It costs a stored 1024×256 projection matrix that every query path must apply, a refit whenever the corpus drifts far enough that the retained variance shifts, and an ordering constraint that says vectors written before the refit cannot be compared with vectors written after. Matryoshka truncation has none of that machinery because the projection is the identity on a prefix.

Truncation also does not make embedding cheaper. Asking text-embedding-3-large for 256 dimensions still runs the full forward pass and still bills the full token rate; the API computes 3072 numbers and returns a slice of them. What shrinks is storage, index memory and the per-comparison cost of search, which is the right win for a fifty-million-chunk corpus and no win at all for a workload whose bill is dominated by embedding fresh documents. A team paying mostly for inference is looking at a smaller model, which is a different decision with different quality consequences.

Nor is a nested model a licence to pick the smallest prefix that still works on a benchmark. The ladder gives you a dial with real positions on it, and where you set the dial depends on what reads the results — a question the next two chapters take up in earnest, one on re-normalising a truncated prefix before anything compares it and one on which ranks the loss actually lands in, once the mechanical prerequisite here is out of the way.

The pilot and the table are the same measurement seen from two sides, and laying one over the other is the most useful thing in this chapter. The table gives bge-large-en-v1.5 a retention of 0.81 at a quarter of its native dimensions. A quarter of 1024 is 256, which is precisely the cut the pilot made, and the pilot's own full-width baseline was 0.79. Multiply the retention by that baseline — eighty-one hundredths of seventy-nine hundredths — and the product is the 0.64 the regression run measured a week after everyone had moved on. Nothing in that arithmetic knows anything about insurance, about policy wordings, or about how the three notebook queries were chosen.

So the pilot's result was neither bad luck nor a corpus that happened to be difficult. It is the number the table predicts for an ordinary encoder cut to a quarter of its width, and the prediction was available before the slice rather than only after it: take the fraction of native width you are considering, read the retention for the kind of model you actually have, multiply it by the recall you measure today, and you have the recall you are about to be living with. The one thing that estimate says nothing about is the vectors themselves. Slicing a unit-length vector produces something that is no longer unit length, and every distance function in your stack has an opinion about that.