MaxSim keeps the per-term evidence that pooling averages into one number

MaxSim is the scoring function a late-interaction retriever uses in place of a single dot product. Both sides of the comparison arrive as sets of L2-normalised 128-dimension token vectors — one per word or word-piece, all rescaled to unit length so that a dot product between any two of them is the cosine of the angle between them. The question it answers is how well a document's tokens, taken individually, satisfy the query's tokens, taken individually. That is a different question from the one a pooled vector answers, which is how close the average of everything the document says is to the average of everything the query asks.

The gap between those two questions is not academic, and it is worth being concrete about where it opens. Consider a support corpus — the whole set of documents an organisation has embedded and stored — holding a 900-token API reference for multipart uploads, containing a single clause explaining that EntityTooSmall is returned when any non-final part is under 5 MiB, and a 400-token best-practices page about retrying multipart uploads, thorough on backoff and idempotency and silent on that error code. A user types EntityTooSmall error retrying multipart upload. Only one of those two documents can answer them, and only one scoring function reliably says so.

Three query pieces, worked in full#

Start with three query tokens rather than the whole query, because the entire operation is visible at that size. The tokeniser splits EntityTooSmall into three word-pieces, entity, ##too and ##small, the double hash marking a piece that continues the previous one rather than starting a new word.

Take the first of them against the API reference. MaxSim computes a similarity between the vector for entity and the vector for every one of the reference's 900 stored tokens, which is a row of 900 numbers. Then it takes the largest one. The largest is 0.94, and it comes from the reference's own literal entity token. Every other entry in that row — all the rest of them, covering everything else the document says about parts, uploads, retries and HTTP status codes — is smaller, and not one of them contributes anything at all to the score. That is the operation: nine hundred comparisons, one survivor.

Doing the same for ##too gives 0.91 against the reference's ##too, and ##small gives 0.93 against its ##small. Those three maxima add to a little under two and eight tenths.

Now the best-practices page, which contains no EntityTooSmall anywhere. Each of the three rows is shorter, since the page has 400 tokens rather than 900, and each row's largest entry is a near-miss rather than a hit: entity finds request at 0.31, ##too finds too at 0.44, and ##small finds smaller at 0.52. Those three maxima add to a little over one and a quarter. The difference between the two sums is 1.51, and it was produced by three rows of a matrix out of a query that has nine content tokens in it.

Two design choices are doing the work there, and each is the right one for a reason that can be stated.

The maximum, rather than an average, is taken over document tokens because a query term is satisfied by its single best piece of evidence. If a document says EntityTooSmall once, it says it; a second mention adds no information that the first did not already provide, and an average over document tokens would make the second mention raise the score while a longer document that says it once would be punished for its other 890 tokens, every one of them irrelevant to the term under consideration. The maximum makes the score for a query term independent of how much else the document contains.

The sum, rather than an average, is taken over query tokens because each query token is a separate requirement. A user asking for an error code and a procedure is asking for both, and a document that satisfies one and misses the other should lose exactly the amount it missed by — once. Summing independent maxima does that: a strong match on EntityTooSmall cannot be pulled down by a weaker match on retrying, because the two terms never meet. Pooling gets both of these wrong in the same operation, because it collapses the document's 900 tokens into one point before the query is ever seen, and it compares that point to a query that has also been collapsed.

Scoring a pair, one query token at a time#

The general form adds nothing to the three-row case except more rows, and the whole function is three lines.

import torch

def maxsim(q: torch.Tensor, d: torch.Tensor) -> torch.Tensor:
    """q: (nq, 128) query token vectors. d: (nd, 128) document token vectors.
    Both L2-normalised, so the dot product is a cosine."""
    sims = q @ d.T                       # (nq, nd): every query/document token pair
    return sims.max(dim=1).values.sum()  # best match per query token, summed

The (nq, nd) matrix is the object to hold in mind: one row per query token, one column per document token, every entry a cosine. max(dim=1) collapses each row to its best entry, and .sum() adds the surviving column of maxima.

Scoring a batch of candidates needs one more detail, and it is the detail that most first implementations get wrong. Documents have different token counts, so the batch has to be padded, and padding vectors participate in the maximum unless they are excluded explicitly.

