Learn AI Create the Future — with Pranay
Main site

Attention Mechanics

Scaled dot-product attention

Match every query against every key, turn the matches into shares that add up to one, and mix the values in those proportions - with one division that stops the mix collapsing onto a single winner.

Read these first

On this page 6
  1. The three steps
  2. Why the volume knob is needed
  3. The honest part
  4. Where you have already seen the shape of this
  5. Remember this
  6. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

Scaled dot-product attention scores every word against every other word. It turns those scores into shares, then mixes the words in those proportions.

Think of a crowded room where several friends are all talking to you at once. Your attention is one whole thing, and you have to split it among them. Whoever is saying something relevant gets most of it, and the rest get a sliver.

Now imagine someone turns the room's master volume up far too high. The loudest person drowns out everyone else completely, and you hear nothing but them. Turn it down too far and everyone sounds identical, so you hear a useless mush.

Attention has exactly this volume problem, and it has a fix built into its name. That fix is the word "scaled".

The three steps

Step one: score. Take one word's query and compare it against every word's key. The comparison is a running total of how well the two line up, slot by slot. A big total means a good match.

Step two: share out. Those scores can be any size at all, including negative. They need to become shares of one whole, so they can be used as mixing proportions. The step that does this is called softmax. It turns any list of numbers into positive fractions that add up to one.

Step three: mix. Multiply each word's value by its share, and add them all up. That sum is the output for this word.

   query for "it"
        |
        v
   score against every key   ->  [ 0.1,  4.0,  0.0,  0.2 ]
        |
        v
   turn down the volume      ->  [ 0.0,  1.4,  0.0,  0.1 ]
        |
        v
   softmax into shares       ->  [ 0.14, 0.57, 0.14, 0.15 ]   (adds up to one)
        |
        v
   mix the values in those shares
        |
        v
   one new vector for "it"

Why the volume knob is needed

Here is the honest version, and it is worth reading twice.

The score is a running total across many slots. Wide models use dozens or hundreds of slots per head. More slots means more things being added together. So the totals drift further from zero on their own. Not because the match is better, but because the sum is longer.

Softmax reacts badly to large numbers. Feed it scores that are far apart and it hands almost the whole share to the single biggest one. Attention stops mixing and starts picking a winner.

Worse, a winner-takes-all softmax is nearly flat. Nudging the scores barely changes the shares. And a model learns by watching how the answer changes when you nudge things. With nothing changing, learning grinds to a halt.

So attention divides every score by a fixed number that grows with the number of slots. The totals come back to a sensible size, and mixing survives.

The honest part

This is confusing for almost everyone the first time. The mechanism is not hard; the reason for the division is.

Break it into two true statements and it gets easier. One: wider heads produce bigger raw scores for no meaningful reason. Two: softmax turns big scores into all-or-nothing behaviour. The division cancels the first, which prevents the second.

Where you have already seen the shape of this

  • A search engine ranking pages, then showing the top few rather than one.
  • A recommendation feed blending several of your interests instead of picking one.
  • Photo apps ranking faces by similarity and returning several possible matches.

Remember this

  • Score every query against every key, share them out with softmax, then mix the values.
  • Softmax makes the shares add up to one whole, so the output is always a proper blend.
  • The division by a fixed number keeps scores at a sensible size. The blend then does not collapse onto one winner.

What to learn next

Developer — Code and libraries.

Setup

bash
pip install numpy

The scaling problem, measured rather than asserted

Most explanations assert that scores grow with head width. Here it is, measured.

scaling.py
import numpy as np

rng = np.random.default_rng(0)
np.set_printoptions(precision=3, suppress=True)

def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)    # keeps exp() out of overflow range
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

print("how big do raw dot products get as the head width d grows?")
print(f"{'d':>5} {'std of q.k':>12} {'std of q.k/sqrt(d)':>20}")
for d in (4, 16, 64, 256):
    q = rng.normal(size=(20000, d))
    k = rng.normal(size=(20000, d))
    raw = np.einsum("nd,nd->n", q, k)          # one dot product per row pair
    print(f"{d:>5} {raw.std():>12.2f} {(raw / np.sqrt(d)).std():>20.2f}")

