Softmax — Module 11.1

How This Lesson Fits the Module & Volume

Logits are scores; a language model must emit a probability distribution over the next token. Softmax is the standard map from logits to that distribution.

You already met softmax inside Vol. 10 attention. Here it sits on the vocabulary axis so sampling and related decoding strategies have a valid categorical distribution to draw from.

Learning Objectives

By the end of this lesson, students should be able to:

Write the softmax formula and state its output properties.
Explain numerical stability via the max-subtraction trick.
Apply softmax over the vocabulary dimension for next-token prediction.
Contrast softmax-for-attention with softmax-for-decoding.
Show how temperature is implemented as logits / T before softmax.
Use torch.softmax / F.softmax correctly on dim=−1.

Definition

For a logit vector z ∈ R^V, softmax is defined by

softmax(z)_i = exp(z_i) / Σ_j exp(z_j).

The result is a probability vector: every entry is positive (strictly, if logits are finite) and the entries sum to 1. In practice we compute softmax(z − max(z)) for stability.

What Softmax Guarantees

Logits

Any real scores

Exp

All positive

Normalize

Divide by sum

Distribution

Ready to sample

Use site	Softmax over	Meaning
Attention (Vol. 10)	Key positions	Weights mixing values
LM head (this lecture)	Vocabulary	P(next token \| context)
Classification head	Class labels	P(class \| input)

Softmax vs Argmax vs Sampling

Softmax alone

Produces full distribution.
Does not choose a token.
Needed for entropy / top-p.

Argmax (greedy)

Picks the mode.
Equivalent to sampling at T→0.
Can be repetitive.

Sample from softmax

Draws one token stochastically.
Enables diversity.
Core of the Sampling lecture.

Code: Softmax and Temperature

import torch
import torch.nn.functional as F

logits = torch.tensor([2.0, 1.0, 0.1])

probs = F.softmax(logits, dim=-1)
print(probs)                 # sums to 1
print(probs.sum())           # tensor(1.)

# Temperature T: sharper (T<1) or flatter (T>1)
T = 0.5
probs_sharp = F.softmax(logits / T, dim=-1)
T = 2.0
probs_flat = F.softmax(logits / T, dim=-1)
print(probs_sharp, probs_flat)

# Mask a forbidden token with -inf before softmax
masked = logits.clone()
masked[2] = float("-inf")
print(F.softmax(masked, dim=-1))  # last entry ~0

Strengths and Tradeoffs

Strengths

Clean probabilistic semantics for categorical tokens.
Differentiable; pairs with cross-entropy.
Masks via −∞ work cleanly.

Tradeoffs

Peakiness can hide near-ties among top tokens.
Full V softmax is costly at huge vocabularies.
Floating overflow without the max trick.

Common Misconception

“Softmax always picks the best token.” Softmax only normalizes. Choosing still requires argmax or sampling. Students often print softmax output and assume the model already “decided” without an explicit decode step.

Knowledge Check

Short Answer: Write softmax(z)_i in words. Answer: exp(z_i) divided by the sum of exp(z_j) over j.
True/False: Softmax outputs always sum to 1 (finite logits). Answer: True.
Multiple Choice: For LM decoding, softmax is over: (a) batch, (b) vocabulary, (c) layers. Answer: (b).
Short Answer: Why subtract max(z) before exp? Answer: Numerical stability / avoid overflow.
True/False: Softmax by itself selects a token ID. Answer: False.
Multiple Choice: Logit −∞ after softmax becomes: (a) 1, (b) ~0, (c) NaN always. Answer: (b).
Short Answer: How does temperature enter the formula? Answer: Softmax(logits / T).
True/False: Attention and LM heads both can use softmax. Answer: True.
Multiple Choice: As T → 0, softmax becomes: (a) uniform, (b) one-hot on the max, (c) undefined always. Answer: (b).
Short Answer: What lecture uses the distribution to draw a token? Answer: Sampling.

Key Takeaways

Softmax turns logits into a categorical distribution over V tokens.
It normalizes; it does not decode by itself.
Temperature and masks are applied in logit space before softmax.
Same math as attention softmax, different axis/meaning.
Next: Sampling.

Trainer’s Guide

Hands-on idea: Give three logits and have students compute softmax by hand, then with PyTorch; compare entropy at T=0.5 vs T=2.

Discussion prompt: When might you use sparsemax or sampled softmax instead of full-vocab softmax?

Recap: Softmax converts vocabulary logits into probabilities for next-token prediction. Continue with Sampling.