Attention, Part 1: Self-Attention and Multi-Head Attention

The one idea every modern AI language model is built on, explained from zero: why attention was invented, how a word decides which other words to look at, why the scores are divided by √d, how several heads work together, and what real attention heads inside Qwen2.5 actually learn. Every claim checked in code.

Every AI chatbot you have used (ChatGPT, Gemini, Llama, Qwen, Gemma, DeepSeek) is built around one idea called attention. It is the part of the model that lets each word look at the other words around it and decide which ones matter.

This series explains attention from zero, and then every important variation of it that today's models use. No background is assumed. Whenever a technical word appears, there is a box like this one that explains it:

First: words become numbers#

A computer cannot read words. So before anything else, a language model turns text into tokens, and every token into a list of numbers.

From here on, "a token" and "the token's vector" mean the same thing: attention only ever works with the numbers.

The problem attention solved#

Before 2014, translation models read a sentence one word at a time and squeezed everything into a single vector, a kind of running summary. Then they wrote the translation from that one summary.

The problem: one fixed-size summary has to hold the whole sentence, whether it has five words or fifty. Long sentences get blurry.

Here is how the paper that introduced attention describes that problem, in its very first section:

In 2014, Bahdanau, Cho and Bengio had a simple idea: let the model look back at every word, every time it writes a word, and learn how much to pay attention to each one. To write "cat" it looks mostly at "chat"; to write "black", mostly at "noir". That weighting is attention.

The first attention equations#

Their version works like this. The encoder leaves one vector $h_j$ for each source word $j$ (they call these annotations). When the decoder is about to write target word $i$, it does three things:

where:

• $s_{i-1}$ is the decoder's current state (what it has written so far);
• $a$ is a small learned network, the alignment model, that scores how well source word $j$ fits the next target word;
• $e_{ij}$ is that score, and $\alpha_{ij}$ is the score turned into a weight by softmax (the same softmax you will meet below);
• $T_x$ is the number of source words;
• $c_i$ is the context vector: a weighted mix of all the source words, made fresh for every target word.

The paper also drew the weights as pictures, and they show something nobody had programmed:

In 2017 the paper Attention Is All You Need went further. It threw the word-by-word reading away and made attention the main way words share information. That design is called the transformer, and every modern chatbot is one. When the words of one text attend to each other, it is called self-attention.

Two small tools: the dot product and softmax#

Attention is built from two simple tools. It is worth meeting them first.

Tool 1: the dot product measures "how similar"#

As an equation, for two vectors $a$ and $b$ with $d$ numbers each:

where $a_1$ is the first number of $a$, $b_1$ the first number of $b$, and so on.

For example, $[1, 2] \cdot [3, 4] = 1 \times 3 + 2 \times 4 = 11$.

Why is this useful? Because it tells you how much two vectors point the same way:

So "is this word relevant to that word?" can become "is the dot product of their vectors big?"

Tool 2: softmax turns scores into percentages#

Dot products can be any number: big, small, or negative. Attention needs weights instead: positive numbers that add up to 1, like percentages.

where $s_1, \dots, s_n$ are the scores, and $e \approx 2.718$. Raising $e$ to a power makes every score positive; dividing by the total makes them add up to 1.

Self-attention, step by step#

Now we can build attention. It answers one question for every token: of the tokens I am allowed to see, which ones should I take information from, and how much?

Each token's vector is turned into three new vectors, each made by multiplying it by a different grid of numbers:

• a query (Q): what am I looking for?
• a key (K): what do I contain? (used for matching)
• a value (V): what will I hand over if I am chosen?

A library is a good picture. Your query is the question you bring. Every book has a key (its title and topic) and a value (its contents). You compare your question with every title, then read mostly from the books that match best.

The steps:

