Positional encoding and multi-head attention

On Wednesday a refund went wrong. A customer in Kothrud had written: “Please refund the delivery fee, not the cake. The cake was lovely, the rider was two hours late.” The support agent, trusting the label Arjun’s new attention-based router had put on the email, refunded the full ₹850 cake instead of the ₹60 delivery fee.

Arjun dug in. He fed the router two emails:

• “Refund the delivery fee, not the cake.”
• “Refund the cake, not the delivery fee.”

The router gave both exactly the same label, with exactly the same confidence: 0.87.

“How?” he asked Priya. “Those mean opposite things.”

“To you they do. To your attention layer they are the same bag of words. Same words, same counts. Nothing in the attention formula says which word came first.” She drew two things on the whiteboard: a clock, and a tasting panel of three people. “Two fixes today. The clock tells each word where it is. The panel gives the model more than one way of looking.”

Why the router could not tell the emails apart

Recall last lesson. A token is a small piece of text the model reads as one unit. Each token becomes an embedding (a list of numbers that captures its meaning). Self-attention then lets every token blend in information from every other token: each token has a query (what I am looking for), a key (what I offer) and a value (what I hand over), and the blend weights come from query-key matches.

Now shuffle the tokens. Every query, key and value moves with its token. The scores shuffle, the weights shuffle, and each token ends up with the exact same blend as before, only in a different row. Mathematicians call this permutation equivariance. In plain English: attention treats a sentence as a bag of words. If the router then averages the token vectors to make its decision, the order disappears completely, and “refund the cake, not the fee” equals “refund the fee, not the cake”.

You can watch it happen with the SelfAttention class from the last lesson:

import torch
torch.manual_seed(0)
x = torch.randn(3, 8)                     # three tokens: A, B, C
attn = SelfAttention(d_model=8, d_k=4)
out, _ = attn(x)
out_shuffled, _ = attn(x[[2, 0, 1]])      # same tokens, order C, A, B
print(torch.allclose(out[[2, 0, 1]], out_shuffled))   # True

True means: same blends, new order. The model has no idea which word came first. We need to tell it.

Arjun’s first attempt: write the position number down

“Easy,” said Arjun. “I’ll add one more number to each embedding: its position. Word 0 gets a 0, word 1 gets a 1, and so on.”

Priya let him try it. Two things went wrong within the hour.

First, the numbers were the wrong size. Embedding values in his model sat roughly between -1 and 1. A long complaint email had 300 tokens, so its last word carried the number 299. That single giant number drowned out every meaningful one, and training became unstable.

So he divided by the email’s length: position 3 in a 10-word email became 0.3. That fixed the size, and broke something subtler. Now 0.3 meant “word 3” in a 10-word email and “word 30” in a 100-word email. The model could never learn a simple rule like “the word right after not is the thing being excluded”, because “right after” was a different number in every email.

“So I want something that stays small,” Arjun said slowly, “but where ‘one step later’ means the same thing everywhere.”

“Now you’re designing it,” said Priya. “Look at the clock.”

Priya’s clock: a position written in many hands

“What time is it?” Priya asked.

“Ten past three.”

“How do you know it is not ten past four?”

“The hour hand.”

“And how do you tell 3:10 from 3:11?”

“The minute hand. The hour hand barely moves in a minute.”

“So one hand is not enough. A fast hand tells nearby moments apart, but it comes back round and repeats. A slow hand never repeats within a day, but it barely moves between neighbouring moments. Read several hands together and every moment gets a unique, bounded reading. All the angles stay between the same limits, however late it gets.”

That is the trick in the original transformer. Each position gets a “clock reading”: pairs of sine and cosine waves spinning at different speeds. Sines and cosines always stay between -1 and 1, so the numbers never blow up. And moving one step forward turns every hand by a fixed amount, so “one step later” means the same thing at position 5 as at position 500.

Where the analogy breaks: a real clock repeats every 12 hours. With enough hands at very different speeds, these patterns do not repeat within any text length you will meet.

Following one word to two different places

Let us make the clock real with a tiny 4-number version. The formula (coming in a moment) gives two hands:

• Fast hand (dimensions 1 and 2): turns 1 radian per position. A radian is a unit of angle, about 57 degrees.
• Slow hand (dimensions 3 and 4): turns 0.01 radians per position.

Each hand is written as a sine and a cosine of its angle. Here are the readings:

