Distributions, Softmax, and the Chain Rule of Words

Every time a language model generates a token, it does the same thing. It looks at everything it's seen so far, does an enormous amount of matrix math, and then produces a list of numbers. One number for every token in its vocabulary. That list might have 32,000 entries. It might have 200,000. From that wall of numbers, a single word appears on your screen.

The list of numbers is a probability distribution. If you want to understand how language models work at any depth beyond "it predicts the next word," you have to be comfortable with what that means. Not in a measure-theory sense. Just the practical mechanics: what a distribution is, how you build one from raw model outputs, how temperature reshapes it, and how the chain rule of probability turns next-token prediction into a maximum-likelihood objective for full sequences.

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 honestly the prerequisites here are lighter than those posts. If you know what a vector is, you're good.

What is a probability distribution?

Let's start at the bottom. 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]$. Six numbers. All non-negative. They sum to 1. Done. 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.

This isn't a metaphor. The model literally constructs a categorical distribution at every decoding step. It doesn't "decide" on a token and then report its confidence. It builds a full probability distribution over every possible next token, and then we sample from it. Sampling is a separate step that happens after the distribution is constructed.

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 they have to be converted into a valid distribution before we can do anything with them.

Softmax: turning logits into probabilities

This is where softmax comes in. 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. Valid distribution. Done.

Okay but why $e^x$? Why not just, like, take the absolute value of each logit and divide by the sum? 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.

Second, and this is more subtle, $e^x$ amplifies differences. The gap between $e^3$ and $e^1$ is much larger than the gap between $3$ and $1$. This means softmax is opinionated: it pushes probability mass toward the larger logits and away from the smaller ones. A logit that's a few points higher than its neighbors will grab a disproportionate share of the probability.

Third, exponentials 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."

Finally, softmax is smooth and differentiable everywhere, which makes gradient-based learning well behaved. (The exact softmax Jacobian is $\partial p_i / \partial z_j = p_i(\delta_{ij} - p_j)$, not "the derivative is itself," which is a property of $e^x$ alone.) These reasons keep lining up: positivity, ranking preservation, multiplicative ratios, shift invariance, smoothness. That's a lot of useful properties for one operation.

If you want to see the amplification rather than just take it on faith, here are the same logits run through both normalizers at once. The leader logit on the right breathes between roughly 1.5 and 5. The "abs and divide" panel barely registers; the softmax panel snaps mass onto the winner.

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]$.

Then we sum them: $7.39 + 2.72 + 1.65 + 0.37 + 20.09 = 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

Here's where it gets interesting. What if we want to control how peaky or flat the distribution is? That's what the temperature parameter $T$ does. 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.

The edge cases are intuitive. As $T \to 0$, the distribution collapses to a spike on the highest logit. 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; in practice APIs like the OpenAI and vLLM ones treat temperature=0 as a special case meaning "greedy argmax decoding."

One thing to note: for any positive $T$, the highest-logit token stays the highest-probability token. Temperature reshapes the sharpness of the distribution but never reorders it.

Play with the slider below to see this in action:

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 makes the model more deterministic and repetitive, high temperature makes it more creative and chaotic.

Every LLM API exposes a temperature parameter. The "creativity dial" framing is vague but not wrong. What it actually does is rescale the logits before softmax, and every downstream effect on generation follows from that one operation.

The numerical stability trick

Before we move on, there's a practical detail that matters if you ever implement softmax yourself. The exponential function blows up fast. $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 numerically identical results but the largest exponent is now 0, so nothing overflows. Every production softmax implementation does this.

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."""
    z = np.array(logits) / 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.118  0.016  0.006  0.000  0.860]
 
print(softmax(logits, temperature=2.0))
# [0.253  0.154  0.120  0.056  0.417]

The chain rule and next-token prediction

Modern language models predict tokens, not words. (Tokenization is Arc 3; for now treat each $x_t$ as one position in the output sequence.)

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:

The point of the factorization isn't that generation happens one token at a time. Generation happens that way because of the architecture. What the factorization buys you is that it makes next-token prediction a principled maximum-likelihood objective for full sequences. Maximizing $\log q_\theta(x_1, \ldots, x_n)$ is, by the chain rule, the same as maximizing $\sum_t \log q_\theta(x_t \mid x_{<t})$. So the loss summed over per-position next-token predictions is exactly the joint-sequence log-likelihood. That equivalence is what justifies training a sequence model with a per-position loss in the first place. (The next post derives that loss in detail.)

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. Prompt tokens don't contribute to the loss when you're scoring or training a completion; only the $y_t$ factors do.

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

Let me tie this together with what happens during a single decode step inside a real language model.

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.

Those logits are raw, unconstrained real numbers. Some might be 15.0, others might be -8.0. They carry no probabilistic meaning yet.

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. One thing to flag: 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.

The entire conversation you have with a chatbot is just this loop running thousands of times: logits, softmax, sample, append, repeat.

This is genuinely random. At $T=1$, even if you fed the model the exact same prompt twice, the next token is drawn fresh from the distribution each time. The same set of probabilities can produce different words on different runs. Over many draws, the empirical frequencies should converge to the underlying probabilities; any finite sample will be noisy. That convergence is what makes a probability "real" in any operational sense.

Here's that idea made concrete. A small fixed distribution sits at the top. Below it, the visualization keeps drawing one token at a time. The faint pills mark each token's true probability; the opaque bars are the empirical frequencies from the draws so far.

This is also why evaluation requires many samples. A model that assigns 70% probability to the right answer will still be wrong 30% of the time. If you run it once and get the wrong answer, that tells you almost nothing about how good the distribution is. The distribution is the thing, not any single draw from it.

Misconceptions

"Softmax outputs are confidence scores." They're not, at least not in a calibrated sense. If a model assigns 90% probability to a token, that doesn't mean it's "right" 90% of the time in some frequentist sense. Neural network softmax outputs are famously poorly calibrated. They tend to be overconfident. A model might assign 95% probability to a wrong answer and 2% to the right one. The distribution tells you the model's relative preferences among tokens, but not necessarily how reliable those preferences are.

"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, though most APIs recommend tuning either temperature or top-p rather than both at once. 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 probability distribution over the vocabulary is the most fundamental object in language modeling. Everything we'll do from here, attention, loss functions, sampling strategies, alignment, builds on top of it. A few handles to carry forward:

• A probability distribution is a list of non-negative numbers that sum to 1. Language models produce one at every decode step.
• Softmax converts raw logits into a valid distribution. The max-subtract trick keeps it numerically stable.
• Temperature controls the sharpness of the distribution. Low $T$ concentrates mass, high $T$ spreads it.
• The chain rule factors any joint sequence distribution into a product of next-token conditionals. That's what makes per-position next-token prediction a principled maximum-likelihood objective for full sequences.
• Softmax outputs aren't calibrated confidences. They're relative preferences, not reliable probabilities.

The next post covers cross-entropy and loss, where we'll see how to measure whether a model's probability distributions are any good, and why "negative log probability of the correct token" turns out to be the right thing to minimize.

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. Introduces the online softmax trick as part of a fused attention kernel, the reason long-context attention fits in memory today.
• 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.