RLHF Math Deep Dive

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) ) ]  ≥ 0

It 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.’

The partition function is the catch

If π* is known in closed form, why not just use it? Because Z(x) = Σ_y π_ref(y|x) exp(r(x,y)/β) sums over every possible response — every sequence of tokens up to the context length. That set is astronomically large, so Z(x) is intractable and you cannot sample from π* directly or even evaluate its normalized density. This single obstacle is why RLHF needs an algorithm rather than a lookup. Every practical method is, at heart, a different strategy for reaching π* without ever forming Z(x). The two dominant strategies — PPO and DPO — attack it from opposite directions: one samples its way around the sum, the other algebraically cancels it out.

How PPO approximates the objective

PPO optimizes the objective the direct way: by sampling. Rather than summing over all responses, it draws them from the current policy, scores them, and takes noisy gradient steps that raise the probability of above-average outcomes — a Monte-Carlo assault on the same reward-minus-KL objective. Concretely it builds a per-token reward: the dense KL penalty −β(log π_θ − log π_ref) at every step, plus the terminal reward-model score at the end, exactly the decomposition from earlier. A value network and generalized advantage estimation reduce variance, and a clipped surrogate keeps each update from moving the policy too far. Because it samples fresh, on-policy data, PPO stays valid even as the policy moves — but it is machinery-heavy: a live reward model, a value head, a reference model, and a delicate sampling loop all in memory at once. It never forms Z(x) because it only compares sampled responses against a running baseline.

How DPO reaches the same optimum

DPO takes the algebraic road. It inverts the optimal-policy equation: solving π* ∝ π_ref exp(r/β) for the reward gives r(x,y) = β log(π*(y|x)/π_ref(y|x)) + β log Z(x). In other words, the policy’s own log-ratio against the reference is an implicit reward. Substituting this into a preference model that compares two responses makes the intractable β log Z(x) term cancel — it is identical for both responses to the same prompt — leaving a simple supervised classification loss on preference pairs, with no reward model, no sampling, and no RL loop. DPO trades PPO’s on-policy machinery for stability and simplicity, at the cost of learning only from a fixed offline dataset. The full cancellation is in the companion DPO article; the point here is that both methods target the very same π* — they merely differ in how they dodge the partition function.