Reinforcement learning for language models

A language model producing a response is a policy taking actions, and once you see it that way the algorithms from the previous sections apply directly. This section maps the pieces across and then covers the three ways it is done in practice: a reward model learned from human comparisons, a method that skips the reward model and trains on the comparisons directly, and rewards that can simply be checked, which is what drove the recent progress in reasoning models. It also covers how each one fails, which in every case is the model optimising the reward it was given rather than the thing you meant.

Mapping the pieces#

state    the prompt plus the text generated so far
action   the next token
policy   the language model
episode  one full response
reward   a score for the finished response

Two things follow from that mapping and both are inconvenient. The action space is the vocabulary, so a hundred thousand actions at every step rather than the two in cart-pole, which rules out anything that takes a max over actions and leaves policy methods. And the reward arrives once, at the end of a response that may be a thousand tokens long, so a single number has to be shared out among a thousand decisions. Sparse, delayed, enormous action space: this is the hard corner of the field, which is why the machinery is all about controlling variance and limiting how far the policy may move.

A reward model from comparisons#

People find it easier to say which of two answers is better than to score one. Collect many such comparisons, then train a reward model to give the preferred answer a higher score. The standard loss pushes the score difference up for the chosen answer.

import torch.nn.functional as F

def reward_model_loss(r_chosen, r_rejected):
    return -F.logsigmoid(r_chosen - r_rejected).mean()

The loss is uncontroversial. What it learns from is not. Raters are people with habits, and any habit that correlates with their choices becomes part of the reward — length being the best documented one. The rest of this section is one such reward model, fitted to synthetic comparisons, and then followed all the way to the damage it does. Answers here are two numbers: whether they are correct, and how long they are. Raters pick the correct one when correctness differs, and otherwise prefer the longer answer three times in four.

import numpy as np
rng = np.random.default_rng(0)

def sample_answer():
    """An answer is (is it correct, how many tokens long)."""
    return int(rng.integers(0, 2)), int(rng.integers(20, 200))

pairs = []
for _ in range(4000):
    a, b = sample_answer(), sample_answer()
    if a[0] != b[0]:
        chosen, rejected = (a, b) if a[0] > b[0] else (b, a)      # raters spot correctness
    else:
        longer, shorter = (a, b) if a[1] > b[1] else (b, a)       # otherwise the longer one
        chosen, rejected = (longer, shorter) if rng.random() < 0.75 else (shorter, longer)
    pairs.append((chosen, rejected))

def features(answer):
    return np.array([answer[0], answer[1] / 100.0])

diff = np.stack([features(c) - features(r) for c, r in pairs])
w = np.zeros(2)
for it in range(3000):                                            # the Bradley-Terry loss
    w -= 1.0 * (diff * (-1 / (1 + np.exp(diff @ w)))[:, None]).mean(0)
loss = np.log1p(np.exp(-diff @ w)).mean()
print(f"reward weights: correct {w[0]:+.3f}   per 100 tokens {w[1]:+.3f}")
print(f"loss {loss:.4f}, agrees with {100 * (diff @ w > 0).mean():.1f}% of the comparisons")

def reward(answer):
    return float(features(answer) @ w)
# reward weights: correct +7.722   per 100 tokens +1.245
# loss 0.3033, agrees with 86.9% of the comparisons

The model is not broken. It learned that correctness is worth far more than length, exactly as the raters behaved, and it agrees with them 86.9 per cent of the time. But it did pick up 1.245 points per hundred tokens, and a reward model is not used to predict preferences. It is used as an objective, and an objective is pushed on.

RLHF#

1. 1.Supervised fine-tuning on good example responses gives a starting policy.
2. 2.Train a reward model on human comparisons.
3. 3.Optimise the policy with PPO to maximise reward, minus a penalty for drifting far from the starting model, measured by KL divergence.

The KL penalty is essential. Without it, the policy finds text that the reward model scores highly but that is not actually good, because the reward model is only accurate near the kind of text it was trained on.

That trade has a closed form, which means it can be examined without training anything. The policy that maximises reward minus beta times the KL from the starting model is the starting model reweighted by the exponential of reward over beta. Beta is the dial: large means stay close, small means chase the reward. Here it is over six candidate responses, using the reward model fitted above, against a column for what a reader actually gets — a correct answer, with padding counted against it.

