A supervised language-model example often contains more than the text you want the model to produce. It may include a system message, a user request, separators, and an assistant answer. If you compute next-token loss over the entire sequence, the model is trained to predict all of those tokens, not just the assistant response.

That may be intentional for some training objectives. For instruction fine-tuning, though, developers often want the prompt to provide context while only selected response tokens contribute to the supervised loss. A loss mask makes that distinction explicit.

This article builds a practical mental model for prompt-token masking, shows where the mask sits in the causal language-model objective, and explains the boundary mistakes that can quietly train the wrong behavior.

Separate input context from supervised targets

Consider a simplified training record:

User: Reset my API key.
Assistant: Open Settings, then choose Rotate key.

After formatting and tokenization, a causal language model receives one token sequence. At each position it predicts the next token from the tokens to its left.

Without masking, the training objective includes predictions throughout the sequence. The model may receive loss for predicting tokens in User: Reset my API key. as well as tokens in the assistant answer.

With response-only masking, prompt tokens still enter the forward pass. The model can attend to them when predicting the answer. Their target positions simply do not contribute to the aggregated training loss.

That distinction is the central idea:

visible to the model != included in the loss

Masking a prompt from the loss does not mean deleting it from the input or blocking attention to it.

See where the loss mask applies

Let a tokenized sequence be:

x0 x1 x2 x3 x4 x5

A causal model uses the prefix through x_t to predict x_(t+1). If the first four tokens belong to the prompt and the final two belong to the assistant response, the conceptual prediction targets are:

input position:   x0  x1  x2  x3  x4
target token:     x1  x2  x3  x4  x5

Suppose x4 is the first assistant token. The prediction made from the prefix ending at x3 is therefore the first prediction that should count toward a response-only objective:

loss included?:   no  no  no  yes yes

For an included target position t, the usual token loss is negative log-likelihood:

L_t = -log p(x_t | x_<t)

A binary mask m_t can select the supervised targets:

L = sum(m_t * L_t) / sum(m_t)

This formula is a teaching model. Frameworks often implement causal shifting internally and may represent ignored labels with a sentinel value rather than multiplying a loss tensor by zero. The important invariant is that the mask must align with target tokens after the causal shift.

The first response token exposes off-by-one bugs

The boundary between prompt and answer deserves special attention because causal prediction is shifted by one position.

Imagine the formatted sequence ends the prompt with a separator and then begins the assistant response:

... <assistant> Open Settings

To learn that Open is an appropriate first response token, the model must receive loss on the prediction of Open from the preceding context. If code starts the loss one token too late, Open becomes visible context for predicting Settings, but the model is never directly trained on the first response token.

The opposite mistake includes one prompt token as a target. That usually will not crash training. It simply changes the objective, which makes the bug easy to miss.

A useful unit test uses a tiny token sequence with a known boundary and prints three aligned rows:

tokens:   [A, B, C, D, E]
roles:    [P, P, P, R, R]
loss:     [0, 0, 0, 1, 1]

Then verify how your training library performs the next-token shift. Do not assume a mask indexed against input tokens automatically has the correct target alignment.

Chat templates make boundaries part of the objective

Real instruction data is usually serialized through a chat template rather than simple User: and Assistant: strings. The template may add role markers, message terminators, beginning-of-sequence tokens, or other control tokens.

Those tokens are not cosmetic. They become part of the sequence the model sees, so you need a deliberate rule for which of them receive loss.

For example, an assistant role marker can be treated as prompt-side context while the answer content is supervised. An assistant end-of-message token may need to be supervised if the model is expected to learn when to stop the response. The correct choice depends on the model’s template and the behavior the training objective is meant to teach.

This is why building masks by searching decoded text for a literal string is fragile. Tokenization can split text unexpectedly, the same text can occur inside user content, and templates can change. Prefer boundaries produced by structured message roles or by a tokenizer/template API that exposes the relevant token spans when such support exists.

Multi-turn conversations need a policy, not a single split

A conversation can contain several assistant messages:

system -> user -> assistant -> user -> assistant

