Logits and softmax: how a model turns scores into the next token
For each next-token prediction, the model outputs one raw score, a logit, for every token in its vocabulary. Logits can be any real number, positive or negative, so they can't be compared as probabilities directly. Softmax exponentiates each logit and divides by the sum, giving probabilities between 0 and 1 that add up to 1, and you subtract the largest logit before exponentiating, or large values overflow.
Logits
At every step, an LLM produces a vector of logits, one per token in its vocabulary, which can be over 100,000 tokens. "cat" might get 3.2, "dog" 1.5, "airplane" 0.8. These numbers mean something relative to each other, and nothing on their own: you can't make a decision from "3.2".
Softmax
Softmax puts all the logits on one scale between 0 and 1:
softmax(xᵢ) = exp(xᵢ) / Σ exp(xⱼ)
Exponentiation makes every value positive and exaggerates the gaps, and dividing by the sum makes them add up to 1. Then the model can select the next token from these probabilities: the highest one with greedy decoding, or a sample from the distribution. See sampling.
Numerical stability
exp(1000) overflows to infinity, and infinity divided by infinity is nan.
Subtracting the largest logit first doesn't change the result (it cancels out in the division) and keeps every exponent at or below 0.
Temperature
Temperature divides the logits before softmax: below 1 sharpens the distribution (the top token gets more probability), above 1 flattens it.
import math
def softmax(logits: list[float], temperature: float = 1.0) -> list[float]:
scaled = [x / temperature for x in logits]
m = max(scaled)
exps = [math.exp(x - m) for x in scaled]
total = sum(exps)
return [e / total for e in exps]
tokens = ["cat", "dog", "airplane"]
logits = [3.2, 1.5, 0.8]
for t in (0.5, 1.0, 2.0):
probs = softmax(logits, t)
print(f"temperature {t}: " + ", ".join(f"{w} {p:.3f}" for w, p in zip(tokens, probs)))
print("stable with huge logits:", [round(p, 3) for p in softmax([1000.0, 999.0, 998.0])])
At temperature 0.5, "cat" dominates, and at 2.0 the other tokens gain ground.
Where softmax appears
- At the output, turning logits into next-token probabilities.
- Inside attention, turning similarity scores between tokens into weights.
- In classifiers, turning class scores into probabilities.
Inside attention, softmax over very long sequences is computed in streaming fashion: a running maximum and a running sum, rescaled as new values arrive, so the whole row never has to sit in memory at once. That "online softmax" is a streaming aggregate, the same family as a running mean, and it's one of the ideas behind FlashAttention.
Log probabilities
APIs often return log probabilities instead of probabilities, because multiplying many small probabilities underflows, and adding their logs doesn't. A log probability of -0.1 is a probability of about 0.9, and -5 is about 0.007.
Check your understanding
Explain what temperature does in terms of logits, and why subtracting the maximum doesn't change the result.
From the floor
In the AI data structures lecture, I used the same three tokens: the raw scores 3.2, 1.5 and 0.8 are the values the model actually generates, but there's no decision-making in them until softmax shrinks them onto one scale between 0 and 1, and then the token with the highest probability, "cat", is selected, and I showed that without subtracting the maximum, the exponent overflows.