Part 1 · Encoder
Every size in the Transformer is a design choice the authors fixed before training. The paper's base model uses the red values (§3, Table 3). A 512-wide matrix is impossible to read, so this notebook runs the whole forward pass with the blue sizes. They keep the same proportions: d_k = d_model / h and d_ff = 4 × d_model.
Every shape note shows both: this run and paper. The homework asks you to work out the paper's shapes by hand.
First time here? Keep the recommended sizes so you can focus on the forward pass. Once it clicks, change them and watch every shape follow.
Everything that follows is computed from whatever you type here, live in your browser. Try a sentence with a repeated word, or try “dog bites man” and then “man bites dog”.
Up to 8 words, so every matrix stays readable.
The model can't read letters, so the sentence is cut into tokens and each token is replaced by its number in a fixed vocabulary.
The paper used byte-pair encoding with a shared vocabulary of about 37,000 tokens for English–German, and a 32,000 word-piece vocabulary for English–French §5.1. Byte-pair encoding often splits rare words into pieces. To keep things readable we split on whole words here, and the IDs are illustrative.
sentence → tokens → vocabulary IDs
The embedding layer is a learned table with one row per vocabulary entry and d_model columns. In the paper that's 37,000 × 512, about 18.9 million numbers. In this run each row has {d} numbers. Each token ID simply picks out its row §3.4.
Stack those rows and you get the matrix E: one row per token, d_model columns.
each ID selects one row of the embedding table
The paper multiplies every embedding by √d_model before going further §3.4. In this run that's √{d} ≈ {sqd}. In the paper it's √512 ≈ 22.63.
The paper doesn't say why. A common reading is that it brings the learned embeddings, which start out small, up to a size comparable with the positional encodings added next, so word meaning isn't drowned out by position.
Attention compares every token with every other token all at once. There is no recurrence and no convolution, so nothing in the maths knows which word came first §3.5.
The card on the right proves it with your own sentence. It runs one full encoder layer on the embeddings without positional information, once in your word order and once shuffled. Each word comes out with exactly the same vector.
Before looking at the paper's answer, try designing one yourself. A good position code should be:
Bounded: small numbers, so it doesn't drown out the word's meaning.
Unique: every position gets a different code.
Consistent: position 3 means the same thing in a short sentence and a long one.
Smooth: nearby positions get similar codes, and “3 steps apart” looks the same everywhere.
Unlimited: it works for sentences longer than any seen in training.
Flip through the candidates on the card. The binary counter is the key clue. Its lowest bit flips every step and higher bits flip more and more slowly, like clock hands at different speeds. Sine and cosine waves are the smooth version of that counter, and they tick every box.
Each position gets its own fixed vector, the same width as an embedding. Its numbers come in pairs. Each pair is a sine and a cosine of the same angle, and each pair turns at its own speed §3.5:
PE[pos, 2i] = sin(pos / 100002i/d_model)
PE[pos, 2i+1] = cos(pos / 100002i/d_model)
Think of each pair as a clock hand, with sine and cosine as the coordinates of its tip. The first hand spins fast and tells neighbours apart. Later hands crawl and tell apart positions that are far away. No two positions put all hands in the same place, so every position gets a unique fingerprint.
The word never enters the formula. Only the position and the column do.
Why a sine and a cosine? One number alone is ambiguous: sin(1) and sin(2.14) are both 0.84. The pair pins the hand down. It also makes “k steps later” a fixed rotation of every pair, a simple multiplication the model can learn. The card shows this rotation with real numbers.
Fill in the formula for every row (position) and column. In this run d_model = {d}, so there are {pairs} pairs with divisors {divs}.
At the paper's size there are 256 pairs, with wavelengths growing geometrically from 2π to 10000·2π positions §3.5. The canvas below the matrix shows the real 512-wide version for 50 positions.
The positional encoding is added, number by number, to the scaled embeddings. This is why PE has exactly d_model columns: addition needs identical shapes.
Watch the repeated words: rows that were identical in E are now different. The result X is the input to encoder layer 1. The paper also applies dropout (10%) to this sum during training §5.4.
Each head owns three learned weight matrices, WQ, WK and WV, each d_model × d_k. Multiplying X by each gives every token three new vectors §3.2.2:
Query: what this token is looking for.
Key: what this token can be matched on.
Value: what this token passes along if it's picked.
The weights are the same for every token and every sentence. Only X changes, which is why each token gets its own Q, K and V. In a trained model the weights are learned. Here they're random.
Each query is compared with each key using a dot product: multiply matching entries and add them up. A large score means “this key matches what I'm looking for”.
Doing it for all pairs at once is a single multiplication, Q · Kᵀ. Transposing K turns its tokens into columns.
Divide every score by √d_k: √{dk} ≈ {sqdk} in this run, √64 = 8 in the paper. Then apply softmax to each row: take ex of every entry and divide by the row total. Each row becomes a set of positive weights adding up to 1 §3.2.1.
Multiply the attention weights by V. Each output row is a weighted average of all the value vectors, mixed in the proportions that token's row asked for.
This is the moment information actually moves between words. In a trained model, “sat” might take most of its new vector from “cat”, so it now carries information about who is sitting.
Attention(Q, K, V) = softmax(QKᵀ / √d_k) · V Eq. 1
All heads run in parallel, each with its own weights, so each can learn to look for a different kind of relationship. Their outputs are placed side by side, which is called concatenation.
Concatenation alone doesn't combine anything: each head's numbers sit in their own block of columns. A final learned matrix WO mixes them, so every output number blends every head. It also maps the result back into X's space, ready to be added to X.
Add (residual connection): the layer's input X is added back to the attention output. Attention produces an update, not a replacement, so each token keeps its own identity and position. It also gives gradients a direct path through all the layers, which makes deep stacks trainable.
Norm (layer normalization): each token's row is rescaled to mean 0 and standard deviation 1 across its numbers, then multiplied by a learned scale γ and shifted by a learned β. This keeps values on a steady scale from layer to layer §3.1.
output = LayerNorm(x + Sublayer(x))
Attention moved information between tokens. The feed-forward network now processes each token on its own with the same small two-layer network §3.3, Eq. 2.
FFN(x) = max(0, x·W1 + b1)·W2 + b2
First it expands each token from d_model to d_ff numbers: {d} → {dff} in this run, 512 → 2048 in the paper. Then ReLU replaces every negative number with 0. Without ReLU, two multiplications in a row would collapse into one, and the network could only learn straight-line relationships.
The second multiplication brings each token back to d_model numbers so it can be added to the residual. Row “cat” of the result depends only on row “cat” of the input. No information crosses between tokens here.
This is where most of a layer's parameters live: about 2.1 million of the 3.15 million in each paper-sized layer.
The same pattern as after attention: add the FFN's input back in, then normalize each row. The result is the output of encoder layer 1.
Layer 1's output becomes layer 2's input, and so on up to the last layer §3.1. The paper uses N = 6. Every layer has the same steps but its own weights: WQ, WK, WV, WO, W1, W2, γ and β.
Each pass lets every token gather a little more context from the others. Press play to send your sentence up the stack. Every layer on the card is computed for real, with its own random weights.
Out comes Z: one vector per token, each now describing that word in the context of the whole sentence. Compare it with step 3, where every word had a fixed vector that knew nothing about its neighbours.
Z isn't turned back into words. It's handed to every layer of the decoder, which turns it into keys and values for encoder–decoder attention. That's Part 2.
Now do the whole pass yourself. This Jupyter notebook gives you a 3-word sentence with d_model 4 and 2 heads. You write down every matrix from the token IDs to Z, and each one is checked as you go, with hints when an entry is off.
It comes as a zip. Unzip it, then open the notebook in Jupyter, VS Code or Google Colab (File → Upload notebook). It only needs NumPy. Finished Part 2 as well? The encoder + decoder version carries on from Z and translates the sentence.
Choose small sizes, pick whether to stop after the encoder or carry on through a decoder layer, generate, and print. You get every number you need, an empty box for each answer, a hint under every step and an answer key at the back. The sizes are capped so one full forward pass stays a few pages of pen work.
Part 2 · Decoder
The encoder turned your sentence into Z, one context-aware vector per word. The decoder's job is narrower than it looks. At every position it answers a single question: given the English sentence and the German words so far, what is the next German word?
Everything new in this part follows from that job. The decoder reads the translation shifted right. Its self-attention is masked, so no position can peek at the word it has to predict. Its cross-attention takes questions from the German side and answers from Z. And it ends with a linear layer and softmax that turn vectors back into words.
This part reuses your Part 1 sentence and sizes: d_model {d}, h {h}, d_k {dk}, d_ff {dff}, N {N}. Change them in Part 1 and everything here follows.
The decoder works with two inputs. One is Z, the encoder's output for your sentence from Part 1. The other is the German sentence it should produce. During training the correct translation is known, so we hand it to the model. Type it here or pick an example.
Up to 7 words. The weights are random, so any words work. A real translation just makes the pictures easier to read.
The decoder is trained to predict the next word. So it reads the translation moved one place to the right, behind a start token ⟨s⟩, and at every position it must produce the word that comes next. The last position must produce the end token ⟨/s⟩, which means “done” §3.1.
Read the card column by column. Click a position to see what it has seen and what it must predict.
Exactly the encoder's steps 3 to 9. Each token picks its row of the embedding table, the rows are multiplied by √d_model, and the positional encoding is added §3.4, §3.5.
The paper uses one vocabulary shared by English and German, and one embedding table for both stacks (and, as you'll see in chapter 5, for the output layer). Position k also gets the same PE vector in both stacks.
The first sublayer is self-attention, like the encoder's: Q, K and V are all computed from the same matrix, Y0. So the German words look at each other here, not at the English ones.
The decoder has its own weight matrices, separate from the encoder's. Each layer of each stack learns its own.
The scores are computed as in the encoder, S = Q · Kᵀ ÷ √d_k: every German position against every other one. Then comes the one new idea. Every entry above the diagonal, where the key's position is later than the query's, is replaced with −∞ before the softmax §3.2.3.
So row “{t1}” can still see ⟨s⟩ and itself, but nothing after it.
Softmax works row by row, as before. Because e−∞ = 0, a masked entry gets exactly zero weight. The allowed entries share the whole weight, so every row still adds up to 1.
The first row can only see itself, so its single weight is 1. Each row below sees one more word.
Tick the box on the card to remove the mask and see where the model would look if it were allowed to cheat.
From here it is the encoder's steps 13 and 14. Each row of A · V is a weighted average of value vectors, but only of positions at or before that row, because the later ones have weight 0.
The heads are then concatenated and mixed by their own WO, which brings the width back to d_model.
The same residual connection and layer normalization as in the encoder. Y0 is added back to the attention output, and each row is rescaled to mean 0 and standard deviation 1 §3.1.
The result, Y1, holds each German position mixed with the German words before it. It is the input to the next sublayer.
This sublayer is the bridge between the two halves. The paper calls it encoder–decoder attention §3.2.3. The maths is the same as every attention so far. What changes is where the inputs come from:
Queries come from the decoder, Y1: one per German position. Each asks “which English words do I need for my next word?”
Keys and values come from the encoder output Z: one per English word.
Each German query is compared with each English key by a dot product, all at once with Q · Kᵀ. For the first time the score table is not square: rows are German positions, columns are English words.
Scale by √d_k and softmax each row, exactly like the encoder. There is no mask here. The whole English sentence is known before the decoder starts, so every position may look at every English word.
Each row is a soft alignment: how much each English word matters for this German position's next word. In a trained model the row that has to produce “katze” would lean on “cat”. With random weights the pattern is arbitrary.
Multiplying by V mixes the English words' value vectors into each German position, in the proportions its row asked for. This is the moment information crosses from the encoder into the decoder.
Then the heads are concatenated and mixed by WO as usual.
Add Y1 back and normalize each row. Y2 now carries two kinds of context: the German words so far, from masked self-attention, and the English sentence, from cross-attention.
The same two-layer network as the encoder's, with its own weights §3.3. Each position is expanded to d_ff numbers, passed through ReLU, then compressed back to d_model. Each position is processed on its own, so nothing moves between positions here.
The decoder layer has three sublayers, so it has three Add & Norms. This last one gives the output of decoder layer 1. Like every layer, it hands back exactly the shape it was given.
Layer 1's output goes into layer 2, and so on §3.1. The paper uses 6 decoder layers, each with its own weights for all three sublayers.
Notice the second input on the card. Every decoder layer's cross-attention reads the same Z, the encoder's final output. Press play to send your translation up the stack.
The last decoder layer gives one d_model-wide vector per position. A linear layer turns each vector into one score, called a logit, per vocabulary word. The paper ties this layer to the embedding table: its weight matrix is the embedding table, transposed §3.4.
So each logit is a dot product between the decoder's vector and one word's embedding. It asks “how much does what I want to say look like this word?”
The paper scores all ~37,000 vocabulary entries. To keep the card readable we score only {V} of them: the words of your translation and ⟨/s⟩.
Softmax each row and you get a probability for every word at every position. During training the loss is minus the log of the probability given to the correct next word, averaged over the positions. This is cross-entropy; the paper also smooths the targets a little (label smoothing 0.1) §5.4.
During training the correct translation is known, so the whole shifted sentence goes in at once and all {m} predictions come out of a single pass. The input is always the true previous words, even where the model would have guessed differently. This is called teacher forcing.
The mask is what makes this honest. {m} predictions are computed side by side, yet each one has only seen the words to its left, exactly as it would when translating.
When translating new text there is no German sentence to feed in. The encoder runs once. Then the decoder starts from ⟨s⟩ alone, predicts a word, appends it to its own input and runs again, until it predicts ⟨/s⟩.
Step through it on the card. Only the last row's prediction is used at each step. And because of the mask, every earlier row comes out exactly as it did on the step before, so real implementations keep those rows' keys and values instead of recomputing them (a “KV cache”).
One row per German position went in. Out came one probability distribution per position over the whole vocabulary. During training that gives the loss; when translating, the last row gives the next word.
The card lists every step and its shape. Notice which steps have n rows (the encoder's side) and which have m (the decoder's side), and that the only place they meet is cross-attention.
Now translate a sentence yourself. This Jupyter notebook takes “the cat sat” through one encoder layer and then “⟨s⟩ die katze saß” through one decoder layer, with d_model 4 and 2 heads. You write every matrix, from the token IDs to the probabilities of the next word, including the mask. Each one is checked as you go, with hints when an entry is off.
It comes as a zip. Unzip it, then open the notebook in Jupyter, VS Code or Google Colab (File → Upload notebook). It only needs NumPy. Prefer paper? The worksheet maker in Part 1 can now add a decoder layer too.