Multiple-negatives ranking loss pulls one pair together and pushes the batch apart

Two training runs over the identical 6,200 pairs, identical base model, identical learning rate of 2e-5, identical three epochs. The first used a batch size of 16 and finished at recall@50 of 0.71 on the frozen real-query set. The second used 256 and finished at 0.79. Nothing else differed, and the second run took slightly less wall-clock time because the GPU was better utilised. Batch size is normally a throughput knob with a mild regularisation side effect, and here it moved the headline metric by eight points, which is more than most people get from changing the base model.

The reason is in the loss function. Each training example is a pair — a query and a passage that answers it — and a pair on its own cannot express an ordering, because ordering requires something to rank below. Multiple-negatives ranking loss manufactures that something out of the batch. For a batch of B pairs it treats the B passages as B candidate answers for every one of the B queries, and asks the model to identify, for each query, which of the B passages was its partner. The other B−1 passages are assumed irrelevant and pushed away.

The loss is cross-entropy over the batch#

Written out, the computation is a similarity matrix and one call to cross-entropy, with the labels being the diagonal. Cross-entropy is the standard loss for a pick-the-right-one-of-B question: it is small when the model concentrates its probability on the option the label names and grows as that probability falls.

import torch, torch.nn.functional as F

def mnrl(q_emb, p_emb, scale=20.0):
    """q_emb, p_emb: (B, d), L2-normalised. Row i of each is one pair."""
    scores = scale * (q_emb @ p_emb.T)          # (B, B) cosine similarities
    labels = torch.arange(len(scores), device=scores.device)
    return F.cross_entropy(scores, labels)      # correct "class" for row i is column i

One property of that computation decides almost everything that follows, and it is worth deriving before any numbers are read. A softmax over a row of scores divides the exponential of each score by the sum of the exponentials of all of them. Add the same constant to every score in the row and each exponential picks up the same multiplying factor — the exponential of a sum is a product of exponentials — so that factor appears once in the numerator and once in every term of the denominator, and it cancels. Cancelling is permitted because the factor is the same positive number in both places and cancelling a common positive factor from a fraction leaves the fraction unchanged. The softmax therefore depends only on the differences between scores, never on their levels. A row whose entries are all shifted upward is exactly the row it was before.

That single fact is why the matrix below has to be read as gaps rather than as levels, why the scale factor exists, and why the objective can never learn a threshold. Each of those is worked through in its place.

Take a batch of four and read the matrix. Suppose the queries are a thermal-derate question, a billing question, a firmware-rollback question and a cable-spec question, and the passages are their four answers. The model produces a 4×4 matrix of cosine similarities scaled by 20; the objective wants each row's diagonal entry to dominate that row's softmax, the normalisation that turns a row of raw scores into probabilities adding to one.

              p0(derate) p1(billing) p2(rollback) p3(cable)
q0 (derate)      12.4        3.1          9.8         2.2
q1 (billing)      2.7       15.1          1.9         0.8
q2 (rollback)     8.9        2.0         11.2         1.4
q3 (cable)        1.1        0.6          2.3        14.7

Row 0 has its diagonal at 12.4 and an off-diagonal at 9.8, which is a softmax probability of about 0.90 for the correct passage — a small loss and therefore a small gradient. Row 1 is nearly solved at 15.1 against a best distractor of 2.7. All the learning in this batch is concentrated in rows 0 and 2, the pair of thermal and firmware passages that genuinely resemble each other.

Reading the same four rows as gaps says why, and says it in a form that transfers to a matrix you have not seen. Row 1's diagonal of 15.1 stands 12.4 above its best distractor at 2.7. Row 3's diagonal of 14.7 stands the same 12.4 above its best distractor at 2.3. Those two rows sit at different levels and have identical gaps, and they are equally solved — the difference-only property above guarantees it, without any arithmetic on the individual entries. Row 2's diagonal of 11.2 stands 2.3 above its own best distractor at 8.9, and row 0's 12.4 stands a shade over two and a half points above the 9.8 beside it. Those are the two narrow rows, and they are narrow at quite different heights. The level of a row says how confidently the encoder scores that subject area; the gap says how much the loss still has to teach.

