Attention Is All You Need — Detailed Summary

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, Illia Polosukhin | Google Brain / Google Research / University of Toronto | 31st Conference on Neural Information Processing Systems (NIPS 2017), Long Beach, CA, USA | arXiv:1706.03762

Per-section summary organized by paper headings. Each section includes paragraph-level bullet points and exact quantitative results where the paper provides them.

Abstract

• Dominant sequence transduction models at the time of writing rely on complex recurrent (RNN) or convolutional (CNN) neural networks arranged in an encoder-decoder configuration.
• The best performing such models additionally connect encoder and decoder via an attention mechanism layered on top of recurrence.
• The paper proposes a new architecture — the Transformer — that is based solely on attention, dispensing entirely with recurrence and convolutions.
• Removing recurrence makes the model significantly more parallelizable and reduces training wall-clock time substantially compared to RNN/CNN baselines.
• On the WMT 2014 English-to-German translation task the Transformer achieves 28.4 BLEU, improving over the previous best (including ensembles) by more than 2 BLEU.
• On WMT 2014 English-to-French, a single Transformer model establishes a new single-model state-of-the-art BLEU of 41.0 after training for only 3.5 days on eight P100 GPUs — a small fraction of the cost of the best prior models.
• The architecture also generalizes well, succeeding on English constituency parsing both with and without large training data.

1. Introduction

Background and the dominance of recurrence:

• Recurrent neural networks, in particular LSTMs and gated recurrent units, have been firmly established as state of the art for sequence modeling and transduction tasks such as language modeling and machine translation.
• Substantial subsequent effort has continued to push the boundaries of recurrent encoder-decoder architectures.

The sequential-computation bottleneck:

• Recurrent models factor computation along the symbol positions of input and output sequences: a hidden state \(h_t\) is produced as a function of the previous hidden state \(h_{t-1}\) and the input at position \(t\).
• This inherently sequential dependence forecloses parallelization within a training example and becomes critical at longer sequence lengths where memory constraints further limit batching across examples.

Mitigations that fall short:

• Recent work has improved computational efficiency via factorization tricks and conditional computation, also improving model quality in some cases.
• However, the fundamental constraint of sequential computation along the position axis remains.

Rising importance of attention:

• Attention mechanisms have become an integral part of compelling sequence modeling and transduction models across many tasks, because they permit modeling of dependencies without regard to distance between positions.
• In nearly all prior cases, attention is used in conjunction with a recurrent network rather than as a replacement for one.

The Transformer proposal:

• The authors propose the Transformer, a model architecture that eschews recurrence entirely and instead relies wholly on an attention mechanism to draw global dependencies between input and output.
• The Transformer permits significantly more parallelization and reaches a new state of the art in translation quality after as little as twelve hours of training on eight P100 GPUs.

2. Background

Reducing sequential computation via convolutions:

• Several earlier models — Extended Neural GPU, ByteNet, and ConvS2S — share the goal of reducing sequential computation by using convolutional neural networks as the basic building block, computing hidden representations in parallel across input and output positions.
• In these models the number of operations required to relate signals from two arbitrary positions grows with the distance between them: linearly for ConvS2S and logarithmically for ByteNet.
• This growth makes it harder to learn dependencies between distant positions.

The Transformer's constant-distance attention:

• In the Transformer the number of operations needed to relate any two positions is reduced to a constant number, although at the cost of reduced effective resolution due to averaging attention-weighted positions.
• The authors counteract this loss of resolution with Multi-Head Attention, described in Section 3.2.

Self-attention in prior literature:

• Self-attention (also called intra-attention) is an attention mechanism that relates different positions of a single sequence in order to compute a representation of that sequence.
• Self-attention has been used successfully in reading comprehension, abstractive summarization, textual entailment, and task-independent sentence representation learning.

End-to-end memory networks:

• End-to-end memory networks are based on a recurrent attention mechanism rather than sequence-aligned recurrence and have performed well on simple language question-answering and language modeling tasks.

