Handout #18 — Instruction Tuning and Reward Models

Demonstrations teach a model what to generate; preferences teach a scorer which response is better. We will design an unseen-task evaluation and turn comparisons into a reward-model loss and working code.

1Instruction tuning: what counts as generalization?

1.1 Same loss family, different data

Instruction tuning applies response-token SFT across instruction–context–response examples from many tasks. Handout #17 developed the loss; here we ask what transfers across tasks.

After training on sentiment and entity extraction, testing new sentiment reviews measures new-example performance; testing another review domain measures domain transfer. Holding out summarization entirely tests unseen-task-type transfer.

1.2 Count diversity and report the right average

A thousand differently worded sentiment prompts need not provide a thousand different reasoning skills. Task counts depend on taxonomy, and training mixtures determine which skills receive most updates. Compare mixtures at a stated token/example budget rather than attributing every improvement to task count alone.

1.3 A demonstration is neither a verifier nor a knowledge guarantee

For “Explain space elevators to a six-year-old,” a useful demonstration must be understandable, relevant, and careful about what is hypothetical. SFT increases the likelihood of its target text; the loss does not independently check those properties. Several different explanations can all be good, even though only one appears as the target.

Always-confident demonstrations may fail to teach uncertainty. In controlled closed-book QA, Gekhman et al. identify hallucination risks when demonstrations introduce knowledge the model does not reliably possess. This is not a claim that SFT cannot teach facts. Evaluate both answerable and unanswerable questions.

2From demonstrations to comparisons

2.1 What an annotator actually supplies

Keep the same prompt: “Explain space elevators to a six-year-old.” Candidate A explains a proposed very long cable that a vehicle could climb to space and notes that it has not been built. Candidate B explains gravity without answering the question. Under a rubric emphasizing relevance, understandable language, and factual care, an annotator prefers A.

# Simplified comparison record; no numerical reward label is required.
{
    "prompt_id": "space-elevator-01",
    "prompt": "Explain space elevators to a six-year-old.",
    "chosen": "Imagine a very long cable ... It is a proposed idea ...",
    "rejected": "Gravity pulls things toward Earth ...",
    "rubric_version": "relevance-clarity-factual-care-v1"
}

The record orders A and B for this prompt and rubric; it supplies neither absolute quality scores nor a factual guarantee. Record which policy generated the candidates.

Connect this to slide 18’s space-elevator example

2.2 Preference is conditional on the evaluation procedure

Two accurate explanations may trade off brevity and detail. Before labeling, specify the audience, what makes an answer useful, how factual errors are handled, and whether ties or “cannot judge” are allowed. Randomize candidate order and audit position/length effects rather than treating interface choices as harmless.

Keep comparisons from the same prompt in one split to prevent leakage. Preserve disagreement: scalar scores impose one ordering for a fixed prompt/rubric, whereas preferences can be heterogeneous or cyclic. The binary loss below assumes strict preferences. Store ties and “cannot judge” separately; using them in training requires an explicit modeling choice, not arbitrary tie-breaking.

Candidate A: detailed steps

Candidate B: brief outline

3An episode and its expected reward

3.1 Two views of the same response

View a complete response as one action, or each next token as an action. In the token view, the state is the prompt plus generated prefix; appending a token forms the next state. The slides use terminal reward, received at EOS or another stopping condition.

For a toy response A cable climbs EOS, the states progress from $x$ to $(x,\text{A})$ to $(x,\text{A cable})$. The language model supplies the next-token distributions and thereby the sequence distribution:

Suppose $R_\phi(x,y)$ is fixed for the moment. The objective written on slides 16–17 averages over both prompts and responses:

$R_\phi$ scores a completed response; $p_\theta$ determines how often it is generated. The learned score reflects its training data and rubric, not guaranteed overall quality.

Averaging scores of sampled responses estimates this expectation. Changing the prompt mix changes the objective; policy-gradient estimation comes next lecture.

4Build a probability model for a preference

4.1 The score difference is the logit

One shared reward model maps a prompt–response pair to a scalar. Let $r_A=R_\phi(x,A)$ and $r_B=R_\phi(x,B)$. The Bradley–Terry model uses their difference as a binary logit:

If A is the observed winner, its negative log-likelihood is

For arbitrary A/B order, let $z=1$ when A wins and $z=0$ when B wins. Then $\ell=\operatorname{BCEWithLogits}(\Delta,z)$. Alternatively store chosen/rejected in a fixed order and use a target of one. Do not reverse the order without changing the label.

