← Concept library

NLP Foundations

Loss Masking in Fine-Tuning

Supervised fine-tuning uses the same shifted next-token loss as pretraining, computed over the exact same kind of sequence; the only change is that most of the sequence is excluded from the loss entirely.

intermediate · 6 min read

A supervised fine-tuning example is a full conversation: a system prompt, a user turn, an assistant turn, maybe several more turns after that. If you computed cross-entropy over every token in that sequence the way pretraining does (see causal-lm-pretraining-task), the model would be trained just as hard to predict the user's next token as the assistant's, which is both wasted signal (the model will never be asked to generate the user's side of a conversation) and actively counterproductive (it nudges the model toward imitating user phrasing, typos, and formatting rather than producing high-quality assistant responses).

Masking at the label, not the architecture

The fix does not change the loss function, the architecture, or the causal mask at all. It changes which positions are allowed to contribute to the loss. Every framework implementing this uses an ignore index, conventionally -100 in PyTorch and Hugging Face's CrossEntropyLoss, applied to the label tensor at positions that should not be trained on:

tokens: [ system... | user: "..." | assistant: "..." | <eos> ]
labels: [ -100, -100, ..., -100  | -100, -100, ...     | t1, t2, ..., tn, <eos> ]

Positions labelled -100 are skipped by the loss computation entirely; their gradient contribution is exactly zero, not merely small. Everything else about training is unchanged: the model still processes the full sequence in one causal forward pass, still predicts a next-token distribution at every position (see next-token-prediction-cross-entropy), it is simply that most of those per-position predictions are thrown away before backpropagation, keeping only the ones over assistant tokens.

Including the stop token is not optional

The assistant's turn should include its terminating token (<eos> or an equivalent turn-end marker) in the unmasked, loss-contributing region. If the end-of-turn token is masked out along with the surrounding structure, the model never receives a training signal for when to stop generating, and at inference it tends to ramble past where a response should end, or requires an external max-token cutoff to behave. Getting the mask boundary exactly one token off, ending it just before the stop token instead of just after, is a common, quiet source of that failure mode.

Multi-turn conversations

For a conversation with several back-and-forth turns, the standard approach masks every user and system turn and keeps every assistant turn unmasked, not just the final one. Training on only the last assistant turn while masking earlier assistant turns as if they were context-only would throw away a large fraction of the available supervision signal and bias the model toward performing well only when it is the very last turn in a conversation, a pattern that does not generalise to how the model is actually queried in production multi-turn use.

Interaction with packing

When fine-tuning examples are packed together the way pretraining documents are (see sequence-packing), the loss mask has to be reconstructed correctly for every packed conversation independently, aligned exactly with that conversation's own document-boundary attention mask. Getting the loss mask right but the attention mask wrong (or vice versa) produces a subtle bug: the model can still attend across conversation boundaries even though only the intended tokens contribute to the loss, so a leaked, unrelated conversation can quietly influence a completion without ever showing up as a labelled training signal you would notice by inspecting the loss curve.

When it falls down

  • Off-by-one masking around the stop token silently removes the model's only training signal for turn termination, producing a model that does not know when to stop.
  • Masking prompts too aggressively can starve formatting. If special role tokens or system-prompt structure are always masked, the model never receives gradient on producing that structure itself, which matters if the model is ever expected to generate structured output resembling the prompt template (for example, in agentic settings where the model writes its own tool-call scaffolding).
  • All-assistant-turns masking assumes every assistant turn is equally high quality. If some assistant turns in the training data are weaker than others (older model outputs, unedited drafts), unmasking all of them trains on that variance uniformly, which a real curation pipeline needs to account for separately from the masking mechanism itself.
  • Loss masking and attention masking are two different mechanisms that must agree, especially under packing; fixing one without checking the other leaves a real bug in place.

Further reading

Sign in to save and react.
Share Copied