Reading Attention Is All You Need: A Practical Guide for Implementers

Fifteen pages, eight equations, and a title that turns out to be an accurate description of the whole thing. The first pass is hard not because the math is exotic but because the paper assumes you already hold three shapes in your head at once: batch, sequence, and feature. Once those click, section 3.2 is a matrix multiply with a softmax in the middle.

What follows is an order that works. Read every equation as a statement about tensor shapes, then write the few lines of code that satisfy it. Multi-head attention — the part everyone remembers and most people mis-implement — is smaller than you expect.

Start With Shapes, Not Symbols

The paper writes Q, K, and V as matrices in RT×d_k and leaves the batch dimension implicit. In code you carry it everywhere, so fix the notation before reading further. For the base configuration, d_model is 512, h is 8, and d_k = d_v = d_model / h = 64. The feed-forward inner dimension is 2048. Write those numbers at the top of your notes.

• x: (B, T, d_model) — a batch of sequences of embedded tokens.
• W_Q, W_K, W_V: (d_model, d_k) each, or one fused (d_model, 3·d_model) projection.
• Q, K, V: (B, T, d_k) per head after the split.
• Scores: (B, T, T) — how much each token attends to each other token.
• Output: (B, T, d_model) after concatenating heads and projecting.

A useful habit: annotate shapes above every symbol in the margin of the PDF. When an equation stops being ambiguous, it stops being intimidating.

Do the arithmetic once on paper and the numbers stop being abstract. A single batch element with T = 512 tokens produces, per head, a 512×512 score matrix — 262,144 floats that describe nothing more than which tokens look at which. Across 8 heads that is roughly 2 million scores for one sequence, and the softmax runs row-wise over the last axis, so each of the 512 query tokens gets its own distribution over 512 key positions. That is the entire "attention" you keep hearing about. Squint at the row and you are looking at a probability vector, not a mystery.

Scaled Dot-Product Attention in Six Lines

Attention(Q, K, V) = softmax(QKᵀ / √d_k)V. That is the whole mechanism. QKᵀ produces a T×T grid of scores, the softmax turns each row into a distribution over positions, and multiplying by V gives a weighted average of value vectors. Each token asks a question; every token answers; the answers get averaged by relevance.

Here it is, with masking included because you will need it immediately:

  def attention(q, k, v, mask=None):

    scores = q @ k.transpose(-2, -1) / math.sqrt(q.size(-1))

    if mask is not None:

      scores = scores.masked_fill(mask == 0, float("-inf"))

    weights = scores.softmax(dim=-1)

    return weights @ v, weights

Notice what the function does not contain: no loops over the sequence, no recurrence, no gating. The transpose is the only fiddly part, and it exists purely so the last two axes line up for the matmul. If you keep returning weights alongside the output, you get a free interpretability hook — plot a row and you can literally see a token's attention distribution.

Why the Square Root Matters

Each score is a dot product of two d_k-dimensional vectors whose components are roughly unit variance, so the score has variance around d_k. At d_k = 64, raw scores routinely land in the tens. Feed numbers that large into a softmax and it saturates: one weight goes to ~1, the rest to ~0, and the gradient through that row collapses to almost nothing. Dividing by √d_k drags the variance back to roughly 1, which keeps the softmax in the responsive part of its curve where gradients still flow.

You can verify this in about ten seconds. Compute a random 64-dimensional dot product, print its standard deviation, and compare it to a 512-dimensional one — the gap is exactly why the paper writes √d_k and not d_k. Scale by d_k instead and training will stall in a way that looks like a learning-rate problem but isn't.

Multi-Head Attention Is Just a Reshape

The famous "multi-head" block is the same attention function run h times in parallel with different learned projections. The only twist is that each head is smaller: d_k = d_model / h = 64, so the total compute stays roughly constant whether you use one wide head or eight narrow ones. The split is a reshape and a transpose, nothing more:

  def split_heads(x, h):

    B, T, d_model = x.shape

    return x.view(B, T, h, d_model // h).transpose(1, 2)

After attention, you transpose back, reshape to (B, T, d_model), and pass through the output projection W_O of shape (h·d_v, d_model). Two things trip people up here. First, the reshape must split the last axis into (h, d_k) — getting the order wrong silently scrambles which dimensions belong to which head. Second, `transpose` alone doesn't move memory; call `.contiguous()` before a `.view()` if your framework complains.

The Parts Everyone Skips: Positions and Feed-Forward

Attention is permutation-equivariant: shuffle the input tokens and the output shuffles with them. The paper fixes this by adding sinusoidal positional encodings to the embeddings before the first block — fixed, not learned, which is why longer sequences at inference time don't break the way they would with learned position tables. In code this is one addition of a precomputed (1, T, d_model) tensor, and forgetting it is the single most common reason a from-scratch transformer trains to garbage.

The other skipped piece is the position-wise feed-forward network: two linear layers with a ReLU between them, applied identically to every token independently. It expands 512 → 2048 → 512 and holds roughly two-thirds of the model's parameters, despite getting a single sentence in the paper.

Five Bugs to Check Before You Trust Your Implementation

1. Masking after the softmax. Zeroing weights post-softmax leaves a distribution that no longer sums to 1. Mask with -inf before.
2. Scaling by d_k. It must be √d_k, and per head, not per d_model.
3. Wrong reshape order. Check that split_heads really splits the feature axis into (h, d_k).
4. Missing positional encodings. Without them the model cannot tell "dog bites man" from "man bites dog."
5. Causal mask in the wrong place. Decoder self-attention needs the upper-triangular mask; encoder attention does not.

A Reading Order That Actually Works

Read 3.2, then 3.2.1, then 3.3. Stop and implement attention. Only after that go back to 3.1 for the encoder/decoder stack, 3.4 for embeddings, and 3.5 for positions. Reproduce the base model hyperparameters from Table 3 in your own config file, and read the ablation in section 3.2.2 last — it only makes sense once you have something to ablate. Two hours with a debugger beats two weeks of rereading. The title isn't a slogan; it's a spec.

Photo: BookBabe / Pixabay