Easy Research

Temperature Scaling

Temperature scaling changes how concentrated a softmax distribution is. Divide the model's logits by one positive temperature, then apply softmax over the classes for each example.

Your task

Implement temperature_scale(logits, temperature).

logits has shape (batch, n_classes), with at least one class per row.
temperature is a finite scalar T > 0, shared across all examples and classes.
Inputs are finite floating-point values; the scaled logits must remain representable in the input dtype.

Return probabilities of shape (batch, n_classes), preserving the logits' dtype and leaving the input unchanged. For each row, compute:

scaled = logits / temperature
shifted = scaled - max(scaled)
probabilities = exp(shifted) / sum(exp(shifted))

The maximum and sum are over the class axis, separately for each row. Subtracting the maximum makes softmax numerically stable without changing its mathematical result. Return probabilities, not scaled logits or class IDs.

For logits=[[2, 0]] and temperature=2, the scaled row is [1, 0] and the result is approximately [[0.731059, 0.268941]]. Each row sums to one, up to floating-point rounding.

What changes with temperature?

T=1 gives the ordinary softmax of the input logits.
T>1 makes the distribution less concentrated.
0<T<1 makes it more concentrated.

Equal logits always produce a uniform distribution. A positive temperature preserves class ranking, so the highest-logit class does not change in exact arithmetic. Do not divide probabilities by temperature or normalize across the batch.

Scaling versus calibration

Guo et al. (2017) study temperature scaling as a calibration method. Here the temperature is supplied: you do not fit it, change model weights or evaluate calibration. Applying an arbitrary temperature does not by itself guarantee well-calibrated probabilities.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

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.

Detecting runtime…