RLHF with PPO

Week 6 of the residency. When you cannot write a loss function, you learn one from human preferences, then stop the model exploiting it.

The last week: preference optimisation.

Every week so far had a loss function sitting there waiting. Predict the missing word. Predict the digit. Predict the next character. Week 6 starts from a task where there’s no such thing.

Write down the loss for a good summary. You cannot. There’s no target string to compare against, and two perfectly good summaries share almost no tokens.

What you can do is show a person two summaries and ask which they prefer. That’s cheap and reliable, and it gives you comparisons rather than targets.

So: collect preferences, fit a model that predicts them, and then optimise the language model against that learned model. Three stages, and the project implements the third.

The Reward Model

There is the running example one last time. The reward model takes a piece of text and returns a number, trained so that the summary a human preferred scores higher than the one they rejected. It never sees an absolute rating, only which of a pair won.

That number is now the thing being maximised. Which creates the problem the rest of the week is about.

You are optimising a learned approximation of what you wanted. It is wrong in places nobody has looked.

LoRA

Stage three needs four models at once.

• The policy being trained
• A frozen reference it must not drift too far from
• The reward model
• A value model, estimating how good a partial generation is

Holding four full copies of a language model is not something a residency budget does.

LoRA is what makes it fit. Instead of learning a full update to a weight matrix, learn two thin matrices whose product has the same shape. The base weights stay frozen and shared between all four roles, and each role carries only its own small low-rank update.

The PPO Loop

Four models, an advantage calculation and a buffer. Before any of that makes sense, the shape underneath it does.

Generating a sequence is an episode. The policy picks a token, that changes the state, and eventually a reward arrives from the reward model. Everything in the overview above exists to turn one number at the end of a sequence into a learning signal for every token in it.

PPO’s contribution is a way of taking that signal without letting a single update move the policy somewhere unrecoverable.

The ratio compares how likely the new policy is to produce a sequence against how likely the policy that generated it was. Clip that ratio and once the policy has moved far enough on a sample, that sample stops pushing. A batch can then be reused for several steps without the policy running away from the data that produced it.

The KL penalty is a different constraint and it is easy to conflate them. Clipping bounds how far one update moves the policy from the policy that collected the batch. The KL penalty bounds how far training moves it from the model you started with. You can clip perfectly and still walk somewhere useless, one small safe step at a time.

What follows shows that happening.

A Toy Language

To watch reward hacking you need something with grammar, and a reward model that is slightly wrong about what is good.

The language has eight tokens, and each one usually follows the one before it, so its sentences are mostly ascending runs. A small model trained on samples of it is the reference policy, standing in for the supervised model you start RLHF from.

The reward model likes token 3. That is all. Think of it as a preference model that correctly noticed people enjoy a particular thing and has no opinion about anything else, which is roughly how real reward models fail.

import sys

sys.path.insert(0, "..")
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from _style import COLOURS, MUTED, figure, style_axes

VOCAB, LENGTH, START = 8, 6, 0

rng = np.random.default_rng(0)
grammar = np.full((VOCAB, VOCAB), 0.02)
for i in range(VOCAB):
    grammar[i, (i + 1) % VOCAB] = 0.70     # the usual next token
    grammar[i, (i + 2) % VOCAB] = 0.16     # sometimes it skips one
grammar /= grammar.sum(1, keepdims=True)


def sample_corpus(n):
    out = np.zeros((n, LENGTH), int)
    for i in range(n):
        previous = START
        for t in range(LENGTH):
            previous = rng.choice(VOCAB, p=grammar[previous])
            out[i, t] = previous
    return torch.from_numpy(out)


class Bigram(nn.Module):
    def __init__(self):
        super().__init__()
        self.logits = nn.Parameter(torch.zeros(VOCAB, VOCAB))

    def forward(self, previous):
        return self.logits[previous]


corpus = sample_corpus(4000)
torch.manual_seed(0)
reference = Bigram()
optimiser = torch.optim.Adam(reference.parameters(), lr=0.1)
shifted = torch.cat([torch.full((len(corpus), 1), START), corpus[:, :-1]], dim=1)

for _ in range(400):
    loss = F.cross_entropy(reference(shifted).reshape(-1, VOCAB), corpus.reshape(-1))
    optimiser.zero_grad()
    loss.backward()
    optimiser.step()

for p in reference.parameters():
    p.requires_grad_(False)

REWARDED = 3
print(f"reference model trained, cross-entropy {loss.item():.4f}")
print(f"the reward model gives one point per token {REWARDED}, so the maximum is {LENGTH}")

reference model trained, cross-entropy 1.0112
the reward model gives one point per token 3, so the maximum is 6

Two measurements matter from here on. Reward is what training maximises. Fluency is how likely a sequence is under the language the model started from, which nothing in training looks at.

Here is what the reference policy scores on both, and the kind of sentence it produces.

def reward(sequences):
    return (sequences == REWARDED).float().sum(1)


@torch.no_grad()
def generate(model, n, generator):
    sequences = torch.zeros(n, LENGTH, dtype=torch.long)
    log_probs = torch.zeros(n, LENGTH)
    previous = torch.full((n,), START)
    for t in range(LENGTH):
        lp = F.log_softmax(model(previous), -1)
        nxt = torch.multinomial(lp.exp(), 1, generator=generator).squeeze(1)
        sequences[:, t] = nxt
        log_probs[:, t] = lp.gather(1, nxt[:, None]).squeeze(1)
        previous = nxt
    return sequences, log_probs


