The language model as a policy

Here is a small experiment you will run in this chapter’s lab. A tiny transformer, about 100,000 parameters, has been trained on roughly 22,000 worked examples of two-digit addition: “37+48=” followed by “85”. Ask it a fresh problem and let it answer greedily (the top-scoring digit at every step), and it is right about 80% of the time. Let it sample instead, and it is right only a little over half the time. Most of its mistakes are near misses: in about four out of five wrong answers exactly one digit is off, “86” instead of “85”.

How do you make it better? One option is more of the same: more worked examples. The other is to let the model try, check each answer, and nudge it toward the attempts that worked. The checking is cheap: a one-line program can verify a sum. That second option is reinforcement learning, and it is how reasoning models, chat assistants and coding agents are post-trained. To do it you need to see the language model as what Part I and Part II called a policy: something that picks actions in states and can be improved with a reward. This chapter builds that view carefully. It also answers a question that every later chapter depends on. When we do RL on a language model, we almost always add a leash, a penalty for moving away from where the model started. Why, and what exactly does the leash do?

Logits, softmax, and a distribution over the next token

A language model reads a sequence of tokens (pieces of words, digits, punctuation; a real model’s vocabulary has about 100,000 of them, Llama 3’s tokenizer about 128,000) and outputs one real number per vocabulary entry, the logits $z \in \mathbb{R}^{V}$, where $V$ is the vocabulary size. The softmax turns them into a probability distribution over the next token:

$$\pi_\theta(a \mid x, y_{<t}) = \frac{e^{z_a}}{\sum_b e^{z_b}}, \qquad z = f_\theta(x, y_{<t}).$$

Here $x$ is the prompt, $y_{<t}$ the tokens the model has written so far, and $f_\theta$ is the transformer. This is the softmax policy of Chapter 7, with a neural network computing the logits from the state. Nothing new is needed to call it a policy: it maps a situation to a distribution over actions.

A worked example. After the prompt “37+48=”, suppose the model’s logits for the three most plausible first digits are $z = (2, 1, 0)$ for “8”, “7” and “9” (pretend the other tokens have logits so low they don’t matter). Then

$$\pi(\text{8}) = \frac{e^2}{e^2 + e^1 + e^0} = \frac{7.389}{11.107} = 0.665, \qquad \pi(\text{7}) = 0.245, \qquad \pi(\text{9}) = 0.090.$$

A response is a product; its log is a sum

The model writes a response $y = (y_1, \dots, y_{\lvert y \rvert})$ one token at a time, feeding each sampled token back in, until it emits an end-of-sequence token <eos>. By the chain rule of probability, the probability of the whole response is the product of the next-token probabilities:

$$\pi_\theta(y \mid x) = \prod_{t=1}^{\lvert y \rvert} \pi_\theta(y_t \mid x, y_{<t}), \qquad \log \pi_\theta(y \mid x) = \sum_{t=1}^{\lvert y \rvert} \log \pi_\theta(y_t \mid x, y_{<t}).$$

Continue the example. If the model then puts probability 0.9 on “5” after “8”, and 0.98 on <eos> after “85”, then

$$\pi(\text{"85<eos>"} \mid x) = 0.665 \times 0.9 \times 0.98 = 0.587, \qquad \log \pi = -0.408 - 0.105 - 0.020 = -0.534 .$$

Everyone works with the log-prob, the sum, rather than the product, for two reasons. The first is numerical. A 1,000-token response whose tokens each have probability 0.9 has probability $0.9^{1000} \approx 1.7 \times 10^{-46}$, which is exactly zero in 32-bit floating point. Its log-prob, $-105.4$, is a perfectly ordinary number. The second reason matters more for RL: the gradient of a sum is a sum. The policy gradient of Chapter 7 needs $\nabla_\theta \log \pi_\theta(y \mid x)$, and that is a sum of per-token gradients, each computed by the same backward pass you would use for ordinary training.

Temperature

When you sample from a language model you can set a temperature $T > 0$. The sampler divides the logits by $T$ before the softmax:

$$\pi_T(a \mid s) = \frac{e^{z_a / T}}{\sum_b e^{z_b / T}} \;\propto\; \pi_1(a \mid s)^{1/T}.$$

(Here $s = (x, y_{<t})$ is the state, the prompt plus the tokens so far; the MDP section below makes this official.) The second form says that temperature raises every probability to the power $1/T$ and renormalizes. With the logits $(2, 1, 0)$ from above:

$T$	$\pi(\text{8})$	$\pi(\text{7})$	$\pi(\text{9})$
0.5	0.867	0.117	0.016
1 (the model as trained)	0.665	0.245	0.090
2	0.506	0.307	0.186

As $T \to 0$ the sampler becomes greedy (always the top token). As $T \to \infty$ it becomes uniform. You have seen this before in statistical mechanics: $\pi_T$ is a Boltzmann distribution with energies $E_a = -z_a$ at temperature $T$, and the softmax denominator is the partition function. Lower temperature freezes the system into its ground state, the most likely token.

Two other knobs truncate the distribution before sampling. Top-k keeps only the $k$ most likely tokens. Top-p or nucleus sampling (Holtzman et al., 2020) keeps the smallest set of tokens whose probabilities add up to at least $p$. Both cut off the long tail of unlikely tokens, where a single bad draw can derail a response. They matter for RL in one way: the policy you sampled from is the truncated, temperature-scaled one, and every importance ratio and log-prob in training should refer to that policy. That is one reason RL rollouts are often sampled at $T = 1$ with no truncation (DeepSeek-R1-Zero used temperature 1, for example), and why some systems record the truncation mask and reapply it in training (Chapter 20).

Three stages: pretraining, SFT, RL

A modern assistant is trained in three stages. Each one can be written as an objective over the same policy $\pi_\theta$.

Pretraining. Take trillions of tokens of text and maximize the likelihood of each next token given the ones before it:

$$\mathcal{L}_\text{pretrain}(\theta) = -\sum_{\text{documents } d}\; \sum_t \log \pi_\theta(d_t \mid d_{<t}).$$

This is maximum likelihood, and through the RL lens it is behavior cloning (Chapter 10) where the “expert” is everyone who ever wrote on the internet. The result, the base model, can continue any text plausibly. It has no particular wish to be helpful, correct or brief.

Supervised fine-tuning (SFT). Collect demonstrations, pairs of a prompt $x$ and a good response $y^\star$ written by a person (or a stronger model), and run the same loss on them, counting only the response tokens:

$$\mathcal{L}_\text{SFT}(\theta) = -\sum_{(x, y^\star)} \log \pi_\theta(y^\star \mid x) = -\sum_{(x, y^\star)} \sum_t \log \pi_\theta(y^\star_t \mid x, y^\star_{<t}).$$

Still behavior cloning, now of a narrower expert. The model learns the format and style of an assistant. The SFT model is usually the starting point of RL and the reference model $\pi_\text{ref}$ that the leash ties it to.

Reinforcement learning. Now the model writes its own responses, something scores them, and we maximize the expected score:

$$J(\theta) = \mathbb{E}_{x \sim \mathcal{D}}\, \mathbb{E}_{y \sim \pi_\theta(\cdot \mid x)} \big[ r(x, y) \big].$$

The reward $r(x, y)$ can come from a model trained on human preferences (Chapter 12), from an AI judge, or from a program that checks the answer (Chapter 14, Chapter 15).

Put the gradients side by side and the stages look strikingly similar. By the log-derivative trick of Chapter 7:

$$\nabla_\theta \big(-\mathcal{L}_\text{SFT}\big) = \sum_{(x, y^\star)} \nabla_\theta \log \pi_\theta(y^\star \mid x), \qquad \nabla_\theta J = \mathbb{E}_{x}\,\mathbb{E}_{y \sim \pi_\theta} \big[ r(x, y)\, \nabla_\theta \log \pi_\theta(y \mid x) \big].$$

Both push up the log-probability of some responses. SFT pushes up responses someone else wrote, all with weight 1. RL pushes up the model’s own samples, weighted by how good they were, and with a baseline subtracted (Chapter 7) it pushes the bad ones down. Halfway between sits a simple, widely used method: sample several responses, keep only the correct ones, and fine-tune on them. It goes by names like rejection-sampling fine-tuning, expert iteration or STaR (Zelikman et al., 2022). For each prompt, that is REINFORCE with a 0/1 reward and no baseline, up to a scale factor: the fraction of samples that were correct (Exercise 11.1 asks you to check this).

Why RL after SFT?

If RL and SFT use the same machinery, why bother with RL at all? Why not just collect more demonstrations? There are four good reasons.

1. RL can learn from mistakes. SFT only ever sees good responses. It raises their probability, and the softmax lowers everything else a little, spread over the whole vocabulary. It never looks at what this model actually gets wrong. RL samples the model’s actual answers, so if the model likes to answer “86” to “37+48=”, RL sees “86”, sees reward 0, and with a baseline pushes exactly that answer down. This is the negative example that SFT never gets.

2. It is often easier to judge than to demonstrate. A person can tell which of two summaries is better more easily than they can write a better one. A unit test can check a program that nobody has written yet. A math checker confirms “85” without knowing how to solve harder problems. RL needs only a judgment of each response, which lets you train on tasks where good demonstrations are expensive or don’t exist.

3. RL trains on the model’s own distribution. In SFT the model always continues a human-written prefix. At generation time it continues its own, and one odd early token can lead it into territory no demonstration covers. Chapter 10 called this covariate shift and showed that behavior cloning’s errors can compound with the length of the episode. For text it has a name, exposure bias (Ranzato et al., 2015): the model is never exposed to its own mistakes during training. RL fixes this by construction: every training response is one the model itself wrote, mistakes included, and it is the consequences of those mistakes that the reward judges.

