Softmax, Probability, and Sampling by Hand

In Linear Projections, Transposes, and Similarity, you learned how attention scores are computed as dot products between query and key vectors. Those scores are logits — unbounded real numbers that can be positive, negative, or zero. But downstream systems (the sampler, the loss function, the human reading the output) need a valid probability distribution: every value between 0 and 1, and the sum exactly 1.

This lesson teaches you how softmax performs that transformation, why a numerical trick is required to keep it working on real hardware, and how categorical sampling turns the distribution into an actual token choice.

---

Logits vs. Probabilities

Logits are the raw, unbounded scores a model produces — the output of a linear projection with no activation function applied. They live in $(-\infty, \infty)$. A logit of -5.0 is perfectly valid; so is 1000.0.

Probabilities must satisfy two rules:

1. Non-negativity: $0 \le p_i \le 1$ for every entry
2. Normalization: $\sum_i p_i = 1$ (the total probability mass is exactly 1, or 100%)

Softmax is the standard function that maps a vector of logits to a vector of probabilities.

---

The Softmax Function

For a vector of logits $z = [z_1, z_2, …, z_n]$, the softmax output $p = [p_1, p_2, …, p_n]$ is defined as:

p_i = exp(z_i) / (exp(z_1) + exp(z_2) + ... + exp(z_n))

Or using summation notation:

p_i = exp(z_i) / Σ_j exp(z_j)

What exp() Does: Two Key Effects

1. Non-negativity: The exponential function $e^z$ is strictly positive for every real input $z$. Even a large negative logit like -5.0 becomes a small positive number:

exp(-5.0) ≈ 0.0067

No matter how negative the logit, exp() maps it to a value in $(0, \infty)$. This guarantees $p_i > 0$.

2. Nonlinear amplification: exp() grows exponentially (hence the name). Larger logits are accentuated much more than smaller ones:

exp(2.0) ≈ 7.39
exp(3.0) ≈ 20.09

The raw difference is $3.0 - 2.0 = 1.0$ (a 1.5× ratio). After exponentiation, the ratio becomes $20.09 / 7.39 ≈ 2.7\times$. This is intentional — it makes the model “more confident” about higher scores.

3. The denominator normalizes: The sum in the denominator adds up all exponentials, then each numerator is divided by that total. This forces the outputs to sum to exactly 1.0.

---

Worked Example: Softmax by Hand on [1.0, 2.0, 3.0]

Let’s compute softmax for a tiny 3-element vector step by step.

Input logits: z = [1.0, 2.0, 3.0]

Step 1: Calculate exponentials

exp(1.0) = 2.7182818...  →  2.7183
exp(2.0) = 7.3890561...  →  7.3891
exp(3.0) = 20.0855369... → 20.0855

Step 2: Sum of exponentials

sum = 2.7183 + 7.3891 + 20.0855 = 30.1929

Step 3: Normalize each element

p_1 = 2.7183 / 30.1929 = 0.0900   (9.00%)
p_2 = 7.3891 / 30.1929 = 0.2447   (24.47%)
p_3 = 20.0855 / 30.1929 = 0.6652  (66.52%)

Step 4: Verify the sum

0.0900 + 0.2447 + 0.6652 = 0.9999

The tiny gap from 1.0000 is rounding error from using 4 decimal places. With full precision the sum is exactly 1.0.

Result: The highest logit (3.0) receives ~66.5% probability. The lowest (1.0) receives ~9%. The middle gets ~24.5%.

---

Numerical Stability: The Max-Subtraction Trick

The Problem: Floating-Point Overflow

In 32-bit floating point (IEEE 754 float32), the maximum finite value is approximately $3.4 \times 10^38$. If a logit is large, say $z = 1000.0$, then exp(1000) overflows to +inf (infinity).

If all logits are large and positive, every exp(z_i) becomes +inf. The softmax formula then computes inf / inf, which produces NaN (Not a Number). Inference crashes!

Even without overflow, very large exponentials lose precision because float32 only has ~7 decimal digits of precision.

The Solution: Shift Invariance

Softmax has a crucial mathematical property: adding the same constant to every logit does not change the output probabilities.

For any constant $c$:

exp(z_i - c) / Σ_j exp(z_j - c)
= [exp(z_i) * exp(-c)] / [exp(-c) * Σ_j exp(z_j)]
= exp(z_i) / Σ_j exp(z_j)

The exp(-c) factor cancels out completely!

The Standard Practice: Subtract the Maximum

Choose $c = \max(z)$. Then:

• The largest shifted logit becomes $\max(z) - \max(z) = 0$
• exp(0) = 1.0 (exactly, no rounding)
• All other shifted logits are $\le 0$, so their exponentials fall in $(0, 1.0]$

