From next-token prediction to assistants that reason

1One model per task

Until the late 2010s, natural language processing mostly built one model per problem. A sentiment classifier, a question-answering system and a translation model each had their own architecture, their own labelled dataset and their own training run. If you had trained a good model and then needed to solve a different problem, you trained another model.

What was shared was the bottom layer. Word embeddings such as word2vec (Mikolov et al., 2013) and GloVe (Pennington et al., 2014) learned one vector per word from unlabelled text, and task models used those vectors as input features. This was already transfer learning, but a shallow kind: each word got a single vector regardless of context, so bank in “river bank” and “bank account” looked the same, and everything above the embeddings still had to be learned from scratch for every task.

Between 2015 and 2020 that changed. The field converged on a different recipe: train one large language model on a very large amount of unlabelled text, keep it, and adapt it to each task. That recipe, “train once broadly, adapt many times cheaply”, is the embryo of today’s split between pre-training and post-training.

2Pre-train, then adapt

The quiet start: semi-supervised sequence learning (2015)

Dai and Le (2015) pre-trained an LSTM as a language model, or as a sequence autoencoder, and used its weights to initialise a supervised text classifier. The pre-trained networks trained more stably and generalised better on several classification benchmarks. The idea that a language model need not be the final product, but can be a pre-training stage for supervised learning, was already on the table.

It did not take over immediately. LSTMs process a sequence one step at a time, which makes training hard to parallelise and expensive to scale. Without strong evidence that pre-training the whole network paid off, there was little reason to abandon the familiar pattern of pre-trained embeddings plus a task-specific model.

Contextual representations: ELMo (2018)

ELMo (Peters et al., 2018) attacked the one-vector-per-word limitation. It trained a bidirectional LSTM language model, one direction predicting the next word and the other the previous word, so every word’s internal representation depends on its context. Different layers capture different information, so ELMo’s final representation is a task-specific weighted combination of all layers, with the weights learned for each task.

ELMo showed that a large pre-trained language model learns representations that transfer across many tasks, and it moved the field from static to contextual representations. But it was still feature-based transfer: the pre-trained model produced features that were fed into a separate task model.

Fine-tuning the whole model: ULMFiT (2018)

ULMFiT (Howard and Ruder, 2018) took the other road: if we already have a general language model, fine-tune that model on the target task. Its recipe had three stages: pre-train a language model on general text; continue language-model training on the target task’s own text, so the model adapts to that domain’s distribution; then add a classifier and train on the labels. The middle stage anticipates today’s continued pre-training and staged adaptation.

Fine-tuning a whole network on a small dataset risks catastrophic forgetting: large updates can erase what pre-training learned. ULMFiT used discriminative fine-tuning (different learning rates for different layers), slanted triangular learning-rate schedules and gradual unfreezing (unfreeze the top layer first, then the next, and so on). The specific tricks matter less today; the tension between adapting and forgetting is still central to post-training.

Its most striking result was data efficiency: the paper reports that with only 100 labelled examples, ULMFiT matched the performance of training from scratch on 100 times more data. That is strong evidence that the pre-trained model already held most of what the task needed, and fine-tuning mainly had to adapt it.

A task-agnostic Transformer: GPT-1 (2018)

GPT-1 (Radford et al., 2018) combined pre-train-then-fine-tune with the Transformer (Vaswani et al., 2017), which processes all positions in parallel and made far larger pre-training runs practical. It pre-trained a decoder-only Transformer on a corpus of books and then fine-tuned it on each task.

Its key design choice was to change the input, not the model. Tasks with very different structures were turned into token sequences through task-specific input transformations: for entailment, premise and hypothesis concatenated with a delimiter token; for similarity, both orderings of the two sentences; for multiple choice, the context paired with each candidate answer. The same network handled them all, with only a small linear output layer added. This is an early form of “write the task as text”; today users state the task in natural language, whereas GPT-1 used a fixed format chosen by the designers.

During fine-tuning GPT-1 also kept a language-modelling loss as an auxiliary objective, so the model would not lose its language-modelling ability, the same worry ULMFiT faced, solved differently. The paper’s ablations found the auxiliary loss helped on the larger datasets but not on the smaller ones, so it is not GPT-1’s main contribution.

More important was an observation about zero-shot behaviour. The authors measured how well the pre-trained model, before any fine-tuning, could solve a few tasks using simple heuristics, at checkpoints throughout pre-training. Performance rose fairly steadily as pre-training continued. Their hypothesis, which still shapes the field, was that to predict the next token well, a language model is forced to learn abilities that downstream tasks need, partly because unlabelled text already contains examples of those tasks. An LSTM showed similar behaviour with higher variance, which they took as a hint that the Transformer’s inductive bias helps transfer. On other tasks the zero-shot heuristics were weak, so this did not make fine-tuning unnecessary.

