Distributions, Softmax, and the Chain Rule of Words

At each generation step, an autoregressive language model produces a logit for every token in its vocabulary. The decoder uses these scores to choose the next token.

Applying softmax turns the logits into a probability distribution. Temperature can reshape that distribution, and the chain rule connects these next-token probabilities to the probability of a sequence.

This post is 1.3 in the series. I'll assume you've read 1.1 (Vectors and Spaces) and 1.2 (Norms, Dots, and Similarity), though this post needs less linear algebra. If you know what a vector is, you're good.

What is a probability distribution?

A discrete probability distribution is a list of non-negative numbers that sum to 1. Each number represents the probability of one possible outcome.

Roll a fair die. The distribution is $[1/6, 1/6, 1/6, 1/6, 1/6, 1/6]$. All six entries are non-negative and sum to 1. Roll the die many times and the empirical frequencies should get closer to those probabilities; any finite sample will be noisy.

Now think about a language model. At each step, the model needs to pick the next token from its vocabulary. (Tokens, not words. A token is the model's unit of output, often a word piece or a chunk of bytes. We'll get into tokenization in Arc 3, but for now think of each token as one possible next thing the model can emit.) If the vocabulary has 50,000 tokens, then the distribution is a list of 50,000 numbers, each one representing the probability of that particular token being next. The constraints are the same: every number is $\geq 0$, and they all sum to 1.

That constraint is rigid, but it admits a lot of different shapes. Here's the same six-token mini-vocabulary wearing four different distributions. Watch the bars rearrange themselves and notice the sum in the corner: it's locked at 1.000 through every morph.

This specific setup, where you're picking one option from a set of discrete choices, has a name: the categorical distribution. It's a generalization of the coin flip. A coin flip is a categorical distribution over two outcomes. A die roll is categorical over six. A language model's output is categorical over its entire vocabulary.

The decoder chooses from this distribution by sampling or another rule. The distribution and the selection rule are separate parts of generation.

But there's a problem. Neural networks don't naturally output valid probability distributions. They output arbitrary real numbers, some positive, some negative, some huge, some tiny. These raw outputs are called logits, and softmax converts them into a valid distribution. Greedy decoding can also select the largest logit directly.

Softmax: turning logits into probabilities

Given a vector of logits $z = [z_1, z_2, \ldots, z_n]$, softmax converts them into a probability distribution:

Each output $p_i$ is between 0 and 1, and all the $p_i$ values sum to 1.

Why exponentiate? Taking the absolute value of each logit and dividing by the sum might seem sufficient. That would also give you non-negative numbers that sum to 1.

A few reasons. First, the exponential function maps any real number to a strictly positive number. $e^{-100}$ is tiny but still positive. $e^{100}$ is enormous. It preserves the ordering of the inputs: if $z_i > z_j$, then $e^{z_i} > e^{z_j}$, so the higher logit always gets the higher probability.

Absolute-value normalization loses ordering: logits −5 and +5 receive the same weight, and shifting every logit changes the distribution. Softmax preserves order and is invariant to a shared shift.

Exponentials also turn additive logit differences into multiplicative probability ratios. For any two tokens $i$ and $j$,

A logit gap of $1$ means the higher-logit token is $e \approx 2.72\times$ as likely as its rival. A gap of $2$ is $e^2 \approx 7.39\times$. A gap of $5$ is about $148\times$. This is the right way to read logits: as log-relative preference scores, where only differences matter.

That last sentence has a precise consequence: adding the same constant to every logit doesn't change the distribution at all.

The same factor $e^c$ shows up in every numerator and in the denominator and cancels. This invariance is why the numerical-stability trick below works, and it's why people say "logits are only defined up to an additive constant."

Softmax is also smooth and differentiable, so gradients can pass through it during training. Its Jacobian is $\partial p_i / \partial z_j = p_i(\delta_{ij} - p_j)$.

The comparison below applies both normalizers to the same changing logits.

Let's run through a concrete example. Say we have five logits coming out of a language model: $[2.0, 1.0, 0.5, -1.0, 3.0]$, corresponding to tokens ["the", "cat", "dog", "ran", "sat"].

First we exponentiate each: $[e^2, e^1, e^{0.5}, e^{-1}, e^3] = [7.39, 2.72, 1.65, 0.37, 20.09]$.

Using the unrounded exponentials, their sum is approximately $32.21$.

Then we divide each by the sum: $[0.229, 0.084, 0.051, 0.011, 0.624]$.

So "sat" (logit 3.0) gets 62.4% of the probability mass. "the" (logit 2.0) gets 22.9%. And "ran" (logit -1.0) is down at 1.1%. The distribution is valid, it sums to 1, and the ordering from the logits is preserved.

Temperature: sharpening and flattening

The temperature parameter $T$ controls how concentrated the distribution is. We divide the logits by $T$ before applying softmax:

When $T = 1$, you get the standard softmax. When $T < 1$, you're dividing by a number less than 1, which increases the magnitude of the logits. This amplifies the differences between them, making the distribution sharper. The winning logit grabs even more probability mass. When $T > 1$, you're shrinking the logits toward zero, squashing the differences, and the distribution flattens out toward uniform.

As $T \to 0^+$, the distribution concentrates on the highest logit (or shares mass equally among tied maxima). As $T \to \infty$, every token gets equal probability (uniform distribution). Temperature is a knob that smoothly interpolates between "always pick the most likely token" and "pick any token at random." Strictly speaking, $T = 0$ is a divide-by-zero, so the formula doesn't apply; a decoder can instead define temperature=0 as a separate greedy-argmax mode, with its own tie-breaking rule.

For any positive $T$, the highest-logit token stays the highest-probability token. Temperature reshapes the sharpness of the distribution but never reorders it.

Notice how at $T = 0.5$, "sat" dominates with over 85% of the probability. At $T = 2.0$, the distribution is much more spread out, and even "ran" has a meaningful probability. This is why temperature matters for generation: low temperature concentrates sampling on the highest-scoring tokens, while high temperature gives lower-scoring tokens more chance.

The numerical stability trick

Large positive logits can overflow when exponentiated. $e^{100}$ is already about $2.69 \times 10^{43}$, and if your logits are large, you'll overflow to infinity.

The fix is simple: subtract the maximum logit from all logits before exponentiating. Since:

(the $e^{\max(z)}$ cancels out of numerator and denominator), this gives the same result in exact arithmetic but the largest exponent is now 0, so nothing overflows. Very small probabilities can still underflow, which is one reason to compute log probabilities directly.

In real training and inference code, you also see frameworks skip the explicit softmax altogether and work in log space. PyTorch's log_softmax and cross_entropy both fuse the max-subtract trick with a $\log\sum e^{z_j}$ computed in a numerically safe way, which avoids ever materializing tiny probabilities or huge exponentials. CrossEntropyLoss in particular expects raw logits and is equivalent to LogSoftmax followed by negative log-likelihood; you almost never apply softmax yourself before the loss.

import numpy as np
 
def softmax(logits, temperature=1.0):
    """Numerically stable softmax with temperature."""
    if temperature <= 0:
        raise ValueError("temperature must be positive")
    z = np.array(logits, dtype=float) / temperature
    z = z - np.max(z)         # stability trick: shift so max is 0
    exp_z = np.exp(z)
    return exp_z / exp_z.sum()
 
# Try it
logits = [2.0, 1.0, 0.5, -1.0, 3.0]
print(softmax(logits, temperature=1.0))
# [0.229  0.084  0.051  0.011  0.624]
 
print(softmax(logits, temperature=0.5))
# [0.117  0.016  0.006  0.000  0.862]
 
print(softmax(logits, temperature=2.0))
# [0.253  0.154  0.120  0.056  0.417]

The chain rule and next-token prediction

The chain rule of probability says any joint distribution over a sequence factorizes into a product of conditionals:

where $x_{<t}$ is shorthand for $(x_1, \ldots, x_{t-1})$. This is an exact identity. It's a theorem of probability and it holds for any joint distribution over an ordered sequence, language or otherwise.

The model doesn't know the true conditionals $P(x_t \mid x_{<t})$. It learns approximations. I'll write the learned distributions as $q_\theta(x_t \mid x_{<t})$, where $\theta$ is the model's parameters. The model's probability for the whole sequence is the product of its own learned conditionals:

By the chain rule, maximizing $\log q_\theta(x_1, \ldots, x_n)$ is the same as maximizing $\sum_t \log q_\theta(x_t \mid x_{<t})$. Training with the sum of per-position next-token log losses therefore maximizes joint-sequence likelihood. The next post derives that loss.

If you want a probability for a complete finite sequence rather than an open-ended stream, you have to mark where it ends. The standard fix is an end-of-sequence token $\langle\text{eos}\rangle$ in the vocabulary, so the final factor $q_\theta(\langle\text{eos}\rangle \mid x_{\le n})$ closes the sequence. (Equivalently, you can condition on the sequence length $n$ explicitly.) Without one of these, the probabilities don't sum to 1 over the space of all sequences of all lengths.

For prompt-conditioned generation, distinguish modeling the full joint $q_\theta(x_1, \ldots, x_n)$ from modeling a completion conditioned on a prompt. If the first $k$ tokens are a prompt $c$ and the rest are the generated continuation $y$, the quantity you actually care about at inference and in completion-style training is

Same factorization, different domain. For this conditional completion objective, only the $y_t$ factors contribute to the loss. Other training recipes also score prompt tokens.

