Softmax with temperature

Background

Softmax turns a vector of logits into a probability distribution. A temperature $T > 0$ rescales the logits first, controlling how sharp the distribution is:

$$\text{softmax}(z; T)_i = \frac{e^{z_i / T}}{\sum_j e^{z_j / T}}$$

Low $T$ sharpens toward the argmax (more confident); high $T$ flattens toward uniform. $T = 1$ is plain softmax.

Problem statement

Implement softmax_temperature(z, T) returning the temperature-scaled softmax of the logit vector z.

Input

• z — array-like of logits.
• T — float $> 0$, the temperature.

Output

Returns an np.ndarray of probabilities (same length as z) that are non-negative and sum to 1.

Examples

Example 1 — equal logits give a uniform distribution

Input:  z = [1, 1, 1], T = 1
Output: [0.33333333 0.33333333 0.33333333]

Example 2 — plain softmax (T=1)

Input:  z = [2, 1, 0], T = 1
Output: [0.66524096 0.24472847 0.09003057]

Explanation: $e^2 : e^1 : e^0$ normalised.

Constraints

• Divide the logits by T before exponentiating.
• The output must sum to 1.
• Return an np.ndarray.

Notes

• As $T \to \infty$ the output approaches uniform; as $T \to 0^+$ it approaches a one-hot at the largest logit. Temperature is the knob behind "sampling temperature" in language models.