Softmax with Temperature

Start from brute force

The obvious approach types in the textbook formula: exp(xᵢ / T) / Σ exp(xⱼ / T). In exact arithmetic it is correct, and it costs O(n) time and O(n) space for the exponentials. In floating point it breaks at both ends. Large inputs overflow: Python’s math.exp(1000) raises OverflowError, and JavaScript’s Math.exp(1000) is Infinity, so the ratio becomes NaN. A small temperature makes this easy to hit, since a logit of 20 at T = 0.01 becomes 2000. Very negative inputs underflow instead: for [-1000, -1001] every exponential is 0.0 and the division is by zero.

The approach

Two facts carry the solution. First, temperature divides the logits before softmax: z = x / T scales every gap between logits by 1/T, so T < 1 widens the gaps (sharper) and T > 1 shrinks them (flatter), while the ranking is unchanged for any positive T. Second, softmax ignores a constant subtracted from every input, because the factor exp(−m) cancels between numerator and denominator. Choosing m = max(z) makes the largest shifted value exactly 0, so no exp can overflow, and it makes the sum at least 1, so it cannot underflow to 0.

The reference does exactly this: z = [x / temperature for x in logits], m = max(z), exps = [math.exp(x - m) for x in z], then each entry divided by total = sum(exps). The result is non-negative, sums to 1 up to rounding, and keeps the order of the logits.

Walk through an example

For the second example, [2, 1, 0] at T = 0.5, the gaps double before exponentiating, and the last two rows run the same logits at T = 2 for contrast.

Complexity

The reference makes a constant number of passes over the n logits (scale, max, exponentiate, sum, normalize): O(n) time and O(n) space for z and exps. The naive formula has the same O(n) cost; its problem is numerical, not speed. Sampling a token from the result is one more O(n) pass.

Edge cases to test

• Large logits: softmax_temp([1000, 999], 1.0) returns [0.7311…, 0.2689…]. Without the shift, Python’s math.exp(1000) raises OverflowError, and in JavaScript the result is NaN.
• Very negative logits: [-1000, -1001] gives the same [0.7311…, 0.2689…]. Unshifted, both exponentials underflow to 0.0 and the division raises ZeroDivisionError.
• temperature = 0 is not handled: Python raises ZeroDivisionError, and TypeScript returns NaN for every entry. A caller that uses T = 0 to mean greedy decoding must branch to argmax before calling.
• A negative temperature flips the ranking: softmax_temp([2, 1, 0], -1.0) returns [0.0900…, 0.2447…, 0.6652…], so the lowest logit is now the most likely. Validate temperature > 0.
• A tiny temperature approaches argmax and splits ties evenly: at T = 0.01, [2, 1, 0] gives about [1.0, 3.7e-44, 1.4e-87], while [2, 2, 0] gives about [0.5, 0.5, 6.9e-88].
• A masked logit of -inf gets probability exactly 0 ([-inf, 1.0, 2.0] gives [0.0, 0.2689…, 0.7311…]), but a +inf logit or an all--inf row gives NaN, because inf − inf is NaN.

Common mistakes

• Multiplying by T instead of dividing — that inverts the effect.
• Dropping the max-subtraction and overflowing on large logits.
• Passing T = 0 straight through: Python raises ZeroDivisionError and TypeScript returns NaN. The limit as T → 0 is argmax (shared evenly among tied maxima), so branch to argmax instead of dividing.

Follow-ups an interviewer asks next

How do you sample a token from the result?

Draw u uniformly from [0, 1) and walk the cumulative sum of the probabilities until it passes u; that index is the sampled token, at O(n) per draw. For many draws from one distribution, precompute the cumulative sums once and binary-search each draw in O(log n), or build an alias table for O(1) draws.

How do temperature, top-k and top-p fit together?

Temperature reshapes the whole distribution. Top-k keeps only the k most likely tokens, and top-p (nucleus sampling) keeps the smallest set whose probabilities add up to at least p. A common order is temperature first, then truncation, then renormalize and sample. Truncation cuts the long tail that a high temperature would otherwise make likely.

Does the max shift still matter on a GPU in half precision?

Even more. float16 tops out at 65504, so exp overflows once its input passes about 11. Fused GPU softmax kernels apply the same shift, and the online softmax used in FlashAttention keeps a running maximum and rescales the partial sums as each block of scores arrives.

Why is temperature used in knowledge distillation?

Dividing both teacher and student logits by a T above 1 softens the teacher’s distribution, so the small probabilities it gives wrong classes, which carry information about how classes relate, become large enough to learn from. Hinton, Vinyals and Dean’s 2015 distillation paper did this and scaled the soft-target loss by T² so its gradients keep a comparable size.

Where this shows up

Temperature is the sampling dial in every LLM API: keep it low for factual and code tasks where you want the top token, raise it for creative diversity. The key insight is that it reshapes the probability distribution — it is randomness control, not a “creativity” knob — which is what lets one model serve both careful and exploratory generation.

Reference solution

import math

def softmax_temp(logits, temperature):
    z = [x / temperature for x in logits]     # scale by temperature
    m = max(z)                                # subtract max for stability
    exps = [math.exp(x - m) for x in z]
    total = sum(exps)
    return [e / total for e in exps]          # normalise to sum 1


print(softmax_temp([5, 5, 5, 5], 2.0))        # -> [0.25, 0.25, 0.25, 0.25]
print(softmax_temp([2, 1, 0], 0.5))           # -> ~[0.867, 0.117, 0.016]

function softmaxTemp(logits: number[], temperature: number): number[] {
  const z = logits.map((x) => x / temperature);     // scale by temperature
  const m = Math.max(...z);                          // subtract max for stability
  const exps = z.map((x) => Math.exp(x - m));
  const total = exps.reduce((a, b) => a + b, 0);
  return exps.map((e) => e / total);                // normalise to sum 1
}

console.log(softmaxTemp([5, 5, 5, 5], 2.0));         // -> [0.25, 0.25, 0.25, 0.25]
console.log(softmaxTemp([2, 1, 0], 0.5));            // -> ~[0.867, 0.117, 0.016]