Softmax in LLMs, explained in 5 lines

Module 06 · LLMs · 22 concepts

In depth / 15 lines

Softmax, in more detail

Softmax turns a row of raw scores into a set of probabilities.

It exponentiates every score and then divides each by the total.

Exponentiating makes everything positive, and dividing makes the results sum to one.

The exponential is what gives softmax its character.

A score two points higher becomes about seven times more likely, not two points more likely.

So small differences in logits become large differences in probability.

The top choice usually takes a large share of the mass, which is why models sound decisive.

Only the gaps between scores matter, since shifting them all equally cancels out.

In practice the largest value is subtracted first to stop the exponential overflowing.

Temperature divides the scores before the exponential and reshapes the outcome.

A temperature below one widens the gaps and concentrates probability on the leader.

A temperature above one narrows them and gives unlikely tokens a real chance.

The same function is used inside attention to turn raw scores into weights.

It also ends every multi-class classifier, which is where most people meet it first.

It cannot be told apart from confidence, though it is only a normalised score.

Want the full technical article? Wikipedia ↗

Deep dive

What is underneath Softmax

Turning scores into probabilities

A model's last layer produces one unbounded real number per class, and not one of them is a probability.

\sigma(z)_i = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}}

$z_i$ : the logit for class $i$ , any real number at all.
$K$ : how many classes there are.
$\sigma(z)_i$ : the probability this class ends up with.

Exponentiating makes every value positive, and dividing by the total makes them sum to one.

That is the entire definition: positive, then normalised.

Why exponentials rather than just dividing

Dividing raw scores by their sum breaks the moment one score is negative, because a negative probability is not a probability.

The exponential is also monotonic and smooth, so the largest logit always takes the largest share and the gradient exists everywhere.

It has one more consequence, and it is the useful one.

\frac{\sigma(z)_i}{\sigma(z)_j} = e^{z_i - z_j}

Only differences between logits matter. Adding the same constant to all of them changes nothing at all.

The shift that stops it overflowing

That last property is not a curiosity, it is what makes softmax computable at all.

\sigma(z)_i = \frac{e^{z_i - m}}{\sum_{k=1}^{K} e^{z_k - m}} \qquad m = \max_j z_j

Subtracting the largest logit leaves the answer identical and forces the biggest exponent to be exactly zero.

Without it a logit of 1000 overflows to infinity in float32, the denominator becomes infinity, and the whole row comes back as NaN.

Every library does this inside its own implementation, which is why you rarely see it written and why it bites the first time you write softmax yourself.

One pass, in order

Take the vector of logits for one position.
Find its maximum and subtract that from every element.
Exponentiate each element.
Add them up.
Divide each exponential by that sum.

The derivative, which is why it pairs with cross-entropy

\frac{\partial \sigma_i}{\partial z_j} = \sigma_i \left( \delta_{ij} - \sigma_j \right)

$\delta_{ij}$ : one when $i$ and $j$ are the same index, zero otherwise.

Every output depends on every logit, so this is a full matrix and not a single number.

Fed straight into cross-entropy it collapses to something far simpler, and that collapse is why the two are almost always implemented as one operation.

Test yourself

One question on this card. Every option explains itself once you pick, right or wrong.

Question 1

What does softmax guarantee about its output?

LLMs

22 concepts in this module, then a graded test with every answer explained.

Study this module →