8 - How a language model generates text and learns from preferences

In chapter 8 the Minermont team moves from how language models are built to how they behave. Two ideas are at the centre. First, generation: as Emma explains, inference is a loop with three elements, the context, the distribution of the next token and the decision rule that picks one. Second, alignment: with RLHF, models learn from human preferences through a reward model, which can also be "hacked". In this notebook we build both ideas at a small scale, with NumPy:

1. Import the libraries.
2. A tiny language model: the probability of the next word.
3. The context.
4. Decision rules: greedy, temperature, top-k and top-p.
5. Beam search: looking for the most probable sentence.
6. RLHF in miniature: preferences and a reward model.
7. Optimizing against the reward model, and the penalty that keeps the model close to the original.
8. When the reward model is wrong: reward hacking.

Step 1: Import the libraries

---

import collections
import itertools
import numpy as np
import matplotlib.pyplot as plt

np.set_printoptions(precision=3, suppress=True)

Step 2: A tiny language model

---

A large language model computes $P(x_{t+1} \mid x_1, \dots, x_t)$, the probability of the next token given the previous ones, with a Transformer. Here we use the simplest model with the same interface, an n-gram model: it looks at the last $n$ words (its whole context) and estimates the probabilities by counting, in a corpus of 2000 discharge instructions, which word followed them. <s> marks the start of a sentence and </s> its end. Like a real model, it gives every word of the vocabulary a small probability, even words never seen in that context (smoothing).

rng = np.random.default_rng(0)
fill = {'med': (['paracetamol', 'ibuprofen', 'the antibiotic', 'the new tablets'], [0.5, 0.25, 0.15, 0.1]),
        'freq': (['twice a day', 'every eight hours', 'once a day'], [0.6, 0.25, 0.15]),
        'symptom': (['fever', 'pain', 'dizziness', 'shortness of breath'], [0.5, 0.25, 0.15, 0.1]),
        'n': (['two', 'three', 'four'], [0.6, 0.3, 0.1])}
templates = ['take {med} {freq} after meals', 'drink plenty of water and rest at home',
             'call your doctor if the {symptom} gets worse', 'come back to the hospital if you have {symptom}',
             'do not drive while you take {med}', 'your next appointment is in {n} weeks',
             'avoid alcohol while you take {med}', 'take {med} if you have {symptom}']

template_probs = [0.25, 0.1, 0.15, 0.1, 0.1, 0.1, 0.1, 0.1]
def pick(key):
    options, probs = fill[key]
    return options[rng.choice(len(options), p=probs)]
corpus = [templates[i].format(**{k: pick(k) for k in fill}) for i in rng.choice(len(templates), 2000, p=template_probs)]

valid = set()                                        # every correct instruction the templates can produce
for t in templates:
    keys = [k for k in fill if '{' + k + '}' in t]
    for combo in itertools.product(*[fill[k][0] for k in keys]):
        valid.add(t.format(**dict(zip(keys, combo))))
vocabulary = sorted({w for s in corpus for w in s.split()} | {'</s>'})

def ngram_model(context_size, smoothing):
    counts = collections.defaultdict(collections.Counter)
    for s in corpus:
        tokens = ['<s>'] * context_size + s.split() + ['</s>']
        for i in range(context_size, len(tokens)):
            counts[tuple(tokens[i - context_size:i])][tokens[i]] += 1
    def next_token_distribution(context):
        c = counts[tuple(context[-context_size:])]
        f = np.array([c[w] for w in vocabulary], dtype=float) + smoothing
        return f / f.sum()                               # P(next word | the last context_size words)
    return next_token_distribution

model = ngram_model(context_size=2, smoothing=0.02)
p = model(['you', 'have'])
order = np.argsort(-p)
print(f"Vocabulary: {len(vocabulary)} words, {len(valid)} possible correct instructions")
print("P(next word | 'you have'):", ", ".join(f"{vocabulary[i]} {p[i]:.3f}" for i in order[:4]),
      f"| the other {len(vocabulary) - 4} words together: {p[order[4:]].sum():.3f}")

