Papers
Attention Is All You Need
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, Illia Polosukhin · NeurIPS 2017 (31st Conference on Neural Information Processing Systems, then called NIPS), 2017
The problem: recurrent networks are sequential by construction
At the time of the paper, the dominant models for sequence transduction (such as machine translation) were recurrent or convolutional networks with an encoder-decoder structure, and the best ones still connected the encoder and decoder through an attention mechanism. Recurrent networks, including LSTM and GRU, process the sequence along the symbol positions: the hidden state at position t is computed from the state at the previous position and the input at t. According to the authors, this inherently sequential nature prevents parallelization within a single training example, which matters more for longer sequences, because memory limits how many examples can be batched together. Factorization tricks and conditional computation improved efficiency, but the fundamental constraint of sequential computation remained. Convolutional models (Extended Neural GPU, ByteNet, ConvS2S) reduce sequentiality by computing representations for all positions in parallel, but the number of operations needed to relate two distant positions grows with the distance: linearly in ConvS2S and logarithmically in ByteNet, which makes it harder to learn long-range dependencies. Attention, in turn, already made it possible to model dependencies regardless of the distance between positions, but it was almost always used together with a recurrent network. The paper's proposal is the Transformer, an architecture that dispenses with recurrence and convolutions and relies on attention alone to relate input and output.
Self-attention: query, keys, and values
The Transformer keeps the encoder-decoder structure, but each of the two stacks is made of 6 identical layers. Each encoder layer has two sublayers: multi-head self-attention and a feed-forward network applied to each position separately and identically. Around each sublayer there is a residual connection followed by layer normalization, and all outputs have dimension 512. The decoder adds a third sublayer, which attends to the encoder output, and masks the self-attention so that the prediction for position i depends only on the known outputs at positions before i, preserving auto-regressive generation. Attention is described as a mapping from a query and a set of key-value pairs to an output: a weighted sum of the values, where the weight comes from a compatibility function between the query and each key. The paper uses scaled dot-product attention: the dot product of the query with all the keys is computed, divided by the square root of the key dimension, and a softmax is applied to obtain the weights. The authors suspect that, for large dimensions, the dot products grow in magnitude and push the softmax into regions with very small gradients, and the division serves to compensate for that. In self-attention, queries, keys and values come from the same place, the output of the previous layer, and each position can attend to all positions in that layer.
Multi-head attention and position information
Instead of a single attention function over full-dimension vectors, the paper projects queries, keys, and values h times with different, learned linear projections, applies attention in parallel to each projected version, concatenates the results, and projects again. The authors claim this lets the model attend, at the same time, to information from different representation subspaces at different positions, something a single head would inhibit by averaging. In the base configuration there are h = 8 heads of dimension 64 each, so the total cost is similar to that of a single full-dimension head. In the model-variation experiments, a single head was 0.9 BLEU worse than the best setting, and quality also drops with too many heads. Since there is no recurrence or convolution, the model cannot see order on its own, so the paper adds positional encodings of the same dimension to the input embeddings. They use sines and cosines of different frequencies, with wavelengths in a geometric progression. The authors' hypothesis is that this makes it easier to learn to attend by relative positions, since, for a fixed offset, the encoding at pos + k is a linear function of the encoding at pos. They also tested learned positional embeddings and got nearly identical results; they preferred the sinusoidal version because it may allow extrapolating to sequences longer than those seen in training.
Why self-attention: the per-layer cost table
The 'Why Self-Attention' section compares self-attention, recurrent and convolutional layers on three criteria: computational complexity per layer, the amount of computation that can be parallelized (measured by the minimum number of sequential operations) and the maximum path length between two positions, since shorter paths make it easier to learn long-range dependencies. In Table 1, with n the sequence length, d the representation dimension, k the kernel size and r the neighborhood size in restricted attention, the values are as follows. Self-attention: complexity O(n² · d), sequential operations O(1), maximum path O(1). Recurrent: O(n · d²), O(n) and O(n). Convolutional: O(k · n · d²), O(1) and O(log_k n). Self-attention restricted to a neighborhood r: O(r · n · d), O(1) and O(n/r). The authors note that self-attention is faster than the recurrent layer when n is smaller than d, which often happens with the sentence representations used in translation, such as word-piece and byte-pair. For very long sequences, they propose as future work restricting attention to a neighborhood, at the cost of increasing the maximum path length. The trade-off is clear: full attention connects all positions with a constant number of sequential operations, but the n² term makes the cost grow quickly with sequence length. As a side benefit, the authors point to potentially more interpretable models, since several heads seem to learn distinct tasks.
Results in translation and the legacy stated by the authors
The models were trained on one machine with 8 NVIDIA P100 GPUs. The base model took about 12 hours (100 thousand steps) and the big model 3.5 days (300 thousand steps). On the WMT 2014 English-to-German translation task, the big model reached 28.4 BLEU, more than 2 BLEU above the best previous results, including ensembles; even the base model surpassed all previously published models and ensembles, at a much lower training cost. On English-to-French translation, the paper's abstract reports 41.8 BLEU as a new state of the art for a single model, after 3.5 days of training on eight GPUs, a fraction of the cost of the best models in the literature (the body of the text cites 41.0 for that same model, an internal inconsistency in the paper). The authors estimate the cost in FLOPs by multiplying training time, number of GPUs, and an estimate of each GPU's sustained capacity. In the variations of Table 3, bigger models were better and dropout helped against overfitting. The Transformer was also applied to English constituency parsing, where it performed well even without task-specific tuning. In the conclusion, the authors describe themselves as the authors of the first transduction model based entirely on attention and plan to extend it to modalities other than text, investigate local, restricted attention for large inputs, and make generation less sequential. The paper does not describe later uses; any legacy beyond that comes from sources external to it.
See the ideas at work
Practice the designs that use this paper in the System Design school.