M3.3 — Attention and Self-Attention

Learning objectives

After this lesson, you should be able to:

• Explain queries, keys and values and where they come from.
• Compute scaled dot-product attention on a small example.
• Explain why the scores are divided by √d.
• Distinguish self-attention from cross-attention.
• Explain what a causal mask does.
• Describe multi-head attention and why attention cost grows quadratically with sequence length.

The one-sentence idea

Attention lets every position in a sequence take a weighted average of the others, with weights from how well their keys match its query.

Match with queries and keys, then read from the values.

The problem attention solves

After M3.2 a model holds a sequence of vectors, one per token. But a token’s meaning depends on its neighbours. In “pick up the cup next to the bowl”, the word “cup” needs to know which cup and where. In an image, a patch showing part of a handle needs to connect to the patch showing the cup. The model needs a way for every element of a sequence to look at the others and take in what is relevant.

Attention does this with a simple recipe: each position scores how relevant every other position is to it, turns the scores into weights with softmax (M3.1), and takes a weighted average of the other positions’ information. The weights depend on the content, which is what makes it flexible.

Queries, keys and values

Each position produces three vectors from its own embedding, using three learned matrices (M1.1):

Query

What this position is looking for. Q = X W_Q

Key

What this position offers to be matched against. K = X W_K

Value

The information this position passes on if it is attended to. V = X W_V

A position’s query is compared with every key by a dot product (a similarity score, as in M1.1). A high score means “this key matches what I am looking for”. The scores are scaled, passed through softmax to become weights, and used to mix the values. In one formula:

Here d is the size of the query and key vectors. Everything in this formula is a matrix multiplication, a softmax or a division, so it is built entirely from operations you already know.

A worked example

To keep the numbers readable, take three tokens with two-dimensional embeddings, and let the three matrices be identity so that Q = K = V = X. A real model learns the matrices, but the mechanics are the same.

import numpy as np

def softmax(z):
    e = np.exp(z - z.max(axis=-1, keepdims=True))     # subtract the max for numerical safety
    return e / e.sum(axis=-1, keepdims=True)

X = np.array([[1.0, 0.0],
              [0.0, 1.0],
              [1.0, 1.0]])
d = X.shape[1]

scores = X @ X.T / np.sqrt(d)          # (3, 3): every token against every token
weights = softmax(scores)              # each row sums to 1
out = weights @ X                      # weighted average of the values

print(scores.round(3))
print(weights.round(3))
print(out.round(3))

[[0.707 0.    0.707]
 [0.    0.707 0.707]
 [0.707 0.707 1.414]]
[[0.401 0.198 0.401]
 [0.198 0.401 0.401]
 [0.248 0.248 0.503]]
[[0.802 0.599]
 [0.599 0.802]
 [0.752 0.752]]

Read row by row. Token 0 is most similar to itself and to token 2, so it gives them weight 0.401 each and token 1 only 0.198. Its output [0.802, 0.599] is the weighted mix of all three value vectors, so it now contains information from tokens it attended to. The third token matches everything, and, being the most similar to itself, it gives itself the largest weight (0.503). The output has the same shape as the input, (3, 2): attention returns one new vector per position, now informed by the rest of the sequence.

This is called self-attention because the queries, keys and values all come from the same sequence. When the queries come from one sequence and the keys and values from another, it is cross-attention, which is a natural way to let, for example, a sequence of action queries read from image and text tokens.

Why divide by the square root of d

Dot products grow as vectors get longer. If scores become large, softmax saturates: one weight approaches 1 and the rest approach 0, and gradients through it nearly vanish, so training stalls. Dividing by √d keeps the scores at a stable scale. The effect of scale on softmax is easy to see:

print(softmax(np.array([1.0, 2.0, 3.0])).round(3))
print(softmax(np.array([10.0, 20.0, 30.0])).round(3))

[0.09  0.245 0.665]
[0. 0. 1.]