def log_prob_of(model, sequences):
    previous = torch.cat([torch.full((len(sequences), 1), START), sequences[:, :-1]], dim=1)
    lp = F.log_softmax(model(previous), -1)
    return lp.gather(2, sequences[..., None]).squeeze(-1)


@torch.no_grad()
def fluency(sequences):
    """How likely these are under the language the model started from."""
    return log_prob_of(reference, sequences).sum(1).mean().item()


generator = torch.Generator().manual_seed(9)
samples, _ = generate(reference, 512, generator)
print(f"reference: reward {reward(samples).mean():.2f}, fluency {fluency(samples):.2f}")
print(f"  it says things like {samples[0].tolist()} and {samples[1].tolist()}")

reference: reward 0.84, fluency -6.19
  it says things like [4, 0, 1, 2, 6, 7] and [7, 0, 1, 2, 4, 5]

Now PPO, with the KL coefficient as a dial. At zero there’s nothing holding the policy near where it started; turn it up and drifting gets expensive. Each run below trains for 120 iterations and then reports reward, fluency, and a sample.

def train(kl_coefficient, clip=0.2, iterations=120, batch=512, lr=0.05):
    torch.manual_seed(1)
    policy = Bigram()
    policy.logits.data = reference.logits.data.clone()   # start from the reference
    optimiser = torch.optim.Adam(policy.parameters(), lr=lr)
    generator = torch.Generator().manual_seed(2)

    for _ in range(iterations):
        sequences, old_log_prob = generate(policy, batch, generator)
        scores = reward(sequences)
        with torch.no_grad():
            reference_log_prob = log_prob_of(reference, sequences)

        advantage = scores - scores.mean()
        advantage = advantage / (advantage.std() + 1e-8)

        for _ in range(4):        # reusing the batch is what clipping makes safe
            new_log_prob = log_prob_of(policy, sequences)
            ratio = (new_log_prob - old_log_prob).exp().sum(1)
            kl = (new_log_prob - reference_log_prob).sum(1)
            clipped = torch.min(ratio * advantage,
                                ratio.clamp(1 - clip, 1 + clip) * advantage)
            loss = -(clipped - kl_coefficient * kl).mean()
            optimiser.zero_grad()
            loss.backward()
            optimiser.step()
    return policy


print(f"{'KL coefficient':>16} {'reward':>8} {'fluency':>9}   a sample")
for coefficient in (0.0, 1.0, 3.0, 6.0, 20.0):
    policy = train(coefficient)
    generator = torch.Generator().manual_seed(9)
    samples, _ = generate(policy, 512, generator)
    label = "none" if coefficient == 0 else f"{coefficient:g}"
    print(f"{label:>16} {reward(samples).mean():>8.2f} {fluency(samples):>9.2f}   "
          f"{samples[0].tolist()}")

generator = torch.Generator().manual_seed(9)
samples, _ = generate(reference, 512, generator)
print(f"{'(reference)':>16} {reward(samples).mean():>8.2f} {fluency(samples):>9.2f}   "
      f"{samples[0].tolist()}")

  KL coefficient   reward   fluency   a sample
            none     6.00    -23.08   [3, 3, 3, 3, 3, 3]
               1     6.00    -23.08   [3, 3, 3, 3, 3, 3]
               3     3.00    -13.73   [1, 3, 1, 3, 1, 3]
               6     3.00    -10.13   [2, 3, 2, 3, 2, 3]
              20     1.99    -12.55   [7, 2, 3, 1, 2, 3]
     (reference)     0.84     -6.19   [4, 0, 1, 2, 6, 7]

Conclusion

With no KL penalty the policy scores a perfect 6 out of 6. It does this by saying 3 3 3 3 3 3.

That is the best possible output according to the reward model, and it is not a sentence. Fluency under the original language collapses from about −6 to −23. The policy found the reward model’s blind spot and moved in.

Nothing in the reward number says so. Reward rose monotonically the whole way, so a run watched through that metric looks like a success. It is why RLHF papers plot reward against distance from the reference rather than reward alone.

Turn the penalty up and something better than a compromise appears. The policy settles on 2 3 2 3 2 3: half the maximum reward, and most of the fluency back.

Look at why that sequence in particular. In this language, 2 is usually followed by 3. So the policy found a way to say the rewarded token as often as the grammar allows, rather than as often as arithmetic allows. It is obeying the language and pursuing the reward at the same time.

That is the objective of RLHF in one sequence. Not maximum reward. Not the original model. The best the reward can be served without leaving the language behind.

• A learned reward is an approximation, and hard optimisation finds its mistakes
• Reward alone cannot tell you this is happening
• Clipping bounds one update; the KL penalty bounds the whole run
• Read the samples, not the metric

One implementation note that cost me real time. The PPO ratio has to score token ids, never re-tokenized text. Recording a generated state as its decoded string and re-encoding it looks equivalent and is not: about one state in twenty does not survive decode-then-encode, sometimes coming back with a different number of tokens. The two sides of the ratio were scoring different sequences, so the ratio was not 1 before any gradient step, and PPO’s whole safety argument rests on it being 1 there.

The full project is on GitHub.