03-03 Softmax and Cross-Entropy

Why this matters

Transformer training optimizes next-token probabilities via softmax outputs and cross-entropy loss.

Intuition first (no jargon)

Softmax converts logits to a normalized distribution, and cross-entropy measures target surprise.

Worked example

• Inputs: logits = [3, 1, -2], targetId = 1
• Shapes: logits [V] -> probs [V]
• One computed step:
  • Stable shift: subtract max 3, so shifted logits are [0, -2, -5]
  • exp(shifted) is approximately [1.000, 0.135, 0.007]
  • Sum is approximately 1.142, so probs[1] = 0.135 / 1.142 = 0.118
  • Cross-entropy is -ln(0.118) = 2.14 nats

Code walkthrough

js
export function softmaxStable(logits) {}
export function crossEntropy(probs, targetId) {}

• softmaxStable: logits -> probabilities
• crossEntropy: target probability -> scalar loss

Your task

Implement stable softmax and cross-entropy.

• Subtract max logit before exponentiation.
• Ensure probabilities sum close to 1.
• Loss is -Math.log(probs[targetId]).

Common mistakes

• No max-logit shift before Math.exp
• No epsilon clamp before Math.log
• Out-of-range targetId

Precision note

One-token cross-entropy here is measured in nats.

Hints

• Guard against Math.log(0) with epsilon.
• Keep numerical stability explicit.
• Validate targetId range.

Check your thinking

1. Why subtract max logit?
2. Why does confident wrong prediction produce high loss?
3. What does near-zero loss imply?

Stretch (optional)

Add meanCrossEntropy for a mini-batch.

Likely test focus

• Stable probabilities for large logits.
• Probability sum near 1.
• Correct loss on known examples.

What should improve

You can now compute the objective used to evaluate and train token predictors.

Bridge to next lesson

Gradient update step in the next lesson.