Position	Fast sin	Fast cos	Slow sin	Slow cos
0	0.000	1.000	0.000	1.000
1	0.841	0.540	0.010	1.000
2	0.909	-0.416	0.020	1.000
3	0.141	-0.990	0.030	1.000
44	0.018	1.000	0.426	0.905

Look at positions 1, 2 and 3: the fast hand changes a lot, the slow hand barely moves. Now look at position 44. On the fast hand it is almost exactly position 0 again (0.018 and 1.000), because the hand has gone round seven times. Only the slow hand (0.426 and 0.905) gives it away. Together, they are unique.

Now suppose the token cake has the embedding [0.5, 0.2, -0.1, 0.3]. The model adds the position reading to it:

• cake at position 1: [0.5 + 0.841, 0.2 + 0.540, -0.1 + 0.010, 0.3 + 1.000] = [1.341, 0.740, -0.090, 1.300]
• cake at position 3: [0.5 + 0.141, 0.2 - 0.990, -0.1 + 0.030, 0.3 + 1.000] = [0.641, -0.790, -0.070, 1.300]

Same word, two different starting vectors. In “refund the cake, not the fee”, cake sits early; in “refund the fee, not the cake”, it sits late. Now their queries and keys differ, attention produces different blends, and the router can finally learn that the thing after not is the thing to keep.

Opening the hood: the sinusoidal formula

For a token at position $pos$ in an embedding of size $d_{model}$, the original transformer fills dimension pair $2i$ and $2i+1$ with:

$$P{E}_{(pos,2i)}=\sin \left(\frac{pos}{10000^{2i/d_{model}}}\right),P{E}_{(pos,2i+1)}=\cos \left(\frac{pos}{10000^{2i/d_{model}}}\right)$$

The divisor $10000^{2i/d_{model}}$ sets each hand’s speed. For the first pair ($i=0$) it is $10000^0=1$, the fast hand. For our 4-number example the second pair has $i=1$, giving $10000^{2/4}=100$: it turns $1/100=0.01$ radians per step, the slow hand. In a real 512-number model there are 256 hands, from one turn every 6.3 tokens down to one turn every 60,000 or so.

Then the position vector is added to the token’s embedding:

$$x_{pos}=\text{embedding}(\text{token})+P{E}_{pos}$$

Why add rather than tack the numbers on the end? Adding keeps the vector the same size, and in a space of hundreds of dimensions the model has plenty of room to keep “what” and “where” mostly separate. It works well in practice.

There is a lovely bonus property. Moving forward by $k$ positions rotates every hand by a fixed angle, and rotation is a simple, fixed calculation. So the reading for position $pos+k$ can be computed from the reading for $pos$ in the same way at any $pos$. That gives attention an easy route to learning relative rules like “the word two steps back”, which is exactly what Arjun’s divided-by-length attempt could not offer.

What models use today

The sinusoidal clock was the 2017 answer. You will meet two others constantly:

• Learned positions. BERT and GPT-2 skip the formula and learn one position vector per slot, like any other weight. Simple and effective, with a hard limit: BERT learned 512 slots and GPT-2 learned 1,024. There is no vector for slot 1,025.
• RoPE (rotary position embedding). Used by the Llama family and many recent open models. Instead of adding a clock reading to the embedding, RoPE rotates each query and key by an angle proportional to its position. When a query at position 10 meets a key at position 7, the two rotations partly cancel and the score depends on the gap, 3. Attention becomes naturally about relative distance.

The goal is always the same: make order visible to attention.

The second problem: one pair of eyes

With positions fixed, the router improved. But Priya had another email ready:

“Refund my eggless cake, it was cold.”

“What does it need to know?” she asked.

“Which noun it refers to: cake. And probably refund, because that is what the customer wants done to it.”

“And cold?”

“What was cold. It, so cake again.”

“Now remember, each token has one query and gets one row of weights that must add to 1. If it puts 0.7 on cake, it can only spread 0.3 over everything else, including refund. One attention pattern has to compromise between every relationship the word cares about.”

Priya’s tasting panel

When Crumb tests a new Diwali hamper, Meera does not ask one person “is it good?”. She asks a panel. One taster judges sweetness. One judges packaging. One judges price. Each gives full attention to their own question. Their notes get combined into one verdict.

Multi-head attention does the same. It runs several attention calculations, called heads, side by side. Each head has its own query, key and value matrices, so each learns its own idea of what counts as a match. Head 1 can put 0.9 of its attention on cake while head 2 puts 0.9 on refund. No compromise needed.