Novelty claim:

• To the best of the authors' knowledge, the Transformer is the first transduction model relying entirely on self-attention to compute representations of its input and output — without using sequence-aligned RNNs or convolutions.

3. Model Architecture

Encoder-decoder framing:

• Most competitive neural sequence transduction models follow an encoder-decoder structure: an encoder maps an input sequence \((x_1, \ldots, x_n)\) to a sequence of continuous representations \(z = (z_1, \ldots, z_n)\), and a decoder generates an output sequence \((y_1, \ldots, y_m)\) one symbol at a time.
• Generation is auto-regressive: at each step the model consumes the previously generated symbols as additional input.

Transformer's instantiation:

• The Transformer follows this overall architecture using stacked self-attention and point-wise, fully connected layers for both the encoder and the decoder, depicted in Figure 1.

3.1 Encoder and Decoder Stacks

Encoder:

• The encoder is a stack of \(N = 6\) identical layers.
• Each layer contains two sub-layers: a multi-head self-attention mechanism and a position-wise fully connected feed-forward network.
• A residual connection wraps each sub-layer, followed by layer normalization: \(\text{LayerNorm}(x + \text{Sublayer}(x))\).
• To enable residual addition, all sub-layers and embedding layers produce outputs of dimension \(d_{model} = 512\).

Decoder:

• The decoder is also a stack of \(N = 6\) identical layers.
• Each decoder layer adds a third sub-layer: multi-head attention over the output of the encoder stack.
• Like the encoder, residual connections around each sub-layer are followed by layer normalization.
• The self-attention sub-layer in the decoder is masked so that position \(i\) can only attend to positions less than \(i\), combined with the offset of output embeddings by one position; together these ensure that predictions for position \(i\) depend only on outputs at positions less than \(i\), preserving auto-regressive behavior.

3.2 Attention

Definition:

• An attention function maps a query and a set of key-value pairs to an output, where query, keys, values, and output are all vectors.
• The output is a weighted sum of the values, where the weight assigned to each value is computed by a compatibility function of the query with the corresponding key.

3.2.1 Scaled Dot-Product Attention

• The authors call their particular attention "Scaled Dot-Product Attention."

• Queries and keys have dimension \(d_k\); values have dimension \(d_v\).

• The attention is computed as

  \[\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{Q K^{T}}{\sqrt{d_k}}\right) V\]

• In matrix form, attention is computed on a set of queries simultaneously, packed into matrix \(Q\), with keys and values also packed into matrices \(K\) and \(V\).

• The two most commonly used attention functions are additive (Bahdanau-style) attention and dot-product (multiplicative) attention. The Transformer uses dot-product attention up to the \(1/\sqrt{d_k}\) scaling.

• Additive attention computes the compatibility function using a feed-forward network with a single hidden layer; while comparable in theoretical complexity, dot-product attention is much faster and more space-efficient in practice because it can leverage highly optimized matrix multiplication code.

• For small values of \(d_k\) the two mechanisms perform similarly; for large \(d_k\), additive attention outperforms unscaled dot-product attention.

• The authors hypothesize that for large \(d_k\) the dot products grow large in magnitude, pushing the softmax into regions with extremely small gradients; to counteract this they scale the dot products by \(1/\sqrt{d_k}\).

3.2.2 Multi-Head Attention

• Instead of performing a single attention function with \(d_{model}\)-dimensional keys, values, and queries, the authors found it beneficial to linearly project the queries, keys, and values \(h\) times with different learned linear projections to \(d_k\), \(d_k\), and \(d_v\) dimensions respectively.

• Attention is performed on each of these projected versions in parallel, yielding \(d_v\)-dimensional outputs that are concatenated and once again projected to produce the final values.

  \[\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) W^{O}\] \[\text{head}_i = \text{Attention}(Q W^{Q}_i, K W^{K}_i, V W^{V}_i)\]

• Multi-head attention allows the model to jointly attend to information from different representation subspaces at different positions; with a single head, averaging inhibits this.

