Skip to content

18 min read

Batches of consecutive chunks train the model to reject relevant text

On this page

The training curve was the healthiest one the team had produced: loss from 2.94 down to 0.31 over three epochs, no spikes, no divergence, evaluation loss tracking training loss within 0.04. Recall@50 on the frozen real-query set went from 0.61 before training to 0.55 after. The model had learned something thoroughly and what it learned made retrieval worse.

The vocabulary for what went wrong is short. An in-batch negative is a passage the loss treats as a wrong answer for a given query purely because it was paired with a different query in the same batch, and a false negative is one of those that in fact answers the query perfectly well. Nothing in the objective can tell the two apart, since the only evidence it has about relevance is which row of the file a passage arrived on. So the composition of a batch is part of the objective, and a loader that decides composition by convenience is deciding what the model is taught.

The batch that settled it

Dumping a single batch settled it in four minutes. Of 32 rows, 14 came from three documents — five consecutive chunks of the thermal management section of one service manual, five from a firmware release note, four from a warranty policy. The generated questions from chapter 3 had been written one per chunk and appended to a JSONL file in corpus order, the loader was reading it with shuffling limited to a 1,000-row buffer, and 1,000 consecutive rows of that file cover about forty documents.

text
batch[0]  q: "why does the cabinet throttle when it is warm"
          p: chunk 41182  doc 7734 "Thermal derate engages at 71 °C intake..."
batch[1]  q: "how hot can the intake get before it slows"
          p: chunk 41183  doc 7734 "Derate is released after 120 s below 66 °C..."
batch[2]  q: "what temperature stops the throttling"
          p: chunk 41184  doc 7734 "Hysteresis prevents oscillation between..."
batch[3]  q: "invoice arrived twice this month"
          p: chunk 88010  doc 9912 "Duplicate invoices are issued when..."
...

Rows 0, 1 and 2 are three questions about the same subject with three passages that all answer all three. Under the loss from the previous chapter, row 0's query is pushed towards chunk 41182 and away from 41183 and 41184, which both genuinely answer it. Row 1's query is pushed away from 41182 and 41184. The gradient is not noisy; it is pointed precisely backwards, and it is strongest exactly where the model was closest to being right, because the softmax concentrates its gradient on the highest-scoring off-diagonal entry and here that entry is a correct answer.

That last clause is the whole mechanism and it is worth taking in four steps, because each one is checkable and skipping any of them leaves a reader who can recite the rule and cannot predict the next case. First, the softmax over a row gives each candidate a share of the row's probability, and those shares are set by the gaps between scores. Second, a candidate that scores close to the correct passage takes a large share, and the loss reduces it by pushing that candidate's score down — so the size of the correction a candidate receives rises with how well it scored. Third, in a same-document batch the highest-scoring off-diagonal entry is a chunk of the same document, which is the most similar text in the corpus to the passage that actually answers the query. Fourth, therefore the largest single correction in the step is applied to a passage that should have been ranked high, in the direction of ranking it lower. The loss is not diluted by the collision. It is aimed by it.

Nothing about that arrangement required carelessness. Corpus-ordered pair files are the default output of every reasonable pipeline: a generation job that iterates documents and writes a row per chunk, a SQL export with an ORDER BY doc_id that somebody added to make the file diffable, a streaming dataset that yields shards in the order they were written, a parallel worker pool where each worker owns a document and appends its questions contiguously. Ordinary supervised training tolerates all of these because a batch of thirty-two images with correlated labels is a mild statistical inconvenience. Here the batch composition is part of the objective, so an ordering that would be harmless in a classifier is a corruption of the loss.

The difference between those two cases is one property, and naming it is what makes the rule transferable. In a classifier, each example carries its own label and the batch is only a device for averaging gradients, so correlation between rows costs some efficiency and nothing else. Here, each example's label is defined by the other rows. Correlation between rows is therefore not a statistical inconvenience but a change to what the labels say, and no amount of averaging recovers a label that was wrong before the average was taken.

Why the loss curve stays beautiful

The loss measures agreement with the labels it was given, and the labels say chunk 41183 is wrong for this query. The model can satisfy that instruction and does. Driving the loss to 0.31 means it has learned to separate three near-identical passages about thermal derate from one another and to attach each query only to its own — a genuinely hard discrimination that the training set rewarded and that no user has ever wanted. Evaluation loss agrees because the held-out split was drawn from the same file with the same ordering, so it contains the same structure.

