Which Tokens Should Carry the Loss: Answer Supervision and Label Boundaries

The label mask decides what is scored, not what is visible#

A training sequence contains task rules, user material, and a demonstration answer. All of it participates in the computation, but not all of it should be treated as output the model is asked to imitate. Consider a long experiment record that must be read before producing a short JSON object. If every position carries a loss, training also demands that the model predict the experiment record itself. That is not a programming error; it changes how the supervision signal is distributed.

The loss mask determines which targets enter the loss. It does not determine which inputs the model can see.

Ignoring framework parameter names for a moment, the objectives are:

There is also a “completion-only” framing. It describes a prompt/completion segmentation, and it is not always equivalent to “all assistant messages”: if the completion contains several roles, or if historical assistant replies sit inside the prompt, the two rules select different positions. So define what should be trained first, then map that definition onto the trainer’s parameters. Reading a switch name is not a substitute for inspecting the actual labels.

1
0 [BOS]
2
1 [USER]
3
2 [material]
4
3 [END]
5
4 [ASSISTANT]
6
5 [answer A]
7
6 [answer B]
8
7 [END]

1
input:  BOS USER material END ASSISTANT answer A answer B END
2
labels:  -100 -100 -100 -100 -100   answer A answer B END

Ignoring prompt loss does not hide the prompt#

Setting the user material’s labels to the ignore value does not remove the corresponding input embeddings. At the answer positions, causal attention can still read the material.

Moreover, backpropagation from the answer loss can still pass through the representations at prompt positions and the parameters that produce them. “The prompt does not carry loss” therefore does not mean “the prompt did not participate in learning”, and certainly not “the prompt’s embeddings receive no gradient”.

Three things need to be kept apart:

• Input present — is the text fed to the model at all?
• Attention visible — can the current output read it?
• Target scored — is the model directly asked to predict the token at that position?

Supervising only the answer typically preserves the first two and selectively keeps the third. Blocking the whole prompt with an attention mask does not achieve the same effect: it would leave an extraction task without the evidence it needs.

End markers, framework switches, and truncation#

The extraction task needs the model to stop once the JSON is complete. If label filtering deletes every real end marker, the model loses direct supervision for stopping. But “put all EOS back in” is not a safe answer either. Some templates distinguish the end of a message from the end of a turn; some role headers are supplied by the prompt while other control tokens belong to the generated range. The decision should follow the target model’s template and inference stopping rule. If padding and EOS share an ID, handle them by their actual positions rather than ignoring every occurrence of that ID. The most direct acceptance check is to print the last few valid labels in readable form and confirm that both the answer content and the expected end marker are present, with no padding mixed in.

Using TRL 0.23.1 as a fixed reference, assistant_only_loss=True applies to assistant replies in conversational data, and the documentation requires the template to be able to return the corresponding token mask. Prompt-completion data has a separate completion_only_loss segmentation convention; the two can be combined, but not every input format can be dropped into either (TRL SFT Trainer). The Transformers 4.57.1 template interface also offers return_assistant_tokens_mask, whose support depends on generation-region markers in the template (tokenizer template interface). These documents express intent; they do not replace output inspection. Even with a switch set to true, spot-check the supervision scope of every turn and confirm that role headers, end markers, truncation, and the first answer token behave as intended. The article uses these fixed versions to explain semantics and does not claim that any local training project installed them or ran the corresponding configuration.

A convenient-looking alternative is to find the assistant marker in the text and treat everything after it as supervised. It runs into several problems: multiple boundaries in a multi-turn record; user material that quotes the same words; template markers that are hard to distinguish from ordinary prose; and the fact that a string tokenized on its own may not have the same boundaries as the same string inside its full context. A more reliable starting point is structured messages plus the generation-region information the template provides, or an explicit position mapping maintained in a custom pipeline. Whichever mechanism is used, it needs to be tested with material that contains boundary-style text.

Truncation can also leave a sample with no valid label at all. Suppose the maximum length is short, long material fills the window, and the whole answer is cut away: the input tensor still holds many tokens, but the loss mask is all zeros. Such a sample cannot be averaged in with the rest as ordinary training data; it produces ineffective updates, or non-finite loss if handled carelessly. It should be excluded explicitly, resegmented, or handled by another defined policy, with the count recorded. A second case is harder to notice: only half the answer survives, so valid targets exist but the end marker is gone. A program that checks “at least one label exists” cannot know the sample no longer matches the complete task. Checks therefore have to cover both the number of valid targets and the completeness of the target sequence.

Sample weights and a minimal label check#

Suppose two samples have effective answer lengths of 2 and 8, with mean token losses of 1 and 3. Averaging over valid tokens gives

while averaging over samples gives

The first gives every target position equal weight, so longer answers contribute more; the second gives every record’s average loss equal weight. Which one applies should match the task and the comparison being asked. If the goal is to compare full-sequence supervision with answer-only supervision, the two loss values cannot simply be compared: they score different sets of positions, so the numbers are not the same measurement. A better approach fixes the data and the training budget, computes validation metrics on the same answer positions, and then looks at the extraction task result.

The snippet below checks only the teaching sequence and the internal shift convention. It does not depend on a training framework and does not generate a real chat template:

1
# Each item is a custom teaching token ID, not real tokenization output.
2
ids = [0, 1, 2, 3, 4, 5, 6, 3]
3
answer_mask = [0, 0, 0, 0, 0, 1, 1, 1]
4
labels = [token if keep else -100
5
          for token, keep in zip(ids, answer_mask)]
6

7
# Internal shift: the output at position t predicts labels[t + 1].
8
pairs = [(t, labels[t + 1])
9
         for t in range(len(ids) - 1)
10
         if labels[t + 1] != -100]
11
assert pairs == [(4, 5), (5, 6), (6, 3)]

It verifies the position relationship between the first answer and the end marker. Real data must also cover multi-turn records, left and right padding, empty answers, templates that do not support masking, truncation, and packing. Passing this small check is not evidence that a whole trainer is correct.

The Chinese original and the limits of this note#

The full Chinese article carries the eight-position token listing, the label table, the pairing table, the Python snippet, the multi-turn and truncation walkthroughs, and the reference list with the TRL 0.23.1 and Transformers 4.57.1 links — material summarized but not reproduced here. The original is at the Chinese source article; the next note covers how instruction data is selected and generated. The main limit is unchanged: the label conventions described here were checked against a hand-made teaching sequence, and no training run establishes that any supervision scope performs better.