Supervised Instruction Fine-Tuning

Turn demonstrations into a masked likelihood objective and fix the chat template that later chapters reuse.

A pretrained model assigns probabilities to continuations; an assistant must recognize roles, stop copying the prompt, and produce a response that satisfies the instruction. Supervised fine-tuning (SFT) closes that gap with the cheapest available signal: demonstrations. Each example is a prompt x and a desired response y, and training maximizes the likelihood of y given x.

This chapter builds the two pieces every later posttraining stage depends on: a response-only loss mask and a chat template with role boundaries. Both are small pieces of code with outsized consequences, because a mask or template mistake is silent — the loss decreases either way. Design rule: the mask is the specification; make it explicit data, never an index convention.

Response likelihood

Let \pi_\theta be the model’s conditional distribution over tokens. The SFT loss applies the language-model objective to response positions only:

\mathcal{L}_{\mathrm{SFT}}(\theta) = -\frac{1}{|y|}\sum_{t \in y} \log \pi_\theta(y_t \mid x, y_{<t}).

The indicator m_t = \mathbb{1}[t \in y] carries the specification. Training on the full sequence instead asks the model to reproduce user prompts, which spends capacity on the wrong target and teaches the model to imitate the user rather than the assistant. The choice is invisible in the loss curve — both objectives decrease smoothly — and visible only in behavior.

Response-only masking

Masked cross-entropy reduces to three steps: normalize logits with a softmax, select the negative log-probability of each target token, and average only over supervised positions. Implementing it directly makes the denominator explicit. The loss is a mean over |y| response tokens, not over sequence length, so padding changes neither the value nor the gradient scale.

import numpy as np

rng = np.random.default_rng(8)

V = 12   # vocabulary size
T = 6    # sequence length

tokens = np.array([3, 5, 8, 2, 9, 0])
is_response = np.array([0, 0, 0, 1, 1, 0])
is_padding = np.array([0, 0, 0, 0, 0, 1])
loss_mask = is_response * (1 - is_padding)

logits = rng.normal(size=(T, V))                    # stand-in for model output


def log_softmax(z):
    z = z - z.max(axis=-1, keepdims=True)
    return z - np.log(np.exp(z).sum(axis=-1, keepdims=True))


def masked_cross_entropy(logits, targets, loss_mask):
    log_probs = log_softmax(logits)
    token_nll = -log_probs[np.arange(len(targets)), targets]
    return (token_nll * loss_mask).sum() / max(loss_mask.sum(), 1.0)


full_loss = masked_cross_entropy(logits, tokens, np.ones(T))
response_loss = masked_cross_entropy(logits, tokens, loss_mask)
print("full-sequence loss:", round(float(full_loss), 4))
print("response-only loss:", round(float(response_loss), 4))
print("loss mask:         ", loss_mask)

full-sequence loss: 2.8008
response-only loss: 3.0798
loss mask:          [0 0 0 1 1 0]

The two numbers differ because they average over different token sets: the full-sequence value mixes prompt and padding positions into a quantity that will be reported as “the loss.” Annotations: is_response marks the positions belonging to the assistant turn (1); is_padding marks positions that must contribute nothing (2); the loss mask is the product, so a padded response position is excluded (3). The max(..., 1.0) guard keeps an all-masked batch defined instead of dividing by zero.

Chat templates

Role information reaches the model only through tokens. A chat template is a deterministic function from a message list to token IDs: each role gets a registered special token, and the response mask is derived from the same serialization, so the token stream and the loss mask cannot disagree. The special tokens belong in the vocabulary before training begins; adding them later shifts every token ID and invalidates saved checkpoints.

Template bugs are silent. A missing end-of-turn token lets one response blur into the next prompt; masking user text as response trains the model to imitate users. Both produce a smoothly decreasing loss, which is why the mask is checked against the serialized text below instead of being assumed.

SPECIALS = ["<|system|>", "<|user|>", "<|assistant|>", "<|tool|>", "<|end|>"]
CHARS = "abcdefghijklmnopqrstuvwxyz .,!?'"

id_of = {tok: i for i, tok in enumerate(SPECIALS)}
for ch in CHARS:
    id_of[ch] = len(id_of)
token_of = {i: tok for tok, i in id_of.items()}


