Softmax with Temperature - PyTorch Basics

Softmax with Temperature

What you'll build: softmax_temperature(logits, temperature), a numerically stable softmax over the last dimension with temperature scaling.

Background

Dividing the logits by a temperature T before the softmax controls how peaked the distribution is. Language models sample with T < 1 for focused text and T > 1 for more diverse text; knowledge distillation softens a teacher's predictions with T > 1 so the student also learns from the wrong-but-plausible classes.

Formula

For logits z (softmax over the last dimension):

s = z / T
softmax(s)[k] = exp(s[k] - max(s)) / sum_j exp(s[j] - max(s))

Subtracting max(s) changes nothing mathematically (it cancels in the ratio) but keeps exp from overflowing: exp(1000) is inf in float64.

Temperature Effects

Temperature	Effect
T = 1	standard softmax
T < 1	sharper, more confident
T → 0	approaches one-hot at the argmax
T > 1	softer, closer to uniform
T → ∞	approaches uniform

The ranking of the classes never changes, only how much mass the top classes get.

Inputs

• logits: float array of any shape; normalise over the last axis.
• temperature: positive float (default 1.0).

Expected Return Format

An array of the same shape as logits, with non-negative entries that sum to 1 along the last axis.

Constraints

• Must stay finite for large logits (e.g. [1000, 1001, 1002]).
• PyTorch track: build it from tensor ops (max, torch.exp, sum); don't call F.softmax or torch.softmax.

Follow-up

• Top-k and nucleus (top-p) sampling are applied after temperature. Sketch how.
• In distillation the soft-target loss is multiplied by T². Why?

Input: softmax_temperature(np.array([2.0, 1.0, 0.1]), temperature=1.0)
Output: array([0.659 , 0.2424, 0.0986])

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

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

import numpy as np

def softmax_temperature(logits: np.ndarray, temperature: float = 1.0) -> np.ndarray:
    """
    Compute softmax with temperature scaling.

    Args:
        logits: Input logits (any shape, softmax along the last axis)
        temperature: Positive temperature (default 1.0)

    Returns:
        Probabilities with the shape of logits, summing to 1 along the last axis
    """
    # Your code here
    pass