Scaled dot-product attention

Why divide by √d_k?

Roll one die and you get a number from 1 to 6. Roll 64 dice and add them up, and the total can land much further from its average. Totals of many random numbers spread out more as you add more of them. They don't spread in step with the count, though. The spread grows with the square root of the count. Four times as many dice gives about twice the spread.

A dot product is that kind of total. It multiplies $d_k$ pairs of numbers and adds up the results. In a fresh model those numbers are roughly random, typically about ±1 each. Some are positive and some negative. So the typical size of a score grows like $\sqrt{d_k}$:

• with $d_k = 4$, scores are typically around ±2;
• with $d_k = 64$, around ±8;
• with $d_k = 512$, around ±23.

Big scores cause two problems.

One token takes everything. Softmax turns score gaps into share gaps, and it does this fiercely. Each extra point of score gives a token about 2.7 times as much weight. A gap of 20 points gives it hundreds of millions of times as much. So nearly all the weight lands on one token, and the blend in step 4 becomes a plain copy.

Learning stalls. When softmax is pushed that far, it is saturated. Small changes to the scores barely move the weights, so the slopes that training follows are almost zero. The model can't adjust where it looks.

Dividing every score by $\sqrt{d_k}$ cancels the growth. The typical score goes back to about ±1, however long the keys are.

In plain words: longer keys make bigger scores, by a factor of about $\sqrt{d_k}$. Dividing by $\sqrt{d_k}$ undoes exactly that, so softmax stays gentle and trainable.

The code below tests this. For three key lengths, it makes one random query and ten random keys. It prints how spread out the scores are, and the biggest weight softmax gives any one token. It does this with and without the scaling.

import numpy as np

rng = np.random.default_rng(0)
T = 10
for d_k in [4, 64, 512]:
    q = rng.normal(size=d_k)
    K = rng.normal(size=(T, d_k))
    for label, scores in [("unscaled", K @ q), ("scaled", K @ q / np.sqrt(d_k))]:
        w = np.exp(scores - scores.max())
        w /= w.sum()
        print(f"d_k {d_k:>3}, {label:<8}: score spread (std) {scores.std():6.2f}   largest weight {w.max():.2f}")

"std" is the standard deviation, the typical distance from the average. Without scaling, it grows with the key length: 6.42 at $d_k = 64$ and 22.56 at 512. That's close to $\sqrt{64} = 8$ and $\sqrt{512} \approx 22.6$. One token then grabs most of the weight: 0.95 and 0.75 here. With scaling, the spread stays at 1 or below. The largest weight stays between 0.15 and 0.25, so the attention is shared out and can still learn. The exact numbers depend on the random draw. At $d_k = 4$, for example, this draw happens to give a small spread of 0.63.

Masks: switching some pairs off

Sometimes a token must not look at certain other tokens. The next lesson's causal mask is the main example: when a model writes text, a token can't look at words that come after it. A mask is a table of True and False, one entry per pair of tokens. True means "allowed". False means "switched off".

To switch a pair off, set its score to minus infinity, written $-\infty$, before the softmax. Softmax raises e to each score, and e to the power of minus infinity is exactly 0. So that pair gets a weight of exactly 0. The other weights still add up to 1.

Setting the score to 0 wouldn't work. A score of 0 becomes $e^0 = 1$, which still earns a share of the weight.

import math
scores = [2.0, 1.0, float("-inf")]
exps = [math.exp(s) for s in scores]
print([round(e / sum(exps), 3) for e in exps])

Batched attention in PyTorch

Real models handle several sequences at once, called a batch. The tensors gain an extra dimension at the front. A batch of $B$ sequences, each $T$ tokens long, gives queries of shape (B, T, d_k).

That doesn't change the code. The @ operator multiplies the last two dimensions and repeats across any dimensions in front. k.transpose(-2, -1) swaps the last two dimensions, which is the $K^\top$ in the formula. So the code reads almost word for word like the formula. The function also takes an optional mask, as described above.

import math
import torch

def scaled_dot_product_attention(q, k, v, mask=None):
    d_k = q.shape[-1]
    scores = q @ k.transpose(-2, -1) / math.sqrt(d_k)          # (..., T, T)
    if mask is not None:
        scores = scores.masked_fill(~mask, float("-inf"))       # forbidden pairs get -inf
    weights = torch.softmax(scores, dim=-1)                     # each row sums to 1
    return weights @ v, weights

torch.manual_seed(0)
B, T, d_k = 2, 4, 8
q, k, v = torch.randn(B, T, d_k), torch.randn(B, T, d_k), torch.randn(B, T, d_k)
out, w = scaled_dot_product_attention(q, k, v)
print("output:", tuple(out.shape), "  weights:", tuple(w.shape))

# check one entry against a plain Python loop
b, i = 1, 2
scores = [float((q[b, i] * k[b, j]).sum()) / math.sqrt(d_k) for j in range(T)]
exps = [math.exp(s - max(scores)) for s in scores]
manual = sum(e / sum(exps) * v[b, j] for j, e in enumerate(exps))
print("matches a loop:", torch.allclose(out[b, i], manual, atol=1e-5))

mask = torch.tensor([[True, True, False, False]] * T)          # every token may only see tokens 0 and 1
_, w_masked = scaled_dot_product_attention(q, k, v, mask)
print("weights on masked tokens:", w_masked[0, :, 2:].abs().max().item())

Here's what happened:

• The batch has 2 sequences of 4 tokens, with 8 numbers per query, key and value. The output has shape (2, 4, 8): one blended value per token. The weights have shape (2, 4, 4): one 4 × 4 table per sequence.
• The second part redoes one token's output with a plain Python loop, the same way as the worked example. It matches.
• The mask lets every token see only tokens 0 and 1. The largest weight left on tokens 2 and 3 is 0.0.

masked_fill(~mask, ...) fills in the entries where the mask is False. The ~ flips True and False. (As before, torch is the course's in-browser PyTorch subset. Real PyTorch also accepts -inf in masked_fill exactly like this.)

In real code

Real PyTorch ships this whole recipe as one ready-made function:

import torch.nn.functional as F
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)

is_causal=True switches on the causal mask from the next lesson.

The function is fast because of how it runs on the graphics chips (GPUs) that train models. It uses programs called fused kernels, such as FlashAttention. "Fused" means several steps run as one, with no stops in between. FlashAttention works through the tokens in small tiles. It never stores the full T × T table of weights. For long texts, that table is huge, so skipping it saves a lot of memory and time.

The result is the same as your short function, up to tiny rounding differences. FlashAttention is not an approximation. It's the same arithmetic done in a smarter order.