4. RL can go beyond the demonstrator. The best SFT can do is reproduce the distribution of its demonstrations, warts and all. RL’s objective has no such ceiling. A model can discover a solution strategy, or a style of reasoning, that no demonstration contained, as long as it samples it sometimes and the reward recognizes it. AlphaGo’s journey from imitating human games to beating the people it imitated is the classic example (Chapter 10).

Two ways to see the MDP

Chapters 3–9 framed RL as an agent moving through states. To use them, you have to say what the states, actions, transitions and rewards of text generation are. There are two natural answers. They are not rival theories: they describe the same process at two levels of detail, and different algorithms work at different levels.

View 1: the token-level MDP

Treat each token as one decision.

MDP piece	For a language model
state $s_t$	the prompt plus the tokens so far, $(x, y_{<t})$
action $a_t$	the next token $y_t$, one of $V$ vocabulary entries
policy $\pi(a \mid s)$	the model’s next-token distribution
transition $P(s' \mid s, a)$	deterministic: append the token, $s_{t+1} = (x, y_{\le t})$
episode end	the model emits <eos> (or hits a length limit)
reward	usually 0 for every token except the last, which gets $r(x, y)$
discount	$\gamma = 1$: the episode is short and finite, and nothing is gained by discounting

Three things make this MDP unusual. The state is the full history, so it is Markov by construction (Chapter 3). The transition function is known exactly and is trivial, which removes one of RL’s usual difficulties (learning or sampling the environment) and leaves only the policy’s own randomness. And the reward is sparse: a single number at the very end. With $\gamma = 1$ and one terminal reward, the return from every token is the same number, $G_t = r(x, y)$. The hard question becomes credit assignment: which of the 300 tokens in a response made it good? A critic $V(x, y_{<t})$ that predicts the eventual reward from a half-finished response is one answer, used by PPO-based RLHF (Chapter 8, Chapter 12).

View 2: the contextual bandit

Zoom out until the whole response is one action.

Bandit piece	For a language model
context	the prompt $x$
arm (action)	the entire response $y$
policy	$\pi_\theta(y \mid x) = \prod_t \pi_\theta(y_t \mid x, y_{<t})$
reward	$r(x, y)$, then the episode is over

This is the contextual bandit of Chapter 2, with an astronomically large set of arms: with 100,000 tokens and responses of 1,000 tokens there are $10^{5000}$ of them. Nobody keeps a table of arms. But the bandit view is often the more honest one, because rewards really are given to whole responses, not to tokens.

Same gradient, different bookkeeping

The two views agree on the most important thing. The bandit’s REINFORCE gradient is

$$r(x, y)\, \nabla_\theta \log \pi_\theta(y \mid x) \;=\; \sum_t r(x, y)\, \nabla_\theta \log \pi_\theta(y_t \mid x, y_{<t}),$$

and the right-hand side is exactly the token-level REINFORCE gradient, $\sum_t G_t \nabla_\theta \log \pi_\theta(a_t \mid s_t)$, with $G_t = r(x, y)$ for every token. Where the views part ways is everything you add on top:

• What baseline or advantage each token gets. Token level: a different advantage per token, from a critic. Bandit level: one advantage for the whole response, shared by its tokens.
• Where importance ratios and clipping live. Per token ($\pi_\theta(y_t \mid s_t)/\pi_\text{old}(y_t \mid s_t)$) or per response ($\pi_\theta(y \mid x)/\pi_\text{old}(y \mid x)$).
• Where the KL penalty enters. As a small per-token reward, or as one number per response.

Method	Reward and advantage	Ratio, clip, KL	Chapter
PPO-based RLHF (InstructGPT, Ouyang et al. 2022)	token level: value head, per-token GAE	per token; KL penalty per token in the reward	12
RLOO (Ahmadian et al. 2024)	bandit: one reward per response, leave-one-out baseline	KL folded into the response’s reward	16
GRPO (Shao et al. 2024)	bandit: one group-normalized advantage per response	per token: ratio, clip and KL term	15
GSPO (Qwen team, 2025)	bandit	per response: one sequence-level ratio	16
DPO (Rafailov et al. 2023)	derived in the bandit view, no sampling at all	none (offline)	13

GRPO is a hybrid: its reward and advantage are per response, while its PPO machinery is per token. GRPO and its variants are the workhorses of 2025–26 reasoning-model training, so this hybrid is the view you will meet most often. Exercise 11.3 works both views on a two-token example.

The KL-regularized objective

Here is the objective that RLHF optimizes, and that almost every method in Part III either optimizes, approximates or modifies:

$$\max_{\pi}\; \mathbb{E}_{x \sim \mathcal{D}} \Big[\, \mathbb{E}_{y \sim \pi(\cdot \mid x)}\big[r(x, y)\big] \;-\; \beta\, \mathrm{KL}\big(\pi(\cdot \mid x) \,\Vert\, \pi_\text{ref}(\cdot \mid x)\big) \Big],$$

with $\mathrm{KL}(\pi \Vert \pi_\text{ref}) = \sum_y \pi(y \mid x) \log \frac{\pi(y \mid x)}{\pi_\text{ref}(y \mid x)}$. In words: get as much reward as you can, but pay $\beta$ for every nat of divergence (a nat is the unit of KL computed with natural logarithms) from the reference model $\pi_\text{ref}$ (the SFT model, frozen). We call this the KL-regularized objective and $\beta$ the KL coefficient.

The best possible policy, in closed form

You have already solved this problem. In Chapter 10 you maximized $\mathbb{E}_\pi[r] - \beta\,\mathrm{KL}(\pi \Vert \pi_0)$ over distributions on a finite set of actions, and found the Gibbs distribution. In the bandit view, each prompt $x$ is one such problem, and its “actions” are the responses $y$. So the optimum is

$$\pi^*(y \mid x) = \frac{1}{Z(x)}\, \pi_\text{ref}(y \mid x)\, e^{r(x, y)/\beta}, \qquad Z(x) = \sum_y \pi_\text{ref}(y \mid x)\, e^{r(x, y)/\beta},$$

and for every other policy, $J(\pi) = \beta \log Z(x) - \beta\, \mathrm{KL}(\pi \Vert \pi^*)$ (the Gibbs identity), so nothing beats $\pi^*$. In words: take the reference model and tilt it, multiplying the probability of each response by $e^{r/\beta}$. Large $\beta$: barely tilt, stay at $\pi_\text{ref}$. Small $\beta$: tilt hard, toward the highest-reward response. A response with $\pi_\text{ref}(y \mid x) = 0$ stays impossible.

∂ The optimum in four lines (a recap of Chapter 10)

Fix a prompt and drop $x$ from the notation. Define $\pi^*$ by the formula above. Taking logs, $r(y) = \beta \log \frac{\pi^*(y)}{\pi_\text{ref}(y)} + \beta \log Z$. Substitute this into the objective for an arbitrary $\pi$:

$$\begin{aligned} J(\pi) &= \sum_y \pi(y)\, r(y) - \beta \sum_y \pi(y) \log \frac{\pi(y)}{\pi_\text{ref}(y)} \\ &= \sum_y \pi(y) \Big[\beta \log \frac{\pi^*(y)}{\pi_\text{ref}(y)} + \beta \log Z - \beta \log \frac{\pi(y)}{\pi_\text{ref}(y)}\Big] \\ &= \beta \log Z - \beta \sum_y \pi(y) \log \frac{\pi(y)}{\pi^*(y)} \;=\; \beta \log Z - \beta\,\mathrm{KL}(\pi \Vert \pi^*). \end{aligned}$$

(The third line uses $\sum_y \pi(y) = 1$.) KL is non-negative and zero only when the two distributions are equal, so $J$ is maximized exactly at $\pi = \pi^*$, where it equals $\beta \log Z$. Chapter 10 has the Lagrange-multiplier route and the free-energy interpretation.

A worked example. Here is a toy language model whose responses are always exactly two tokens long, from the vocabulary $\{a, b\}$. You will meet it again in the demo and the lab. The reference model and the rewards are:

response	$\pi_\text{ref}(y_1)$	$\pi_\text{ref}(y_2 \mid y_1)$	$\pi_\text{ref}(y)$	$r$
aa	0.7	0.6	0.42	0
ab	0.7	0.4	0.28	1
ba	0.3	0.5	0.15	0
bb	0.3	0.5	0.15	2

The best response, “bb”, is the least likely one, and it starts with the less likely first token. The reference’s expected reward is $0.28 + 2(0.15) = 0.58$. With $\beta = 1$:

• Tilted weights $\pi_\text{ref}(y) e^{r(y)}$: $0.42$, $\;0.28e = 0.761$, $\;0.15$, $\;0.15e^2 = 1.108$. Their sum is $Z = 2.439$.
• $\pi^* = (0.172,\; 0.312,\; 0.062,\; 0.454)$ for aa, ab, ba, bb.
• $\mathbb{E}_{\pi^*}[r] = 0.312 + 2(0.454) = 1.221$, and $\mathrm{KL}(\pi^* \Vert \pi_\text{ref}) = 0.329$ nats.
• $J(\pi^*) = 1.221 - 0.329 = 0.892 = \log Z$. Check: $\ln 2.439 = 0.892$.

Compare the alternatives at the same $\beta$. Copying the reference gives $J = 0.58$. Always answering “bb” gives reward 2 but costs $\mathrm{KL} = \ln(1/0.15) = 1.897$, so $J = 0.103$. The tilted policy is far better than either.

The optimum, one token at a time

Now look at the same $\pi^*$ through the token-level lens. What is its probability for the first token? Add up the responses: $\pi^*(a) = 0.172 + 0.312 = 0.484$ and $\pi^*(b) = 0.516$. The reference put 0.7 on $a$; the optimum prefers $b$. But the token $b$ itself earns no reward. It is tilted up because of what can follow it: after $b$, the model can write “bb”, worth 2.

Here is the precise version. Define the soft value of a prefix as the log-sum-exp of the rewards of what can follow it, working backwards from the end. For our two-token responses:

$$V(y_1) = \beta \log \sum_{y_2} \pi_\text{ref}(y_2 \mid y_1)\, e^{r(y_1 y_2)/\beta}, \qquad V_0 = \beta \log \sum_{y_1} \pi_\text{ref}(y_1)\, e^{V(y_1)/\beta} = \beta \log Z.$$

$V(y_1)$ is a “soft maximum” of the rewards reachable after $y_1$: it lies between their $\pi_\text{ref}$-weighted average (large $\beta$) and their maximum (small $\beta$). In physics language, $V(y_1)$ is minus the free energy of the subtree below $y_1$, at temperature $\beta$ with energies $-r$. Then the optimum’s next-token probabilities are the reference’s, tilted by the soft value of where each token leads:

$$\pi^*(y_1) = \pi_\text{ref}(y_1)\, e^{(V(y_1) - V_0)/\beta}, \qquad \pi^*(y_2 \mid y_1) = \pi_\text{ref}(y_2 \mid y_1)\, e^{(r(y_1 y_2) - V(y_1))/\beta}.$$

With $\beta = 1$: $V(a) = \ln(0.6 + 0.4e) = 0.523$ and $V(b) = \ln(0.5 + 0.5e^2) = 1.434$. So $\pi^*(a) \propto 0.7 \cdot e^{0.523} = 1.181$ and $\pi^*(b) \propto 0.3 \cdot e^{1.434} = 1.258$, giving $(0.484, 0.516)$ after dividing by their sum, $Z = 2.439$, as before. The second-token probabilities come out as $\pi^*(b \mid a) = 0.644$ and $\pi^*(b \mid b) = 0.881$.

∂ Why the token-level product equals the sequence-level π* (any length)

Let $s_t = (x, y_{<t})$ and let $s_{t+1} = (x, y_{\le t})$ be the state after appending $y_t$. Define soft values backwards: at a finished response, $V(x, y) = r(x, y)$; at any unfinished prefix,

$$V(s) = \beta \log \sum_a \pi_\text{ref}(a \mid s)\, e^{V(s, a)/\beta},$$

where $(s, a)$ is the state after appending $a$ (a finished response if $a$ is <eos>). Define $\pi^*(a \mid s) = \pi_\text{ref}(a \mid s)\, e^{(V(s, a) - V(s))/\beta}$. It sums to 1 over $a$ by the definition of $V(s)$. Multiply along a response $y$ of length $n$:

$$\prod_{t=1}^{n} \pi^*(y_t \mid s_t) = \prod_{t=1}^{n} \pi_\text{ref}(y_t \mid s_t) \cdot \exp\Big(\frac{1}{\beta}\sum_{t=1}^{n} \big[V(s_{t+1}) - V(s_t)\big]\Big) = \pi_\text{ref}(y \mid x)\, e^{(r(x, y) - V(s_1))/\beta}.$$

The sum telescopes, leaving the value of the finished response, $r(x, y)$, minus the value at the start, $V(s_1)$. Since the product sums to 1 over all responses, $e^{V(s_1)/\beta} = \sum_y \pi_\text{ref}(y \mid x) e^{r(x,y)/\beta} = Z(x)$. That is the sequence-level $\pi^*$.

Why not just compute $\pi^*$ and be done? Because $Z(x)$ and the soft values are sums over every possible continuation, and there are $10^{5000}$ of them. Everything in Part III is a way around that sum. RL methods (Chapters 12, 15, 16) climb the KL-regularized objective with sampled responses. DPO (Chapter 13) notices that $Z(x)$ cancels when you compare two responses to the same prompt. Best-of-$n$ sampling (Chapter 17) approximates a tilted distribution by brute force.

ScenarioWorked example from the text (2 tokens)Likely answers are good answers (3 tokens)Reward model with a blind spot (3 tokens)β

temperature T

show π_ref at temperature T

Token view (tree) and response view (bars): the same policy, seen two ways

reference π_ref   optimum π*   π_ref sampled at temperature T  · edge labels: π_ref → π* for that token; V = soft value of the node

Reward–KL plane: every policy is a point (KL from π_ref, expected reward)

π* as β varies (the frontier)   π_ref at temperature T as T varies   true reward of π* (the reward model is wrong on "cc")

Optimum π* at this β

E_π*[r]

E_π*[true reward]

KL(π* ‖ π_ref)

J(π*) = E[r] − β·KL

β·log Z = V at the root

Other policies, same β

J(π_ref) = E_ref[r]

π_ref at T: E[r]

π_ref at T: KL from π_ref

π_ref at T: J = E[r] − β·KL

What to try: (1) In the worked example at β = 1, check the numbers against the text: π* = 0.172, 0.312, 0.062, 0.454; V(a) = 0.52, V(b) = 1.43; J(π*) = β log Z = 0.892. Note the first-token edge to b rises from 0.30 to 0.52 although b earns nothing by itself. (2) Sweep β from right to left and watch the accent dot run along the frontier, from π_ref (KL = 0) toward the greedy answer bb. (3) Move the temperature slider. In the worked example, lowering T moves the green dot down: the likeliest responses are not the best ones, and temperature can't know that. Switch to 'Likely answers are good answers': now lowering T helps, but its curve still stays below the frontier, and J(π_ref at T) never beats J(π*). (4) In 'Reward model with a blind spot', the reward model wrongly scores the rare response cc at 1.6 (true value 0). Lower β slowly: the true reward (dashed red) rises, peaks near β ≈ 0.4, then collapses as π* piles onto cc.

Why a leash?

Why not set $\beta = 0$ and maximize reward outright? Four reasons, from most to least important.

1. The reward model is only trustworthy near its training data. A reward model (Chapter 12) is trained on responses sampled from (roughly) $\pi_\text{ref}$, labeled by people. On responses like those it is reasonably accurate. On responses unlike anything it saw, its scores are guesses, and an optimizer will find the guesses that are too high. The demo’s third scenario is the smallest possible version: the reward model is right everywhere except on one rare response, and a weak leash lets the policy find it. With real reward models, Gao, Schulman and Hilton (2022) measured the same shape: as the KL from the initial policy grows, the true reward first rises, then falls, while the learned reward keeps climbing. This is reward overoptimization, Goodhart’s law with a learned reward, and Chapter 19 is about it. The KL term keeps the policy where the reward is still a good proxy for what we want.

2. It preserves what the reward doesn’t measure. With $\beta \to 0$, $\pi^*$ puts all its mass on the single highest-reward response per prompt. Every other skill the reference had, fluency, variety, knowledge the reward never asks about, is up for grabs. The leash keeps the policy a reweighting of the reference, so it keeps the reference’s breadth except where the reward argues otherwise. Korbak, Perez and Buckley (2022) push this view all the way: the KL-regularized optimum is a Bayesian posterior, with $\pi_\text{ref}$ as the prior and $e^{r/\beta}$ as the likelihood.

3. It stops degenerate exploits. Many reward hacks are text that no sensible model would write: repeated phrases, odd formatting, keyword stuffing. Such text has very low probability under $\pi_\text{ref}$, so reaching it costs a lot of KL.

4. It keeps training stable. A policy far from its starting point produces long, strange rollouts and huge gradients. Staying near a known-good model is a safety margin.

How strong is the leash in practice? It varies a lot, and the trend has been toward weaker leashes when the reward is a reliable checker. DeepSeekMath’s GRPO run used a KL coefficient of 0.04, and DeepSeek-R1-Zero used 0.001. Many 2025 recipes for RL with verifiable rewards (DAPO and Dr. GRPO among them) drop the KL term entirely, arguing that a verifier can’t be fooled the way a learned reward model can and that reasoning models should move far from the reference. Others keep it, or periodically reset $\pi_\text{ref}$ to the current policy. Chapter 16 treats the KL term as one of the knobs of the algorithm zoo. One sanity check the demo shows: whatever you do, a policy’s reward is only meaningful together with how far it has moved. When comparing methods, compare them on the reward–KL plane, not on reward alone.

Estimating the KL from samples

The objective needs $\mathrm{KL}(\pi_\theta \Vert \pi_\text{ref})$, a sum over all responses, which we can’t compute. Two facts make it manageable.

The KL of a sequence is a sum of per-token KLs. Because both models factor into next-token probabilities, the log-ratio of a response is a sum over its tokens, and so

$$\mathrm{KL}\big(\pi_\theta(\cdot \mid x) \,\Vert\, \pi_\text{ref}(\cdot \mid x)\big) = \mathbb{E}_{y \sim \pi_\theta}\Big[\sum_t \log \frac{\pi_\theta(y_t \mid s_t)}{\pi_\text{ref}(y_t \mid s_t)}\Big] = \mathbb{E}_{y \sim \pi_\theta}\Big[\sum_t \mathrm{KL}\big(\pi_\theta(\cdot \mid s_t) \,\Vert\, \pi_\text{ref}(\cdot \mid s_t)\big)\Big].$$

The second equality is the chain rule for KL: take the expectation over each token given its prefix. (In the worked example at $\beta = 1$: first-token KL 0.101, plus second-token KL 0.121 after $a$ or 0.328 after $b$, weighted by $\pi^*(a) = 0.484$ and $\pi^*(b) = 0.516$, gives $0.101 + 0.059 + 0.169 = 0.329$, the sequence KL from before.) That is why it is natural to charge the KL per token.

We only need samples. During RL you already have responses sampled from $\pi_\theta$, and for each token both log-probabilities, $\log \pi_\theta(y_t \mid s_t)$ and $\log \pi_\text{ref}(y_t \mid s_t)$ (one extra forward pass through the frozen reference). Write the sampled token’s log-ratio as

$$\ell_t = \log \frac{\pi_\theta(y_t \mid s_t)}{\pi_\text{ref}(y_t \mid s_t)} .$$

John Schulman’s note “Approximating KL divergence” (2020) compares three per-token estimators of $\mathrm{KL}(\pi_\theta \Vert \pi_\text{ref})$ built from it:

Estimator	Formula	Unbiased?	Can be negative?
$k_1$	$\ell_t$	yes	yes
$k_2$	$\tfrac12 \ell_t^2$	no (small bias near $\pi_\text{ref}$)	no
$k_3$	$e^{-\ell_t} - 1 + \ell_t \;=\; \frac{\pi_\text{ref}}{\pi_\theta} - 1 - \log \frac{\pi_\text{ref}}{\pi_\theta}$	yes	no

• $k_1$ is the definition of KL with the expectation replaced by a sample. Unbiased, but a single token’s value is negative whenever the sampled token is less likely under $\pi_\theta$ than under $\pi_\text{ref}$, although KL is never negative. Its variance can be large compared with the KL itself.
• $k_2$ is always non-negative and usually low-variance, but biased. Its expectation is a different divergence that agrees with KL to second order when $\pi_\theta \approx \pi_\text{ref}$, so the bias is small near the reference and grows as the policy moves.
• $k_3$ adds $u - 1$, with $u = \pi_\text{ref}/\pi_\theta$, to $k_1$‘s $-\log u$. Since $\mathbb{E}_{y \sim \pi_\theta}[u] = \sum_y \pi_\theta(y)\frac{\pi_\text{ref}(y)}{\pi_\theta(y)} = 1$, the added term has mean zero, so $k_3$ is still unbiased. It is a control variate, like the baselines of Chapter 7. And since $\log u \le u - 1$ for all $u > 0$ (the log lies below its tangent line at 1), every single $k_3$ value is $\ge 0$.

A worked example. One position, three tokens, $\pi_\theta = (0.6, 0.3, 0.1)$ and $\pi_\text{ref} = (0.4, 0.4, 0.2)$. The exact KL is $0.6 \ln 1.5 + 0.3 \ln 0.75 + 0.1 \ln 0.5 = 0.0877$.

sampled token	prob. under $\pi_\theta$	$\ell$	$k_1$	$k_2$	$k_3$
1	0.6	0.405	0.405	0.082	0.072
2	0.3	−0.288	−0.288	0.041	0.046
3	0.1	−0.693	−0.693	0.240	0.307
mean under $\pi_\theta$			0.0877	0.0858	0.0877
std. dev.			0.405	0.055	0.074

$k_1$ and $k_3$ hit the exact value on average; $k_2$ is 2% low. A single $k_1$ sample has a standard deviation almost five times the KL itself, while $k_3$‘s is about the size of the KL.

move π_θ away from π_ref: sharper flatter

tokens per estimate N

PlayDraw one batchReset

Each curve is what one sampled token contributes, as a function of its log-ratio ℓ = log π_θ/π_ref. Dots on the axis: the 8 tokens, sized by how often π_θ samples them. k₁ = ℓ   k₂ = ℓ²/2   k₃ = e^(−ℓ) − 1 + ℓ

	mean of estimates	bias	std. dev.	single tokens < 0
k₁
k₂
k₃

What to try: (1) Press Play with the default settings (π_θ a little flatter than π_ref, N = 16). All three histograms center near the dashed exact KL, but k₁'s is several times wider, and about half of the single k₁ values are negative. k₃ is never negative. (2) Drag the slider far to the right (π_θ much flatter, KL ≈ 1). k₂'s solid line separates from the dashed one: it is biased upward by about 45%. (3) Now drag far to the left (π_θ much sharper than π_ref). k₃'s histogram becomes skewed: most batches land below the true KL, with rare large values pulling the mean back. Its std. dev. overtakes k₁'s. (4) Set N = 1 and N = 256 and compare the spreads: they shrink like 1/√N for all three.

The penalty in the reward, or a term in the loss?

There are two places to put the KL, and they are not equivalent.

In the reward. Subtract $\beta\, k_1$ from the reward of every token, treat it as a fixed number (no gradient flows through it), and let the usual policy-gradient machinery do the rest. Per token, with the task reward at the end:

$$r_t = -\beta \log \frac{\pi_\theta(y_t \mid s_t)}{\pi_\text{ref}(y_t \mid s_t)} \;+\; \big[t = \lvert y \rvert\big]\, r(x, y).$$

Here $[t = \lvert y \rvert]$ is 1 for the last token and 0 otherwise. This is how InstructGPT-style PPO does it (Chapter 12), and RLOO does the same at the level of whole responses, with $R(x, y) = r(x, y) - \beta \sum_t \ell_t$. The score-function gradient of this shaped reward is an unbiased estimate of the gradient of the KL-regularized objective (when the samples come from the current policy). Exercise 11.6 shows why.

In the loss. GRPO (Shao et al., 2024) instead keeps the reward clean and adds $\beta$ times the average per-token $k_3$ to the loss, then differentiates it directly. That is convenient: the advantage stays a simple function of the task reward. But the gradient of an estimate is not, in general, an estimate of the gradient. For one position, differentiating $k_3$ and averaging over $y_t \sim \pi_\theta$ gives

$$\mathbb{E}_{y_t \sim \pi_\theta}\big[\nabla_\theta k_3\big] = \nabla_\theta\, \mathrm{KL}\big(\pi_\text{ref}(\cdot \mid s_t) \,\Vert\, \pi_\theta(\cdot \mid s_t)\big),$$

the gradient of the KL with its arguments swapped. Near $\pi_\text{ref}$ the two KLs agree to second order, so this still pulls the policy toward the reference, and in practice it works. But it is not the gradient of the objective written on the page, and the difference grows as the policy moves away. Tang and Munos (2025) point out this and a second pitfall (ignoring how the KL of later tokens depends on earlier ones) in several open-source implementations. Shah et al. (2025) compare the placements systematically. In their on-policy experiments, $k_1$ in the reward gave unbiased gradients and trained stably, $k_3$ in the loss (GRPO’s choice) was biased but stable, and $k_3$ used as a reward led to collapse. Exercise 11.6 derives the expected gradient of each estimator used as a loss.

Log-probs in practice: the shift and the mask

Every algorithm in Part III starts by computing, for a batch of sampled responses, the log-probability of each response token under the policy (and under the reference, and sometimes under the old policy). It looks trivial. It is also where much of the debugging time in real RL code goes. Here is one row of a TinyLM batch: the 7-column prompt <bos> 3 7 + 4 8 = (shorter prompts in the batch would be padded on the left so that all prompts end in the same column, $L = 7$), followed by a 4-token response budget, of which “85<eos>” used three:

column:        0      1   2   3   4   5   6  |  7    8     9      10
ids:         <bos>    3   7   +   4   8   =  |  8    5   <eos>   <pad>
                       prompt                |        response
logits col t predicts the token in col t+1:  col 6 -> "8", col 7 -> "5",
                                             col 8 -> <eos>, col 9 -> <pad>
response mask:                               |  1    1     1      0

Three rules:

1. Shift by one. The logits at column $t$ are the model’s prediction for the token in column $t + 1$. So the log-probs of the response tokens come from logits columns $L-1$ to the second-to-last column, matched with token columns $L$ to the end.
2. Mask the prompt. Prompt tokens are given, not chosen by the policy. They are not actions, and they must get no gradient (if they did, you would be training the model to write prompts).
3. Mask the padding. Tokens after <eos> were never generated. Exclude them, and do it with where, not by multiplying by 0: a padding position can hold $-\infty$, and $0 \times (-\infty)$ is NaN.

logits   = model(ids)                                 # (B, L_total, V)
logp     = log_softmax(logits[:, :-1], dim=-1)        # column t predicts token t+1
tok_logp = gather(logp, ids[:, 1:])                   # (B, L_total - 1): log-prob of the actual next token
tok_logp = tok_logp[:, L - 1:]                        # keep only the response tokens
seq_logp = where(mask, tok_logp, 0).sum(dim=-1)       # log π(y | x), one number per response

Real systems add details (attention masks so padding is never attended to, position ids that skip padding, packing several sequences into one row), but these three rules are where the bugs live. One more decision hides in the last line: when you turn per-token quantities into a loss, do you sum over a response’s tokens, average them per response, or average over all tokens in the batch? The choice changes how much long and short responses count, and Chapter 16 shows it creates a length bias. The lab’s Task 1 is exactly this code, with tests that catch each of the three bugs.

Summary

• A language model is a policy: logits → softmax → next-token distribution. A response’s probability is the product of its token probabilities, and its log-prob is the sum of their logs.
• Temperature divides the logits by $T$: $\pi_T \propto \pi^{1/T}$, a Boltzmann distribution. Low $T$ trades diversity for the model’s favorite answers; it does not know about the reward. Top-k and top-p truncate the tail.
• Pretraining and SFT are behavior cloning; RL maximizes a reward on the model’s own samples. The REINFORCE gradient is SFT on your own samples, weighted by reward. RL adds negative examples, learns from judgments instead of demonstrations, fixes exposure bias, and can exceed the demonstrator, but it can only reinforce what the model samples.
• Two MDP views: the token-level MDP (state = prompt + prefix, action = token, deterministic transitions, reward at the end, $\gamma = 1$) and the contextual bandit (action = whole response). They give the same REINFORCE gradient; algorithms differ in where advantages, ratios and the KL live.
• The KL-regularized objective $\mathbb{E}[r] - \beta\,\mathrm{KL}(\pi \Vert \pi_\text{ref})$ is maximized by the tilted reference $\pi^* = \pi_\text{ref}\, e^{r/\beta}/Z$. Token by token, $\pi^*$ tilts each token by the soft value of what can follow it: the optimum looks ahead.
• The leash exists because learned rewards are only trustworthy near the data (overoptimization), and to keep the reference’s breadth and block degenerate exploits. With reliable verifiers, many recent recipes weaken or drop it.
• Per-token KL estimators from sampled tokens: $k_1 = \ell$ (unbiased, noisy, can be negative), $k_2 = \ell^2/2$ (biased, low variance), $k_3 = e^{-\ell} - 1 + \ell$ (unbiased, non-negative, usually low variance, but heavy-tailed when $\pi_\theta$ is much sharper than $\pi_\text{ref}$). $k_1$ in the reward gives the right gradient; $k_3$ in the loss (GRPO) gives the gradient of the reversed KL.
• Computing log-probs: shift by one, mask the prompt, mask the padding (with where).

Exercises

Exercise 11.1Log-probs, products, and rejection samplingPen & paper●○○

(a) A model answers “37+48=” with “86<eos>”, with token probabilities 0.665, 0.08 and 0.97. Compute the response’s probability and log-prob. Which token contributes most to the log-prob?

(b) Show from the chain rule of probability that $\log \pi(y \mid x) = \sum_t \log \pi(y_t \mid x, y_{<t})$.

(c) Rejection-sampling fine-tuning: for each prompt, sample $N$ responses from $\pi_\theta$, keep the correct ones, and take an SFT step on them (average gradient of $\log \pi_\theta$ over the kept responses). Show that, for a single prompt, this gradient points in the same direction as the REINFORCE estimate $\frac{1}{N}\sum_i r^{(i)} \nabla_\theta \log \pi_\theta(y^{(i)} \mid x)$ with $r \in \lbrace 0, 1 \rbrace$. What is the scale factor between them, and what is missing compared with REINFORCE with a baseline?

Solution

(a) $\pi = 0.665 \times 0.08 \times 0.97 = 0.0516$. $\log \pi = \ln 0.665 + \ln 0.08 + \ln 0.97 = -0.408 - 2.526 - 0.030 = -2.964$ (and $e^{-2.964} = 0.0516$). The unlikely second digit “6” contributes 85% of the total. That’s typical: a response’s log-prob is dominated by its few surprising tokens.

(b) For any random sequence, $P(y_1, \dots, y_n) = P(y_1)\, P(y_2 \mid y_1) \cdots P(y_n \mid y_1, \dots, y_{n-1})$, the chain rule, applied conditionally on $x$. A language model defines exactly these conditionals, so $\pi(y \mid x) = \prod_t \pi(y_t \mid x, y_{<t})$, and the log of a product is the sum of the logs.

(c) Let $K$ be the set of kept (correct) responses, $\lvert K \rvert = k$. The SFT step’s gradient is $g_\text{RFT} = \frac{1}{k}\sum_{i \in K} \nabla \log \pi(y^{(i)})$. Since $r^{(i)} = 1$ for $i \in K$ and $0$ otherwise, REINFORCE gives $g_\text{RF} = \frac{1}{N}\sum_i r^{(i)} \nabla \log \pi(y^{(i)}) = \frac{1}{N}\sum_{i \in K} \nabla \log \pi(y^{(i)}) = \frac{k}{N}\, g_\text{RFT}$. Same direction, scale factor $k/N$ (the fraction correct), so rejection sampling takes relatively larger steps on hard prompts where few samples are correct. What’s missing is the baseline: wrong answers get weight 0 rather than a negative weight, so rejection sampling never pushes a specific mistake down. Its only route to lowering wrong answers is the indirect one through normalization.

I've done this one

Exercise 11.2Temperature, entropy and heat capacityPen & paper●●○

Let $\pi_T(a) = e^{z_a/T} / \sum_b e^{z_b/T}$.

(a) Show that $\pi_T(a) \propto \pi_1(a)^{1/T}$.

(b) For logits $z = (2, 1, 0)$, find the temperature at which the top token has probability 0.95.

(c) Write $\kappa = 1/T$ (the inverse temperature). Show that the entropy $H(\pi_T) = -\sum_a \pi_T(a) \log \pi_T(a)$ satisfies $\frac{dH}{d\kappa} = -\kappa\, \mathrm{Var}_{\pi_T}(z)$. Conclude that entropy increases with temperature, strictly unless all logits are equal.

Hint

For (c), write $\log \pi_T(a) = \kappa z_a - \log Z(\kappa)$, so $H = \log Z - \kappa\, \mathbb{E}[z]$. You need two facts about the log-partition function from statistical mechanics: $\frac{d}{d\kappa}\log Z = \mathbb{E}[z]$ and $\frac{d}{d\kappa}\mathbb{E}[z] = \mathrm{Var}(z)$.

Solution

(a) $\pi_1(a) = e^{z_a}/Z_1$, so $\pi_1(a)^{1/T} = e^{z_a/T}/Z_1^{1/T}$. The factor $Z_1^{1/T}$ is the same for every $a$, so $\pi_1(a)^{1/T} \propto e^{z_a/T} \propto \pi_T(a)$.

(b) $\pi_T(\text{top}) = \frac{1}{1 + e^{-\kappa} + e^{-2\kappa}} = 0.95$, so $e^{-\kappa} + e^{-2\kappa} = 1/0.95 - 1 = 0.05263$. With $u = e^{-\kappa}$: $u^2 + u - 0.05263 = 0$, so $u = \frac{-1 + \sqrt{1 + 0.2105}}{2} = 0.0501$, $\kappa = -\ln 0.0501 = 2.99$, and $T = 0.334$. Check: logits $(6, 3, 0)$ give $1/(1 + e^{-3} + e^{-6}) = 1/1.0522 = 0.950$.

(c) $\log \pi_T(a) = \kappa z_a - \log Z$ with $Z(\kappa) = \sum_b e^{\kappa z_b}$. Then $H = -\sum_a \pi_T(a)(\kappa z_a - \log Z) = \log Z - \kappa\, \mathbb{E}[z]$. Differentiate: $\frac{d \log Z}{d\kappa} = \frac{\sum_b z_b e^{\kappa z_b}}{Z} = \mathbb{E}[z]$, and $\frac{d\,\mathbb{E}[z]}{d\kappa} = \sum_a z_a \frac{d\pi_T(a)}{d\kappa} = \sum_a z_a \pi_T(a)(z_a - \mathbb{E}[z]) = \mathrm{Var}(z)$. So

$$\frac{dH}{d\kappa} = \mathbb{E}[z] - \mathbb{E}[z] - \kappa\, \mathrm{Var}(z) = -\kappa\, \mathrm{Var}_{\pi_T}(z) \le 0 .$$

Entropy falls as $\kappa$ rises, so it rises with $T = 1/\kappa$. It is strict unless $\mathrm{Var}(z) = 0$, which (since $\pi_T$ gives every token positive probability) means all logits are equal. In thermodynamic language, $\mathrm{Var}(z)$ is proportional to a heat capacity, which is never negative.

I've done this one

Exercise 11.3Two views of one tiny modelPen & paper●●○

Use the worked example’s reference model: $\pi_\text{ref}(a) = 0.7$, $\pi_\text{ref}(a \mid a) = 0.6$, $\pi_\text{ref}(a \mid b) = 0.5$, with rewards $r(aa) = 0$, $r(ab) = 1$, $r(ba) = 0$, $r(bb) = 2$.

(a) Suppose the first-token distribution is a softmax with logits $(\theta_a, \theta_b)$ (and the second-token distributions don’t depend on these parameters). In the bandit view, compute the exact policy gradient $\partial J / \partial \theta_b$ at $\pi = \pi_\text{ref}$, where $J = \sum_y \pi(y) r(y)$.

(b) In the token-level view, compute $\partial J/\partial \theta_b$ as $\mathbb{E}[\sum_t G_t\, \partial_{\theta_b} \log \pi(y_t \mid s_t)]$ and check that you get the same number. Then redo it with a per-token baseline: the value of the state before the first token, $V(x) = \mathbb{E}[r]$. Does the answer change?

(c) Use the chain rule for KL to compute $\mathrm{KL}(\pi^* \Vert \pi_\text{ref})$ at $\beta = 1$ from the token-level numbers in the text, and check that it matches the sequence-level 0.329.

Hint

For a softmax over two tokens, $\partial \log \pi(b) / \partial \theta_b = 1 - \pi(b)$ and $\partial \log \pi(a)/\partial \theta_b = -\pi(b)$. Only the first token depends on $\theta_b$, so only $t = 1$ terms survive in (b).

Solution

(a) $J = \pi(a)\,[0.6 \cdot 0 + 0.4 \cdot 1] + \pi(b)\,[0.5 \cdot 0 + 0.5 \cdot 2] = 0.4\,\pi(a) + 1.0\,\pi(b)$. With $\partial \pi(b)/\partial \theta_b = \pi(a)\pi(b)$ and $\partial \pi(a)/\partial\theta_b = -\pi(a)\pi(b)$:

$$\frac{\partial J}{\partial \theta_b} = (1.0 - 0.4)\, \pi(a)\pi(b) = 0.6 \times 0.7 \times 0.3 = 0.126 .$$

(b) Every token’s return is $G_t = r(y)$. Only $t = 1$ involves $\theta_b$. So the gradient is $\mathbb{E}[r(y)\,\partial_{\theta_b} \log \pi(y_1)]$:

• responses starting with $a$ (probability 0.7, mean reward 0.4): $\partial_{\theta_b} \log \pi(a) = -0.3$, contributing $0.7 \times 0.4 \times (-0.3) = -0.084$;
• responses starting with $b$ (probability 0.3, mean reward 1.0): $\partial_{\theta_b} \log \pi(b) = 0.7$, contributing $0.3 \times 1.0 \times 0.7 = 0.210$.

Total $0.126$, the same. With the baseline $V(x) = 0.58$, the weights become $0.4 - 0.58 = -0.18$ and $1.0 - 0.58 = 0.42$: $0.7 \times (-0.18)(-0.3) + 0.3 \times 0.42 \times 0.7 = 0.0378 + 0.0882 = 0.126$. Unchanged in expectation, as Chapter 7 proved; only the variance of the single-sample estimates changes.

(c) First-token KL: $0.484 \ln\frac{0.484}{0.7} + 0.516 \ln\frac{0.516}{0.3} = -0.178 + 0.280 = 0.101$. Second token after $a$: $0.356 \ln\frac{0.356}{0.6} + 0.644 \ln\frac{0.644}{0.4} = -0.186 + 0.307 = 0.121$. After $b$: $0.119 \ln\frac{0.119}{0.5} + 0.881 \ln\frac{0.881}{0.5} = -0.171 + 0.499 = 0.328$. Total: $0.101 + 0.484 \times 0.121 + 0.516 \times 0.328 = 0.101 + 0.059 + 0.169 = 0.329$. It matches.

I've done this one

Exercise 11.4The k₃ estimator: why it works, and where it strugglesPen & paper●●○

Samples $y \sim \pi_\theta$; $u = \pi_\text{ref}(y)/\pi_\theta(y)$; $k_3 = u - 1 - \log u$.

(a) Prove $k_3 \ge 0$ for every sample, with equality only when $u = 1$.

(b) Prove $\mathbb{E}_{y \sim \pi_\theta}[k_3] = \mathrm{KL}(\pi_\theta \Vert \pi_\text{ref})$. Which assumption about the supports of the two distributions do you need?

(c) Take $\pi_\theta = (0.99, 0.01)$ and $\pi_\text{ref} = (0.5, 0.5)$. Compute the exact KL, the values of $k_1$ and $k_3$ for each token, and their standard deviations.

(d) You estimate the KL by averaging $k_3$ over $N = 16$ sampled tokens. What is the probability that no rare token appears, and what is the estimate in that case? Compare with $k_1$.

Hint

(a) Study $f(u) = u - 1 - \log u$ on $u > 0$: its derivative and its minimum. (b) Write the expectation as a sum weighted by $\pi_\theta(y)$ and simplify $\pi_\theta(y)\, u$.

Solution

(a) $f(u) = u - 1 - \log u$ has $f'(u) = 1 - 1/u$, negative for $u < 1$ and positive for $u > 1$, and $f''(u) = 1/u^2 > 0$. So $f$ is strictly convex with its unique minimum at $u = 1$, where $f(1) = 0$. Hence $k_3 \ge 0$, with equality only at $u = 1$. (Equivalently: $\log u \le u - 1$, the log lies below its tangent at 1.)

(b) $\mathbb{E}[k_3] = \mathbb{E}[u] - 1 - \mathbb{E}[\log u]$. First, $\mathbb{E}[u] = \sum_y \pi_\theta(y)\frac{\pi_\text{ref}(y)}{\pi_\theta(y)} = \sum_{y:\, \pi_\theta(y) > 0} \pi_\text{ref}(y)$. This equals 1 if $\pi_\text{ref}$ puts no mass where $\pi_\theta$ has none. For softmax policies every token has positive probability, so the condition holds. Second, $-\mathbb{E}[\log u] = \sum_y \pi_\theta(y) \log\frac{\pi_\theta(y)}{\pi_\text{ref}(y)} = \mathrm{KL}(\pi_\theta \Vert \pi_\text{ref})$. So $\mathbb{E}[k_3] = 1 - 1 + \mathrm{KL} = \mathrm{KL}$. (If $\pi_\theta$ had zero mass on some tokens where $\pi_\text{ref}$ has mass, $\mathbb{E}[u] < 1$ and $k_3$ would underestimate.)

(c) $\mathrm{KL} = 0.99 \ln 1.98 + 0.01 \ln 0.02 = 0.6763 - 0.0391 = 0.637$.

• Common token: $\ell = \ln(0.99/0.5) = 0.683$, so $k_1 = 0.683$ and $k_3 = 0.505 - 1 + 0.683 = 0.188$.
• Rare token: $\ell = \ln(0.01/0.5) = -3.912$, so $k_1 = -3.912$ and $k_3 = 50 - 1 - 3.912 = 45.09$.

Check the means: $k_1$: $0.99(0.683) + 0.01(-3.912) = 0.637$. $k_3$: $0.99(0.188) + 0.01(45.09) = 0.186 + 0.451 = 0.637$. Both unbiased. Standard deviations: $k_1$: $\sqrt{0.99 \cdot 0.01}\,(0.683 + 3.912) = 0.0995 \times 4.595 = 0.457$. $k_3$: $0.0995 \times (45.09 - 0.188) = 4.47$, ten times larger.

(d) $P(\text{no rare token in 16}) = 0.99^{16} = 0.851$. In that case the $k_3$ estimate is $0.188$, less than a third of the true 0.637, and the $k_1$ estimate is $0.683$, 7% high. So 85% of the time $k_3$ underestimates by a factor of 3.4. When exactly one rare token appears, the estimate jumps to $(45.09 + 15 \times 0.188)/16 \approx 3.0$, almost five times too high. Unbiased on average, but in a single batch it is usually far too low. The same thing happens with TinyLM in the lab when the sampler is much sharper than the reference. The lesson: $k_3$ is excellent when $\pi_\theta$ is close to $\pi_\text{ref}$ or broader than it, and treacherous when $\pi_\theta$ has become much more confident, which is what RL tends to do. Exact per-position KL over the whole vocabulary avoids the problem at the cost of more computation (some recent systems do this).

I've done this one

Exercise 11.5Derive the tilted optimum, then build it token by tokenPen & paper●●○

(a) Without looking back, derive $\pi^*(y) = \pi_\text{ref}(y) e^{r(y)/\beta}/Z$ as the maximizer of $\mathbb{E}_\pi[r] - \beta\,\mathrm{KL}(\pi \Vert \pi_\text{ref})$ over distributions on a finite set, using a Lagrange multiplier for $\sum_y \pi(y) = 1$.

(b) Redo the worked example at $\beta = 0.5$: compute $V(a)$, $V(b)$, $\beta \log Z$, the first-token probabilities of $\pi^*$, and $\pi^*(bb)$.

(c) What happens to $\pi^*(y_1)$ as $\beta \to 0$ and as $\beta \to \infty$? Express your answer in terms of $\max_{y_2} r(y_1 y_2)$.

Hint

(a) Differentiate $\sum_y \pi(y) r(y) - \beta \sum_y \pi(y)[\log \pi(y) - \log \pi_\text{ref}(y)] + \lambda(1 - \sum_y \pi(y))$ with respect to one $\pi(y)$. (c) The soft value $\beta \log \sum_{y_2} \pi_\text{ref}(y_2 \mid y_1) e^{r/\beta}$ is a “soft max”.

Solution

(a) Setting the derivative to zero: $r(y) - \beta[\log \pi(y) + 1 - \log \pi_\text{ref}(y)] - \lambda = 0$, so $\pi(y) = \pi_\text{ref}(y)\, e^{r(y)/\beta} e^{-1 - \lambda/\beta}$. The last factor doesn’t depend on $y$; normalization fixes it to $1/Z$. The objective is concave (linear term plus $-\beta$ times a convex function), so this stationary point is the global maximum. (Chapter 10 gives the full argument.)

(b) With $\beta = 0.5$, $e^{r/\beta} = e^{2r}$:

• $V(a) = 0.5 \ln(0.6 + 0.4 e^{2}) = 0.5 \ln(3.556) = 0.634$.
• $V(b) = 0.5 \ln(0.5 + 0.5 e^{4}) = 0.5 \ln(27.80) = 1.662$.
• $Z = 0.7 \cdot 3.556 + 0.3 \cdot 27.80 = 2.489 + 8.340 = 10.83$, and $\beta \log Z = 0.5 \ln 10.83 = 1.191$.
• $\pi^*(a) = 2.489/10.83 = 0.230$ and $\pi^*(b) = 0.770$.
• $\pi^*(b \mid b) = 0.5 e^4 / 27.80 = 27.30/27.80 = 0.982$, so $\pi^*(bb) = 0.770 \times 0.982 = 0.756$.

Check against the sequence level: $0.15 e^{4}/10.83 = 8.19/10.83 = 0.756$. (The demo shows the same numbers at β = 0.5.)

(c) As $\beta \to \infty$, expand $e^{r/\beta} \approx 1 + r/\beta$: the soft value tends to the plain average, $V(y_1) \to \mathbb{E}_{\pi_\text{ref}}[r \mid y_1]$, and $V_0 \to \mathbb{E}_{\pi_\text{ref}}[r]$. Both stay finite, so $(V(y_1) - V_0)/\beta \to 0$ and $\pi^*(y_1) \to \pi_\text{ref}(y_1)$. As $\beta \to 0$, the soft max becomes a hard max: $V(y_1) \to \max_{y_2} r(y_1 y_2)$ (over continuations with $\pi_\text{ref} > 0$), and $\pi^*(y_1) \propto \pi_\text{ref}(y_1) e^{V(y_1)/\beta}$ concentrates on the first token(s) with the largest $\max_{y_2} r(y_1 y_2)$. Here: $b$, because $\max r$ after $b$ is 2 and after $a$ only 1. The first token’s fate is decided by the best thing that can follow it, not the average.

I've done this one

Exercise 11.6Which KL gradient does your loss compute?Pen & paper●●●

Consider one position with a softmax policy $\pi_\theta$ over a finite vocabulary and a fixed reference $\pi_\text{ref}$. Tokens are sampled $y \sim \pi_\theta$ and $\ell = \log \pi_\theta(y) - \log \pi_\text{ref}(y)$. Write $s(y) = \nabla_\theta \log \pi_\theta(y)$ for the score. In this setting $\mathrm{KL}(\pi_\theta \Vert \pi_\text{ref})$, the one in the objective, is called the reverse KL, and $\mathrm{KL}(\pi_\text{ref} \Vert \pi_\theta)$, with the arguments swapped, the forward KL.

(a) Show that the true gradient of the reverse KL is $\nabla_\theta \mathrm{KL}(\pi_\theta \Vert \pi_\text{ref}) = \mathbb{E}_{y \sim \pi_\theta}[\ell\, s(y)]$.

(b) “KL in the reward”: treat $\ell$ as a fixed number and use the score-function estimator $\ell\, s(y)$. Show it is unbiased for the gradient in (a).

(c) “KL in the loss”: compute $\mathbb{E}_{y \sim \pi_\theta}[\nabla_\theta k]$ when you differentiate the estimator itself, for $k = k_1$, $k = k_2$ and $k = k_3$. Which KL, if any, does each one differentiate?

Hint

Use $\mathbb{E}_{y \sim \pi_\theta}[s(y)] = \sum_y \nabla_\theta \pi_\theta(y) = 0$ repeatedly. For $k_3$, note $\nabla_\theta\, \frac{\pi_\text{ref}(y)}{\pi_\theta(y)} = -\frac{\pi_\text{ref}(y)}{\pi_\theta(y)}\, s(y)$. For the forward KL, $\mathrm{KL}(\pi_\text{ref} \Vert \pi_\theta) = \sum_y \pi_\text{ref}(y)\log\pi_\text{ref}(y) - \sum_y \pi_\text{ref}(y) \log \pi_\theta(y)$.

Solution

(a) $\mathrm{KL} = \sum_y \pi_\theta(y)\,\ell(y)$, and both factors depend on $\theta$. Product rule: $\nabla \mathrm{KL} = \sum_y \nabla\pi_\theta(y)\,\ell(y) + \sum_y \pi_\theta(y) \nabla \ell(y)$. The first sum is $\sum_y \pi_\theta(y) s(y) \ell(y) = \mathbb{E}[\ell s]$. In the second, $\nabla \ell(y) = s(y)$ (the reference doesn’t depend on $\theta$), so it is $\mathbb{E}[s] = 0$. Hence $\nabla \mathrm{KL}(\pi_\theta \Vert \pi_\text{ref}) = \mathbb{E}[\ell\, s]$.

(b) Its expectation is $\mathbb{E}[\ell\, s]$, the right-hand side of (a). Unbiased. This is why ”$-\beta\, k_1$ as a per-token reward, no gradient through it” implements the KL-regularized objective correctly on-policy. (The full sequence version needs the KL of later tokens to be included in each token’s return, which reward-to-go does automatically; that’s the sequential pitfall of Tang and Munos.)

(c) Differentiate through the sampled value, keeping $y$ fixed:

• $k_1 = \ell$: $\nabla k_1 = s(y)$, and $\mathbb{E}[s] = 0$. Zero expected gradient. As a loss term, $k_1$ does nothing on average (it only adds noise).
• $k_2 = \tfrac12 \ell^2$: $\nabla k_2 = \ell\, s(y)$, with expectation $\mathbb{E}[\ell s] = \nabla \mathrm{KL}(\pi_\theta \Vert \pi_\text{ref})$. The correct reverse-KL gradient, although $k_2$‘s value is biased. (This is the argument REINFORCE++ makes for $k_2$ in the loss.)
• $k_3 = u - 1 + \ell$ with $u = \pi_\text{ref}/\pi_\theta$: $\nabla k_3 = -u\, s + s = (1 - u)\,s$. Expectation: $\mathbb{E}[s] - \sum_y \pi_\theta(y)\frac{\pi_\text{ref}(y)}{\pi_\theta(y)} s(y) = -\sum_y \pi_\text{ref}(y)\, \nabla \log \pi_\theta(y)$. From the hint, that is $\nabla \mathrm{KL}(\pi_\text{ref} \Vert \pi_\theta)$. The gradient of the forward KL, with the arguments swapped.

So, used as a loss: $k_1$ is useless, $k_2$ is right, and $k_3$ (GRPO’s choice) optimizes a different divergence. Near $\pi_\text{ref}$ both KLs have the same second-order expansion (Fisher metric), so the three options that do something ($k_1$ in the reward, $k_2$ or $k_3$ in the loss) push toward the reference in about the same way, which is why GRPO’s choice works in practice. Far from the reference they differ: the forward KL is “mass-covering” (it heavily penalizes putting tiny probability on tokens $\pi_\text{ref}$ likes), the reverse KL is “mode-seeking”.

I've done this one

Exercise 11.7Sampling temperature as a KL coefficientPen & paper●●●

(a) Fix $T \in (0, 1)$ and define the sequence-level power distribution $p_T(y \mid x) \propto \pi_\text{ref}(y \mid x)^{1/T}$. Show that $p_T$ is the KL-regularized optimum $\pi^*$ for the reward $r(x, y) = \log \pi_\text{ref}(y \mid x)$, and find $\beta$ as a function of $T$.

(b) Ordinary sampling at temperature $T$ applies $\pi^{1/T}$ to each next-token distribution separately. Show by example that this is not the same distribution over responses as $p_T$. Use a reference with $\pi_\text{ref}(a) = \pi_\text{ref}(b) = 0.5$, $\pi_\text{ref}(\cdot \mid a)$ uniform over two tokens, and $\pi_\text{ref}(\cdot \mid b)$ putting all its mass on one token, at $T = 0.5$.

(c) Explain the difference using soft values.

Hint

(a) Write $\pi_\text{ref} e^{r/\beta}$ with $r = \log \pi_\text{ref}$ and match exponents. (b) There are three responses with nonzero probability: “aa”, “ab” and “bc” (say). (c) Under $p_T$, which first token has the larger soft value?

Solution

(a) $\pi_\text{ref}(y)\, e^{\log \pi_\text{ref}(y)/\beta} = \pi_\text{ref}(y)^{1 + 1/\beta}$. This is $\propto \pi_\text{ref}^{1/T}$ when $1 + 1/\beta = 1/T$, that is $\beta = \frac{T}{1 - T}$. For example $T = 0.5$ corresponds to $\beta = 1$ with “reward = the reference’s own log-likelihood”. Lowering the temperature is a KL-regularized optimization toward likely text: no outside reward involved. ($T \to 1$ gives $\beta \to \infty$, no change; $T \to 0$ gives $\beta \to 0$, the single most likely response.)

(b) Reference response probabilities: $aa = 0.25$, $ab = 0.25$, $bc = 0.5$.

• Power distribution at $T = 0.5$: square and renormalize: $0.0625, 0.0625, 0.25$, divided by $0.375$: $(0.167, 0.167, 0.667)$.
• Token-level temperature: the first token is $0.5^2$ vs $0.5^2$, still $(0.5, 0.5)$; after $a$ it stays uniform; after $b$ it stays deterministic. Responses: $(0.25, 0.25, 0.5)$, unchanged from the reference.

Different distributions. Token-level sharpening can’t see that choosing $b$ leads to a confident continuation.

(c) Under $p_T$ (as an instance of $\pi^*$ with $r = \log \pi_\text{ref}$, $\beta = 1$), the soft value of a first token is $V(y_1) = \log \sum_{y_2} \pi_\text{ref}(y_2 \mid y_1)\, e^{\log \pi_\text{ref}(y_1 y_2)}$, which is larger after $b$ (one continuation with all the mass) than after $a$ (mass split two ways). The power distribution looks ahead, exactly like the tilted optimum in the text; per-token temperature doesn’t. Karan and Du (2025) sample from the sequence-level power distribution of base models with an MCMC method, and report reasoning gains approaching those of RL-trained models without any training, which is one reason this distinction has attracted recent attention.

I've done this one

Exercise 11.8Why the reward model needs a leashConcept●●○

A reward model was trained on comparisons between responses sampled from the SFT model.

(a) Explain in two or three sentences why its errors are likely to be largest on responses that the SFT model rarely produces, and why an RL optimizer with $\beta = 0$ will find those responses.

(b) In the demo’s “blind spot” scenario, the reference gives the response “cc” probability 0.02 and the reward model scores it 1.6 instead of its true 0. The best honest response, “bb”, has reference probability 0.24 and reward 1.0. Using the tilted form $\pi^* \propto \pi_\text{ref} e^{r/\beta}$, find the $\beta$ at which $\pi^*$ puts equal mass on “bb” and “cc”. What happens to the true reward for $\beta$ well below that value?

(c) Give two reasons, other than reward-model errors, to keep $\beta > 0$, and one reason recent RL-with-verifiers recipes give for setting $\beta = 0$.

(d) In one sentence each: what does PPO’s clip keep $\pi_\theta$ close to, and what does the KL term keep it close to? Why would you want both?

Hint

(a) Think of the reward model’s score as true reward plus an error, and ask what happens when you take the maximum over many responses with noisy scores (Chapter 5’s maximization bias). (b) Set $\pi_\text{ref}(bb)\, e^{r(bb)/\beta} = \pi_\text{ref}(cc)\, e^{r(cc)/\beta}$ and solve for $\beta$. (c) and (d): reread “Why a leash?” and the callout “Two different leashes”.

Solution

(a) The reward model learned from examples near $\pi_\text{ref}$‘s typical outputs; on unusual responses it is extrapolating, and nothing in its training penalized wrong scores there. An optimizer looks for the highest score anywhere, and among many responses with noisy scores, the highest is very likely one whose error is large and positive (the same “max of noisy estimates” effect as maximization bias in Chapter 5). With $\beta = 0$ nothing stops it from moving all its mass there.

(b) Equal mass when $0.24\, e^{1.0/\beta} = 0.02\, e^{1.6/\beta}$, i.e. $e^{0.6/\beta} = 12$, so $\beta = 0.6/\ln 12 = 0.241$. Above that, “bb” dominates the two; below it, “cc” does, and as $\beta \to 0$, $\pi^*(cc) \to 1$ and the true reward falls toward 0, below even the reference’s 0.46. In the demo the true reward peaks around $\beta \approx 0.42$ at about 0.63 and is 0.03 at $\beta = 0.1$. (The peak comes well before the crossover at 0.241: below $\beta \approx 0.42$, “bb” stops gaining mass, at about 0.5, while “cc”, already at 18%, keeps growing.)

(c) Keep $\beta > 0$ to preserve what the reward doesn’t measure (fluency, diversity, general knowledge: as $\beta \to 0$, $\pi^*$ collapses onto one response per prompt), to make degenerate exploits expensive (text with tiny $\pi_\text{ref}$ costs a lot of KL), and for stability. Reason to drop it: a programmatic verifier is much harder to fool than a learned reward model, and for long reasoning the model should move far from the SFT model; recipes like DAPO argue the KL term just slows that down.

(d) PPO’s clip keeps $\pi_\theta$ close to $\pi_\text{old}$, the policy that collected the current batch, so each update is a trustworthy small step. The KL term keeps it close to $\pi_\text{ref}$, the fixed starting model, which changes what the final answer is. One makes each step safe; the other decides how far the whole journey may go.

I've done this one

Exercise 11.9Explore the demosExplore the demo●○○

(a) In “Tilting a toy language model”, worked scenario, set $\beta = 0.5$ and check your answers to Exercise 11.5(b) against the readouts and the tree.

(b) Still in the worked scenario, find a temperature $T < 1$ at which “π_ref at temperature T” has the same KL from $\pi_\text{ref}$ as $\pi^*$ at $\beta = 2$ (KL ≈ 0.08). Compare the two expected rewards. Then push $T$ to its maximum. Why can’t temperature reach the frontier in either direction?

(c) Switch to “Reward model with a blind spot”. Find the $\beta$ that maximizes the true reward (dashed red). What is the KL there? How much proxy reward does $\pi^*$ have at that point, and how much at $\beta = 0.1$?

(d) In “Three ways to estimate a KL”, set the slider to the far left (sharper) and $N = 16$. After a few thousand batches, what fraction of the $k_3$ histogram lies below the true KL? Which estimator has the smallest standard deviation now?

Solution

(a) At $\beta = 0.5$: $V(a) = 0.63$, $V(b) = 1.66$, $\beta \log Z = 1.191$, $\pi^*(a) = 0.23$, $\pi^*(bb) = 0.756$, matching Exercise 11.5. The slider’s nearest setting is $\beta = 0.501$, so the readouts may differ in the third decimal ($\beta \log Z = 1.190$, $\pi^*(bb) = 0.755$).

(b) Around $T \approx 0.48$ the KL reaches about 0.08, and the expected reward there is about 0.40, far below $\pi^*$‘s 0.887 at the same KL (and below the reference’s 0.58: lowering $T$ favors “aa”, the likeliest response, which earns 0). At the maximum, $T = 3.16$, the KL is only about 0.05 and the reward about 0.70, while $\pi^*$ at that KL earns about 0.81. Temperature reshapes the distribution using only $\pi_\text{ref}$‘s own probabilities; it has no information about the reward, so it can’t move mass specifically toward “ab” and “bb”. $\pi^*$ is by construction the best policy at each KL (the Gibbs identity), so every other policy lies on or below the frontier.

(c) The true reward peaks at about 0.63 near $\beta \approx 0.42$, where $\mathrm{KL} \approx 0.5$ nats and the proxy reward is about 0.90. (The peak is flat: anywhere from $\beta \approx 0.37$ to $0.5$ the true reward is within 0.01 of it.) At $\beta = 0.1$ the proxy reward is about 1.58 while the true reward is about 0.03. More “reward”, much worse model: overoptimization in miniature.

(d) Roughly 70% of the batch estimates fall below the true KL (the mean still matches it, pulled up by a long right tail). $k_2$ has the smallest standard deviation but is now clearly biased low; $k_1$ has a smaller standard deviation than $k_3$ in this regime, the reverse of the “flatter” side.

I've done this one

Exercise 11.10Lab: the plumbing of LLM RLCode●●○

Do the lab 11-llm-as-policy (download below). In policy.py:

• Task 1a response_logprobs and Task 1b sequence_logprob: the shift, the gather, and the mask.
• Task 2 sample_with_temperature, used by the provided generate to run the token-level MDP on TinyLM.
• Task 3 kl_estimators: $k_1$, $k_2$, $k_3$.
• Task 4 tilted_optimum and kl_objective: $\pi^*$ over all responses of the toy LM, in log space.
• Task 5 token_level_optimum: the same $\pi^*$ from soft values.

Then run run.py. In panel 2, at which temperatures do $k_1$ and $k_3$ disagree, and which one agrees with the exact (full-vocabulary) value? Connect your answer to Exercise 11.4.

Hint

Task 1a: if your test says the log-probs are off, write out which logits column predicts which token column for the 4-column example in test_task1_shift_on_a_hand_example. Task 1b: torch.where(mask.bool(), tok_logp, torch.zeros_like(tok_logp)). Task 4: np.errstate(divide="ignore") around np.log(ref_probs), then subtract the max before exponentiating.

Solution

The key lines (full code in solutions/policy.py):

# Task 1a
logp = F.log_softmax(logits[:, :-1], dim=-1)               # column t predicts token t+1
tok = logp.gather(-1, ids[:, 1:].unsqueeze(-1)).squeeze(-1)
return tok[:, prompt_len - 1:]
# Task 1b
return torch.where(mask.bool(), tok_logp, torch.zeros_like(tok_logp)).sum(-1)
# Task 2
if temperature == 0: return logits.argmax(-1)
return torch.multinomial(F.softmax(logits / temperature, -1), 1, generator=generator).squeeze(-1)
# Task 3
l = logp - ref_logp; return l, 0.5 * l**2, torch.exp(-l) - 1 + l
# Task 4
logw = np.log(ref_probs) + rewards / beta; m = logw.max(); w = np.exp(logw - m)
return w / w.sum(), m + np.log(w.sum())
# Task 5
V = beta * np.logaddexp.reduce(np.log(second) + rewards / beta, axis=1)
V0 = beta * np.logaddexp.reduce(np.log(first) + V / beta)
return first * np.exp((V - V0) / beta), second * np.exp((rewards - V[:, None]) / beta), V, V0

In run.py’s panel 2, the two estimators agree with each other and with the exact value for $T \ge 0.75$ (within their error bars). At $T = 0.25$, $k_3$ is far below the exact value, by much more than its error bar, while $k_1$ matches it; at $T = 0.5$ both are noisy and $k_3$ is again on the low side. A sampler at low temperature is much sharper than the $T = 1$ model, so the tokens where $\pi_\text{ref}/\pi_\theta$ is huge are almost never drawn: exactly the situation of Exercise 11.4(d). At high temperature the sampler is flatter than the reference, the ratio stays bounded, and $k_3$ behaves well.

I've done this one

Further reading

• Schulman (2020), Approximating KL Divergence. Two pages that define $k_1$, $k_2$, $k_3$, with a small numerical comparison. The source of the estimator names used everywhere.
• Ouyang et al. (2022), Training language models to follow instructions with human feedback. InstructGPT: the pretraining → SFT → RL pipeline and the KL-penalized reward, in the paper that made it standard. Section 3.5 states the KL-penalized RL objective.
• Rafailov, Hejna, Park & Finn (2024), From r to Q*: Your Language Model is Secretly a Q-Function. Works out the token-level MDP view of the KL-regularized objective that this chapter sketched with soft values, and connects it to DPO.
• Korbak, Perez & Buckley (2022), RL with KL penalties is better viewed as Bayesian inference. Short and clarifying: $\pi^*$ as a posterior with $\pi_\text{ref}$ as the prior, and why plain reward maximization collapses the distribution.
• Shah et al. (2025), A Comedy of Estimators: On KL Regularization in RL Training of LLMs, and Tang & Munos (2025), On a few pitfalls in KL divergence gradient estimation for RL. Where the KL goes (reward or loss) and which estimator you differentiate change the gradient you actually get; the first paper also tests the consequences on 4–8B models.
• Holtzman et al. (2020), The Curious Case of Neural Text Degeneration. Why pure sampling and greedy decoding both fail for open-ended text, and nucleus (top-p) sampling.
• Lambert, Reinforcement Learning from Human Feedback (online book). A practitioner’s tour of the whole pipeline, with a chapter on regularization; good to read alongside Part III.

Next, Chapter 12 puts all of this to work: a reward model learned from human comparisons, and PPO climbing the KL-regularized objective, one token at a time.