Self-attention

Monday’s report from the support inbox was not good. Crumb’s little ticket router, a model Arjun had trained to send each email to the right queue (refunds, late deliveries, allergy questions), had misrouted 41 of last week’s 310 emails. Customers asking for refunds had waited a day in the wrong queue. Two allergy questions had landed in “general feedback”.

Priya pulled one of the failures up on the big screen:

“My eggless cake arrived cold. Can I get a refund for it?”

“Your router filed this under allergy,” she said. “Why?”

Arjun winced. “It saw eggless. Eggless usually means an allergy question.”

“And what does it mean at the end?”

“The cake. What else would it be?”

“Clear to you. You read the whole sentence and connected it back to cake, and you understood that eggless here is a product name, not a medical worry. Your router looks at each word on its own. Today we build the trick every modern language model uses to read words in context. It is called self-attention.”

Why one vector per word is not enough

Start with how Arjun’s router sees text. A token is a small piece of text, often a word or part of one, that a model reads as one unit. Each token becomes an embedding: a list of numbers that captures what the token means. You met embeddings in Module 07, where similar meanings landed close together on a map.

The catch is that the router looks up the same embedding for a word wherever it appears. It gets the same numbers in “can I get a refund for it” and “it was raining”. Cold means the same thing in “my cake arrived cold” and “do you have cold coffee”. Bank is identical next to river and next to loan.

Words do not carry their full meaning on their own. They borrow it from their neighbours. So the question for today is very practical: how can each word’s numbers soak up the right amount of meaning from the other words in its sentence?

In Module 29 you saw one answer, the RNN (recurrent neural network), which reads one token at a time and passes a running memory forward like a note handed down a line of people. By the end of a long email, the note is smudged, and nothing can be computed in parallel because each step waits for the last. The 2017 paper Attention Is All You Need dropped the note-passing entirely. It let every token look at every other token directly, all at the same time. That design is the transformer, the architecture behind every LLM you have called through an API, and self-attention is its heart.

Arjun’s first attempt: everyone gets an equal say

Arjun’s first idea was reasonable. “If words need context, let each word’s new vector be the average of every word in the sentence. Then every word knows about every other word.”

He tried it on three tokens from the email: eggless, chocolate, cake. To keep the numbers readable, he gave each a two-number embedding:

Token	Embedding
eggless	[1, 0]
chocolate	[0, 1]
cake	[1, 1]

emb = {"eggless": [1, 0], "chocolate": [0, 1], "cake": [1, 1]}
avg = [sum(v[i] for v in emb.values()) / 3 for i in range(2)]
for word in emb:
    print(word, "->", [round(a, 2) for a in avg])

eggless -> [0.67, 0.67]
chocolate -> [0.67, 0.67]
cake -> [0.67, 0.67]

He stared at it. Every word had become the same grey blur. “Well, now they all know about each other,” he said. “They also all mean the same thing.”

“Right,” said Priya. “Context is useful, but not every word matters equally to every other word. What does cake need to know about?”

“Eggless, definitely. Chocolate, a bit. The full stop at the end, not at all.”

“So instead of an equal average?”

“A weighted average. Each word decides how much to take from each other word.”

That is the core idea of attention, discovered in four lines of failed code. The rest of this lesson answers one question: where do the weights come from?

Second attempt: weights from similarity

Arjun’s next idea: words that are similar should listen to each other. A standard way to measure how similar two lists of numbers are is the dot product: multiply matching positions and add them up. A big dot product means the two lists point in a similar direction.

For cake = [1, 1]:

• cake · eggless = $1\times 1+1\times 0=1$
• cake · chocolate = $1\times 0+1\times 1=1$
• cake · cake = $1\times 1+1\times 1=2$

Cake gives its biggest weight to… itself. Every word is most similar to itself, so every word mostly listens to itself. That is barely better than no context at all.

Priya added a second problem on the whiteboard. “Think about it and cake. It desperately needs to find cake. Does cake need to find it?”

“Not really. Cake is fine on its own.”

“But a dot product is symmetric. it · cake is exactly the same number as cake · it. Similarity treats the relationship as two-way. Relevance is often one-way.”

Priya’s kitchen: questions, labels and goods

Priya drew the Crumb kitchen at 5 a.m. “Every baker has three things,” she said.

• A question they want answered. The person making plum cake asks, “Who has the soaked fruit?”
• A label on their apron saying what they can offer. “Soaked fruit.” “Whipped cream.” “Oven 3 is free.”
• The actual goods they hand over when asked.