If pre-training already gives a model general representations and some problem-solving ability, does post-training teach new abilities, or does it select and release abilities the model already has?

Bidirectional pre-training: BERT (2018)

At the end of 2018 BERT (Devlin et al., 2019) pre-trained a bidirectional Transformer encoder with masked language modelling (predict randomly masked tokens from both sides) and a next-sentence-prediction task, then fine-tuned it on a wide range of language-understanding benchmarks with strong results. BERT made “pre-train a large model, then fine-tune for each task” the mainstream paradigm in NLP. It still needed a separate fine-tuning run per task, though, and an encoder does not generate text, so the unified task interface was not there yet.

3Scale and prompting

If one model can be adapted to so many tasks, and pre-training alone already produces some zero-shot ability, do we need to fine-tune for every task at all?

GPT-2: language models as unsupervised multitask learners (2019)

The GPT-2 paper’s title states its thesis (Radford et al., 2019). Web text is full of implicit task demonstrations: an interview transcript pairs “Q:” with a question and “A:” with its answer; an article followed by “TL;DR:” is followed by a summary. A model that predicts such text well may have to learn what follows “Q: … A:”. So train a large enough model on enough text, and tasks may be learned without being labelled as tasks.

GPT-2 had 1.5 billion parameters, more than ten times GPT-1, and was trained on WebText, about 40 GB of text from roughly 8 million web documents. The authors trained several model sizes and evaluated them zero-shot, without any fine-tuning; performance improved roughly log-linearly with model size across tasks. That is an early signal of the scaling behaviour later studied systematically (Kaplan et al., 2020).

GPT-2 clearly acquired task-relevant behaviour from pre-training alone. It also showed why pre-training is not the whole story: asked a question, a base model may continue the document instead of answering it, for example by writing more questions. It has the knowledge, but not the habit of using it the way a user wants. That gap is a central motivation for instruction tuning.

T5: everything is text-to-text (2019)

T5 (Raffel et al., 2020) introduced no new architecture. Its contribution was a unified format: every task is “text in, text out”, with a task prefix such as translate English to German: or summarize:. One model and one loss then serve translation, summarisation, classification and question answering. T5 was still fine-tuned on downstream tasks, but its systematic study also found that performance kept improving with scale, up to its largest, 11-billion-parameter model.

GPT-3: in-context learning (2020)

GPT-3 (Brown et al., 2020) had 175 billion parameters, two orders of magnitude more than GPT-2, and the paper trained eight sizes from 125 million parameters upwards. Across many tasks, zero-shot, one-shot and few-shot performance generally improved with size, and the gap between zero-shot and few-shot often widened as models grew.

Few-shot in-context learning means putting a handful of worked examples in the prompt, with no weight updates, and letting the model continue the pattern. Given

the model is very likely to produce fromage. The examples change the model’s output distribution towards what we want without changing a single parameter. Arguably, the most important scaling result of GPT-3 is not that larger models generate better text but that they get better at using prompts to do tasks.

Supervised fine-tuning is the natural next step. In-context learning shows that a few demonstrations can temporarily steer a fixed model towards a task; SFT writes that “do what I ask” behaviour into the parameters, permanently and for every prompt. The rest of post-training is built on top of it.

4What SFT optimises

Supervised fine-tuning (SFT), also called instruction tuning, continues training a pre-trained model on a dataset of prompt–response pairs, so that it learns to follow instructions and answer in the intended way. It usually comes straight after pre-training and is the bridge from general language ability to useful behaviour.

The training signal is the same as in pre-training: predict the next token. What changes is the data. Suppose we have $N$ pairs of a prompt $x$ and a reference response $y=(y_1,\dots ,y_T)$:

The factorisation is just the chain rule of probability: the model generates one token at a time, each conditioned on the prompt and on the tokens before it. Maximum likelihood chooses the parameters that make the observed responses as probable as possible.

Each term is the cross-entropy between the one-hot distribution on the correct token and the model’s predicted distribution over the vocabulary, so SFT is token-level classification, exactly as in pre-training. Averaged over the data, minimising this loss is equivalent to minimising the KL divergence $\mathrm{K}\mathrm{L}\left(p_{\mathrm{d}\mathrm{a}\mathrm{t}\mathrm{a}}\Vert P_{\theta }\right)$, because the remaining entropy term of the data does not depend on $\theta$. This “forward” KL heavily penalises the model for giving low probability to anything in the data, so SFT tends to cover the demonstrations rather than pick the best of them, a point that matters in §7.

5Data, templates and loss masks

