matryoshka-slice

Truncate embeddings to a shorter prefix dimension and keep the scores comparable across widths.

import { defineSpace, sliceTo, rank, calibrateThreshold, applyThreshold } from 'matryoshka-slice';

const space = defineSpace({
  id: 'text-embedding-3-large@2024-01',
  sourceDim: 3072,
  metric: 'innerProduct',
});

// The policy is required. Under cosine the two options rank identically,
// under inner product they do not, so there is nothing safe to default to.
const query = sliceTo(space, queryVector, 512, 'renormalize');
const corpus = docs.map((d) => ({ id: d.id, vector: sliceTo(space, d.vector, 512, 'renormalize') }));

const { hits, separable, boundaryRisk } = rank(query, corpus, { limit: 10 });

hits[0].score.atWidth;        // the exact score at 512 dimensions
hits[0].score.full;           // { lower, upper }: provably contains the 3072 dimension score
hits[0].uncertainAgainst;     // ids this hit cannot be ordered against from a prefix alone
boundaryRisk;                 // ids outside the top 10 that could belong inside it

Cutting the array is one line, and it loses the number that mattered

A matryoshka trained model emits unit vectors, so v.slice(0, 512) looks free. It is not. The prefix of a unit vector is shorter than the vector, by an amount that is different for every document:

doc A   ||v|| = 1.0    ||v[0:512]|| = 0.94    kept 88 percent of its energy
doc B   ||v|| = 1.0    ||v[0:512]|| = 0.61    kept 37 percent of its energy

Cosine divides that number out of both sides, which is why nothing looks wrong. Recall holds, the test suite is green, and the retained norm never appears in a score. Inner product does not divide it out. It sits in the score as a multiplier, so doc A is boosted against doc B for a reason that has nothing to do with the query.

This module keeps sourceNorm, retainedNorm and discardedNorm on every sliced vector, because once slice() has run they cannot be recovered, and every guarantee below is built out of them.

The reordering, in numbers you can check by hand

Eight dimensions, truncated to four. The query, doc A and doc B all live on two coordinates:

query   [0.707, 0, 0, 0,   0.707, 0, 0, 0]
docA    [0.100, 0, 0, 0,   0.995, 0, 0, 0]
docB    [0.500, 0, 0, 0,   0.000, 0, 0, 0]

At the native width the inner products are 0.774 for doc A and 0.354 for doc B. Doc A wins. Truncate to four dimensions and the prefix inner products are 0.071 and 0.354. Doc B wins. Nothing threw, nothing returned NaN, and a recall metric computed against truncated ground truth agrees with itself.

What score() returns for the same pair at width four:

docA    atWidth 0.071    full [-0.633, 0.774]
docB    atWidth 0.354    full [ 0.354, 0.354]

rank() still orders by atWidth, because that is the number a real index computes. What it adds is that both hits come back with the other one listed in uncertainAgainst, and separable is false. Ask for the top result only and boundaryRisk names doc A: the document truncation pushed out.

The interval is not a heuristic and not a confidence interval. Write the native width inner product as the prefix part plus the tail part. Cauchy Schwarz bounds the tail part by the product of the two discarded norms, and both of those were measured before the tails were thrown away. So full provably contains the native width score, it shrinks as the width grows, and at the native width it collapses to a point. Doc B's tail is exactly zero, so its interval is a point already.

Padding a short vector up to a long one is the version with no symptom at all

A 768 dimension document and a 1536 dimension query do not have the same shape, and the fix that suggests itself is to pad the document with zeros. That produces a defined score. It is <q[0:768], x> in the numerator and ||q[0:1536]|| * ||x|| in the denominator, which is the correct 768 dimension score multiplied by ||q[0:768]|| / ||q[0:1536]||.

That factor is below one whenever the query has any energy past dimension 768. For a query with energy spread evenly it is about sqrt(768/1536), roughly 0.71. It is the same factor for every short document, because it depends only on the query. So the short documents do not scatter, they sink together, by an amount that looks like a plausible relevance gap. zeroPadDeflation(query, 768) computes it on your own data if you want to see the size of it before deciding it is small.

