Probability, policies & gradients

The vocabulary every later method reuses

Post-training is, underneath the slogans, two things: probability and optimization. Every method in this explainer — SFT, RLHF, direct-preference methods, reasoning RL — is a way of reshaping a probability distribution by following a gradient. Before we meet the specific mathematical measures the next chapter is about, it pays to assemble the primitives they’re made of: what a distribution over tokens is, how a network produces one, what it means to sample from it, why we call the model a policy, how we summarize its behavior with an expectation, and how gradient ascent nudges all of it. None of this is new mathematics — but having it explicit here means the later objectives read as rearrangements of pieces you already own.

This is a focused refresher aimed at post-training. The pre-training explainer covers some of the same ground (softmax, sampling, gradient descent) from the angle of building a model; here we line the same tools up for the angle of steering one.

A probability distribution over tokens

Say the model has read some text and is about to predict what comes next. Its job at that point is to produce a probability distribution probability distribution An assignment of a non-negative probability to every possible outcome that sums to one. A language model produces one over the vocabulary at each step — its bet on the next token. See in glossary → over the next token: a number for every entry in the vocabulary, saying how likely that token is to come next. Two rules make it a distribution: every number is non-negative, and they sum to one. Nothing is left over; the model is always making a complete bet over all $|V|$ possibilities.

It helps to think of the next token as a random variable random variable A quantity whose value is an outcome of a random process — e.g. the next token, which isn't fixed until it's drawn from the model's distribution. See in glossary → $Y$ — an outcome that isn’t fixed until it’s drawn — whose distribution the model specifies. Write $\pi(y \mid x)$ for the probability the model assigns to token $y$ given context $x$ — read aloud as “pi of $y$ given $x$,” where the bar $\mid$ is spoken as “given.” The whole of generation is: produce this distribution, commit to one token, append it, and repeat.

Logits and softmax

The network doesn’t emit probabilities directly. Its final layer produces one raw score per vocabulary token — the logits logits The raw, pre-softmax scores the model produces — one per vocabulary token, per position. Bigger logit = the model finds that token more likely; the actual value can be any real number, positive or negative. Applying softmax across the vocabulary turns logits into a probability distribution that sums to 1. Sampling then picks one token from that distribution. See in glossary → — which can be any real number, positive or negative. To turn that vector of scores into a valid distribution, we apply the softmax softmax Function that turns any vector into a probability distribution (positive, sums to 1) by exponentiating and normalizing. See in glossary → : exponentiate every logit (making it positive), then divide by the total (making them sum to one).

Exponentiating means a small lead in logits becomes a larger lead in probability — softmax amplifies differences. Dividing a logit by a temperature before the exponential lets us dial that sharpening up or down: low temperature concentrates mass on the top token, high temperature flattens toward uniform. The widget below lets you push logits around and watch the distribution respond.

Conditional probability and the chain rule

The bar $\mid$ in $\pi(y \mid x)$ means given: it’s a conditional probability, the probability of $y$ conditioned on having already seen $x$. A full response isn’t one token, though — it’s a sequence $y = (y_1, y_2, \ldots, y_T)$. The probability of the whole sequence is built by the chain rule: predict each token given everything before it, and multiply.

Here $y_{<t}$ is shorthand for all tokens before position $t$. This single product is the spine of the next chapter: the model’s score for an entire response is just its per-token conditional probabilities multiplied together. (And because multiplying many small numbers underflows, we’ll soon work with the log of this product — a sum of log-probabilities — but that’s a convenience, not a new idea.)

Sampling vs. greedy

A distribution is a menu, not a choice. To actually generate text we must pick a token, and there are two basic ways. Greedy decoding takes the argmax greedy decoding (argmax) Always picking the single most probable next token (the argmax of the distribution). Deterministic — it explores nothing, which is why RL relies on sampling instead. See in glossary → — the single most probable token — every step; it’s deterministic. Sampling draws a token in proportion to its probability, so a token the model gives 30% lands about 30% of the time; it’s stochastic stochastic Involving randomness, so the outcome varies from run to run rather than being fixed. Sampling is stochastic — the same prompt can yield different responses; greedy decoding is deterministic. See in glossary → , and it’s what gives generation variety.

The distinction matters enormously in post-training. Reinforcement learning improves by trying things — it can only discover that a response is good if the model ever produces it. A model that always decodes greedily explores nothing; the randomness of sampling is the exploration that RL feeds on. Temperature is the knob between the two extremes — toward 0 it approaches greedy, higher it widens the exploration.

A model as a policy

We’ve quietly earned a powerful reframing. A model that maps a context to a distribution over next tokens is exactly what reinforcement learning calls a policy policy In RL, the thing that chooses actions — here, the language model itself, viewed as a distribution over next tokens given the context. RL post-training optimizes the policy. See in glossary → : a rule that, given a state (the tokens so far), assigns probabilities to actions (the next token). Generating a full response is rolling that policy forward step by step — sample a token, append it, sample the next — to produce a trajectory trajectory The sequence of states and actions in a rollout. For text generation, the tokens generated one after another, each conditioned on those before it. See in glossary → .

