LM7-T4 — Masking the Loss

Learning objectives

After this lesson, you should be able to:

• Construct input and target sequences shifted by one with −100 over the prompt and padding.
• Compute the masked mean loss and explain the normalisation.
• Verify with PyTorch that ignored positions get exactly zero gradient.
• Choose between token-level and sequence-level averaging and describe the effect of partial prompt weights.

The idea: read everything, learn only from the answer

The model processes the whole sequence in one forward pass. It produces, at every position, a prediction of the next token. The loss is a sum over positions of the cost of each prediction. Masking removes some positions from that sum. The model still reads the prompt (it is part of the input and shapes every later hidden state) but is not scored on predicting it.

The shift

Position i predicts the token at position i + 1. So the target sequence is the input sequence moved one place to the left. A position is counted if its target is a response or end token. This means that the last prompt token is counted (its target is the first response token), but the first prompt token is not even a target of anything. Off-by-one errors here are the most common bug in instruction-tuning code: if you mask the labels but forget to shift, the model is trained to predict the wrong token at every counted position.

Ignore index

PyTorch’s cross-entropy takes ignore_index, by default −100: a position whose target equals it contributes nothing to the loss or the gradient. We put −100 in the labels for every position we want to skip. Because −100 can never be a token id, there is no ambiguity.

Padding

Padding positions are always ignored. They are artefacts of batching; predicting them teaches nothing and wastes capacity. They are masked in the labels; in a causal model they sit after the real tokens, so they cannot affect real positions, and no attention mask is needed for correctness.

The end token

The end-of-turn token is part of the response and is counted. If it is masked out, the model never learns to emit it and generation runs on to the length limit (Lesson LM7-T3).

The average

The loss is the mean over counted positions. Two ways to take it differ when examples have different response lengths: a token-level mean over all counted positions in the batch weights every token equally, so long responses count more; a sequence-level mean first averages within each example, then across the batch, so every example counts equally. PyTorch’s default with ignore_index is token-level. For instruction data with very different response lengths it is a design choice,. Token-level is simple and stable; sequence-level prevents long answers from dominating.

The mathematics: the masked mean

For one sequence with input tokens s₀, …, s_n−1 padded to length n, define the labels ℓ_t = s_t for response and end positions and ℓ_t = −100 elsewhere. The model input is s_0…n−2 and the target is ℓ_1…n−1. With the per-position loss c_t = −ln p(ℓ_t+1 | s_0…t), the loss over a batch is

1[·]

indicator: 1 if the condition holds, else 0.

ℓ_b,t+1

the (shifted) target of position t in sequence b.

c_b,t

the cross-entropy cost at that position.

denominator

the number of counted positions in the whole batch; it varies from batch to batch.

The gradient of ignored positions is exactly zero: they enter neither the sum nor the count. If no position in the batch is counted, the expression is 0/0 and PyTorch returns NaN; guard against batches with no response (for example, an example truncated before its response starts).

Numbers. A sequence has 12 prompt tokens with mean loss 3.0, 6 response tokens with mean loss 0.8, 1 end token with loss 2.6, and 4 padding tokens with loss 5.0. The unmasked mean over 23 positions is (36 + 4.8 + 2.6 + 20)/23 = 2.757; the masked mean over the 7 counted positions is (4.8 + 2.6)/7 = 1.057. A drop from 2.757 to 1.057 is purely a change of what is averaged, not an improvement of the model.

Token versus sequence average. Two sequences have 2 and 20 counted tokens and mean losses 1.0 and 2.0. Token-level mean: (2·1.0 + 20·2.0)/22 = 1.909. Sequence-level mean: (1.0 + 2.0)/2 = 1.5. The long example dominates the first by weight 20/22.

Partial weights. A middle path gives prompt positions a small weight w (for example 0.1) instead of 0; the numerator adds w · Σ_prompt c_t and the denominator adds w times the number of prompt positions. Some studies find a small benefit for short responses; the standard recipe uses w = 0.

A worked example: build, mask, verify the gradients

The code builds padded labels for one sequence, computes the masked loss with ignore_index, checks it against the manual formula, and verifies that ignored positions receive no gradient at all.

import torch
import torch.nn.functional as F

