Python softmax and temperature

A classifier often produces raw scores called logits. Softmax turns those scores into positive numbers that sum to one, which makes them usable as a probability distribution. The trick is to subtract the largest score before exponentiating so large values do not overflow. You will also score confident mistakes with log loss and turn probabilities into a review threshold.

Exponentiate and normalize

Softmax applies exp to every logit, then divides by the total. exp makes all values positive and preserves their order. The division makes the results sum to one. A larger logit always receives a larger probability, but the output is not a claim that the model is correct.

import math

def softmax(logits):
    exps = [math.exp(x) for x in logits]
    total = sum(exps)
    return [x / total for x in exps]

print([round(p, 3) for p in softmax([2, 0])])

Positive values that sum to one

Subtract the maximum for stability

exp(1000) is too large for a float, even though the final proportions are ordinary. Subtracting the maximum logit shifts every value without changing the ratios. This is the standard stable implementation; it is not an approximation.

import math

def softmax(logits):
    largest = max(logits)
    exps = [math.exp(x - largest) for x in logits]
    total = sum(exps)
    return [x / total for x in exps]

print([round(p, 3) for p in softmax([1000, 1000])])

Shift first, then exponentiate

Use temperature to change sharpness

Divide every logit by a positive temperature before softmax. A temperature below one makes the largest score dominate; a temperature above one flattens the distribution. Temperature changes uncertainty in the output, not the underlying evidence.

import math

def softmax(logits, temperature=1.0):
    scaled = [x / temperature for x in logits]
    largest = max(scaled)
    exps = [math.exp(x - largest) for x in scaled]
    total = sum(exps)
    return [x / total for x in exps]

print([round(p, 3) for p in softmax([2, 0], 0.5)])
print([round(p, 3) for p in softmax([2, 0], 2.0)])

Low temperature is sharper

Read the distribution, not just the winner

The largest probability is the argmax, but the other values carry information. A distribution of [0.34, 0.33, 0.33] is very different from [0.98, 0.01, 0.01], even though both have the same winner. Check that the values are finite and sum to approximately one.

probs = [0.7, 0.2, 0.1]
winner = max(range(len(probs)), key=probs.__getitem__)
print(winner)
print(round(sum(probs), 3))

Argmax plus a sanity check

Log loss punishes confident mistakes

For one example, log loss is the negative log of the probability assigned to the correct class. A correct prediction with probability 0.9 contributes little; a confident wrong prediction contributes a lot. Average the per-example losses, and clip exact zeros before taking the logarithm.

import math

for probability in [0.9, 0.5, 0.01]:
    loss = -math.log(probability)
    print(probability, round(loss, 3))

Confidence changes the penalty

Turn probabilities into a review threshold

A probability becomes an action only after you choose a threshold. If a wrong automatic answer is expensive, send uncertain cases to review. Choose the threshold on validation data using the relative costs of review and error, not on the test set.

def action(probability, threshold=0.8):
    return "accept" if probability >= threshold else "review"

for probability in [0.95, 0.72, 0.40]:
    print(probability, action(probability))

A confidence gate, not a guess

Try it yourself

Change the values, predict the output, then run the code. The Python runtime is downloaded the first time you press Run, so the first run can take a few seconds.