Attention Is All You Need
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, Illia Polosukhin | Google Brain / Google Research / University of Toronto | NIPS 2017 | arXiv:1706.03762
Problem
Dominant sequence transduction models for tasks such as language modeling and machine translation are built on recurrent neural networks (LSTMs, GRUs), often augmented with attention. Recurrent models factor computation sequentially along the position axis: hidden state \(h_t\) depends on \(h_{t-1}\), which forecloses parallelization within a training example and becomes a critical constraint at longer sequence lengths where memory limits also reduce batching across examples. Convolutional alternatives (Extended Neural GPU, ByteNet, ConvS2S) reduce sequential computation but still require operation counts that grow with the distance between positions (linearly or logarithmically), making long-range dependencies harder to learn.
Core Insight
A model based entirely on attention — with no recurrence and no convolutions — can match or exceed the quality of recurrent/convolutional encoder-decoders while reducing the path length between any two positions to \(O(1)\) and dramatically increasing training parallelism. Multi-Head Attention restores the representational resolution lost when a single attention head averages over many positions.
Method
The Transformer is an encoder-decoder model. The encoder is a stack of \(N=6\) identical layers; each layer has two sub-layers — a multi-head self-attention mechanism and a position-wise fully connected feed-forward network — each wrapped by a residual connection and layer normalization. The decoder is also \(N=6\) identical layers and adds a third sub-layer that performs multi-head attention over the encoder output; the decoder's self-attention is masked so that position \(i\) can only attend to positions less than \(i\), preserving auto-regressive generation.
The atomic attention operation is Scaled Dot-Product Attention:
\[\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{Q K^{T}}{\sqrt{d_k}}\right) V\]
The \(1/\sqrt{d_k}\) scaling keeps dot products from growing large at large \(d_k\) and pushing softmax into low-gradient regions. Multi-Head Attention projects Q, K, V \(h\) times into lower-dimensional subspaces, runs attention in parallel, and concatenates: \(h=8\) heads with \(d_k = d_v = d_{model}/h = 64\). Attention appears in three roles: encoder-decoder attention (queries from decoder, keys/values from encoder), encoder self-attention, and masked decoder self-attention.
Each layer's position-wise feed-forward network is \(\text{FFN}(x) = \max(0, xW_1 + b_1)W_2 + b_2\) with \(d_{ff} = 2048\). Token embeddings (input and output) share weights with the pre-softmax linear projection and are scaled by \(\sqrt{d_{model}}\). Because there is no recurrence or convolution, sinusoidal positional encodings at multiple frequencies are added to the embeddings to inject position information:
\[PE_{(pos, 2i)} = \sin(pos/10000^{2i/d_{model}}), \quad PE_{(pos, 2i+1)} = \cos(pos/10000^{2i/d_{model}})\]
The authors verified that learned positional embeddings perform nearly identically; sinusoidal was chosen for its potential to extrapolate to longer sequences than seen in training.
The Section 4 ("Why Self-Attention") argument compares layer types on three axes — per-layer complexity, sequential operations, and maximum path length — and shows that self-attention is the only family with \(O(1)\) on both sequential ops and path length.
| Layer Type | Complexity per Layer | Sequential Ops | Max Path Length |
|---|---|---|---|
| Self-Attention | \(O(n^2 \cdot d)\) | \(O(1)\) | \(O(1)\) |
| Recurrent | \(O(n \cdot d^2)\) | \(O(n)\) | \(O(n)\) |
| Convolutional | \(O(k \cdot n \cdot d^2)\) | \(O(1)\) | \(O(\log_k n)\) |
| Restricted Self-Attention | \(O(r \cdot n \cdot d)\) | \(O(1)\) | \(O(n/r)\) |
Experimental Setup
| Component | Value |
|---|---|
| Tasks | WMT 2014 English-German (~4.5M pairs, ~37k shared BPE vocab); WMT 2014 English-French (~36M pairs, 32k word-pieces) |
| Hardware | 1 machine x 8 NVIDIA P100 GPUs |
| Base model | \(N=6\), \(d_{model}=512\), \(d_{ff}=2048\), \(h=8\), \(d_k=d_v=64\), \(P_{drop}=0.1\), \(\epsilon_{ls}=0.1\) |
| Big model | \(N=6\), \(d_{model}=1024\), \(d_{ff}=4096\), \(h=16\), \(P_{drop}=0.3\) (0.1 for En-Fr) |
| Optimizer | Adam (\(\beta_1=0.9\), \(\beta_2=0.98\), \(\epsilon=10^{-9}\)); warmup 4000 steps then inverse-sqrt decay |
| Batching | ~25k source + ~25k target tokens per batch, length-bucketed |
| Training schedule | Base: 100k steps (~12 h); Big: 300k steps (~3.5 days) |
| Inference | Beam size 4, length penalty \(\alpha=0.6\); checkpoint averaging (last 5 for base, last 20 for big) |
| Metric | BLEU on newstest2014 (test) and newstest2013 (development) |
Headline Quantitative Results
WMT 2014 BLEU vs. selected baselines:
| Model | EN-DE (BLEU) | EN-FR (BLEU) |
|---|---|---|
| GNMT + RL | 24.6 | 39.92 |
| ConvS2S | 25.16 | 40.46 |
| MoE | 26.03 | 40.56 |
| Transformer (base) | 27.3 | 38.1 |
| Transformer (big) | 28.4 | 41.0 |
- The big Transformer beats the previous best (including ensembles) on En-De by more than 2 BLEU, setting a new state of the art.
- On En-Fr it sets a new single-model state of the art at less than 1/4 the training cost of the previous best.
- Even the base model surpasses all previously published models and ensembles at a fraction of their training cost.
Training cost:
- Base (En-De): \(3.3 \times 10^{18}\) FLOPs.
- Big (En-De): \(2.3 \times 10^{19}\) FLOPs.
Ablations (En-De newstest2013):
- Single-head attention is 0.9 BLEU worse than the \(h=8\) best setting; too many heads also degrades quality.
- Reducing \(d_k\) alone hurts quality, suggesting the dot-product compatibility function is not always sufficient.
- Bigger models perform better; dropout is decisive for the big model.
- Sinusoidal vs. learned positional encodings: nearly identical results.
Generalization beyond MT:
- A 4-layer Transformer with \(d_{model}=1024\) on English constituency parsing (Penn Treebank WSJ; semi-supervised variant with ~17M sentences) beats all prior reported models except the Recurrent Neural Network Grammar in the WSJ-only setting, with minimal tuning.
Limitations
- Quadratic complexity in sequence length: Per-layer self-attention is \(O(n^2 \cdot d)\), problematic for very long sequences.
- Restricted attention not yet evaluated: A neighborhood-of-size-\(r\) variant is proposed but not measured; it raises max path length to \(O(n/r)\).
- Inference still sequential: Auto-regressive decoding remains inherently sequential at generation time, even though training is highly parallel.
- One domain (NMT) plus one auxiliary task (parsing): Modalities such as images, audio, and video are left as future work.
Open Problems
- Scale self-attention to very long sequences without paying \(O(n^2)\) — the authors explicitly propose investigating restricted/local self-attention.
- Reduce the sequentiality of auto-regressive generation at inference time.
- Extend Transformer-style architectures to non-text modalities (images, audio, video).
- Investigate richer compatibility functions than dot product, motivated by the observation that reducing \(d_k\) alone hurts quality.