The Softmax Function

Softmax turns a vector of scores into a probability distribution, the standard output layer for multi-class classification and attention weights.

Definition

Given a vector of scores (logits) z1..zK, softmax outputs p_i = e^{z_i} / sum_j e^{z_j}. Every output lies in (0, 1) and the outputs sum to 1, so the result is a valid probability distribution over K classes. The largest logit gets the largest probability, but all classes receive nonzero mass. Softmax is the multi-class generalization of the sigmoid.

Numerical stability

Kronos motion — neural operator

Exponentiating large logits can overflow. The standard fix subtracts the maximum logit from every score before exponentiating: this leaves the output unchanged mathematically because the constant cancels in numerator and denominator, but keeps the exponentials bounded. Every production implementation does this shift.

python

import numpy as np
def softmax(z):
    z = z - np.max(z)          # numerical stability
    e = np.exp(z)
    return e / np.sum(e)

Pairing with cross-entropy

Softmax is almost always trained with the cross-entropy loss. The gradient of that combination is remarkably clean: for a true class one-hot y, the gradient of the loss with respect to the logits is simply p - y. This tidy form is numerically stable and is why frameworks fuse softmax and cross-entropy into a single operation rather than computing them separately.

Temperature

Dividing logits by a temperature T before softmax controls sharpness. T > 1 flattens the distribution toward uniform; T < 1 sharpens it toward a one-hot. Temperature is used to sample from language models with adjustable randomness, and in knowledge distillation to expose the soft structure a large model assigns across classes.

Beyond classification

Softmax also normalizes attention scores inside transformers: raw compatibility scores between a query and every key are turned into weights that sum to one, so attention forms a weighted average over values. The same property that makes softmax a good classifier output, a normalized distribution, makes it the natural way to allocate attention across a sequence.

Outputs a probability distribution summing to one.
Subtract the max for numerical stability.
Cross-entropy gradient reduces to p minus y.
Temperature tunes sharpness; also used in attention.