Where the analogy breaks: Meera assigns each taster a job. Nobody assigns a head a job. Researchers poking at trained models have found heads that seem to track the previous word, or link pronouns to nouns, but those roles emerge from training. Many heads do things nobody has neatly explained.

Opening the hood: the multi-head formula

Each head $i$ is an ordinary attention calculation with its own learned matrices:

$${\text{head}}_i=\text{Attention}(X{W}_i^Q,X{W}_i^K,X{W}_i^V)$$

The heads’ outputs are concatenated (placed end to end) and mixed by one more learned matrix, $W^O$, so the panel’s notes become one verdict:

$$\text{MultiHead}(X)=\text{Concat}({\text{head}}_1,\dots ,{\text{head}}_h)W^O$$

Here is the clever accounting. The original transformer had $d_{model}=512$ and $h=8$ heads. Rather than giving each head full 512-number queries, it split the space: each head works in $d_k=512/8=64$ dimensions. Count the weights. One big head needs $W^Q$, $W^K$, $W^V$ of size $512\times 512$ each, plus $W^O$: $4\times 512\times 512=1,048,576$ numbers. Eight heads of size 64 need $8\times 3\times 512\times 64=786,432$ for their Q, K and V, plus the same $512\times 512=262,144$ for $W^O$: also $1,048,576$. Eight viewpoints for the price of one.

Building both pieces in PyTorch

Step 1: the clock. This function returns one row of readings per position.

import math
import torch

def sinusoidal_positions(n_positions: int, d_model: int) -> torch.Tensor:
    pos = torch.arange(n_positions).unsqueeze(1)          # column of positions: (n, 1)
    i = torch.arange(0, d_model, 2)                       # even dimension indices: 0, 2, 4...
    speed = torch.exp(-math.log(10000.0) * i / d_model)   # 1 / 10000^(2i/d), one per hand
    pe = torch.zeros(n_positions, d_model)
    pe[:, 0::2] = torch.sin(pos * speed)                  # even dimensions: sine
    pe[:, 1::2] = torch.cos(pos * speed)                  # odd dimensions: cosine
    return pe

pe = sinusoidal_positions(n_positions=50, d_model=4)
print(pe[[0, 1, 2, 3, 44]].round(decimals=3))

tensor([[ 0.0000,  1.0000,  0.0000,  1.0000],
        [ 0.8410,  0.5400,  0.0100,  1.0000],
        [ 0.9090, -0.4160,  0.0200,  1.0000],
        [ 0.1410, -0.9900,  0.0300,  1.0000],
        [ 0.0180,  1.0000,  0.4260,  0.9050]])

The same table you traced by hand. The exp(-log(10000) * ...) line is a numerically safe way to compute $1/10000^{2i/d}$.

Step 2: rerun the shuffle test with positions. Positions belong to slots, not words, so you add the same position table after shuffling.

torch.manual_seed(0)
x = torch.randn(3, 8)
attn = SelfAttention(d_model=8, d_k=4)
pos = sinusoidal_positions(3, 8)
out, _ = attn(x + pos)
out_shuffled, _ = attn(x[[2, 0, 1]] + pos)
print(torch.allclose(out[[2, 0, 1]], out_shuffled))   # False: order now matters

Step 3: multi-head attention. The standard trick computes every head’s Q, K and V with one big matrix, then reshapes so each head gets its own slice.

import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        assert d_model % n_heads == 0, "heads must split d_model evenly"
        self.n_heads, self.d_k = n_heads, d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)  # all heads' Q, K, V at once
        self.out = nn.Linear(d_model, d_model, bias=False)      # W^O mixes the heads

    def forward(self, x):                        # x: (batch, tokens, d_model)
        B, T, D = x.shape
        q, k, v = self.qkv(x).split(D, dim=-1)   # each (B, T, D)
        # (B, T, D) -> (B, T, heads, d_k) -> (B, heads, T, d_k)
        q, k, v = (t.view(B, T, self.n_heads, self.d_k).transpose(1, 2) for t in (q, k, v))
        scores = q @ k.transpose(-2, -1) / self.d_k ** 0.5      # (B, heads, T, T)
        y = F.softmax(scores, dim=-1) @ v                        # (B, heads, T, d_k)
        y = y.transpose(1, 2).reshape(B, T, D)                   # glue heads back together
        return self.out(y)

