Softmax

Raise $e$ to each score, then divide by the total. Every weight comes out positive, and they add up to 1.

$e^2 \approx 7.389$, $e^1 \approx 2.718$, $e^0 = 1$. Bigger scores grow much faster: a gap of 1 in score becomes a factor of $e \approx 2.7$.

The total is $7.389 + 2.718 + 1 = 11.107$. The weights keep the order of the scores, and the biggest score gets most of the weight.

The factor $e^c$ cancels. So $[102, 101, 100]$ gives the same weights as $[2, 1, 0]$. Code uses this: it subtracts the largest score first, so $e^{1000}$ never has to be computed (it would overflow to infinity).

$\operatorname{softmax}(z / \tau)$. Dividing by a small $\tau$ stretches the gaps, so the top score takes nearly everything; a large $\tau$ shrinks them toward a uniform spread. Dividing attention scores by $\sqrt{d_k}$ is a temperature of $\sqrt{d_k}$.

$e^{-\infty} = 0$, so that entry gets no weight, and the others are normalized among themselves. This is how the causal mask works.

weights = scores.softmax(dim=-1)    # each row sums to 1

For attention scores (one row per query, one column per key), each query gets its own distribution over the keys.

Use $e^{\ln 2} = 2$. Write the three weights.

Exponentiate to $[2, 1, 2]$, total 5, divide.

• $\operatorname{softmax}(z)_i = e^{z_i}/\sum_j e^{z_j}$: positive, sums to 1, keeps the order.
• Shift-invariant, so subtract the max for safety.
• Temperature: $\operatorname{softmax}(z/\tau)$.
• A score of $-\infty$ gets weight exactly 0.