RLHF from scratch: SFT, a reward model, and PPO on a KL leash

Pretraining gives a language model knowledge. It does not give it manners. A base model has read the internet and can continue any text fluently, but ask it "give me three tips to sleep better" and it may cheerfully reply with five tips to eat healthier and an invitation to subscribe to a newsletter — it treats your instruction as a document to keep writing. RLHF is how that raw capability gets turned into an assistant you can actually talk to. I built an interactive walkthrough of all three stages, and the last one runs on a leash for a reason.

Stage 1 — SFT: imitate demonstrations

First, supervised fine-tuning. Humans write a few thousand ideal prompt → answer demonstrations, and you fine-tune the base model to copy them with ordinary next-token cross-entropy.

demos = [(prompt, human_written_answer), ...]     # ideal behaviour
sft   = finetune(base, demos, loss="cross_entropy")
ref   = freeze(copy(sft))   # the reference policy for later

Now it follows instructions and uses the right format. Critically, you freeze a copy of the SFT model as the reference policy — Stage 3 will measure drift against it.

Stage 2 — a reward model, trained by ranking

SFT can only imitate answers humans bother to write, and it cannot express comparative judgements like "this one is a bit safer" or "this tone is nicer." So you stop writing answers and start ranking them: sample two completions per prompt, have a human pick the better one. Judging is far cheaper than authoring and captures nuance you could never fully write down.

Then you distil those preferences into a reward model with the Bradley–Terry loss — push the chosen answer's score above the rejected one:

def rm_loss(x, y_chosen, y_rejected):
    return -log_sigmoid(rm(x, y_chosen) - rm(x, y_rejected))
train(rm, prefs, loss=rm_loss)   # rm(x,y) now approximates human preference

The demo makes this stage hands-on: you are the labeller. Each pair you rank runs one real Bradley–Terry gradient step on a little weight vector over quality features (helpful, honest, harmless), and you watch the weights and the loss update as the model infers what you value — from your clicks, not a written rulebook.

Stage 3 — PPO, and the leash

Now optimize the policy to produce answers the reward model scores highly. But the reward model is only a proxy — trained on limited data — and if you push too hard the policy escapes the region it understands and produces high-scoring nonsense. That is reward hacking. The fix is a KL-divergence penalty back to the frozen reference:

def rlhf_reward(x, y):
    r  = reward_model(x, y)                   # want this HIGH
    kl = log pi_theta(y|x) - log pi_ref(y|x)  # drift from the SFT reference
    return r - beta * kl                      # the leash: penalize drift

Every token of drift costs reward. beta sets the leash length: large beta is a short leash (safe, smaller gains), small beta is a long leash (higher proxy reward, more risk). PPO then optimizes that signal with clipped updates so the policy never lurches too far in one step.

The single most important RLHF picture

The demo's Stage 3 lets you set beta and run a PPO trajectory, plotting three quantities: the proxy reward the RM assigns, the KL drift from the reference, and the true quality you actually want. As training pushes harder, the proxy reward keeps climbing — but true quality rises, peaks, and then falls.

as KL rises:
  proxy reward:  up, up, up   (always climbs)
  true quality:  up, peak, down   (over-optimization)

That is Goodhart's law for RLHF: maximizing the proxy eventually hurts the goal. Set beta ≈ 0 in the demo and you watch true quality collapse even as the proxy reward looks fantastic — the policy is gaming the reward model off-distribution. Keep beta in the sweet spot and reward climbs and quality stays near its peak. The KL leash, plus early stopping, is what keeps you on the good side of that peak.

The whole recipe

sft    = finetune(base, human_demos)          # 1) imitate
rm     = train(prefs, bradley_terry_loss)     # 2) score preferences
policy = ppo(sft, reward=rm, ref=sft, beta)   # 3) optimize on a KL leash

Three moves: SFT teaches it to answer, the reward model captures what humans prefer, and PPO chases that reward without going off the rails. This is the pipeline that turned raw GPT-style pretraining into the assistants people use. It works, but it is heavy — three models and an unstable RL loop — which is exactly why DPO rewrites the math to train straight on preference pairs, no separate reward model or PPO required. That is the next one.

Step through SFT, rank the pairs yourself to train the reward model, then set the KL leash and watch the reward-vs-quality trade-off play out:
https://dev48v.infy.uk/ai/days/day34-rlhf.html