Figure 1 from Attention Is All You Need, Vaswani et al. (2017), the paper this implements.
A from-scratch PyTorch implementation of the original Transformer from Vaswani et al., 2017, written to be read alongside the paper rather than to be the fastest or the shortest version of it. It trains an encoder-decoder model for English to German translation on WMT14, and it has since collected 251 stars and 57 forks on GitHub under an MIT license.
Written to be read
The design goal was that someone with the paper open in one window and the repository in the other should be able to move between them without translating between two vocabularies.
Every component gets its own file, named after what the paper calls it: multi-head attention, positional encoding, the position-wise feed-forward network, add-and-norm, the encoder block, the decoder block. Nothing is folded together for brevity. That folding is exactly what makes most implementations hard to follow.
Inside those files every tensor operation carries its resulting shape as a comment. Splitting into heads is annotated at each step: batch by sequence by model dimension, then batch by sequence by heads by head dimension, then batch by heads by sequence by head dimension after the transpose. Reading attention is mostly tracking which axis is which. Writing the shapes down turns that from an exercise into something you can follow.
Where the implementation departs from the paper, the comment says so. The dropout on attention weights is marked as optional and not in the paper. That is the sort of detail that quietly diverges between implementations and then confuses anyone comparing them.
Reproducing the training recipe on one GPU
The paper’s base model trains with a batch size far larger than a single consumer GPU holds. Rather than quietly train at a smaller batch and produce different results, the repository uses gradient accumulation to get there.
The shipped configuration runs a batch of 48 and accumulates gradients until it has seen 2048 examples before stepping the optimizer. That reaches the paper’s effective batch on roughly 10 GB of VRAM. The rest follows the paper too: Adam with betas of 0.9 and 0.98, epsilon at 1e-9, label smoothing at 0.1, and the base dimensions of 512 with six blocks, eight heads, a feed-forward width of 2048 and dropout at 0.1.
The learning rate schedule is the paper’s formula written out directly rather than approximated with an off-the-shelf scheduler. It takes the model dimension to the power of minus a half, times the smaller of the step count to the power of minus a half and the step count times the warmup steps to the power of minus one and a half, over 4000 warmup steps. Three lines, and no question about whether the schedule matches.
Batching by length
Padding is where a naive translation dataloader wastes most of its compute. Batch a two-word sentence with a forty-word one and the short example carries thirty-eight tokens of padding through every layer.
The dataset is sorted by length up front. A custom batch sampler chunks that sorted order into contiguous batches, then shuffles the batches rather than the examples. Each batch therefore holds sentences of similar length and needs almost no padding, and shuffling at the batch level still varies what the model sees each epoch. Preprocessing, tokenization and dataloader construction all happen automatically on the first run, using HuggingFace datasets and tokenizers, with a choice of BPE or word-level.
Configurations for debugging, not just for training
Alongside the full training configuration there are two that exist only to check the machinery works.
Both set the dataset to a couple of examples and train for a thousand epochs, one on GPU and one on CPU. A correct implementation should memorize two sentences completely. If it cannot, the bug is in the model or the training loop, not the data or the hyperparameters. Knowing that before launching a run over four and a half million sentence pairs saves a lot of time.
That separation of concerns is the useful habit here: prove the thing can overfit before asking whether it can generalize.
Running and using it
Training is driven from a single configuration file. It covers the run name and output path, dataset size and test proportion, sequence length, vocabulary size and tokenizer type, batch size and accumulation steps, epochs, optimizer and scheduler parameters, the model dimensions, and label smoothing.
Metrics go to Weights and Biases, including BLEU computed during training rather than only at the end. A sample translation is printed after every epoch, so progress shows up as text and not only as a loss curve. The logs for the pretrained models are public. Those models download with a script, and a small Streamlit app loads a checkpoint and its tokenizer to translate in the browser with greedy decoding.
Everything runs on CPU as well as GPU, which matters more than it sounds for a learning repository: the point is that someone can clone it and watch a transformer train without first finding a GPU.
Related reading
- Transformer paper from scratch: a walkthrough, a step-by-step post building the same architecture.
View the code on GitHub → · Read the paper