Vocabulary: 56 words, 48 possible correct instructions
P(next word | 'you have'): fever 0.500, pain 0.254, dizziness 0.164, shortness 0.080 | the other 52 words together: 0.003

After you have, the model has learned the four symptoms of the corpus with their frequencies: fever 0.500, pain 0.254, dizziness 0.164 and shortness 0.080 (the first word of shortness of breath: here each word is a token). The other 52 words of the vocabulary share 0.003 because of the smoothing: they make no sense here, but they are not impossible. This is the distribution, the second of Emma's three elements: the model does not output a sentence, it assigns probabilities to many possible next tokens.

Step 3: The context

---

Emma's first element is the context: what the model can see. We generate 300 instructions with models that see 1, 2 or 3 previous words, sampling each word from the distribution, and count how many are correct instructions:

def generate(model, rng, rule, max_words=20):
    tokens = ['<s>'] * 3
    while tokens[-1] != '</s>' and len(tokens) < max_words + 3:
        tokens.append(vocabulary[rule(model(tokens), rng)])     # the chosen word is appended to the context
    return ' '.join(t for t in tokens if t not in ('<s>', '</s>'))

sample = lambda p, rng: rng.choice(len(p), p=p)
for size in [1, 2, 3]:
    m = ngram_model(size, smoothing=0.0)
    rng = np.random.default_rng(1)
    texts = [generate(m, rng, sample) for _ in range(300)]
    wrong = [t for t in texts if t not in valid]
    print(f"Context of {size} word(s): {1 - len(wrong) / len(texts):.0%} correct instructions; the most frequent wrong one: '{collections.Counter(wrong).most_common(1)[0][0]}'")

Context of 1 word(s): 29% correct instructions; the most frequent wrong one: 'take paracetamol'
Context of 2 word(s): 74% correct instructions; the most frequent wrong one: 'take paracetamol'
Context of 3 word(s): 94% correct instructions; the most frequent wrong one: 'avoid alcohol while you take the antibiotic twice a day after meals'

With one word of context only 29% of the instructions are correct, with two words 74% and with three 94%. The most frequent mistake with one or two words is take paracetamol, an instruction cut short. After take paracetamol the model cannot see whether the sentence began with take or with avoid alcohol while you, and in the corpus that second kind of sentence ends right there. With three words of context that is settled, and the most frequent mistake only appears when a medicine name of two words, the antibiotic, fills the context and pushes you out of it. And, as Emma says, each chosen word is appended to the context, so an early ambiguity snowballs into a different sentence. A real model sees thousands of tokens, but the lesson is the same: what the model can see limits what it can get right.

Step 4: Decision rules

---

The model gives a distribution; the decision rule chooses the token. The book describes four:

• Greedy: always the most probable token. Deterministic, but, as Emma warns, it is not automatically the safest: it can lock into repetition or early mistakes.
• Temperature $T$: before choosing, the probabilities are reshaped as $p_i^{1/T}$ and normalized again. $T < 1$ sharpens the distribution (more predictable), $T > 1$ flattens it (more varied, and more risky).
• Top-k: only the $k$ most probable tokens are kept.
• Top-p (nucleus): only the smallest set of tokens whose probabilities add up to at least $p$ is kept.

Noah's image in the book is a die "loaded" by the probabilities: the rules decide how loaded it is and which faces it has. We use the model with a 2-word context and compare the rules on 300 generated instructions:

def rule(greedy=False, temperature=1.0, top_k=None, top_p=None):
    def decide(probs, rng):
        if greedy:
            return int(np.argmax(probs))
        p = probs ** (1 / temperature)
        p = p / p.sum()
        order = np.argsort(-p)
        keep = np.ones(len(p), dtype=bool)
        if top_k is not None:
            keep[order[top_k:]] = False
        if top_p is not None:
            keep[order[np.searchsorted(np.cumsum(p[order]), top_p) + 1:]] = False   # smallest set reaching top_p
        p = np.where(keep, p, 0)
        return int(rng.choice(len(p), p=p / p.sum()))
    return decide