def maxsim_batch(q, D, mask):
    """q: (nq, dim). D: (b, nd, dim) right-padded. mask: (b, nd) bool, True = real token."""
    sims = torch.einsum("qd,bnd->bqn", q, D)          # (b, nq, nd)
    sims = sims.masked_fill(~mask[:, None, :], -1.0)  # padding must never win a max
    return sims.max(dim=2).values.sum(dim=1)          # (b,) one score per document

Leaving out the masked_fill produces a system that runs, returns plausible rankings, and quietly favours whichever documents were padded most, since a zero vector scores 0.0 against every query token while genuinely unrelated real tokens often score slightly negative. The bug is invisible in a unit test on equal-length inputs and shows up as a mild, unexplained preference for short documents in production.

Running the two documents through the real thing shows where the margin lives. The query tokenises to nine content tokens plus the [Q] marker, and is padded to query_maxlen=32 with [MASK] positions, all of which are scored.

Before reading the margin, check that the table accounts for everything. Nine content rows plus 22 mask positions plus the [Q] marker is exactly 32, which is query_maxlen, so there is no fourth category of query position sitting outside the table. Each column then adds without a residual: 7.92 and 13.64 give 21.56, and 6.48 and 14.52 give 21.00. Nothing has been rounded away and nothing unlisted is contributing, which is what licenses reading the margin as a decomposition rather than as a summary.

The reference wins by 0.56 out of roughly 21.5, and that margin comes apart into three pieces that close on it exactly. The three subword pieces of EntityTooSmall contribute 1.51 in the reference's favour, worked out above. The remaining six generic content tokens run the other way by seven hundredths in total, that being the best-practices page's stronger showing on retry, ##ing, multi, ##part and upload, partly given back by error going the reference's way. And the mask block runs the other way by a little under nine tenths. Take the second and third away from the first and you are left with 0.56, with nothing unexplained. The identifier is the only evidence that distinguishes the two documents, and MaxSim gives it three independent slots to speak in rather than one ninth of a share in an average.

Two things in that table are worth carrying forward. The [MASK] positions contribute about two thirds of the total score for both documents, which is why a raw MaxSim value is not interpretable on its own and why a score threshold tuned on one query set will not transfer to another; scores are only comparable within a single query. And the margin between a good answer and a mediocre one is small in absolute terms, which means that compression schemes that perturb the stored vectors slightly are perturbing exactly the quantity the ranking depends on.

What the average could not have kept#

The dilution is easier to see as arithmetic than as intuition. Mean pooling produces the centroid of the document's token vectors, the centroid being the average position of a set of points and therefore the place a document lands once every token has had an equal vote, so a token's influence on the final embedding is 1/n. In the 900-token reference, the clause about EntityTooSmall occupies about 20 tokens, giving it roughly 2.2% of the centroid's mass; whatever direction those tokens point in, they move the document's point by a rounding error. The best-practices page has no such clause, and its 400 tokens all point in broadly the same direction, so its centroid is a tight, confident summary of "retrying multipart uploads" — and a query that is 70% about retrying multipart uploads lands nearer to it.

Set that 2.2% beside the MaxSim table and the contrast is exact. Under pooling, the decisive clause gets a share of the document's representation proportional to its length, and it is short, so it gets almost none. Under MaxSim it gets three rows out of nine content rows, because the share is set by how many query tokens it answers rather than by how much of the document it occupies. The two functions weight the same clause by two different quantities, and only one of those quantities has anything to do with the query.

This is the shape of the failure worth remembering, because it is the answer to the question of which queries will show a difference. A pooled vector rewards documents whose overall topic matches the query's overall topic, and it has no mechanism for representing that a document contains one specific fact. Queries that carry a single constraint, or several constraints that all point the same way, are served perfectly well by a centroid. Queries that carry two independent constraints are the ones where the centroid has to sit somewhere between them, and the document nearest that midpoint tends to be one that is moderately about both and precisely about neither.

The summation over query tokens rather than an average is a deliberate choice with a consequence. A sum means every query token, including the [MASK] padding, adds a positive quantity, so scores grow with query length and cannot be compared across queries of different lengths; the fixed query_maxlen makes the length constant and hides the issue rather than solving it. An average would make scores comparable and would also make a query's single decisive token worth 1/32 of the result, which is the dilution the whole design exists to avoid. The choice buys discrimination within a query at the cost of any global calibration.

