The Annotated Transformer

transformerattentionnlppytorchdeep-learningannotated

Abstraction: Line-by-line annotated implementation of Attention Is All You Need

Key points:

Complete working implementation of the original Transformer in ~400 lines of PyTorch; processes 27,000 tokens/second on 4 GPUs
Architecture: 6-layer encoder-decoder, d_model=512, 8 attention heads, d_k=d_v=64, d_ff=2048, with residual connections and layer normalization
Scaled dot-product attention divides by sqrt(d_k) to prevent large gradients in softmax for high-dimensional keys
Positional encoding uses sine/cosine functions at different frequencies; allows model to extrapolate to longer sequences than seen in training
Training on WMT 2014 EN-DE (4.5M pairs, 37k BPE vocab); big model achieves 28.4 BLEU, +2 over prior SOTA, trained 3.5 days on 8 P100 GPUs
Adam optimizer with warmup schedule: learning rate increases linearly for first 4000 steps, then decays as inverse sqrt of step number

Connections: Pytorch · Transformers · Attention Mechanism · Positional Encoding