Overflow is now mathematically impossible — no exponential can exceed 1.0.

Worked Example: [1001.0, 1002.0, 1003.0] with Max Subtraction

Original logits: z = [1001.0, 1002.0, 1003.0]

Step 1: Find the maximum

c = max(z) = 1003.0

Step 2: Subtract c from every element (shift)

z' = [1001.0 - 1003.0, 1002.0 - 1003.0, 1003.0 - 1003.0]
   = [-2.0, -1.0, 0.0]

Step 3: Exponentiate (all values safely in (0, 1])

exp(-2.0) = 0.1353353... → 0.1353
exp(-1.0) = 0.3678794... → 0.3679
exp(0.0)  = 1.0000000... → 1.0000

Step 4: Sum

sum = 0.1353 + 0.3679 + 1.0000 = 1.5032

Step 5: Normalize

p_1 = 0.1353 / 1.5032 = 0.0900
p_2 = 0.3679 / 1.5032 = 0.2447
p_3 = 1.0000 / 1.5032 = 0.6652

Result: Identical to the [1.0, 2.0, 3.0] case — zero overflow risk, identical probabilities.

Key insight: Every production transformer implementation (PyTorch, TensorFlow, JAX, llama.cpp, vLLM, TensorRT-LLM) applies this max-subtraction trick before computing softmax. It is not optional — it is required for numerical correctness.

---

Greedy Decoding (Argmax) vs. Categorical Sampling

Once we have a probability distribution $p$, how do we pick a token?

Argmax (Greedy Decoding)

chosen_index = argmax(p)

Always selects the single highest-probability token. Deterministic. Zero variety. If you run the same prompt twice, you get the exact same output.

Categorical Sampling

Treats the probabilities as weights on a roulette wheel. The algorithm:

1. Build cumulative intervals (CDF):

   Given p = [0.0900, 0.2447, 0.6652]:

   • Interval 0 (token 0): [0.0000, 0.0900)
   • Interval 1 (token 1): [0.0900, 0.3347) (0.0900 + 0.2447 = 0.3347)
   • Interval 2 (token 2): [0.3347, 1.0000) (0.3347 + 0.6652 = 1.0000)

2. Draw a uniform random number: draw r from a uniform distribution on [0, 1]

3. Find which interval contains $r$: That token is selected.

Example: If $r = 0.20$:

• $0.20$ falls in [0.0900, 0.3347) → Token 1 is chosen, even though Token 2 had the highest single probability (66.5%).

Over many draws:

• Token 0 chosen ~9.0% of the time
• Token 1 chosen ~24.5% of the time
• Token 2 chosen ~66.5% of the time

This matches the probability distribution exactly — that is the definition of categorical sampling.

---

Where Softmax Appears in Modern LLMs

There are two essential places softmax appears inside every decoder-only transformer:

1. Attention Weights (Inside Every Layer)

attn_weights = softmax(Q @ K^T / sqrt(d_k))

For each query token, this normalizes the attention scores across all context tokens so they form a probability distribution (sum to 1). The resulting weights determine how much each value vector contributes to the output.

2. Output Head / Next-Token Generation (Final Layer)

probs = softmax(logits)

At the very end of the model, the vocabulary-sized logits vector (shape [vocab_size]) is passed through softmax to produce a probability distribution over all tokens in the vocabulary $V$. This distribution is then used for either:

• Greedy decoding: argmax(probs)
• Sampling: categorical draw from probs
• Top-k / Top-p / Temperature: modified sampling strategies (covered in From Logits to the Next Token)

---

Hands-On Exercises & Checked Solutions

Exercise 1: Hand-Computing Softmax on [0.0, 1.0, 2.0]

Problem: Compute softmax probabilities for z = [0.0, 1.0, 2.0] step by step. Explain why the result matches the [1.0, 2.0, 3.0] case from the worked example.

Solution:

Step 1: Exponentials

exp(0.0) = 1.0000
exp(1.0) = 2.7183
exp(2.0) = 7.3891

Step 2: Sum

sum = 1.0000 + 2.7183 + 7.3891 = 11.1074

Step 3: Normalize

p_0 = 1.0000 / 11.1074 = 0.0900
p_1 = 2.7183 / 11.1074 = 0.2447
p_2 = 7.3891 / 11.1074 = 0.6652

Step 4: Verify

0.0900 + 0.2447 + 0.6652 = 0.9999 ≈ 1.0000

Why it matches [1.0, 2.0, 3.0]: The vector [0.0, 1.0, 2.0] is exactly [1.0, 2.0, 3.0] minus 1.0 from every element. Softmax is shift-invariant: subtracting the same constant (here, 1.0) from all logits produces identical probabilities. This is the same principle as the max-subtraction trick!

