The Annotated Transformer — Visual Explainer

The core idea

The whole architecture is a few `forward()` methods.

Reading the Transformer paper and implementing it are two different things. The paper compresses the model into a handful of equations; the gap between those equations and a working model is the tensor bookkeeping the math leaves out — shapes, transposes, masks, scaling constants. That gap is exactly where most people get stuck.

The Annotated Transformer closes the gap by interleaving the paper's own text with a complete, runnable PyTorch version, so the abstract math and the concrete tensor operations sit side by side. The surprise is how small the code is: scaled dot-product attention is five lines, a full encoder–decoder layer is three calls, and the entire model is built by one make_model() function. This page reads the paper through that lens — every equation, shown as the code that runs it.

If you can read a forward() method, you can read the Transformer. The math was the spec; the code is the machine.

PART A

The skeleton — a paper you can run

What the document is, and the shape of the code inside it.

Step 1 · The document is the program

The paper goes in; a working translator comes out.

The Annotated Transformer is not an article about the Transformer — it is the Transformer, as a runnable notebook. The paper's prose sits in the margins; between the paragraphs, the actual PyTorch that implements each idea executes top to bottom. Run the whole document and you get a translator. Our running example, the same sentence the paper uses, will follow us through all twelve steps:

Our running example. Every line of code on this page will be shown doing something to this sentence.

The original 2017 paper packs this whole machine into a few pages of equations. Implementing it means filling in everything the equations leave implicit — what shape each tensor is, when to transpose, where to mask, which constant to divide by. The Annotated Transformer's move is to put the code right next to the prose: read a sentence of the paper, then read the lines that make it true.

So the layout of the document, top to bottom, is the layout of this lesson: a stack of small classes, each one a sentence of the paper made executable. Here is the shape of that stack — the names we'll meet, in the order the document builds them:

Each sentence of the paper has a class that implements it. This page walks that mapping, left to right.

From the document: “This post presents an annotated version of the paper in the form of a line-by-line implementation. … This document itself is a working notebook, and a completely usable implementation.”

Step 2 · The EncoderDecoder class

The top-level model is one tiny class that wires five parts together.

You saw the mapping from paper to code. The very first class in the document is the skeleton everything else hangs off — so we start there.

The paper opens with one sentence: “most competitive models have an encoder–decoder structure.” In code that becomes a class with five members and three short methods. The whole machine is just two towers — a reading tower and a writing tower — and this class holds them, plus the embeddings that feed each side and the generator that turns the final vectors back into words.

class EncoderDecoderannotated-transformer · Part 1

class EncoderDecoder(nn.Module):
    """A standard Encoder-Decoder. Base for this and many other models."""
    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
        self.encoder, self.decoder = encoder, decoder
        self.src_embed, self.tgt_embed = src_embed, tgt_embed
        self.generator = generator              # vectors → word scores

    def forward(self, src, tgt, src_mask, tgt_mask):
        # read the source, then write the target conditioned on it
        return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)

    def encode(self, src, src_mask):
        return self.encoder(self.src_embed(src), src_mask)

    def decode(self, memory, src_mask, tgt, tgt_mask):
        return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)

encode reads the whole English sentence into a block of vectors called memory; decode writes German one piece at a time, looking back at that memory. That is the entire top-level control flow.

Now unwrap each tower in three zooms — each picture is the same machine, drawn with one more level of detail. Each tower is six identical floors (clones(layer, N=6)), and each floor does only two or three things.

Zoom 1 · two towers

The box is just two towers: one reads the English, one writes the German.

Zoom 2 · six floors each

Each tower is six identical floors (N = 6). The reading tower's final output is handed to every floor of the writing tower.

Zoom 3 · inside one floor

Dissect one floor: an encoder floor does only [self-attention → feed-forward]; a decoder floor adds one block in between.

The trick that makes the code short

Six identical floors are not written six times. One helper — clones(layer, N), a one-line nn.ModuleList of deep copies — builds the whole stack. The same idea repeats everywhere: define one small block, clone it. That is why the model is a few hundred lines, not a few thousand.

Per-layer costs, from Table 1 of the paper. O(1) sequential ops is the parallelism win; O(1) path length is what makes long-range word links easy — and why self-attention, not recurrence, became the block to implement.
Layer type	Complexity	Sequential ops	Max path length
Recurrent	O(n·d²)	O(n)	O(n)
Convolutional	O(k·n·d²)	O(1)	O(log_k n)
Self-attention	O(n²·d)	O(1)	O(1)

