Instruction Tuning Training › Loss Masking
Language AI HandbookAll chapters Instruction Tuning Training › Loss MaskingOne section · Sections in this chapter3 / 8 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Standard language modeling calculates loss for all tokens. In instruction tuning, this behavior wastes gradient capacity on recreating instructions instead of learning to produce helpful responses. Loss masking is the technique that redirects the training signal to where it belongs.
The concept is straightforward once you understand how gradient flow works. Every token that contributes to the loss function produces a gradient update that adjusts the model's parameters. If both the instruction tokens and the response tokens contribute equally to the loss, then the model is simultaneously being trained to reproduce instructions and to generate appropriate responses. These two objectives are related but distinct. The model already knows how to predict text fluently from pre-training. What it needs to learn during instruction tuning is specifically how to transition from a received instruction to an appropriate response, not how to reproduce the instruction itself.
Think of loss masking like a professional training simulation. When training a customer service representative, you play recorded customer calls and ask the trainee to practice the representative's responses. You would not grade the trainee on how accurately they reproduced what the customer said. The customer's words are the context that the trainee receives; the trainee's words are what you are training. Loss masking makes the same distinction: the instruction is input context, and the response is what you are training the model to produce.
Why Mask the Prompt
Consider an instruction tuning example:
User: Explain photosynthesis in simple terms.
Assistant: Photosynthesis is the process plants use to convert sunlight into food...
Without loss masking, the model receives gradient updates for every token, including the instruction "Explain photosynthesis in simple terms." But recreating user instructions is not the goal. The model should learn to generate helpful responses, not memorize prompts. Training on prompt tokens dilutes the learning signal with information the model already understands from pre-training. The pre-trained model is already excellent at predicting the token "Explain" given the token "User:" in its context. Spending gradient updates to reinforce this existing capability is a waste.
Loss masking addresses this by zeroing out the loss for prompt tokens. Only the assistant's response contributes to parameter updates. The result is that every gradient step moves the model's parameters specifically in the direction of better response generation, making instruction tuning more sample-efficient and more precisely targeted.
The effect is particularly significant for long prompts. In retrieval-augmented generation scenarios, the prompt may contain hundreds or thousands of tokens of retrieved context. Without loss masking, the model would spend the large majority of its training budget on predicting tokens that appear in documents it was given as context, which is irrelevant to its goal of creating accurate responses. With loss masking, all of that context is ignored in the loss calculation and only the response tokens shape the parameters.
import torch
def compute_masked_loss(logits, labels, mask):
"""
Compute cross-entropy loss only on unmasked positions.
Args:
logits: Model outputs of shape (batch, seq_len, vocab_size)
labels: Target token IDs of shape (batch, seq_len)
mask: Binary mask where 1 = compute loss, 0 = ignore
"""
# Flatten for cross-entropy computation
batch_size, seq_len, vocab_size = logits.shape
# Shift for next-token prediction
shift_logits = logits[:, :-1, :].contiguous()
shift_labels = labels[:, 1:].contiguous()
shift_mask = mask[:, 1:].contiguous()
# Compute per-token loss
loss_fn = torch.nn.CrossEntropyLoss(reduction="none")
per_token_loss = loss_fn(
shift_logits.view(-1, vocab_size), shift_labels.view(-1)
).view(batch_size, -1)
# Apply mask and average over valid tokens
masked_loss = per_token_loss * shift_mask
total_loss = masked_loss.sum() / shift_mask.sum()
return total_lossThe necessary detail in this implementation is how the reduction works. Standard cross-entropy with reduction="mean" averages over all positions in the sequence. By using reduction="none" first, we get a per-token loss, which we then multiply by the mask (zeroing out prompt positions) and divide by the total number of active mask positions. This ensures the averaged loss shows only the response tokens, not the entire sequence length. If we divided by sequence length instead of by the mask sum, the loss would be artificially diluted by all the zero-loss prompt positions, making the effective gradient much smaller than intended.
Creating Loss Masks
Loss masks must align with tokenized sequences. This requires tracking where the prompt ends and the response begins. In practice, this tracking happens during data preprocessing, where the tokenizer produces an attention mask alongside the token IDs, and you augment this with a separate loss mask that follows the same alignment:
def create_loss_mask(input_ids, response_start_idx):
"""
Create a loss mask that zeros out prompt tokens.
Args:
input_ids: Tokenized sequence
response_start_idx: Index where assistant response begins
Returns:
Binary mask tensor
"""
seq_len = len(input_ids)
mask = torch.zeros(seq_len)
mask[response_start_idx:] = 1.0
return mask
# Example usage
prompt = "User: What is machine learning?\nAssistant:"
response = " Machine learning is a subset of AI..."
full_text = prompt + response
# Simulate tokenization (actual implementation uses a tokenizer)
prompt_tokens = 12 # Number of tokens in prompt
response_tokens = 8 # Number of tokens in response
mask = create_loss_mask(
input_ids=list(range(prompt_tokens + response_tokens)),
response_start_idx=prompt_tokens,
)Loss mask pattern: Prompt tokens (masked): [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] Response tokens (active): [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
A subtle but important detail is that finding the response start index requires care when using sub-word tokenizers. You cannot simply count characters in the prompt and divide by some average token length. The correct approach is to tokenize the prompt and the full sequence separately, then compute response_start_idx = len(tokenized_prompt). This is exact and handles cases where the prompt-response boundary falls inside a multi-character token correctly.
For multi-turn conversations, the masking logic extends naturally. Only tokens that are part of assistant turns contribute to the loss; user turns and system messages are masked throughout. This means that in a long conversation, the active loss positions may be scattered through the sequence, interspersed with masked regions corresponding to each user message.
Visualizing Masked Loss
The following visualization shows how loss masking focuses training on the response portion of each example:
Impact of Loss Masking
Loss masking materially affects what the model learns during instruction tuning. Without masking, gradients flow from both prompt and response tokens, diluting the signal that teaches response generation. With masking, all gradient information comes from response tokens, making training more efficient and focused.
The effect varies considerably across different types of instruction-following examples. In short question-answering examples, the prompt might constitute 40% of the tokens, so masking reclaims 40% of the gradient budget for response learning. In code generation tasks where the prompt includes a detailed specification and the response is a single function, the ratio might be 30% prompt to 70% response. But in long-context applications where the prompt contains retrieved documents or multi-turn conversation history, prompts can easily constitute 60% to 70% of all tokens in the sequence.
The effect is particularly pronounced in scenarios with long prompts relative to responses. In long-context applications or multi-turn conversations, prompts can comprise 60% to 70% of tokens. Without masking, most of the learning signal would be spent on recreating context rather than learning to respond appropriately. The model would implicitly be trained to be a very good text predictor for the kinds of text that appear in instructions, which is exactly what it already learned to do during pre-training. Loss masking gives instruction tuning a different objective from continued pre-training on another distribution.
Worked Example: Computing Masked Loss Step by Step
Let us trace through a concrete numerical example to see exactly what happens during a masked loss computation. Suppose we have a simple vocabulary of five tokens: "the", "cat", "sat", "mat", "on". Our instruction-response sequence consists of three prompt tokens followed by two response tokens.
The prompt is "the cat sat" (token IDs 0, 1, 2). The response is "on mat" (token IDs 3, 4). The full sequence is [0, 1, 2, 3, 4] with loss mask [0, 0, 0, 1, 1].
During training, the model processes the full sequence autoregressively. For the shifted prediction task, the input is [0, 1, 2, 3] and the target is [1, 2, 3, 4]. The loss mask applied to the shifted targets is [0, 0, 1, 1] (we shift the mask along with the labels).
Suppose the model produces the following logits for the four prediction positions:
- Position 0 (predicting token 1 = "cat"): logits = [0.1, 2.0, 0.3, 0.1, 0.1], cross-entropy = 0.15. Mask = 0, so contribution = 0.
- Position 1 (predicting token 2 = "sat"): logits = [0.2, 0.4, 1.8, 0.2, 0.2], cross-entropy = 0.22. Mask = 0, so contribution = 0.
- Position 2 (predicting token 3 = "on"): logits = [0.1, 0.1, 0.3, 1.9, 0.3], cross-entropy = 0.18. Mask = 1, contribution = 0.18.
- Position 3 (predicting token 4 = "mat"): logits = [0.2, 0.2, 0.2, 0.5, 1.7], cross-entropy = 0.31. Mask = 1, contribution = 0.31.
The total masked loss is:
where the denominator is the sum of mask values (2), not the total sequence length (4). If we had used the total sequence length, the loss would be , which is half as large. The smaller loss would produce smaller gradients, making training half as efficient. The denominator must be the number of active (unmasked) positions.
This numerical example shows why the implementation detail in compute_masked_loss matters: total_loss = masked_loss.sum() / shift_mask.sum() divides by the mask sum, not by the sequence length. Getting this wrong silently reduces gradient magnitudes and can require doubling your learning rate or training steps to compensate.