The objective, written down carefully
Reinforcement learning from human feedback fine-tunes a language model π_θ (the policy) to maximize a learned scalar reward r(x,y) — how good response y is for prompt x — while staying close to a frozen reference π_ref (usually the supervised-fine-tuned model you started from). The full objective is:
max_θ E_{x~D, y~π_θ(·|x)} [ r(x,y) ] − β · KL( π_θ(·|x) || π_ref(·|x) )Read it left to right. Sample a prompt x from data distribution D; sample a response y from the current policy; score it with the reward model; and subtract a penalty proportional to how far the policy has drifted from the reference, measured by the Kullback–Leibler divergence. The coefficient β > 0 sets the exchange rate between chasing reward and staying anchored. Two forces, one knob — and everything below follows from this single line, with no separate ‘RL magic’ underneath it.
What the KL term actually measures
The KL divergence between the policy and the reference, for a fixed prompt, is the expected log-ratio of the two distributions over responses:
KL(π_θ || π_ref) = E_{y~π_θ} [ log ( π_θ(y|x) / π_ref(y|x) ) ] ≥ 0It is zero exactly when the two policies are identical and grows as they diverge. Crucially the expectation is taken under π_θ, the model being trained — this is the reverse KL, and its asymmetry has teeth. Reverse KL is mode-seeking: it heavily punishes the policy for placing mass where π_ref is tiny (the log-ratio blows up), but not for ignoring regions the reference liked. In RLHF that is the behavior you want: the leash yanks hard the moment the policy wanders into text the base model would essentially never produce — precisely where a reward model’s scores are least trustworthy.
Why the KL leash is not optional
Delete the KL term and the objective collapses to ‘maximize reward,’ which sounds fine until you remember r(x,y) is a learned approximation trained on finite human comparisons. It is accurate on responses the reference produces and increasingly unreliable the further you leave it. An unconstrained optimizer will find the seams — degenerate outputs, repetitive flattery, keyword-stuffed text, weird formatting — that score absurdly high because the reward model never saw anything like them and extrapolates badly. This is reward hacking: the policy exploits errors in the proxy instead of satisfying the true preference behind it.
The KL penalty is the guardrail. By charging the policy for every nat of divergence from π_ref, it confines optimization to the region where the reward model is valid, and preserves the fluency, diversity, and world knowledge already baked into the reference. The objective is not ‘get maximum reward’ but ‘get more reward without straying off the manifold of sensible language.’ That constraint is why RLHF produces usable models rather than confident nonsense.
Per-token versus sequence-level KL
The objective defines the KL over whole sequences, but a response is generated token by token, and the reward model usually only scores the completed sequence. So how is the penalty applied during training? The key identity is that a sequence log-probability factorizes, so the sequence log-ratio is a sum of per-token log-ratios:
log ( π_θ(y|x) / π_ref(y|x) ) = Σ_t log ( π_θ(y_t | x, y_<t) / π_ref(y_t | x, y_<t) )Take the expectation under π_θ and the left side is the sequence-level KL — so summing a per-token penalty recovers exactly the sequence-level quantity. PPO exploits this: at every step it subtracts β(log π_θ − log π_ref) as a dense, per-token cost, and only adds the terminal reward r(x,y) at the final token. Spreading the KL across tokens gives a signal at each decoding step rather than one lump at the end, sharply improving credit assignment. The two views are equivalent in expectation; the per-token form is just far easier to learn from.
The closed-form optimal policy
Here is the payoff for writing the objective precisely: it has an exact solution. Ask which distribution π maximizes the per-prompt objective. The answer is the KL-regularized optimal policy:
π*(y|x) = (1 / Z(x)) · π_ref(y|x) · exp( r(x,y) / β )
where Z(x) = Σ_y π_ref(y|x) · exp( r(x,y) / β ) is the partition function.The optimal policy is the reference distribution reweighted by an exponential tilt in the reward, then renormalized: high-reward responses get their probability multiplied up, low-reward ones pushed down, with the strength of the reshaping controlled by 1/β. Notice π* can only put mass where π_ref already had some — if the reference assigns a response probability zero, no amount of reward resurrects it. That is the leash again, visible in the solution itself: the optimum is structurally unable to leave the reference’s support.
Deriving it in three lines
The derivation is short and worth seeing, because it makes the KL leash’s role unmistakable. Divide the per-prompt objective by β and rewrite it as a single expectation:
(1/β) J(π) = E_{y~π} [ r(x,y)/β − log(π(y|x)/π_ref(y|x)) ]
= E_{y~π} [ log ( π_ref(y|x) exp(r/β) / π(y|x) ) ]
= log Z(x) − KL( π(·|x) || π*(·|x) )The last step inserts Z(x) to turn the tilted reference into the normalized π*. Now read the result: log Z(x) is independent of π, and a KL divergence is ≥ 0 with equality only when its arguments match. So the objective is maximized precisely when π = π*. The messy reward-maximization problem was, in disguise, the problem of getting as close as possible to one target — the tilted reference. No Lagrange multipliers needed; the objective is a KL to the optimum plus a constant.
What the optimal policy means
The form π* ∝ π_ref · exp(r/β) is a Boltzmann distribution: the shape statistical physics assigns to states with energy −r at temperature β. Reading β as a temperature makes the whole trade-off intuitive. As β → ∞ the exponent vanishes, the tilt disappears, and π* → π_ref: infinite regularization means never move. As β → 0 the exponential becomes razor-sharp and π* collapses onto the single highest-reward response in the reference’s support — a greedy, reward-maximizing argmax. Real RLHF lives in between: small β risks reward hacking and mode collapse; large β barely improves on the reference. The coefficient is not a minor hyperparameter — it is the axis the whole method slides along.
A worked numeric example
Make it concrete. Fix a prompt with three candidate responses A, B, C, reference π_ref = [0.50, 0.30, 0.20], and reward scores r = [1, 2, 0]. Compute the tilted weights π_ref · exp(r/β) and renormalize by Z.
β = 1: weights = [0.50·e^1, 0.30·e^2, 0.20·e^0] = [1.359, 2.217, 0.200]
Z = 3.776 → π* = [0.360, 0.587, 0.053]
β = 0.5: weights = [0.50·e^2, 0.30·e^4, 0.20·e^0] = [3.695, 16.38, 0.200]
Z = 20.27 → π* = [0.182, 0.808, 0.010]
β = 10: weights = [0.50·e^0.1, 0.30·e^0.2, 0.20·e^0] = [0.553, 0.366, 0.200]
Z = 1.119 → π* = [0.494, 0.327, 0.179]Response B, the reward winner, climbs from 0.30 to 0.587 at β=1, to 0.808 when the leash loosens to β=0.5, but only to 0.327 when it tightens to β=10 — where the policy stays glued to the reference. One prompt, one knob, and you watch the temperature interpolate from ‘copy the reference’ to ‘chase the reward.’