Stable log_softmax

Log-softmax returns the natural logarithms of softmax probabilities. This is useful when a loss needs log probabilities rather than probabilities:

import jax
import jax.numpy as jnp

x = jnp.array([1.0, 2.0, 3.0])
log_probs = jax.nn.log_softmax(x)
print(jnp.round(log_probs, 4))  # approximately [-2.4076, -1.4076, -0.4076]
print(jnp.isclose(jnp.exp(log_probs).sum(), 1.0))  # True

The log probabilities do not themselves sum to zero. Their exponentials sum to approximately one. For four equal scores, each result is -log(4), about -1.386294.

Why log after softmax can fail

A small probability can underflow to zero before you take its logarithm. Then log returns negative infinity, even when the correct log probability is finite:

x = jnp.array([-1000.0, 0.0], dtype=jnp.float32)
print(jnp.log(jax.nn.softmax(x)))  # [-inf, 0.0]
print(jax.nn.log_softmax(x))       # [-1000.0, 0.0]

This is a numerical difference, not just a speed preference. Runtime depends on the backend and compilation, so the lesson makes no fixed speedup claim.

Keep the calculation in shifted coordinates

Mathematically, log_softmax(x) equals x - logsumexp(x). But computing those two large quantities separately can lose small differences through rounding. For a manual implementation, first let shifted = x - max(x), then compute shifted - log(sum(exp(shifted))). That avoids both the probability underflow above and subtracting a large rounded normalization constant.

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.