Worked example. Suppose the model assigns the following conditionals to a five-token sequence (these are the model's $q_\theta$ values, not ground truth):

The joint probability of the sequence under the model is the product, $0.05 \times 0.02 \times 0.08 \times 0.20 \times 0.45 \approx 7.2 \times 10^{-6}$. Sequence joint probabilities are usually small for the same reason: many factors less than 1 multiplied together. In practice we work in log space, where the product becomes a sum:

Per-position log probabilities just add. That additivity is what cross-entropy loss accumulates across positions, and it's why "average loss per token" is a meaningful quantity to report and compare across sequences of different lengths.

import numpy as np
 
def sequence_log_prob(token_probs):
    """
    Given a list of P(w_t | w_1, ..., w_{t-1}) values,
    compute the log probability of the full sequence.
    
    The chain rule says:
    P(sequence) = prod(token_probs)
    log P(sequence) = sum(log(token_probs))
    """
    return sum(np.log(p) for p in token_probs)
 
# "the cat sat on the" example
token_probs = [0.05, 0.02, 0.08, 0.20, 0.45]
log_prob = sequence_log_prob(token_probs)
print(f"log P(sequence) = {log_prob:.2f}")    # -11.84
print(f"P(sequence) = {np.exp(log_prob):.6e}")  # 7.20e-06
print(f"perplexity = {np.exp(-log_prob / len(token_probs)):.1f}")  # 10.7

What actually happens at inference

The model receives a sequence of token IDs. It passes them through an embedding layer, then through many layers of attention and feed-forward networks. At the end, the final hidden state for the last token position gets projected through a linear layer into a vector of size $|V|$, where $V$ is the vocabulary. This vector is the logits.

Softmax shows up twice in transformer inference: there's a softmax inside every attention layer (turning query–key dot products into attention weights over values) and a separate softmax at the very end (turning final logits into a distribution over the vocabulary). This post is about the second one. The attention softmax has its own post in Arc 4.

Then we apply softmax (possibly with a temperature parameter), and now we have a valid probability distribution over the entire vocabulary. We can sample from it (randomly draw a token according to the probabilities), or we can take the argmax (greedy decoding), or we can use more sophisticated strategies like top-k or nucleus (top-p) sampling, which we'll cover in Arc 5.

The sampled token gets appended to the input, and the whole process repeats for the next position. That's one step of the chain rule being evaluated. In efficient implementations, the model doesn't recompute the entire prompt from scratch on every step. It caches the attention keys and values from previous tokens (the KV cache) and only runs the new token through the network. The chain-rule view stays exact; the runtime view is much faster. The KV cache gets its own post in Arc 5.

When sampling with fresh random draws, the same distribution can produce different tokens. For repeated independent draws from a fixed distribution, empirical frequencies converge to those probabilities; a finite sample is noisy. A fixed random seed can reproduce the same draws.

The visualization below draws from a fixed distribution. Faint bars mark the target probabilities; opaque bars show empirical frequencies as draws accumulate.

Misconceptions

"A token probability is the probability that an answer is correct." A probability of 0.9 describes the model’s next-token distribution. Factual correctness of a whole answer is a different event and needs its own evaluation. Softmax alone does not guarantee that predictions are calibrated against observed outcomes.

"You always pick the highest probability token." Greedy decoding (always taking the argmax) is just one of many strategies, and it isn't always the right one. For open-ended generation it tends to produce repetitive text, so chat and creative-writing systems usually sample from a modified distribution: temperature reshapes the whole distribution, top-k restricts sampling to the $k$ most likely tokens, and top-p (nucleus sampling) keeps the smallest set of tokens whose cumulative probability exceeds $p$. These are often combined, and changing both can make their effects harder to isolate. For coding, extraction, classification, structured output, or anything that wants a single canonical answer, low-temperature or constrained or greedy decoding is often preferred. The point is that the model produces a full distribution; what to do with it is a separate design choice that depends on the task.

What's next

The next post derives cross-entropy and loss from the probabilities a model assigns to observed tokens.

Additional reading (and watching)

• Bridle, J. S. (1990). Training Stochastic Model Recognition Algorithms as Networks Can Lead to Maximum Mutual Information Estimation of Parameters. NIPS 1989. The softmax function as used in neural networks originates here.
• Goodfellow, I., Bengio, Y., & Courville, A. (2016). Deep Learning, Ch. 3 (Probability and Information Theory). MIT Press. See also Bishop, C. M. (2006). Pattern Recognition and Machine Learning, Ch. 2. Springer. Canonical treatments of discrete distributions, softmax, and numerical stability.
• Jurafsky, D. & Martin, J. H. (2024). Speech and Language Processing, 3rd ed., Ch. 3. The standard reference for the chain rule factorization of sequence probability and n-gram language models.
• Guo, C., et al. (2017). On Calibration of Modern Neural Networks. ICML 2017. Demonstrates that modern neural networks are poorly calibrated and that simple temperature scaling fixes most of the gap without changing predictions.
• Dao, T., et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. Uses online softmax in an IO-aware tiled attention algorithm, avoiding storage of the full attention matrix.
• Holtzman, A., et al. (2020). The Curious Case of Neural Text Degeneration. ICLR 2020. Introduces nucleus (top-p) sampling and motivates why greedy and pure-temperature sampling both fail in different ways.