def apply_template(messages):
    """Serialize messages to IDs, marking assistant spans as response."""
    ids, mask = [], []
    for msg in messages:
        role = f"<|{msg['role']}|>"
        supervised = msg["role"] == "assistant"
        ids.append(id_of[role]); mask.append(int(supervised))
        for ch in msg["text"]:
            ids.append(id_of[ch]); mask.append(int(supervised))
        ids.append(id_of["<|end|>"]); mask.append(int(supervised))
    return ids, mask


messages = [
    {"role": "user", "text": "name three colors"},
    {"role": "assistant", "text": "red, green, blue"},
]
ids, mask = apply_template(messages)

assistant_span = [token_of[i] for i, m in zip(ids, mask) if m]
recovered = "".join(t for t in assistant_span if not t.startswith("<|"))
print("supervised span:", " ".join(assistant_span))
print("recovered text: ", recovered)
assert recovered == "red, green, blue"

supervised span: <|assistant|> r e d ,   g r e e n ,   b l u e <|end|>
recovered text:  red, green, blue

The role token opens the supervised span (1), and the end-of-turn token is supervised as well (2): the model must learn when to stop, not only what to say. The assertion (3) is the round-trip check that catches template and mask disagreement. Decoding only the supervised span must recover exactly the assistant text; if it also recovers user text or truncates the response, the mask is wrong regardless of what the loss curve reports.

The template defined here, including the <|tool|> role reserved now, is the exact substrate for tool calling in Chapter 11: tool calls and tool results serialize as additional message types over the same boundaries.

Masking experiment

Train two identical models on the same demonstrations, one with full-sequence loss and one with response-only loss, then evaluate both on three inputs: held-out demonstrations, the same instructions paraphrased, and the general suite from Chapter 07. The training loss is not the comparison of interest, since it falls for both. The behavioral differences are:

• Full-sequence training fits user text as well as responses, which shows up at generation time as models that continue or imitate the prompt instead of answering it.
• The paraphrase evaluation separates instruction-following from phrasing memorization: if accuracy drops sharply when an instruction is reworded, the model keyed on the demonstration wording rather than the task.

Record response length as well. SFT data skews short and formulaic, and models inherit that skew.

Regression checks

Freeze an evaluation set from before posttraining and rerun it after every SFT checkpoint: perplexity on held-out text, calibration of next-token probabilities, and the generation checks from Chapter 07. The fine-tuning distribution (short, formatted responses) differs from the pretraining distribution, and nothing in the SFT loss reports the cost of that shift. An SFT result is the pair of numbers, instruction-following metric and general-suite score, not the first alone.

Shared SFT serialization contract

The NumPy example above makes the objective transparent; the reusable path derives the mask while it serializes the role boundaries. This prevents a later chat-template edit from silently leaving the loss mask behind. The assistant end-of-turn marker is supervised, while role markers and user content are not.

from proof_lm.posttraining import ChatMessage, serialize_chat

posttraining_messages = [
    ChatMessage("user", "state a theorem"),
    ChatMessage("assistant", "the theorem is valid"),
]
serialized = serialize_chat(
    posttraining_messages,
    lambda text: [ord(character) for character in text],
    {"bos": 0, "eot": 1, "role_user": 2, "role_assistant": 3},
)
print("serialized tokens:", len(serialized.token_ids))
print("supervised tokens:", sum(serialized.loss_mask))
print("assistant span:", serialized.spans[-1])
assert not any(serialized.loss_mask[: serialized.spans[-1].start])
assert all(serialized.loss_mask[serialized.spans[-1].start + 1 :])

serialized tokens: 40
supervised tokens: 21
assistant span: RoleSpan(role='assistant', start=18, end=40, supervised=True)

Summary

• SFT is the language-model loss restricted to response tokens; the loss mask is the specification, and it is data derived from the template, never an index convention.
• Masked cross-entropy averages over supervised tokens only, so padding changes neither the loss value nor the gradient scale.
• The chat template, with role and end-of-turn special tokens registered up front, is shared with tool calling in Chapter 11.
• Evaluate on held-out and paraphrased instructions and rerun the general suite; the loss measures likelihood of demonstrations, which is compatible with having learned response phrasing instead of the task.

Exercises

Use the exercises to test the chapter’s invariants and connect the derivations to the reusable implementation. Solutions are hidden in the notebook source and are available through the course tooling when needed.

[P8.1] Response-only loss mask

Given prompt tokens [1, 2] and response tokens [3, 4], write the loss mask for response-only SFT and explain what changes when padding is appended.

[P8.2]

Use the shared chat serializer to explain why a response-only mask is derived from role spans rather than hard-coded token indices. State which token is supervised at the end of an assistant turn.