This module refuses the comparison instead:

• score() throws WIDTH_MISMATCH on two different stored widths, and the message names padding as the thing it is not doing.
• sliceTo() and reslice() throw INVALID_DIM on any request to widen.
• assertUniform() fails a mixed width corpus at index build time, which is where you want to hear about it, rather than on whichever query first touches a narrow document.
• auditWidths() counts the widths without throwing, for when you want to look before you decide.

The one sanctioned move is alignWidths(a, b), which cuts the wider side down to the narrower. Cutting is symmetric: both sides lose the same coordinates and the result is a real score at the narrower width. Padding is not symmetric, and that asymmetry is the whole bug.

reslice() carries the original sourceNorm and the original discarded tail through the second cut, so a vector cut 1536 to 768 and then 768 to 256 reports the same discarded energy as one cut straight to 256, and the native width bounds survive the round trip.

A threshold is a number plus the width it was measured at

0.82 at 3072 dimensions and 0.82 at 512 dimensions are not the same operating point, and the direction of the difference is a property of your corpus, not of the two widths. Truncated scores are computed over fewer terms and spread out more, so a cut off in the tail can admit several times as many documents at the narrow width as it did at the wide one.

On a sample of 80 pairs at 64 dimensions, a cosine cut off of 0.181 admits 12 of them. Carried over unchanged to 8 dimensions, the same number admits 25. Nothing about that comparison errors.

So ScoreThreshold carries the width it is valid at, and applyThreshold() throws THRESHOLD_MISMATCH rather than comparing across widths. calibrateThreshold() fits a value for the narrow width against a sample of real pairs, matching the accept count rather than the score, because the accept count is the operating point and the score is not. On that same sample the fitted value is 0.418.

It reports what the fit could not fix. 16 of the 80 pairs still change verdict between the two widths: matching the count is not the same as matching the set, and a single reassuring agreement number would hide that. Calibration is refused outright when the sample is smaller than 24 pairs, when the source threshold accepts none of the sample, and when it accepts all of it, since none of those cases contain a boundary to fit against.

A calibrated threshold also carries the native width value it reproduces, so applyThreshold() can return sourceWidthVerdict: 'indeterminate' for a pair whose interval straddles the original cut off. A declared threshold reports 'unavailable' instead, because a number typed in for one width says nothing about any other.

Known limitations

The bound is tight only when the tails are aligned. residualBound is ||q_tail|| * ||d_tail||, which is attained when the discarded halves point the same way and is loose when they are near orthogonal, as they usually are in high dimensions. For a query and a document that each keep 90 percent of their energy in the prefix, the bound is around 0.1 wide on a cosine, which is enough to make many real pairs indeterminate even though the truncated ranking is fine. separable: false means the order is unproven, not that it is wrong.

Nothing here measures relevance. The module can tell you truncation may have reordered two documents. It cannot tell you the native width order was the better one. If your evaluation says the 512 dimension index answers questions well, that is evidence this module does not have.

Calibration is quantile matching on the sample you supply. It reproduces an accept count on that sample and nothing more. A sample drawn from a different distribution than production traffic gives a threshold fitted to the wrong distribution, and the reported agreement will look fine because it is computed on the same sample.

Reports are per corpus, not per query. sliceReport() summarises retained energy over a set of vectors. The score shift for a specific query depends on that query's own discarded tail, which is why the per pair bound lives on Score and not in the report.

Vectors are held as Float64Array and copied on every slice. That is roughly eight bytes per stored dimension plus one allocation per sliceTo or reslice. For a large index this is a scoring and analysis library, not a storage format.

No approximate nearest neighbour index. rank() scores every candidate you hand it. The intended use is rescoring a candidate set, or auditing a truncation decision before it ships, not serving as the first stage retriever over millions of documents.

Test

npm install
npm test   # 113 tests: retained energy, bound containment, cross width refusal, threshold transfer

The suite includes the naive implementations as reference functions, so the tests show plain truncation reordering documents, zero padding deflating short vectors by a fixed factor, and a carried over threshold changing the accept rate, each of them returning an ordinary finite number while it happens.

License

MIT