Softmax

Softmax is a function that turns a vector of raw scores, called logits, into a probability distribution: every output lies between 0 and 1, and the outputs sum to exactly 1. For a vector z, it computes softmax(z)ᵢ = e^(zᵢ) / Σⱼ e^(zⱼ), exponentiating each score and dividing by the sum of all the exponentials. Softmax is the standard output layer of neural networks that pick one class out of many, it turns similarity scores into weights inside the attention mechanism of transformers, and it converts a language model's scores over its vocabulary into next-token probabilities. Its name reflects that it is a smooth, differentiable version of the argmax that would put all probability on the single largest score.

How Softmax Works#

The exponential makes every output positive and magnifies differences between scores, so the largest logit takes a disproportionate share of the probability. Softmax preserves the order of the logits, so the highest score always gets the highest probability. Only differences between logits matter: adding the same constant to every logit leaves the output unchanged. Libraries use that property for numerical stability by subtracting the maximum logit before exponentiating, since e¹⁰⁰⁰ overflows 32-bit floating point.

A temperature T controls how peaked the distribution is: softmax(z / T). Values of T below 1 sharpen the distribution toward the top class, and values above 1 flatten it toward uniform. The same formula appears in statistical mechanics as the Boltzmann distribution, where T is literally temperature. The PyTorch snippet below shows both effects:

import torch

logits = torch.tensor([2.0, 1.0, 0.1])

# Divide logits by a temperature before softmax: low sharpens, high flattens
for temperature in [0.5, 1.0, 2.0]:
    probs = torch.softmax(logits / temperature, dim=0)
    print(f"T={temperature}: {[round(p, 3) for p in probs.tolist()]}")

# Adding a constant leaves softmax unchanged, so libraries subtract the max to avoid overflow
big = logits + 1000
naive = big.exp() / big.exp().sum()
print(f"Naive: {naive.tolist()}, torch.softmax: {[round(p, 3) for p in torch.softmax(big, dim=0).tolist()]}")

At T = 1 the logits 2.0, 1.0, and 0.1 become [0.659, 0.242, 0.099]. Halving the temperature sharpens this to [0.864, 0.117, 0.019], and doubling it flattens it to [0.502, 0.304, 0.194]. With 1,000 added to every logit, the naive formula overflows and prints [nan, nan, nan], while torch.softmax still returns [0.659, 0.242, 0.099].

Softmax in Training#

Classifiers are trained with softmax and cross-entropy loss together. The loss is −log(pᵧ), the negative log of the probability assigned to the correct class y, and its gradient with respect to the logits is simply p − y, the predicted probabilities minus the one-hot label. Frameworks therefore fuse the two: PyTorch's CrossEntropyLoss takes raw logits and computes log-softmax internally, which is more stable than taking the log of separately computed probabilities. A model trained this way outputs logits, and softmax is applied only when probabilities are needed at inference.

Where Softmax Is Used#

• Image classification: The final layer of a classifier produces one logit per class, and softmax turns them into the confidence scores a user sees.
• Attention: Transformers compute softmax(QKᵀ / √dₖ) over the similarity scores between a query and every key, so each token's attention weights sum to 1. Dividing by √dₖ keeps the scores from growing with dimension and saturating the softmax.
• Text generation: A large language model (LLM) applies softmax over a vocabulary of tens or hundreds of thousands of tokens at every step, and the sampling temperature is the T above. Nucleus sampling then draws from the smallest set of tokens whose probabilities reach a threshold.
• Reinforcement learning: Softmax over action values gives a policy that usually picks the best action while still exploring the others, with temperature trading exploration against exploitation.

YOLO26 classification models apply softmax in their classification head at inference, and the probs attribute of each result holds the full distribution over the 1,000 ImageNet classes:

from ultralytics import YOLO

# Load a YOLO26 classification model trained on ImageNet
model = YOLO("yolo26n-cls.pt")
results = model("https://ultralytics.com/images/bus.jpg")

# probs holds the softmax output of the classification head
probs = results[0].probs
for index, conf in zip(probs.top5, probs.top5conf.tolist()):
    print(f"{results[0].names[index]}: {conf:.3f}")
print(f"Sum over all 1,000 classes: {probs.data.sum().item():.3f}")

It prints the five most likely classes with their probabilities, and the sum over all 1,000 classes prints as 1.000 because the head returns a softmax distribution.

Softmax vs. Related Functions#

• Sigmoid: Squashes each score into 0 to 1 independently, so outputs need not sum to 1. Sigmoid suits binary and multi-label problems, such as an image containing both a person and a backpack, and object detectors commonly score each class with it. For two classes, softmax reduces exactly to a sigmoid of the difference between the two logits.
• Argmax: Returns the index of the largest score, a hard decision with no gradient. Softmax keeps the same ranking but stays differentiable, so it can be trained, and taking the argmax of softmax probabilities gives the same answer as the argmax of the logits.
• Log-softmax: The logarithm of softmax, computed directly for stability, which cross-entropy loss consumes.
• Temperature scaling: A calibration method that fits a single temperature on validation data after training, so the softmax probabilities match how often the model is actually right. It leaves the predicted class unchanged because dividing every logit by the same positive T preserves their order.