• Hyperparameters: \(h = 8\) parallel attention heads; \(d_k = d_v = d_{model}/h = 64\).

• Because the dimension of each head is reduced, total computational cost is similar to that of single-head attention with full dimensionality.

3.2.3 Applications of Attention in our Model

• Encoder-decoder attention: Queries come from the previous decoder layer; the memory keys and values come from the output of the encoder. This allows every decoder position to attend over all input positions, mimicking the classic encoder-decoder attention of seq2seq models.
• Encoder self-attention: All keys, values, and queries come from the output of the previous encoder layer; each encoder position can attend to all positions in the previous layer.
• Decoder self-attention: Permits each decoder position to attend to all previous positions in the decoder up to and including the current one. To preserve the auto-regressive property, leftward information flow is masked out by setting all values in the input of the softmax that correspond to illegal connections to \(-\infty\) (equivalent to forbidding those keys).

3.3 Position-wise Feed-Forward Networks

• In addition to attention sub-layers, each encoder and decoder layer contains a fully connected feed-forward network applied to each position separately and identically.

• The FFN consists of two linear transformations with a ReLU activation in between:

  \[\text{FFN}(x) = \max(0, x W_1 + b_1) W_2 + b_2\]

• While the linear transformations are the same across different positions, they use different parameters from layer to layer.

• An equivalent description is two convolutions with kernel size 1.

• Dimensionality: input/output \(d_{model} = 512\); inner-layer \(d_{ff} = 2048\).

3.4 Embeddings and Softmax

• Similarly to other sequence transduction models, the Transformer uses learned embeddings to convert input tokens and output tokens to vectors of dimension \(d_{model}\).
• The usual learned linear transformation and softmax are used to convert the decoder output to predicted next-token probabilities.
• The model shares the same weight matrix between the two embedding layers and the pre-softmax linear transformation.
• In the embedding layers, those weights are multiplied by \(\sqrt{d_{model}}\).

3.5 Positional Encoding

• Since the model contains no recurrence and no convolution, it must inject some information about the relative or absolute position of tokens.

• The authors add positional encodings to the input embeddings at the bottoms of the encoder and decoder stacks; these have the same dimension \(d_{model}\) as the embeddings so that the two can be summed.

• The authors use sinusoidal positional encodings:

  \[PE_{(pos, 2i)} = \sin(pos / 10000^{2i/d_{model}})\] \[PE_{(pos, 2i+1)} = \cos(pos / 10000^{2i/d_{model}})\]

• Each dimension of the positional encoding corresponds to a sinusoid; the wavelengths form a geometric progression from \(2\pi\) to \(10000 \cdot 2\pi\).

• This function was chosen because it would allow the model to easily learn to attend by relative positions: for any fixed offset \(k\), \(PE_{pos+k}\) can be represented as a linear function of \(PE_{pos}\).

• The authors also experimented with learned positional embeddings and found nearly identical results; the sinusoidal version was chosen because it may allow extrapolation to sequence lengths longer than those seen in training.

4. Why Self-Attention

Desiderata for the comparison:

• The authors compare self-attention layers to recurrent and convolutional layers along three axes: (1) total computational complexity per layer, (2) the amount of computation that can be parallelized, measured by the minimum number of sequential operations required, and (3) the path length between long-range dependencies in the network.

Per-layer complexity:

• A self-attention layer connects all positions with a constant number of sequentially executed operations, \(O(1)\), while a recurrent layer requires \(O(n)\) sequential operations.
• Self-attention is faster than recurrent layers in terms of computational complexity when the sequence length \(n\) is smaller than the representation dimensionality \(d\) — common for state-of-the-art models in machine translation, where word-piece or byte-pair tokenizations keep \(n\) modest.

Path length and long-range dependencies:

• Shorter paths between any combination of positions in the input/output sequences make it easier to learn long-range dependencies, a key challenge in sequence transduction.
• Maximum path length is summarized in Table 1: self-attention \(O(1)\), recurrent \(O(n)\), convolutional \(O(\log_k n)\), restricted self-attention \(O(n/r)\).

