Self-attention: every word looks at every other word

Transformer on-ramp · Part 4 of 8

A transformer reads a sentence by letting every word ask every other word one question: “how much should I pay attention to you?” That one idea, self-attention, is what separates it from everything before.

Three roles for every token

For each token, the model builds three vectors from its embedding, using three separate learned projections:

• a query: what am I looking for?
• a key: what do I offer?
• a value: what do I actually contribute if someone attends to me?

To update one token, take its query, compare it with every token’s key, turn the comparison scores into weights that add up to 1, and use those weights to blend the values. In one line:

Attention(Q, K, V) = softmax( Q·K? / ?dk ) · V

Why divide by √d_k?

A score is a dot product of two vectors of length d. If each number has spread 1, the dot product has spread about √d: the bigger the vectors, the bigger the scores. Large scores make softmax almost a hard choice and leave it with almost no gradient to learn from. Dividing by √d keeps the scores in a sensible range at every size. Measured with random vectors (64 keys per query, 2,000 queries):

Without scaling, at d = 1,024 the biggest weight is already 0.95 for random inputs: the layer is making a near-hard choice before it has learned anything. With scaling the picture is the same at every size, so one learning rate works as the model grows.

Try it yourself

import torch
torch.manual_seed(0)
for d in (16, 64, 256, 1024):
    q = torch.randn(2000, d); k = torch.randn(2000, 64, d)   # 2,000 queries, 64 keys each
    s = torch.einsum('qd,qkd->qk', q, k)                    # raw scores
    raw, scaled = s.softmax(-1), (s / d**0.5).softmax(-1)
    print(d, f"spread {s.std():.2f}",
          f"top weight raw {raw.max(-1).values.mean():.3f}",
          f"scaled {scaled.max(-1).values.mean():.3f}")

Open in Colab A small notebook with the tiny worked example and this experiment. Runs free in your browser.

Takeaways

• Each token makes a query, a key and a value from its embedding.
• Query–key scores become weights (softmax); the weights blend the values.
• Dividing by √d_k stops scores from growing with vector size and saturating softmax.
• Without it, at d = 1,024 the top weight is already 0.95 on random inputs.

Next in the series (Part 5): Self-attention from scratch in PyTorch