ProblemNest
Sign In
You are coding as a Guest. Sign in with your RoleNest account to permanently track your streak, earn XP, and climb the Campus Leaderboard!
Sign In with RoleNest
🔥Stable Softmax & Temperature (LLM Sampling)Medium
MediumAI & Machine Learning•Acceptance: 58.1%

Stable Softmax & Temperature (LLM Sampling)

Real-World Engineering Context
Directly implements the final token sampling layer of ChatGPT and Llama-3, translating raw transformer unnormalized log-probabilities into categorical next-token predictions.
Given a 1D array of float `logits` and a positive `temperature`, compute the numerically stable Softmax probability distribution. Each element should be rounded to 4 decimal places: $$P(i) = \frac{\exp((\text{logits}[i] - \max(\text{logits})) / T)}{\sum_j \exp((\text{logits}[j] - \max(\text{logits})) / T)}$$

Sample Test Cases

Input: [[10,10],1]
Expected: [0.5,0.5]
Input: [[1000,1000],1]
Expected: [0.5,0.5]

Constraints

  • 1 <= logits.length <= 1000
  • -1000.0 <= logits[i] <= 1000.0
  • 0.1 <= temperature <= 5.0