Attention Is All You Need

Vaswani et al., 2017

Explained as a program, because that's the only way I could understand it.

Download attention_from_scratch.py — everything below runs from it. NumPy only.

I'm an AI Engineer. I've shipped models, done the MLOps, and I still couldn't read this paper. The maths wasn't the problem — the framing was. Nobody told me the core idea is something I've written a thousand times.

Attention is a dictionary lookup

You know this:

d = {"cat": vec_a, "dog": vec_b}
result = d["cat"]        # exact match, or KeyError

Attention is that, except the match is fuzzy and the result is a blend:

def attention(query, keys, values):
    scores  = [dot(query, k) for k in keys]   # how well does query match each key?
    weights = softmax(scores)                 # turn into percentages summing to 1
    return sum(w * v for w, v in zip(weights, values))   # weighted blend

That's the whole paper's beating heart. Three lines.

TermMeansLibrary analogy
QueryWhat am I looking for?The question you ask
KeyWhat do I advertise about myself?The book's title
ValueWhat do I hand over if picked?The book's contents

Why this beat the RNN

An RNN is a for loop:

state = init
for token in sentence:       # strictly sequential — can't parallelize
    state = update(state, token)

Two problems. Token 50 sees token 1 only through 49 rounds of lossy state-squeezing. And your GPU sits idle, because step n needs step n−1 to finish first.

Attention is a matrix multiply. Every token talks to every other token in one shot. Path length between any two tokens: 1. And it parallelizes perfectly. That's the title — throw out the recurrence, keep only attention, and it works better and trains faster.

The one formula

Attention(Q, K, V) = softmax(QKT / √dk) V

Read it right-to-left, as pipeline stages:

StepShapeWhat it does
Q(n, dk)n queries stacked as rows — batch your lookups
QKT(n, n)Score matrix. Cell [i][j] = how much token i should care about token j
/ √dkscalarNumerical hygiene — see below
softmaxrow-wiseEach row becomes percentages summing to 1
@ V(n, dv)Blend the values using those percentages

Why divide by √dk?

Dot products of dk-dimensional vectors grow with dk. Big numbers into softmax → one weight becomes 0.9999, the rest ≈ 0 → gradients die → training stops. Dividing keeps the variance ≈ 1.

It's numerical hygiene, like normalising before a float comparison. Not deep.

Don't take my word for it. Here's the actual output:

d_k = 4    → raw score std:  1.24  → max softmax weight: 0.6125
d_k = 64   → raw score std:  5.73  → max softmax weight: 0.8319
d_k = 512  → raw score std: 25.12  → max softmax weight: 1.0000  ← dead

after dividing by √d_k, every case sits at std ≈ 1 

Multi-head: eight questions at once

One attention pass computes one weighted average — it can only track one kind of relationship. Eight heads track eight (syntax, coreference, position…) simultaneously.

The model doesn't get wider: dmodel=512 with 8 heads means each head works in 64 dims. Same total compute, split eight ways. Like sharding a query across 8 workers with different indexes, then merging.

Position: the thing attention can't do

Attention is a set operation. A weighted sum doesn't care about order — so "cat ate food" and "food ate cat" produce identical output. An RNN got order for free from its loop. We threw the loop away.

The fix: add a unique sine/cosine wave pattern to each position before attention. Not learned — just computed.

Proof, from the running code:

without positional encoding, output sets identical? True   ← order invisible
with    positional encoding, output sets identical? False  ← order now matters

That single + pe is the only thing giving a transformer any notion of word order.

The causal mask: how GPT is stopped from cheating

Set forbidden positions to −∞ before the softmax, so they come out as exactly 0. Not 0 — because 0 is a perfectly ordinary score that softmax would happily give weight to.

[[1.    0.    0.    0.    0.   ]   ← "the" sees only itself
 [0.29  0.71  0.    0.    0.   ]
 [0.283 0.36  0.357 0.    0.   ]
 [0.317 0.267 0.261 0.155 0.   ]
 [0.228 0.134 0.172 0.153 0.314]]  ← "it" sees everything before it
Remove the mask → BERT. Keep it → GPT. Same code otherwise. That's the whole architectural difference between the two model families everyone talks about.