Step 3 · The Embeddings class

A lookup table turns words into vectors — times one curious constant.

The skeleton needs something to read. Before our sentence can enter the bottom floor, there's a conversion step: networks compute with numbers, not letters.

Each word is swapped for its embedding — a learned list of 512 numbers (the d_model dimension) that encodes the word's meaning, where words used similarly end up with similar lists. In code this is just one nn.Embedding lookup table. So “I am a student” enters the machine as four strips of 512 numbers each:

Four words → four vectors of d_model = 512. Cell shade stands in for the numeric value; the real lists are learned during training.

class Embeddingsannotated-transformer · Part 1

class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        self.lut = nn.Embedding(vocab, d_model)   # the lookup table
        self.d_model = d_model

    def forward(self, x):
        return self.lut(x) * math.sqrt(self.d_model)  # ← the constant

Why the * math.sqrt(d_model)? The paper says embeddings are scaled by √d_model = √512 ≈ 22.6. It keeps the embeddings' magnitude in the same range as the positional encodings added next — a one-character detail the equations hide, but the code must spell out.

One detail that makes the whole tower work: only the bottom floor sees raw embeddings. Every floor above consumes the floor below's output — four vectors in, four vectors out, always the same shape. Same shape in, same shape out is exactly what lets you stack six identical floors (and clone one layer six times).

Two ingredients are still missing. Nothing here records word order (we fix that in Step 8 with PositionalEncoding), and at the far end we'll need to turn a vector back into a word (Step 11, the Generator). Park both questions — first, the heart of the machine: the five-line attention function.

PART B

Each equation, as the code that runs it

Feel the mechanism · compute it by hand · then read the five lines that do it

Step 4 · What attention() must compute

When the machine reads “its”, what does “its” mean?

Embeddings done, the heart of the model is the attention function. Before reading its five lines, feel what they have to compute — the rest of Part B builds straight to that code.

Take a sentence from the paper's own appendix: “The Law will never be perfect, but its application should be just …”. When the machine reads its — possessive of what, exactly? A useful representation of “its” must contain some “Law”.

That is self-attention: while encoding a word, the model scores every other word for relevance and blends the useful ones into this word's representation. Here, “its” pulls in “Law” and “application”. Try it:

Hover / tap an underlined word

The Law will never be , but application should be — this is what we are , in my opinion .

Arc thickness = attention weight. The paper's appendix shows two trained heads attending sharply from its to Law and to application — resolving the pronoun with no grammar rules provided.

Second exhibit · a nine-word reach

It is in this spirit that a majority of American have passed new since 2009 the registration or voting more difficult .

Also from the appendix: in encoder layer 5 of 6, heads attend from making across nine words to more and difficult — completing the phrase “making … more difficult”.

Arc weights here are illustrative; the qualitative patterns are taken from the paper's appendix visualizations of real trained heads.

Step 5 · Pass 2 — compute it by hand

Six small steps. Real numbers. No magic left.

You've felt what attention does. Now do it yourself, with actual numbers, so there's nothing left to take on faith.

We shrink the sentence to two words — Attention works — and follow word #1 (“Attention”) through one attention pass, at the real model's dimensions. Six rows; each one is a single arithmetic move.

Make a query, a key, and a value for each word

Multiply each word's 512-number embedding by three learned matrices — W^Q, W^K, W^V — to get three 64-number vectors. The query asks “what am I looking for?”, the key advertises “what do I contain?”, and the value is “what I hand over if you pick me.”

x₁ “Attention” · 512 → q₁ 64 k₁ 64 v₁ 64 · x₂ “works” · 512 → q₂ 64 k₂ 64 v₂ 64

Score: dot word #1's query against every key

A dot product multiplies two vectors position by position and adds it all up — a big result means the directions align. It's our relevance meter: how well does each key answer q₁'s question?

q₁ · k₁ = 96 q₁ · k₂ = 80 ← raw scores

Scale: divide by 8

8 = √64, the square root of the key length. With 64 positions summed, dot products grow large — and huge scores park the next step (softmax) in a flat region where learning gradients die. Dividing by √d_k keeps the numbers in the healthy zone.

