Attention Mechanism Explained

Attention is a soft dictionary lookup. Each position asks a question, every position advertises what it contains, and the answers get mixed in proportion to how well they match. It replaced recurrence because the lookup is a matmul, so all positions can be processed in parallel.

The idea

Project the input into three roles per position:

Then

• QKᵀ computes every query–key similarity at once, giving a (T, T) score matrix.
• The softmax runs along the key axis, so each row is a distribution over positions that sums to 1.
• Multiplying by V gives each position a weighted average of every value vector.

The 1/√d_k scale is not cosmetic. If query and key entries are independent with unit variance, a dot product of d_k terms has variance d_k. Measured over 20,000 random pairs with d_k = 64: raw dot-product variance 64.59, scaled variance 1.009. Without the scale, logits of magnitude 8 saturate the softmax, where the largest weight goes to 1 and gradients to the others vanish.

Causal masking makes attention autoregressive. Set the scores above the diagonal to −∞ before the softmax, and each position can only see itself and earlier positions. Setting them to 0 instead would let the future leak in at full weight, since e⁰ = 1.

Cost is O(T²d): quadratic in sequence length, linear in width.

Worked example

Two tokens with d_k = 2:

QKᵀ = [[1, 1], [0, 1]], and dividing by √2 gives [[0.7071, 0.7071], [0, 0.7071]].

• Row 1 softmax: (0.5, 0.5), so the output is the midpoint of the two values: (2.0, 3.0).
• Row 2 softmax: (0.3302, 0.6698), so token 2 keeps more of its own value: (2.3395, 3.3395).

Apply a causal mask and row 1 becomes (1.0, 0.0), so its output is exactly V₁ = (1.0, 2.0) — token 1 cannot see token 2 at all. That is the mechanism that makes next-token prediction possible.

In code

import numpy as np

Q = np.array([[1., 0.], [0., 1.]])
K = np.array([[1., 0.], [1., 1.]])
V = np.array([[1., 2.], [3., 4.]])

scores = Q @ K.T / np.sqrt(2)
w = np.exp(scores - scores.max(1, keepdims=True))
w = w / w.sum(1, keepdims=True)
print(np.round(w, 4))        # [[0.5    0.5   ] [0.3302 0.6698]]
print(np.round(w @ V, 4))    # [[2.     3.    ] [2.3395 3.3395]]

mask = np.triu(np.ones((2, 2)), k=1).astype(bool)
w_masked = np.where(mask, -np.inf, scores)
w_masked = np.exp(w_masked - w_masked.max(1, keepdims=True))
w_masked = w_masked / w_masked.sum(1, keepdims=True)
print(np.round(w_masked @ V, 4))    # [[1. 2.] [2.3395 3.3395]]

Check yourself

1. Why divide the scores by √d_k, and what happens to the softmax if you skip it at d_k = 512?
2. The mask sets future scores to −∞ rather than 0. Why does 0 fail?
3. Which axis must the softmax normalize over, and what does each row sum to after masking?

Key takeaways

• Attention weights positions by query–key similarity and returns a weighted average of values.
• The 1/√d_k scale keeps scores in the range where softmax retains usable gradients.
• A causal mask turns a full attention matrix into a strictly backward-looking one.