Attention, Step by Step

Attention is what makes an input token's vector depend on the tokens around it — every token reads every other, then rewrites itself as a blend of what it found.

Three parts, each one fixing what the one before it cannot do. What stands at the end is the causal multi-head self-attention of the original transformer, exactly as GPT-2 runs it — and still what the attention layers of today's models compute, under a decade of efficiency work.

Part 1: Self-attention

This picks up where the pretraining data pipeline left off: text split into tokens, each token embedded, and a position embedding summed in. Those input embeddings — $x^{(i)} = t^{(i)} + p^{(i)}$, a token vector plus a position vector — are the input to everything below: one sequence of $T$ tokens stacked into $X$, [T, d_in]. Real training feeds a batch of sequences at once, [batch_size, T, d_in], but attention never mixes one sequence with another. Each is projected, scored and blended on its own, so every shape here drops the batch dimension.

A real sequence is a window, not a sentence. The pipeline cuts one wherever the sliding window lands — 1,024 tokens wide in GPT-2 — so it can start and end mid-thought. Attention does not care. It runs the same over any $T$ tokens.

The self is load-bearing. Attention began in translation, where a decoder read across to the source sequence — cross-attention. Self-attention has one sequence only: queries, keys, and values all come from the same tokens.

1.1: Self-attention input

1.2: Self-attention with no trainable parameters

Attention scores. Calculate the dot product between every pair of input embeddings, which comes out high when the two vectors point in the same direction.

Attention weights. Each row of scores is normalized using softmax to obtain weights that sum up to 1. This makes them interpretable as probabilities or relative importance.

Context vectors. Each token's context vector is the weighted sum of all the input embeddings, obtained by multiplying each input embedding by its corresponding attention weight.

Things to notice:

• The attention score matrix is symmetric. Hover or tap any score to see the same value at its transposed position. The score for the pair $(i, j)$ equals the score for $(j, i)$, because $x^{(i)} \cdot x^{(j)} = x^{(j)} \cdot x^{(i)}$. The matrix cannot represent a one-way relationship: if token $i$ is relevant to token $j$, then token $j$ is exactly as relevant to token $i$. The position embeddings do not break this symmetry — they make the input embeddings depend on position. Switch them off and a repeated token produces identical rows.

• Nothing here is trainable. The output is fully determined by the input embeddings, so gradient descent has nothing to update. A token also scores highest against itself — a vector's dot product with itself is its squared length — so the diagonal usually dominates and each context vector stays close to the input embedding it came from.

1.3: Self-attention with trainable parameters

Introduce a trainable weight matrix $W$, project every input through it, and take the dot products in that projected space. Training now has something to update, but the score matrix symmetry survives untouched — the score for the pair $(i, j)$ is $x^{(i)} W W^\top x^{(j)\top}$, and $W W^\top$ is symmetric whatever $W$ is.

Breaking this symmetry takes two different projections. Score $x^{(i)} W_q$ against $x^{(j)} W_k$: the pair $(i, j)$ now runs through $W_q W_k^\top$, the pair $(j, i)$ through $W_k W_q^\top$. Two matrices, and with them two roles for a token to play — a query (what am I looking for?) when it is the one asking, a key (what do I contain?) when it is the one being asked.

Nothing so far forces a third matrix. $W_v$ separates what a token is matched on from what it hands over — its value. Queries are searched against keys, and what comes back are values — the names come from database lookup, except a database returns the one value whose key matched, while attention blends all of them by how well each matched.

d_out is the size knob: each of the three matrices is [d_in, d_out], so attention costs 3 × d_in × d_out weights — 1.8M per layer at GPT-2's 768 — and whatever width you leave is the width everything below runs at, part 3 included.

1.4: Scaling and softmax normalization

Scores are divided by $\sqrt{d_k}$, the square root of the key width. Dot products grow with that width, and large scores drive softmax toward a one-hot row, where gradients shrink to near zero and training stalls. The division holds them in a range where they keep training. Flip the switch here and little changes — the divisor stays small at a single-digit $d_k$. One GPT-2 head runs at $d_k = 64$, where the scaling stops being optional. This is why the mechanism is called scaled dot-product attention.

The softmax then normalizes each row to sum to 1, as in step 1.2 — except the scores now come through $W_q$ and $W_k$. That is the attention weight matrix used to build context vectors.

1.5: Context vectors

The attention weights say how much of each value vector to take. Each token's context vector $z^{(i)}$ is the weighted sum of every token's value, and all $T$ of them stack into $Z$, [T, d_out] — the output of the attention layer. Softmax left the weights positive and summing to 1, which is what makes that sum a weighted average: each dimension of $z^{(i)}$ stays within the range that same dimension spans in $V$.

