Reward Model Math

From language model to reward model: the value head

A reward model is not trained from scratch — it is a pretrained language model with its head swapped out. A normal LM ends in a projection to vocabulary logits: the last hidden state h_t ∈ ℝ^d becomes a distribution over the next token. For a reward model you discard that [d, vocab] unembedding and bolt on a tiny value head: a single linear layer w ∈ ℝ^d (plus bias) that maps a hidden state to one number.

You feed the full sequence — prompt concatenated with response — through the transformer, take the hidden state at the final token (the position that has attended to everything before it), and project it: r(x, y) = w · h_last + b. Shapes: the transformer maps tokens → H: [N, d], you select h_last: [d], and the head produces a scalar [1]. The whole backbone is fine-tuned along with the head. Everything the pretrained model knows about language is reused; training only has to learn the far smaller thing — a direction in hidden space that correlates with human preference.

The Bradley-Terry preference model

Humans are unreliable at absolute scoring (‘rate this 7.3/10’) but reliable at comparison (‘A is better than B’). So the data is pairs: a prompt x, two responses, and a label saying which won. The Bradley-Terry model (1952) is the bridge from a latent scalar score to the probability of such a comparison. It posits that the odds of y_w beating y_l are the ratio of their exponentiated strengths:

P(y_w > y_l | x)  =  exp(r(x, y_w)) / ( exp(r(x, y_w)) + exp(r(x, y_l)) )
             =  σ( r(x, y_w) - r(x, y_l) )

where  σ(z) = 1 / (1 + exp(-z))   is the logistic sigmoid

The second line is the key simplification: divide top and bottom by exp(r_w) and the pair-probability collapses to a sigmoid of the reward difference alone. This is exactly binary logistic regression where the ‘feature’ is r_w - r_l. Preference probability depends on the gap between rewards, never their individual magnitudes — a fact that governs everything that follows.

The pairwise ranking loss

Fitting Bradley-Terry by maximum likelihood gives the loss the whole RM is trained on. For a dataset D of triples (x, y_w, y_l) with y_w the human-chosen response, minimise the negative log-likelihood of the observed preferences:

L(θ)  =  - E_(x, y_w, y_l) ~ D [ log σ( r_θ(x, y_w) - r_θ(x, y_l) ) ]

Read it plainly: push the chosen response’s reward above the rejected one’s, and the wider the correct margin the lower the loss. Because σ saturates, the penalty for getting a pair confidently wrong (large negative margin) grows nearly linearly, while a pair already ranked correctly with a big margin contributes almost nothing — the gradient concentrates on the pairs the model still gets wrong or is unsure about. Both responses pass through the same network with shared weights θ, so one training step nudges the reward surface for both at once. In practice this is one forward pass per response, one subtraction, and a logsigmoid — computed as -log σ(z) = softplus(-z) to avoid overflow.

Why only relative reward is learned

Look again at the loss: it depends on r_w - r_l and nothing else. Add any constant c to every reward the model outputs and the difference is unchanged: (r_w + c) - (r_l + c) = r_w - r_l. The loss is shift-invariant. There is no term anywhere that pins down where zero sits, so the training signal simply cannot determine absolute reward — only reward differences are identifiable.

The practical consequences are sharp. First, an RM output of +4.2 means nothing on its own; only r(A) - r(B) is meaningful. Second, two correctly-trained RMs can disagree wildly in absolute scale and offset while ranking every pair identically. Third — and this is why it matters downstream — whatever consumes the reward (a PPO loop) must not depend on the offset; it typically normalizes rewards to zero mean per batch. You can think of the RM as learning a potential surface defined only up to a constant, like altitude measured without a fixed sea level: the slopes and relative heights are real, the absolute number is a free gauge.

A worked example

Take one prompt and two candidate responses. The RM, in a forward pass, emits r(x, y_w) = 2.0 for the human-preferred response and r(x, y_l) = 1.0 for the other. The margin is z = r_w - r_l = 1.0. Then:

σ(z)  = 1 / (1 + e^-1.0)  = 0.731   # model says 73% chance chosen > rejected
loss   = -log(0.731)          = 0.313   # per-pair NLL

gradient of loss w.r.t. z:  d L / d z = σ(z) - 1 = -0.269
  -> descent moves z UP by 0.269 * lr : widen the gap

