How to Implement the Transformer from Attention Is All You Need, Step by Step

How do you implement a Transformer from scratch?

You implement a Transformer from scratch by writing its components in the order data flows through them, testing each one on small inputs, and only then joining them into an encoder and a decoder. The components are scaled embeddings, positional encoding, scaled dot-product attention, multi-head attention, masks, the feed-forward layer, residual connections with layer normalization, the output projection, a label-smoothed loss, a warmup learning-rate schedule and a decoding rule. TensorTonic has a coding problem with tests for most of these parts, and its Transformer paper track in the Research section walks through Attention Is All You Need in stages from tokenization to the encoder block.

This guide follows the original paper, which describes an encoder-decoder model for translation. If you want to build a decoder-only GPT instead, the guide on how to build an LLM from scratch covers next-token training, sampling and the KV cache, and this guide links to it where the two overlap.

What is the original Transformer from Attention Is All You Need?

The original Transformer is an encoder-decoder model that Vaswani and colleagues introduced in 2017 to translate sentences from one language to another. The encoder reads the whole source sentence with self-attention and produces one vector per source token. The decoder then writes the target sentence one token at a time, attending to the tokens it has already written and, through cross-attention, to the encoder output.

The paper replaced the recurrent and convolutional layers that translation models used at the time with attention alone, which let every position be processed in parallel during training. Its base model was trained on the WMT 2014 English-German and English-French datasets, and the table below lists the settings the paper gives for that base model.

You do not need these sizes for your own implementation. Keep them as named parameters, run your first experiments with something much smaller, such as two layers and a d_model of 128, and scale up once everything works.

What do you need to know before implementing a Transformer?

You need comfortable Python, NumPy array shapes and broadcasting, matrix multiplication, and a basic idea of how a neural network is trained with gradient descent. PyTorch helps once you train the full model, and the PyTorch study plan teaches it from tensors upward.

• For NumPy, you should understand shapes, broadcasting and indexing, which the NumPy study plan covers.
• For linear algebra, you should be comfortable with the matrix multiplication lesson, since attention is mostly batched matrix products.
• For training, you should know how the chain rule turns a loss into gradients for every weight.
• For reading the paper itself, the guide on how to read an ML paper and implement it explains how to turn its equations into a table of symbols and shapes before you write code.

What are the steps to implement the Transformer from scratch?

The steps follow a sentence pair through the model. Source and target tokens become scaled embeddings with position information, attention mixes information between positions, the encoder and decoder layers wrap attention with feed-forward layers and normalization, and the output projection turns decoder vectors into scores over the vocabulary. Training then adds the loss and the learning-rate schedule, and decoding turns a trained model back into sentences. Each stage below has a short explanation and a table that pairs every topic with a page to read and a problem to solve.

Stage 1: How do you turn tokens into scaled embeddings?

You turn tokens into embeddings by splitting each sentence into subword tokens, mapping each token ID to a learned vector of size d_model, and multiplying that vector by the square root of d_model. The paper used byte-pair encoding with a vocabulary shared between source and target for English-German, so one embedding table can serve both sides. Sentences in a batch have different lengths, so you pad the shorter ones with a padding token and remember where the padding is for the masks later.

The scaling by the square root of d_model is easy to forget. Embedding weights are usually initialized with small values, and the scaling keeps the token signal from being swamped by the positional encoding that is added next.

Stage 2: How does positional encoding work in the Transformer?

Positional encoding works by adding a fixed vector to each token embedding that depends only on the position of the token in the sentence. The paper uses sines for even dimensions and cosines for odd dimensions, with PE(pos, 2i) = sin(pos / 10000^(2i / d_model)) and PE(pos, 2i + 1) = cos(pos / 10000^(2i / d_model)), so each pair of dimensions oscillates at a different frequency. Attention treats its input as a set, and without these vectors the model could not tell the first word of a sentence from the last.

Compute the encoding table once for the longest sentence you expect and slice it for each batch. The paper also tried learned position embeddings and reported nearly identical results, and it chose the sinusoidal version because it may extend to sentences longer than those seen in training.

Stage 3: How do you implement scaled dot-product attention?

You implement scaled dot-product attention by multiplying the queries by the transposed keys, dividing by the square root of d_k, applying a mask, taking a softmax over the key axis and multiplying the result by the values. In the paper's notation this is softmax(QK^T / sqrt(d_k)) V. The division keeps the scores from growing with the head dimension, which would otherwise push the softmax towards a single position and shrink its gradients.

The softmax must run over the last axis, which indexes keys. Running it over the query axis still produces numbers of the right shape, and the only sign of the bug is a model that never learns, so test that every row of the attention weights sums to one.

