Vibe Engines Learn
CODING CHALLENGE · N°04

Softmax

Easy AI EngineeringLLMDecoding

The function at the end of every classifier and language model: turn a vector of raw scores (logits) into a probability distribution. The catch every AI engineer learns the hard way — do it the numerically stable way so big logits do not overflow.

The problem

Given a list of logits (raw real-valued scores), return the softmax: a list of the same length where each value is exp(logit) / Σ exp(logits). The result is a probability distribution — every value in [0, 1] and the whole thing sums to 1. Subtract the max logit before exponentiating so large values do not overflow (this does not change the result).

EXAMPLE 1
Input logits = [0, 0]
Output [0.5, 0.5]
equal scores → uniform
EXAMPLE 2
Input logits = [2, 1, 0]
Output [0.665, 0.245, 0.090]
higher logit → higher probability
EXAMPLE 3
Input logits = [1000, 1000]
Output [0.5, 0.5]
the stable version must not overflow
CONSTRAINTS
  • 1 ≤ len(logits)
  • The output must sum to 1 (within floating-point tolerance).
  • Subtract max(logits) before exp — the numerically stable softmax.
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(logits) → a probability distribution over the logits. Subtract the max first for numerical stability, exponentiate, then divide by the sum.

HINTS — 4 IDEAS
  1. Probabilities must be positive and sum to 1 — exponentiate, then normalize by the total.
  2. Find m = max(logits) and exponentiate logit - m. Shifting by a constant cancels out in the ratio but stops exp from overflowing.
  3. Sum the exponentials once, then divide each by that sum.
  4. Sanity check: all-equal logits give a uniform distribution.
CPython · WebAssembly
Approach, complexity & discussion — open after you solve

The approach

Exponentiate each logit and divide by the sum of all the exponentials. The one non-obvious step is numerical stability: subtract the max logit before exponentiating — softmax(x) = softmax(x − max x) — which gives the identical result but keeps the exponentials from overflowing on large inputs.

Complexity

Time O(n) — one pass for the max, one to exponentiate and sum, one to normalize; space O(n) for the output (O(1) extra if written in place).

Common mistakes

  • Skipping the max-subtraction — exp() of a large logit overflows to inf and the result becomes nan.
  • Dividing by n (the count) instead of the sum of the exponentials.
  • Treating it element-wise — softmax is defined over the whole vector, so the denominator couples every entry.

Where this shows up

Softmax is the output layer of essentially every classifier and language model (a distribution over the next token) and sits inside attention (a distribution over the keys). The max-subtraction trick is a classic “implement it from scratch” interview check precisely because forgetting it is the bug that silently produces nan in real training code.

▶  Watch it explained

Explore the topic

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

LLM Engineering ML Foundations Interview Prep