The gaps also give the softmax probabilities directly. A gap of two point three means the distractor's exponential is about a tenth of the diagonal's, so the diagonal takes roughly nine tenths of the row — which is the 0.90 the paragraph above quotes. Row 0's slightly wider gap pushes its own share a little past that, so the two narrow rows are near neighbours rather than the same number. The two wide rows, at a gap of 12.4, leave their distractors a few parts in a million and are solved to within rounding. That concentration is the mechanism worth holding onto: the gradient magnitude for an example is governed by how close its best distractor came, so the batch's contents decide how much any given pair teaches.

Two properties follow from the matrix being square and shared. Every column is used B times, once per row, so a batch of 256 yields 256 × 255 = 65,280 negative comparisons from 256 examples, which is why the objective is sample-efficient. And because the query encoder and passage encoder are usually the same weights in a fine-tune, the gradient moves both sides of every comparison at once; the derate query vector moves towards its passage, the passage moves towards the query, and both move away from the firmware pair. This is what "moving the vector space" means concretely. It is not a reordering layer bolted on top of fixed embeddings, in the way the rerankers of stage 8.14 are; it changes where every document in your corpus will land when you re-encode it, including documents that were never in the training set, because the change is in the encoder's parameters rather than in any stored vector.

The first of those two properties is worth stating in its general form, because the instance hides how fast it grows. B examples produce B rows, each row carrying B−1 off-diagonal entries, so the batch supplies B times B−1 comparisons. Comparisons therefore grow with the square of the batch while the examples that produce them grow linearly, and the ratio of comparisons to examples is simply B−1. That is the whole efficiency argument, and it is algebra rather than a measured result: it holds for any batch size, any corpus and any encoder. It is also, read the other way, the reason chapter 5's same-document collisions are dangerous out of proportion to how often they occur.

The change is measurable in the space itself rather than only in the metric. Before training, the thermal derate chunk and the firmware rollback chunk sat at cosine 0.81 — close, because they share register, product names and sentence structure, which is what a general-purpose encoder keys on when it has no reason to care about the difference. After three epochs they sat at 0.42, while the derate chunk and its own paraphrase in a different manual moved from 0.72 to 0.88. Nothing in the training data mentioned either of those pairs directly; the geometry moved because the parameters did. It is worth logging a handful of such pairs before and after every run, because they say what the model learned in a form a metric cannot: a fine-tune that separates everything from everything, mean pairwise similarity across a random sample of chunks collapsing from 0.55 to 0.15, has learned to spread the corpus rather than to organise it, and that shows up as instability on queries outside the training distribution.

A variant worth knowing is the symmetric form, which adds the transposed term — for each passage, identify its query among the batch's queries — and averages the two losses. It costs one extra cross-entropy call over a matrix you already computed and tends to help slightly when queries and passages are of similar length and register, which is the case for question-to-question or title-to-title data and much less so for a three-word query against a 300-word chunk. The asymmetric form in the code above is the default for retrieval for that reason.

The lineage is contrastive representation learning, where the same expression appears under the name InfoNCE and the "negatives" are other images in the batch rather than other passages. The scale factor of 20 is the reciprocal of a temperature of 0.05, and it exists because raw cosine similarities live in [−1, 1], a range over which the softmax is nearly flat and the gradient nearly zero; multiplying by 20 sharpens the distribution enough for the correct answer to be distinguishable. Sentence Transformers ships this default and it is rarely worth changing, though very large batches sometimes prefer a slightly lower scale because the softmax is already competing against more alternatives.

Why a multiplication is the right operation there follows from the difference-only property, and is worth spelling out because "sharpens the distribution" is an assertion until it is. Cosine similarity between two normalised vectors cannot leave the range from minus one to one, so the widest gap any row can present is two units, and a realistic gap between a correct passage and a plausible distractor is a small fraction of that. The softmax cares about gaps; multiplying every entry of the matrix by a constant multiplies every gap by that constant, and leaves every ordering untouched. So the scale factor is the one operation that changes what the loss can resolve without changing what it prefers. That is a design decision rather than a requirement: the requirement is that gaps be large enough for the exponential to separate them, and a fixed multiplier is one way of meeting it — the temperature is simply the same knob written as a divisor.