96 ÷ 8 = 12 80 ÷ 8 = 10

Softmax: turn scores into weights

Softmax squashes any list of scores into positive weights that sum to 1 — by exponentiating each score and dividing by the total. Note what it did to our gap: a 12-vs-10 score lead became an 88-vs-12 weight lead.

softmax(12, 10) = 0.88 0.12

Weigh each value by its weight

Multiply every word's value vector by its weight. The losing word isn't erased — it's shrunk toward zero, still faintly present.

0.88 ×v₁ · 0.12 ×v₂

Sum — and you have z₁

Add the weighted values. The result z₁ is the new representation of “Attention”: mostly its own content, plus a measured dose of “works”.

0.88·v₁ + 0.12·v₂ = z₁ 64

Drag the scores — watch row 4 happen live

score s₁ 12

score s₂ 10

weight w₁

88.1%

weight w₂

11.9%

Softmax exaggerates gaps: a small score lead becomes a big weight lead. Push the scores equal and the weights split 50/50; open a gap of 4+ and the winner takes nearly everything.

Takeaway

A token mostly keeps itself, but measurably blends in what it needs from others. That blend — recomputed fresh for every word, in every layer — is the entire trick.

Checking row 4: e¹²/(e¹² + e¹⁰) = 1/(1 + e⁻²) ≈ 0.88. The same pass runs for word #2 with q₂ — every word gets its own z.

Step 6 · The five-line attention() function

Your six steps, written out as five lines of PyTorch.

You did six steps for one word. Stack every word's q, k and v into matrices Q, K, V — one row per word — and the paper's one-line formula falls out. Here it is, with the code beside it:

The paper's equation Attention(Q, K, V) = softmax(QKᵀ / √d_k) V

becomes one matmul + softmax + matmul scores = Q @ Kᵀ / √d_k ; p = softmax(scores) ; out = p @ V

def attentionannotated-transformer · the keystone

def attention(query, key, value, mask=None, dropout=None):
    "Compute 'Scaled Dot Product Attention'"
    d_k = query.size(-1)                                  # 64
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)   # rows 2–3
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)      # block forbidden positions
    p_attn = scores.softmax(dim=-1)                       # row 4
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn           # rows 5–6

Read it against your worksheet: line 3 is rows 2–3 (score, then ÷ √d_k); line 6 is row 4 (softmax); the last line is rows 5–6 (weigh the values, sum). The whole mechanism is one short function — and it returns p_attn too, so you can see the weights (that is the visualization at the very bottom of the document).

