Softmax: Turning Scores into Weights

The problem softmax solves

In attention, you'll compute a raw compatibility score between a query and each key — a dot product. You'll end up with something like

meaning "token 1 looks most relevant, token 2 somewhat, token 3 least." But scores aren't directly usable as mixing weights. You want to build a weighted average of value vectors, and for that you need weights that are:

1. non-negative (you can't have $-0.4$ of a token), and
2. sum to exactly 1 (so the output stays on the same scale as the inputs, regardless of sequence length).

Softmax is the standard way to convert an arbitrary real-valued vector into exactly that: a probability distribution.

The formula

For a vector $\mathbf{s} = [s_1, s_2, \dots, s_n]$:

Read it as a two-step recipe:

1. Exponentiate each score. This forces everything positive and amplifies differences — $e^x$ grows fast, so a slightly bigger score becomes a much bigger number.
2. Normalize by the total. Dividing by the sum guarantees the outputs add to 1.

The sum-to-1 property is basically free from the construction: if $Z = \sum_j e^{s_j}$, then

The name is "soft max" because the output is a smoothed version of "put all the weight on the largest element." Hard max would give $[1, 0, 0]$; softmax gives a differentiable approximation to it, which is what lets gradients flow during training.

Worked example, by hand

Take $\mathbf{s} = [2, 1, 0]$.

Step 1 — exponentiate:

Step 2 — sum:

Step 3 — divide:

Verify: $0.6652 + 0.2447 + 0.0900 = 0.9999 \approx 1$ (rounding).

So a score gap of 1.0 between the top two entries translates into weights of roughly 67% and 24%. Note the nonlinearity: the scores were evenly spaced ($2, 1, 0$) but the weights are not ($0.67, 0.24, 0.09$) — each step down in score multiplies the weight by about $1/e \approx 0.37$.

Scaling the scores: sharpening and flattening

Now the key insight for understanding attention. Multiply all scores by a constant $\beta$ before applying softmax:

$\beta$ is often written as $1/T$ where $T$ is the temperature. Large $\beta$ (low temperature) = sharp. Small $\beta$ (high temperature) = flat.

Sharpen it: $\beta = 3$, so $\mathbf{s} = [6, 3, 0]$

Almost all the weight is on the first element. The distribution is approaching hard max — nearly one-hot.

Flatten it: $\beta = 0.25$, so $\mathbf{s} = [0.5, 0.25, 0]$

Now it's much closer to uniform $[0.333, 0.333, 0.333]$. As $\beta \to 0$, every score becomes 0, every exponential becomes 1, and you get exactly the uniform distribution — the scores stop mattering at all.

Summary of the three cases, same underlying scores:

| $\beta$ | scores | weights | |---|---|---| | $0.25$ | $[0.5, 0.25, 0]$ | $[0.42, 0.33, 0.25]$ | | $1$ | $[2, 1, 0]$ | $[0.67, 0.24, 0.09]$ | | $3$ | $[6, 3, 0]$ | $[0.95, 0.05, 0.002]$ |

Why this matters for transformers. Attention computes $\text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)$. That $\sqrt{d_k}$ in the denominator is exactly a $\beta < 1$ flattening factor. A dot product of two $d_k$-dimensional vectors with unit-variance components has variance $d_k$, so raw scores grow like $\sqrt{d_k}$ as the model gets wider. Without the division, scores for $d_k = 64$ would be large enough to push softmax into the near-one-hot regime, where the gradient with respect to every score is nearly zero and learning stalls. Dividing by $\sqrt{d_k}$ keeps the score magnitudes roughly constant as you change the head dimension. That's the entire justification — it's a temperature correction, not deep magic.

The same knob shows up again at the output of the model: when sampling the next token, "temperature 0.7" means $\beta = 1/0.7 \approx 1.43$, sharpening the distribution toward the likeliest tokens; temperature 1.5 flattens it and produces more surprising text.

One more property: shift invariance

Adding the same constant $c$ to every score changes nothing:

Check it on our example by subtracting the max, $c = -2$, giving $[0, -1, -2]$:

Identical. This is not a curiosity — it's how softmax is actually implemented. $e^{1000}$ overflows to inf in float32, so every real implementation subtracts the max first:

import numpy as np

def softmax(s, beta=1.0):
    s = beta * np.asarray(s, dtype=np.float64)
    s = s - s.max()              # shift-invariance: safe, changes nothing
    e = np.exp(s)
    return e / e.sum()

w = softmax([2, 1, 0])
print(w, w.sum())                # [0.665 0.245 0.090] 1.0
print(softmax([2, 1, 0], beta=3))    # sharp
print(softmax([2, 1, 0], beta=0.25)) # flat

Contrast this with scale invariance, which softmax does not have — that's the whole point of the previous section. Shifting is free; scaling changes everything.

Also worth knowing

• Only differences matter. Since shifts are free, softmax depends solely on the gaps between scores. A gap of $\Delta$ between two entries means their weight ratio is exactly $e^{\Delta}$. Gap of 1 → ratio $2.7$; gap of 5 → ratio $148$.
• Masking is $-\infty$. In causal attention, positions a token isn't allowed to see get score $-\infty$ (in practice a large negative number like $-10^9$). Then $e^{-\infty} = 0$, the weight vanishes, and the remaining weights still sum to 1 automatically. No special-case code needed.
• Outputs are never exactly 0 or 1. Every unmasked position always gets a nonzero weight. Attention is always a blend, just sometimes a very lopsided one.

Try it yourself

Take the score vector

1. Compute $\text{softmax}(\mathbf{s})$ by hand (calculator is fine) and confirm the four weights sum to 1. Notice what happens to the two tied scores.
2. Without recomputing anything from scratch, predict the weight ratio between the first entry and the last. (Hint: use the gap.) Then check it against your numbers.
3. Recompute with $\beta = 2$ and $\beta = 0.5$. Rank the three distributions from flattest to sharpest and write a one-sentence explanation of why, in the form you'd say it out loud to a colleague.
4. Confirm all of it with the softmax function above, then mask out the second element (set its score to $-10^9$) and verify the remaining three weights still sum to 1.

If you can do step 3 out loud without looking at notes, you have the piece of softmax that attention actually depends on.