Temperature Scaling - LLM Generation & Decoding

Temperature Scaling

What you'll build

Turn a vector of next-token logits into probabilities with a temperature T: divide the logits by T, then take a numerically stable softmax.

Background

At every generation step the model outputs one logit per vocabulary token. Temperature is the first knob every decoding API exposes. Dividing by T < 1 stretches the gaps between logits, so the distribution gets sharper (more deterministic); T > 1 shrinks the gaps and flattens it (more diverse, more mistakes). It never changes the ranking of tokens, only how peaked the distribution is.

Formula

z_i   = logits_i / T
p_i   = exp(z_i - max(z)) / sum_j exp(z_j - max(z))

Subtracting max(z) doesn't change the result (it cancels in the ratio) but keeps exp from overflowing.

Temperature	Effect
T -> 0	all mass on the argmax (greedy decoding)
T < 1	sharper than the model's distribution
T = 1	the model's own distribution
T > 1	flatter; T -> inf approaches uniform

Inputs

• logits: float array (vocab_size,) (a list also works)
• temperature: float > 0, default 1.0

Expected Return Format

probs: float array (vocab_size,), non-negative, sums to 1.

Edge cases and conventions

• Very small T (e.g. 0.01) makes the scaled logits huge; the max subtraction keeps the result finite.
• T = 0 itself is undefined (division by zero); real APIs special-case it as greedy argmax. You don't need to handle it here.

Constraints

NumPy only, vectorised.

Follow-up

• Show that temperature scaling preserves the argmax and the ranking.
• Temperature is also used to calibrate a classifier after training (Guo et al., 2017). How would you fit T on a validation set?

Input: temperature_scaling(np.array([2.0, 1.0, 0.1]), temperature=0.5)
Output: [0.8638, 0.1169, 0.0193]

Input: temperature_scaling(np.array([2.0, 1.0, 0.1]), temperature=2.0)
Output: [0.5017, 0.3043, 0.194]

import numpy as np

def temperature_scaling(logits, temperature=1.0):
    """
    Apply temperature scaling to logits and return probabilities.

    Args:
        logits: Raw logits from the model (vocab_size,)
        temperature: Temperature T > 0

    Returns:
        probs: softmax(logits / T), shape (vocab_size,), sums to 1
    """
    # Your code here
    pass