Softmax and Temperature

From Scores to Probabilities

Neural networks produce raw scores called logits — unbounded real numbers that can be positive, negative, or zero. A classification network with four output classes might produce logits like $[2.0, 1.0, 0.1, -1.0]$. These numbers tell us something about relative preferences, but they are not probabilities: they don't sum to 1, and some are negative. To interpret them as a probability distribution (positive values that sum to 1), we need a function that normalises them. Softmax is that function.

Let's break down every piece of this formula.

$z_i$ is the raw logit for class $i$. It can be any real number: positive, negative, or zero. This is the input to softmax — the raw score that the network assigns to class $i$ before any normalisation.

$e^{z_i}$ is the exponential of the logit. The exponential function makes all values strictly positive, since $e^x > 0$ for every real number $x$. Crucially, it also amplifies differences: larger logits become exponentially larger. A logit of 2 maps to $e^2 \approx 7.4$, while a logit of 4 maps to $e^4 \approx 54.6$ — a difference of 2 in logit space becomes a factor of roughly 7.4 in exponential space.

$\sum_{j=1}^{K} e^{z_j}$ is the normalisation constant. We sum the exponentials of all $K$ logits, then divide each individual exponential by this sum. This guarantees that all outputs sum to 1, giving us a valid probability distribution.

Output range : each $\text{softmax}(z_i) \in (0, 1)$ (strictly between 0 and 1, never exactly 0 or 1) and $\sum_{i=1}^{K} \text{softmax}(z_i) = 1$.

What happens at the extremes is instructive. If one logit is much larger than the rest (e.g., $z_1 = 10$, all others $\approx 0$), then $\text{softmax}(z_1) \approx 1$ and all other outputs are $\approx 0$. Softmax approaches a hard argmax — nearly all the probability mass lands on the winner. If all logits are equal ($z_i = c$ for all $i$), then every exponential is the same, so $\text{softmax}(z_i) = 1/K$ for all $i$ — a perfectly uniform distribution. And if you add a constant $c$ to all logits, $\text{softmax}(z_i + c) = \text{softmax}(z_i)$, because the $e^c$ factor cancels between numerator and denominator. Only relative differences between logits matter.

Let's see this step by step in code.

import numpy as np

def softmax(z):
    # Subtract max for numerical stability (doesn't change result)
    z_stable = z - np.max(z)
    exp_z = np.exp(z_stable)
    return exp_z / np.sum(exp_z)

# Example: 4-class classification logits
logits = np.array([2.0, 1.0, 0.1, -1.0])

print("Step-by-step softmax:")
print(f"  Logits:        {logits}")
print(f"  Subtract max:  {logits - np.max(logits)}")
print(f"  Exponentials:  {np.exp(logits - np.max(logits)).round(4)}")
print(f"  Sum of exp:    {np.exp(logits - np.max(logits)).sum():.4f}")
print(f"  Softmax:       {softmax(logits).round(4)}")
print(f"  Sum:           {softmax(logits).sum():.6f}")
print()

# All equal -> uniform
equal = np.array([1.0, 1.0, 1.0, 1.0])
print(f"  Equal logits {equal} -> softmax {softmax(equal).round(4)} (uniform)")

# One dominant -> near argmax
dominant = np.array([10.0, 0.0, 0.0, 0.0])
print(f"  Dominant {dominant} -> softmax {softmax(dominant).round(6)} (near argmax)")

Why Exponentials?

A natural question arises: why use $e^{z_i}$ specifically? Why not square the logits ($z_i^2$), take absolute values ($|z_i|$), or use some other function to make things positive before normalising? There are three compelling reasons the exponential is the right choice.

First, positivity . We need all values to be positive to form a valid probability distribution. The exponential satisfies this: $e^x > 0$ for all $x$, including negative values. Squaring also makes things positive, but...

Second, monotonicity . The exponential is monotonically increasing — larger logits always produce larger exponentials, which means larger probabilities. This preserves the ranking: if the network thinks class A is more likely than class B (higher logit), class A gets a higher probability. Squaring would break this: $z^2$ maps $z = -5$ to 25 and $z = 2$ to 4, reversing the ranking entirely. A function like $|z|$ has the same problem.