With $r_A=1.2$, $r_B=0.8$, and A preferred, $\Delta=0.4$, $P(A\succ B)\approx0.598688$, and $\ell\approx0.513015$ nats. The scores 1.2 and 0.8 are not probabilities; the difference determines a probability for this comparison.

4.2 Check the direction, not just the scalar loss

For the example with A winning, the gradients are approximately −0.401312 and +0.401312. If the two scores are optimized as independent variables, gradient descent raises A’s score and lowers B’s. A neural reward model instead updates shared parameters through the chain rule, so that individual-score movement is not guaranteed. At learning rate 0.1, independent score updates give approximately 1.240131 and 0.759869, widening the difference to 0.480262.

4.3 Disagreement need not mean the optimizer failed

Suppose the same pair receives seven A-win labels and three B-win labels. A single score difference cannot make every observation certain. Without other examples or regularization, the average binary cross-entropy is minimized at $\sigma(\Delta)=0.7$, so $\Delta=\log(7/3)\approx0.847298$. Its nonzero loss, about 0.610864, reflects the conflicting labels.

For unanimous, separable comparisons, the unregularized loss favors ever-larger margins, with no finite optimum in a freely adjustable difference. Evaluate held-out accuracy and loss, including disagreement slices: identical rankings can conceal very different probability calibration.

5Make the code implement the intended comparison

5.1 Where did the sigmoid go?

BCEWithLogitsLoss combines the sigmoid and binary cross-entropy using a numerically stable formulation. Feed it the raw score difference. For chosen/rejected pairs, the equivalent stable expression is -F.logsigmoid(chosen - rejected).mean(). See the PyTorch 2.12 API, matching the tested lab environment.

# score_a and score_b are [batch] scalars from the SAME reward model.
assert score_a.shape == score_b.shape == labels.shape
logits = score_a - score_b
loss = torch.nn.functional.binary_cross_entropy_with_logits(
    logits, labels.to(dtype=logits.dtype))  # 1 iff A was preferred.
# No sigmoid() before this loss.

Computing -log(sigmoid(delta)) directly can underflow when the labeled winner scores far below the loser. At $\Delta=-1000$ and target 1, the stable loss is approximately 1000 and its derivative with respect to $\Delta$ is approximately −1. That is a useful corrective gradient, not a reason to discard the pair.

5.2 One score per complete prompt–response sequence

A reward model often adds a scalar head to a contextual LM. Score the prompt and response together, using the same model for both candidates: relevance depends on the instruction.

If the chosen convention scores the last non-padding hidden state, find that state explicitly. hidden[:, -1] selects padding for shorter right-padded sequences. mask.sum(1)-1 is valid only for nonempty, contiguous right padding; it gives the wrong index under left padding. Apply the backbone’s attention mask too: correct pooling does not repair representations that already attended to padding.

Keep tokenization, EOS, and truncation consistent; reject empty sequences and check that truncation preserves the distinguishing text. Follow the model’s pooling convention—last-token pooling is not required by the loss.

6Quiz preparation: a 20-point practice set

20 points; allow 20–25 minutes. Attempt the five questions before revealing solutions. Exact log expressions are acceptable; no code execution is required. These cover slides 1–23, not a specific upcoming quiz.

Review §1 for task splits, §3 for expectations, and §4–§5 for losses/code. The final challenge is optional.

labels = torch.ones_like(score_b)
logits = score_b - score_a
loss = F.binary_cross_entropy_with_logits(torch.sigmoid(logits), labels)

a = score_a.squeeze(-1)
b = score_b
labels = torch.ones_like(b)
assert a.shape == b.shape == labels.shape
loss = F.binary_cross_entropy_with_logits(a - b, labels)

7Optional CPU lab and real-code reading

7.1 Run six small checks before loading an LM

Run 18_reward_model_checks.py with python3 18_reward_model_checks.py. Tested on CPU with PyTorch 2.12.0; numerical checks use float64. Synthetic scores/features test mechanics without downloading models or data.

1. check_loss_and_gradients: the three loss forms agree, gradient signs are correct, and one free-score step widens the winning margin.
2. check_order_offset_and_scale: swapping and relabeling agrees; offsets cancel, scaling changes probabilities, and mismatched shapes are rejected.
3. check_numerical_stability: extreme logits stay finite and double sigmoid changes both loss and gradients.
4. check_disagreement: seven-versus-three labels have a stationary point at probability 0.7, not zero loss.
5. check_padding: last-valid pooling works for left and right padding and rejects an empty sequence.
6. check_shared_scorer: a shared scorer learns feature comparisons through a DataLoader, validates without changing state/gradients, aggregates unequal batches correctly, and has zero task gradient for its common output bias.