mha = MultiHeadAttention(d_model=64, n_heads=8)
x = torch.randn(2, 5, 64)       # a batch of 2 emails, 5 tokens each, 64 numbers per token
print(mha(x).shape)             # torch.Size([2, 5, 64])
print(sum(p.numel() for p in mha.parameters()))   # 16384 = 4 x 64 x 64

The reshape is the only new idea. view(B, T, 8, 8) cuts each 64-number vector into 8 slices of 8. transpose(1, 2) moves the head dimension forward, so the attention maths runs on all 8 heads at once, each producing its own $T\times T$ grid of weights. At the end, the reverse reshape glues the 8 slices back into 64 numbers, and self.out mixes them. The parameter count, 16,384, is the same as a single 64-wide head with its output matrix: the accounting from above, in code.

In real projects you would replace the scores-softmax-multiply lines with F.scaled_dot_product_attention(q, k, v), which does the same maths with much faster kernels. Module 32 explains why it is faster.

Where positions and heads break

Running past the end of the table. Arjun’s router used learned positions with 128 slots, because his training emails were short. The first 400-token complaint crashed it with an IndexError: there was no slot 128. Learned tables cannot stretch. Sinusoidal and RoPE schemes do not crash, but quality still drops on text much longer than anything seen in training, unless the model is extended with care. When a provider advertises a long context window, part of what they did was teach the model to handle positions it never saw early on.

Positions are soft, not exact. A position signal tells the model roughly where words are. It does not make the model good at exact counting tasks like “what is the 37th word?”. Models still stumble on those.

Many heads turn out to be spare. Researchers have pruned trained models and found that a good fraction of heads can be removed with only a small drop in quality. Heads overlap, and some do very little. That is one reason later designs share keys and values across groups of heads to save memory, which you will meet in Module 32.

Head count must divide the model size. 512 with 8 heads works; 512 with 6 heads does not. The assert in the code catches it.

Try it yourself

1. Using sinusoidal_positions(50, 4), find the position between 1 and 49 whose fast hand (first two numbers) is closest to position 0’s. Is its slow hand also close?
2. In MultiHeadAttention(d_model=768, n_heads=12), the size used by BERT-base and GPT-2 small, how many numbers does each head’s query have per token?

Show the answer

1. Position 44, with fast readings of about (0.018, 1.000), almost identical to position 0’s (0.000, 1.000). Its slow hand reads (0.426, 0.905) against position 0’s (0.000, 1.000), clearly different. That is exactly why you need several hands. (Position 19 and position 25 are also fairly close on the fast hand, because 19 and 25 are near multiples of 6.28.)
2. $768/12=64$ numbers per head.

Common mistakes

• Adding positions to words instead of slots. Shuffle the words, then add the same position table. A position describes where a token is, not which token it is.
• Using raw position numbers. Large numbers swamp the meaning; dividing by length makes “one step later” mean something different in every text. Use a bounded, consistent scheme.
• Assuming more heads is always better. More heads means narrower heads. Past a point, each is too thin to be useful.
• Expecting free long context. A model trained on 2,000-token texts will not handle 20,000 tokens well because you passed it more. Check the model’s documented context length.

Recap

• Attention alone ignores word order: shuffle the input and you get the same blends, shuffled.
• Positional encoding adds a “where am I” signal to each token; the classic version is a clock of sine and cosine hands at many speeds, bounded and consistent at every position.
• Learned position tables and RoPE are the common modern choices, with the same goal.
• Multi-head attention runs several attention patterns in parallel, each with its own queries, keys and values, and mixes them with $W^O$.
• Splitting $d_{model}$ across heads keeps the cost the same as one big head.

Meera wanted to know about the ₹850 refund. “Will it happen again?”

“Less likely,” said Arjun. “Before, the router saw each email as a bag of words: refund, cake, fee, not, all thrown in together. Now every word carries a little clock reading that says where it sits, so ‘not the cake’ and ‘not the fee’ look different. And it reads each email with several sets of eyes at once: one can follow what it refers to while another tracks what the customer wants done.”

“So it understands now?”

“It notices more,” he said carefully. “Priya says there’s one more trick before this thing could write replies instead of only labelling them. Right now, if we trained it to write, it could peek at the answer.”

Go deeper

• The Illustrated Transformer, especially the sections on positional encoding and multi-headed attention.
• Attention Is All You Need, sections 3.2.2 and 3.5.
• Hugging Face LLM course