There is no universal rule saying only the final assistant message should receive loss. Two reasonable objectives are different:

  • all-assistant loss trains on every assistant response in the example;
  • final-response loss treats earlier turns as context and supervises only the final assistant response.

The first uses more target tokens from each conversation. The second can be useful when each training record is designed around one final behavior and earlier assistant messages are supplied only as context.

Whichever policy you choose, encode it explicitly. A data pipeline that silently assumes “everything after the first assistant marker is a target” can accidentally include later user messages in the loss.

Normalize by supervised tokens when comparing examples

Masking changes how many target tokens remain in each sequence. That affects how losses should be aggregated.

Suppose one batch contains two examples:

example A: 20 prompt tokens + 5 supervised tokens
example B: 5 prompt tokens  + 20 supervised tokens

If you average a per-sequence loss and then average sequences equally, the five supervised tokens in example A can receive more aggregate weight per token than the twenty in example B. If your intended objective is the mean loss over supervised tokens, sum losses over valid target positions and divide by the number of valid target positions.

Distributed and gradient-accumulation setups need the same care. Averaging local batch means can differ from a true global token mean when workers or microbatches contain different numbers of supervised tokens. The exact reduction depends on the training framework, so verify its contract rather than assuming the denominator.

Masking changes the training objective, not just efficiency

Prompt masking is sometimes described as a way to avoid “wasting loss” on instructions. That framing is incomplete. The choice changes what distribution the model is optimized to predict.

Training on every token can be appropriate when the dataset is ordinary language-model text or when reproducing the complete serialized conversation is part of the objective. Response-only loss focuses supervised capacity on assistant-side targets while retaining prompts as conditioning context.

Neither objective is automatically better for every dataset. If you are continuing pretraining on domain text, masking arbitrary prefixes would discard useful language-model targets. If you are adapting a chat model from instruction-response pairs, supervising user text may spend optimization effort modeling a role the deployed model is not expected to generate.

The practical question is therefore: which tokens represent behavior the model should learn to generate? The loss mask should follow that answer.

Watch for empty and truncated targets

Loss masking creates a few edge cases that ordinary full-sequence training can hide.

If preprocessing produces an example with no supervised tokens, the denominator in a token-mean objective becomes zero. The example should usually be rejected or handled explicitly rather than allowed to produce an undefined loss.

Truncation is another common source of silent damage. If long prompts are truncated from the end, they may remove the assistant response entirely. If truncation preserves the response by cutting the prompt, it can remove context needed to make the answer valid. Inspect the distribution of supervised-token counts after tokenization and truncation, not only the raw text lengths before preprocessing.

Packing multiple conversations into one training sequence adds another constraint: each example’s role boundaries and attention semantics must remain correct after packing. A correct loss mask cannot repair accidental cross-example attention or a target that crosses an example boundary.

Validate the mask before spending training compute

A small preprocessing audit catches many expensive mistakes. Before a full run, decode or otherwise inspect a sample of tokenized examples together with their loss masks.

Check that:

  • user and system content is included as context but excluded from loss when that is your policy;
  • the first intended assistant token is supervised;
  • intended stop or end-of-message tokens are handled consistently;
  • later user turns are not accidentally supervised;
  • truncated examples still contain valid targets;
  • every retained example has at least one supervised token;
  • the reported token count used for loss normalization matches the mask.

Also compare a few training examples against the exact chat template used at inference. Fine-tuning one serialization and serving another can create a distribution mismatch even when the loss mask itself is correct.

Use the simplest objective that matches deployment

Loss masking is useful because it lets one token sequence play two roles: some tokens condition the model, while selected tokens define the supervised behavior. Once that distinction is clear, the implementation becomes easier to reason about.

Start from the behavior you want the model to generate, mark those target tokens explicitly, and test the causal shift at the first response token. Then inspect truncation, multi-turn boundaries, and loss normalization before scaling the run.

If full-sequence next-token training already matches your task, keep it. A mask is not a mandatory ingredient of fine-tuning. It is a precise tool for making the supervised objective match the parts of a conversation the model is expected to produce.

Artificial Intelligence LLM Fine-Tuning Model Training