def pair_loss(chosen, rejected):
    # One scalar score per sequence; both tensors must have identical shapes.
    if chosen.shape != rejected.shape:
        raise ValueError('chosen/rejected score shapes must match')
    return -F.logsigmoid(chosen - rejected).mean()

A global additive bias cancels from every score difference, explaining the last check. Weight decay or an explicit centering penalty can affect that bias; the pure pair loss cannot identify it.

7.2 PyTorch: one training pass and one validation pass

Each DataLoader item pairs chosen/rejected feature vectors for one prompt. A shared linear scorer maps each to a scalar; these are synthetic features, not token IDs.

18_reward_model_checks.py · run_pairs · complete teaching function

def run_pairs(model, loader, optimizer=None):
    """One CPU pass; optimizer supplied = train, otherwise evaluate."""
    training = optimizer is not None
    model.train(training)  # Controls dropout/normalization behavior.
    loss_sum, correct, total = 0.0, 0, 0
    with torch.set_grad_enabled(training):  # Separately controls autograd.
        for chosen, rejected in loader:
            if training:
                optimizer.zero_grad(set_to_none=True)
            chosen_scores = model(chosen).squeeze(-1)
            rejected_scores = model(rejected).squeeze(-1)
            expected = (chosen.shape[0],)
            if chosen_scores.shape != expected or rejected_scores.shape != expected:
                raise ValueError('the scorer must return one scalar per sequence')
            loss = pair_loss(chosen_scores, rejected_scores)
            if training:
                loss.backward()
                optimizer.step()
            n = chosen_scores.numel()
            loss_sum += loss.detach().item() * n
            correct += (chosen_scores.detach() > rejected_scores.detach()).sum().item()
            total += n
    if total == 0:
        raise ValueError('the loader must contain at least one comparison')
    return {'loss': loss_sum / total, 'accuracy': correct / total, 'pairs': total}

Pass an optimizer to train; omit it to evaluate. Autograd mode is restored on exit; model train/eval mode persists. This excerpt uses variables from check_shared_scorer; the download contains the complete program:

from torch.utils.data import DataLoader, TensorDataset

train_loader = DataLoader(
    TensorDataset(chosen, rejected), batch_size=2, shuffle=False)
optimizer = torch.optim.SGD(scorer.parameters(), lr=0.2)
train_metrics = run_pairs(scorer, train_loader, optimizer)
valid_metrics = run_pairs(scorer, held_loader)  # Separate held-out pairs.

Read the two mode controls separately. model.train(training) switches behavior such as dropout. torch.set_grad_enabled(training) decides whether operations are recorded for backward. Calling eval() alone does not disable gradients, and disabling gradients alone does not switch dropout off. See the PyTorch mode explanation and gradient-mode API.

Keep the graph for backward; detach for logging. loss.detach().item() is a Python number, not a differentiable loss. Clear gradients before each ordinary update; deliberate accumulation needs its own normalization.

Average comparisons, not batches. Weight each batch-mean loss by its pair count; the last batch may be smaller. Accuracy counts pairs, with ties marked incorrect. Training metrics use pre-update predictions along the pass; validation uses a fixed model. The lab checks combined-batch parity and rejects non-scalar reward heads.

check_shared_scorer trains on batches of two and one, evaluates new synthetic pairs, and verifies unchanged validation state/gradients. This checks the loop, not language-task transfer.

7.3 Read the actual reward-training loss

reward_trainer.py · compute_loss · lines 238–247 (verbatim)

        rewards_chosen = model(
            input_ids=inputs["input_ids_chosen"],
            attention_mask=inputs["attention_mask_chosen"],
            return_dict=True,
        )["logits"]
        rewards_rejected = model(
            input_ids=inputs["input_ids_rejected"],
            attention_mask=inputs["attention_mask_rejected"],
            return_dict=True,
        )["logits"]

The same model scores both candidates, and each receives an attention mask. The field named logits holds scalar rewards here, not vocabulary logits. Margin and centering below are optional implementation extensions beyond the slides’ basic loss. Trace where the sequences and masks are constructed before interpreting those outputs.

reward_trainer.py · compute_loss · lines 248–255 (verbatim)

        # calculate loss, optionally modulate with margin
        if "margin" in inputs:
            loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - inputs["margin"]).mean()
        else:
            loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean()

        if self.args.center_rewards_coefficient is not None:
            loss += self.args.center_rewards_coefficient * torch.mean((rewards_chosen + rewards_rejected) ** 2)

With no margin and centering disabled, the else branch is exactly §4’s chosen/rejected loss. The optional margin and centering term change that objective; do not silently include them in the offset-invariance argument. The pairwise term remains offset-invariant, but the centering term does not.