None of this works with token vectors taken from an off-the-shelf encoder, and that is the belief most likely to cost somebody a wasted week. The reasoning that leads to it is reasonable on its face: any transformer emits per-token vectors, MaxSim is nine lines of tensor code, so late interaction looks like a scoring function you can bolt onto vectors you already know how to produce. What that reasoning leaves out is where the token vectors get their meaning. ColBERT is trained end to end through the MaxSim function itself, with a contrastive objective over positive and negative passages — an objective that pushes a query's vectors towards the passage known to answer it and away from passages that do not — so the gradients that shape each token's vector are the gradients of the maximum operation. That training is what teaches a token to be a good witness for the query terms it should answer and a poor witness for terms it should not, and it is why common tokens end up matching many things at moderate similarity instead of dominating. Running MaxSim over the raw per-token outputs of a general-purpose embedding model produces something that looks like late-interaction retrieval and ranks noticeably worse than the pooled version of the same model, which is a cheap experiment to get wrong and then draw the wrong conclusion from.

Concretely, the query classes where late interaction is worth expecting a gain look like this: an identifier plus a concept (ERR_STREAM_PREMATURE_CLOSE during graceful shutdown), a version or product qualifier plus a task (enable RLS on a view in Postgres 15), a negated or narrowing modifier plus a topic (websocket reconnect without losing subscription state), and anything where the answer occupies one clause of a document about something broader. The classes where it will not show a gain are single-topic chunks of a few hundred tokens, where the centroid already is the clause, and queries that are one coherent request with no second constraint to be diluted by.

A sum of maxima can be satisfied piecemeal#

The most common misreading is to treat MaxSim as a cheap cross-encoder. It is not, and the difference is visible in the code above: query tokens never attend to each other, document tokens never attend to the query, and the only operation is a dot product between two vectors that were computed in complete ignorance of one another. All of the compositional work that a cross-encoder does with cross-attention — resolving that "without" negates the clause after it, noticing that the number in the query has to be the number attached to that particular parameter, deciding that the query's subject is the document's object — is unavailable. Late interaction preserves per-term evidence, and per-term evidence is not the same thing as understanding the query.

The scoring function also has no notion of conjunction, which is the failure mode that surprises people who came to it expecting something keyword-like. A sum of independent maxima can be satisfied piecemeal: a 5,000-token changelog that mentions EntityTooSmall in one entry from March, retry in an unrelated entry from July and multipart in a third has a strong witness for every query token and will score well, despite no single passage in it answering anything. There is no proximity term, no requirement that the winning document tokens be near each other, and no penalty for a document that satisfies each constraint in a different context. Longer documents have more chances to win each maximum, so raw MaxSim carries a mild length bias that chunking tends to remove.

Neither is there any term weighting of the kind BM25 gives you for free, BM25 being the classical lexical ranking function that scores a document by which of the query's literal words it contains and how unusual they are. Every query token contributes on the same scale, so the and EntityTooSmall each own one slot in the sum, and a rare identifier is not up-weighted for being rare. The model can learn to make common tokens match many things at moderate similarity, which flattens their contribution across candidates, and that is a softer instrument than an explicit IDF — the inverse document frequency weight that makes a word appearing in twelve documents count for far more than one appearing in a million. The asymmetry runs one way only: the maximum is over document tokens for each query token, never the reverse, so a document is not rewarded for having every one of its tokens explained by the query. That is the right choice for retrieval, since documents are long and queries are not, and it is also why the function tells you nothing about whether a document is only about the query.

Everything above scores one query against one document, and the code takes two tensors that are already in memory. Retrieval requires choosing which documents get scored at all, and the collection that produced these two candidates holds five million documents and four billion token vectors, which no amount of tensor-level efficiency will let you sweep per query. Something has to select a few thousand documents before the function in this chapter ever runs.

That selection machinery is a retrieval system in its own right, with its own recall and its own failure modes, and the single knob inside it determines both the tail latency and the amount of work the rest of this topic asks you to measure.