How the KL-regularized RLHF objective collapses into a classification loss on log-probability ratios, how to implement it without silent bugs, the family of variants it spawned, and what is lost when you never sample from your own model.

Chapter 16 trained a language model to maximize a learned reward while staying close to the model it started from. The machinery was heavy. Four networks sat in memory (policy, reference, reward model, and value model), every training step began by sampling responses, a critic had to be learned at the same time as the policy, and a dozen hyperparameters each had a way of breaking the run (§16.4, §16.6).

You have already met the answer to the problem PPO was solving. Section 15.7 stated it, with a one-line balance argument for why it has its shape, and §16.4 recapped it as the target PPO chases. If we stop insisting that the policy be a neural network, and allow any distribution over responses at all, the best policy can be written down exactly,

Take the reference model's distribution, multiply each response's probability by an exponential of its reward, and renormalize. Both earlier chapters used this formula without proving it. This section is where it is proved, and where the facts the rest of the chapter leans on are taken apart one at a time. We derive it twice, once from the Gibbs variational principle and once with Lagrange multipliers (the second turns §15.7's balance argument into a proof), check when the derivation needs care, read off what the formula says about $\beta$, reward, and KL, see what changes if the KL is replaced by another divergence, and then explain why, despite having the answer in closed form, nobody can use it directly. The obstacle is one number per prompt, the normalizer $Z(x)$. The next section is about the trick that makes that number disappear, which is Direct Preference Optimization.

The problem PPO is solving

Fix the notation we will use for the whole chapter. A prompt $x$ is drawn from a dataset $\mathcal{D}$ of prompts. A response $y$ is a whole sequence of tokens $y_1, \dots, y_{|y|}$, and $|y|$ is its length. A policy $\pi(y \mid x)$ is a distribution over complete responses, which for a language model is the product of its next-token probabilities, $\pi(y\mid x) = \prod_{t=1}^{|y|} \pi(y_t \mid x, y_{<t})$, as in §6.4. When a subscript on $y$ is a position $t$, as here, it means a token. Otherwise, $y_1, y_2, \dots$ and $y_w, y_l$ name different whole responses. The reference policy $\pi_{\text{ref}}$ is fixed, usually the supervised fine-tuned model. The reward $r(x, y)$ is a number that scores a complete response, and $\beta > 0$ sets how strongly we are tied to the reference. The KL-regularized objective of §15.7 and §16.4 is

In plain words, produce responses that score well on average, but pay $\beta$ for every nat of divergence from the reference model. The KL here is over whole responses. By the chain rule of KL divergence it equals the expected sum of the per-token KLs along a sampled response, which is the per-token form PPO implements.

Two observations turn this into a problem we can solve on paper.

First, let the policy be anything. PPO optimizes over the parameters $\theta$ of a network. Here we optimize over every conditional distribution $\pi(\cdot \mid x)$, one probability vector per prompt. This is a relaxation, and the answer tells us what a network with unlimited capacity would converge to.

Second, the prompts decouple. The objective is an average over prompts of a term that depends only on that prompt's distribution $\pi(\cdot\mid x)$. Changing the policy on one prompt cannot help or hurt another, so we can maximize each term separately. From here until the end of the derivations, fix one prompt $x$ and drop it from the notation. We want the probability vector $\pi$ over responses that maximizes

subject to $\pi(y) \ge 0$ and $\sum_y \pi(y) = 1$. The sum runs over every possible response, a huge but finite set once lengths are capped.

A four-response example

Before any derivation, look at numbers small enough to check by hand. Suppose a prompt has only four possible responses. The reference model likes the first one most, and the reward model likes the last one most.

With $\beta = 1$ the unnormalized weights in the last column sum to $Z = 2.284$, so the optimal policy is $(0.175, 0.357, 0.144, 0.323)$. The rare but excellent $y_4$ more than triples its probability, the bland favorite $y_1$ loses more than half of its probability, and nothing is set to zero. Now halve $\beta$. Each weight becomes $\pi_{\text{ref}}(y)\,e^{2r(y)}$, the exponential gaps double, and $y_4$ takes $0.633$ of the mass.

The cell computes the same thing for several values of $\beta$, in log space so that small $\beta$ cannot overflow (the log-sum-exp trick of §0.4).

At $\beta = 10$ the policy barely moves, its expected reward is 0.640 against the reference's 0.600, and it pays a KL of only 0.002 nats. At $\beta = 0.1$ it puts essentially all its mass on $y_4$, collects the maximum reward of 2, and pays a KL of 2.301 nats, which is almost exactly $\log(1/0.1) = 2.303$, the price of turning a 10 percent response into a certainty. Every value of $\beta$ in between gives a different compromise, and the derivations below show that each one is the best compromise available at its price.

import numpy as np

def optimal_policy(pi_ref, r, beta):
    w = pi_ref * np.exp(r / beta)          # overflows, and is not normalized per row
    return w / w.sum()

Derivation one, the Gibbs variational principle

Section 15.7 stated the identity this derivation produces and used it, without proof, to grade training runs, and §16.4 sketched it in one line. Here it is proved with every step written out, because the rearrangement it produces along the way, the reward expressed through the policy, is what DPO is built on, and because the maximum value it produces, $\beta\log Z$, is the quantity the rest of this section is about. The cleanest derivation guesses the answer and then proves that nothing beats it. Define the tilted distribution

It is a valid distribution, because every weight is nonnegative and dividing by $Z$ makes them sum to one. Taking logs and rearranging gives a formula for the reward in terms of $\pi^\star$, which we will use again and again,

Now take any policy $\pi$ and rewrite its objective. Pull $\beta$ out of both terms and put them under one sum,

The bracket is a single logarithm, since $\log\pi - \log\pi_{\text{ref}} - r/\beta = \log\big(\pi / (\pi_{\text{ref}}\, e^{r/\beta})\big)$. Replace the denominator using $\pi_{\text{ref}}(y)e^{r(y)/\beta} = Z\,\pi^\star(y)$,

The first sum is a KL divergence and the last sum is 1, so

In plain words, the objective of any policy equals a constant that does not depend on the policy, minus $\beta$ times how far the policy is from $\pi^\star$. By Gibbs' inequality (§0.5) a KL divergence is never negative and is zero only when the two distributions are equal. So $J_x$ is maximized by $\pi = \pi^\star$ and by nothing else, and the maximum value is

This identity, with its maximizer, is called the Gibbs variational principle (in probability it is also known as the Donsker-Varadhan formula). Statistical physicists will recognize the pieces. With energy $-r$ and temperature $\beta$, $\pi^\star$ is the Boltzmann distribution relative to the prior $\pi_{\text{ref}}$, and $-\beta\log Z$ is the free energy. The quantity $\beta\log Z$ is called the soft value of the prompt, and we will write $V^\star(x) = \beta \log Z(x)$. It is a soft maximum of the rewards in exactly the sense of log-sum-exp. As $\beta \to 0$ it approaches the largest reward among responses the reference can produce, and as $\beta \to \infty$ it approaches the reference's average reward $\mathbb{E}_{\pi_{\text{ref}}}[r]$ (Exercise 17.1b proves both).

None of the 5,000 random policies reaches $J(\pi^\star) = 0.826$, and the identity $J = \beta\log Z - \beta D_{\mathrm{KL}}(\pi\Vert\pi^\star)$ holds for every one of them to rounding error.

import numpy as np

def objective(pi, pi_ref, r, beta):
    return float(np.sum(pi * r))           # forgot the KL penalty

def soft_value(pi_ref, r, beta):
    return float(np.max(r))                # only right as beta -> 0

Derivation two, Lagrange multipliers

The first derivation needed us to guess $\pi^\star$. The second one finds it without guessing, which is how you would discover it at a whiteboard. Treat each $\pi(y)$ as a separate variable and enforce $\sum_y\pi(y) = 1$ with a Lagrange multiplier $\lambda$ (nonnegativity will take care of itself). The Lagrangian (we write it $\mathcal{J}$, since it is maximized) is

Differentiate with respect to one coordinate $\pi(y)$. The product rule on $\pi\log\pi$ gives $\log\pi + 1$, so

Solve for $\log \pi(y)$ and exponentiate,

The last factor is the same for every $y$. The constraint $\sum_y \pi(y) = 1$ fixes it, $e^{-(\beta+\lambda)/\beta} = 1/Z$, and we recover $\pi^\star$.

Read the stationarity condition again, rearranged as $r(y) - \beta\big[\log\frac{\pi(y)}{\pi_{\text{ref}}(y)} + 1\big] = \lambda$. The left side is the marginal value of a little extra probability on $y$, its reward minus the marginal KL cost, and the condition says it must be the same number $\lambda$ for every response. That is exactly the balance argument of §15.7 (if two responses had different marginal values, moving probability from the worse to the better would raise $J$), and $\lambda$ is the common value. At the solution it equals $\beta\log Z - \beta$, which you can check by substituting $\pi^\star$.

Two loose ends make this a proof rather than a calculation. The first is the constraint $\pi(y) \ge 0$, which we did not enforce. A maximum could in principle sit on the boundary, with some $\pi(y) = 0$, where the derivative above does not exist. It cannot. The derivative of $-\beta\,\pi\log\pi$ is $-\beta(\log\pi + 1)$, which tends to $+\infty$ as $\pi \to 0$. So at any point with $\pi(y) = 0$ and $\pi_{\text{ref}}(y) > 0$, moving a sliver of probability $\delta$ onto $y$ from any other response gains about $\beta\,\delta\log(1/\delta)$ from that term, which beats every finite loss of order $\delta$. No boundary point is optimal, the maximizer is interior, and an interior maximum must satisfy the stationarity condition.

The second is whether this stationary point is a maximum. The second derivatives are $\partial^2\mathcal{J}/\partial\pi(y)^2 = -\beta/\pi(y) < 0$ and every mixed partial is zero, so the Hessian is diagonal with negative entries and the objective is strictly concave on the simplex. For a strictly concave function on a convex set, any local maximum is the unique global one, so the single stationary point is the maximizer. Both derivations agree, as they must.

One case needs care. If $\pi_{\text{ref}}(y) = 0$, the KL term is infinite for any $\pi$ that gives $y$ positive probability, so $\pi^\star(y) = 0$ no matter how large $r(y)$ is. (In the Lagrangian, drop those responses before differentiating, and run the argument above on the support of $\pi_{\text{ref}}$.) The optimal policy can reweight what the reference can say, but it can never say something the reference cannot. For language models this is less restrictive than it sounds, because softmax probabilities are never exactly zero, but it is the precise sense in which the KL penalty keeps the policy near the reference.

A last technical remark, for completeness. The derivations treat the set of responses as finite, which it is once lengths are capped, so $Z$ is a finite sum and always exists. Without a length cap, $Z$ is an infinite series, and it is finite whenever the reward is bounded above, since then $Z \le e^{\sup r/\beta}$. The Gibbs derivation goes through unchanged, because it never differentiates and only uses that $Z$ is finite and KL is nonnegative.

What the formula says

Write the optimal policy in log space,

The tuned model's log-probabilities are the reference's log-probabilities plus the reward divided by $\beta$, minus a constant per prompt. This is called an exponential tilt of $\pi_{\text{ref}}$. Five consequences follow, and each one comes back later in the chapter.

Ratios do not need $Z$. For any two responses,

The normalizer cancels in every ratio. Hold on to this, because it is the whole idea of DPO.

Only reward differences matter. Adding any function $f(x)$ of the prompt alone to the reward multiplies every weight by $e^{f(x)/\beta}$, which $Z$ absorbs, so $\pi^\star$ is unchanged. A reward model trained with the Bradley-Terry loss is only defined up to exactly such a shift (§16.3), so this is reassuring. The shift that reward models cannot pin down is also one the optimal policy cannot see.

$\beta$ is a temperature. As $\beta\to\infty$ the tilt vanishes and $\pi^\star \to \pi_{\text{ref}}$. As $\beta \to 0$ the mass concentrates on the highest-reward response among those the reference can produce. In between, $\beta$ converts reward into log-probability at an exchange rate of one nat per $\beta$ units of reward.

It is a Bayesian posterior. Read $\pi_{\text{ref}}$ as a prior over responses and $e^{r(y)/\beta}$ as a likelihood, proportional to the probability of observing "this response is good" given $y$. Then $\pi^\star$ is the posterior and $Z$ is the evidence. Korbak, Perez, and Buckley (2022) argue that this is the right way to think about KL-regularized fine-tuning, as conditioning a prior on evidence rather than as maximizing reward, and it explains why the KL term is not a hack added to stabilize RL but part of what defines the target.

The reward can be read off the policy. Rearranged, the log-space formula is $r(y) = \beta\log\big(\pi^\star(y)/\pi_{\text{ref}}(y)\big) + \beta\log Z$. Knowing the optimal policy and the reference, we know the reward up to the per-prompt constant $\beta\log Z$. Section 17.2 builds DPO on exactly this sentence.

The price of reward, in nats

How much reward does a nat of KL buy? Take the expectation of $\log(\pi^\star/\pi_{\text{ref}}) = r/\beta - \log Z$ under $\pi^\star$,

Write $\kappa = 1/\beta$ for the inverse temperature and $A(\kappa) = \log\sum_y\pi_{\text{ref}}(y)e^{\kappa r(y)}$, which is $\log Z$. Then $\pi^\star$ is an exponential family in $\kappa$, and its log-normalizer has the familiar derivatives $A'(\kappa) = \mathbb{E}_{\pi^\star}[r]$ and $A''(\kappa) = \operatorname{Var}_{\pi^\star}(r)$ (differentiate under the sum once and then again). So as $\kappa$ grows, which means as $\beta$ shrinks,

Both the reward and the KL grow as $\beta$ shrinks, and the ratio of the two derivatives is the slope of the curve traced by $(D_{\mathrm{KL}}, \mathbb{E}[r])$,

In plain words, at the optimum for a given $\beta$, one more nat of divergence buys exactly $\beta$ more units of reward. This is the Lagrangian picture of the problem. $\pi^\star_\beta$ gets the highest reward of any policy with the same KL (if some $\pi$ had the same KL and more reward, it would have a larger $J$), so the curve is the reward-KL frontier, and $\beta$ is the price that picks a point on it. Large $\beta$ makes KL expensive and stops near the origin. Small $\beta$ makes it cheap and slides far out along a frontier that flattens, because each extra nat buys less and less. When the reward is a learned model with errors, the far end of that flat stretch is where reward hacking lives (§16.5).

The measured slope matches $\beta$ at every point, to the accuracy of the finite differences. The widget below draws the whole frontier for an eight-response prompt, along with the policy that every practitioner reaches for first, best-of-$n$, which samples $n$ responses from the reference and keeps the one with the highest reward (§18.2). Its KL from the reference is often quoted as $\log n - (n-1)/n$. Beirami et al. (2024) showed that this formula is only an upper bound, and the widget computes the exact value for this small example, where the best-of-$n$ distribution is known in closed form. Sort the responses by reward and let $F$ be the reference CDF in that order. The best of $n$ independent draws is response $i$ with probability $F(i)^n - F(i-1)^n$.

The bound also gives a feel for real KL budgets. Best-of-16 sits at most $\log 16 - 15/16 = 1.84$ nats from the reference, and best-of-64 at most 3.17 nats, and $n$ in this range is where best-of-$n$ is commonly used to improve outputs. Spread over a 300-token response, 3.17 nats is about 0.01 nats per token. Sequence-level KL is the number that measures how far a policy has moved, and a policy can change its behavior a great deal while every individual token distribution barely moves. Keep this scale in mind when §17.2 finds DPO moving single pairs by tens of nats.

import numpy as np

def beta_for_kl(pi_ref, r, kl_target, tol=1e-8):
    return 0.1                              # the default everyone uses, whatever the budget

What is special about the KL

Nothing in the Lagrange derivation used much about the KL, so it is fair to ask what happens with a different penalty. Interviewers ask it, and the answer separates the parts of the closed form that DPO truly needs from the parts that come with the KL.

An $f$-divergence measures the gap between two distributions as $D_\varphi(\pi\Vert\pi_{\text{ref}}) = \sum_y\pi_{\text{ref}}(y)\,\varphi\big(\pi(y)/\pi_{\text{ref}}(y)\big)$ for a convex function $\varphi$ with $\varphi(1) = 0$. (The literature writes $f$. We write $\varphi$ because $f(x)$ already names a function of the prompt in this section.) Our KL is the case $\varphi(u) = u\log u$, since $\sum_y\pi_{\text{ref}}\frac{\pi}{\pi_{\text{ref}}}\log\frac{\pi}{\pi_{\text{ref}}} = \sum_y\pi\log\frac{\pi}{\pi_{\text{ref}}}$. The forward KL of §0.5, $D_{\mathrm{KL}}(\pi_{\text{ref}}\Vert\pi) = \sum_y\pi_{\text{ref}}\log\frac{\pi_{\text{ref}}}{\pi}$, is the case $\varphi(u) = -\log u$.

Put $D_\varphi$ in place of the KL and repeat the Lagrange calculation. The derivative of $\pi_{\text{ref}}(y)\,\varphi\big(\pi(y)/\pi_{\text{ref}}(y)\big)$ with respect to $\pi(y)$ is $\varphi'\big(\pi(y)/\pi_{\text{ref}}(y)\big)$, by the chain rule, so stationarity reads

with $\lambda$ chosen so the probabilities sum to one. For $\varphi(u) = u\log u$, $\varphi'(u) = \log u + 1$, its inverse is $e^{v - 1}$, and this is the tilt again. For the forward KL, $\varphi'(u) = -1/u$, and

In plain words, the forward-KL optimum also reweights the reference toward high reward, but by a factor that grows like $1/(\lambda - r)$ rather than exponentially. There is no closed form for $\lambda$. The sum on the right falls from $+\infty$ to below 1 as $\lambda$ rises from $\max r$ to $\max r + \beta$ (at $\lambda = \max r + \beta$ every term is at most its $\pi_{\text{ref}}$), so $\lambda$ is a root in that interval, found by bisection. The objective is again strictly concave (the second derivative of $-\beta\pi_{\text{ref}}\log(\pi_{\text{ref}}/\pi)$ is $-\beta\pi_{\text{ref}}/\pi^2$), so this is the unique maximum.

At $\beta = 1$ the forward-KL optimum is $(0.180, 0.246, 0.116, 0.457)$ with $\lambda = 2.219$. At $\beta = 0.1$ the difference is stark. The tilt gives $y_1$ a probability of about $8\times10^{-9}$, while the forward-KL policy still gives it 0.020. As $\beta \to 0$, the losers' probabilities under the forward KL shrink only in proportion to $\beta$, about $\beta\pi_{\text{ref}}(y)/(\max r - r(y))$, while under the reverse KL they shrink like $e^{-(\max r - r(y))/\beta}$. This is the mode-seeking versus mass-covering contrast of §0.5, now visible in the optimum itself. The forward KL charges infinitely for letting any supported response die, so it never does.

Two properties carry over to every $f$-divergence, and three belong to the KL alone.

• Carries over, the one-constant-per-prompt structure. Rearranged, stationarity says $r(x, y) = \beta\varphi'\big(\pi(y\mid x)/\pi_{\text{ref}}(y\mid x)\big) + \lambda(x)$. The multiplier plays the role of $\beta\log Z$, it depends only on the prompt, and it enters additively, so it cancels in the reward difference of two responses to the same prompt. That is all the next section needs, and Wang et al. (2023) build exactly this $f$-DPO, with forward KL, Jensen-Shannon, and $\alpha$-divergences, reporting that the choice trades some alignment performance for more diverse generations. TRL exposes it as f_divergence_type.
• Carries over, no new support. With $\pi_{\text{ref}}(y) = 0$, the formula gives $\pi(y) = 0$ again.
• KL only, an explicit normalizer. The tilt has $Z = \mathbb{E}_{\pi_{\text{ref}}}[e^{r/\beta}]$, an expectation, which is what makes its soft value, its Bayesian reading, and the reward-KL calculus above so clean. Other divergences give $\lambda$ only as the root of an equation.
• KL only, a sum over tokens. $\beta\log\frac{\pi(y)}{\pi_{\text{ref}}(y)}$ is a sum of per-token log-ratios, so the implied reward decomposes over tokens. That is what makes the per-token soft values of the callout below, and PPO's per-token penalty, possible. The forward KL's $\varphi'(\pi/\pi_{\text{ref}}) = -\pi_{\text{ref}}(y)/\pi(y)$ is a product of per-token ratios. For whole responses whose log-ratios reach tens of nats (§17.2), that is $e^{\pm 40}$, which makes such losses numerically delicate on long sequences.
• KL only, estimable from the policy's own samples. $D_{\mathrm{KL}}(\pi\Vert\pi_{\text{ref}})$ is an expectation under $\pi$, which is why PPO can penalize it on its rollouts (§15.7). The forward KL is an expectation under the reference, which would need reference samples at every step.

import numpy as np

def forward_kl_policy(pi_ref, r, beta):
    a = np.log(pi_ref) + r / beta                       # the reverse-KL tilt, not the forward-KL optimum
    w = np.exp(a - a.max(-1, keepdims=True))
    return w / w.sum(-1, keepdims=True)

The partition function problem

So we have the answer. Why does anyone run PPO? Look again at the normalizer,

The sum is over every possible response. With a vocabulary of $V$ tokens and responses of up to $T$ tokens there are about $V^T$ of them. Llama 3's tokenizer has about 128,000 tokens, so responses of 100 tokens number about $128{,}000^{100} \approx 10^{511}$. Enumerating them is not slow, it is impossible, and every prompt has its own $Z(x)$.

Could we estimate $Z(x)$ instead? It is an expectation under the reference, $Z = \mathbb{E}_{y\sim\pi_{\text{ref}}}[e^{r(y)/\beta}]$, so averaging $e^{r/\beta}$ over $N$ samples from the reference gives an unbiased estimate. The trouble is the variance. The average is dominated by the rare samples with the largest rewards, and the smaller $\beta$ is, the more extreme that domination becomes. A clean way to measure it is the effective sample size of the weights $w_i = e^{r(y_i)/\beta}$,

which equals $N$ when all weights are equal and 1 when a single sample carries all the weight. A worked example makes the problem concrete. Suppose rewards under the reference are Gaussian, $r \sim \mathcal{N}(0, s^2)$, a convenient model for rewards that have been standardized to unit spread. The moment generating function of a Gaussian gives everything in closed form,

With unit reward spread and $\beta = 0.1$, the fraction of useful samples is $e^{-100} \approx 4\times 10^{-44}$. We would need about $e^{100} \approx 3\times10^{43}$ samples from the reference just to get one effective sample for one prompt.

At $\beta = 1$ the estimate of $\log Z$ is right and over a third of the samples are useful. At $\beta = 0.25$ the effective sample size has collapsed to about one, and the estimate is wrong by almost 2 nats. At $\beta = 0.1$ it is off by 20 nats, a factor of $e^{20}$ in $Z$, while looking perfectly confident. (The predicted ESS for the last rows is far below 1. The measured one cannot go below 1, which is exactly the problem, since one sample is carrying everything.) For reference, $\beta = 0.1$ is the default value in the DPO paper and in many recipes.

And estimating $Z$ would not be enough. To use $\pi^\star$ we must sample from it, and token by token that requires the soft value of every prefix the model might reach, each of which is another sum of the same kind (the callout below spells this out). We also need a policy for prompts that were not in the training set, which means a network $\pi_\theta$ that has learned to approximate $\pi^\star$ everywhere. So the real problem was never "what is the optimal policy". It is "how do we train $\pi_\theta$ toward $\pi^\star$ without ever computing $Z$".

Three ways around the partition function

Every practical method for KL-regularized preference learning is a way of reaching $\pi^\star$ without $Z$. There are three families.

Optimize $J$ directly with reinforcement learning. The policy gradient of §15.3 needs only samples from $\pi_\theta$ and their rewards, with the KL folded into the reward as a per-token penalty. $Z$ never appears, because we never normalize anything over all responses. The price is everything that made Chapter 16 heavy, on-policy sampling during training, a critic or another baseline, and a reward model that is queried on the policy's own outputs, which it may not have been trained on.

Fit $\pi_\theta$ to $\pi^\star$ by weighted maximum likelihood. Minimize $D_{\mathrm{KL}}(\pi^\star\Vert\pi_\theta)$. Up to a constant this is $-\mathbb{E}_{y\sim\pi^\star}[\log\pi_\theta(y)] = -\mathbb{E}_{y\sim\pi_{\text{ref}}}\big[\tfrac{e^{r(y)/\beta}}{Z}\log\pi_\theta(y)\big]$, supervised learning on reference samples weighted by exponentiated reward. This is reward-weighted regression (Peters and Schaal 2007) and, with advantages in place of rewards, advantage-weighted regression (Peng et al. 2019). In practice $Z$ is replaced by the sum of the weights in a batch, and the method inherits exactly the effective sample size problem we just measured.

Change variables so that $Z$ cancels. We saw that ratios of $\pi^\star$ do not involve $Z$, and that the reward can be read off the policy up to $\beta\log Z$. Preference data only ever compares two responses to the same prompt, so a loss written on preferences might never need $Z$ at all. That is Direct Preference Optimization, the subject of the next section.

In PyTorch

The closed form is the optimum over all policies. A neural policy is a parameterized subset, and the honest check is that gradient ascent on $J$ through a softmax parameterization, with autograd and no knowledge of the formula, finds the same distribution. Here the policy is a free vector of logits, which can represent any distribution with full support, so it should converge to $\pi^\star$ exactly.

Notice what this toy hides. We could compute $J$ exactly because there were four responses. With a real model there are $10^{511}$, so the expectation must be estimated from samples of $\pi_\theta$, and that is the road back to policy gradients.

Read more

• Rafailov, Sharma, Mitchell, Ermon, Manning & Finn (2023). Direct Preference Optimization: Your Language Model is Secretly a Reward Model. Section 4 and Appendix A.1 derive the closed-form optimum exactly as the first derivation here does, as the first step of DPO.
• Korbak, Perez & Buckley (2022). RL with KL penalties is better viewed as Bayesian inference. The posterior view of $\pi^\star$, and why the KL term defines the target rather than stabilizing RL.
• Peng, Kumar, Zhang & Levine (2019). Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learning. Derives the same exponential tilt with Lagrange multipliers and fits it by weighted regression, the second route around $Z$.
• Beirami, Agarwal, Berant, D'Amour, Eisenstein, Nagpal & Suresh (2024). Theoretical guarantees on the best-of-n alignment policy. Shows that the popular $\log n - (n-1)/n$ is an upper bound on the KL of best-of-$n$, and how loose it can be.
• Rafailov, Hejna, Park & Finn (2024). From r to Q*: Your Language Model is Secretly a Q-Function. The token-level version of the closed form, with soft values at every prefix.
• Wang, Jiang, Yang, Liu & Chen (2023). Beyond Reverse KL: Generalizing Direct Preference Optimization with Diverse Divergence Constraints. The optimum under $f$-divergence penalties from the KKT conditions, and the $f$-DPO losses built on it.
• Gao, Schulman & Hilton (2023). Scaling Laws for Reward Model Overoptimization. Measures what happens far out along the reward-KL frontier when the reward is a learned proxy.

Watch next

▶

▶