Restricted self-attention for very long sequences:

• To improve computational performance for tasks involving very long sequences, self-attention could be restricted to consider only a neighborhood of size \(r\) in the input sequence centered around each output position — increasing the maximum path length to \(O(n/r)\).

Convolutions:

• A single convolutional layer with kernel width \(k < n\) does not connect all pairs of positions; doing so requires a stack of \(O(n/k)\) convolutional layers, increasing the longest paths.
• Convolutional layers are generally more expensive than recurrent layers by a factor of \(k\); separable convolutions reduce complexity considerably to \(O(k \cdot n \cdot d + n \cdot d^2)\) — comparable to a self-attention plus point-wise feed-forward layer in the Transformer.

Interpretability:

• As a side benefit, self-attention can yield more interpretable models. The authors report that individual attention heads clearly learn to perform different tasks, with many appearing to exhibit behavior related to syntactic and semantic structure of sentences (such as anaphora).

Table 1 (paraphrased): Layer-type comparison

5. Training

5.1 Training Data and Batching

• WMT 2014 English-German: about 4.5 million sentence pairs, encoded with byte-pair encoding (BPE) and a shared source-target vocabulary of roughly 37,000 tokens.
• WMT 2014 English-French: about 36 million sentence pairs, with tokens split into a 32,000 word-piece vocabulary.
• Sentence pairs were batched together by approximate sequence length; each training batch contained a set of sentence pairs with roughly 25,000 source tokens and 25,000 target tokens.

5.2 Hardware and Schedule

• The authors trained their models on one machine with 8 NVIDIA P100 GPUs.
• For the base models using the hyperparameters described, each training step took about 0.4 seconds; base models were trained for 100,000 steps, or approximately 12 hours.
• For the big models each step took about 1.0 seconds; big models were trained for 300,000 steps, or approximately 3.5 days.

5.3 Optimizer

• Adam with \(\beta_1 = 0.9\), \(\beta_2 = 0.98\), \(\epsilon = 10^{-9}\).

• Learning rate is varied over training:

  \[lrate = d_{model}^{-0.5} \cdot \min(\text{step\_num}^{-0.5}, \text{step\_num} \cdot \text{warmup\_steps}^{-1.5})\]

• This corresponds to linearly increasing the learning rate for the first \(warmup\_steps = 4000\) training steps, then decreasing it proportionally to the inverse square root of the step number.

5.4 Regularization

• Residual Dropout: Dropout applied to the output of each sub-layer before it is added to the sub-layer input and normalized; also applied to the sums of embeddings and positional encodings in both encoder and decoder stacks. Base model uses \(P_{drop} = 0.1\).
• Label Smoothing: During training the authors employ label smoothing of value \(\epsilon_{ls} = 0.1\). This hurts perplexity (because the model learns to be less confident) but improves accuracy and BLEU score.

6. Results

6.1 Machine Translation

• On WMT 2014 English-to-German, the big Transformer model outperforms the best previously reported models (including ensembles) by more than 2.0 BLEU, establishing a new state-of-the-art BLEU score of 28.4.
• The configuration of this big model is listed at the bottom of Table 3 (large variant).
• Training took 3.5 days on 8 P100 GPUs; notably, even the base model surpasses all previously published models and ensembles, at a fraction of the training cost of any of the competitive models.
• On WMT 2014 English-to-French, the big model achieves a BLEU score of 41.0, outperforming all previously published single models at less than 1/4 the training cost of the previous state-of-the-art model.
• The big En-Fr model uses dropout rate \(P_{drop} = 0.1\) (rather than 0.3 as in En-De).

Table 2 (paraphrased): BLEU comparison with selected baselines

Training cost (FLOPs):

• Transformer (base, En-De): \(3.3 \times 10^{18}\) FLOPs.
• Transformer (big, En-De): \(2.3 \times 10^{19}\) FLOPs.
• The authors compare these training costs to those of strong baselines and argue that the Transformer achieves better quality at substantially lower training cost.