Stage 4: How do you implement multi-head attention?

You implement multi-head attention by projecting the input into queries, keys and values, splitting each projection into h heads of size d_model / h, running scaled dot-product attention on every head in parallel, concatenating the heads and applying one more linear projection. With the paper's settings, 8 heads of 64 dimensions each make up the 512-dimensional model, so multi-head attention costs about the same as one full-width head while letting each head attend to different positions.

Most bugs in this stage come from reshaping. Reshape from (batch, length, d_model) to (batch, length, heads, d_head) and then transpose to (batch, heads, length, d_head), and write the shape of every tensor beside it while you work. The same function serves all three attention sublayers in the model, because only the source of the queries, keys and values and the mask change between them.

Stage 5: How do padding masks and causal masks work?

Padding masks and causal masks both work by setting chosen attention scores to negative infinity, or a very large negative number, before the softmax so that those positions receive zero weight. A padding mask hides the padding tokens among the keys, and it is used in encoder self-attention, decoder self-attention and cross-attention. A causal mask hides every later position in decoder self-attention, so that the prediction for target position i depends only on the target tokens before it.

The causal mask is what makes teacher forcing possible. During training the decoder receives the whole target sentence shifted right by one position, starting with a start token, and the mask lets it predict every target token in one parallel pass without seeing the answer. Combine the two masks with a logical and in the decoder, and broadcast them across the head dimension.

Stage 6: What is the position-wise feed-forward layer?

The position-wise feed-forward layer is a two-layer network applied to each token vector separately and identically, with a ReLU between the layers. The paper writes it as FFN(x) = max(0, xW1 + b1)W2 + b2, with an inner width d_ff of 2048 for a d_model of 512. Attention moves information between positions, and the feed-forward layer transforms each position on its own, so the two sublayers do different jobs.

Stage 7: How do residual connections and layer normalization fit together?

Residual connections and layer normalization fit together around every sublayer, and the paper computes LayerNorm(x + Dropout(Sublayer(x))). This arrangement is called post-norm, because normalization comes after the residual addition. Layer normalization normalizes each token vector over its features and then applies a learned scale and shift, which keeps activations in a stable range as the stack gets deeper.

Most later Transformers use pre-norm instead, which computes x + Sublayer(LayerNorm(x)) and adds one final layer norm after the last layer. Pre-norm keeps a clean path for gradients through the residual stream and usually trains without the long warmup that post-norm needs. Implement post-norm first if you want to reproduce the paper, then make the position of the norm a flag so you can compare the two.

Stage 8: How do you build the encoder layer?

You build the encoder layer from two sublayers, multi-head self-attention followed by the feed-forward layer, each wrapped in a residual connection and layer normalization. Self-attention in the encoder is bidirectional, so every source token can attend to every other source token, and the only mask it needs is the padding mask. The encoder is six of these layers stacked, and its output is one vector per source token that the decoder reads through cross-attention.

A good test for the encoder layer is that its output has the same shape as its input. That property is what lets you stack layers freely, and it also means a padded source sentence should give the same outputs at its real positions as the same sentence without padding.

Stage 9: How do you build the decoder layer with cross-attention?

You build the decoder layer from three sublayers, which are masked self-attention over the target tokens, cross-attention over the encoder output and the feed-forward layer, each with its own residual connection and layer normalization. In cross-attention the queries come from the decoder, while the keys and values come from the final encoder output, so each target position can look at the whole source sentence. The same encoder output feeds the cross-attention of all six decoder layers.

Cross-attention uses the source padding mask and no causal mask, since the whole source sentence is known before decoding starts. Its attention matrix is rectangular, with one row per target position and one column per source position, so test it with source and target sentences of different lengths.

Stage 10: How do the output projection and weight tying work?

The output projection works by multiplying each final decoder vector by a matrix of size d_model by vocabulary size, which gives one score for every token in the target vocabulary, and a softmax turns those scores into probabilities. Weight tying means the paper uses the same weight matrix for the source embeddings, the target embeddings and this pre-softmax projection. Tying all three needs a shared vocabulary, and it removes a large share of the parameters in a model with a big vocabulary.

In PyTorch you tie weights by pointing the projection weight at the embedding weight, so both names refer to one tensor. Remember that the embedding lookup multiplies by the square root of d_model and the projection does not.

Stage 11: How do you compute the loss with label smoothing?

You compute the loss with label smoothing by replacing each one-hot target with a distribution that keeps most of the probability on the correct token and spreads the rest over the other tokens, then taking the cross-entropy between that distribution and the model's predictions. The paper used a smoothing value of 0.1 and reported that it made perplexity worse, because the model learns to be less certain, while it improved accuracy and BLEU. Padding positions in the target must be left out of the loss, both from the sum and from the count used to average it.

