Attention Is All You Need

“Attention Is All You Need” was written by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser and Illia Polosukhin, and presented at the 31st Conference on Neural Information Processing Systems (NIPS 2017) in Long Beach. The paper lists the authors at Google Brain, Google Research and the University of Toronto (Polosukhin without an affiliation, with a footnote that the work was done at Google Research), and a footnote states that all eight contributed equally, in random order. It proposed the Transformer, an encoder–decoder for machine translation that uses no recurrence and no convolution, only attention and position-wise feed-forward layers, and reported state-of-the-art translation quality at a fraction of the training cost of earlier systems. The paper concerned text only. Its architecture reached computer vision a few years later and is now a standard backbone there. The transformer article describes the architecture as it is used today, and the self-attention article the operation at its core.

Problem

In 2017, neural machine translation was dominated by encoder–decoder models built from recurrent networks, usually LSTMs or gated recurrent units. The paper’s introduction names their basic limitation: a recurrent network computes the hidden state at position $t$ from the state at $t-1$, so the computation within one training sequence cannot be parallelized. Factorization tricks and conditional computation had made such models faster, but the sequential dependency remained.

Attention already played a central part. Bahdanau, Cho and Bengio (2015) had let a recurrent decoder attend over all encoder states when producing each output word, so dependencies could be modeled regardless of distance. The paper observes that, with few exceptions, attention was still used together with a recurrent network.

The other route to parallel computation was convolution. The paper discusses the Extended Neural GPU, ByteNet (Kalchbrenner et al., 2016) and ConvS2S (Gehring et al., 2017), which compute representations for all positions at once. In these models, however, the number of operations needed to relate two positions grows with their distance: linearly for ConvS2S, logarithmically for ByteNet’s dilated convolutions. The authors argue, following work on gradient flow in recurrent networks, that long paths make distant dependencies harder to learn.

Contribution

The authors describe the Transformer as, to their knowledge, the first transduction model that computes representations of its input and output entirely with self-attention, without sequence-aligned recurrence or convolution. The paper’s main elements are:

• An attention-only encoder–decoder. Six encoder and six decoder layers, each built from multi-head attention and a position-wise feed-forward network, with residual connections and layer normalization around every sub-layer.
• Scaled dot-product attention. Dot-product attention with the scores divided by $\sqrt{d_k}$. The authors suspected that large dot products push the softmax into regions with tiny gradients, and explain in a footnote that the dot product of two random vectors with unit-variance components has variance $d_k$.
• Multi-head attention. Eight heads, each attending over its own learned projections of width 64, so that the model can attend to different representation subspaces at once at about the cost of a single full-width head.
• Sinusoidal positional encodings. Because nothing in the model depends on order, the authors added fixed sines and cosines of geometrically spaced wavelengths to the input embeddings, hypothesizing that they would make relative positions easy to learn and might allow extrapolation to longer sequences.
• A comparison of layer types. The paper’s Table 1 compares self-attention, recurrent and convolutional layers by computation per layer, minimum number of sequential operations and maximum path length between positions. Self-attention costs $O(n^2 d)$ per layer for $n$ positions of width $d$, but needs $O(1)$ sequential operations and connects any two positions by a path of length $O(1)$; a recurrent layer costs $O(n d^2)$ with $O(n)$ sequential operations and path length. Self-attention is therefore cheaper than recurrence whenever $n < d$, which the authors note is usually the case for sentences in word-piece or byte-pair form.

The footnote on contributions credits Uszkoreit with proposing to replace RNNs with self-attention and Shazeer with proposing scaled dot-product attention, multi-head attention and the parameter-free position representation.

Method

The architecture, its masks and its equations are covered in the transformer and self-attention articles; this section records the training setup as published.

The models were trained on WMT 2014 English–German, about 4.5 million sentence pairs with a shared byte-pair vocabulary of about 37,000 tokens, and on WMT 2014 English–French, 36 million sentences with a 32,000 word-piece vocabulary. Each batch held about 25,000 source and 25,000 target tokens. Training ran on one machine with eight NVIDIA P100 GPUs: the base model ($d_{\text{model}} = 512$, 65 million parameters) for 100,000 steps, about 12 hours, and the big model ($d_{\text{model}} = 1024$, 16 heads, 213 million parameters) for 300,000 steps, about 3.5 days.

The optimizer was Adam with $\beta_1 = 0.9$, $\beta_2 = 0.98$ and $\epsilon = 10^{-9}$, and a learning rate that rose linearly over 4,000 warmup steps and then decayed with the inverse square root of the step number. Regularization combined dropout on every sub-layer output and on the embedding sums (rate 0.1 for the base model; 0.3 for the big model, except 0.1 for English–French), and label smoothing with $\epsilon_{ls} = 0.1$ in the cross-entropy loss, which the authors note made perplexity worse but improved accuracy and BLEU. At test time they averaged the last checkpoints (5 for base, 20 for big) and decoded with beam search.