Notice the two implementation details the equation never mentions. key.transpose(-2, -1) is the ᵀ in QKᵀ — a tensor reshape, not a math symbol. And masked_fill(…, -1e9) is how you forbid a position: set its score to nearly −∞, so after softmax its weight is effectively zero (we'll use this in the decoder, Step 9). QKᵀ dots every query against every key at once, producing the all-pairs score table:

One matrix multiply scores all pairs at once — no loops, no waiting. This is the shape GPUs are built for, and a big part of why training parallelizes (Step 2's key property, now in matrix form).

Map every line of code to a worksheet row

Q = XW^Q, K = XW^K, V = XW^V → row 1 (built one level up, in MultiHeadedAttention)
matmul(query, key.transpose(-2,-1)) / sqrt(d_k) → rows 2–3
scores.softmax(dim=-1) → row 4
matmul(p_attn, value) → rows 5–6

Real dimensions: X is n × 512; W^Q, W^K, W^V are 512 × 64, so Q, K, V are n × 64 and the score table is n × n. The code is shape-agnostic — it works for one head or eight, because matmul broadcasts over the leading batch and head axes.

Step 7 · Multi-head attention

The real model runs eight of these at once.

The attention function runs one head. The MultiHeadedAttention class wraps it to run eight at once — and the wrapping is almost all reshaping.

Look back at Step 4: “its” needed both “Law” and “application”. A single head has one weight budget that sums to 1, so it must split — averaging the two relationships into a muddier blend of both. The paper's fix: run h = 8 independent attention units in parallel, each with its own learned W^Q, W^K, W^V. Each head can develop its own habit — one tracks pronouns, another tracks word order — then the eight results are concatenated and projected back to 512 numbers by one more learned matrix, W^O.

Pick a head, watch the pattern switch

Looks at itself

Each position attends mostly to its own token — keeping identity intact.

Habits illustrative — the appendix shows real heads specializing in syntax-like and semantics-like roles. With one head, averaging inhibits this.

class MultiHeadedAttentionannotated-transformer · forward()

def __init__(self, h, d_model, dropout=0.1):
    assert d_model % h == 0
    self.d_k = d_model // h           # 512 // 8 = 64
    self.h = h                        # 8 heads
    self.linears = clones(nn.Linear(d_model, d_model), 4)   # W^Q, W^K, W^V, W^O

def forward(self, query, key, value, mask=None):
    nbatches = query.size(0)
    # 1) project, then SPLIT 512 into 8 heads × 64 via view + transpose
    query, key, value = [
        lin(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
        for lin, x in zip(self.linears, (query, key, value))
    ]
    # 2) one call to the SAME five-line attention() — now over 8 heads
    x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)
    # 3) "concat": glue the heads back into 512, then the final W^O
    x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)
    return self.linears[-1](x)

“Eight heads” is mostly a reshape. One view(…, h, d_k) slices the 512-wide vector into 8 lanes of 64, transpose moves the head axis up front, and the very same attention() from Step 6 runs over all of them. The “Concat” in the paper's formula is literally a view back to 512. Four nn.Linear layers cover W^Q, W^K, W^V and the output W^O.

X → 8 × (Q, K, V) → 8 z-matrices → concat → ×Wᴼ → Z. Each head works in d_k = d_v = 512/8 = 64 dimensions, so eight heads cost about the same as one full-width head.

From the ablations: a single head is 0.9 BLEU worse than the 8-head base; 32 heads is also worse than 8. More habits help — up to a point.

Step 8 · class PositionalEncoding

Order, restored — and precomputed once into a buffer.

Attention is done. But it pays the debt from Step 3: nothing computed so far cares about word order — shuffle “I am a student” into “student a am I” and every dot product in Step 5 comes out identical, because attention compares contents, not positions.

The fix is almost embarrassingly simple: before the first floor, add a position-dependent vector to each word's embedding — same length, 512 numbers — so that “am, the 2nd word” enters the tower as a slightly different vector than “am, the 5th word” would. The paper builds these vectors from sines and cosines of different wavelengths. Here's a toy version with d_model = 4, interleaving sin and cos exactly like the paper:

pos dim 0
sin(pos/1) dim 1
cos(pos/1) dim 2
sin(pos/100) dim 3
cos(pos/100) 0 · I 0.00 1.00 0.00 1.00 1 · am 0.84 0.54 0.01 1.00 2 · a 0.91 −0.42 0.02 1.00 3 · student 0.14 −0.99 0.03 1.00

Read it column by column. Dims 0–1 spin fast — they change completely between neighboring words. Dims 2–3 barely move — they encode coarse position. The real model has 512 dims spanning wavelengths from 2π up to 10000·2π.

PE(pos, 2i) = sin(pos / 10000^(2i/d_model)), cos on odd dims. Fast waves distinguish neighbors; slow waves give the big picture.

class PositionalEncodingannotated-transformer · Part 1

def __init__(self, d_model, dropout, max_len=5000):
    pe = torch.zeros(max_len, d_model)                  # table: 5000 positions × 512
    position = torch.arange(0, max_len).unsqueeze(1)
    div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
    pe[:, 0::2] = torch.sin(position * div_term)        # even dims → sin
    pe[:, 1::2] = torch.cos(position * div_term)        # odd dims  → cos
    self.register_buffer("pe", pe)                      # not a parameter: never trained

def forward(self, x):
    x = x + self.pe[:, : x.size(1)].requires_grad_(False)   # just ADD it on
    return self.dropout(x)

Two implementation choices the math hides. The 10000^2i/d divisor is computed in log space (exp(… · −log 10000 / d)) for numerical stability, and the whole table is precomputed once and stashed with register_buffer — so it moves with the model to the GPU but is never updated by gradient descent. forward is a single +.

Why sines instead of just learning a vector per position? Two nice properties. For any fixed offset k, PE(pos+k) is a linear function of PE(pos) — so “three words to the left” is an easy, learnable relationship. And waves keep going: sinusoids may extrapolate to sentences longer than any seen in training. Learned position embeddings scored nearly identically in the paper's tests; the authors chose sinusoids for that extrapolation chance.

Step 9 · Assembling one floor in code

FFN, residual, norm — and an EncoderLayer is three lines.

All the heavy parts exist as classes. Now we glue them into a floor. A floor needs two more small pieces and one wrapper — and then the whole layer's forward is just three calls.

First the second sub-layer, the feed-forward network: two nn.Linear layers with a relu() between them, applied to each position separately. Then the wrapper that keeps a 12-floor building trainable — a residual connection (add the input back to the output, a highway for signal and gradients) followed by LayerNorm (rescale the 512 numbers to a standard spread):

FFN · SublayerConnection · EncoderLayerannotated-transformer · three forward()s

class PositionwiseFeedForward(nn.Module):   # FFN(x) = max(0, xW1+b1)W2+b2
    def forward(self, x):
        return self.w_2(self.dropout(self.w_1(x).relu()))   # 512 → 2048 → 512

class SublayerConnection(nn.Module):        # "Add & Norm" around any sub-layer
    def forward(self, x, sublayer):
        return x + self.dropout(sublayer(self.norm(x)))     # residual + norm + dropout

class EncoderLayer(nn.Module):              # one of the six floors
    def forward(self, x, mask):
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))  # Step 7
        return self.sublayer[1](x, self.feed_forward)                       # the FFN