This is the chapter's stumble, and it is a belief that a defence works rather than a slip. A held-out split is the standard protection against a model memorising its training rows, and it does protect against that. It cannot protect against a training set whose labels are wrong, because the held-out rows were manufactured by the same procedure and carry the same wrong labels. The quantity being watched — evaluation loss — stays healthy across the whole range of the dangerous property, from a file in strict corpus order to a globally shuffled one, because it is computed inside the same construction either way. A measurement that does not move when the fault moves is not a weak instrument. It is not an instrument.

The metric that catches this is the recall on the frozen real-query set, and the reason it is the only one that catches it is that it is the only signal in the system that was not produced by the same pipeline. That is a general property worth carrying: a corrupted training set corrupts every metric computed from data that shares its construction, and the only immune measurement is one built by a different process, which in practice means the hand-collected set from chapter 3.

The damage also concentrates in the worst possible place. The chunks that end up in a batch together under document-ordered loading are the chunks of one document, which are by construction the most semantically similar passages in your entire corpus. Those neighbourhoods are the ones where fine distinctions matter — which of five thermal chunks answers "what releases the derate" — and a fine-tune that has been trained to scatter them is worse there than the base model, while being unchanged in the regions the batches never touched.

Audit the batch, not the sampler configuration

The check is a loop over the dataloader that reports two things: how often two rows in a batch come from the same source document, and how similar the in-batch negatives actually are under the current model.

python
import torch
from collections import Counter

def audit_batches(loader, model, doc_of, n_batches=25):
    same_doc, high_sim, total_neg = 0, 0, 0
    hist = Counter()
    for i, batch in enumerate(loader):
        if i >= n_batches:
            break
        docs = [doc_of[cid] for cid in batch["chunk_id"]]
        for d, c in Counter(docs).items():
            if c > 1:
                same_doc += c - 1
                hist[c] += 1
        with torch.no_grad():
            q = model.encode(batch["anchor"], convert_to_tensor=True,
                             normalize_embeddings=True)
            p = model.encode(batch["positive"], convert_to_tensor=True,
                             normalize_embeddings=True)
        s = q @ p.T
        pos = s.diag().unsqueeze(1)
        off = s - torch.eye(len(s), device=s.device) * 1e4
        # a negative scoring within 0.05 of the positive is a false-negative suspect
        high_sim += int((off > pos - 0.05).sum())
        total_neg += len(s) * (len(s) - 1)
    return {"same_doc_collisions_per_batch": same_doc / n_batches,
            "suspect_negative_rate": high_sim / total_neg,
            "collision_sizes": dict(hist)}

The collision count is defined by one line of that function and the definition is worth reading rather than assuming. For a document appearing c times in a batch it adds c - 1, not c, because one row per document is the intended state and costs nothing; only the extra copies create comparisons between passages of a single source. Applied to the dumped batch, three documents contributing five, five and four rows give four plus four plus three, which is eleven collisions in that one batch — close to the mean reported below, as it should be for a batch drawn from the same misconfigured loader.

Run over the failing configuration this reported 12.4 same-document collisions per 32-row batch and a suspect-negative rate of 0.041, meaning about 4% of all in-batch negatives scored within 0.05 of the true positive. Four percent sounds small until it is multiplied through: at batch 32 that is roughly 40 wrongly-signed comparisons per batch, every step, aimed at the tightest part of the space.

That multiplication is worth doing rather than accepting, because the count it produces is larger than the rate suggests for a structural reason. Each of the 32 rows is compared against the 31 passages that are not its own, so the batch holds thirty-two times thirty-one comparisons — a little under a thousand of them, from only thirty-two examples. A rate of 0.041 applied to that count gives the 40. In general a batch of B rows supplies B times B−1 comparisons, so the comparisons grow with the square of the batch while the examples that produce them grow linearly. A false-negative rate that is flat in percentage terms therefore produces a wrong-comparison count that grows quadratically, which is why a rate small enough to ignore at a batch of eight is not ignorable at a batch of several hundred.

Sampling strategySame-doc collisions per 32-row batchSuspect negativesrecall@50 after 3 epochs
File order, 1k shuffle buffer12.44.1%0.55
Global shuffle0.90.7%0.74
Global shuffle + max 1 chunk per doc0.00.5%0.77
Global shuffle + max 1 per doc + version grouping0.00.2%0.79

The first row and the last differ by 24 recall points on the same data, the same model and the same loss. None of that is hyperparameter tuning.

Read down the recall column as gaps and the twenty-four points divide very unevenly. Global shuffling alone accounts for nineteen of them. Capping same-document chunks adds three more. Version grouping adds the last two. The section title below states that ordering as advice; the table is where it is earned, and it is the reason to spend an afternoon on a shuffle before spending a week on a grouping key.

