Scaled Dot-Product Attention in transformers

Module 05 · Transformers · 19 concepts

In depth / 15 lines

Scaled Dot-Product Attention, in more detail

This is the exact arithmetic behind every attention layer, and it is four steps.

First, take the dot product of a query with every key to get raw scores.

The dot product is large when two vectors point the same way, which is the relevance signal.

Second, divide every score by the square root of the key dimension.

That division is the scaling in the name, and it is not cosmetic.

With 512 dimensions the dot products grow large simply because many terms are being summed.

Large values pushed into softmax produce a distribution that is nearly all on one token.

The gradients then vanish and the layer learns almost nothing.

Dividing by the square root of the dimension keeps the scores in a workable range.

Third, apply softmax so the weights are positive and sum to one.

In a decoder a mask is applied first, setting future positions to negative infinity.

Softmax turns those into zero, which is how a token is prevented from seeing ahead.

Fourth, multiply the weights by the value vectors and add them up.

All of it is two matrix multiplications with a softmax in between, which suits a GPU perfectly.

The whole sequence is computed at once, which is why transformers train so much faster than RNNs.

Want the full technical article? Wikipedia ↗

Deep dive

What is underneath Scaled Dot-Product Attention

The formula

One equation carries the entire transformer, and everything else in the architecture is arranged around it.

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

$Q$ : the queries, one row per position: what each position is looking for.
$K$ : the keys, one row per position: what each position offers.
$V$ : the values, one row per position: what each position actually contributes.
$d_k$ : the width of a single query or key vector.

Read it right to left

$QK^{\top}$ is every query dotted with every key, producing an $n \times n$ table of raw compatibility scores.

Dividing by $\sqrt{d_k}$ rescales those scores, for a reason that has its own part below.

Softmax turns each row of that table into weights that are positive and sum to one, so every position distributes exactly one unit of attention across all positions.

Multiplying by $V$ takes, for each position, a weighted average of every value in the sequence.

The output is the same shape as the input: one vector per position, each now a blend of the whole sequence.

Why the square root is there

The dot product of two independent random vectors of width $d_k$ , each component with unit variance, has variance $d_k$ .

\mathrm{Var}(q \cdot k) = \sum_{i=1}^{d_k} \mathrm{Var}(q_i k_i) = d_k

So the scores grow in size as the head gets wider, and softmax on large inputs saturates: one weight approaches 1 and the rest approach 0.

A saturated softmax has almost no gradient, which stops the layer learning at all.

\mathrm{Var}\!\left( \frac{q \cdot k}{\sqrt{d_k}} \right) = 1

Dividing by $\sqrt{d_k}$ puts the variance back to one whatever the width, which keeps the softmax in the range where it has a usable gradient.

The mask

\mathrm{softmax}\!\left( \frac{QK^{\top}}{\sqrt{d_k}} + M \right) V

$M$ : a matrix that is zero where attention is allowed and negative infinity where it is not.

Negative infinity, not zero: it goes in before the softmax, so the exponential sends it to exactly zero afterwards.

In a decoder this is upper-triangular, which is what stops position $i$ reading anything after it.

One pass, in order

Project the input three times to get queries, keys and values.
Multiply queries by transposed keys to get the score table.
Divide every score by the square root of the head width.
Add the mask, then softmax each row.
Multiply those weights by the values.

What it costs

O(n^2 d) \text{ time}, \qquad O(n^2) \text{ memory}

The score table is $n \times n$ , so doubling the sequence length quadruples both.

Everything from FlashAttention to sliding-window attention exists to attack one or other of those two terms.

Test yourself

One question on this card. Every option explains itself once you pick, right or wrong.

Question 1

Why are the dot products divided by the square root of the key dimension?

Transformers

19 concepts in this module, then a graded test with every answer explained.

Study this module →