“Attention is all you need.” Great paper title. Terrible first explanation. I read the paper, read a few blog posts, and still couldn’t have told you what a query or a key actually is — only that there were three matrices and a softmax floating around somewhere.
What finally made it line up for me was 3Blue1Brown’s Attention in transformers. This is my write-up of that mechanism — one continuous story from intuition through a fully worked three-token example, then masking, the low-rank value trick, the real GPT-3 parameter counts, and multi-head attention — in the order I wish someone had handed me.
[!NOTE] If tokens and embeddings are fuzzy for you, read how GPT actually works first — it sets up why directions in vector space carry meaning. I’ll lean on that here.
What problem does attention actually solve?
Right after the embedding step in how GPT actually works, a token’s vector is just a context-free lookup. Take the word mole:
- American shrew mole — the animal
- One mole of carbon dioxide — the chemistry unit
- Take a biopsy of the mole — the skin growth
All three start with the exact same vector. That’s obviously wrong — meaning depends on what’s around it. Attention is the step where the surrounding tokens get to reach in and update the mole vector toward the right sense.
And these updates travel far. Picture feeding in a whole mystery novel that ends with:
Therefore the murderer was …
To get the next word right, that final vector — which started life just meaning was — has to have absorbed information from thousands of tokens back. Attention is the plumbing that moves that information across the whole sequence.
How do tokens ask each other questions?
The example I keep coming back to is “a fluffy blue creature.” The adjectives fluffy and blue should update the noun creature.
Imagine each noun broadcasting a question: “are there any adjectives in front of me?” And each adjective broadcasting an answer: “yes — I’m an adjective, I describe nouns.” Where a question and an answer match, information flows. Make that matching numerical and you’ve got the entire mechanism. Three players:
| Role | What it means | Matrix |
|---|---|---|
| Query | what a token is looking for | $W_Q$ |
| Key | what a token offers | $W_K$ |
| Value | the info actually passed along | $W_V$ |
What is a query?
For each token, multiply its embedding by the query matrix $W_Q$ to get a query vector:
The query lives in a much smaller space than the embedding. In GPT-3 the embedding is 12,288-dimensional but the query/key space is only 128. So $W_Q$ is a $128 \times 12{,}288$ matrix that squeezes each token’s meaning down into a compact “here’s what I’m hunting for.” The query for creature roughly says “I’m a noun, any adjectives ahead of me?”
What is a key?
In parallel, multiply every embedding by the key matrix $W_K$ to get key vectors:
Same 128-dimensional space. Think of keys as potential answers to the queries — you want fluffy’s key to land close to creature’s query.
“Close” means large dot product. For every (key, query) pair you compute $\vec{K}_i \cdot \vec{Q}_j$: large positive = strong match (token $j$ attends to token $i$), near zero or negative = unrelated.
How do dot products become attention scores?
Lay all those dot products out as a grid. If you have $T$ tokens, you get a $T \times T$ matrix of raw scores. In matrix form, every key dotted with every query is one product:
Entry $(i, j)$ is key $i$ dotted with query $j$ — “how much should token $j$ listen to token $i$?”
flowchart TB
subgraph GRID["Attention scores (dot products)"]
direction LR
R1["fluffy·creature = large"]
R2["blue·creature = large"]
R3["the·creature ≈ 0"]
end
GRID --> SCALE["divide by √d_k"]
SCALE --> SM["softmax each column"]
SM --> AP["Attention pattern<br/>(weights sum to 1 per column)"]
Why do we divide by $\sqrt{d_k}$?
Raw dot products grow with the dimension $d_k$. If each coordinate of query and key has variance near 1, their dot product sums $d_k$ terms and has variance near $d_k$. Big scores shove softmax into a nearly one-hot corner where gradients vanish. The fix is to scale each score down:
In GPT-3, $d_k = 128$, so the divisor is:
How does softmax turn scores into weights?
Raw scores can be negative and they don’t sum to 1. Softmax fixes that — applied down each column of the score matrix, so for every token $j$ the weights over all source tokens $i$ land in $[0, 1]$ and add up to 1.
For one column with scaled scores $s_1, s_2, \ldots, s_T$, the first step is exponentiate each score:
Then sum them into a normalizer:
Each weight is one exponent divided by that sum:
The full matrix of weights is the attention pattern.
What is a value?
Scores tell you how much to mix. Values are what you mix. Multiply each embedding by the value matrix to get value vectors:
Then each token’s updated output is a weighted sum of all value vectors, using the attention weights from its column:
That weighted sum is the heart of attention — you don’t pick one token; you blend them.
What is the full attention formula?
Once you’ve built it piece by piece, the original paper’s one-liner is just shorthand:
Read it left to right: dot all keys with all queries, scale, softmax each column, blend the values. That’s the whole mechanism for one head.
Can we walk through a complete three-token example?
Yes — and I think this is the step that finally made the formula feel real. I’ll use three tokens — fluffy, blue, creature — with tiny 2-dimensional query/key/value vectors so every arithmetic step fits on the page. (Real GPT-3 uses $d_k = 128$; the logic is identical.)
Step 1 — the inputs. After multiplying by $W_Q$, $W_K$, and $W_V$:
Here $d_k = 2$, so $\sqrt{d_k} = \sqrt{2} \approx 1.414$.
Step 2 — dot products for the creature column. Token creature is the one hunting for adjectives, so we build column 3 of the score matrix. Key fluffy dotted with query creature:
Key blue dotted with query creature:
Key creature dotted with query creature:
So the raw scores for creature’s column are $[3,\, 2,\, 4]$.
Step 3 — scale by $\sqrt{d_k}$. Divide each raw score by $1.414$:
Step 4 — softmax, one line per step. Exponentiate each scaled score:
Add them for the normalizer:
Divide each exponent by $Z$ to get the attention weights:
Check: $0.284 + 0.140 + 0.576 = 1.000$. Good — they sum to 1.
Step 5 — weighted sum of values. Blend the three value vectors using those weights:
First term:
Second term:
Third term:
Add them:
That’s it. Creature didn’t copy any single neighbor — it mixed information from fluffy, blue, and itself, weighted by how well each key matched its query. The biggest weight (0.576) went to creature itself, but meaningful chunks came from the adjectives. Run the same five steps for every column and you update all three tokens.
$K^{\top} Q$ decides where to read. Softmax turns scores into a distribution. Multiplying by $V$ decides what information comes back.
Why can’t the model peek at future words? (masking)
During training the model predicts the next token at every position at once — way more efficient than one prediction per pass. But that opens a cheating door: when predicting the token after position 5, it must not look at positions 6, 7, 8…
The fix is masking. Before softmax, set any score where a later token would influence an earlier one to $-\infty$:
Softmax then turns those entries to exactly zero. Each token attends only to itself and tokens before it.
flowchart LR
A["Raw score grid"] --> B["Set 'future' entries<br/>to −∞<br/>(upper triangle)"]
B --> C["Softmax →<br/>those become 0"]
C --> D["Each token attends<br/>only to itself<br/>and earlier tokens"]
In GPT-style models masking is always on — it’s what makes generation strictly left-to-right.
There’s a hidden cost too. The attention pattern has one row and one column per token, so its size is the square of the context length:
Double the context, quadruple the grid. That quadratic blow-up is the real reason long context windows are hard.
What is the low-rank value trick?
Scores tell you how much to mix; values tell you what to mix. Naively, the value map goes from embedding space back to embedding space, so one full matrix would be:
~151 million for one value matrix in one head — absurdly more than the ~1.5M each for query and key. So GPT-3 factors the value map into two smaller matrices: a value-down ($12{,}288 \to 128$) and a value-up ($128 \to 12{,}288$):
In linear-algebra terms that constrains the map to a low-rank transformation. The payoff: all four maps in one head now cost about the same.
How many parameters are in one GPT-3 attention head?
GPT-3 numbers — embedding 12,288, key/query/value dim 128:
| Matrix | Shape | Parameters |
|---|---|---|
| Query $W_Q$ | $128 \times 12{,}288$ | 1,572,864 |
| Key $W_K$ | $128 \times 12{,}288$ | 1,572,864 |
| Value-down | $128 \times 12{,}288$ | 1,572,864 |
| Value-up | $12{,}288 \times 128$ | 1,572,864 |
| Per head | ≈ 6.3 million |
All four come out the same size — which is the whole point of that low-rank factorization.
[!NOTE] Self- vs cross-attention: everything above is self-attention — queries, keys, and values all come from one sequence. Cross-attention (translation, speech-to-text) is nearly identical; the only difference is keys and values come from a different sequence than the queries.
What is multi-head attention?
One head learns one kind of relationship — say adjective→noun. Real language is richer than that, so a transformer runs many heads in parallel, each with its own $W_Q, W_K, W_V$:
- GPT-3 uses 96 heads per block.
- 96 × ~6.3M ≈ 600 million parameters per multi-head attention block.
flowchart TB
E["Token embeddings"] --> H1["Head 1<br/>(grammar?)"]
E --> H2["Head 2<br/>(subject–object?)"]
E --> H3["…"]
E --> H96["Head 96<br/>(coreference?)"]
H1 --> CAT["Combine all heads"]
H2 --> CAT
H3 --> CAT
H96 --> CAT
CAT --> OUT["Updated embeddings"]
Each head makes its own attention pattern and its own value mix; the results get concatenated and projected back into the embedding space. Different heads end up tracking different things — grammar, coreference (“Harry” → Harry Potter vs Prince Harry depending on nearby words), physical implications (“they crashed the car”), and so on.
For tensors: input $X$ has shape $(B, T, d)$; per head, $Q, K, V$ have shape $(B, h, T, d_h)$; scores have shape $(B, h, T, T)$; the mixed result has shape $(B, h, T, d_h)$. Concatenating all $h$ heads restores $(B, T, d)$.
How does attention fit into the full transformer?
A transformer doesn’t run attention once. Data flows through attention, then an MLP (feed-forward) block, then attention again, then another MLP — dozens of times. Each round gives the now-more-nuanced embedding another chance to be refined by its now-more-nuanced neighbors. That depth is where the model’s apparent “reasoning” quietly piles up.
flowchart LR
A["Embeddings"] --> B["Multi-head<br/>Attention"]
B --> C["MLP /<br/>Feed-Forward"]
C --> D["Multi-head<br/>Attention"]
D --> E["MLP /<br/>Feed-Forward"]
E --> F["… ×96 layers …"]
F --> G["Final embeddings<br/>→ predict next token"]
Attention retrieves and mixes information across tokens, but it can’t do all the computation by itself — the surrounding MLPs, residual paths, layer normalization, and depth matter just as much. Standard attention also stores a $T \times T$ score matrix, so memory grows quadratically with context length.
Where does this leave you?
The thing I didn’t expect is how small the core idea turns out to be: dot product → scale → softmax → weighted sum, repeated across heads and layers. Queries, keys, and values are just “what I’m looking for / what I offer / what I’ll hand over.” Masking keeps it causal. The low-rank value trick keeps parameter counts sane. The quadratic grid is the price you pay for letting every token talk to every other token.
After sitting with this, the original Attention Is All You Need paper finally reads like a summary of something I already understand instead of a wall of notation.
If you haven’t read the setup, start with how GPT actually works.
Watch: Attention in transformers — 3Blue1Brown · full text version on 3blue1brown.com · part of the Neural Networks series.
Support the writing
If this post helped, a coffee keeps the deep dives, paper picks, and news digests coming.
Buy Me a Coffee