The same ordering of scores gives a soft blend in the first case and a hard choice in the second. The scaling keeps attention in the soft regime at the start of training, where it can still learn.

Masks: looking only at the past

A model that generates text or actions one step at a time must not peek at the future, or it could cheat during training. A causal mask sets the scores for later positions to minus infinity before the softmax, so their weights become exactly 0:

mask = np.triu(np.ones((3, 3), dtype=bool), 1)     # True above the diagonal: the future
masked = np.where(mask, -np.inf, scores)
print(softmax(masked).round(3))

[[1.    0.    0.   ]
 [0.33  0.67  0.   ]
 [0.248 0.248 0.503]]

The first token can only see itself, the second can see the first two, and the last sees all three. The weights are lower-triangular. A model that reads a full input at once, such as one that encodes an image and a question, does not need this mask.

Multi-head attention and cost

One attention pattern may not be enough: a position may need to track what something is, where it is, and what it is next to. Multi-head attention splits the embedding into several smaller pieces, runs attention on each independently with its own learned matrices, then concatenates the results and mixes them with a final matrix. Each head can learn to look for something different.

rng = np.random.default_rng(0)
T_, d_model, heads = 3, 8, 2
d_head = d_model // heads

X8 = rng.normal(size=(T_, d_model))
Wq, Wk, Wv = (rng.normal(size=(d_model, d_model)) for _ in range(3))

def split(M):                                       # (T, d_model) -> (heads, T, d_head)
    return M.reshape(T_, heads, d_head).transpose(1, 0, 2)

Q, K, V = split(X8 @ Wq), split(X8 @ Wk), split(X8 @ Wv)
w = softmax(Q @ K.transpose(0, 2, 1) / np.sqrt(d_head))     # (heads, T, T)
out = (w @ V).transpose(1, 0, 2).reshape(T_, d_model)       # concatenate the heads
print(w.shape, out.shape)

(2, 3, 3) (3, 8)

Each head produced its own 3 × 3 table of weights, and the concatenated output is back to (3, 8), the same shape as the input, so attention blocks can be stacked.

One cost to remember: the weight table is T × T, so doubling the sequence length quadruples its size. A sequence of 196 image patches needs 38,416 weights per head, and 1,000 tokens need a million. This is why long inputs, such as many camera frames, are expensive and why efficient attention is an active research area.

import torch
import torch.nn.functional as F

torch.manual_seed(0)
q, k, v = (torch.randn(1, 2, 3, 4) for _ in range(3))      # (batch, heads, tokens, d_head)

manual = torch.softmax(q @ k.transpose(-2, -1) / 4 ** 0.5, dim=-1) @ v
print(torch.allclose(F.scaled_dot_product_attention(q, k, v), manual, atol=1e-6))

causal = torch.tril(torch.ones(3, 3, dtype=torch.bool))
manual_c = torch.softmax((q @ k.transpose(-2, -1) / 4 ** 0.5).masked_fill(~causal, float("-inf")), dim=-1) @ v
print(torch.allclose(F.scaled_dot_product_attention(q, k, v, is_causal=True), manual_c, atol=1e-6))

True
True

Common misconceptions

Engineering takeaways

1. Each position makes a query, a key and a value with three learned matrices.
2. Attention is softmax(QKᵀ/√d) V.
3. Scores are dot-product similarities, and softmax turns each row into weights that sum to 1.
4. Dividing by √d keeps softmax from saturating.
5. Self-attention uses one sequence. Cross-attention takes queries from one sequence and keys and values from another.
6. A causal mask sets future positions to −∞ so their weight is 0.
7. Multi-head attention runs several attentions in parallel and concatenates them.
8. The weight table is T × T, so cost grows quadratically with length.

Attention is a content-based lookup: ask a question with the query, match it against keys, and read the values.

Knowledge check

Try each question first, then reveal the explanation.

What’s next

You now have the core operation. Next we see how attention is combined with other layers into the transformer blocks that are stacked to make full models.

Back to the VLA roadmap prerequisites