5Attention Is All You Need
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser and Illia Polosukhin, NeurIPS, 2017
Linked to 12 chapters: supports 5 · contradicts 2 · mechanism 9 · pairing 1
Eight authors at Google Brain and Google Research described a sequence model with no recurrent connections and no convolutions in it. Every competitive translation system of the previous three years had read a sentence one position at a time, each hidden state waiting on the one before it: a recurrent network cannot compute position twelve before it has computed position eleven, so a single training example cannot be spread across processors. Vaswani and his co-authors kept the encoder–decoder shape that machine translation had settled on and replaced its recurrence with attention alone, "dispensing with recurrence and convolutions entirely". What they gained was arithmetic that runs all at once. A base model trained for twelve hours on eight graphics processors beat every published result on English-to-German; a larger version trained for three and a half days reached 28.4 BLEU, more than two points above the best prior score including ensembles, at roughly a tenth of the training cost. The parts they assembled — six layers, eight attention heads, a 512-dimensional representation, a feed-forward network at every position, a residual connection and a normalisation around every sub-layer, and sinusoids added to the embeddings to stand for order — are still the parts.
The cost of doing one position at a time
Recurrent networks factor computation along the positions of the sequence. The hidden state at step t is a function of the state at t − 1 and the input at position t, so the steps line up with the positions and have to be taken in order. Vaswani and his co-authors say plainly what that forbids: no parallelisation within a training example, which bites hardest at long sequences, because memory limits how many examples can be batched together to make up for it. Factorisation tricks and conditional computation had made recurrent models cheaper in the year before; the constraint itself remains.
Attention was not their invention. Since Bahdanau, Cho and Bengio in 2014 it had been ordinary for a decoder to consult a learned, weighted view of every encoder position, which lets a dependency be modelled without regard to the distance between the two positions involved. But in all but a few cases, the authors observe, attention had been bolted onto a recurrent network rather than used in place of one.
The other escape route was convolution. The Extended Neural GPU, ByteNet and ConvS2S all computed positions in parallel out of convolutional blocks, at the cost that the number of operations needed to relate two positions grows with the distance between them — linearly for ConvS2S, logarithmically for ByteNet — so a distant dependency is harder to learn than a near one. Self-attention makes that number constant. The price is resolution, since an attention output is an average over weighted positions, and multi-head attention is the authors' answer to it.
The encoder and the decoder stacks
The encoder is six identical layers. Each has two sub-layers: multi-head self-attention over the previous layer's output, then a fully connected feed-forward network applied at every position. Each sub-layer is wrapped, so that the layer emits LayerNorm(x + Sublayer(x)) — the sub-layer's output added back to its own input and then normalised. For that addition to be defined, every sub-layer and every embedding layer produces vectors of the same width, 512.
The decoder is six identical layers of three sub-layers. The first is self-attention over what the decoder has produced so far, masked so that no position may attend to a later one; combined with shifting the output embeddings one position to the right, the masking is what keeps the model auto-regressive, since the prediction at position i can then depend only on outputs before i. The second sub-layer is attention whose queries come from the decoder and whose keys and values come from the encoder's output, so every decoder position can look at every input position. The third is the feed-forward network again.
Those are the three uses of attention in the model, and the only three: encoder self-attention, masked decoder self-attention, and the encoder–decoder attention that joins them.
Attention
An attention function maps a query and a set of key–value pairs to an output, all of them vectors. The output is a weighted sum of the values, and the weight on each value is computed by a compatibility function between the query and that value's key. Everything else is a choice of compatibility function.
The choice here is the dot product, scaled. Take the dot product of the query with every key, divide each result by the square root of the key width, pass the whole set through a softmax, and use the numbers that come out as weights on the values. Packed into matrices — queries as the rows of Q, keys as K, values as V — the operation is softmax(QKᵀ / √dₖ)V: one matrix multiply, one scale, one softmax, one matrix multiply.
Two justifications are given. The dot product is preferred to additive attention, which computes compatibility with a small feed-forward network, because the two have similar theoretical complexity while the dot product is nothing but highly optimised matrix multiplication, and so faster and more space-efficient in practice. The scaling is there because unscaled dot-product attention had been found to lose to additive attention at large key widths, and the authors suspect why: if the components of a query and a key are independent with mean zero and variance one, their dot product has variance equal to the key width, so the numbers entering the softmax grow with that width and push it into regions where its gradients are extremely small. Dividing by the square root of the width puts the variance back at one.
Rather than one attention over 512-dimensional queries, keys and values, the model runs eight. Each head has its own learned linear projections down to 64 dimensions for queries, keys and values; each runs the same scaled dot-product attention; the eight 64-dimensional outputs are concatenated and projected once more. Because each head is an eighth as wide, the total cost is about that of a single head at full width. The reason for wanting eight is that multi-head attention lets the model attend jointly to information from different representation subspaces at different positions. "With a single attention head, averaging inhibits this."
Position, without recurrence
The rest of a layer is small. The position-wise feed-forward network is two linear transformations with a rectified linear unit between them, applied separately and identically at every position: 512 in, 2,048 in the middle, 512 out. The same parameters serve every position within a layer and differ from layer to layer, which the authors remark is another way of describing two convolutions with kernel size one. Input and output tokens become vectors through learned embeddings, and the decoder's output becomes next-token probabilities through a linear map and a softmax; one weight matrix serves all three.
That leaves order. Nothing described so far can tell one arrangement of the same tokens from another, because attention is a sum over a set. So the authors add a positional encoding to the input embeddings at the bottom of each stack, of the same width so that the two can simply be summed. Their encoding is fixed rather than learned: each dimension is a sinusoid, sine on the even dimensions and cosine on the odd, with wavelengths in a geometric progression from 2π to 10000 · 2π. The reason given is that for any fixed offset k, the encoding at position pos + k is a linear function of the encoding at pos, so relative position should be easy to attend by. A version trained with learned positional embeddings scored almost the same — 25.7 against 25.8 BLEU on the development set — and the sinusoids were kept on the hypothesis that they might extrapolate to sequences longer than any seen in training.
Why self-attention
The case is made on three counts: the computation a layer costs, how much of it can be parallelised — measured as the minimum number of operations that must be taken in sequence — and the longest path a signal must travel between two positions.
| Layer type | Complexity per layer | Sequential operations | Maximum path length |
|---|---|---|---|
| Self-attention | O(n² · d) | O(1) | O(1) |
| Recurrent | O(n · d²) | O(n) | O(n) |
| Convolutional | O(k · n · d²) | O(1) | O(log_k n) |
| Self-attention, restricted | O(r · n · d) | O(1) | O(n/r) |
The third column is the one the argument turns on. A dependency between two positions gets harder to learn as the forward and backward signals between them get longer, and a self-attention layer connects every position to every other in one step. The second column is the one the training budget turns on: a recurrent layer needs as many sequential operations as there are positions, a self-attention layer needs one, so a whole sentence can be computed at once. The first column is the bill, and it is quadratic in the sequence length — which the authors treat as harmless, because self-attention costs less per layer than recurrence whenever the sequence is shorter than the representation is wide, and with word-piece or byte-pair vocabularies a sentence is a few dozen tokens against a width of 512. For longer inputs they suggest restricting each position's attention to a neighbourhood, at the cost of a path length that grows again, and leave it as future work. They add one benefit they do not measure: attention distributions can be read off the trained model, individual heads visibly learn different jobs, and many appear to track syntactic and semantic structure.
Training, and what it scored
The English-to-German models were trained on the 4.5 million sentence pairs of WMT 2014 with a shared byte-pair vocabulary of about 37,000 tokens, the English-to-French models on 36 million sentences with a 32,000-token word-piece vocabulary. Batches held roughly 25,000 source and 25,000 target tokens, grouped by length. Everything ran on one machine with eight NVIDIA P100 processors: the base model took 0.4 seconds a step for 100,000 steps, twelve hours in all; the big model took a second a step for 300,000 steps, three and a half days. The optimiser was Adam, on a schedule that raises the learning rate linearly through 4,000 warm-up steps and then decays it as the inverse square root of the step number. Dropout was applied to each sub-layer's output before it is added back and to the sum of the embeddings and the positional encodings, and label smoothing of 0.1 was used, which the authors note hurts perplexity, because the model learns to be less sure of itself, while improving accuracy and BLEU.
The big model scored 28.4 BLEU on English-to-German, more than two points above anything previously reported, ensembles included. The base model's 27.3 was already above every published result, at 3.3 × 10¹⁸ floating-point operations of training against the 7.7 × 10¹⁹ of the ConvS2S ensemble it beat. On English-to-French, Table 2 records 41.8 for the big model, ahead of every published single model at less than a quarter of the previous best system's training cost; the text of the results section gives 41.0 for the same run.
Table 3 varies the base model one component at a time and reports development-set BLEU against a baseline of 25.8. Holding the total computation fixed and changing only how it is divided among heads, a single head scores 0.9 below the best setting and thirty-two heads score below sixteen; eight and sixteen tie at the top. Cutting the key width from 64 to 16 costs 0.7, from which the authors draw a conclusion against their own mechanism: determining compatibility is not easy, and something more elaborate than a dot product may do better. Size helps throughout the range tried — two layers instead of six loses 2.1 points, widening the representation to 1,024 gains 0.2, widening the inner feed-forward layer to 4,096 gains 0.4. Removing dropout costs 1.2, and replacing the sinusoids with learned position embeddings costs 0.1, which is nothing.
The architecture was not particular to translation. A four-layer, 1,024-wide model trained on English constituency parsing, whose output obeys strong structural constraints and runs much longer than its input, reached 91.3 F1 on Section 23 of the Penn Treebank from the 40,000 Wall Street Journal sentences alone, above the Berkeley parser's 90.4; trained semi-supervised on about 17 million sentences it reached 92.7, behind only a recurrent neural network grammar.
The paper ends on a plan rather than a result: to extend the model to inputs and outputs other than text — images, audio, video — to investigate local and restricted attention for inputs too large to attend over in full, and to make generation less sequential.
Since 2017
The encoder–decoder pair came apart first. Within a year two models had each kept one half — a stack of encoder layers trained to fill in masked tokens, and a stack of decoder layers trained to predict the next token. For generation the decoder-only half won, and the cross-attention sub-layer that had joined the two stacks dropped out of the dominant design. What survived intact is everything inside a layer: masked self-attention over queries, keys and values, several heads of it, the position-wise feed-forward network with its wider middle, and the residual connection that carries one representation up through the whole stack.
Then the model was made larger, which the paper neither proposed nor tested. Its biggest configuration held 213 million parameters and trained in under four days. The architecture turned out to improve smoothly as parameters, data and computation were raised together, and by 2020 published fits of that relation were being used to set model sizes and training budgets before a run began. Nothing in the paper anticipates this: its argument for the design is that it trains faster and translates better, and its measure of success is BLEU.
Three choices were later revised. Layer normalisation moved from after the residual addition to before the sub-layer, which made deep stacks trainable without the warm-up schedule the paper needed. The sinusoids gave way to relative and rotary position encodings, which express the offset between two positions inside the attention computation instead of adding a vector at the bottom of the stack. And the number of heads producing keys and values was cut below the number producing queries, to shrink what must be held in memory while generating.
The quadratic term changed character. At sentence length it is the smaller cost, and the paper says so. Once contexts reached thousands and then hundreds of thousands of positions it became the binding constraint on training and inference alike, and restricted attention — which the paper had named as future work and set aside — turned into a permanent branch of the subject.
The last sentence of the paper has held up best. Applying the same encoder to fixed-size patches of an image, with the patches treated as a sequence, worked in 2020; the same treatment of audio frames and of video followed. What the authors offered as a plan for other modalities turned out to be a description of what they had built: a general way of letting every element of a set consult every other, with order supplied from outside rather than wired in.