“Each baker compares their question with everyone’s apron label. Good matches get most of their attention. Then they collect goods, mostly from the best matches and a little from the rest.”

“Why three separate things?” Arjun asked. “Why not one?”

“Because they do different jobs. The plum cake baker’s question is about fruit. Her apron label says ‘plum cake station’. And the goods she hands over, if someone asks her, are finished cakes. Question, label and goods are three different facts about the same person.”

In a transformer, those three things have names:

• Query (Q): what this token is looking for.
• Key (K): what this token advertises about itself.
• Value (V): the information this token hands over when someone attends to it.

This fixes both problems from the second attempt. Cake’s query can ask “what kind am I?” while eggless’s key says “I describe a kind”, so they match even though the words are not similar. And because the query and key are different vectors, it can strongly match cake without cake strongly matching it.

Where the analogy breaks: bakers ask one helper for goods. Attention never picks only one. It always takes a weighted blend of everyone, including the token itself, with weights that add up to exactly 1.

Following “cake” through attention, number by number

Here is the worked example. Real models use vectors with hundreds of numbers; we use two. Where do queries, keys and values come from? From the embeddings, multiplied by three learned tables of numbers. You will see that step shortly. For now, pretend it already happened and these are the results:

Token	Query $q$	Key $k$	Value $v$
eggless	[1, 0]	[2, 0]	[10, 0]
chocolate	[0, 1]	[0, 1]	[0, 10]
cake	[1, 1]	[1, 1]	[5, 5]

The values are chosen so you can read them: the first number is “how much eggless-ness”, the second “how much chocolate-ness”. Follow cake, whose query is [1, 1].

Step 1: Score. Dot cake’s query with every key.

• with eggless: $1\times 2+1\times 0=2$
• with chocolate: $1\times 0+1\times 1=1$
• with cake: $1\times 1+1\times 1=2$

Step 2: Scale. Divide by $\sqrt{d_k}$, the square root of the key length. Our keys have 2 numbers, so divide by $\sqrt{2}\approx 1.414$. The scores become 1.414, 0.707, 1.414. (Why this step exists is coming up; for now, follow along.)

Step 3: Softmax. Softmax turns any list of numbers into positive weights that add up to 1, giving bigger numbers a bigger share. It raises $e$ (about 2.718) to the power of each score, then divides each result by their total:

• $e^{1.414}\approx 4.11$
• $e^{0.707}\approx 2.03$
• $e^{1.414}\approx 4.11$

The total is 10.25, so the weights are $4.11/10.25\approx$ 0.40, $2.03/10.25\approx$ 0.20 and 0.40.

Step 4: Blend. Multiply each value by its weight and add them up:

$$0.40\times [10,0]+0.20\times [0,10]+0.40\times [5,5]=[4+0+2,0+2+2]\approx [6.0,4.0]$$

Cake went in as a neutral [5, 5]. It came out as [6.0, 4.0]: tilted towards eggless. Its new vector now says “an eggless kind of cake”, which is exactly the meaning a router needs to tell an eggless-cake refund from an egg allergy question.

Run the same four steps for the other two tokens and you get:

Token	Weights on (eggless, chocolate, cake)	Output
eggless	0.58, 0.14, 0.28	[7.2, 2.8]
chocolate	0.20, 0.40, 0.40	[4.0, 6.0]
cake	0.40, 0.20, 0.40	[6.0, 4.0]

Three words went in. Three context-aware words came out, each tilted by what matters to it. Compare that with Arjun’s grey blur of [0.67, 0.67] for everyone.

In a real model, nobody chooses these numbers. Training nudges the query, key and value tables until useful matches like “describing word ↔ the thing it describes” appear on their own.

Opening the hood: the formula and its matrices

Now the precise version, with every term defined.

Where Q, K and V come from

Each token starts as an embedding vector $x$. The model holds three weight matrices, $W^Q$, $W^K$ and $W^V$. A matrix is a grid of numbers; multiplying a vector by it produces a new vector. These three are learned during training (the slow process of adjusting a model’s numbers to reduce its mistakes):

$$q=x{W}^Q,k=x{W}^K,v=x{W}^V$$

Same input, three different “views” of it. That is the apron analogy in maths: one baker, three facts. Stack every token’s vector as a row of a matrix and you get the matrices $Q$, $K$ and $V$ for the whole sentence at once.

