The Softmax Function, Line by Line

Softmax turns any list of numbers into a probability distribution: exponentiate each one, divide by the total. The version every real implementation uses subtracts the largest value first, which changes nothing about the answer and prevents a crash that would otherwise happen at a logit of about 89 in fp32 and about 12 in fp16.

The formula, and what each part does

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

Two jobs, one line. The exp makes every number positive and amplifies differences — a gap of 1 in the inputs becomes a factor of e = 2.718 in the outputs. The division makes the results sum to 1, which is what turns them into probabilities.

The inputs are called logits. They can be any real numbers, positive or negative, of any magnitude. The outputs are strictly between 0 and 1 and always sum to exactly 1.

Worked on three numbers

Take three logits: z = [2.0, 1.0, 0.1]. Exponentiate each:

exp(2.0) = 7.389056
exp(1.0) = 2.718282
exp(0.1) = 1.105171
           ---------
sum      = 11.212509

p_1 = 7.389056 / 11.212509 = 0.659001
p_2 = 2.718282 / 11.212509 = 0.242433
p_3 = 1.105171 / 11.212509 = 0.098566
                             --------
                             1.000000

Look at what the exponential did. The first logit is twice the second, but its probability is 2.7 times as large. The third logit is only 1.9 below the first, and its probability is 6.7 times smaller. Small differences in logits are large differences in probability, and that non-linearity is why a model that is “slightly more confident” produces overwhelmingly different sampling behaviour.

The max-subtraction trick

Softmax is invariant to adding a constant to every input. Subtract c from every logit and the exp(-c) factor appears in the numerator and in every term of the denominator, so it cancels:

exp(z_i - c)               exp(z_i) * exp(-c)          exp(z_i)
------------------- = ---------------------------- = -----------------
sum_j exp(z_j - c)     sum_j exp(z_j) * exp(-c)      sum_j exp(z_j)

Since it is free, choose c = max(z). Then the largest input becomes 0, exp(0) = 1, and every other term is between 0 and 1. Nothing can overflow. Same three numbers:

c = max(z) = 2.0
z - c = [0.0, -1.0, -1.9]

exp( 0.0) = 1.000000
exp(-1.0) = 0.367879
exp(-1.9) = 0.149569
            --------
sum       = 1.517448

p_1 = 1.000000 / 1.517448 = 0.659001
p_2 = 0.367879 / 1.517448 = 0.242433
p_3 = 0.149569 / 1.517448 = 0.098566

Identical to six decimal places, because it is identical in exact arithmetic. This is not an approximation or a stability “trade-off”. It is the same function, evaluated somewhere the floating-point representation can cope with.

The overflow it prevents, in real numbers

The largest finite value in fp32 is about 3.4028e38. Since ln(3.4028e38) = 88.72, any logit above about 88.7 makes exp overflow to infinity. In fp16 the largest finite value is 65504, and ln(65504) = 11.09 — a logit of 12 overflows.

Naive softmax on z = [100.0, 99.0, 98.0] in fp32:

  exp(100) = 2.688e43   -> inf   (max fp32 is 3.4028e38)
  exp(99)  = 9.889e42   -> inf
  exp(98)  = 3.637e42   -> inf
  sum      = inf
  p        = inf / inf  -> nan, nan, nan

With max subtraction, c = 100:

  exp( 0) = 1.000000
  exp(-1) = 0.367879
  exp(-2) = 0.135335
  sum     = 1.503214

  p = [0.665241, 0.244728, 0.090031]

The naive version returns three NaNs; the shifted version returns the correct distribution. And notice that the answer only ever depended on the differences between the logits — [100, 99, 98] and [0, -1, -2] are the same distribution — which is the shift-invariance property, doing useful work.

Underflow is the benign direction. exp(-1000) rounds to exactly 0, and 0 is the right answer to seven hundred decimal places, so nothing is lost. The asymmetry is why the trick subtracts the maximum rather than, say, the mean: pushing everything down is safe, pushing anything up is not.

Temperature is a division before the exp

Temperature is not a separate mechanism. It divides the logits before softmax sees them:

p_i = exp(z_i / T) / sum_j exp(z_j / T)

z = [2.0, 1.0, 0.1]

T = 1.0   ->  [0.6590, 0.2424, 0.0986]
T = 0.5   ->  logits [4.0, 2.0, 0.2]
              exp    [54.598, 7.389, 1.221]  sum 63.209
              p      [0.8638, 0.1169, 0.0193]
T = 2.0   ->  logits [1.0, 0.5, 0.05]
              exp    [2.7183, 1.6487, 1.0513]  sum 5.4183
              p      [0.5017, 0.3043, 0.1940]

Halving the temperature took the top token from 66% to 86%; doubling it took it down to 50%. As T approaches 0 the distribution approaches a spike on the argmax, which is what “temperature 0” means — though implementations special-case it rather than dividing by zero, and temperature 0 still is not fully deterministic in practice. As T grows large every logit approaches 0 and the distribution approaches uniform.

The rest of the sampler — top-p and top-k — operates on the output of this function, and the order those steps run in changes the result.

Where else softmax appears

The vocabulary softmax is the famous one, and it is not the one that runs most often. A transformer applies softmax three or four times per layer, over much smaller sets, and the arithmetic is identical each time.

• Attention weights. Every row of the score matrix goes through softmax so the weights over the keys sum to 1. For a 32-layer model with 32 heads over a 4,096-token sequence, that is 32 * 32 * 4096 = 4,194,304 separate softmaxes of length 4,096 per forward pass — against one softmax of length 128,000 for the output. The attention softmaxes dominate by orders of magnitude, which is why the running-maximum version of this trick is what makes tiled attention kernels possible at all.
• The expert router. In a mixture-of-experts layer a small network scores each expert and a softmax turns those scores into mixing weights, from which the top few are kept. The router softmax runs over 8 to 128 values and decides which billions of parameters are touched.
• Classifier heads. Any fine-tuned classification head — a safety filter, an intent router, a reranker scoring relevance — ends in a softmax over a handful of labels. This is the one place the output probabilities get read directly by a threshold, and therefore the one place calibration matters.

The cost of the vocabulary softmax is worth one line, because it is the largest single tensor in the model’s forward pass. The output projection is a 4096 x 128,000 matrix — 524 million parameters, 7.5% of a 7B model in one layer — and producing the logits costs 2 * 4096 * 128,000 = 1.05 GFLOP per token, against 13.9 GFLOP for the entire rest of the model. The exponentials themselves are cheap; getting to them is not.

Four properties worth knowing

• It never outputs exactly 0 or exactly 1. exp of any finite number is strictly positive, so every token in the vocabulary keeps a nonzero probability. In floating point it can underflow to 0, which is a representation artefact, not the function.
• It is shift-invariant but not scale-invariant. Adding 5 to every logit changes nothing; multiplying every logit by 5 sharpens the distribution enormously. This is exactly why a raw logit value on its own tells you nothing.
• It preserves order. The largest logit always gets the largest probability, so greedy decoding by argmax over logits and argmax over probabilities are the same operation. Computing the softmax to pick the maximum is wasted work.
• Its derivative is cheap and that is why it won. Paired with cross entropy, the gradient with respect to the logits collapses to p - y — predicted probability minus the one-hot target. One subtraction per logit, no exponentials in the backward pass.