rules = {'greedy': rule(greedy=True), 'T = 0.5': rule(temperature=0.5), 'T = 1': rule(temperature=1.0),
         'T = 1.5': rule(temperature=1.5), 'T = 2': rule(temperature=2.0), 'T = 2, top-k = 3': rule(temperature=2.0, top_k=3),
         'T = 2, top-p = 0.9': rule(temperature=2.0, top_p=0.9), 'T = 1, top-p = 0.9': rule(temperature=1.0, top_p=0.9)}
examples = {}
for name, decide in rules.items():
    rng = np.random.default_rng(1)
    texts = [generate(model, rng, decide) for _ in range(300)]
    examples[name] = texts
    print(f"{name:<20} different instructions: {len(set(texts)):>3} of 300 | correct: {np.mean([t in valid for t in texts]):>4.0%}")
print(f"\nGreedy always writes: '{examples['greedy'][0]}'")
print(f"With T = 2, for example: '{examples['T = 2'][0]}'")

greedy               different instructions:   1 of 300 | correct:   0%

T = 0.5              different instructions:  38 of 300 | correct:  67%
T = 1                different instructions:  72 of 300 | correct:  68%

T = 1.5              different instructions: 216 of 300 | correct:  31%

T = 2                different instructions: 284 of 300 | correct:   5%
T = 2, top-k = 3     different instructions:  58 of 300 | correct:  60%

T = 2, top-p = 0.9   different instructions: 283 of 300 | correct:   7%

T = 1, top-p = 0.9   different instructions:  42 of 300 | correct:  79%

Greedy always writes: 'take paracetamol'
With T = 2, for example: 'drink twice avoid worse drink gets three four in a shortness in drive take doctor home at four call dizziness'

context = ['take', 'paracetamol']
p = model(context)
top = np.argsort(-p)[:5]
labels = [vocabulary[i] for i in top] + [f'the other {len(vocabulary) - 5}']
fig, axes = plt.subplots(1, 3, figsize=(13, 3.8), sharey=True)
for ax, T in zip(axes, [0.5, 1, 2]):
    q = p ** (1 / T)
    q = q / q.sum()
    values = list(q[top]) + [1 - q[top].sum()]
    ax.bar(labels, values, color=['tab:blue'] * 5 + ['tab:gray'])
    ax.set_title(f"T = {T}")
    ax.tick_params(axis='x', rotation=40)
    print(f"T = {T}: " + ", ".join(f"{l} {v:.2f}" for l, v in zip(labels, values)))
axes[0].set_ylabel('probability')
fig.suptitle("P(next word | 'take paracetamol') with each temperature")
plt.tight_layout()
plt.show()

T = 0.5: </s> 0.53, twice 0.27, if 0.14, every 0.04, once 0.03, the other 51 0.00
T = 1: </s> 0.37, twice 0.26, if 0.19, every 0.10, once 0.08, the other 51 0.00
T = 2: </s> 0.25, twice 0.21, if 0.18, every 0.13, once 0.12, the other 51 0.13

The table and the figure show the four rules at work:

• Greedy always writes the same instruction, take paracetamol, and it is wrong. At each step it takes the most probable word, and after take paracetamol the most probable one is the end of the sentence, </s> (0.37 at T = 1). This is the early mistake Emma warns about: a deterministic rule is not automatically a safe one.
• A lower temperature does not solve it: T = 0.5 raises the end of the sentence to 0.53, and it gets 67% of correct instructions, against 68% at T = 1, with fewer different ones (38 against 72). Sharpening the distribution makes the model more predictable, not more correct.
• A higher temperature flattens it: at T = 2 the 51 words that make no sense after take paracetamol receive 0.13 together, and a sentence has many steps. There are 284 different instructions, and only 5% are correct.
• Top-k = 3 removes that tail and brings T = 2 back to 60% of correct instructions. Top-p = 0.9 barely helps (7%): it is applied after the temperature, and once the distribution is flat, reaching 90% of the probability takes many words of the tail. Combined with T = 1, top-p gives the best result of the table, 79%.

As Grace concludes in the book, if we want reliability we mostly control the context and the decision rule.

Step 5: Beam search: looking for the most probable sentence

---