Self-attention is self_attn(x, x, x) — query, key and value are all the same x; that triple-x is exactly what makes it self-attention. Each sub-layer is wrapped by a SublayerConnection, so the floor is literally two lines of plumbing. (One detail: this code norms first for simplicity, slightly different from the paper's “norm last”.)

LayerNorm(x + Sublayer(x))

this wrapper appears around every sub-layer, in all 12 floors

The dashed teal arcs are the residual highways. Every sub-layer in the machine — all attention blocks and all FFNs — gets this exact wrapper.

The DecoderLayer is the same idea with one extra sub-layer: self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask)) — cross-attention, whose query is the German being written but whose key and value are the encoder's memory m. All sub-layer outputs are d_model = 512, which is what keeps the floors stackable.

PART C

Build it, train it, run it

One make_model() call, a copy-task you can train, and the decode loop in code

Step 10 · greedy_decode, the loop

Generation is a for-loop that calls decode() one word at a time.

Every class is built. Now we run the model the way inference does — and it's a short loop you can read in full.

The encoder runs once over “I am a student” and hands its final K and V to every decoder floor. There, cross-attention works exactly like the attention you computed in Step 5 — except Q comes from the decoder (“what do I need right now?”) while K and V come from the encoder (“here's the source sentence”).

Generation is a loop: emit a word, feed it back in as input, emit the next — until the model produces ⟨eos⟩, a special end-of-sentence token that means “I'm done.” Run the loop yourself:

Step the decoder — watch the German emerge

Source — encoded once, attended at every step

I am a student

Target — written so far

· · · · ·

ich

0.62

ein

0.21

danke

0.09

Step 1 of 5 — top-3 candidates for the first word.

Arc = cross-attention into the current step. Notice the decoder leaning on “I” while writing ich, on “am” while writing bin — the source stays available the whole time. Probabilities are toy values for our running example.

def greedy_decodeannotated-transformer · inference

def greedy_decode(model, src, src_mask, max_len, start_symbol):
    memory = model.encode(src, src_mask)              # read the source ONCE
    ys = torch.zeros(1, 1).fill_(start_symbol)        # begin with ⟨start⟩
    for i in range(max_len - 1):
        out = model.decode(memory, src_mask, ys,
                           subsequent_mask(ys.size(1)))   # can't see ahead
        prob = model.generator(out[:, -1])            # scores for the next word
        next_word = torch.max(prob, dim=1)[1].item()  # greedy: take the top one
        ys = torch.cat([ys, torch.empty(1,1).fill_(next_word)], dim=1)   # feed it back
    return ys

The whole inference algorithm — encode once, then loop: decode, pick the top word with generator, append it, repeat. The stepper above is exactly this loop, one iteration per click.

One rule makes this loop honest: masking. While writing word 3, the decoder's self-attention may only look at words 1–2. That sounds obvious during generation — the future doesn't exist yet — but during training the whole correct German sentence is sitting right there in the input. In code, subsequent_mask is a torch.triu upper-triangle of zeros; back in the attention function (Step 6), masked_fill(mask==0, -1e9) sets future scores to −∞ so softmax gives them weight 0. That one line is what stops the model from copying the answer instead of learning to predict it.

The greyed upper-right triangle is the mask. One more trick pairs with it: the target sequence is fed in “shifted right” — each position receives the words before it and is graded on predicting the next one.

Three attentions, one mechanism: encoder self-attention (Step 4), decoder masked self-attention (this mask), and encoder→decoder cross-attention (Q from decoder; K, V from encoder output).

Step 11 · class Generator & make_model

A vector becomes a word — and the whole model is one function call.

The loop called model.generator to score the next word. That's the last class — and once it exists, every piece is built.

The Generator is tiny: one nn.Linear from 512 to the vocabulary size, then log_softmax. The decoder's top floor outputs one 512-number vector per position; the Linear assigns every word a raw score (a logit); softmax (the same one from Step 5) turns those into probabilities. Then you pick.

class Generator · def make_modelannotated-transformer · the factory

class Generator(nn.Module):                  # vectors → word probabilities
    def forward(self, x):
        return log_softmax(self.proj(x), dim=-1)  # Linear(512 → vocab) + softmax

def make_model(src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1):
    c = copy.deepcopy
    attn = MultiHeadedAttention(h, d_model)
    ff   = PositionwiseFeedForward(d_model, d_ff, dropout)
    position = PositionalEncoding(d_model, dropout)
    model = EncoderDecoder(                       # ← every class from this page, wired up
        Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
        nn.Sequential(Embeddings(d_model, src_vocab), c(position)),
        nn.Sequential(Embeddings(d_model, tgt_vocab), c(position)),
        Generator(d_model, tgt_vocab))
    return model

This is the whole assembly. Every default here is a paper number: N=6 floors, d_model=512, d_ff=2048, h=8 heads, dropout=0.1. c = copy.deepcopy clones one attention and one FFN into all the floors — the Step 2 trick again. Eleven steps of classes, built by one call.

Watch it on a toy six-word German vocabulary at the moment the model should say “bin”. The training target is a one-hot vector — all zeros except a single 1 on the correct word:

Untrained modelrandom junk — ≈ uniform

ich

0.21

bin

0.17

ein

0.15

Student

0.19

danke

0.14

⟨eos⟩

0.14

Trained modelafter 4.5M sentence pairs

ich

0.02

bin

0.90

ein

0.03

Student

0.02

danke

0.01

⟨eos⟩

0.02

The targetone-hot for “bin”

ich

bin

ein

Student

danke

⟨eos⟩

Training is the journey from the left column to the middle one. The loss function — cross-entropy, here computed against a label-smoothed target (the paper spreads a little mass off the correct word, ε_ls = 0.1) — measures how far the predicted bars are from the target; every update nudges the matrices of Steps 5–9 so the right bar grows. The document proves the whole thing works on a tiny copy task first: feed in [1,2,3,4,5,6,7,8,9,10], and after 20 epochs the trained model decodes the same sequence straight back — a runnable sanity check before the 4.5-million-pair real dataset.

One last decision remains at inference time: how do you pick? Greedy decoding (the loop in Step 10) always takes the tallest bar — fast, but a single early misstep can't be undone. Beam search keeps several candidate sentences in flight:

Beam of 2 on our toy step: keep both “ich” (0.62) and “ein” (0.21), grow both, and keep whichever full sentence scores higher. The paper ran a beam of 4 with length penalty α = 0.6.

Two paper details worth knowing: the final Linear layer shares its weights with the embedding layers (scaled by √d_model). And training used label smoothing 0.1 — the model is taught to never be 100% sure, which worsens perplexity (≈ how surprised the model is by held-out text) but improves BLEU (≈ overlap with reference translations, the field's yardstick).

Step 12 · It runs — and you can read it

A few hundred lines reproduce a state-of-the-art model.

Every class is built and wired by make_model. Train it on the copy task and it works; train the same code on WMT and it reproduces the paper. The point of the document: those results come from code short enough to read end to end.

What the original paper's code achieved · WMT 2014 (BLEU)

Transformer (big) · EN→FR

41.8

prior best ensemble · EN→FR

41.29

Transformer (big) · EN→DE

28.4

prior best ensemble · EN→DE

26.36

0BLEU →45

~5 lines the entire scaled dot-product attention function — the mechanism the whole paper is named for.

1 call make_model() assembles every class on this page, with the paper's defaults baked in.

20 epochs to pass the runnable copy-task check before touching the 4.5M-pair real dataset.

28.4 / 41.8 BLEU on EN→DE / EN→FR for the paper's big model — a new state of the art, from code this short.

Now you can read it. Six questions about the implementation — answer each out loud before opening it. If all six come easily, you can read the Annotated Transformer unaided.

Which line of attention() does the √d_k scaling?

The third line: matmul(query, key.transpose(-2,-1)) / math.sqrt(d_k), with d_k = query.size(-1) = 64. Big scores park softmax in a flat region where gradients vanish; dividing by √64 = 8 keeps learning alive (Step 5, row 3).

How does MultiHeadedAttention turn 1 head into 8?

Mostly reshaping: view(nbatches, -1, h, d_k).transpose(1,2) slices the 512-wide vector into 8 lanes of 64, the same attention() runs over all of them, and a view back to 512 is the “Concat”. Four nn.Linears cover W^Q, W^K, W^V, W^O.

What does masked_fill(mask==0, -1e9) accomplish?

It forbids positions: setting their pre-softmax score to nearly −∞ makes their weight ≈ 0. Paired with subsequent_mask (a triu upper triangle), it stops the decoder from seeing future words — so it must predict, not copy.

Why is self-attention written self_attn(x, x, x)?

Because query, key and value all come from the same input x — that's what makes it self-attention. In the decoder's cross-attention the call is src_attn(x, m, m): query from the decoder, key and value from the encoder's memory.

What does register_buffer do for positional encoding?

It stores the precomputed sin/cos table as part of the module — so it moves to the GPU with the model — but not as a parameter, so gradient descent never updates it. forward is then a single x + pe.

What is the entire greedy_decode algorithm?

Encode the source once into memory, then loop: decode, take the top word from generator, append it, repeat until max_len. Encode-once, then a short for-loop — that's inference.

Read the source — the document at nlp.seas.harvard.edu runs as a notebook: open it and you can execute every class on this page, train the copy task, and print real attention matrices. Best read right after a first pass over “Attention Is All You Need.”

“This post presents an annotated version of the paper in the form of a line-by-line implementation. … This document itself is a working notebook, and a completely usable implementation.” — The Annotated Transformer, Harvard NLP

The profound impact

The page that taught a generation to build Transformers.

The architecture went on to power nearly every modern model — but the jump from a dense paper to working code is steep. The Annotated Transformer became the bridge: the resource many people credit for finally making attention click, by showing the equations as code you can run.

2018 · the original

Sasha Rush's notebook

One annotated PyTorch reimplementation became the standard companion to the paper — copied into countless courses, tutorials and codebases.

2022 · the rewrite

Maintained, not bit-rotted

The Harvard NLP group modernized it for current PyTorch, so the notebook still runs — a living document, not a frozen artifact.

the lineage · BERT · GPT

The block it teaches

The same encoder and decoder stacks became BERT and GPT, and now power vision, audio, code and protein models far past translation.

The paper was the spec; this notebook is how a generation of builders learned to read it.

The whole architecture is a few forward() methods.

The skeleton — a paper you can run

The paper goes in; a working translator comes out.

The top-level model is one tiny class that wires five parts together.

A lookup table turns words into vectors — times one curious constant.

Each equation, as the code that runs it

When the machine reads “its”, what does “its” mean?

Second exhibit · a nine-word reach

Six small steps. Real numbers. No magic left.

Make a query, a key, and a value for each word

Score: dot word #1's query against every key

Scale: divide by 8

Softmax: turn scores into weights

Weigh each value by its weight

Sum — and you have z₁

Your six steps, written out as five lines of PyTorch.

The real model runs eight of these at once.

Order, restored — and precomputed once into a buffer.

FFN, residual, norm — and an EncoderLayer is three lines.

Build it, train it, run it

Generation is a for-loop that calls decode() one word at a time.

A vector becomes a word — and the whole model is one function call.

Untrained modelrandom junk — ≈ uniform

Trained modelafter 4.5M sentence pairs

The targetone-hot for “bin”

A few hundred lines reproduce a state-of-the-art model.

The page that taught a generation to build Transformers.

Sasha Rush's notebook

Maintained, not bit-rotted

The block it teaches

The whole architecture is a few `forward()` methods.