The whole formula

Everything you did by hand fits on one line:

$$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_k}}\right)V$$

Read it from the inside out:

• $Q{K}^{\top }$ computes every query-key dot product at once. ($K^{\top }$ means $K$ flipped so its rows become columns.) For $n$ tokens this is an $n\times n$ grid of scores. Row $i$ holds token $i$’s scores against everyone.
• Dividing by $\sqrt{d_k}$ scales the grid.
• Softmax runs along each row, so each row of weights adds to 1.
• Multiplying by $V$ blends the values with those weights.

The softmax itself, for a row of scores $s_1,\dots ,s_n$, is:

$$\text{softmax}(s{)}_j=\frac{e^{s_j}}{{\sum }_{m=1}^n{e}^{s_m}}$$

That is “raise $e$ to each score and divide by the total”, exactly as you did above.

Why the square root?

Here is the reason you skipped earlier. Suppose the numbers in queries and keys are roughly random, averaging 0 with a spread of about 1. A dot product adds up $d_k$ such products, so its spread grows with the length: its variance (spread, squared) is about $d_k$. Real models often use $d_k=64$, so raw scores swing by about $\sqrt{64}=8$ either way.

Watch what that does to softmax. Raw scores of [8, 0, 0] give weights of about [0.9993, 0.0003, 0.0003]: all attention on one token. Divided by 8, they become [1, 0, 0] and the weights are a much gentler [0.58, 0.21, 0.21]. When softmax is that spiky, the gradients (the signals training uses to decide which way to adjust each number, from Module 27) shrink towards zero, and learning stalls. Dividing by $\sqrt{d_k}$ puts the spread back near 1, whatever the vector length.

Building it in PyTorch, one step at a time

PyTorch is the Python library for neural networks you used in Module 27. A tensor is PyTorch’s name for a grid of numbers of any shape. Install it with uv add torch.

Step 1: reproduce the worked example. Put the table into tensors and run the formula directly.

import math
import torch

Q = torch.tensor([[1., 0.], [0., 1.], [1., 1.]])    # eggless, chocolate, cake
K = torch.tensor([[2., 0.], [0., 1.], [1., 1.]])
V = torch.tensor([[10., 0.], [0., 10.], [5., 5.]])

scores = Q @ K.T / math.sqrt(K.shape[-1])   # @ is matrix multiplication
weights = torch.softmax(scores, dim=-1)     # softmax along each row
print(weights.round(decimals=2))
print((weights @ V).round(decimals=1))

tensor([[0.5800, 0.1400, 0.2800],
        [0.2000, 0.4000, 0.4000],
        [0.4000, 0.2000, 0.4000]])
tensor([[7.2000, 2.8000],
        [4.0000, 6.0000],
        [6.0000, 4.0000]])

The same numbers you worked out by hand. dim=-1 means “along the last dimension”, which is across each row. That detail matters, as you will see in the mistakes section.

Step 2: learn the projections. In a real layer, $Q$, $K$ and $V$ come from the embeddings through learned weights. PyTorch’s nn.Linear is a learned matrix multiplication, so three of them give you $W^Q$, $W^K$ and $W^V$.

import torch.nn as nn

class SelfAttention(nn.Module):
    def __init__(self, d_model: int, d_k: int):
        super().__init__()
        # Three learned projections: queries, keys and values.
        self.W_q = nn.Linear(d_model, d_k, bias=False)
        self.W_k = nn.Linear(d_model, d_k, bias=False)
        self.W_v = nn.Linear(d_model, d_k, bias=False)
        self.d_k = d_k

    def forward(self, x):                        # x: (tokens, d_model)
        Q, K, V = self.W_q(x), self.W_k(x), self.W_v(x)
        scores = Q @ K.T / math.sqrt(self.d_k)   # (tokens, tokens)
        weights = torch.softmax(scores, dim=-1)  # each row sums to 1
        return weights @ V, weights

d_model is the length of each embedding coming in; d_k is the length of each query, key and value. They do not have to match.

Step 3: run it on a fake sentence.

torch.manual_seed(0)
x = torch.randn(3, 8)                 # 3 tokens, each an 8-number embedding
attn = SelfAttention(d_model=8, d_k=4)
out, weights = attn(x)
print(out.shape)                      # torch.Size([3, 4])
print(weights.sum(dim=-1))            # all ones: each row of weights sums to 1