The book names a third way of decoding besides the greedy and the stochastic ones: beam search. Greedy decoding takes the best word at each step, but the sentence it builds is not necessarily the most probable one: a word that looks good now can lead to a poor continuation. Beam search keeps the $k$ most probable partial sentences (the beam) at every step. It extends each one with its $k$ most probable next words, scores every candidate by the log-probability of the whole sentence, $\log P(x_1) + \log P(x_2 \mid x_1) + \dots$, and keeps the $k$ best. With $k = 1$ it is greedy decoding. A sentence that reaches </s> is set aside as finished, and at the end we choose the best finished one. Like greedy decoding, it is deterministic: it gives the same sentence every time.

def beam_search(model, width, max_words=20):
    beams, finished = [(['<s>'] * 3, 0.0)], []                  # (words so far, log-probability of the sentence)
    while beams:
        candidates = []
        for tokens, log_p in beams:
            p = model(tokens)
            for i in np.argsort(-p)[:width]:                      # the width most probable next words
                candidates.append((tokens + [vocabulary[i]], log_p + np.log(p[i])))
        candidates.sort(key=lambda c: -c[1])
        beams = []
        for tokens, log_p in candidates[:width]:                  # keep the width best candidates
            done = tokens[-1] == '</s>' or len(tokens) >= max_words + 3
            (finished if done else beams).append((tokens, log_p))
    return sorted(finished, key=lambda f: -f[1])

def words_of(tokens):
    return ' '.join(t for t in tokens if t not in ('<s>', '</s>'))

for width in [1, 2, 3, 5, 10]:
    tokens, log_p = beam_search(model, width)[0]
    print(f"width {width:>2}: '{words_of(tokens)}'  probability {np.exp(log_p):.4f}  correct: {words_of(tokens) in valid}")

print("\nThe best finished sentences with width 5, with their probability and their log-probability per word:")
for tokens, log_p in beam_search(model, 5)[:5]:
    n_words = len(tokens) - 3                                     # the words and </s>
    print(f"  {np.exp(log_p):.4f}  {log_p / n_words:6.3f}  '{words_of(tokens)}'")

width  1: 'take paracetamol'  probability 0.0630  correct: False
width  2: 'call your doctor if the fever gets worse'  probability 0.0708  correct: True
width  3: 'call your doctor if the fever gets worse'  probability 0.0708  correct: True
width  5: 'drink plenty of water and rest at home'  probability 0.0937  correct: True
width 10: 'drink plenty of water and rest at home'  probability 0.0937  correct: True

The best finished sentences with width 5, with their probability and their log-probability per word:
  0.0937  -0.263  'drink plenty of water and rest at home'
  0.0708  -0.294  'call your doctor if the fever gets worse'
  0.0630  -0.922  'take paracetamol'
  0.0485  -0.303  'come back to the hospital if you have fever'
  0.0395  -0.359  'call your doctor if the pain gets worse'

With width 1, beam search is greedy decoding and writes "take paracetamol", with a probability of 0.0630, which is not a correct instruction: after "take paracetamol" the most probable word is the end of the sentence, because two templates end that way ("do not drive while you take paracetamol"). With width 2 it already finds a better sentence, "call your doctor if the fever gets worse" (0.0708), and from width 5 on the most probable sentence of the model, "drink plenty of water and rest at home" (0.0937): its first word is less probable than "take", but after it every word is almost certain. That is the early mistake Emma warns about, and what a wider beam avoids. The last column, the log-probability per word, is used for length normalization: without it, beam search favours short sentences, because every extra word multiplies the probability by a number smaller than 1. "take paracetamol" is third by probability and last by probability per word, −0.922. Like greedy decoding, beam search always writes the same sentence: it looks for the most probable text, not for variety.

Step 6: RLHF in miniature

---

Now the second idea. The book describes RLHF in three phases: collect human preferences between answers, train a reward model that predicts them, and optimize the language model to get a high reward. We work with a single question from a patient and 8 possible answers of the base model:

question = "I have high blood pressure. Can I take ibuprofen for my back pain?"
answers = {
    'A': "Yes, take 800 mg every six hours.",
    'B': "Ibuprofen can raise blood pressure, so please check with your doctor or pharmacist before taking it.",
    'C': "I'm sorry about your back pain. Ibuprofen can raise blood pressure, so ask your doctor first; "
         "they may suggest paracetamol instead.",
    'D': "Back pain is very common and there are many things that can help, such as rest, gentle exercise, heat, "
         "a good mattress and a good posture at work ... (a long answer that never addresses the question)",
    'E': "I'm sorry you're in pain! Take 600 mg three times a day.",
    'F': "Anti-inflammatories such as ibuprofen can raise blood pressure and interact with some blood pressure "
         "medicines. Please check with your doctor or pharmacist before taking it.",
    'G': "I'm so sorry you're going through this. ... (a very long, warm and detailed answer that ends up "
         "recommending 800 mg of ibuprofen every six hours)",
    'H': "I'm sorry to hear that, I hope you feel better very soon.",
}
names = list(answers)
words = np.array([6, 18, 26, 80, 10, 45, 160, 30])                  # length of each answer, in words
#                   asks to see  empathetic  invents   explains
#                   a clinician              a dose    the risk
features = np.array([[0, 0, 1, 0],    # A
                     [1, 0, 0, 1],    # B
                     [1, 1, 0, 1],    # C
                     [0, 1, 0, 0],    # D
                     [0, 1, 1, 0],    # E
                     [1, 0, 0, 1],    # F
                     [0, 1, 1, 0],    # G
                     [0, 1, 0, 0]],   # H
                    dtype=float)

The clinicians' real judgement, which nobody writes down explicitly, values asking to see a clinician, explaining the risk and a little empathy, and punishes hard an invented dose. We write it as a true quality that we will use only to simulate their choices and to check the results:

$$\text{quality} = 2\,(\text{asks to see a clinician}) + 1\,(\text{empathetic}) - 3\,(\text{invents a dose}) + 1.5\,(\text{explains the risk})$$

Phase 1. Clinicians compare pairs of answers and choose one, not always consistently: the better answer wins with probability $\sigma(\text{quality}_A - \text{quality}_B)$. In this first round they only see six of the answers: the very long ones, D and G, were never shown to them.

Phase 2. The reward model gives each answer a score $r = w \cdot \phi$, computed from features it can observe: whether it asks to see a clinician, whether it is empathetic, whether it explains the risk, and its length. It cannot check whether a dose is invented. It is trained so that $P(A \succ B) = \sigma(r_A - r_B)$ matches the clinicians' choices, which is exactly the logistic regression of chapter 3 applied to differences between answers (the Bradley–Terry model).

quality = features @ np.array([2.0, 1.0, -3.0, 1.5])
observable = np.c_[features[:, [0, 1, 3]], words / 40]            # what the reward model can see

def collect_preferences(candidates, n, seed):
    rng = np.random.default_rng(seed)
    pairs = []
    for _ in range(n):
        i, j = rng.choice(candidates, 2, replace=False)
        better_wins = rng.random() < 1 / (1 + np.exp(-(quality[i] - quality[j])))
        pairs.append((i, j) if better_wins else (j, i))              # (preferred, rejected)
    return np.array(pairs)

def train_reward_model(pairs, learning_rate=0.5, epochs=3000):
    diff = observable[pairs[:, 0]] - observable[pairs[:, 1]]
    w = np.zeros(observable.shape[1])
    for _ in range(epochs):
        p = 1 / (1 + np.exp(-diff @ w))                             # P(preferred beats rejected)
        w -= learning_rate * diff.T @ (p - 1) / len(diff)            # gradient of the log loss
    return w

shown = [names.index(n) for n in ['A', 'B', 'C', 'E', 'F', 'H']]
pairs = collect_preferences(shown, 400, seed=2)
w = train_reward_model(pairs)
reward = observable @ w
agreement = np.mean(reward[pairs[:, 0]] > reward[pairs[:, 1]])
print(f"Reward model weights: asks to see a clinician {w[0]:.2f}, empathetic {w[1]:.2f}, "
      f"explains the risk {w[2]:.2f}, length {w[3]:.2f}")