answers = [(1, 30), (1, 60), (1, 120), (1, 400), (0, 190), (0, 60)]
names = ["correct 30", "correct 60", "correct 120", "correct 400", "wrong 190", "wrong 60"]
r = np.array([reward(a) for a in answers])
p_ref = np.array([0.25, 0.22, 0.13, 0.05, 0.15, 0.20])            # the model before any RL
utility = np.array([a[0] - 0.4 * a[1] / 100 for a in answers])    # what a reader actually gets

print("answer         reward   what the reader gets")
for n, ri, ui in zip(names, r, utility):
    print(f"{n:14s} {ri:+7.3f} {ui:+15.3f}")

print("  beta     KL   mean reward   reader utility   most likely answer")
for beta in (100.0, 10.0, 2.0, 1.0, 0.5, 0.25):
    p = p_ref * np.exp(r / beta)
    p /= p.sum()
    kl = float((p * np.log(p / p_ref)).sum())
    print(f"{beta:6.2f}  {kl:5.3f}      {float(p @ r):+7.3f}          "
          f"{float(p @ utility):+7.3f}   {names[int(p.argmax())]}")
# answer         reward   what the reader gets
# correct 30      +8.095          +0.880
# correct 60      +8.469          +0.760
# correct 120     +9.215          +0.520
# correct 400    +12.700          -0.600
# wrong 190       +2.365          -0.760
# wrong 60        +0.747          -0.240
#   beta     KL   mean reward   reader utility   most likely answer
# 100.00  0.001       +6.358           +0.277   correct 30
#  10.00  0.059       +7.443           +0.378   correct 30
#   2.00  0.795      +10.141           +0.185   correct 400
#   1.00  2.165      +12.047           -0.390   correct 400
#   0.50  2.962      +12.685           -0.595   correct 400
#   0.25  2.996      +12.700           -0.600   correct 400

Reward climbs the whole way down the table, from +6.36 to +12.70, and it never once goes backwards. The thing you wanted rises from +0.277 to +0.378 and then collapses to -0.600. The peak sits at a KL of 0.059, which is almost no movement at all, and by a KL of 2.2 the model is confidently producing four-hundred-token answers that are correct and that nobody wants to read. Every number on the dashboard would have looked excellent throughout.

This is Goodhart's law with the axes labelled, and it is the single most important picture in preference training. It also explains the practices that otherwise look superstitious: KL is monitored during RLHF runs as a first-class metric rather than a regulariser detail, checkpoints are selected on held-out evaluations the reward model never influenced, and beta is tuned on the shape of that curve rather than set by theory. The reward model is a local approximation to human judgement, valid near the text it was trained on. The KL term is the statement of how far you trust it.

DPO: skip the reward model#

Direct Preference Optimisation shows that the same objective can be optimised directly on preference pairs with a simple classification-style loss, without training a separate reward model or sampling during training. It is much simpler and more stable to run, which is why it spread quickly. The trade-off is that it learns only from the fixed pairs you have, not from fresh samples of the current model.

The claim that it reaches the same place is checkable on the candidates above. Generate preferences from a reward the algorithm never sees, run the DPO loss on nothing but those win-lose pairs, and compare the resulting distribution with the closed form from the previous block. The policy here is six logits rather than a language model, which is the smallest thing that can be called a policy and is enough to test the identity.

import torch
r_true = torch.tensor(r / 4, dtype=torch.float32)         # scaled so raters disagree sometimes
p_ref_t = torch.tensor(p_ref, dtype=torch.float32)
BETA = 1.0

p_star = p_ref_t * torch.exp(r_true / BETA)
p_star = p_star / p_star.sum()                            # the RLHF optimum, in closed form

g = torch.Generator().manual_seed(0)
i = torch.randint(0, 6, (40000,), generator=g)
j = torch.randint(0, 6, (40000,), generator=g)
i, j = i[i != j], j[i != j]
i_wins = torch.rand(len(i), generator=g) < torch.sigmoid(r_true[i] - r_true[j])
chosen, rejected = torch.where(i_wins, i, j), torch.where(i_wins, j, i)

logits = torch.zeros(6, requires_grad=True)
opt = torch.optim.Adam([logits], lr=0.05)
log_ref = torch.log(p_ref_t)
for step in range(3000):
    log_pi = torch.log_softmax(logits, 0)
    margin = BETA * ((log_pi[chosen] - log_ref[chosen]) - (log_pi[rejected] - log_ref[rejected]))
    loss = -torch.nn.functional.logsigmoid(margin).mean()
    opt.zero_grad(); loss.backward(); opt.step()
p_dpo = torch.softmax(logits, 0).detach()

