CODING CHALLENGE · N°52

Softmax with Temperature

Easy AISamplingMath

The dial that controls how "creative" a language model is: temperature scales the logits before the softmax. Low temperature sharpens the distribution toward the top token (greedy, deterministic); high temperature flattens it toward uniform (random, diverse). Implement it — numerically stable. Solve it in Python or TypeScript, with hidden tests.

The problem

Implement softmax_temp(logits, temperature): divide each logit by temperature, then apply the softmax to turn the scaled logits into a probability distribution (positive values summing to 1). Do it in a numerically stable way by subtracting the max before exponentiating. Lower temperature → sharper (more peaked) distribution; higher temperature → flatter.

EXAMPLE 1

Input logits = [5, 5, 5, 5], temperature = 2.0

Output [0.25, 0.25, 0.25, 0.25]

equal logits give a uniform distribution at any temperature

EXAMPLE 2

Input logits = [2, 1, 0], temperature = 0.5

Output ≈ [0.867, 0.117, 0.016]

low temperature concentrates mass on the top logit

CONSTRAINTS

Return a list of probabilities that are all ≥ 0 and sum to 1.
Subtract max(scaled_logits) before exp to avoid overflow — this does not change the result, only its stability.
Temperature < 1 sharpens the distribution (toward argmax); temperature > 1 flattens it (toward uniform).

SOLVE IT YOURSELF

Your turn — write it

Edit the stub, hit Run (or ⌘/Ctrl + Enter), and watch the hidden tests. Stuck? the hints are right above and Reveal solution is one click away.

YOUR TASK

Implement softmax_temp(logits, temperature): scale by dividing by the temperature, subtract the max for stability, exponentiate, and normalise by the sum.

HINTS — 4 IDEAS

First scale: z[i] = logits[i] / temperature.
Subtract the maximum of z from every element before exponentiating — this prevents overflow and leaves the probabilities unchanged.
Exponentiate each shifted value, then divide by the sum of all of them so the result sums to 1.
Sanity check: at high temperature the outputs approach uniform; at low temperature the largest logit dominates.

CPython · WebAssembly

Approach, complexity & discussion — open after you solve

The approach

Divide the logits by the temperature T before applying softmax. T < 1 sharpens the distribution (more deterministic, more peaked), T > 1 flattens it (more diverse), and T → 0 approaches picking the argmax. Keep the standard max-subtraction for numerical stability.

Complexity

Time O(n) over the logits.

Common mistakes

Multiplying by T instead of dividing — that inverts the effect.
Dropping the max-subtraction and overflowing on large logits.
Not special-casing T = 0, which is a division by zero — treat it as argmax.

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.

Finished this one? 0 / 75 Challenges done

Explore the topic

See this alongside everything else on the same subject — handbooks, system designs, challenges and tools, in one place.

ML Foundations