Stage 12: How do you train the Transformer with the warmup schedule?

You train the Transformer with Adam and a learning rate that rises linearly for the first warmup steps and then decays with the inverse square root of the step number. The paper sets lrate = d_model^(-0.5) times min(step^(-0.5), step times warmup_steps^(-1.5)) with 4000 warmup steps, which for a d_model of 512 peaks at about 0.0007 at step 4000. Post-norm Transformers tend to diverge early without this warmup, so implement the schedule as a function of the step and plot it before you train.

The paper batched sentence pairs of similar length together, with about 25,000 source tokens and 25,000 target tokens per batch, which wastes less computation on padding. For the shared parts of the training loop, such as overfitting one batch first and reading the loss curve, the LLM from scratch guide and the guide on a training loss that is not decreasing apply unchanged.

Stage 13: How do you decode with greedy search and beam search?

You decode by running the encoder once on the source sentence and then calling the decoder repeatedly, starting from the start token and appending one predicted token per step until the model emits an end token or reaches a length limit. Greedy search appends the single most likely token at every step. Beam search keeps the k most likely partial sentences at every step, extends each of them, and keeps the k best of the results, and the paper used a beam size of 4 with a length penalty of 0.6.

Beam search compares sentences by the sum of their token log probabilities, which favours short outputs, and the length penalty divides that sum by a function of the length to correct for it. Start with greedy decoding to confirm the model works, then add beam search and check that it raises BLEU on held-out sentences. Sampling methods such as temperature and top-p are covered in the LLM from scratch guide.

How does the original Transformer differ from GPT and BERT?

The original Transformer uses both an encoder and a decoder, while GPT keeps only the decoder and BERT keeps only the encoder. GPT drops cross-attention and trains with causal self-attention to predict the next token. BERT drops the causal mask and trains bidirectional self-attention to fill in masked tokens.

Once your encoder works, BERT is a short step away. The BERT paper track in the Research section continues with a WordPiece lesson, a masked language model lesson, a next sentence prediction lesson and a fine-tuning lesson, and you can practise the objectives in Masked Language Modeling and Next Sentence Prediction. Once your decoder works, the Build a Tiny GPT from Scratch project trains a decoder-only model in a PyTorch notebook.

Can you implement a Transformer without a GPU?

Yes, you can implement a Transformer without a GPU, because every component can be written and tested on small arrays on a laptop. The paper trained its base model on 8 GPUs for about 12 hours, and a model with two layers and a small d_model trained on a copying or reversal task needs only minutes on a CPU. That is enough to confirm the masks, the decoder input shift and beam search all work.

The component problems linked in this guide run in the browser and are checked against tests, so there is nothing to install. Our guide on practising ML without a GPU covers where to find compute when you want to train on a real translation dataset.

What mistakes do people make when implementing a Transformer?

• Taking the softmax over the query axis instead of the key axis gives tensors of the right shape and a model that does not learn.
• Applying the mask after the softmax leaves rows of weights that no longer sum to one, so always mask the scores before the softmax.
• Forgetting the padding mask in cross-attention lets the decoder attend to padding tokens in the source.
• Feeding the decoder the unshifted target lets it copy the token it should predict, which gives a very low training loss and useless translations.
• Leaving padding positions in the label-smoothed loss, or averaging over them, makes the loss depend on how much padding a batch happens to contain.
• Skipping the scaling of embeddings by the square root of d_model, or applying it to the output projection as well, changes the balance between token and position information.
• Training a post-norm model with a high constant learning rate often makes it diverge early in training, which the warmup schedule prevents.
• Using a library module such as nn.Transformer for the core hides the decisions you are trying to learn, so keep it only as a reference to compare outputs against.

Many of these bugs produce plausible numbers and raise no error. Testing each component in isolation, and checking properties such as rows that sum to one and outputs that do not change when padding is added, finds them before they reach the full model.

Practice this

Then work through the Transformer paper track in the Research section, or the Transformers section of the Cracking Deep Learning study plan, which covers layer norm, rotary embeddings, scaled dot-product and multi-head attention, encoder and decoder blocks and ViT patch embeddings. The Cracking NLP study plan continues with Transformers for NLP, including masked language modeling, next sentence prediction and beam search.

Related guides

• How to build an LLM from scratch, one component at a time
• How to read an ML paper and implement it
• 15 ML papers worth implementing from scratch, and in what order
• Training loss not decreasing? Causes and fixes
• NLP interview questions: 35 questions with answers
• LLM interview questions: 43 questions with answers
• How to learn NLP from scratch, one stage at a time