Softmax with temperature
An LLM's temperature parameter divides the logits before the softmax. Implement it in a numerically stable way and see its effect.
Write softmax(logits, temperature = 1) that returns an array of probabilities. Divide each logit by temperature, subtract the largest of those values (so Math.exp does not overflow with large logits), apply Math.exp and normalise so the sum is 1.
Throw an error when temperature is not a number greater than 0. With a lower temperature, the option with the largest logit gains probability; with a higher one, the distribution flattens.
LOGITS holds the logits of four candidate tokens.
Challenges 0/5
- softmax([1, 2, 3]) gives 0.0900 / 0.2447 / 0.6652
- Probabilities sum to 1
- Stays stable with large logits (no NaN)
- Lower temperature gives the top option more probability
- Throws when temperature is 0 or negative
function softmax(logits, temperature = 1) {
// scale by temperature, subtract the max, exp, normalise
const exps = logits.map((x) => Math.exp(x));
const sum = exps.reduce((a, b) => a + b, 0);
return exps.map((e) => e / sum);
}
console.log(softmax(LOGITS, 0.5));Console output appears here (console.log).
Go deeper: the Hugging Face reference →
This in production, with your data? Let's talk for 15 minutes →