---

Exercise 2: Numerical Stability Verification

Problem: Given logits z = [-100.0, -98.0, -97.0], apply the max-subtraction trick and compute the final probabilities. Verify the sum is 1.0.

Solution:

Step 1: Find maximum

c = max(z) = -97.0

Step 2: Shift by subtracting c

z' = [-100.0 - (-97.0),  -98.0 - (-97.0),  -97.0 - (-97.0)]
   = [-3.0, -1.0, 0.0]

Step 3: Exponentiate (all in (0, 1])

exp(-3.0) = 0.049787... → 0.0498
exp(-1.0) = 0.367879... → 0.3679
exp(0.0)  = 1.000000... → 1.0000

Step 4: Sum

sum = 0.0498 + 0.3679 + 1.0000 = 1.4177

Step 5: Normalize

p_0 = 0.0498 / 1.4177 = 0.0351
p_1 = 0.3679 / 1.4177 = 0.2595
p_2 = 1.0000 / 1.4177 = 0.7054

Step 6: Verify sum

0.0351 + 0.2595 + 0.7054 = 1.0000

Observation: Even with very negative logits, max-subtraction produces valid probabilities without underflow (all exponentials are $\ge 0.0498$, well above float32 underflow threshold).

---

Exercise 3: Sampling Simulation

Problem: Using the distribution p = [0.10, 0.60, 0.30]:

1. Construct the cumulative intervals for categorical sampling.
2. Simulate two random draws: $r_1 = 0.08$ and $r_2 = 0.85$. Which token does each select?
3. Contrast with greedy argmax decoding.

Solution:

1. Cumulative intervals (CDF):

• Token 0: [0.00, 0.10)
• Token 1: [0.10, 0.70) (0.10 + 0.60 = 0.70)
• Token 2: [0.70, 1.00) (0.70 + 0.30 = 1.00)

2. Simulate draws:

• $r_1 = 0.08$: Falls in [0.00, 0.10) → Token 0 selected
• $r_2 = 0.85$: Falls in [0.70, 1.00) → Token 2 selected

3. Greedy argmax comparison:

Argmax always picks the highest probability: Token 1 (probability 0.60).

• Draw 1: Sampling picked Token 0; Greedy would pick Token 1.
• Draw 2: Sampling picked Token 2; Greedy would pick Token 1.

Sampling introduces variety — Token 1 is still the most likely outcome (~60% of draws), but other tokens appear proportionally to their probabilities. Greedy decoding is a special case of sampling with temperature = 0.

---

Summary Reference Table

Concept	Formula / Rule	Purpose
Probability distribution	$0 \le p_i \le 1$, $\sum p_i = 1$	Valid downstream input (sampler, loss)
Logits	Unbounded $(-\infty, \infty)$	Raw model output before softmax
Softmax	p_i = exp(z_i) / Σ_j exp(z_j)	Map logits → probabilities
exp() effect 1	$\exp(z) > 0$ for all $z$	Guarantees non-negativity
exp() effect 2	Larger $z$ → exponentially larger $\exp(z)$	Amplifies differences (confidence)
Max-subtraction trick	z' = z - max(z)	Prevents float overflow, exact math
Shift invariance	softmax(z + c) = softmax(z)	Justifies the trick
Argmax (greedy)	argmax_i p_i	Deterministic, zero variety
Categorical sampling	Draw r ~ U(0,1), find CDF interval	Stochastic, matches p over many draws
Attention softmax	softmax(Q @ K^T / sqrt(d_k))	Normalize attention scores per query
Output-head softmax	softmax(logits)	Vocabulary distribution for next token

---

Primary Source Citations

• Softmax definition and properties: Goodfellow, Bengio, & Courville (2016), Deep Learning, MIT Press, Section 6.2.2.3 (Softmax Units) — deeplearningbook.org
• Scaled dot-product attention with softmax: Vaswani et al. (2017), Attention Is All You Need, NeurIPS 2017, Section 3.2.1 — arXiv:1706.03762
• Numerical stability / max-subtraction: Standard implementation in PyTorch (torch.nn.functional.softmax), TensorFlow, JAX, and all production inference engines — see e.g., PyTorch source aten/src/ATen/native/SoftMax.cpp

---

What Comes Next

• Next lesson: Normalization, Computation Cost, and Number Formats — RMSNorm, Big-O growth, and the storage trade-offs among FP32, FP16, BF16, FP8, INT8, and INT4.
• Previous lesson: Linear Projections, Transposes, and Similarity
• Track overview: Start Here: LLM Foundations