Attention from scratch, and why it divides by √d

What I built

This is the second module of llm-lab. It adds attention and a full transformer block on top of the autograd engine from the first module: multi-head attention, the causal mask, LayerNorm, the MLP, residual connections, and two kinds of position encoding, sinusoidal and RoPE. None of it needed a new op in the engine, so there's no new backward code anywhere in this module.

Attention itself is a few lines:

def attention(q, k, v, causal=True):
    T, d = q.shape[-2], q.shape[-1]
    att = (q @ swap_last2(k)) * (1.0 / math.sqrt(d))
    if causal:
        att = att + causal_mask(T)
    return softmax(att, axis=-1) @ v

The causal mask is a constant matrix with -inf above the diagonal, added before softmax. exp(-inf) is exactly 0, so a token gets exactly zero weight on anything after it, and exactly zero gradient from it, with no special handling. RoPE's rotation is a fixed matrix too, so it's one matmul.

Proving it's right

There are 14 tests. Every piece and a full block are gradient-checked in float64, the way the first module settled on. Multi-head attention is compared with a slow version that loops over the heads one at a time. LayerNorm is checked against its closed form, and RoPE is checked to depend only on the distance between two tokens, not on where they are.

The leak test changes a later token and checks that the outputs at earlier positions move by exactly 0.0, not "close to 0". My first version of that demo was misleading. Without the mask I expected a visible leak and got 1e-9. The random weights were too large, softmax was already saturated, and every query was looking at a single token. That was the √d problem showing up before I'd got to it.

Why divide by √d?

Every explanation says it's "so softmax doesn't saturate". I wanted to see what that actually means.

A score q·k is a sum of d products, so without any scaling its spread grows like √d: about 3 at d = 8, and 32 at d = 1024. I measured what that does at initialization, with random queries scored against 32 random keys.

• With no scaling, at d = 1024, each query puts 96% of its weight on one token, picked at random before any training.
• The gradient that gets back through softmax is 5.9e-5, against 0.22 with ÷√d. That's about 3,700x smaller in this run. Across seeds the ratio moved between roughly 3,000x and 9,000x, so the exact number isn't worth much, but it was thousands of times smaller every time.
• Dividing by d goes too far the other way. Attention starts out completely uniform, with every token getting the same weight.

It still trains

This is the part that surprised me. I expected the unscaled version to fail to learn, and it didn't.

The task: 16 random token vectors, and a query that's an exact copy of one of them. A single attention head has to learn Wq and Wk so the copy gets the attention. Plain SGD, 300 steps. Without scaling it still reached 99–100% accuracy, as long as the learning rate was small enough. What changes is which learning rates work:

These are the learning rates that reach at least 95% accuracy over the last 50 steps. 0.03 was the smallest I tried, so a range that starts at 0.03 may go lower.

Without √d, the largest learning rate that still trains drops 10x between d = 16 and d = 512. With √d, 0.1 to 3 works at every size.

My first explanation was wrong

I assumed the saturation was what made training fragile, so I tested it. The last row in the table is a control: no scaling, but Wq and Wk start smaller, so the scores start out exactly like ÷√d and nothing is saturated. If saturation were the cause, its range should have stayed put. It shrank exactly like the plain no-scaling version.

The scale factor does two things. It sets how big the scores are at the start, and it sets how far each update moves them. Without √d, the same learning rate moves the scores further as d grows, so the learning rate has to come down. ÷d is never saturated, and its range moves too, just upward.

So the full answer: without √d, attention starts out saturated, and each update moves the scores further as d grows, so the usable learning rate shrinks. √d keeps both independent of d. That's what lets you tune a small model and make the heads bigger without retuning.

This was one task, one seed, one head and plain SGD. Adam rescales gradients and probably hides part of the effect. The next module trains a small GPT with AdamW, so I'll check there.