Softmax from the definition

Softmax turns a row of scores into probabilities:

softmax(x)_i = exp(x_i) / sum_j exp(x_j)

Evaluating this formula directly can overflow even when the probabilities are well defined:

import torch

x = torch.tensor([[1000.0, 1001.0]], dtype=torch.float32)
e = x.exp()
print(e.tolist())  # [[inf, inf]]
print((e / e.sum(dim=-1, keepdim=True)).tolist())  # [[nan, nan]]

The exponentials exceed float32's finite range. Dividing infinity by infinity then produces NaN. Very negative inputs cause a different failure: all exponentials can round to zero, leaving a zero denominator.

Shift before exponentiating

Subtract the maximum of each row before applying exp:

z = x - x.max(dim=-1, keepdim=True).values
e = z.exp()
probabilities = e / e.sum(dim=-1, keepdim=True)
print([round(v, 6) for v in probabilities[0].tolist()])
# [0.268941, 0.731059]

For real-valued arithmetic, subtracting the same constant from every score does not change softmax. The numerator and denominator both acquire the same factor, which cancels.

For finite inputs, subtracting the maximum makes the largest shifted score zero. Its exponential is one, and the other exponentials are at most one. This prevents large positive scores from overflowing during exponentiation and prevents every exponential from underflowing to zero.

Floating-point rounding still applies. Very small probabilities can become zero, and the shift is not a promise of exact output or correct handling of NaNs and infinities. Low-precision sums and extremely wide rows can require additional care. This exercise uses small rows of finite float32 scores.

Preserve the reduced dimension

x.max(dim=-1, keepdim=True) returns values and indices; use its .values for subtraction. Keeping the last dimension gives one maximum per row in a shape that broadcasts back across that row. Keep the dimension when summing exponentials for the same reason.

In application code, use torch.softmax(x, dim=-1). For log-probabilities, use log_softmax rather than taking the logarithm of probabilities that may have rounded to zero. Implementing the shift here helps explain the numerical-stability idea also used by log-sum-exp.

Turn scores into probabilities

Softmax converts a list of scores, also called logits, into probabilities that sum to one. It exponentiates each score and divides by the sum:

$$p_i=\frac{\exp(z_i/T)}{\sum_j\exp(z_j/T)},\qquad T>0.$$

T is the temperature. With the usual choice T = 1, scores [0, 0] give probabilities [0.5, 0.5]. Scores [0, log(3)] give [0.25, 0.75], because their exponentials are [1, 3].

Lower temperature puts more probability on the largest scores. Higher temperature makes the probabilities more even. A positive temperature preserves the ordering of finite scores.

Subtract the maximum before exponentiating

Large positive scores can make exponentials overflow. Subtracting the largest scaled score avoids large positive exponents without changing the probabilities: the common exponential factor cancels between numerator and denominator.

For a finite, nonempty list of scaled scores, every shifted score is at most zero and at least one is exactly zero. Their exponentials are therefore at most one, with at least one equal to one. For example, [1000, 1000] becomes [0, 0] before exponentiation and still gives [0.5, 0.5].

Choose the axis that contains the alternatives

For logits shaped (batch, classes), apply softmax over classes. Each example then gets its own probability distribution. Summing over the entire tensor would incorrectly make examples compete with each other.

When implementing the division yourself, keep the summed axis as a size-one dimension so the denominator broadcasts back across the intended scores.

If you need log probabilities, use log-softmax or a log-sum-exp formulation. Taking the logarithm after softmax can produce negative infinity when a small probability has already rounded to zero.

For masked softmax, each row needs at least one allowed, finite score. A row containing only negative infinity has no valid choices to normalise.