Answer-only loss computes cross-entropy only on the teacher’s answer tokens, excluding prompt tokens from the training objective. It is used here for distillation so the student learns the reasoning trace and final answer conditioned on the prompt. section 8.6 section 8.6
Implementation starts by shifting tokens into input_ids and target_ids, then slicing away the prompt portion before calling torch.nn.functional.cross_entropy. The prompt-length boundary controls which logits and targets count toward the loss. section 8.6 section 8.6
The workflow uses this computation in three steps:
- Prepare shifted inputs and targets from the token sequence. section 8.6
- Drop logits and targets that belong to the prompt prefix. section 8.6
- Compute cross-entropy on the remaining answer tokens and reuse it for example and dataset averages. section 8.6
Single-example distillation loss
A helper computes loss for one example by building input_ids and target_ids, choosing answer_start = max(prompt_len - 1, 0), and applying cross-entropy to answer_logits and answer_targets. The same logic is then wrapped into an evaluation loop for train and validation examples.
