Multi-head attention and positional encoding

Last lesson wrote one attention head. To build a GPT with it, two things are still missing: multiple heads, and position information.

python heads_positions.py

This lesson's experiments use the model in gpt.py, but need no training; they look at properties of the structure itself.

Multiple heads: looking at several things at once

The information a character needs from earlier characters is often of more than one kind. When writing "尽" (end) in "白日依山尽" ("the white sun sinks behind the mountains"), the model may need to know at once: what scenery came before (白日, sun; 山, mountains), which character this is (in a five-character poem the fifth character rhymes or ends the line), and which words the previous line used.

One attention head has only one set of q and k, so it can score by only one standard and produce one set of weights. Multi-head attention cuts the vector into several segments, runs attention on each separately, and joins the results back together. Each head has its own q, k and v, and can learn to attend to different things.

== 2. 多头注意力：32 维的向量切成 4 个头，每个头 8 维
  q 的形状 (1, 5, 32)
  拆成多个头之后 (1, 4, 5, 8)：(句子数, 头数, 字数, 每个头的维度)
  注意力分数 (1, 4, 5, 5)：每个头都有自己的一张 5×5 的表
  一个注意力层的参数：qkv 3168 个，proj 1056 个，头数多少都不影响这个数

In gpt.py, implementing it just takes a few more reshapes:

q, k, v = self.qkv(x).split(C, dim=2)
# 拆成多个头：(B, T, C) -> (B, 头数, T, 每个头的维度)
q, k, v = (t.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) for t in (q, k, v))

self.qkv is one linear layer that computes q, k and v in one go, which split then separates. The 32-dimensional vector is then viewed as 4 segments of 8 dimensions each, and the "head" dimension is moved forward. The scoring, masking, softmax and weighted sum after that are exactly as in the last lesson, except PyTorch's matrix multiplication computes each head separately and automatically, all 4 heads at once.

Afterwards, the 4 heads' results are joined back into 32 dimensions and passed through a linear layer, proj, to mix them, so the information each head gathered can combine:

out = out.transpose(1, 2).contiguous().view(B, T, C)  # 各个头拼回去
return self.proj(out)

Note the last line of output: the number of parameters doesn't depend on the number of heads. 4 heads of 8 dimensions and 1 head of 32 dimensions use matrices of the same size, just split differently. Multiple heads add no parameters and almost no computation, yet let the model learn several different ways of attending at once.

Our GPT uses 4 heads of 32 dimensions each (128 dimensions in total). LLMs usually have dozens of heads, each of 64 or 128 dimensions.

Attention can't tell order

Look back at last lesson's attention formula: scoring looks only at the content of q and k, the weighted sum only at the content of v, and nowhere is "which position this character is in" used.

That means, to attention, "白日依山尽" and "山依日白尽" may be no different.

The experiment: have the model read these two lines, both ending in "尽", with the first four characters in a different order. Compare how much the model's output at the last position (its scores for the next character) differs:

model.pos_emb.weight.zero_()  # 把位置嵌入清零，等于没有位置信息
a = model(torch.tensor([tok.encode("白日依山尽")]))[0][0, -1]
b = model(torch.tensor([tok.encode("山依日白尽")]))[0][0, -1]

== 1. '白日依山尽' 和 '山依日白尽'，最后一个位置的输出差多少
  1 层：没有位置嵌入 3.7e-08，有位置嵌入 2.8e-03
  2 层：没有位置嵌入 1.9e-02，有位置嵌入 1.5e-02

Look first at the first line, the 1-layer model. Without position information, the two lines' outputs differ by 3.7e-08, nothing but floating-point rounding error: exactly the same. For the last position, the same five characters are in front of it, just in a different order. Attention scores them and takes the weighted sum, and the result doesn't depend on order.

For a language model that's fatal. "白日依山尽" is poetry, "山依日白尽" is nonsense; "我打你" (I hit you) and "你打我" (you hit me) mean opposite things. The model must know the order.

The unexpected second line

The second line is the 2-layer model. By the reasoning above, without position embeddings the two lines' outputs should also be the same. But the experiment shows a difference of 1.9e-02, about as large as the 1.5e-02 with position embeddings.

