LLM Primer

Stable Softmax

Softmax turns logits into probabilities.

For logits z_i:

softmax(zi)=ezi∑jezjsoftmax(z_i) = \frac{e^{z_i}}{\sum_j e^{z_j}}

The direct formula can overflow when logits are large.

logitssubtract maxexpnormalizesame probabilities, safer numbers
Stable softmax shifts logits before exponentiation.

Softmax does not care about the absolute position of all logits. It cares about their differences.

For example, [10, 12, 8] and [-2, 0, -4] have the same gaps between entries. The second version is safer because its largest value is 0, so the largest exponential is e^0 = 1.

The Stable Trick

Subtract the maximum logit before exponentiating:

softmax(zi)=ezi−m∑jezj−msoftmax(z_i) = \frac{e^{z_i - m}}{\sum_j e^{z_j - m}}

where:

m=max⁡jzjm = \max_j z_j

This gives the same probabilities, but safer numbers.

Why does this not change the answer? The same positive factor is multiplied into every numerator and the denominator, so it cancels.

For a constant m:

ezi−m∑jezj−m=ezi/em∑jezj/em=ezi∑jezj\frac{e^{z_i - m}}{\sum_j e^{z_j - m}} = \frac{e^{z_i}/e^m}{\sum_j e^{z_j}/e^m} = \frac{e^{z_i}}{\sum_j e^{z_j}}

What To Watch

Subtracting the maximum is not a model trick. It does not make the model more confident or less confident. It changes the computation path so the same probabilities can be computed without unsafe intermediate exponentials.

This is also why the predicted class does not change. If one logit was largest before the shift, it is still largest after subtracting the same constant from every logit.

MATH-C08-T06-001Exercise: Shift logits

Logits are [10, 12, 8]. What value should be subtracted for stable softmax?

Compute it first, then check your number.

Hint
Subtract the maximum logit.
Solution

The maximum is 12, so subtract 12. This makes the largest shifted logit equal to zero, which keeps the largest exponential at e^0 = 1.

MATH-C08-T06-002Exercise: Shifted logits

After subtracting the maximum from [10, 12, 8], what is the largest shifted logit?

Compute it first, then check your number.

Hint

The maximum entry is subtracted from itself.

Solution

The largest shifted logit is 12 - 12 = 0. Every other shifted logit is less than or equal to zero, so no shifted exponential is larger than 1.

MATH-C08-T06-003Exercise: Same probabilities

Does subtracting the same constant from every logit change the softmax probabilities?

Answer it first, then check.

Hint

The same factor appears in every numerator and in the denominator.

Solution

No. The common factor cancels, so the probabilities stay the same. The shift changes the intermediate numbers, not the probability distribution being computed.

MATH-C08-T06-004Exercise: Why the max

Why subtract the maximum instead of a random logit?

Enter 1 for because it makes the largest shifted value 0, or 2 for because it changes the predicted class.

Compute it first, then check your number.

Hint

The goal is numerical safety, not changing the model decision.

Solution

Subtracting the maximum makes the largest shifted value 0, so the largest exponential is 1. Enter 1.

MATH-C08-T06-005Exercise: Predicted class after shift

Enter 1 if subtracting the same constant from every logit keeps the ordering of the logits the same.

Compute it first, then check your number.

Hint

Compare a > b with a - m > b - m.

Solution

Enter 1. Subtracting the same constant from every logit preserves their differences, so the largest logit remains largest.

Before Moving On

Stable softmax changes the computation, not the probabilities.