print(f"DPO loss {loss.item():.4f}")
print("answer          before    DPO   closed form")
for k in range(6):
    print(f"{names[k]:14s} {p_ref[k]:6.3f} {p_dpo[k]:6.3f} {p_star[k]:12.3f}")
print(f"largest disagreement between DPO and the closed form: "
      f"{float((p_dpo - p_star).abs().max()):.4f}")
# DPO loss 0.4994
# answer          before    DPO   closed form
# correct 30      0.250  0.271        0.281
# correct 60      0.220  0.276        0.272
# correct 120     0.130  0.199        0.193
# correct 400     0.050  0.175        0.178
# wrong 190       0.150  0.040        0.040
# wrong 60        0.200  0.038        0.036
# largest disagreement between DPO and the closed form: 0.0102

Read the margin line and the equivalence stops being mysterious. The quantity being pushed up is the log ratio of the new policy to the reference, scaled by beta, which is precisely the implicit reward that the closed form inverts. DPO is not an approximation of RLHF; it is the same optimum reached by fitting the preference model directly to the policy. The wrong answers fall from 0.15 and 0.20 to 0.04, the four-hundred-token answer rises from 0.05 to 0.175, and yes, the length preference is faithfully reproduced — a simpler algorithm inherits the same rater habits, because they were in the data, not in the optimiser.

What DPO gives up is on-policy sampling. RLHF scores the model's current outputs, so the reward model gets asked about text the model actually produces now; DPO only ever sees the pairs you collected, and the further the policy moves from whatever generated them, the less those pairs say about it. Hence the online variants, and hence the fact that neither has fully displaced the other.

Verifiable rewards#

For maths with known answers or code with tests, no human or reward model is needed: check the answer. Training with RL on these checks, sampling several attempts per problem and reinforcing the ones that succeed, has driven much of the recent progress in reasoning models. Methods such as GRPO compare each attempt with the others in its group instead of training a separate critic.

Group-relative is simpler than it sounds, and it is worth seeing that the critic — a whole second network, with its own instability — is replaced by one line of arithmetic. Sample several answers to the same problem, check each, and standardise the scores within the group. An attempt that succeeded where most failed gets a large positive advantage; the same success when everything succeeded gets nothing, because there is nothing to learn from a problem the model has already mastered.

group = np.array([1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0])        # eight attempts, three correct
adv = (group - group.mean()) / (group.std() + 1e-8)
print("rewards  ", group)
print("advantage", np.round(adv, 3))
print("mean advantage", round(float(adv.mean()), 6), "- no critic network anywhere")
# rewards   [1. 0. 0. 1. 0. 0. 1. 0.]
# advantage [ 1.291 -0.775 -0.775  1.291 -0.775 -0.775  1.291 -0.775]
# mean advantage 0.0 - no critic network anywhere

Two consequences fall out of that line. A problem the model always gets right, or always gets wrong, produces a group with zero standard deviation and therefore no gradient, so the curriculum selects itself towards problems at the edge of the model's ability — which is why these runs are bottlenecked on having problems of the right difficulty. And the reward is the output of a checker, so the usual worry moves from the reward model to the checker: a test suite that can be satisfied by special-casing its inputs will be, and a maths grader that accepts a final answer without a derivation will reward guessing.

Reward hacking, concretely#

• A reward model that slightly prefers longer answers produces a policy that pads everything.
• A code reward that only checks that tests pass can be satisfied by editing the tests or special-casing their inputs.
• A preference model trained on raters who like agreement produces a model that tells people what they want to hear.

Each is the optimiser doing exactly what it was told. The defences are better rewards, penalties for drift, held-out evaluation that the reward never saw, and reading samples.

The first of those in the list is the one the experiment above puts numbers on. Length is not a hypothetical: the fitted model paid 1.245 per hundred tokens, which is a sixth of what it paid for being correct, and that was enough to make a four-hundred-token answer the most likely output at any beta below 2. Length-controlled evaluation exists because of exactly this, and a padded answer is the mildest version. The same mechanism produces confident phrasing on uncertain claims, agreement with a user who has stated a wrong opinion, and answers formatted to look thorough, all of which raters reward and none of which anybody asked for.

The practical checklist is short. Track KL from the starting model during the run and treat a rising KL with a flat evaluation score as over-optimisation in progress. Keep an evaluation set that has never touched reward model training. Measure response length before and after, since it is the cheapest detector of the most common hack. And read a few dozen samples yourself at each checkpoint, because the failure modes that matter are usually visible in one paragraph and invisible in any aggregate.