Where SFT data comes from

• Human demonstrations. Annotators write a prompt, or take one from real users, and write the response a model ideally should give. InstructGPT fine-tuned GPT-3 on such labeller-written demonstrations (Ouyang et al., 2022). This is the purest form of SFT data, and the most expensive.
• Reformatted datasets. Existing supervised datasets can be rewritten as instructions. A sentiment example with sentence $s$ becomes “Is the sentiment of the following review positive or negative? $s$” with the answer “positive”. FLAN built its training mixture this way, with several natural-language templates per dataset (Wei et al., 2022a).
• Synthetic data. A model generates instructions and responses, which are filtered and added to the training set. Self-Instruct bootstrapped instruction data from a small set of seed tasks this way (Wang et al., 2023).
• Domain data. Code, mathematics, medicine or a company’s own documents, when the model needs a specific domain.

Modern SFT mixes all of these. What the mixture covers, its quality and the proportions of each kind largely decide the result. Before training, data is usually cleaned: harmful, incorrect, low-quality or simply undesired examples are filtered out, and the mixture is balanced so no single kind dominates.

Chat templates

Conversations are serialised into a single token sequence with special tokens marking the roles. The exact tokens differ between model families; schematically:

The system prompt sets overall behaviour; the user turn is the prompt $x$; the assistant turn is the response $y$. At inference time the model receives everything up to <|assistant|> and generates the rest, stopping when it produces the end token. Using the same template in training and at inference is essential.

Teacher forcing and the loss mask

Training uses teacher forcing: when predicting token $t$, the model is conditioned on the true previous tokens from the data, not on its own earlier predictions. Because every position’s context is known in advance, all positions are trained in parallel in one forward pass, exactly as in pre-training.

Usually the loss is computed only on the response, the highlighted part above. With $z$ the whole tokenised conversation and a mask $m_t=1$ on assistant tokens and $0$ elsewhere,

The prompt still shapes the prediction, because it is in the context, but the model is not trained to produce prompts; it is trained to answer them. This response-only loss is common but not universal: some pipelines train on whole conversations, and how the loss is normalised (per token, per response or per batch) also varies.

6Why a little data goes far

A striking empirical fact is how little SFT data is needed to change a model’s behaviour substantially, and that the change generalises.

• LIMA fine-tuned a 65-billion-parameter LLaMA model on just 1,000 carefully curated prompt–response pairs, with no reinforcement learning, and obtained strong instruction following that generalised to tasks absent from the training data (Zhou et al., 2023).
• FLAN instruction-tuned a 137-billion-parameter model on more than 60 NLP datasets phrased as instructions. Zero-shot performance on unseen task types improved substantially, and FLAN surpassed zero-shot 175B GPT-3 on 20 of the 25 datasets evaluated (Wei et al., 2022a). Its ablations are just as telling: more task clusters helped, the natural-language instructions mattered, and the benefit appeared only at sufficient scale; for the smaller models tested, instruction tuning hurt performance on held-out tasks.

Pre-training has already shown the model vast amounts of text. It is hard to believe that a thousand examples teach it enough new knowledge to explain such a large change. A widely held view, which the LIMA authors call the superficial alignment hypothesis, is that a model’s knowledge and capabilities are learned almost entirely in pre-training, and alignment mainly teaches it which format and style to use when interacting with users (Zhou et al., 2023).

Figure 2 shows the intuition. Pre-training has seen “the capital of France is Paris” countless times, so Paris is already a candidate. But the base model has also seen the question in quizzes, lists and forum threads, where it is followed by another question or a digression. SFT does not need to teach the fact; it needs to move probability from “continue the document” to “answer the question”. (SFT can push towards other goals too, such as continuing a story; answering is the common case.)

There is direct evidence for this picture. Comparing base and aligned versions of the same models token by token, Lin et al. (2024) found that the two agree on the top-ranked token at most positions, and that the shifts concentrate on stylistic tokens such as discourse markers and safety disclaimers. They also showed that a base model prompted with a few well-chosen in-context examples can approach the aligned model’s behaviour.

The view has limits. SFT can teach new facts, but slowly, and there is evidence that it comes at a cost: Gekhman et al. (2024) found that fine-tuning examples introducing knowledge the model did not have are learned more slowly than consistent ones, and that as they are learned the model becomes more prone to hallucinate. So “SFT mainly elicits” is best read as a description of what SFT does well, not a law.

Why parameter-efficient fine-tuning works

The same view explains why tiny updates suffice. LoRA (Hu et al., 2022) freezes the pre-trained weights $W_0$ and learns a low-rank correction,