The first two columns are worth reading as a pair, because they are two instruments with no step in common. The collision column is a structural count: it loads no model, computes no similarity, and merely asks whether two rows carry the same document id. The suspect column is a score-based measurement that encodes every anchor and passage with the current model and compares similarities. One never looks at text, the other never looks at metadata, and they fall together monotonically down all four rows. That agreement is what licenses treating the document ordering as the cause of the wrong gradients rather than as something that merely travels alongside them. Either column alone would leave the causal claim unsupported.

The same document grouping has to govern the train and evaluation split as well, for a related reason with a different consequence. Splitting the pair file randomly puts chunk 41182 in training and chunk 41183 in evaluation, and since a question about one is very nearly a question about the other, the evaluation is measuring memorisation. Split by document — or by whatever grouping key the near-duplicate analysis produced — so that no document contributes rows to both sides. On this dataset a random split reported held-out recall@10 of 0.91 and a document-grouped split of the identical model reported 0.68, and the second number was the one that matched behaviour on the frozen real-query set.

Twenty-three points between those two numbers, on one model with one set of weights, is worth pausing on. Nothing about the model differed between the two readings. The only thing that changed was which rows were allowed to appear on each side of the split, so the difference is a property of the measurement procedure and not of the thing measured. That is the same defect as the healthy evaluation loss above, arriving through a different door: a number computed from data that shares a construction with the training set reports on the construction.

Gradient accumulation is the trap that catches people who have understood everything above and reach for the obvious way to get a big batch on a small card. Accumulating gradients over eight micro-batches of 32 gives the optimiser the statistics of a batch of 256 and gives the loss the negatives of a batch of 32, because each micro-batch computes its own 32×32 similarity matrix and never sees the other seven. The optimiser step is smoother and the training signal is unchanged, so the recall stays where the 32-row configuration put it while the logs claim an effective batch size of 256. This is why CachedMultipleNegativesRankingLoss exists as a separate mechanism: it caches embeddings so the similarity matrix spans the full batch, which is the property accumulation does not provide.

The reason accumulation cannot deliver the negatives follows from where the matrix is built. Similarities are computed inside a forward pass, from the embeddings that pass produced, and accumulation works by summing the gradients of several separate forward passes before stepping the optimiser. No comparison can therefore span two micro-batches, because at no point do two micro-batches' worth of embeddings exist together to be multiplied. Gradient caching changes exactly that: it holds the embeddings from every sub-batch, forms one matrix across all of them, and only then re-runs the sub-batches to attribute gradients. The distinction is between accumulating results and accumulating inputs, and only the second gives the loss a wider pool.

The three fixes, in order of how much they buy

Global shuffling is the first and does most of the work. Shuffle the pair file once, on disk, before training, rather than relying on a streaming buffer whose window is smaller than a document's chunk count. If the dataset is small enough to fit in memory — and 6,200 or even 200,000 pairs is — load it fully and permute it, which removes the question entirely.

Capping same-document chunks is the second, and it needs a group-aware sampler because shuffling makes collisions rare rather than impossible. The reason they survive shuffling is a counting one. In a 200,000-chunk corpus a document holding 400 chunks owns one row in five hundred, so a batch of a couple of hundred rows contains, on average, a fraction of one of them — which is small, and is not zero, and does not become zero however well the file is shuffled. With a 200,000-chunk corpus and 400 chunks in a large document, a batch of 256 will contain two chunks of that document about a fifth of the time.

python
import random
from collections import defaultdict

def grouped_batches(rows, batch_size, group_key, max_per_group=1, seed=0):
    """rows: list of dicts. group_key(row) -> document or version-family id."""
    rng = random.Random(seed)
    pool = rows[:]
    rng.shuffle(pool)
    batch, counts, deferred = [], defaultdict(int), []
    for row in pool:
        g = group_key(row)
        if counts[g] >= max_per_group:
            deferred.append(row)
            continue
        batch.append(row)
        counts[g] += 1
        if len(batch) == batch_size:
            yield batch
            batch, counts = [], defaultdict(int)
            pool.extend(deferred)   # retry deferred rows in later batches
            deferred = []
    if len(batch) > 1:
        yield batch

Sentence Transformers' batch_sampler="no_duplicates" handles a narrower version of this by refusing to place two rows with identical text in one batch, which catches boilerplate — the legal footer that got chunked separately in 4,000 documents — and does not catch three different chunks of one manual. The two mechanisms are complementary and both are worth having on.

