0012 — Softmax: The Bridge Between Worlds

Convergence · ~25 minutes · Where linear algebra meets probability

🧩 Puzzle Piece: Softmax appears twice in the transformer. First, inside attention: softmax(Q·Kᵀ/√dₖ) converts raw similarity scores into attention weights (a probability distribution over which tokens to attend to). Second, at the output: softmax(logits) converts the final vector into a probability distribution over the vocabulary. Softmax is the mathematical function that takes the outputs of linear algebra (dot products, matrix multiplications) and turns them into the language of probability. It is the bridge.

You have a vector of raw numbers: [2.3, −0.5, 0.0, 4.1]. You want a probability distribution: numbers between 0 and 1 that sum to 1, preserving the relative ordering (4.1 should get the most probability, −0.5 the least). How? Softmax.

The Softmax Function

softmax(z)_i = exp(z_i) / Σⱼ exp(z_j)

Step by step:

exp(zᵢ): exponentiate each component. This makes all values positive (eˣ > 0 for all real x).
Σⱼ exp(zⱼ): sum all the exponentiated values. This is the normalization constant.
Divide: each exp(zᵢ) by the sum. This ensures all outputs sum to 1.

The result is a valid probability distribution — every entry is non-negative and they sum to exactly 1.

Why It's the Perfect Bridge

Property	Why It Matters	Connected Lessons
Always positive	Probabilities can't be negative (Axiom 1)	0007
Sums to 1	Probabilities must sum to 1 (Axiom 2)	0007
Preserves order	If zᵢ > zⱼ then softmax(z)ᵢ > softmax(z)ⱼ — bigger inputs get bigger probabilities	0003
Differentiable	We can compute gradients through it for backpropagation	0002
Amplifies differences	exp() makes big numbers much bigger, separating signal from noise	0010

Softmax is the function that makes attention work. Q·Kᵀ produces a matrix of raw similarity scores — real numbers that could be anything. Softmax converts each row into a probability distribution: "here's how much attention to pay to each token." Without softmax, attention would just be a weighted sum with arbitrary weights. With softmax, it's a weighted average with probabilistically meaningful weights.

Temperature: Controlling Confidence

With a temperature parameter T:

softmax(z/T)_i = exp(z_i/T) / Σⱼ exp(z_j/T)

T → 0: softmax approaches argmax — all probability on the largest value. Deterministic, "cold."
T = 1: standard softmax. The model's natural confidence.
T → ∞: softmax approaches uniform — all tokens equally likely. Maximum randomness, "hot."

This is the "temperature" slider you see in ChatGPT and other LLM interfaces. Lower temperature = more focused, predictable outputs. Higher temperature = more creative, diverse outputs. You're literally adjusting how peaked the probability distribution is.

The Scaling Factor: √dₖ

In the attention formula, softmax gets Q·Kᵀ/√dₖ — the dot products are divided by √dₖ before softmax. Why?

As dₖ (the dimension of key/query vectors) grows, dot products grow in magnitude. Large magnitudes pushed through softmax produce near-one-hot distributions — the model becomes overconfident, and gradients vanish (the softmax saturates). Dividing by √dₖ keeps the variance of the dot products roughly constant regardless of dimension. This is a subtle but critical detail that makes deep transformers trainable.

The √dₖ scaling is a mathematical fix for a practical problem. Without it, large models with high-dimensional keys would produce degenerate attention patterns. With it, attention works across a wide range of model sizes. This kind of detail — understanding not just what the formula is, but why each piece is there — is the essence of mathematical competence.

Softmax in the Full Architecture

Softmax appears at two critical junctions:

Attention weights: softmax(Q·Kᵀ/√dₖ) — converts similarity matrix to attention probabilities. "Which tokens should I pay attention to?"
Output layer: softmax(W_output · h_final) — converts the final hidden state to next-token probabilities. "What's the most likely next word?"

Both are the same function, applied at different scales. The first operates over context positions (1024–128K). The second operates over vocabulary tokens (~50K). Same math, different purpose.

Check Your Understanding

Q1: Input vector z = [1, 2, 3]. After softmax, which entry gets the most probability?

A. The first entry (z=1) B. The third entry (z=3) — but the distribution won't be [0,0,1] C. All entries get equal probability D. It's impossible to determine

Q2: If you increase the temperature T in softmax(z/T), what happens to the output distribution?

A. It becomes more peaked (higher confidence) B. It becomes more uniform (flatter, more exploratory) C. The probabilities sum to more than 1 D. Nothing — temperature doesn't affect softmax

Q3: Why is the division by √dₖ in attention softmax(Q·Kᵀ/√dₖ) necessary as dₖ grows?

A. It's an arbitrary convention from the original paper B. Dot products grow with dimension; without scaling, softmax saturates to near-one-hot, killing gradients during training C. It makes the math notation look more impressive D. It's needed to make the output sum to 1

📖 Recommended Reading Bishop, Pattern Recognition and Machine Learning — Section 4.2 (Probabilistic Generative Models) for the connection between softmax and categorical distributions.

Softmax is deceptively simple — three lines of math — but it's the single most important bridge function in all of deep learning. It's where vector algebra becomes probability, where raw scores become calibrated confidence, where attention weights become meaningful. Master it.

← 0011: KL Divergence Next: 0013 — The Attention Mechanism →