with $B$ initialised to zero so training starts from the pre-trained model exactly. LoRA was motivated by the hypothesis that the weight change during adaptation has a low “intrinsic rank”, building on earlier work which found that fine-tuning can be reparameterised in a surprisingly low-dimensional subspace, which shrinks as pre-trained models grow (Aghajanyan et al., 2021). If the pre-trained model already sits in a very capable region of parameter space, post-training needs to move it only along a few directions. Learning $B$ and $A$ is, in effect, learning which directions to move in. This is consistent with the elicitation view rather than a proof of it.

7Strengths and limits

SFT’s strengths are easy to list. It is simple and stable: supervised learning with a cross-entropy loss. Its signal is direct: people show the model what they want. It is data-efficient, changing behaviour with far less data than pre-training. It is controllable: style, format and domain can be shaped through demonstrations. And instruction tuning brings some generalisation to unseen tasks.

Its limits are what motivate the rest of post-training.

1. It shows a good answer, but not which answer is better. SFT says “imitate this response”. Many prompts have no single correct reply: two answers can both be acceptable while one is more accurate, more concise, more helpful or safer. In the running example, SFT leaves “Paris.” and the full sentence both likely; nothing in its data says which is preferred. And writing an ideal answer is often harder for a person than comparing two answers. Comparisons may be an easier and more direct signal for improving quality.
2. Its ceiling is the demonstrations. The model imitates its training responses. If they are mediocre, biased or inconsistent, it learns those flaws too, and maximum likelihood’s mass-covering behaviour (§4) means it learns to reproduce the whole spread, not just the best.
3. It never sees its own mistakes. With teacher forcing the model always conditions on ground-truth prefixes, but at inference it conditions on its own generations. Small errors can compound into states the training data never covered. This mismatch is known as exposure bias (Ranzato et al., 2016), and SFT has no mechanism for targeting the errors the model itself tends to make.

The next parts change the training signal from “learn a good answer” to “learn a better answer”, and let the model learn from what it generates itself.

8From demonstrations to comparisons

Once a model has learned to answer, the next goal is for it to answer better. Often there is no unique reference answer to write down, but saying which of two answers is better is easy. That is the essential difference between demonstration data and preference data: one provides a reference answer to imitate, the other a preference relation between answers.

Preference data can come from many places. People can write contrasting pairs to target a property, for example a terse reply against a detailed one when detail is wanted. More commonly the model generates candidates and people judge them. In SFT that would be dangerous, because training on a weak model’s own outputs could reinforce its flaws; for comparisons it is natural, and this division of labour, models generate and people evaluate, is exactly what RLHF exploits (Christiano et al., 2017; Stiennon et al., 2020). It has a further advantage: the feedback is about the model’s own outputs, which are the outputs we care about.

Comparisons are not a free lunch. Some comparisons are genuinely hard, people disagree, and a single “A is better than B” carries less information than a whole high-quality demonstration. Their real advantage is that, for many open-ended generation tasks, comparison is a supervision interface that scales more easily.

The data now looks like this, with $y^w$ the preferred (“winning”) and $y^l$ the rejected (“losing”) response:

In the running example a labeller shown $y_4$ (“The capital of France is Paris.”) and $y_3$ (“Paris.”) might prefer the complete sentence, giving the triple $(x,y_4,y_3)$. How do we turn such orderings into a training signal? One idea is to label the better answer 1 and the worse 0 and train a classifier. But a classifier that can tell which of two answers is better cannot necessarily say how good one answer is, and that is what we need: if the model is going to generate answers, have them scored and improve, which is reinforcement learning, it needs a scoring function for single answers.

9The Bradley–Terry reward model

So we want a function $r(x,y)$ that scores a single response, learned from pairwise orderings. Since “better” has no fixed definition and human judgements are noisy, a probabilistic model of preferences is the natural starting point. Suppose responses $y_1$ and $y_2$ have scores $r_1$ and $r_2$. A softmax over the two scores gives the probability that $y_1$ wins:

where $\sigma$ is the logistic sigmoid. This is the Bradley–Terry model of paired comparisons (Bradley and Terry, 1952): a logistic model in the score difference. Following standard maximum likelihood, we maximise the probability of the observed preferences, that is, minimise the negative log-likelihood:

It looks like classification, but the classifier is constrained to the form $\sigma (\Delta r)$, so what is actually learned is the score $r_{\varphi }$.

Each update raises the winner’s score and lowers the loser’s, weighted by $\sigma (-\Delta )$: pairs the model already orders confidently contribute almost nothing, and mistakes contribute most. Although it is trained on pairs, the result is a stand-alone scoring function. In LLM work, $r_{\varphi }$ is usually itself a Transformer, often initialised from the SFT model with the output layer replaced by a scalar head, which reads the whole prompt and response and outputs one number (Ouyang et al., 2022). InstructGPT collected rankings of several responses per prompt and trained on all pairs from each ranking together.

