Probability in LLMs: Next-Token Prediction, Softmax and Perplexity

Everything in the previous three subtopics — distributions, Bayes' theorem's language of belief-updating, cross-entropy — converges into one concrete object: at every single step, an LLM is computing a full probability distribution over its entire vocabulary (every possible next token) and then picking (or sampling) from it.

How does raw model output become a valid probability distribution? The last layer of a language model produces one raw score (called a 'logit') per vocabulary token — an arbitrary real number, could be negative, could be huge, and these numbers don't sum to 1 or stay in [0, 1]. The softmax function fixes this: it exponentiates every logit (making them all positive) and divides each by the sum of all the exponentiated logits (making them sum to exactly 1). The result is a genuine probability distribution — larger logits get proportionally larger probabilities, and everything is non-negative and sums to 1, exactly the two properties a valid distribution needs (from the first subtopic in this unit).

Generating text is then: run softmax to get a distribution over the next token, pick a token (greedily take the highest probability, or sample proportionally to add variety — the temperature parameter controls how 'sharp' vs 'flat' that sampling distribution is), append it to the sequence, and repeat. The whole probability of a generated SENTENCE is the product of each token's conditional probability given everything before it — exactly the joint-probability chain rule from the first subtopic.

Perplexity is how you turn a model's average cross-entropy loss (previous subtopic) into one intuitive number: perplexity = 2^(cross-entropy loss in bits). It's often read as 'the model is, on average, as confused as if it were choosing uniformly among this many options.' A perplexity of 1 means perfect, certain predictions; a perplexity equal to the vocabulary size means the model is no better than guessing uniformly at random. Lower perplexity means the model assigns higher probability to the actual text it's being tested on — it's the single most common intrinsic metric for comparing how well different language models fit real text, and it's directly derived from the cross-entropy this whole unit has been building toward.

import math

def softmax(logits: list[float]) -> list[float]:
    """Turn raw scores into a valid probability distribution:
    all non-negative, sums to 1."""
    max_logit = max(logits)                       # for numerical stability
    exps = [math.exp(x - max_logit) for x in logits]
    total = sum(exps)
    return [e / total for e in exps]

# Raw model scores for 4 candidate next-tokens: 'cat', 'dog', 'car', 'the'
vocab = ["cat", "dog", "car", "the"]
logits = [4.2, 3.9, 1.1, 0.5]
probs = softmax(logits)
for token, p in zip(vocab, probs):
    print(f"P({token!r}) = {p:.3f}")
print(f"sums to: {sum(probs):.3f}")  # always 1.0

def perplexity_from_cross_entropy(bits_per_token: float) -> float:
    """perplexity = 2^(cross-entropy loss in bits) --
    'the model is as confused as choosing among this many options.'"""
    return 2 ** bits_per_token

print(f"Perplexity at 0.1 bits/token loss (confident model): "
      f"{perplexity_from_cross_entropy(0.1):.2f}")
print(f"Perplexity at 3.0 bits/token loss (confused model):  "
      f"{perplexity_from_cross_entropy(3.0):.2f}")