Decoding Temperature

Why this matters

Logits are scores before normalization. Temperature changes the gaps between these scores before softmax turns them into probabilities.

For positive temperature T, calculate softmax(logits / T). A smaller temperature concentrates probability on the highest scores. A larger temperature makes the distribution more uniform. At T = 1, you get the original softmax distribution.

This changes sampling probabilities, not the model's knowledge or factual reliability. It does not guarantee creative or coherent text.

Worked mini-example

import jax
import jax.numpy as jnp

logits = jnp.array([1.0, 2.0, 3.0])
baseline = jax.nn.softmax(logits)        # approximately [0.090, 0.245, 0.665]
sharper = jax.nn.softmax(logits / 0.5)   # approximately [0.016, 0.117, 0.867]
flatter = jax.nn.softmax(logits / 2.0)   # approximately [0.186, 0.307, 0.506]

Index two remains the most likely choice. Temperature changes how often sampling selects other indices; it does not select an index itself.

Common pitfalls

• Dividing probabilities by T after softmax is not temperature scaling. Their sum becomes 1 / T; renormalizing that vector simply recovers the original probabilities.
• From exact positive probabilities, raising each to 1 / T and renormalizing is mathematically equivalent. Working from logits is usually preferable because tiny probabilities may already have rounded to zero.
• Zero temperature is outside this formula. Use argmax for greedy selection. As positive temperature approaches zero, equal maximum logits still share probability.
• Positive scaling preserves score order mathematically. Equal logits stay equal; uniform logits stay uniform.
• Use the stable softmax implementation rather than directly exponentiating large scores. Extremely small temperatures can still overflow the preceding division.

Turn scores into probabilities

Softmax converts a list of scores, also called logits, into probabilities that sum to one. It exponentiates each score and divides by the sum:

$$p_i=\frac{\exp(z_i/T)}{\sum_j\exp(z_j/T)},\qquad T>0.$$

T is the temperature. With the usual choice T = 1, scores [0, 0] give probabilities [0.5, 0.5]. Scores [0, log(3)] give [0.25, 0.75], because their exponentials are [1, 3].

Lower temperature puts more probability on the largest scores. Higher temperature makes the probabilities more even. A positive temperature preserves the ordering of finite scores.

Subtract the maximum before exponentiating

Large positive scores can make exponentials overflow. Subtracting the largest scaled score avoids large positive exponents without changing the probabilities: the common exponential factor cancels between numerator and denominator.

For a finite, nonempty list of scaled scores, every shifted score is at most zero and at least one is exactly zero. Their exponentials are therefore at most one, with at least one equal to one. For example, [1000, 1000] becomes [0, 0] before exponentiation and still gives [0.5, 0.5].

Choose the axis that contains the alternatives

For logits shaped (batch, classes), apply softmax over classes. Each example then gets its own probability distribution. Summing over the entire tensor would incorrectly make examples compete with each other.

When implementing the division yourself, keep the summed axis as a size-one dimension so the denominator broadcasts back across the intended scores.

If you need log probabilities, use log-softmax or a log-sum-exp formulation. Taking the logarithm after softmax can produce negative infinity when a small probability has already rounded to zero.

For masked softmax, each row needs at least one allowed, finite score. A row containing only negative infinity has no valid choices to normalise.