Translation invariance

Adding a constant to every score leaves every preference probability unchanged, because only differences enter $\sigma$. In fact any shift that depends only on the prompt, $r^{\prime }(x,y)=r(x,y)+c(x)$, is invisible to the data. This translation invariance means the reward’s level is not identified, which is why InstructGPT normalised its reward model so that labeller demonstrations scored zero on average before RL (Ouyang et al., 2022). It matters because preference learning cares only about relative differences, but RL optimisers are sensitive to the magnitude of rewards; the level and scale have to be controlled by other means. The same invariance reappears in §14, where subtracting a baseline is harmless, and in §19, where DPO’s normalising constant cancels.

Reward hacking

The reward model’s biggest risk is already visible. We wanted to optimise a human objective that is hard to formalise, so we trained a model to approximate it. But is a higher score always better? Nothing guarantees that there are no responses the reward model rates highly but people do not like. In the running example the reward model gives the flattering $y_6$ the highest score of all, 2.2. Optimising hard against an imperfect proxy finds exactly such responses: this is reward hacking, an instance of Goodhart’s law. Measured against a held-out “gold” reward, optimising a proxy reward model first improves true quality and then degrades it as the optimisation pushes further (Gao et al., 2023; Stiennon et al., 2020). That is why, in practice, we never simply maximise the learned reward; §16 adds a constraint that keeps the policy close to where the reward model is trustworthy.

10Agents, MDPs and returns

A reward model can say how good an answer is, but it does not make the model produce better answers. For that the model must adjust its generation strategy, its output distribution, in response to the scores. That is reinforcement learning (RL). This part builds the RL needed for RLHF from first principles; Sutton and Barto (2018) is the standard reference for all of it.

Agent and environment

An agent interacts repeatedly with an environment. At time $t$ it observes the state $S_t=s_t$, chooses an action $A_t=a_t$ according to its policy $\pi$, and the environment returns a reward $R_{t+1}=r_{t+1}$ and a new state $S_{t+1}=s_{t+1}$:

We usually assume the environment is Markov:

A common misreading is that the future does not depend on the past. The correct reading is that the current state already contains everything about the past that matters for the future, so earlier states add no information. For a language model the state is naturally

the prompt plus the tokens generated so far. Each state literally contains all earlier ones, so the Markov property holds by construction, and the transition is deterministic: appending the chosen token.

The Markov decision process

Formally the problem is a Markov decision process (MDP) $ℳ=(𝒮,𝒜,P,R,\gamma )$: a state space, an action space, transition probabilities $P(s^{\prime }\mid s,a)$ (the same action in the same state may lead to different outcomes), a reward function $R$, and a discount factor $\gamma$. We want the agent to maximise long-run reward, not just the next reward, so we define the discounted return

The recursion on the right will drive almost everything that follows. With $\gamma$ near 0 the agent is myopic; with $\gamma$ near 1 it weighs distant rewards almost as much as immediate ones and can learn to trade short-term for long-term gain. With $\gamma <1$ and bounded rewards the sum is finite even for endless episodes. Language-model episodes end at an end-of-sequence token, and RLHF commonly uses $\gamma =1$.

Policies, trajectories and the objective

The object RL learns is the policy. A stochastic policy gives a distribution over actions in each state:

A language model is a policy already: it computes $P_{\theta }(y_t\mid x,y_{<t})$, a distribution over the next action given the current state. Running a policy in an environment produces a trajectory $\tau =(s_0,a_0,r_1,s_1,a_1,r_2,\dots )$; for a language model, a trajectory is one whole generated response. Its probability is

with ${\rho }_0$ the distribution of starting states (for us, the distribution of prompts). The goal of RL is to maximise the expected return

In the simplest RLHF setting the reward arrives once, at the end, so $J(\theta )={𝔼}_{x~𝒟}{𝔼}_{y~{\pi }_{\theta }(·\mid x)}[r(x,y)]$. In the running example a trajectory is one of the six replies, and $J$ is the average reward-model score of the replies the policy produces.

Compare this with SFT. SFT tells the model what to generate on a fixed dataset. RL lets the model generate, scores what it generated, and changes the policy so that high-scoring outputs become more likely and low-scoring ones less likely. Given a good enough reward, the model can in principle keep improving from its own attempts. Today’s reasoning methods (§20–21) rest on the same idea; they mainly replace the learned reward model with a verifier.

Exploration is not free, though. The model can only discover good responses that it has some chance of generating, so where it starts matters. That is why RL for language models starts from an SFT model rather than from scratch: in the running example the base model produces a correct answer only 22% of the time, the SFT model 76%.

