Sequences Before Attention ended with a recurrent network that could not reverse 40 digits until it was given a way to look back at every input position directly. That way is attention: each position’s new representation is a weighted average of every position’s, with weights computed from the data. The Transformer (Vaswani et al., 2017) is built from attention alone, used in three patterns: self-attention, masked (causal) self-attention, and cross-attention.
The three are one computation. This post derives it from a weighted average, works each pattern through on matrices small enough to check by hand, and then looks at what the patterns do inside real models: a decoder-only language model and an encoder-decoder translation model. The Residual Stream shows where attention sits in a full Transformer block.
The three patterns. Each is the same weighted average; what changes is which sequence supplies the queries, keys and values, and which cells a mask blanks.
- Self, masked and cross-attention are one computation, softmax(QKᵀ/√d)V. On the worked examples, one 6-line function produces all three.
- Dividing by √d keeps the softmax trainable: unscaled, the largest of 16 weights averages 0.96 at d = 1,024 and the softmax’s gradient shrinks to 0.048; scaled, about 0.24 and 0.30 at every width.
- The causal mask hides the future exactly: in Qwen2.5-1.5B-Instruct, changing the last of 17 tokens moved no earlier position’s hidden state at all, 0.0, in any of 28 layers.
- Heads learn patterns, and cross-attention learns alignment: one Qwen head puts 0.78 of its weight on the previous token, on average; a translation model’s cross-attention pairs “Katze” with “cat” and “Küche” with “kitchen”.
- The cost grows as the square of the length: from 2,048 tokens up, each doubling took about 4 times as long.
Table of contents
Open Table of contents
From a weighted average to Q, K and V
Start with the simplest version of “every position looks at every other”. Each input token has a vector , and its new representation is a weighted average of all of them:
The weights should say how relevant token is to token . A dot product measures that, large when two vectors point the same way, and a softmax turns a row of scores into non-negative weights that sum to 1 (Cross-Entropy and the Softmax Gradient derives it):
This already has the shape of attention, but one vector does three jobs: what token is looking for, what token offers to be matched, and what token contributes once matched. Giving each job its own learned projection separates them:
- the query , what token is looking for;
- the key , what token offers to be matched against;
- the value , what token contributes.
Replacing the dot product with and the averaged with gives scaled dot-product attention:
where is the width of a query and a key. Every pattern in this post is this formula; the patterns differ only in which sequence , and are projected from, and in a mask added before the softmax.
Why divide by √d
Suppose the entries of a query and a key are independent, with mean 0 and variance 1. Their dot product is a sum of products , each with mean 0 and variance 1, so
A score’s typical size grows with the width for no reason but the number of terms. The softmax is sensitive to exactly that: scaling every score up pushes the output toward a single 1 and the rest 0, and its gradient, the matrix , toward zero, so a head would stop learning. Dividing by restores unit variance at every width. The check, over 20,000 random pairs per width, and 2,000 random softmaxes over 16 keys:
for d in (4, 16, 64, 256, 1024):
q, k = rng.standard_normal((20_000, d)), rng.standard_normal((20_000, d))
dots = (q * k).sum(1)
rows = {}
for name, s in (("raw", 1.0), ("scaled", np.sqrt(d))):
top, grad = [], []
for _ in range(2_000):
z = np.einsum("d,dk->k", rng.standard_normal(d), rng.standard_normal((d, 16))) / s # 1 query, 16 keys
p = np.exp(z - z.max()); p /= p.sum()
top.append(p.max())
grad.append(np.linalg.norm(np.diag(p) - np.outer(p, p)))
rows[name] = np.mean(top), np.mean(grad)
print(f"{d:4d} {dots.std():10.2f} {(dots / np.sqrt(d)).std():18.2f} "
f"{rows['raw'][0]:20.2f} {rows['scaled'][0]:6.2f} {rows['raw'][1]:6.3f} {rows['scaled'][1]:6.3f}")
d std of q.k std of q.k/sqrt(d) softmax over 16 keys: largest weight gradient norm
raw scaled raw scaled
4 2.00 1.00 0.42 0.23 0.325 0.298
16 4.02 1.00 0.68 0.24 0.279 0.303
64 7.96 0.99 0.84 0.25 0.179 0.304
256 15.98 1.00 0.93 0.25 0.096 0.305
1024 32.51 1.02 0.96 0.24 0.048 0.305
Unscaled, wider heads make the softmax nearly one-hot and its gradient small. Divided by √d, both stay put.
Data: transformer-attention-patterns/scale.pyThe spread of is 2.00, 4.02, 7.96, 15.98 and 32.51 at widths 4 to 1,024, close to , and 1.00 to 1.02 after dividing. Unscaled, the largest of 16 weights climbs from 0.42 to 0.96 and the gradient’s size falls from 0.325 to 0.048, almost seven times smaller; scaled, both are the same at every width.
Self-attention, worked through
Every number below comes from one function, applied three times:
def attention(Q, K, V, mask=None):
scores = Q @ K.T / np.sqrt(K.shape[-1])
if mask is not None:
scores = scores + mask # -inf where a query may not look
weights = np.exp(scores - scores.max(axis=-1, keepdims=True))
weights /= weights.sum(axis=-1, keepdims=True) # softmax, row by row
return scores, weights, weights @ V
Take two tokens with 2-dimensional embeddings, and three projection matrices (in a real model they are learned; here they are chosen small):
Since is the identity, , and . and must have the same width for to exist; ‘s width sets the output’s. The script prints the scores, the weights and the output:
self-attention, 2 tokens
scores QK^T/sqrt(d):
[[0.000 0.707]
[0.707 0.707]]
weights A:
[[0.330 0.670]
[0.500 0.500]]
output AV:
[[0.330 1.000]
[0.500 1.000]]
Row 1 of the output is 33% of token 1’s value plus 67% of token 2’s: token 2 mattered more to token 1 than token 1 did to itself, because of how token 1’s query lined up with the two keys. The softmax turned scores of 0 and 0.707 into weights of 0.330 and 0.670.
Multi-head attention
One set of projections gives one way of relating tokens. Multi-head attention runs of them in parallel, each with smaller projections, then concatenates the heads’ outputs and mixes them with one more learned matrix :
Qwen2.5-1.5B, used below, has 12 query heads of width 128 in each of its 28 layers. Its keys and values come from only 2 heads, each shared by 6 query heads (grouped-query attention), which makes its KV cache six times smaller (The KV Cache).
Masked (causal) self-attention
A decoder generates text one token at a time, and it is trained to predict each next token from the ones before it. If token 2 could attend to token 3, training would let it read the answer. So in decoder self-attention each token must be blind to every later one.
The mask does this before the softmax. A matrix is added to the scores, 0 where a query may look and where it may not; since , the softmax gives those positions exactly zero weight, and the remaining weights still sum to 1:
With a third token added, , and the same projections:
X3 = np.array([[1, 0], [0, 1], [1, -1]], dtype=float)
causal = np.triu(np.full((3, 3), -np.inf), k=1) # -inf above the diagonal, 0 elsewhere
show("causal self-attention, 3 tokens", *attention(X3 @ W_Q, X3 @ W_K, X3 @ W_V, causal))
causal self-attention, 3 tokens
scores QK^T/sqrt(d):
[[ 0.000 -inf -inf]
[ 0.707 0.707 -inf]
[-0.707 0.000 -0.707]]
weights A:
[[1.000 0.000 0.000]
[0.500 0.500 0.000]
[0.248 0.503 0.248]]
output AV:
[[1.000 1.000]
[0.500 1.000]
[0.497 0.752]]
The mask blanks the future before the softmax, so each row’s weights spread over itself and earlier tokens only, and still sum to 1.
Data: transformer-attention-patterns/attention.pyThe weights are lower-triangular. Token 1 can only attend to itself, so its output is its own value; token 2 splits its weight between tokens 1 and 2; only the last token sees the whole sequence. A weight above the diagonal in a decoder would mean the future leaking into the present.
Cross-attention
In self-attention, masked or not, , and come from one sequence. Cross-attention takes the queries from one sequence and the keys and values from another. In an encoder-decoder model the decoder’s states supply the queries and the encoder’s output supplies the keys and values, so each token being generated can look at the whole source:
With the three tokens above as the encoder’s output and two decoder states :
cross-attention, 2 decoder tokens into 3 encoder tokens
scores QK^T/sqrt(d):
[[ 0.354 0.707 -0.354]
[ 0.707 0.000 0.707]]
weights A:
[[0.343 0.488 0.169]
[0.401 0.198 0.401]]
output AV:
[[0.512 0.831]
[0.802 0.599]]
The scores are 2 by 3, two decoder tokens against three encoder tokens: cross-attention does not need the two sequences to have the same length, which is what translation needs. And there is no mask, since the source is complete before the decoder starts.
What real models learn
The worked examples use chosen weights. In a trained model the patterns come from data. Two checks on Qwen2.5-1.5B-Instruct, a decoder-only model: does its mask really hide the future, and what do its heads attend to?
ids = tok(text, return_tensors="pt").input_ids.to(DEVICE)
other = ids.clone()
other[0, -1] = tok(" hot").input_ids[0] # a different last token
with torch.no_grad():
a = model(ids, output_hidden_states=True, output_attentions=True)
b = model(other, output_hidden_states=True)
n = ids.shape[1]
before = max((x[0, :-1] - y[0, :-1]).abs().max().item() for x, y in zip(a.hidden_states, b.hidden_states))
last = (a.hidden_states[-1][0, -1] - b.hidden_states[-1][0, -1]).abs().max().item()
print(f"{n} tokens; changing the last one moves earlier positions' hidden states by at most {before:.1e} "
f"(in every layer), and the last position's by {last:.1f}")
17 tokens; changing the last one moves earlier positions' hidden states by at most 0.0e+00 (in every layer), and the last position's by 60.9
largest attention weight above the diagonal, any layer or head: 0.0e+00
average over all heads: weight on the previous token 0.13, on the first token 0.54
strongest previous token head: layer 1, head 4, average weight 0.78
strongest first token head: layer 14, head 5, average weight 1.00
Changing the last of 17 tokens moved the last position’s hidden state by up to 60.9 and every earlier position’s by exactly 0, in all 28 layers: no information flows backward. No head in any layer put any weight above the diagonal.
Two of Qwen2.5-1.5B-Instruct’s 336 heads. One looks one token back; the other puts nearly all its weight on the first token.
Data: transformer-attention-patterns/causal.pyThe heads specialize. The strongest previous-token head, layer 1 head 4, puts 0.78 of a token’s weight on the token just before it on average, a building block for copying patterns from earlier in the text. Across all heads, the first token gets 0.54 of the weight on average, where spreading weight evenly over the visible tokens would give it about 0.15, and in layer 14 head 5 it gets 1.00. A head with nothing useful to read still has to put its weight somewhere, since the weights sum to 1, and many put it on the first token, which every position can see; this is called an attention sink.
Cross-attention, in Marian’s English-to-German translation model, an encoder-decoder with 6 layers on each side:
with torch.no_grad():
out = model.generate(**src, num_beams=1, do_sample=False, max_new_tokens=40)
a = model(**src, decoder_input_ids=out[:, :-1], output_attentions=True)
en = tok.convert_ids_to_tokens(src.input_ids[0])
de = tok.convert_ids_to_tokens(out[0, 1:]) # the token each decoder step produced
English: The black cat sleeps on the warm mat in the kitchen.
German: Die schwarze Katze schläft auf der warmen Matte in der Küche.
14 German tokens attend over 14 English tokens (including </s>)
layer 0: each German token's most-weighted English token: ▁Die-></s> ▁schwarze->▁The ▁Katze-></s> ▁schläft-></s> ▁auf-></s> ▁der->▁on ▁warmen-></s> ▁Matt-></s> e-></s> ▁in-></s> ▁der->▁in ▁Küche-></s> .-></s> </s>->.
layer 1: each German token's most-weighted English token: ▁Die-></s> ▁schwarze->▁The ▁Katze->▁black ▁schläft->▁cat ▁auf->▁sleep ▁der->▁on ▁warmen->▁the ▁Matt-></s> e-></s> ▁in-></s> ▁der->▁in ▁Küche->▁the .-></s> </s>->.
layer 2: each German token's most-weighted English token: ▁Die-></s> ▁schwarze->▁The ▁Katze->▁black ▁schläft->▁cat ▁auf-></s> ▁der->▁on ▁warmen->▁the ▁Matt-></s> e-></s> ▁in->▁mat ▁der->▁in ▁Küche->▁kitchen .-></s> </s>->.
layer 3: each German token's most-weighted English token: ▁Die->. ▁schwarze->▁The ▁Katze->▁black ▁schläft->▁cat ▁auf->▁sleep ▁der->▁on ▁warmen->▁the ▁Matt->▁warm e->▁mat ▁in->▁mat ▁der->▁in ▁Küche-></s> .-></s> </s>-></s>
layer 4: each German token's most-weighted English token: ▁Die->▁The ▁schwarze->▁black ▁Katze->▁cat ▁schläft->▁sleep ▁auf->▁on ▁der->▁the ▁warmen->▁warm ▁Matt->▁mat e-></s> ▁in->▁in ▁der-></s> ▁Küche->▁kitchen .-></s> </s>-></s>
layer 5: each German token's most-weighted English token: ▁Die-></s> ▁schwarze->▁black ▁Katze->▁cat ▁schläft->▁sleep ▁auf-></s> ▁der->▁warm ▁warmen->▁warm ▁Matt->▁mat e-></s> ▁in-></s> ▁der-></s> ▁Küche->▁kitchen .-></s> </s>-></s>
layer 4 is closest to a word-for-word diagonal
Each German token, as it is generated, weights the English tokens it translates. One sentence, one layer chosen on it, the mean of 8 heads.
Data: transformer-attention-patterns/cross.pyEach row is the step that produced one German token. In layer 4, the step that produced “Katze” weights “cat” most, “schläft” weights “sleep”, “warmen” weights “warm” and “Küche” weights “kitchen”: a word alignment. This model was trained with guided alignment: external word alignments supervised the cross-attention of one head in one layer (by Marian’s default, the last, layer 5), so the alignment is not learned from translation alone; layer 4, shown here, was not supervised directly. Other layers are less clean. Layer 0 puts most weight on the end-of-sentence token, another sink, and layers 1 to 3 mostly run one token behind, a row reading the English word that the previous German token translates.
What attention does not provide
Two properties of the formula need patching from outside it, and one choice in it is not the only option.
Attention has no sense of order. treats its input as a set of rows. Shuffle the rows of and the rows of the output shuffle the same way, and nothing else changes:
6 random tokens, rows shuffled by [5 4 2 1 3 0]: output rows shuffled the same way, largest difference 4.4e-16
The difference, , is rounding. “The cat sat on the mat” and “the mat sat on the cat” would give the same vectors for each word. A Transformer is given positions separately, by adding a positional encoding to the input or, in most current models including Qwen, by rotating the queries and keys by their position (Where Am I? Position in Attention, from Sinusoids to RoPE derives both). A causal mask also leaks some position, since a token can count how many tokens it sees (Haviv et al., 2022, trained models with no positional encoding at all that still learned position this way).
The cost grows as the square of the length. for tokens is an matrix per head. One layer of causal attention at Qwen2.5-1.5B’s shapes, written out as above, in 16-bit on an Apple M4 laptop GPU:
def attend(q, k, v):
k, v = k.repeat_interleave(H // KV, 0), v.repeat_interleave(H // KV, 0)
n = q.shape[1]
scores = q @ k.transpose(-1, -2) / D ** 0.5 # [heads, n, n]
scores = scores.masked_fill(torch.ones(n, n, dtype=torch.bool, device=q.device).triu(1), float("-inf"))
return F.softmax(scores, dim=-1) @ v
tokens time (ms) x previous score matrices (MiB)
512 2.00 6
1024 6.13 3.1 24
2048 17.64 2.9 96
4096 70.37 4.0 384
8192 289.81 4.1 1536
From 2,048 tokens up, doubling the length quadruples the time. The score matrices for 8,192 tokens take 1,536 MiB in one layer.
Data: transformer-attention-patterns/cost.pyFrom 2,048 tokens up each doubling took about 4 times as long (4.0 and 4.1); below that, fixed costs blur the ratio. The score matrices grow the same way, to 1,536 MiB at 8,192 tokens for one layer (one 16-bit matrix per head; the softmax’s output needs as much again). Three later posts deal with this: The KV Cache avoids recomputing past keys and values during generation, FlashAttention computes the same result without ever storing the score matrix, and Attention Alternatives changes the mechanism to escape the square.
Dot products are not the only way to score. Bahdanau et al. (2014), who introduced attention for translation, scored each query-key pair with a small neural network (additive attention). Vaswani et al. chose the dot product because it “is much faster and more space-efficient in practice, since it can be implemented using highly optimized matrix multiplication code”. Citing Britz et al. (2017), they note that the two “perform similarly” for small , while “additive attention outperforms dot product attention without scaling for larger values of ”, and give their reason for as a suspicion: “We suspect that for large values of , the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients”.
Three families of models
How the three patterns are combined defines a model’s family:
- Encoder-only models, such as BERT (Devlin et al., 2019), use unmasked self-attention. Every token sees the whole input in both directions, which suits understanding a fixed text: classifying it or embedding it (Embeddings and Contrastive Learning). They cannot generate left to right.
- Decoder-only models, such as GPT and Qwen, use masked self-attention alone. The same model scores text and generates it, one token at a time. Almost every large language model today is one.
- Encoder-decoder models, such as the original Transformer, T5 (Raffel et al., 2020) and Marian, use all three: the encoder reads the input with self-attention, and each decoder layer applies masked self-attention, then cross-attention into the encoder’s output. They suit tasks with a clear input and output, such as translation.
The projections , , and (and their biases, where a model has them) are attention’s only learned parameters; everything else is fixed arithmetic.
Where this leads
- Where Am I? Position in Attention, from Sinusoids to RoPE gives attention the order it lacks.
- The Residual Stream puts attention into a full Transformer block and follows a token to a prediction.
- The KV Cache and FlashAttention make the quadratic cost affordable in practice.
Takeaways
- Self, masked and cross-attention are one computation, : a weighted average of values, weighted by how well each query matches each key. Only where , and come from, and the mask, change.
- Dividing by keeps scores at unit variance. Without it, at the softmax over 16 keys put 0.96 on one key and its gradient fell to 0.048, against 0.305 scaled.
- The causal mask adds above the diagonal before the softmax. In Qwen2.5-1.5B-Instruct it hides the future exactly: changing the last token moved no earlier position at all.
- Trained heads specialize: a previous-token head (0.78 of its weight one token back, on average), heads that park their weight on the first token, and cross-attention that aligns German words with the English words they translate.
- Attention has no sense of order, so positions are added separately, and its cost grows as the square of the length: each doubling from 2,048 tokens took about 4 times as long.
References
Start with Vaswani et al. §3.2 for scaled dot-product and multi-head attention, then Bahdanau et al. for where attention came from.
- Vaswani et al., Attention Is All You Need (NeurIPS 2017): the Transformer, scaled dot-product attention and the reason for .
- Bahdanau, Cho and Bengio, Neural Machine Translation by Jointly Learning to Align and Translate (2014): attention, introduced as additive scoring over an RNN encoder’s states.
- Devlin, Chang, Lee and Toutanova, BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding (NAACL 2019): the encoder-only family.
- Raffel et al., Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer (JMLR 2020): T5, an encoder-decoder.
- Xiao et al., Efficient Streaming Language Models with Attention Sinks (ICLR 2024): heads that park their weight on the first tokens, and why.
- Olsson et al., In-context Learning and Induction Heads (2022): previous-token heads as a building block of copying.
- Haviv et al., Transformer Language Models without Positional Encodings Still Learn Positional Information (Findings of EMNLP 2022): how a causal mask leaks position.
- Junczys-Dowmunt et al., Marian: Fast Neural Machine Translation in C++ (ACL 2018), and Tiedemann and Thottingal, OPUS-MT (EAMT 2020): the translation model used above; its OPUS-MT-train model notes list it as “transformer-align”.
- Britz, Goldie, Luong and Le, Massive Exploration of Neural Machine Translation Architectures (EMNLP 2017): the comparison of additive and dot-product scoring that Vaswani et al. cite.