Batch size is the supply of negatives#

Once the batch is the negative pool, its size stops being an optimisation detail and becomes a modelling choice. A batch of 8 asks the model a multiple-choice question with 8 options; a batch of 512 asks the same question with 512, and getting it right at 512 requires a finer distinction than getting it right at 8. Since the distinctions that matter in your domain are fine ones — that is why you are fine-tuning — the number of options is close to being the resolution of the training signal.

The recall column is the one to read as gaps, and the gaps shrink as the batch grows. Doubling from 8 to 16 buys three points. Quadrupling from 16 to 64 buys five. Quadrupling again to 256 buys three. Quadrupling a third time to 1024 buys half a point. Each multiplication of the batch buys less than the one before it, which is the shape of a logarithm rather than a line, and it is the shape to expect: a batch of B poses a question with B options, and the extra discrimination demanded by going from B options to twice B is one further distinction however large B already was.

The VRAM column is not a second copy of that reading and should not be fitted to a line. Peak memory carries a fixed model footprint that does not depend on the batch at all, and the similarity matrix itself grows with the square of the batch while the activations grow with the batch, so the per-example marginal cost falls as the batch grows and then, eventually, stops falling. What the column is for is the practical point that the 256-row configuration does not fit on a card most teams have, and the two GradCache rows are how it is made to.

The flattening at the bottom is as important as the climb. Returns on batch size fall off once the pool is large enough that a random extra passage is almost never a plausible distractor, and where that happens depends on how homogeneous your corpus is: a corpus of 200,000 documents about one product line saturates earlier than a general web crawl, because random passages are already somewhat similar. Chasing batch size past the point of saturation buys nothing, and the more productive move at that point is to control which passages are in the pool, which is chapter 6.

When VRAM is the binding constraint, the technique to reach for is gradient caching, exposed in Sentence Transformers as CachedMultipleNegativesRankingLoss. It computes embeddings in small sub-batches without gradients, builds the full similarity matrix from the cached vectors, then re-runs the sub-batches to accumulate gradients. The result is numerically the same loss as the large batch at roughly the memory of the sub-batch, paying about 1.7× in step time. On a single 24 GB card that is the difference between a batch of 48 and a batch of 1,024.

The training loop itself is short#

The code that runs this is unremarkable, which is worth showing precisely because the impression that fine-tuning is a research project is what keeps teams from trying it.

from datasets import Dataset
from sentence_transformers import (SentenceTransformer, SentenceTransformerTrainer,
                                   SentenceTransformerTrainingArguments)
from sentence_transformers.losses import CachedMultipleNegativesRankingLoss

model = SentenceTransformer("BAAI/bge-base-en-v1.5")
train = Dataset.from_dict({"anchor": queries, "positive": passages})  # column order matters

args = SentenceTransformerTrainingArguments(
    output_dir="out/derate-v1",
    num_train_epochs=3,
    per_device_train_batch_size=256,
    learning_rate=2e-5,
    warmup_ratio=0.1,
    bf16=True,
    batch_sampler="no_duplicates",   # chapter 5 is about this line
)

SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=train,
    loss=CachedMultipleNegativesRankingLoss(model, mini_batch_size=32),
).train()

One thing that snippet does not show, and that costs teams a day when they discover it late, is the instruction prefix. Stage 8.6 established that many retrieval encoders were trained with a short string glued to the front of the text, different on the query side than on the document side; that string has to be applied during fine-tuning as well, on the same side, in the same form. Training the anchors without the query prefix and then serving with it produces vectors from a region the fine-tune never visited, and the symptom is a model that scores well in the training harness and poorly in production, which is a long way to travel for a missing eight-character string. The rule is that whatever function produces text for the index is the function that produces the positive column, and whatever produces text for a live query produces the anchor column, with no second implementation anywhere.