Now the instructive cases. If the RM had scored them equal (z = 0): σ(0) = 0.5, loss = -log 0.5 = 0.693 — maximum uncertainty. If it got the pair backwards (z = -1.0, chosen scored lower): σ(-1) = 0.269, loss = 1.313 — over four times the correct-margin loss, and the gradient σ(z) - 1 = -0.731 is much larger, hauling the two rewards apart hard. Notice the gradient magnitude is exactly 1 - σ(z): the model’s own error probability on that pair. Confident-correct pairs (z large) yield near-zero gradient; the learning budget flows to the mistakes.

Training dynamics and data

Training is standard supervised learning over the preference set: sample a batch of pairs, forward both responses, compute the mean pairwise loss, backprop into the shared backbone and head. A few things bite in practice. One epoch is common — reward models overfit fast, memorising annotator quirks and surface features (length, formatting) rather than quality, so held-out preference accuracy is watched closely and training stopped early. Pairs from the same prompt are the unit of signal; the loss never compares responses across different prompts, which is another way to see why absolute scores across prompts are not mutually calibrated. When a prompt has a full ranking of K responses, the common trick (InstructGPT) is to expand it into all C(K,2) pairs and average their losses within the prompt — more sample-efficient and less overfit-prone than treating each pair independently. Learning rates stay small: you are steering a large pretrained model with a weak scalar signal, not reshaping it.

Calibration and reward-model accuracy

How good is a reward model? The honest, model-agnostic metric is preference accuracy: on held-out pairs, how often does r(y_w) > r(y_l) agree with the human label? Strong RMs land around 65–75% — and the ceiling is real, because human annotators only agree with each other perhaps 70–80% of the time. An RM that scored 99% would be fitting noise. Chasing accuracy past inter-annotator agreement is chasing a phantom.

Calibration is the finer question: when Bradley-Terry says σ(r_w - r_l) = 0.8, is the chosen response actually preferred ~80% of the time? Because the model is fit by proper (logistic) likelihood, it is reasonably calibrated in-distribution — the reward gap is a meaningful confidence, so a large gap really is a more decisive preference than a tiny one. The danger is out-of-distribution: on the weird, off-manifold text an RL policy learns to generate, calibration and even ranking can collapse. A reward model is trustworthy where its training pairs lived, and increasingly fictional the further the policy roams from them.

Ensembling reward models

Because a single RM is a noisy, shift-free estimate that degrades off-distribution, a robust move is to train several — different seeds, data shuffles, or even different base checkpoints — and combine them. The simplest ensemble averages the scores: r_ens(x, y) = (1/K) Σ_k r_k(x, y). Averaging cancels the idiosyncratic errors of any one model and gives a smoother, lower-variance reward surface for the policy to optimise against.

The subtler payoff is the disagreement across the ensemble. The variance Var_k[ r_k(x, y) ] is a cheap uncertainty estimate: where the members agree, the reward is trustworthy; where they scatter, you are likely off-distribution and the reward is not to be believed. This powers uncertainty-penalised optimisation — subtract a multiple of the ensemble std from the reward — which attacks reward hacking by taxing the policy for wandering where the RMs are unsure. The cost is real (K forward passes per rollout), so small ensembles of 3–5, or cheaper approximations like reward-head dropout, are typical.

Using the reward model to score RL rollouts

Trained and frozen, the RM becomes the environment the policy plays against. In PPO-style RLHF the loop is: sample a prompt, let the policy generate a response, and score it with the RM — one forward pass yielding the scalar r(x, y) as the terminal reward for that whole generation. The optimised objective is that reward minus a per-token KL penalty that keeps the policy near its supervised-fine-tuned starting point:

reward(x, y) = r_RM(x, y)  -  β * KL( π_θ(y|x) || π_ref(y|x) )

Here the shift-invariance from earlier comes home to roost. Since the RM’s absolute scale is an arbitrary gauge, the raw scores are whitened per batch — (r - mean) / std — before they drive the advantage estimate, so the offset the RM never learned cannot bias the update, and the scale stays stable across training. The RM being merely a learned, imperfect, in-distribution proxy is precisely why the KL leash and normalisation exist: left unchecked, the policy will find inputs that score high on the RM but are junk to humans — reward hacking. That failure, and how reward magnitudes are scaled and clipped, are sibling topics; here the focus has been the RM itself — how the scalar it emits is trained from preferences.