11Policy gradients

How exactly should the policy change in response to rewards? The intuition: make high-return trajectories more likely and low-return ones less likely. Write $R(\tau )=G_0$ for a trajectory’s return. We want to do gradient ascent, $\theta \leftarrow \theta +\alpha {\nabla }_{\theta }J(\theta )$.

This is the basic policy gradient, and it says exactly what the intuition said: ${\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)$ is the direction that makes action $a_t$ more likely, and it is scaled by how good the whole trajectory was. For a language model it is the gradient of the log-likelihood of the sampled response, the same quantity SFT uses, weighted by the reward. SFT is the special case where every demonstration has weight one.

REINFORCE

REINFORCE (Williams, 1992) estimates the expectation by Monte Carlo: sample $M$ trajectories with the current policy, average, and step:

Reward-to-go

There is an obvious improvement. Every action in a trajectory gets the same weight $R(\tau )$, but they should not. If the return was high because of a good choice at step 2, later, unremarkable choices should not be credited for it, and an action certainly cannot affect rewards that arrived before it. So replace the full return by the return from step $t$ onwards. With $R(\tau )={\sum }_k{\gamma }^k{r}_{k+1}$, dropping the rewards that precede $a_t$ leaves ${\sum }_{k\ge t}{\gamma }^k{r}_{k+1}={\gamma }^t{G}_t$:

This is exact, not an approximation: the dropped terms have expectation zero. In practice the factor ${\gamma }^t$ is usually omitted, and with $\gamma =1$, as is common for language models, it is 1 anyway, giving the familiar ${\sum }_t{G}_t{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)$.

REINFORCE samples trajectories and averages ${\sum }_t{G}_t{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)$. But $G_t$ is a single, noisy sample of the future. To ask how good a state or an action is on average, we need the conditional expectation of $G_t$, which is what value functions are.

12Values and Bellman equations

The state-value function is the expected return from state $s$ when following $\pi$, and the action-value function (Q-function) is the expected return after taking action $a$ in state $s$ and following $\pi$ afterwards:

They are related directly. In state $s$ the next action is drawn from $\pi$, so averaging over it gives

In the running example, with the SFT policy and the reward model, $V(x)$ for the prompt is the SFT-weighted average reward, 0.60; it reappears in §16.

Bellman expectation equations

Using $G_t=R_{t+1}+\gamma G_{t+1}$ and the Markov property, the value after one step is the immediate reward plus the discounted value of wherever we land. With $p(s^{\prime },r\mid s,a)$ the probability of next state $s^{\prime }$ and reward $r$:

These are the Bellman equations: a value is defined in terms of the values one step later. They are linear in $V^{\pi }$, so for a known, finite MDP they can be solved exactly.

Optimality

We want the best policy. Define the optimal value functions and an optimal policy that is greedy with respect to $Q^{\star }$:

Acting optimally at every step, the average over the next action becomes a maximum, giving the Bellman optimality equations:

Value iteration and policy iteration

Read the right-hand side of the first equation as an operator $𝒯$ acting on value functions. For $\gamma <1$ it is a contraction in the maximum norm, $\Vert 𝒯V-𝒯U{\Vert }_{\infty }\le \gamma \Vert V-U{\Vert }_{\infty }$, so by the Banach fixed-point theorem repeatedly applying it, $V_{k+1}=𝒯V_k$, converges to its unique fixed point $V^{\star }$ from any start. This is value iteration.

Policy iteration alternates two steps. Evaluation computes $V^{\pi }$ for the current policy by solving the Bellman expectation equation. Improvement makes the policy greedy: ${\pi }^{\prime }(s)\in {argmax}_a{Q}^{\pi }(s,a)$. The new policy is never worse, and strictly better wherever improvement is possible. The reason is short: $V^{\pi }(s)$ is a weighted average of the $Q^{\pi }(s,a)$, so the maximum is at least the average, $Q^{\pi }(s,{\pi }^{\prime }(s))={max}_a{Q}^{\pi }(s,a)\ge V^{\pi }(s)$ for every state.

Both methods need the transition model $p(s^{\prime },r\mid s,a)$ to compute the expectations. Usually we do not have it; we can only act and observe. That is model-free reinforcement learning.

13Monte Carlo and TD learning

Monte Carlo

The most direct idea returns to the definition $V^{\pi }(s)={𝔼}_{\pi }[G_t\mid S_t=s]$ and estimates the expectation by sampling: run the policy from $s$ many times, record the returns and average them. Incrementally, after each episode, move the estimate towards the observed return:

$Q^{\pi }(s,a)$ can be estimated the same way from the returns that follow each visit to $(s,a)$, and the policy improved greedily as before, so evaluation and improvement work without a model. But Monte Carlo has to wait until an episode ends before it can update, and a whole trajectory’s worth of randomness makes $G_t$ high-variance, increasingly so for long episodes.

Temporal-difference learning

The Bellman equation says we need not wait:

After one step we observe $r$ and $s^{\prime }$, and $r+\gamma V(s^{\prime })$ is already a new estimate of $V(s)$. Temporal-difference (TD) learning (Sutton, 1988) moves the old estimate towards it:

${\delta }_t$ is the TD error: the difference between the prediction made before seeing the step and the better-informed prediction made after it. TD learning nudges the old prediction towards the new one. Using an estimate to update an estimate is called bootstrapping.

SARSA and Q-learning

The same idea applies to Q-functions. For a fixed policy the action-value Bellman equation is $Q^{\pi }(s,a)=𝔼[r+\gamma Q^{\pi }(s^{\prime },a^{\prime })]$ with $a^{\prime }~\pi (·\mid s^{\prime })$. Interact once, observe $(S_t,A_t,R_{t+1},S_{t+1},A_{t+1})$, where $A_{t+1}$ is the action the policy actually takes next, and update

This is SARSA (Rummery and Niranjan, 1994), named after the quintuple it uses. The policy that collects the data is the policy being evaluated and improved, so SARSA is on-policy.

To learn the optimal values directly, apply the same sampling to the Bellman optimality equation $Q^{\star }(s,a)=𝔼[r+\gamma {max}_{a^{\prime }}{Q}^{\star }(s^{\prime },a^{\prime })]$:

This is Q-learning (Watkins and Dayan, 1992): a sample-based, model-free version of value iteration on $Q$. The target uses the greedy action ${max}_{a^{\prime }}$ regardless of what the behaviour policy will actually do next, so the policy being learned (greedy) differs from the one collecting data (often exploratory): Q-learning is off-policy.

The whole of this section answers one question: how to estimate the long-run value of a state or action, and use it to improve the policy. Sampling whole trajectories gives Monte Carlo; sampling one step and bootstrapping gives TD; TD on a fixed policy’s Q-function gives SARSA; TD on the optimality equation gives Q-learning.

14Baselines and advantages

Back to policy gradients. $G_t$ measures an absolute level of return, which mixes three things: how good the state was, how good the action was, and luck. That makes the gradient estimate noisy, and it cannot tell whether a high return was earned by the action or inherited from a good state.

The first fix is to replace the sample $G_t$ by its expectation $Q^{\pi }(s_t,a_t)$. Because ${\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)$ depends only on $(s_t,a_t)$ and $𝔼[G_t\mid s_t,a_t]=Q^{\pi }(s_t,a_t)$, the tower property leaves the gradient unchanged while removing the sampling noise in the future:

This is the form of the policy gradient theorem (Sutton et al., 2000), which also holds when the policy is a general function approximator. It still does not separate state from action. So subtract a baseline $b(s_t)$ that measures how good the state is, and weight by the relative return:

For this to be the same gradient, the baseline term must have expectation zero. It does, for any baseline that depends on the state but not on the action:

The natural choice is the state value, $b(s)=V^{\pi }(s)$, the expected return of the state. The weight becomes the advantage,

which says how much better action $a$ is than the policy’s average in that state; it averages to zero over actions. The expected gradient is the same, but the state’s own quality has been removed and the weights are centred near zero, which reduces variance. ($V^{\pi }$ is not exactly the variance-minimising baseline in general, but it is simple and works well.) It is the RL counterpart of the translation invariance in §9: subtracting anything that depends only on the state changes nothing in expectation, but it changes the noise a great deal.

In practice $Q^{\pi }$ and $V^{\pi }$ are unknown, so the formula is not yet an algorithm. How do we estimate the advantage?

15Actor–critic and GAE

If we do not know the advantage, we can learn it. The policy ${\pi }_{\theta }$ is the actor: it acts. A learned value function is the critic: it evaluates, and its evaluation guides the actor’s updates. This is the actor–critic family; with the advantage as the weight it is often called advantage actor–critic, or A2C, the synchronous variant of A3C (Mnih et al., 2016).

Learning both $Q$ and $V$ to subtract them would be cumbersome. The TD error avoids it. If $V$ is the true value function of the policy, the expected TD error is the advantage:

So one learned critic $V_w\approx V^{\pi }$ is enough: ${\delta }_t$ computed with $V_w$ is a one-step advantage estimate. The critic itself is trained by regression towards its targets, as in TD learning.