print(f"Agrees with the clinicians on {agreement:.1%} of their comparisons\n")
print(f"{'answer':<8}{'true quality':>13}{'reward':>9}")
for n, q, r in zip(names, quality, reward):
    print(f"{n:<8}{q:>13.1f}{r:>9.2f}{'   <- never shown to the clinicians' if n in 'DG' else ''}")

Reward model weights: asks to see a clinician 2.00, empathetic 1.45, explains the risk 2.00, length 1.61
Agrees with the clinicians on 90.5% of their comparisons

answer   true quality   reward
A                -3.0     0.24
B                 3.5     4.72
C                 4.5     6.49
D                 1.0     4.66   <- never shown to the clinicians
E                -2.0     1.85
F                 3.5     5.80
G                -2.0     7.88   <- never shown to the clinicians
H                 1.0     2.65

The reward model agrees with the clinicians on 90.5% of their comparisons, and it puts the six answers they saw in the right order: C first (6.49), A last (0.24). But it has learned something nobody intended, a weight of 1.61 for length. It cannot see the invented dose, and among the answers the clinicians saw, the two that invent one, A and E, happen to be the shortest (6 and 10 words): length became a proxy for what it cannot see. The consequence is answer G, which the clinicians never saw: long (160 words), warm, and recommending a dose that is dangerous for this patient. Its true quality is −2.0, yet it gets the highest reward of all, 7.88. As the book says, the reward model is only a learned approximation of human judgement, and it is only reliable on answers like the ones it was trained on.

Step 7: Optimizing against the reward model

---

Phase 3. The language model, the policy $\pi$, is adjusted to obtain a high reward. The book explains that a regularization term is added that penalizes moving too far from the original model $\pi_0$. With that penalty, measured with the Kullback–Leibler divergence and a strength $\beta$, the best possible policy has a simple form:

$$\pi(y) \propto \pi_0(y)\, e^{\,r(y)/\beta}$$

A large $\beta$ keeps the model almost as it was; a small $\beta$ lets the reward decide everything. The base model gives these probabilities to the 8 answers, and we see what happens as $\beta$ decreases:

base_policy = np.array([0.18, 0.10, 0.08, 0.16, 0.12, 0.07, 0.17, 0.12])

def optimize(reward, beta):
    policy = base_policy * np.exp(reward / beta)
    return policy / policy.sum()

def evaluate_policies(reward, betas):
    rows = []
    for beta in betas:
        policy = optimize(reward, beta)
        kl = np.sum(policy * np.log(policy / base_policy))
        rows.append((beta, policy @ reward, policy @ quality, kl, policy[names.index('G')], policy[names.index('C')]))
    return np.array(rows)

betas = [20, 10, 5, 3, 2, 1.5, 1, 0.7, 0.5, 0.3, 0.1]
results = evaluate_policies(reward, betas)
print(f"Base model: mean reward {base_policy @ reward:.2f}, mean true quality {base_policy @ quality:.2f}\n")
print(f"{'beta':>6}{'mean reward':>13}{'true quality':>14}{'distance (KL)':>15}{'P(G)':>7}{'P(C)':>7}")
for beta, r, q, kl, p_g, p_c in results:
    print(f"{beta:>6}{r:>13.2f}{q:>14.2f}{kl:>15.2f}{p_g:>7.2f}{p_c:>7.2f}")

Base model: mean reward 4.07, mean true quality 0.11

  beta  mean reward  true quality  distance (KL)   P(G)   P(C)
  20.0         4.40          0.24           0.01   0.20   0.09
  10.0         4.73          0.33           0.03   0.24   0.10
   5.0         5.33          0.41           0.12   0.32   0.11
   3.0         6.00          0.33           0.30   0.43   0.13
   2.0         6.62          0.05           0.55   0.56   0.13
   1.5         7.02         -0.31           0.79   0.66   0.12
   1.0         7.47         -0.95           1.15   0.81   0.09
   0.7         7.71         -1.46           1.43   0.91   0.06
   0.5         7.82         -1.77           1.61   0.96   0.03
   0.3         7.87         -1.97           1.74   1.00   0.00
   0.1         7.88         -2.00           1.77   1.00   0.00