Third, gradient properties . When softmax is combined with cross-entropy loss (the standard loss for classification, which we'll cover in the next article), the gradient simplifies beautifully to $p_i - y_i$ — the predicted probability minus the true label. This clean gradient comes specifically from the exponential family of distributions and makes optimisation stable and efficient. Other positive-making functions would produce far messier gradients.

Temperature: Controlling Sharpness

Sometimes we don't want the standard softmax distribution. We might want a sharper distribution (more confident, more deterministic) or a flatter one (more uncertain, more exploratory). Temperature scaling gives us this control by dividing the logits by a parameter $\tau$ (tau) before applying softmax.

The parameter $\tau$ is called the temperature , borrowing terminology from statistical mechanics where temperature controls the randomness of particle states. Here's what different temperature values do.

$\tau = 1$ : standard softmax. No change — dividing by 1 is a no-op.

$\tau \to 0^+$ (low temperature): dividing by a tiny positive number makes all logits huge in magnitude, amplifying the differences between them. The softmax approaches a hard argmax — nearly all probability mass concentrates on the single largest logit. The model becomes very confident and deterministic.

$\tau \to \infty$ (high temperature): dividing by a huge number makes all logits approach 0, erasing the differences between them. The softmax approaches a uniform distribution — the model becomes maximally uncertain and random, assigning equal probability to every class.

In summary: $\tau < 1$ (low temperature) produces sharper, more confident distributions. $\tau > 1$ (high temperature) produces flatter, more uncertain distributions. Temperature is a single knob that smoothly interpolates between "always pick the best option" and "pick uniformly at random."

The following plot makes this concrete. We take the same set of five logits and apply softmax at five different temperatures, showing how the probability distribution changes from nearly one-hot (low $\tau$) to nearly uniform (high $\tau$).

import numpy as np
import json
import js

logits = np.array([2.0, 1.0, 0.5, -0.5, -1.0])
classes = ["A", "B", "C", "D", "E"]

def softmax_temp(z, tau):
    z_t = z / tau
    z_t = z_t - np.max(z_t)
    exp_z = np.exp(z_t)
    return exp_z / np.sum(exp_z)

temperatures = [0.2, 0.5, 1.0, 2.0, 5.0]
colors = ["#ef4444", "#f59e0b", "#3b82f6", "#10b981", "#8b5cf6"]

lines = []
for tau, color in zip(temperatures, colors):
    probs = softmax_temp(logits, tau)
    lines.append({"label": f"\u03c4 = {tau}", "data": probs.tolist(), "color": color})

plot_data = [{
    "title": "Softmax with Different Temperatures",
    "x_label": "Class",
    "y_label": "Probability",
    "x_data": classes,
    "lines": lines
}]
js.window.py_plot_data = json.dumps(plot_data)

Temperature in Practice

Temperature is not just a theoretical curiosity — it is used everywhere in modern machine learning, often as one of the most important hyperparameters. Here are three major applications.

LLM sampling. When a large language model generates text, it produces logits over the entire vocabulary at each step. The temperature controls how the next token is sampled from the resulting distribution. High temperature (1.0-1.5) encourages creative, diverse text by spreading probability across many tokens. Low temperature (0.1-0.5) produces more factual, deterministic outputs by concentrating probability on the most likely tokens. Temperature = 0 is equivalent to greedy decoding: always pick the single most probable token, with no randomness at all.

Knowledge distillation. When training a smaller "student" model to mimic a larger "teacher," Hinton et al. (2015) showed that using high temperature (typically $\tau = 4$ to $\tau = 20$) on both teacher and student softmax outputs is crucial. At $\tau = 1$, the teacher's predictions are often nearly one-hot — most of the probability sits on one class — so the student only learns "the answer is class 3." At high temperature, the distribution softens, revealing the relative rankings of all classes. The teacher might show that class 5 is more plausible than class 7, even though both have tiny probability at $\tau = 1$. This "dark knowledge" in the tail of the distribution contains rich information about class similarities that helps the student generalise better.

Contrastive learning. In models like CLIP , a low temperature ($\tau \approx 0.07$) is used in the contrastive loss. The loss pushes matching image-text pairs to have high similarity and non-matching pairs to have low similarity. A low temperature makes this loss sharper, forcing the model to discriminate more aggressively between matching and non-matching pairs. The temperature is a learned parameter in CLIP, starting around 0.07 and adapted during training.

Log-Softmax and Numerical Stability

In practice, we almost always need $\log(\text{softmax}(z_i))$ rather than $\text{softmax}(z_i)$ directly. The reason is that the standard training loss for classification — cross-entropy, covered in the next article — involves the logarithm of the predicted probability. Computing softmax first and then taking the log is numerically dangerous: softmax can produce values extremely close to 0 (e.g., $10^{-45}$), and $\log(0) = -\infty$. Even values that are merely very small can lose precision when stored in floating point.

The solution is to compute log-softmax directly, without ever materialising the intermediate softmax values.

This follows from taking the log of the softmax formula: $\log(e^{z_i} / \sum_j e^{z_j}) = z_i - \log \sum_j e^{z_j}$. The key term is $\log \sum_j e^{z_j}$, known as the log-sum-exp . It can be computed stably by factoring out the maximum logit $m = \max_j z_j$:

After subtracting $m$, all exponents are $\leq 0$, so no exponential overflows. The largest exponential is $e^0 = 1$, which is perfectly safe. This is why PyTorch provides F.log_softmax and F.cross_entropy (which fuses log-softmax with the loss computation internally) — they use this identity under the hood to avoid the dangerous intermediate step.

The following code demonstrates why this matters.

import numpy as np

logits = np.array([100.0, 101.0, 102.0])  # large logits

# Naive: overflow!
try:
    exp_z = np.exp(logits)
    naive_softmax = exp_z / np.sum(exp_z)
    naive_log_softmax = np.log(naive_softmax)
    print(f"Naive log-softmax: {naive_log_softmax}")
except:
    print("Naive approach: overflow or nan!")

# Stable: subtract max, then use log-sum-exp identity
m = np.max(logits)
log_sum_exp = m + np.log(np.sum(np.exp(logits - m)))
stable_log_softmax = logits - log_sum_exp
print(f"Stable log-softmax: {stable_log_softmax.round(4)}")
print(f"Sum of exp(log-softmax): {np.exp(stable_log_softmax).sum():.6f} (should be 1.0)")

Quiz

Test your understanding of softmax and temperature.