Three tokens in, three context-aware vectors out. The weights grid is the famous attention map: row $i$ shows how much token $i$ looked at each token. Before training, these weights are meaningless, because the projections are random. Training is what makes cake look at eggless.

Where self-attention breaks

Self-attention is powerful, but it has sharp edges. Experienced engineers watch for four.

It gets expensive, fast. The score grid is $n\times n$. A 1,000-token support email means 1,000,000 scores. A 100,000-token document means 10,000,000,000 scores, and that is for one attention layer; a real model repeats it in every head of every layer. Double the text and you quadruple the work. This is why long context windows (the maximum text a model can read at once, from Module 01) cost so much, and it is exactly the problem Module 32 attacks with flash attention and sparse attention.

It has no sense of order. Look at the formula again: nothing in it says which token came first. Shuffle the words and every token gets the same blend, in a different row. “Cancel the order, not the refund” and “cancel the refund, not the order” would look identical. The next lesson fixes this.

Attention maps are not explanations. It is tempting to point at a map and say “the model filed this as allergy because it looked at eggless”. But one map shows one layer. A real model has dozens of layers with many attention patterns each, plus other machinery in between. A map is a clue, not a proof.

Softmax never says zero. Every token always gets some weight. In a long document, a thousand tiny irrelevant weights can add up to real noise. Models learn to live with this, but it is one reason very long prompts sometimes confuse them.

And the story’s twist: when Arjun replaced his router with a small pretrained transformer, the eggless misroutes mostly vanished, and misrouted emails fell from 41 a week to 9. Most of the remaining nine were long, rambling emails. To keep costs down he had cut every email off at 256 tokens, and the crucial sentence was sometimes past the cut. Quadratic cost is not an abstract worry; it shapes the choices you make.

Try it yourself

Change the worked example so cake ignores chocolate almost entirely. You only need to change one key. Predict cake’s three weights by hand first, then check with the Step 1 code.

Show the answer

Make chocolate’s key point away from cake’s query, for example [0, -3]. Cake’s scores become 2, -3 and 2. Divided by 1.414 they are 1.414, -2.121 and 1.414. After softmax the weights are about 0.49, 0.01 and 0.49. Cake now blends almost only eggless and itself, and its output becomes about [7.4, 2.6], tilted even further towards eggless. Notice the chocolate weight is tiny but still not zero.

Common mistakes

• Thinking attention picks one word. It always blends every token. When people say “the model attends to cake”, they mean cake got the biggest share.
• Thinking queries and keys are the same thing. They come from different learned matrices. That difference is what lets relevance be one-way.
• Forgetting the scale. Drop $\sqrt{d_k}$ and training gets unstable, because softmax saturates and gradients vanish. The model still runs; it learns badly.
• Softmax along the wrong dimension. Softmax must run along each row (dim=-1) so each query’s weights add to 1. Use dim=0 and the code still runs, with no error, and produces nonsense.
• Reading attention maps as reasons. They show where information flowed in one layer, not why the model decided something.

Recap

• A fixed vector per word ignores context; self-attention lets each token update itself using every other token.
• An equal average blurs everything together; attention uses a weighted average instead.
• Each token gets a query (what I want), a key (what I offer) and a value (what I hand over), made by three learned matrices, so relevance can be one-way.
• Scores are query-key dot products, divided by $\sqrt{d_k}$ to keep softmax calm, then turned into weights that add to 1.
• Each token’s output is the weighted blend of all values; the work grows with the square of the text length.

That evening, Meera stopped by Arjun’s desk. “Priya says you fixed the router. What was wrong with it?”

Arjun thought about how to say it without a whiteboard. “It read every word on its own, like reading a list of ingredients and guessing the dish. Now each word looks at all the others and asks, ‘which of you changes what I mean?’ Cake looks around, sees eggless, and becomes eggless cake: a product, not an allergy. It takes a little from every word, but most from the ones that matter.”

“And it can still get it wrong?”

“Yes, especially on long rambling emails. And there’s a hole Priya pointed out: it doesn’t know word order. ‘Refund the cake, not the delivery’ looks the same as the opposite.” He sighed. “That’s tomorrow.”

Go deeper

• The Illustrated Transformer, Jay Alammar’s visual walkthrough of exactly this calculation.
• Attention Is All You Need, the original 2017 paper. Section 3.2 is this lesson.
• PyTorch tutorials