LLMs · RLHF · Sahaj's blog

Reinforcement learning from human feedback in three stages, with the reward-model step built from scratch: a Bradley–Terry model learning to prefer the better of two answers from pairwise labels.

InstructGPT (Ouyang et al., 2022) popularised a three-stage recipe:

SFT on demonstrations (last post).
Reward model (RM): people compare pairs of model answers. A model learns a scalar score r(x, y) so that preferred answers score higher. The loss is the Bradley–Terry form, $-\log\sigma(r(y_w) - r(y_l))$ .
RL (PPO): optimise the policy to maximise reward, minus a KL penalty that keeps it close to the SFT model, so it doesn’t drift into nonsense that happens to fool the RM.

The idea goes back to Christiano et al. (2017), learning rewards from human comparisons.

To see stage 2 work, I made a toy setup. Each “answer” is a feature vector, and a hidden true quality decides which answer of a pair a simulated labeller prefers, with some label noise. A linear reward model learns only from those pairwise choices.

import torch, torch.nn.functional as F
torch.manual_seed(0)

d = 8
true_w = torch.randn(d)                        # hidden "true quality" direction
def answers(n): return torch.randn(n, d)

A, B = answers(2000), answers(2000)
p_A = torch.sigmoid((A - B) @ true_w)          # labeller prefers higher true quality, noisily
prefer_A = torch.bernoulli(p_A).bool()
chosen = torch.where(prefer_A[:, None], A, B); rejected = torch.where(prefer_A[:, None], B, A)

w = torch.zeros(d, requires_grad=True)
opt = torch.optim.Adam([w], lr=0.05)
for step in range(300):
    loss = -F.logsigmoid(chosen @ w - rejected @ w).mean()      # Bradley-Terry
    opt.zero_grad(); loss.backward(); opt.step()

A2, B2 = answers(5000), answers(5000)
agree = (((A2 - B2) @ w) > 0) == (((A2 - B2) @ true_w) > 0)
print(f"final pairwise loss {loss.item():.3f}")
print(f"reward model ranks unseen pairs like the true quality {agree.float().mean():.1%} of the time")
print(f"cosine(learned w, true w) = {F.cosine_similarity(w.detach(), true_w, dim=0):.3f}")
print(f"label noise: labeller disagreed with true quality on {((p_A > 0.5) != prefer_A).float().mean():.1%} of training pairs")

Output:

final pairwise loss 0.250
reward model ranks unseen pairs like the true quality 98.5% of the time
cosine(learned w, true w) = 0.999
label noise: labeller disagreed with true quality on 11.3% of training pairs

LLMs · RLHF

Sources