1. Score. Compare every query with every key using the dot product. A big score means "this key matches what I am looking for".
2. Scale. Divide the scores by $\sqrt{d}$ (the square root of the vector length). The next section shows why.
3. Hide the future. A model that writes left to right must not peek at words that come later, so those scores are blocked.
4. Softmax. Turn each token's scores into weights that add up to 1.
5. Mix. Each token's new vector is the weighted average of the values.

All five steps fit in one equation, the most important equation in this series:

where:

• $Q$, $K$, $V$ are the queries, keys and values of all the tokens, stacked as rows;
• $QK^\top$ is every query's dot product with every key (the $^\top$, "transpose", just flips $K$ so the multiplication lines up);
• $d$ is the length of each key vector;
• $M$ is the causal mask: $0$ where looking is allowed, $-\infty$ where it is not.

This is the equation from Attention Is All You Need (Vaswani et al., 2017), the paper that introduced the transformer. Here it is in the original:

A tiny example you can check#

Here is the same thing in a few lines of Python, using PyTorch:

import math, torch

def self_attention(x, Wq, Wk, Wv, causal=True):
    """x holds one vector per token. Returns the attention weights A and the new vectors Z."""
    Q, K, V = x @ Wq, x @ Wk, x @ Wv                   # 1. make queries, keys and values
    scores = Q @ K.T / math.sqrt(K.shape[-1])          # 2. score every pair, divided by sqrt(d)
    if causal:                                         # 3. hide the future
        T = x.shape[0]
        future = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
        scores = scores.masked_fill(future, float("-inf"))
    A = torch.softmax(scores, dim=-1)                  # 4. scores -> weights; each row adds up to 1
    return A, A @ V                                    # 5. mix the values

I ran it on four tokens, The cat sat down, with random weights. Here are the attention weights it printed: each row is one token, showing how much it looks at each token in the columns.

attention weights (rows = query token, columns = key token):
     The  1.00  0.00  0.00  0.00
     cat  0.25  0.75  0.00  0.00
     sat  0.31  0.26  0.42  0.00
    down  0.08  0.29  0.09  0.54
   row sums: [1.0, 1.0, 1.0, 1.0]

Three things to notice:

• The can only look at itself, so its row is a single 1.00. The causal mask is working.
• Every row adds up to exactly 1. A token can choose where to look, but its total attention is always 100%.
• There is one number for every pair of tokens. 4 tokens make 16 numbers. 1,000 tokens make 1,000,000, in every head of every layer. Remember this: it is the cost that the whole rest of this series is about.

Why divide by √d?#

Step 2 looks like a small detail. It is not: leave it out and the model cannot learn properly. Here is why, with measurements.

A dot product adds up $d$ multiplications. If the numbers are random, adding more of them makes the total swing more widely: the typical size of a dot product grows like $\sqrt{d}$.

Softmax is very sensitive to big numbers. If one score is much bigger than the others, softmax gives it almost all the weight, and the model effectively stops looking at anything else.

I measured it: 10,000 random queries, each scored against 64 random keys, for four vector lengths.

The spread grows exactly as $\sqrt{d}$: 4, 8, 16, 32. Without the division, at $d = 1{,}024$ one random key already takes 95% of the attention before the model has learned anything. When softmax is that lopsided, its learning signals (gradients) become tiny and training stalls. Dividing by $\sqrt{d}$ cancels the growth: the top weight stays near 0.107 at every size.

Where the √d comes from (the math in four lines)#

Suppose every number in $q$ and $k$ is random, with mean 0 and variance 1, and all are independent. Look at one term of the dot product, $q_i k_i$:

The dot product adds $d$ such independent terms, and variances of independent terms simply add up:

So dividing by $\sqrt{d}$ brings the spread back to exactly 1, whatever the vector length:

where $\mathbb{E}$ means "average value" (expected value) and $\operatorname{Var}$ means variance. That is exactly what my measurement showed: spreads of 4, 8, 16, 32 for $d$ = 16, 64, 256, 1,024, which are $\sqrt{d}$.

The transformer paper gives the same reasoning, in one sentence and one footnote:

Many heads instead of one#