That single run of the policy from prompt to finished response is a rollout rollout A complete generated sample from the policy — for an LLM, one full response to a prompt. RL collects rollouts, scores them, and updates the policy. See in glossary → : one complete response sampled from $\pi_\theta$ — our notation for the model written as a policy, the subscript $\theta$ standing for its parameters parameters The numbers (weights) inside a model that get adjusted during training. A “7B model” has 7 billion of them. See in glossary → (the weights that training adjusts). It’s the unit reinforcement learning works in — in a moment we’ll see that training means generating many rollouts and scoring each. Sampling (not greedy decoding) is what makes every rollout a fresh draw, which is exactly why a collapsed, near-deterministic policy explores so poorly.

Writing the model this way — as a policy $\pi_\theta$ — is what lets us talk about improving it rather than just running it. For now, the one-line version: the language model is the policy, the next token is its action, and post-training is about changing the policy’s odds.

Expectation and Monte-Carlo

Take a single finished response and we can score it: run it past some set of criteria and read off a number — its reward reward A scalar signal saying how good an outcome was. In post-training it can come from a learned reward model, a verifier, or a rule, and is what RL maximizes. See in glossary → — saying how good or bad that response was. What those criteria actually are is a whole topic of its own (a learned reward model trained on human preferences, or a verifier that checks whether an answer is correct), and later chapters are devoted to it. For now, take the reward as given: a function $r(y)$ that scores a response $y$.

Knowing one response’s score isn’t quite what we’re after, though. We want to know whether the policy itself is any good — and since a policy is a distribution over responses, its quality is the average reward across everything it might produce. We can’t get that average by brute force: for any real prompt the number of possible token sequences is astronomical, so we can never generate every response and tabulate every reward.

That average is an expectation — each possible response weighted by how likely the policy is to produce it:

We can’t evaluate that sum, but we can estimate it with the Monte-Carlo method Monte-Carlo estimate Estimating an expectation you can't compute exactly by drawing samples and averaging. Accuracy improves with the number of samples; it underlies every policy-gradient estimator. See in glossary → : draw a sample of responses from the policy, average their rewards, and use that average as a stand-in for the true expectation. The more responses we sample, the closer the estimate lands — as the widget shows.

This is the loop at the heart of RL post-training. To judge how good the current policy is — and which way to nudge it — we let it generate a batch of rollouts rollout A complete generated sample from the policy — for an LLM, one full response to a prompt. RL collects rollouts, scores them, and updates the policy. See in glossary → (the generation step — ordinary model inference run on the full model, and usually the most compute-hungry part of the loop), score each one’s reward, and average them into a Monte-Carlo estimate of the policy’s expected reward. Then we use those same scored rollouts to estimate which direction in weight-space would raise that expected reward, take a small gradient gradient ascent Taking small steps along the gradient to maximize an objective — the same operation as gradient descent with the sign flipped. "Maximize the log-likelihood" and "minimize the negative log-likelihood" describe one update. See in glossary → step on $\theta$ in that direction, and repeat — sample rollouts, score, estimate, update the weights, again and again — each round shifting probability toward responses that score well. Because the estimate comes from a finite sample, how many rollouts we draw and how noisy they are matters a great deal, which is why later chapters spend whole sections on sample count, variance, and baselines.

Parameters and gradient ascent

All of those probabilities are produced by the network’s parameters parameters The numbers (weights) inside a model that get adjusted during training. A “7B model” has 7 billion of them. See in glossary → $\theta$ — its billions of weights. Change $\theta$ and you change every logit, hence every probability, hence the whole policy. Training is the search for a $\theta$ that makes some objective large (or some loss small).

The tool for that search is the gradient: the direction in parameter space that most increases the objective. Gradient descent gradient descent The core training algorithm: repeatedly nudge each parameter a small step in the direction that lowers the loss, as told by the gradient. See in glossary → takes small steps against the gradient to minimize a loss; gradient ascent gradient ascent Taking small steps along the gradient to maximize an objective — the same operation as gradient descent with the sign flipped. "Maximize the log-likelihood" and "minimize the negative log-likelihood" describe one update. See in glossary → takes them along the gradient to maximize an objective. They’re the same operation with a sign flip, and post-training uses both framings interchangeably — “minimize the negative log-likelihood” and “maximize the log-likelihood” describe one update. The widget below shows the mechanic: follow the slope downhill (or up) in small steps until you reach a minimum (or maximum).

Bridge

These are the primitives. The next chapter assembles three specific quantities out of them — the likelihood of a response (a product of conditional probabilities), the KL divergence between two policies (how far one distribution sits from another), and the entropy of a distribution (how spread out it is). Each is just a particular thing to measure about the distributions you now know how to produce, sample, and score — and each becomes a target we move $\theta$ toward or away from.