Softmax and dot products

Attention needs two small tools. The first measures how well a query and a key agree. The second turns those agreement scores into weights. Both are simple, and the rest of the topic depends on them.

Dot products measure agreement

Take a query q = [1, 2, 0] and a key k = [2, 1, 3]:

q · k = 1·2 + 2·1 + 0·3
      = 2 + 2 + 0
      = 4

A key that points the same way as the query gives a big number. A key that points away gives a small or negative one. Attention uses this as a cheap relevance score.

Softmax turns scores into weights

Scores can be any real numbers. Attention needs weights that are positive and add up to 1, so it applies softmax:

weight_i = exp(score_i) / Σ_j exp(score_j)

Each score is pushed through exp, which makes it positive. Then each value is divided by the total, so the results sum to 1. Larger scores get larger shares, but softmax is soft: even a low score keeps some weight.

Take three keys with scores [2.0, 1.0, 0.1]:

exp:   [7.389, 2.718, 1.105]     sum = 11.212
weight: [0.659, 0.242, 0.099]    sums to 1.000

The top score gets about two-thirds of the weight, but the others are not zeroed out.

Subtract the max before exponentiating

Large scores overflow exp. The fix is to subtract the largest score first. This does not change the weights, because the same factor appears on top and bottom.

import numpy as np

def softmax(scores):
    # Subtracting the max leaves the result unchanged but keeps exp() from overflowing.
    e = np.exp(scores - scores.max())
    return e / e.sum()

print(np.round(softmax(np.array([2.0, 1.0, 0.1])), 3))
print(np.round(softmax(np.array([1002.0, 1000.0, 997.0])), 4))

Output:

[0.659 0.242 0.099]
[0.8756 0.1185 0.0059]

Without the subtraction, exp(1002) overflows to inf in float64. The second call still works.

Common mistakes

• Expecting softmax to pick a winner. It spreads weight across all entries. A hard pick would be an argmax, which is not differentiable.
• Forgetting the sum-to-1 rule. The outputs are shares, not raw scores. A score of 8 does not mean an 8 in the output.
• Computing exp on raw scores. Subtract the max first, or large inputs turn into inf and NaN.

Check yourself

1. If all scores are equal, what are the weights for n entries? Each weight is 1/n.
2. Why does softmax of [5, 5] equal softmax of [0, 0]? Adding the same constant to every score cancels out in the division.

Key takeaways

• A dot product is large when two vectors point the same way.
• Softmax makes scores positive and normalizes them to sum to 1.
• Softmax is soft: every entry keeps some weight, unless it is masked out.
• Subtract the max before exp to avoid overflow.

Further learning

Next, see how queries, keys and values use these two tools in queries, keys and values.