Inference:

• For the base models, the authors used a single model obtained by averaging the last 5 checkpoints (written 10 minutes apart); for the big models, the last 20 checkpoints were averaged.
• Beam search of size 4 with length penalty \(\alpha = 0.6\).

6.2 Model Variations

• To evaluate the importance of different components of the Transformer, the authors varied the base model in different ways and measured the change in performance on English-to-German translation on the newstest2013 development set.
• Rows (A): varying the number of attention heads and the attention key/value dimensions while keeping the total amount of computation constant. Single- head attention is 0.9 BLEU worse than the best setting; quality also drops with too many heads.
• Rows (B): reducing the attention key size \(d_k\) alone hurts model quality, suggesting that determining compatibility is non-trivial and that a more sophisticated compatibility function than dot product may be beneficial.
• Rows (C) and (D): bigger models perform better; dropout is very helpful in avoiding over-fitting.
• Row (E): replacing the sinusoidal positional encoding with learned positional embeddings yields nearly identical results.

6.3 English Constituency Parsing

• To evaluate generalization beyond machine translation, the authors trained a 4-layer Transformer with \(d_{model} = 1024\) on the Penn Treebank Wall Street Journal portion (about 40K training sentences) and in a semi-supervised setting using larger high-confidence and BerkleyParser corpora (about 17M sentences).
• Despite minimal tuning, the model performs surprisingly well, yielding better results than all previously reported models with the exception of the Recurrent Neural Network Grammar in the WSJ-only setting.
• This demonstrates that the Transformer generalizes to other tasks with different output structures.

7. Conclusion

• The authors presented the Transformer, the first sequence transduction model based entirely on attention, replacing the recurrent layers most commonly used in encoder-decoder architectures with multi-headed self-attention.
• For translation tasks, the Transformer can be trained significantly faster than architectures based on recurrent or convolutional layers; on both WMT 2014 En-De and En-Fr translation tasks, the model achieves a new state of the art.
• They are excited about the future of attention-based models and plan to apply them to problems involving input and output modalities other than text — images, audio, and video — and to investigate local, restricted attention mechanisms to handle large inputs and outputs efficiently.
• Making generation less sequential is another research goal.
• Code used to train and evaluate the models is available at github.com/tensorflow/tensor2tensor.

Figures

• Figure 1: The Transformer — model architecture. Shows the full encoder-decoder stack: input embeddings plus positional encodings feeding into \(N = 6\) encoder layers (each: multi-head self-attention, add-and-norm, position-wise feed-forward, add-and-norm); output embeddings (shifted right) plus positional encodings feeding into \(N = 6\) decoder layers (each: masked multi-head self-attention, encoder-decoder multi-head attention, position-wise feed-forward, with residual connections and layer norm after each sub-layer); final linear plus softmax to produce output probabilities.
• Figure 2 (left): Scaled Dot-Product Attention. Diagram of the inner attention computation showing MatMul of Q and K, Scale by \(1/\sqrt{d_k}\), optional Mask, Softmax, and final MatMul with V.
• Figure 2 (right): Multi-Head Attention. Diagram showing \(h\) parallel linear projections of V, K, Q feeding into Scaled Dot-Product Attention blocks, whose outputs are concatenated and projected linearly to produce the final output.

Hyperparameter Summary

Limitations and Future Work

• Quadratic self-attention cost in sequence length: Per-layer complexity is \(O(n^2 \cdot d)\), which becomes problematic for very long sequences.
• Restricted attention: The authors propose investigating local / restricted self-attention (neighborhood of size \(r\)) as a remedy, at the cost of path length \(O(n/r)\).
• Other modalities: Extending Transformer-style architectures to images, audio, and video is left as future work.
• Sequential generation: The auto-regressive decoder is still inherently sequential at inference time; making generation less sequential is flagged as an open direction.

Focal References

Cross-Cutting Take-Aways