Neural Mastery
← Practice
Temperature Scaled Softmaxeasy

Temperature Scaled Softmax

easy Inside an LLM 15–20 min

NumPy allowed · plain Python earns 10 bonus points

LLM step 1/4

Logits become a next-token distribution

Softmax turns a model score vector into probabilities.
Open lesson →
01prompt
02tokens
03logits
04probabilities
attention→next token
mechanism65%
layer20%
pattern10%
model5%

Softmax(logits) → probability distribution

Control the sharpness of output probability distributions via temperature-scaled softmax.

Task

Implement temperature_scaled_softmax(logits, temperature). Divide each logit by temperature. Then compute stable softmax probabilities by subtracting the maximum scaled logit before exponentiating. Return probabilities rounded to 4 decimal places. If temperature <= 0.0, raise ValueError. If logits is empty, return [].

Function Signature

temperature_scaled_softmax(logits: list[float], temperature: float) -> list[float]

Examples

Example 1: Standard temperature T=1.0
Input: logits = [2, 1, 0.5] temperature = 1
Output: [0.6285, 0.2312, 0.1402]
Example 2: Low temperature T=0.5 (sharpening)
Input: logits = [2, 1, 0.5] temperature = 0.5
Output: [0.8438, 0.1142, 0.042]
Example 3: Invalid temperature T=0 error
Input: logits = [1, 2] temperature = 0
Output: Raises ValueError

Constraints

  • Temperature <= 0.0 must raise ValueError.
  • Empty logits list must return [].
  • Probabilities must be rounded to 4 decimal places.
  • A plain-Python solution earns 10 bonus points.
Python3Saved ✓
def temperature_scaled_softmax(logits: list[float], temperature: float) -> list[float]:
"""Compute temperature-scaled softmax probabilities over logits.
Args:
logits: Raw token logits.
temperature: Positive scaling factor T > 0.
Returns:
List of probabilities summing to 1.0, rounded to 4 decimals.
"""
# Your implementation here
pass

Case 1: Standard temperature T=1.0
Input: logits = [2, 1, 0.5] temperature = 1
Expected: [0.6285, 0.2312, 0.1402]
Case 2: Low temperature T=0.5 (sharpening)
Input: logits = [2, 1, 0.5] temperature = 0.5
Expected: [0.8438, 0.1142, 0.042]
Case 3: Invalid temperature T=0 error
Input: logits = [1, 2] temperature = 0
Expected: Raises ValueError