Softmax

Motivation

A network head outputs raw scores called logits, one per class or per vocabulary token. Logits are unbounded real numbers, so they are not yet probabilities. Softmax is the function that converts a vector of logits into a proper probability distribution, which is what Sampling then draws from.

The function

softmax (z)_{i} = \frac{e ^{z_{i}}}{\sum _{j} e ^{z_{j}}}

the exponential makes every value positive
the denominator normalises the values so they sum to 1
the output keeps the same ranking as the input, so the largest logit stays the largest probability

Worked example on three logits:

Token	Logit $z$	$e^{z}$	Probability
red	3.0	20.1	0.84
blue	1.0	2.7	0.11
green	0.0	1.0	0.04

A gap of 2 in the logits became a factor of about 7 in the probabilities, which shows how the exponential amplifies differences.

Properties worth remembering

Four properties explain most of its behaviour:

shift invariance: adding a constant to every logit leaves the output unchanged, which is why implementations subtract the maximum logit before exponentiating for numerical stability
soft argmax: it is a smooth and differentiable stand-in for picking the maximum, so gradients can flow through it
only differences matter: the absolute size of a logit is meaningless, and only the gaps between logits carry information
winner takes most: because of the exponential, a moderate lead in logit space becomes a large lead in probability space

Temperature

Dividing the logits by a temperature $τ$ before the softmax rescales all the gaps:

softmax (z / τ)

Temperature	Effect on the distribution	Behaviour
$τ < 1$	sharper, lower entropy	conservative and repetitive
$τ = 1$	the raw model distribution	default
$τ > 1$	flatter, higher entropy	diverse and more random
$τ \to 0$	one-hot	equal to greedy argmax

Temperature changes behaviour only, and it never changes the model weights.

Where it appears in a transformer

Three distinct uses in the same model:

attention weights, where $A = softmax (Q K^{T} / d_{h} + M)$ turns match scores into retrieval probabilities over positions
the output head, where $p_{θ} (x_{t + 1} ∣ x_{\leq t}) = softmax (z_{t})$ turns the final hidden state projection into a next-token distribution
classification heads in general, where the softmax output is compared against the label

Training uses the logits, not the probabilities

A common point of confusion:

CrossEntropyLoss in PyTorch accepts raw logits, because it fuses the log-softmax and the negative log-likelihood into one numerically stable operation
applying softmax yourself and then passing the result to that loss applies the softmax twice and blunts the gradients
softmax is needed explicitly at inference, when you want a distribution to sample or inspect

Code

PyTorch: apply softmax over class columns

import torch
 
logits = torch.tensor([[3.0, 1.0, 0.0], [0.0, 1.0, 3.0]])
probabilities = torch.softmax(logits, dim=-1)
print(probabilities.round(decimals=3))
# [[0.844, 0.114, 0.042], [0.042, 0.114, 0.844]]
print(probabilities.sum(dim=-1))  # [1., 1.]
print(probabilities.argmax(dim=-1))  # [0, 2]

The input has two rows and three class columns. dim=-1 selects the last axis, so each row becomes a distribution over its three classes. The shape stays (2, 3). Using dim=0 would combine different samples within each column.

TensorFlow/Keras

import tensorflow as tf
 
logits = tf.constant([[3.0, 1.0, 0.0], [0.0, 1.0, 3.0]])
probabilities = tf.nn.softmax(logits, axis=-1)
print(tf.reduce_sum(probabilities, axis=-1).numpy())  # [1., 1.]

axis is the TensorFlow argument corresponding to PyTorch's dim. For a training loss configured to accept logits, pass the original raw scores directly to the loss.

JJ

Explorer