Per-Token Loss Masking

What is Per Token Loss Masking?

In Supervised Fine Tuning (SFT), we train a base Large Language Model to act as a helpful conversational assistant.

Each training sample consists of a Prompt (written by the user) and a Target Response (written by an expert assistant):

  Token Stream:   [ "User: What is the capital of France? \n Assistant:" ] [ "The capital of France is Paris." ]
  Target Labels:  [ -100, -100, -100, -100, -100, -100, -100, -100, -100 ] [ "The", "capital", "of", "France", "is", "Paris", "." ]
                  (MASKED OUT! Loss = 0.0)                                 (GRADIENTS COMPUTED HERE ONLY!)

Per Token Loss Masking means we calculate Cross Entropy loss ONLY on the Assistant Response tokens, ignoring prompt tokens completely.

Why Prompt Tokens Must Be Masked

Suppose you do NOT mask prompt tokens, computing loss over the entire sequence (Prompt + Response).

What happens during backpropagation?

1. The model is penalized for failing to predict the user's exact prompt words ("What is the...").
2. The model spends GPU capacity memorizing user prompt phrasing rather than learning how to generate good answers.
3. In multi turn chat, early user turns get re-trained repeatedly, corrupting language representations.

We provide user prompt tokens as given input context, masking their loss so gradient updates focus 100 percent on generating accurate assistant responses.

PyTorch Implementation (ignore_index = -100)

PyTorch nn.CrossEntropyLoss has a built in parameter ignore_index = -100.

Any token assigned label -100 is skipped during loss and gradient calculations:

import torch
import torch.nn as nn

# Loss function ignoring label -100
criterion = nn.CrossEntropyLoss(ignore_index=-100)

# Example: Prompt tokens assigned -100, target tokens assigned real ID
labels = torch.tensor([[-100, -100, -100, 4521, 8932, 19042]])

loss = criterion(logits.view(-1, vocab_size), labels.view(-1))
loss.backward() # Gradients computed ONLY for token IDs 4521, 8932, 19042!

Multi Turn Chat Masking Pattern

In multi turn conversations:

  Turn 1 User:      "Hi"                     ──► Assign Labels = -100 (Masked)
  Turn 1 Assistant: "Hello! How can I help?" ──► Assign Real Token IDs (Compute Loss!)
  Turn 2 User:      "Tell me a joke"          ──► Assign Labels = -100 (Masked)
  Turn 2 Assistant: "Why did the chicken..." ──► Assign Real Token IDs (Compute Loss!)

This allows packing multi turn chat sessions into a single GPU forward pass while updating weights exclusively on assistant responses.

Say this out loud

Per Token Loss Masking sets target labels for input prompt tokens to -100 during Supervised Fine Tuning, excluding them from cross entropy loss. This ensures model gradients update weights exclusively on generating accurate assistant response tokens rather than memorizing user prompt syntax.

Followups to expect

1. What happens if you accidentally mask assistant response tokens? The model receives zero gradient signal for those tokens, failing to learn how to output those response words.
2. What is Data Packing in SFT? Concatenating multiple short instruction examples into a single long 4096 token tensor separated by EOS tokens, using loss masking and attention position resetting to train GPUs at 100 percent efficiency.