Everything else is wrapping

PieceWhat it really is
Residual x + f(x)A skip link. The sublayer proposes a diff, not a replacement. Same reasoning as ResNet — gradients need a short path back.
LayerNormNormalise each token's vector. Numerical hygiene again.
FeedforwardA plain 2-layer MLP, 512→2048→512, applied per position.

Attention mixes information between tokens. The FFN processes it within each token. Mix, process, mix, process — that alternation is the entire stack.

def block(x):
    x = layernorm(x + multi_head_attention(x))   # tokens talk to each other
    x = layernorm(x + feedforward(x))            # each token thinks alone
    return x                                      # shape unchanged → stackable

Stack six. That's the encoder. That's the paper.


The code

Pure NumPy. No torch, no autograd, no training — just the forward pass, so you can watch the numbers move. Every function under 15 lines.

Download the full script — 60 lines plus five runnable demos. Then:

python3 attention_from_scratch.py
import numpy as np

def softmax(x, axis=-1):
    # the `- max` isn't in the paper — it's overflow safety.
    # exp(1000) is inf; exp(1000-1000) is 1. Cancels in the ratio.
    shifted = x - np.max(x, axis=axis, keepdims=True)
    e = np.exp(shifted)
    return e / np.sum(e, axis=axis, keepdims=True)


def scaled_dot_product_attention(Q, K, V, mask=None):
    d_k = Q.shape[-1]

    # (1) SCORE — scores[i][j] = how much token i should care about token j
    scores = Q @ K.T

    # (2) SCALE — keep variance ~1 so softmax doesn't saturate
    scores = scores / np.sqrt(d_k)

    # (3) MASK — -inf, not 0, so softmax gives exactly zero
    if mask is not None:
        scores = np.where(mask, scores, -np.inf)

    # (4) NORMALISE + BLEND
    weights = softmax(scores, axis=-1)   # each row sums to 1.0
    return weights @ V, weights


def positional_encoding(n_positions, d_model):
    pos = np.arange(n_positions)[:, None]
    i = np.arange(d_model)[None, :]
    angle = pos / np.power(10000, (2 * (i // 2)) / d_model)
    pe = np.zeros((n_positions, d_model))
    pe[:, 0::2] = np.sin(angle[:, 0::2])   # even dims
    pe[:, 1::2] = np.cos(angle[:, 1::2])   # odd dims
    return pe


def causal_mask(n):
    return np.tril(np.ones((n, n), dtype=bool))

Now break it

This is the part that actually makes it stick. Reading won't do it. You debug for a living — reverse-engineering from broken behaviour is your native mode, and it's far faster than re-reading section 3.2 for the fifth time.

Break thisWatch for
Delete / np.sqrt(d_k)One softmax weight goes to 1.0000 as dk grows. Gradient dead.
Remove + pe"cat ate food" and "food ate cat" become indistinguishable.
Flip np.trilnp.triuThe decoder now reads the future. It's cheating.
Use 0 instead of -np.inf in the maskMasked positions still get real weight. The mask silently does nothing.
Then read the paper — but only §3.1–3.3 and §3.5. Skip the BLEU tables, skip §6 entirely. You'll find the paper is just the comments you've already read.

On the maths anxiety

Honestly: for AI engineering — shipping systems, not writing papers — this is roughly the ceiling of maths you need, and you've just cleared it. Matrix shapes (that's array indexing), dot product as similarity, softmax as normalisation. That's it.

You don't need to derive backprop. loss.backward() is a library call, the same way you don't hand-roll TCP.

What actually separates AI engineers is different: tokenisation edge cases, KV-cache and memory maths, quantisation tradeoffs, eval design, latency and cost budgeting, retrieval quality. None of it is research maths.

Where to go next

This page gave you one thing: attention as a lookup, in code you can break. That's deliberately narrow. Four other people explain the parts I skipped, and each is better at its own angle than I'd be at copying it.

Order I'd take them in: seq2seq → illustrated transformer → Bloem → annotated. Roughly increasing maths, and each one assumes the previous. If you only read one, read Bloem — it's the closest to how an engineer thinks.