On 6,200 pairs at batch 256 that is 24 steps per epoch, 72 steps total, about eleven minutes on one A10G including evaluation. The Sentence Transformers loss documentation is worth reading once for the full catalogue, because the column layout of the dataset is what selects the behaviour: two columns give in-batch negatives only, and a third column named negative turns each row into a triplet whose explicit negative is added to the pool alongside the in-batch ones. Learning rates above about 5e-5 tend to degrade the base model's general competence quickly on datasets this size, and three epochs on a few thousand pairs is usually where the real-query metric peaks before overfitting starts.

Checkpoint selection matters more than the learning rate at this scale, because the window in which the model is better at your domain and not yet worse at everything else is narrow. Evaluating every 20 steps against the real-query set from chapter 3, rather than against training loss, on this run produced recall@50 of 0.74 at step 24, 0.79 at step 48, and 0.77 at step 72, so the final checkpoint was not the best one and would have been shipped by any script that saves at the end. Save on the evaluation metric, keep the base model's own score in the same table as a row, and treat any run whose best checkpoint is the last one as unfinished rather than as converged.

Those three step numbers are not three arbitrary probes, and noticing why makes the reading stronger. Twenty-four steps is exactly one epoch at this batch size, so the three measurements land on the three epoch boundaries with nothing in between them that mattered. The model improved across epoch one, improved again across epoch two, and gave some of it back across epoch three. A curve that turns over inside the run and is sampled only at the end reports the wrong side of the turn, and the alignment is what tells you the turn is a real feature of the schedule rather than noise between two adjacent evaluations.

The supervision the batch does not contain#

The objective never sees a label meaning "irrelevant". There is no supervision anywhere in this loss saying that passage 3 is a bad answer to query 0; there is only the structural assumption that a passage paired with a different query is not the answer to this one. That assumption is doing enormous work, and it is an assumption rather than a fact.

The chain that produces it is short enough to check link by link, and checking it is what makes the next chapter predictable rather than surprising. A positive pair asserts one relation and says nothing about ordering. A retrieval objective needs an ordering, because retrieval returns a list. An ordering needs something below the top of it. The only material in the training step that could serve is the other rows of the batch. So the negatives are manufactured from batch membership, and batch membership is decided by whatever assembled the file and shuffled it. The loss has no way to distinguish a passage that is in the batch because it is irrelevant from one that is in the batch because it was the next row in a JSONL file.

It also does not learn a threshold, and that is the difference-only property arriving at its consequence. The objective is entirely relative — it cares that the right passage outscores the others in its batch, not that the score is 0.8 rather than 0.4 — so a model fine-tuned this way gives you no calibrated notion of "similar enough", and code that filters on a fixed cosine cutoff will need that cutoff re-derived after every fine-tune. Any absolute score you had memorised from the base model is meaningless against the new one. A loss built from a softmax cannot see the level of a row; a model trained only by that loss is therefore never asked to put its scores anywhere in particular, and there is no mechanism by which it would arrive at a stable scale by accident.

The same property makes the training loss itself a poor number to compare across configurations, which is a stumble worth naming because the belief behind it is so nearly always right. A falling loss is the standard evidence that a model is learning, and here it is evidence about a question whose difficulty the batch size sets. A loss of a given value at a batch of 8 is the score on an eight-option question; the same value at a batch of 512 is the score on a 512-option one. The quantity stays fixed in name and units while the thing it measures moves underneath it, so two runs with different batch sizes can be ranked by their losses in the opposite order to their recall. The real-query metric is the comparison; the loss is diagnostic within one configuration only.

And it does nothing about a query whose answer is genuinely absent from the corpus, since every training row asserts that an answer exists and is exactly one passage. Retrieval systems that must sometimes return nothing need that behaviour built elsewhere, usually as a calibrated stage after retrieval.

The structural assumption is the fragile one, and it is fragile in a specific direction. Everything above holds only while the other B−1 passages in a batch really are irrelevant to each query. Because the gradient concentrates on the highest-scoring off-diagonal entry, a batch that contains a genuinely relevant passage aims the loss at exactly that entry, and the way batches are ordinarily assembled from a document corpus produces such passages constantly. The result is a gradient that is confidently wrong and a loss curve that looks entirely healthy while it happens.