print("\nwhat that does to attention weights over 6 keys (2000 random trials):")
print(f"{'d':>5} {'mean top weight, no scaling':>29} {'with scaling':>14}")
for d in (4, 16, 64, 256):
    q = rng.normal(size=(2000, 1, d))
    K = rng.normal(size=(2000, 6, d))
    logits = np.einsum("nqd,nkd->nqk", q, K)
    hot = softmax(logits).max(axis=-1).mean()
    cool = softmax(logits / np.sqrt(d)).max(axis=-1).mean()
    print(f"{d:>5} {hot:>29.3f} {cool:>14.3f}")

print("\none full attention step, small enough to check by hand:")
Q = np.array([[1., 0.],
              [0., 1.]])
K = np.array([[1., 0.],
              [0., 1.],
              [1., 1.]])
V = np.array([[10., 0.],
              [0., 10.],
              [5., 5.]])
d_k = Q.shape[-1]
logits = Q @ K.T / np.sqrt(d_k)
weights = softmax(logits)
print("logits =\n", logits)
print("weights (each row sums to 1) =\n", weights)
print("row sums =", weights.sum(axis=1))
print("output =\n", weights @ V)
Output
how big do raw dot products get as the head width d grows?
    d   std of q.k   std of q.k/sqrt(d)
    4         2.00                 1.00
   16         4.01                 1.00
   64         7.96                 1.00
  256        16.06                 1.00

what that does to attention weights over 6 keys (2000 random trials):
    d   mean top weight, no scaling   with scaling
    4                         0.574          0.399
   16                         0.776          0.418
   64                         0.890          0.422
  256                         0.946          0.421

one full attention step, small enough to check by hand:
logits =
 [[0.707 0.    0.707]
 [0.    0.707 0.707]]
weights (each row sums to 1) =
 [[0.401 0.198 0.401]
 [0.198 0.401 0.401]]
row sums = [1. 1.]
output =
 [[6.017 3.983]
 [3.983 6.017]]

Reading the output carefully

The first table is the whole argument in four rows. Raw score spread goes 2.00, 4.01, 7.96, 16.06 as d goes 4, 16, 64, 256. Those are the square roots of 4, 16, 64 and 256, to two decimals. Divide by the square root of d and every row reads 1.00. The prediction is not approximately right; it is right.

The second table shows why anyone cares. With no scaling, the mean top attention weight climbs 0.574, 0.776, 0.890, 0.946. A wider head is not paying better attention. It is being pushed toward picking one key and discarding the rest. With scaling, the same four numbers sit at 0.399, 0.418, 0.422, 0.421 — flat, as intended.

The hand-checkable example. Query row one is [1, 0]. It matches key one and key three, but not key two. Both matching logits come out as 0.707, which is one divided by the square root of two. The resulting weights 0.401, 0.198, 0.401 are symmetric. The output [6.017, 3.983] is a genuine blend, not a copy of any single value row.

In PyTorch

python
import torch.nn.functional as F
out = F.scaled_dot_product_attention(q, k, v)      # q, k, v: (batch, heads, tokens, width)

Written against PyTorch 2.5.1, and the signature is unchanged in the current 2.13 documentation:

scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0,
                             is_causal=False, scale=None, enable_gqa=False)

scale defaults to one over the square root of the query width. Passing attn_mask and is_causal=True together raises an error — they are mutually exclusive. Use this function rather than writing the matmuls yourself. It dispatches to a fused kernel that never materialises the full score matrix. That is the difference between fitting in memory and not.

Common mistakes

Dividing by the model width instead of the head width. The scaling constant is built from the width of one head. That is the model width divided by the number of heads. Getting this wrong by a factor of the head count is a quiet bug. It shows up as slow, unstable training.

Softmaxing over the wrong axis. The shares must add to one across keys. That is the last axis of a (tokens, tokens) score matrix. Softmax over the query axis compiles fine and trains to nothing useful. Assert weights.sum(-1) is all ones during development.

Applying the mask after the softmax. Zeroing weights after softmax leaves the rows no longer summing to one. Masks belong on the logits, before softmax, as a large negative number. See causal masking.

Assuming the scale is a hyperparameter to tune. It follows from the width, and the standard value is right for standard initialisation. If your logits still explode, the cause is drifting weights, not the constant — see QK normalisation.

Try it yourself

