Attention is a lookup table you learn

The standard explanation of attention is that the model “pays attention to relevant parts of the input.” This is true in the same way that a database “cares about” your query. It describes the outcome and hides the mechanism, and because it hides the mechanism, everything built on top of it — multi-head, positional encodings, KV caching, the entire inference cost model — arrives later as a list of unmotivated tricks.

Here is the version that generalizes: attention is a dictionary lookup where the keys are learned, the query is learned, and the match is soft.

Start with a dictionary§

An ordinary lookup is exact. You have pairs of keys and values, you arrive with a key, you get one value back:

table = {"paris": "france", "tokyo": "japan"}
table["tokyo"]  # -> "japan"

Two properties make this useless inside a neural network. It is discrete — a key either matches or it does not — so there is no gradient. And the keys are fixed in advance, so nothing about the lookup can be learned.

Fix both. Make keys and queries vectors instead of strings, and replace exact matching with a similarity score. A query now matches every key a little, and the result is a weighted average of all the values rather than one of them:

import numpy as np

def soft_lookup(query, keys, values, temperature):
    scores = keys @ query / temperature       # (n_keys,)
    weights = np.exp(scores - scores.max())
    weights /= weights.sum()                  # softmax
    return weights @ values                   # (d_value,)

That is attention. Everything else is bookkeeping about where the queries, keys and values come from.

Where they come from§

In self-attention, all three are linear projections of the same input sequence. Given a sequence of token representations $X \in \mathbb{R}^{n \times d}$:

$$Q = XW_Q, \quad K = XW_K, \quad V = XW_V$$

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

The three weight matrices are the entire learned content of the layer. $W_K$ decides what each position advertises about itself. $W_Q$ decides what each position is looking for. $W_V$ decides what gets passed along when the two agree.

The $\sqrt{d_k}$ is the temperature from the snippet above, and it is not a tuning knob. If $q$ and $k$ have components with unit variance, their dot product over $d_k$ dimensions has variance $d_k$. Without the scaling, the softmax saturates as the model gets wider, gradients vanish, and the layer stops learning.^[1] The one line of arithmetic in the paper that looks like a detail is the reason wide models train at all.

Why more than one head§

A single attention layer computes one weighted average per position. That is one relationship. Language needs many at once — the subject of this verb, the antecedent of this pronoun, the open bracket this one closes — and a single softmax cannot hold several sharp distributions simultaneously; it has to compromise between them.

So run $h$ of them in parallel on lower-dimensional slices, then concatenate:

The parameter count is unchanged. You are not spending more; you are spending the same budget on twelve narrow questions instead of one wide one. This turns out to be the better trade almost everywhere, and it is the first design decision in this series that you could be asked to defend in a review.

What this costs you§

Every query attends to every key, so the score matrix is $n \times n$. Compute and memory both scale as $O(n^2 d)$ in the sequence length. Doubling the context window quadruples the attention cost.

That single fact sets up most of the rest of this series. It is why context windows were 2,048 tokens for years. It is why FlashAttention mattered enough to change hardware utilization industry-wide. It is why serving a long conversation is expensive in a way that serving a short one is not, and why the KV cache — which we get to in part seven — is the difference between a viable product and a research demo.

Where this leaves us§

Attention is a soft, learned dictionary. Multi-head attention is several of them sharing a budget. The quadratic cost is not an implementation detail but the central economic fact of the architecture.

Next: the rest of the block — residual streams, layer norm placement, and the feed-forward layers that hold two thirds of the parameters and get one tenth of the attention.