Scaled dot-product attention, in one line:

Every row reads the whole sequence. Row $i$ of $Z$ is a weighted sum over every position $j$, the ones after $i$ included. But row $i$ is what the model uses to predict token $i+1$, so the vector doing the predicting already contains $v^{(i+1)}$ — a projection of the token it is meant to guess. Pretraining runs every position at once, so the loss falls on a shortcut that is gone at generation time, where the positions after $i$ do not exist yet.

Part 2: Causal attention

Next-token prediction uses a specialized form of self-attention — causal attention, also known as masked attention. With the causal mask on, every score at $j > i$ is set to $-\infty$ before the softmax, and since $e^{-\infty} = 0$ those weights come out exactly zero, each row still summing to 1 over the positions that are left. The diagonal stays: a token attends to itself.

One side effect is worth knowing. Repeat a word in the text and the two copies get different context vectors, because the second one attends over a longer prefix than the first.

This matrix is also where GPT-2 applies its regularization. On every training step a random subset of these weights is zeroed before they blend the values — attention dropout, covered in the dropout chapter. At inference it is off, so the weights are exactly as described above.

Part 3: Multi-head attention

Causal attention produces one softmax probability distribution per token: a single set of weights over the tokens it can attend to. That one weighting has to carry everything the token needs to read from the sequence. Several distributions let it read the sequence several ways at once, and pretraining is what settles on what each way looks for.

Multi-head attention runs several instances of the causal attention mechanism at once, each with its own $W_q$, $W_k$ and $W_v$. Each instance is a head, and each produces its own distribution per token. Every head reads the whole input and writes head_dim columns of the output: d_out is divided among the heads rather than multiplied by them. They stay independent until their context vectors are concatenated back to the full width.

The concatenation then passes through out_proj, a fully connected linear layer, so that every dimension of the output draws on all the heads rather than one — standard in implementations, though not strictly necessary.

Takeaway

Attention is a weighted average. Each token's context vector $z^{(i)}$ blends the value vectors of the tokens it is allowed to see. The weights come from one dot product per pair — the reader's query against the other token's key. Everything after that — the scaling, the mask, the extra heads — changes how the weights come out, not what happens once they exist.

The learned part is a short list: $W_q$, $W_k$, $W_v$ per head, plus out_proj. The attention weight matrix is not on it. That matrix is rebuilt for every sequence, out of the sequence's own content. A linear layer applies the same numbers to whatever arrives; attention builds a new mixing pattern for every input. Training fixes the projections, not the pattern they produce.

Nothing in the mechanism is decorative. Plain dot products are symmetric, so queries and keys need separate matrices — how much $i$ needs $j$ is not how much $j$ needs $i$. Nothing forces $W_v$ — it is what lets a token hand over something other than its own input embedding. Dot products grow with width, so scores are divided by $\sqrt{d_k}$ to keep the softmax from saturating. Every row would otherwise read past its own position, so everything above the diagonal goes to $-\infty$. A single distribution can weigh the sequence only one way, so heads run several in parallel and out_proj recombines them.

Two limits come with it. The score matrix is $T \times T$, so doubling the sequence quadruples the scoring work — which is what makes long context expensive. And attention itself has no notion of order. Apart from the mask, a weight depends only on the two vectors being compared, never on how far apart they sit or which came first. Order has to be in the input embeddings already, which is what $p^{(i)}$ is for.

That cost buys parallelism. The attention mechanism has no step that waits on a previous one: every score is independent, the softmax runs row by row, and a whole sequence resolves in a few matrix multiplications rather than $T$ steps taken in order. The mask is what makes running all positions at once safe. That parallelism is what lets pretraining keep GPUs and TPUs fully utilized.

With d_out equal to d_in — 768 in GPT-2 — $Z$ leaves in the shape $X$ arrived in, so the layer stacks. Add a feed-forward layer, residual connections and normalization around it and that is a transformer block. GPT-2 stacks twelve. Everything in that wrapper acts on one position at a time: attention is the only place in the model where tokens see each other.

References

Sebastian Raschka, Build a Large Language Model (From Scratch) (Manning, 2024). Chapter 3 builds all three parts in PyTorch, and the code exported by each step here follows its notation closely.

Ashish Vaswani et al., Attention Is All You Need (2017). The paper the mechanism comes from.

The input embeddings this page starts from are where the data pipeline ends — tokenizing, windowing, and embedding the text before any of the above runs.