Change logits / np.sqrt(d_k) to logits / d_k and re-run the second table. Over-dividing flattens the weights toward one over six, which is 0.167. Under-dividing sharpens them toward one. Only the square root holds the middle across every width.

What to learn next

Researcher — Mathematics and papers.

The definition

For $Q \in \mathbb{R}^{T_q \times d_k}$, $K \in \mathbb{R}^{T_k \times d_k}$ and $V \in \mathbb{R}^{T_k \times d_v}$:

$$ \operatorname{Attention}(Q, K, V) = \operatorname{softmax}!\left( \frac{QK^\top}{\sqrt{d_k}} \right) V $$

Here $T_q$ is the number of queries and $T_k$ the number of keys. $d_k$ is the shared query-key width, and $d_v$ the value width. The softmax is applied along the last axis. Each of the $T_q$ rows is then a point on the probability simplex over $T_k$ keys.

Why the denominator is the square root of the width

Assume query and key components are independent with zero mean and unit variance. Then for a single pair,

$$ \mathbb{E}[q \cdot k] = 0, \qquad \operatorname{Var}(q \cdot k) = \sum_{i=1}^{d_k} \operatorname{Var}(q_i k_i) = d_k $$

so the standard deviation is $\sqrt{d_k}$. Dividing by $\sqrt{d_k}$ restores unit variance regardless of width. Vaswani et al. state exactly this in footnote 4 of the original paper.

The consequence for optimisation is the part that matters. The softmax Jacobian is

$$ \frac{\partial p_i}{\partial s_j} = p_i (\delta_{ij} - p_j) $$

with $p = \operatorname{softmax}(s)$ and $\delta_{ij}$ the Kronecker delta. As $p$ approaches a one-hot vector, every entry of this Jacobian approaches zero. Saturated attention passes no gradient to the query and key projections. A head that saturates early tends to stay saturated.

Cost

For a single head over a sequence of length $T$:

  • $QK^\top$: $2 T^2 d_k$ FLOPs, producing a $T \times T$ matrix.
  • Softmax: $O(T^2)$ elementwise work, dominated by the exponential.
  • $AV$: $2 T^2 d_v$ FLOPs.

Memory for the explicit score matrix is $T^2$ per head per layer. That is the binding constraint long before the arithmetic is. Dao et al. (2022), FlashAttention, arXiv:2205.14135, removes it by tiling. The softmax uses running maximum and running sum statistics. The full matrix is never written to high-bandwidth memory. Memory drops to linear in $T$ while the result stays numerically identical to the unfused computation. FlashAttention-2 (arXiv:2307.08691) improves work partitioning. FlashAttention-3 (arXiv:2407.08608) exploits asynchronous copies and low-precision arithmetic on newer hardware.

The trick that makes this possible is online softmax. It is due to Milakov and Gimelshein (2018), arXiv:1805.02867.

Alternatives to the dot product

  • Additive attention (Bahdanau et al., 2015) scores with $v^\top \tanh(W_1 q + W_2 k)$. Comparable quality, but it cannot be expressed as a single matrix multiply. That makes it far slower on modern hardware. This is the entire practical argument for the dot product.
  • Cosine attention normalises both operands, bounding logits to $[-1, 1]$ before a learned scale. Used in Swin Transformer V2 (Liu et al., 2022) for stability at scale. See QK normalisation.
  • Kernelised attention replaces $\exp(q \cdot k)$ with a feature map $\phi(q)^\top \phi(k)$. See Katharopoulos et al. (2020), Linear Transformers, and Choromanski et al. (2021), Performer. The product can then be reassociated, reducing cost to linear in $T$. Quality at scale has been the persistent difficulty.

Softmax is not the only normaliser

The choice of softmax is a design decision, not a derivation. Two documented alternatives:

  • Sparse variants. Sparsemax and $\alpha$-entmax project onto the simplex and produce exact zeros. See Martins and Astudillo (2016) and Peters et al. (2019). A head can then ignore a token outright, rather than assigning it a small weight.
  • Softmax with an extra denominator term lets a row's weights sum to less than one. It is sometimes called softmax-off-by-one. A head can then decline to attend at all. This is related to the attention-sink phenomenon documented by Xiao et al. (2024), arXiv:2309.17453. Models dump surplus attention onto early tokens with no semantic content.

Papers

What to learn next