Softmax: How LLM Scores Become Probabilities

Softmax looks almost too simple to deserve much attention. It is one exponential, one sum and one division. Yet it sits at a crucial boundary inside a language model: before softmax, the model has scores; after softmax, it has a distribution.

I find it useful to treat softmax as a tiny piece of plumbing with a precise contract. It accepts any real numbers, including negative ones. It returns positive values between zero and one. Those values sum to exactly one, while preserving the original ranking.

Raw scores are not probabilities

Imagine that an LLM is choosing among three possible next tokens. Its output layer produces these raw scores, often called logits:

logits = [2.0, 1.0, 0.0]

The first token has the highest score, but the numbers are not probabilities. They do not sum to one, and another model could produce scores such as [12.0, -3.0, 4.5]. Softmax converts this unrestricted score space into a probability distribution.

The whole equation

softmax(x_i) = exp(x_i) / sum(exp(x_j))

The numerator says: exponentiate the score I care about. The denominator says: exponentiate every score and add them together. Dividing by that shared total makes all outputs add up to one.

For [2.0, 1.0, 0.0], the arithmetic is small enough to inspect:

The result is approximately [0.665, 0.245, 0.090]. The first item is still first, the second is still second, and the three values sum to one.

Why the exponential?

Exponentiation does three useful things at once. It makes every value positive. It preserves order: if one score is larger before exponentiation, it remains larger afterwards. It also makes gaps matter more, so the largest score receives a stronger share of the final distribution.

This is why a logit of 2 should not be read as “twice as likely” as a logit of 1. Logits live on a relative score scale. The probability only exists after every competing score has gone through the same exponential-and-normalize calculation.

There is another consequence worth noticing: each probability depends on every score. Raising one logit changes the denominator, so it changes all the probabilities. Softmax models competition, not independent confidence meters.

The same answer, computed safely

Large exponentials can overflow. The standard implementation subtracts the largest score before exponentiating:

def stable_softmax(scores):
    shifted = scores - scores.max(dim=-1, keepdim=True).values
    weights = shifted.exp()
    return weights / weights.sum(dim=-1, keepdim=True)

Our scores become [0.0, -1.0, -2.0]. Their exponentials are approximately [1.000, 0.368, 0.135], which are safer to represent. The probabilities do not change, because adding or subtracting the same constant from every score leaves softmax unchanged.

PyTorch performs this operation directly:

import torch

logits = torch.tensor([2.0, 1.0, 0.0])
probabilities = torch.softmax(logits, dim=-1)

# tensor([0.6652, 0.2447, 0.0900])

The dim=-1 matters. It says which axis contains the competing scores that should form one distribution.

Where softmax appears inside an LLM

At the output, softmax turns one logit per vocabulary token into next-token probabilities. A decoding strategy can then choose the most likely token or sample from that distribution.

Inside self-attention, the same function has a different meaning. Dot products produce compatibility scores between tokens. After any causal mask has replaced future positions with negative infinity, softmax converts each row of scores into attention weights. A masked position receives exactly zero weight because exp(-∞) = 0. Those weights then control how the value vectors are blended.

So the pattern is the same in both places:

scores → softmax → normalized weights

Only the interpretation changes: vocabulary probabilities at the output, attention proportions inside the network.

Temperature changes the sharpness

Temperature applies softmax to logits / T. A temperature below one stretches the gaps and produces a sharper distribution. A temperature above one compresses the gaps and produces a flatter distribution. It does not make the model more knowledgeable; it only changes how strongly the existing score differences influence sampling.

Mathematically, temperature zero would mean division by zero. Software that offers a zero-temperature mode usually treats it as a special instruction to choose the largest logit directly.

The mental model I keep

Softmax is not “the model deciding.” It is the conversion layer between relative scores and usable proportions:

1. Subtract the largest score for numerical safety.
2. Exponentiate, making every value positive.
3. Divide by the total, making the values sum to one.

That is the whole mechanism. Simple arithmetic, but a major change in meaning: arbitrary scores go in; a shared, competitive distribution comes out.