Results

On newstest2014, the big model reached 28.4 BLEU on English–German, more than 2 BLEU above the best previously reported results, including ensembles. The base model’s 27.3 BLEU also exceeded all previously published models and ensembles; the strongest earlier entry in the paper’s Table 2 is a ConvS2S ensemble at 26.36. The authors estimated training cost by multiplying training time, number of GPUs and an assumed sustained throughput per GPU: $3.3 \times 10^{18}$ FLOPs for the base model and $2.3 \times 10^{19}$ for the big model, against $7.7 \times 10^{19}$ for the ConvS2S ensemble and $1.8 \times 10^{20}$ for the GNMT + RL ensemble on English–German.

On English–French, the proceedings version reports 41.0 BLEU for the big model, better than all previously published single models at less than a quarter of the training cost of the previous state of the art (the ConvS2S and GNMT ensembles scored 41.29 and 41.16). The arXiv revision of December 2017 (v5, unchanged in later revisions on this point) gives 41.8 in its abstract and Table 2 but still says 41.0 in the results text. The abstract on the proceedings web page differs from the published PDF: it reports a single 165-million-parameter model with 27.5 BLEU on English–German and 41.1 on English–French.

The ablations in Table 3 varied the base model on the English–German development set. With computation held constant, a single attention head scored 0.9 BLEU below the best setting, and quality also fell with 32 heads. Reducing the key width $d_k$ hurt, which the authors took as a sign that a more sophisticated compatibility function than the dot product might help. Larger models did better, removing dropout cost 1.2 BLEU, and learned positional embeddings performed nearly identically to the sinusoids (25.7 against 25.8 BLEU).

The arXiv version contains two parts that the proceedings PDF lacks. A 4-layer Transformer trained for English constituency parsing on the roughly 40,000 sentences of the Wall Street Journal portion of the Penn Treebank reached 91.3 F1 on Section 23, and 92.7 F1 with about 17 million additional semi-supervised sentences; the authors report that, with little task-specific tuning, it did better than all previously reported models except the Recurrent Neural Network Grammar, although their own comparison table also lists a multi-task model at 93.0. An appendix of attention visualizations shows encoder self-attention heads in layer 5 that follow a long-distance verb dependency and others that appear to resolve pronoun reference, which the authors read as heads learning different tasks.

Impact

The transformer became the model of choice in natural language processing, as Dosovitskiy et al. (2021) put it. Its conclusion already announced the intent to apply attention to images, audio and video, using local, restricted attention to handle large inputs.

In computer vision the ideas arrived in stages. Wang et al. (2018) presented non-local neural networks for video and images and showed that the paper’s self-attention module is a special case of their non-local operation. Parmar et al. (2018), including several of the original authors, built the Image Transformer, which generates images pixel by pixel with self-attention restricted to local neighborhoods. Ramachandran et al. (2019) replaced every spatial convolution in a ResNet-50 with a self-attention layer and reported better ImageNet accuracy than the baseline with fewer FLOPs and parameters.

Two papers then made the transformer itself a vision architecture. DETR (Carion et al., 2020) feeds the features of a convolutional network to a transformer encoder–decoder and predicts a set of objects directly, and its authors state that they use standard transformer implementations. The Vision Transformer, ViT (Dosovitskiy et al., 2021), follows the original encoder as closely as possible, treating image patches as tokens, and reported that with large-scale pre-training it approached or beat state-of-the-art convolutional networks on image classification benchmarks. The transformer article covers both models and their successors.

Limitations

• Quadratic cost. Table 1 itself shows the $O(n^2 d)$ cost of self-attention. The authors proposed restricting attention to a neighborhood of size $r$, at the price of an $O(n/r)$ path length, and left it for future work; images, with many more positions than sentences, made this the main obstacle for vision.
• Order supplied from outside. The model sees positions only through the added encodings. The sinusoids’ advantage for longer sequences was a hypothesis, not tested in the paper.
• Narrow evaluation. The published results cover two translation pairs and, in the arXiv version, one parsing task.
• Training recipe. The model depended on a specific warmup schedule, label smoothing and checkpoint averaging; the paper does not isolate the effect of the warmup.

What Came After

The authors released their code in the tensor2tensor library. Later work kept the overall design but changed its details: relative position representations, pre-norm placement of layer normalization, and encoder-only and decoder-only variants, all described in the transformer article. In vision, the attention-only encoder of 2017 now underlies image classifiers, detectors and segmentation models, usually with image patches or convolutional features in place of word pieces.