Softmax Function: Formula, Examples, and Temperature

The softmax function turns a vector of real numbers (logits) into a probability distribution: every output is between 0 and 1 and all outputs add up to 1. Larger inputs get larger probabilities.

This process creates a “soft,” differentiable version of argmax. Instead of picking only the maximum value, it gives the largest value the biggest share while keeping the others. It is the standard activation function used when you need to model mutually exclusive classes.

Softmax formula

The core calculation is straightforward. For a vector of logits z, the probability for the i-th element is:

softmax(z)_i = exp(z_i) / sum over j of exp(z_j)

The exponential function ensures that all outputs are positive. Because the denominator is the sum of all exponentials, the outputs naturally sum to 1.

A key property of this formula is invariance to shifts. Adding the same constant to every logit does not change the result. This allows for numerical tricks, such as subtracting the maximum logit before exponentiating, which prevents overflow without altering the final probabilities.

Worked example

Let’s look at a concrete example to see how the math works in practice.

Logits: [2.0, 1.0, 0.1]
Exponentials: [7.389, 2.718, 1.105]
Sum of exponentials: 11.212
Probabilities: [0.659, 0.242, 0.099]

Notice how the exponential amplifies differences. The first logit is only 1.0 higher than the second, but its resulting probability is almost three times larger. This sensitivity is why softmax is effective at highlighting the dominant signal in a noisy vector.

Where softmax is used

The softmax function is a staple in modern machine learning architectures. Its primary use cases include:

Multi-class classification: It is the standard output layer for classifiers where classes are mutually exclusive. Each class gets one output neuron.
Attention mechanisms: In Transformer models, softmax is applied over scaled dot products. This turns attention scores into weights that sum to 1, allowing the model to focus on specific parts of the input sequence.
Large language models: In a Large language model, LLM, the final layer produces a logit for every token in the vocabulary. Softmax converts these logits into next-token probabilities, determining which word comes next in the generated text.

Softmax temperature

You can control the “sharpness” of the output distribution by introducing a temperature parameter T. The modified formula is:

softmax(z / T)

T < 1: Sharpening. The distribution becomes more confident, closer to always picking the top option.
T > 1: Flattening. The distribution becomes more random, reducing the gap between the highest and lowest probabilities.

This is exactly what the “temperature” setting in LLM APIs controls when sampling tokens. Here is how different temperatures affect the example logits [2.0, 1.0, 0.1]:

Temperature (T)	Output Probabilities	Effect
0.5	[0.864, 0.117, 0.019]	Sharpened; top class dominates
1.0	[0.659, 0.242, 0.099]	Standard distribution
2.0	[0.502, 0.304, 0.194]	Flattened; more random choices

Softmax vs sigmoid

While both functions map inputs to probabilities, they serve different purposes.

Sigmoid gives an independent probability for each output. It does not force the outputs to sum to 1. This makes it ideal for multi-label classification, where multiple labels can be true at once (e.g., an image tagged both “beach” and “sunset”).

Softmax makes the classes compete. It forces the outputs to sum to 1, making exactly one class the expected outcome. This is required for multi-class classification where only one label is correct.

Note that for two classes, softmax is equivalent to a sigmoid applied to the difference of the two logits. However, as the number of classes grows, the competitive nature of softmax becomes distinct from independent sigmoid outputs.

Numerical stability

A common issue with the softmax formula is overflow. If logits are large, exp(z_i) can exceed the maximum representable floating-point number, resulting in NaN (Not a Number).

To fix this, subtract the maximum logit from every element in the vector before exponentiating. Because of the invariance property mentioned earlier, this shift does not change the final result but keeps the numbers manageable.

In practice, libraries often combine log-softmax with the loss function. For example, PyTorch’s nn.CrossEntropyLoss takes raw logits and applies log-softmax internally. This is more numerically stable and efficient than applying softmax, taking the log, and then computing the loss separately. You should not apply softmax yourself before passing logits to these combined loss functions.

Softmax in Python

Here is a short implementation of the softmax function with temperature and the numerical stability trick.

import numpy as np

def softmax(z, T=1.0):
    z = np.asarray(z, dtype=float) / T
    z = z - z.max()        # numerical stability
    e = np.exp(z)
    return e / e.sum()

print(softmax([2.0, 1.0, 0.1]))  # [0.659 0.242 0.099]

This code divides by the temperature, shifts the values for stability and normalizes the result.

FAQ

What is the difference between softmax and sigmoid?

Softmax produces a probability distribution where all outputs sum to 1, making it suitable for mutually exclusive multi-class problems. Sigmoid produces independent probabilities for each output, making it better for multi-label classification where multiple labels can be true simultaneously.

Why do we subtract the maximum logit in softmax?

Subtracting the maximum logit before exponentiating prevents numerical overflow. Since exponentials of large numbers can exceed floating-point limits, this shift keeps the values small while preserving the relative differences and final probabilities.

How does temperature affect softmax output?

Temperature scales the logits before applying the softmax function. A temperature below 1 sharpens the distribution, making the model more confident in its top choice. A temperature above 1 flattens the distribution, making the output more random and uniform.