One attention calculation learns one pattern. But a word often needs several kinds of information at once: the word just before it, the subject of the sentence, the matching bracket in some code. So models run several attention calculations side by side.

As equations:

where:

• $X$ holds the vectors of all the tokens;
• $W_q^{(i)}, W_k^{(i)}, W_v^{(i)}$ are head $i$'s own query, key and value matrices;
• $h$ is the number of heads;
• Concat places the heads' results side by side into one long vector per token;
• $W_o$ is the output matrix that mixes the heads back together.

Each head works with shorter vectors, so $h$ heads cost about the same as one big head:

where $d_{\text{model}}$ is the length of each token's vector and $h$ the number of heads. The original transformer used exactly these numbers: 512 per token, 8 heads, 64 per head. Total work is $h \times d_{\text{head}} = d_{\text{model}}$, the same as one head of full size.

class MultiHeadAttention(torch.nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.h, self.d = n_heads, d_model // n_heads          # number of heads, numbers per head
        self.Wq = torch.nn.Linear(d_model, d_model, bias=False)
        self.Wk = torch.nn.Linear(d_model, d_model, bias=False)
        self.Wv = torch.nn.Linear(d_model, d_model, bias=False)
        self.Wo = torch.nn.Linear(d_model, d_model, bias=False)

    def forward(self, x):                                     # x: (batch, tokens, d_model)
        B, T, _ = x.shape
        split = lambda t: t.view(B, T, self.h, self.d).transpose(1, 2)   # cut into h heads
        q, k, v = split(self.Wq(x)), split(self.Wk(x)), split(self.Wv(x))
        scores = q @ k.transpose(-2, -1) / math.sqrt(self.d)             # one score grid per head
        mask = torch.triu(torch.ones(T, T, dtype=torch.bool, device=x.device), 1)
        A = torch.softmax(scores.masked_fill(mask, float("-inf")), -1)
        heads = A @ v
        return self.Wo(heads.transpose(1, 2).reshape(B, T, self.h * self.d))   # join and mix

Proof that it is right#

Code you write yourself is only trustworthy if it agrees with a trusted version. PyTorch has two built-in attention functions, so I gave them exactly the same weights as the class above and compared the results:

our MHA vs F.scaled_dot_product_attention: max |diff| = 1.7e-16
our MHA vs torch.nn.MultiheadAttention:     max |diff| = 1.7e-16

The largest difference is $1.7 \times 10^{-16}$, that is 0.00000000000000017. That is the smallest rounding error a computer can make with these numbers: the three versions compute exactly the same thing.

What real attention heads learn#

Random weights show how attention works. To see what it does, we have to look inside a trained model. I used Qwen2.5-0.5B, a small open model with 24 layers and 14 heads in each layer: 336 heads in total.

I gave it this sentence and read out the attention weights of every head:

The cat sat on the mat because it was tired.

Three kinds of head stood out:

1. Attention sinks. In 68% of all heads, more than half of the attention goes to the very first token, whatever it is.

2. A previous-token head. Layer 8, head 7 puts on average 84% of each token's attention on the token just before it: the clean diagonal in the middle picture. Heads like this help with copying and spotting patterns.

3. A pronoun head. In layer 5, head 5, the word "it" puts 83% of its attention on "cat", the thing "it" refers to. One sentence could be luck, so I tested the same head on six sentences it had never seen:

Six out of six. One honest note: in every one of these sentences, the right word is also the main noun near the start, so this head might be finding "the main noun" rather than truly understanding who "it" is. Either way, nobody programmed this. It appeared on its own during training.

The same patterns, found by researchers#

These findings are not special to Qwen. Two well-known papers found the same kinds of heads in other models.

The cost that every other part tries to cut#

Attention has two costs, and both grow with the length of the text:

1. Work. Every token compares itself with every earlier token. Twice as much text means four times as many comparisons.
2. Memory. When the model writes a reply one token at a time, it keeps every earlier token's keys and values in memory so it does not have to recompute them.

With plain multi-head attention, every head in every layer stores a key and a value for every token. The memory per token is:

where the 2 counts one key and one value, and $d_{\text{head}}$ is the length of each head's vectors. For Llama 2 7B that is $2 \times 32 \times 32 \times 128 \times 2 = 524{,}288$ bytes, or 512 KiB per token. At its full 4,096-token length, one conversation needs 2 GiB just for this memory.

The transformer paper itself lists this cost, comparing attention with the older layer types:

Every idea in the rest of the series cuts one of these two costs:

How attention changed everything#

It is hard to overstate what these two papers started. A short timeline of what happened next:

Use cases: where you meet this today#

Everything in this part runs, almost unchanged, inside products you use:

• Translation. The original use case. Attention lets the model align words across languages, as Figure 3 of Bahdanau et al. showed above.
• Chatbots and writing assistants. ChatGPT, Gemini, Llama, Qwen and DeepSeek are stacks of the multi-head attention built in this part (plus the efficiency tricks of Parts 2 to 4). Each word they write is one round of exactly this attention over the conversation so far.
• Search and understanding. Encoder transformers such as BERT read a query and documents and match their meaning, not just their words. The heads Clark et al. found (objects, determiners, coreference) are part of why this works.
• Code assistants. The same attention links a variable to where it was defined, or a closing bracket to its opening one, exactly like the previous-token and pronoun heads above.
• Long conversations. Attention sinks matter in practice: some systems that keep a chat going for a very long time keep the first few tokens in memory for exactly this reason (Part 3 tests it).

The translation results that started it all:

Summary#

• Models turn text into tokens and tokens into vectors (lists of numbers).
• Attention was invented so a model could look back at every word instead of squeezing a sentence into one summary.
• Self-attention: each token makes a query, a key and a value. Queries are compared with keys (dot products), softmax turns the scores into weights, and the weights mix the values: $\operatorname{softmax}(QK^\top/\sqrt{d} + M)\,V$.
• The causal mask stops tokens from seeing the future; every row of weights adds up to 1.
• Dividing by √d keeps scores from growing with vector length. Without it, the top weight reached 0.95; with it, 0.107.
• Multi-head attention runs several heads side by side. Our version matched PyTorch's to $1.7 \times 10^{-16}$.
• Real heads specialise: in Qwen2.5-0.5B, 68% are attention sinks, one follows the previous token (84%), and one links pronouns to nouns (6 out of 6 test sentences). BERT (Clark et al.) and Llama-2 (Xiao et al.) show the same patterns.
• The √d comes from simple statistics: the dot product of two random vectors has variance $d$, so dividing by $\sqrt{d}$ brings it back to 1.
• The transformer won because attention needs no sequential steps (fast, parallel training), at the price of $O(n^2)$ work.
• Attention's work grows with the square of the text length, and its memory (the KV cache) with the length. Parts 2 to 4 are about cutting both.

pip install torch transformers
python part1_attention.py      # writes results/part1.json
python part1_coref_check.py    # the six pronoun sentences

from transformers import AutoModelForCausalLM, AutoTokenizer
tok = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B", attn_implementation="eager")
ids = tok("The cat sat on the mat because it was tired.", return_tensors="pt").input_ids
attentions = model(ids, output_attentions=True).attentions   # one tensor per layer: (batch, heads, tokens, tokens)

References#

1. D. Bahdanau, K. Cho, Y. Bengio. Neural Machine Translation by Jointly Learning to Align and Translate. 2014.
2. A. Vaswani et al. Attention Is All You Need. NeurIPS 2017.
3. G. Xiao et al. Efficient Streaming Language Models with Attention Sinks. ICLR 2024.
4. K. Clark, U. Khandelwal, O. Levy, C. Manning. What Does BERT Look At? An Analysis of BERT's Attention. 2019.
5. Qwen Team. Qwen2.5-0.5B.