torch.manual_seed(0)
V, END, PAD = 20, 2, 0
prompt, response = [5, 6, 7, 8, 9], [11, 12, 13]
n = 10
seq = prompt + response + [END]
labels = [-100] * len(prompt) + response + [END]
seq += [PAD] * (n - len(seq)); labels += [-100] * (n - len(labels))
inputs = torch.tensor(seq[:-1])                 # model input: positions 0 .. n-2
targets = torch.tensor(labels[1:])              # what each position must predict
print(seq)                                      # [5, 6, 7, 8, 9, 11, 12, 13, 2, 0]
print(targets.tolist())                         # [-100, -100, -100, -100, 11, 12, 13, 2, -100]

logits = torch.randn(1, n - 1, V, requires_grad=True)      # stand-in for model(inputs)
loss = F.cross_entropy(logits.flatten(0, 1), targets.flatten(), ignore_index=-100)

counted = targets != -100
manual = F.cross_entropy(logits[0][counted], targets[counted])
print(int(counted.sum()), torch.allclose(loss, manual))    # 4 True

loss.backward()
grad = logits.grad[0]
print(bool((grad[~counted] == 0).all()), bool((grad[counted] != 0).any()))   # True True: ignored positions get no gradient

# a batch with nothing to learn from gives NaN, so guard against it
none = torch.full((1, 9), -100)
print(F.cross_entropy(logits.detach().flatten(0, 1), none.flatten(), ignore_index=-100).isnan().item())   # True

# token-level versus sequence-level averaging
tok = (2 * 1.0 + 20 * 2.0) / 22
seqm = (1.0 + 2.0) / 2
print(round(tok, 3), seqm)                       # 1.909 1.5

The first print shows the shift: the first four targets are −100 (they would predict prompt tokens) and the fifth is 11, the first response token, predicted from the last prompt token 9. The count is 4 because the target list contains three response tokens and the end token.

Engineering implications and limitations

Write a unit test for the labels. For a known example, assert the exact labels (the one in this lesson is a good fixture). Most bugs are visible there: no shift, shift in the wrong direction, prompt not masked, end token missing, padding counted.

Mask in the data pipeline. Build labels when you build the example, where the prompt length is known, not later from the token ids (a response token can equal a prompt token).

Use the same ignore value everywhere. −100 is PyTorch’s default; a custom value needs passing to the loss everywhere it is used, including evaluation.

Truncation can erase the response. If an example is longer than the context and the end is cut off, the labels may have no counted position or no end token. Drop such examples or truncate the prompt instead of the response.

Report the counted-token share. Log the fraction of positions that carry the loss. A sudden change signals a data or template bug.

Evaluate with the same mask. Validation loss should be computed on response tokens only, so that it measures what training optimises. A validation loss computed over prompts too is not comparable.

Gradient accumulation and token counts. With accumulation, divide by the total number of counted tokens across the micro-batches, not by the number of micro-batches, if you want an exact token-level mean.

For robots. When the response is a structured command, you may want the loss only on the fields that matter (the action and its parameters) and not on boilerplate. The same mechanism gives you that: set the labels of the boilerplate to −100.

Common misconceptions

Engineering takeaways

1. Inputs are tokens 0…n−2, targets are labels 1…n−1; labels are −100 for prompt and padding, token ids for response and end token.
2. Masked mean: sum of costs at counted positions divided by their number; ignored positions get exactly zero gradient.
3. The first counted position is the last prompt token, which predicts the first response token.
4. Token-level and sequence-level means differ when lengths vary (1.909 versus 1.5 in the example); PyTorch’s default is token-level.
5. Guard against batches with no counted positions (NaN), test the labels on a fixed example, and evaluate with the same mask.

References and further reading

See Chapter 7 of the book (creating data loaders for an instruction dataset, where the padding tokens are replaced by the ignore index −100, and the discussion of masking the instruction in the loss).

• PyTorch documentation: torch.nn.functional.cross_entropy (the ignore_index argument).
• Build a Large Language Model (From Scratch), Chapter 7: Finetuning to Follow Instructions (primary book reference)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

Next: evaluating whether the tuned model follows instructions, and what contamination does to the answer.

Related content

Nothing below is required to finish this lesson.

• Before
  LM7-T3: Special Tokens and Conversation Templates

  Use role markers and end-of-turn tokens so the model knows where to stop.

• Lab
  LM7-L2: Implement Loss Masking

  Build the label sequence with ignore-index masks and compute the masked loss.