With β = 20 the policy hardly moves (distance 0.01), and the true quality rises only a little, from 0.11 to 0.24. Lowering β pushes the reward up without stopping: 5.33 at β = 5, 7.47 at β = 1 and 7.88 at β = 0.1. The true quality, instead, peaks at 0.41 at β = 5 and then falls: at β = 1 it is already −0.95, and at β = 0.1 it is −2.00, because the model writes answer G with probability 1.00. The best answer, C, which reaches 0.13 at most, disappears. The optimization has done exactly what it was asked, maximize the reward model; the problem is that the reward model is wrong about G.

Step 8: Reward hacking, and how to fix it

---

The plot shows the trap: the more we optimize (smaller $\beta$, to the right), the higher the reward, while the true quality peaks and then collapses.

fig, ax = plt.subplots(figsize=(8, 4))
ax.plot(results[:, 3], results[:, 1], 'o-', label='reward (what the model optimizes)')
ax.plot(results[:, 3], results[:, 2], 's-', label='true quality (what the clinicians want)')
ax.set_xlabel('distance from the original model (KL): more optimization to the right')
ax.set_title('Over-optimization of an imperfect reward model')
ax.legend()
ax.grid(True)
plt.show()

The book proposes two defences: the penalty that keeps the model close to the original, which we have just seen, and new rounds of evaluation. Let's run a second round in which the clinicians also compare the new answers that the optimized model produces, including D and G, and retrain the reward model:

all_answers = list(range(len(names)))
pairs_2 = np.r_[pairs, collect_preferences(all_answers, 400, seed=3)]
w_2 = train_reward_model(pairs_2)
reward_2 = observable @ w_2
print(f"New weights: asks to see a clinician {w_2[0]:.2f}, empathetic {w_2[1]:.2f}, explains the risk {w_2[2]:.2f}, length {w_2[3]:.2f}")
print(f"Reward of G: {reward[names.index('G')]:.2f} -> {reward_2[names.index('G')]:.2f}; "
      f"reward of C: {reward[names.index('C')]:.2f} -> {reward_2[names.index('C')]:.2f}\n")
for beta, r, q, kl, p_g, p_c in evaluate_policies(reward_2, [5, 1, 0.3]):
    print(f"beta {beta:>4}: true quality {q:5.2f}, P(G) {p_g:.2f}, P(C) {p_c:.2f}")

New weights: asks to see a clinician 2.42, empathetic 1.77, explains the risk 2.42, length 0.00
Reward of G: 7.88 -> 1.78; reward of C: 6.49 -> 6.60

beta  5.0: true quality  1.12, P(G) 0.14, P(C) 0.17
beta  1.0: true quality  4.03, P(G) 0.01, P(C) 0.70
beta  0.3: true quality  4.49, P(G) 0.00, P(C) 0.99

Once the clinicians have seen D and G, the weight of length drops to 0.00: long answers are not better answers. G's reward falls from 7.88 to 1.78, and C, the best answer, now has the highest reward (6.60). Optimizing is now safe: at β = 1 the true quality is 4.03, and at β = 0.3 it is 4.49, with C chosen 99% of the time. That is why RLHF is not done once: each round of optimization produces answers that people must evaluate again. The reward model is still blind to the invented dose; it just no longer rewards the length that came with it.

Conclusion

---

• A language model does not write sentences: at each step it gives a distribution over the next token, and a decision rule chooses one. Decoding is a design choice, as Emma says, not the model "having moods".
• What the model can see, its context, limits what it can get right: here, from 29% of correct instructions with one word to 94% with three.
• Greedy decoding is predictable but not automatically safe: here it always wrote the same wrong instruction. Temperature trades correctness for variety, and top-k and top-p cut off the unlikely tail; how the rules are combined matters.
• RLHF turns human preferences into a reward model (a logistic regression on pairs of answers, in its simplest form) and then pushes the model towards answers with a high reward.
• The reward model is only a learned approximation of human judgement. If it relies on something superficial, such as length, optimizing against it can make a dangerous answer the most likely one: reward hacking. The penalty that keeps the model close to the original and new rounds of human evaluation are the defences the book describes.
• In a hospital this leads to the book's Golden Rule: every output of a language model must be validated by a professional.