One step is short-sighted, though. It uses one real reward and then trusts $V(s_{t+1})$ for the entire future; while the critic is inaccurate, its error goes straight into the advantage. The fix is to use more real rewards before handing over to the critic. The $k$-step advantage uses $k$ rewards and then bootstraps, and it telescopes into a sum of TD errors:

With $k=1$ this is the one-step TD error of actor–critic; as $k$ runs to the end of the episode it becomes the Monte Carlo return minus the baseline. Generalized advantage estimation (GAE) (Schulman et al., 2016) does not choose: it takes an exponentially weighted average of all of them,

In code GAE is a single backward pass, ${\widehat{A}}_t={\delta }_t+\gamma \lambda {\widehat{A}}_{t+1}$, and the critic is regressed onto ${\widehat{A}}_t+V(s_t)$. By discounting distant TD errors GAE offers a continuous bias–variance trade-off between one-step actor–critic and Monte Carlo; it does not remove the trade-off. Intermediate values are typical; §23 lists the values reported in the papers.

Through all of this the skeleton of the policy gradient never changed: ${\nabla }_{\theta }J=𝔼[{\sum }_t{\Psi }_t{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)]$. Only the weight ${\Psi }_t$ evolved: $R(\tau )$, then $G_t$, then $Q$, then $Q-V$, then the TD error, then GAE. Each step keeps the gradient (nearly) the same while making the estimate less noisy and easier to compute online.

16From TRPO to PPO

A vanilla policy-gradient step has two practical problems. First, the gradient is only valid at the current policy: the data were sampled from ${\pi }_{\theta }$, so once $\theta$ moves the samples no longer describe the new policy, and each batch can safely be used for only one small step. Second, the step size is hard to choose: a small change in parameters can be a large change in behaviour, and one bad step can collapse a policy that then generates only bad data from which it cannot recover.

A surrogate objective with importance weights

If the new policy stays close to the old one, we can use the states visited by the old policy and correct the action distribution with an importance ratio. This gives the surrogate objective

which can be estimated from old samples, and whose gradient at $\theta ={\theta }_{\mathrm{o}\mathrm{l}\mathrm{d}}$ equals the policy gradient. It is only trustworthy near ${\theta }_{\mathrm{o}\mathrm{l}\mathrm{d}}$.

TRPO: a trust region

Trust region policy optimization (Schulman et al., 2015) makes “near” precise. Its theory bounds the true improvement from below by the surrogate minus a penalty on the KL divergence between old and new policies; the practical algorithm maximises the surrogate subject to a limit on the average KL:

The constraint is measured in the space of behaviours, not parameters, which is what makes it robust. TRPO solves it with a second-order approximation, conjugate gradients and a line search. It works, but it is complicated and awkward to combine with large shared networks.

PPO: clip the ratio

Proximal policy optimization (Schulman et al., 2017) keeps the idea and drops the machinery. With the per-sample ratio

Read it case by case. If an action was good (${\widehat{A}}_t>0$), the objective rewards raising its probability, but only until the ratio reaches $1+ϵ$; beyond that the clipped term is flat and the sample contributes no gradient. If it was bad (${\widehat{A}}_t<0$), lowering its probability helps only until $1-ϵ$. The minimum makes the objective a pessimistic bound: clipping can remove an incentive to move further, but never hides a penalty. If the ratio has already moved the wrong way, for instance a bad action became more likely, the unclipped term still applies in full.

Because the clip limits how far one batch can move the policy, PPO can safely run several epochs of minibatch gradient steps on each batch of samples, which is much more sample-efficient than vanilla policy gradients. With a shared actor–critic network the full objective adds a value-function loss and, optionally, an entropy bonus $\mathscr{H}$ that encourages exploration:

The PPO paper also describes a variant that replaces clipping with a KL penalty whose weight adapts to hit a target KL; the clipped version performed better in its experiments.

The KL penalty to a reference model

For language models there is a second, different KL. PPO’s clip keeps each update close to the previous policy. RLHF also keeps the whole optimisation close to a fixed reference model ${\pi }_{\mathrm{r}\mathrm{e}\mathrm{f}}$, usually the SFT model (Ziegler et al., 2019; Stiennon et al., 2020; Ouyang et al., 2022):

This is the guard against reward hacking promised in §9: the reward model was trained on outputs of models like ${\pi }_{\mathrm{r}\mathrm{e}\mathrm{f}}$, so it is only reliable near them, and the penalty also keeps the text fluent and diverse. In implementations the penalty is usually folded into a per-token reward, with the reward model’s score added at the last token:

Then PPO runs as usual: a value head estimates $V(s_t)$ at each token, GAE turns the per-token rewards into advantages, and the clipped objective updates the policy.