Softmax: Turning Logits into Probabilities

Softmax is the function that turns a model's raw scores into a probability distribution. It is tiny, it is everywhere, and it has a numerical trap that bites almost everyone who implements it from scratch at least once. This post is the explainer I wish I'd had before I spent an afternoon chasing NaNs in a training loop.

What it does, intuitively

A language model's final layer produces a vector of logits — one score per token in the vocabulary. Those scores are unbounded: they can be -3.2, 11.7, 0.4, anything. We can't sample from unbounded numbers as if they were probabilities. Softmax fixes that:

softmax(z_i) = exp(z_i) / sum_j exp(z_j)

After softmax, every output is positive and they all sum to 1. So softmax(z) is a proper probability distribution over the vocabulary, and the token with the highest logit gets the highest probability. That's the whole job.

A concrete example

Suppose the model is choosing the next word and the logits for three candidate tokens are:

cat:    2.0
dog:    1.0
rocket: -1.0

Step by step:

exp(2.0)  = 7.389
exp(1.0)  = 2.718
exp(-1.0) = 0.368
sum       = 10.475

P(cat)    = 7.389 / 10.475 = 0.705
P(dog)    = 2.718 / 10.475 = 0.260
P(rocket) = 0.368 / 10.475 = 0.035

"cat" had the highest logit, so it ends up with ~70% of the probability mass. Notice the exponential makes the gap nonlinear: a 1-point logit lead becomes a roughly 2.7x probability lead, because exp(1) ≈ 2.718. That's why softmax is "winner-takes-most" — it sharpens differences rather than preserving them.

The numerical trap (and the fix you should always use)

Here's the bug: exp(1000) overflows to infinity, and inf / inf = NaN, which poisons the whole batch. In a real LLM, logits can easily exceed 50 or 100 once the model is confident. The standard fix is the max-subtraction trick:

import numpy as np

def softmax(z):
    z = z - np.max(z)          # shift so the max is 0
    e = np.exp(z)
    return e / np.sum(e)

Subtracting the maximum from every element doesn't change the result — exp(z_i - c) / sum(exp(z_j - c)) simplifies to the original because the exp(-c) cancels top and bottom — but it guarantees the largest exponent is exp(0) = 1, so nothing overflows. Every serious library (PyTorch, JAX, TensorFlow) does this internally; if you ever hand-roll it, do too. I've seen production inference code break on exactly this after a model update made logits larger.

Why temperature matters

You'll hear "temperature" constantly in LLM settings. All it does is scale the logits before softmax:

softmax(z / T)

• T = 1: no change.
• T > 1: flattens the distribution → more random, more diverse, sometimes incoherent.
• T < 1 (e.g. 0.7): sharpens it → the model becomes more confident and repetitive, often more "on topic."

At T → 0 softmax degenerates to argmax (always pick the top token). At T → ∞ it becomes uniform random. So temperature is literally a knob on how much the model explores vs. commits. For code generation I usually want T around 0.2–0.4; for creative writing, 0.8–1.0.

Where softmax actually lives in a Transformer

Softmax appears in two critical places:

1. Attention weights. Inside each attention head, the Q·K scores are divided by sqrt(d_k) and softmaxed so each query distributes its attention as a probability over all keys (each row sums to 1). This is what lets "it" attend to the right noun.
2. The output layer. The final vocabulary logits are softmaxed to get the next-token distribution we sample from.

Both are the same function; only the scale and meaning differ.

A subtlety people miss: softmax isn't symmetric

Because the exponential is sensitive to absolute offsets, adding a constant to all logits does nothing (the max-trick proves it), but adding a constant to only some logits changes the output. This is why calibrating logits matters: two models can "agree" on ranking yet produce very different sampling behavior if one is systematically shifted.

My take

Softmax is the unglamorous hinge of the entire LLM stack. It's the boundary between "the model computed a score" and "the model made a probabilistic choice." Spend an hour really understanding the max-trick and temperature and you'll debug sampling problems an order of magnitude faster than someone who treats it as a black box. And if you're implementing it yourself: always subtract the max first. Always.