Probability, policies & gradients

How models represent probabilities and learn from them

Post-training changes the probabilities of the model’s possible responses. Optimization is the process of adjusting its parameters to improve a chosen numerical objective. To understand those adjustments, it helps to follow how the model produces probabilities, chooses tokens, and receives a score for the resulting response.

The same mathematical tools also apply to pre-training. In post-training, the objective emphasizes the behavior of a model that already has broad language capabilities.

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 distributionprobability distributionAn 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 vocabularyvocabularyThe fixed set of tokens a model knows about. Modern LLMs have ~32k–200k entries.See in glossary →, the model’s fixed list of possible tokens, 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 variablerandom variableA 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 logitslogitsThe 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 vectorvectorAn ordered list of numbers, such as the measurements describing one example or the learned values representing a token.See in glossary →, or ordered list of scores, into a valid distribution, we apply the softmaxsoftmaxFunction that turns any vector into a probability distribution (positive, sums to 1) by exponentiating and normalizing.See in glossary →: raise $e$, a mathematical constant approximately equal to 2.718, to the power of 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 probabilityconditional probabilityThe probability of an outcome given that some other information or event is already known.See in glossary →, 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$. The product gives the probability of an entire response: the model’s score for an entire response is just its per-token conditional probabilities multiplied together. Multiplying many small numbers can underflowunderflowA nonzero numerical result becoming too small in magnitude for the chosen number format to represent normally, potentially rounding to zero.See in glossary →, falling below the computer’s representable range. A logarithmlogarithmThe exponent needed to obtain a number from a chosen base. The natural logarithm uses the base e, approximately 2.718; logarithms turn products into sums.See in glossary → avoids that problem by turning products into sums. The natural logarithm reverses exponentiation: if $e^a = b$, then $\log b = a$.

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 argmaxargmaxThe input or index at which a function or list reaches its largest value.See in glossary → (the single most probable token) every step; it’s deterministicdeterministicProducing the same result from the same inputs and settings, without random choices.See in glossary →, giving the same result for the same inputs and settings. Sampling draws a token in proportion to its probability, so a token the model gives 30% lands about 30% of the time; it’s stochasticstochasticInvolving 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 →, meaning it involves randomness, 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

A model that maps a context to a distribution over next tokens is exactly what reinforcement learning calls a policypolicyIn 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 trajectorytrajectoryThe 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 rolloutrolloutA 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 parametersparametersThe 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. Training generates and scores many such rollouts. 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 rewardrewardA single-number feedback signal in reinforcement learning that says how good an outcome was. The agent tries to collect as much reward as possible over time.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), but in either case the result is 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 methodMonte-Carlo estimateEstimating 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 batchbatchThe group of training examples used for one gradient estimate. Bigger batches reduce gradient noise but use more memory and compute per step.See in glossary →, or group, of rolloutsrolloutA 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-ascentgradient ascentTaking 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—a small parameter update intended to increase the objective, 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, because it affects how reliable each proposed update is.

Parameters and gradient ascent

All of those probabilities are produced by the network’s parametersparametersThe 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 descentgradient descentThe 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 ascentgradient ascentTaking 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).

Three measurements for training

Three useful measurements follow from these probabilities: 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). A training objective can use these measurements to favor desired responses, limit changes from a reference model, or preserve variety.