The grouping key is the third fix and the one that requires thought about your corpus. Documents are the obvious group, and they are not always the right one. A manual that exists as v3.1, v3.2 and v3.2-de is three documents whose chunks are near-identical; a product family with twelve SKUs may have twelve datasheets that differ in four numbers. Grouping by document leaves those collisions intact, and they behave exactly like same-document collisions. The workable key is usually derived rather than declared: cluster the corpus once with the base model, or hash normalised text shingles — the overlapping runs of consecutive words a document is cut into, so that two texts sharing most of their shingles are near-copies — to detect near-duplicates, and group by the resulting cluster id. On a versioned technical corpus this moved the suspect-negative rate from 0.5% to 0.2% and bought the last two recall points in the table.

Which key is right is a design decision and not a requirement, and the two are worth separating. The requirement is that no two passages in a batch answer each other's queries. Document identity is one proxy for that, chosen because it is free — every row already carries a document id. A near-duplicate cluster is a second proxy, chosen because it also catches translations and versions, and it costs one clustering pass over the corpus. Neither proxy is the requirement itself, which is why the suspect-negative rate never reaches zero in the table: the last column measures the requirement, and the first measures the proxy.

A corpus that is genuinely about one narrow subject puts a floor under all three fixes, and it is worth recognising rather than fighting. If every document describes one product line, a batch of 512 random chunks contains many passages that plausibly answer many queries, and no sampler removes that because the collisions are semantic rather than structural. The practical response is to cap batch size where the suspect-negative rate starts climbing — measure it at 64, 128, 256 and 512 with the audit above and stop one step before the curve bends — and to spend the remaining budget on negatives chosen for their content rather than on more random ones. Multi-tenant corpora behave the same way for a different reason: forty customers with the same onboarding guide under forty tenant ids produce forty near-identical chunks that no document key groups together, and there the grouping key has to be content-derived.

Directed noise, and the residue no sampler removes

This is not label noise in the usual sense. Random label noise degrades a model gradually and roughly symmetrically, and can be partly absorbed by lowering the learning rate or training for fewer steps. This noise is directed — it always pushes apart things that belong together, never the reverse — so more training makes it worse and a lower learning rate only slows the descent into the same place. There is no optimiser setting that fixes it.

The asymmetry is what makes the usual remedies inapplicable, and it can be stated without any arithmetic. Random errors point in every direction, so averaging many of them leaves something close to nothing, and the standard mitigations work by buying more averaging: smaller steps, fewer of them, more examples per step. A signed error survives averaging exactly because it has a sign. Averaging a thousand pushes that all point the same way gives one push that points that way, scaled down. Every knob that trades speed for stability therefore slows the arrival at the wrong answer without changing which answer it is.

It is also not fully eliminable, and the honest version of the fix says so. Once batch sizes reach several hundred over a corpus about one subject, some fraction of in-batch negatives will be relevant to some query no matter how the batch was assembled, because your corpus contains many passages that answer "how do I reset the controller". At batch 1,024 over a homogeneous 200,000-document corpus, a hand-check of 500 sampled negatives found 1.1% that a judge called relevant. That residue is one reason the returns to batch size flatten where they do, and it is why the score-based suspect filter in the audit above is worth keeping in the training loop rather than only in diagnostics — masking out any in-batch negative that scores above the positive minus a margin costs one line and removes the worst of the remainder.

That residue is also the clean statement of what a sampler can and cannot do. A sampler decides batch membership from metadata, so it can remove every collision that has a structural signature — same document, same version family, same normalised text. It has no access to meaning, so a passage from an unrelated document that happens to answer the query is invisible to it, and remains invisible however the grouping key is refined. The score-based mask reaches those because it consults the model rather than the metadata, and it is the only one of the two mechanisms whose coverage improves as the model does.

The three fixes are cheap, they are checkable in a loop over the dataloader, and between them they moved this run's recall@50 by twenty-four points without touching the model, the data or the loss. What remains after them is the semantic residue, and the ordinary reason a random batch teaches so little.

Having stopped the batch from poisoning the gradient, the batch is still filled at random, which means most of what the loss pushes against is a passage about invoices when the query was about thermal derate. Separating those two was never in doubt. The comparisons that would actually teach the model something are the ones it currently gets wrong with confidence, and putting those in the batch on purpose is a different piece of machinery.

Adapting the retriever to your domain — read the current version at https://daviddominguez.dev/learn/en/adapting-the-retriever-to-your-domain/batches-of-consecutive-chunks-train-the-model-to-reject-relevant-text

Back to top