I didn't expect this when I wrote the experiment either. The cause is the causal mask:

• In layer 1, each position can see only itself and the characters before it. In "白日依山尽", position 2 sees "白日"; in "山依日白尽", position 2 sees "山依". So layer 1's outputs at positions 2, 3 and 4 differ between the two lines.
• In layer 2, the last position looks at layer 1's outputs at the earlier positions, and those outputs already differ.

In other words, the causal mask itself leaks order: how many characters a position can see indirectly tells it which position it's in. With two or more layers, the model can extract some position information from that.

This isn't a quirk of our code. A 2022 paper studied exactly this (Haviv et al., "Transformer Language Models without Positional Encodings Still Learn Positional Information"), finding that causal language models with no positional encoding at all still learn position information, performing only a little worse than those with positional encoding.

Still, position information "guessed" through the mask is indirect and fuzzy. Real models all add position information explicitly.

Position embeddings

The simplest approach, and the one GPT-2 uses: give each position a vector too, and add it to the character's vector.

self.pos_emb = nn.Embedding(cfg.block_size, cfg.n_embd)
...
pos = torch.arange(start, start + T, device=idx.device)
x = self.drop(self.tok_emb(idx) + self.pos_emb(pos))

pos_emb is a table of 128 rows: row 0 is the vector for position 0, row 1 for position 1, and so on. These vectors are learned in training like any other parameters.

After adding, the same character has a different vector in different positions. "白" in position 1 and in position 4 gets two different vectors, so attention can tell the order apart. In the experiment above, adding position embeddings to the 1-layer model changed the difference between the two lines' outputs from 3.7e-08 to 2.8e-03: the model can tell the two lines apart (this model hasn't been trained, so the difference is small and not yet meaningful).

This approach has one limitation: the table has only 128 rows, so the model can handle at most 128 positions. For positions not seen in training, it has no vector.

Rotary position embedding

As of September 2026, most mainstream open-source LLMs (Llama, Qwen, DeepSeek and others) don't use GPT-2-style position embeddings, but rotary position embedding (RoPE, from the 2021 RoFormer paper).

The idea: instead of adding the position to the character's vector, "rotate" q and k by an angle according to their positions before computing the attention scores. The later the position, the more it rotates. That way, when the q and k of two positions take a dot product, the result depends only on how far apart they are, not on their absolute positions.

"月" comes one character after "明", and whether the pair appears at the start or the end of a poem, the relationship between them is the same. Relative position fits how language works better than absolute position, and generalises more easily to text longer than what was seen in training.

This course's GPT sticks with the simplest position embeddings, because they work perfectly well in our experiments and are easy to understand. Exercise 3 lets you try RoPE yourself.

Exercises

1. Add another case to heads_positions.py: a 3-layer model without position embeddings. How much do the two lines' outputs differ?
2. Change the model gpt.py trains from 4 heads to 1 and to 8 (keeping n_embd at 128), train each for 1,000 steps with next lesson's train.py, and compare validation loss.
3. Challenge: referring to the RoFormer paper or an open-source model's code, add rotary position embedding to CausalSelfAttention, remove pos_emb, and see how training goes.

Self-check

1. How does multi-head attention work? Does it add parameters?

Each character's vector is cut into several segments, attention is run on each separately (each with its own q, k, v and weights), and the results are joined back together and mixed by a linear layer. The number of parameters is the same as with one head; only the way the matrices are split differs. The benefit is that each head can learn to attend to different things.

2. Why does attention need position information?

Attention's scoring and weighted sum look only at the vectors' content, not their positions. In a one-layer model, reordering the earlier characters leaves the last position's output completely unchanged. But the meaning of language depends on order, so position information has to be added.

3. In the experiment, why could the 2-layer model without position embeddings tell the two lines' order apart?

The causal mask lets each position see only itself and earlier characters, so in the first layer different positions see different numbers of characters, and their outputs depend on order. When the second layer reads those outputs, it indirectly gets position information. Research has found that causal language models without positional encoding can still learn position information, but real models still add positional encoding explicitly.