Attention Is All You Need

A sequence transduction model with an encoder–decoder structure that relies entirely on attention (self-attention and encoder–decoder attention plus position-wise feed-forward layers), dispensing with recurrence and convolutions.

RNNs compute hidden states sequentially, \(h_t\) from \(h_{t-1}\), which precludes parallelization within a training example, critical for long sequences when memory limits batching. Factorization tricks and conditional computation helped, but the sequential constraint remained.

On WMT 2014 English–German, 28.4 BLEU, over 2 BLEU better than the best previous results including ensembles; on English–French, a single-model state of the art of 41.0 BLEU after 3.5 days on eight P100 GPUs, a small fraction of earlier training costs.

They compute representations in parallel, but the number of operations to relate two positions grows with their distance (linearly for ConvS2S, logarithmically for ByteNet). In the Transformer it is constant, at the cost of reduced resolution from averaging attention-weighted positions, which multi-head attention counteracts.

An attention mechanism relating different positions of a single sequence to compute a representation of that sequence (also called intra-attention). The Transformer was the first transduction model relying entirely on it, without sequence-aligned RNNs or convolutions.

Two sub-layers: multi-head self-attention, then a position-wise feed-forward network. Each sub-layer is wrapped as \(\text{LayerNorm}(x + \text{Sublayer}(x))\), a residual connection followed by layer normalization; all outputs have dimension \(d_{\text{model}} = 512\). The encoder stacks \(N = 6\) such layers.

It has a third sub-layer, multi-head attention over the encoder output, and its self-attention is masked so position \(i\) can’t attend to later positions; with outputs shifted right by one, predictions for position \(i\) depend only on outputs before \(i\).

A mapping from a query and a set of key–value pairs to an output: a weighted sum of the values, where each value’s weight comes from a compatibility function of the query with the corresponding key.

\(\text{Attention}(Q, K, V) = \text{softmax}\big(\frac{QK^\top}{\sqrt{d_k}}\big) V\), with queries and keys of dimension \(d_k\) and values of dimension \(d_v\), packed as rows of matrices \(Q\), \(K\), \(V\).

If query and key components are independent with mean 0 and variance 1, \(q \cdot k = \sum_{i=1}^{d_k} q_i k_i\) has variance \(d_k\). For large \(d_k\) dot products get large, pushing the softmax into regions with extremely small gradients; scaling keeps the variance at 1.

Similar theoretical complexity, but dot-product attention is much faster and more space-efficient in practice with optimized matrix multiplication. Without scaling, additive attention outperforms it for large \(d_k\); the scaling fixes that.

Project queries, keys and values \(h\) times with different learned matrices, run attention in parallel on each projection, concatenate and project again: \(\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) W^O\), \(\text{head}_i = \text{Attention}(Q W_i^Q, K W_i^K, V W_i^V)\).

Heads can jointly attend to information from different representation subspaces at different positions; a single head averages them. With \(h = 8\) and \(d_k = d_v = d_{\text{model}}/h = 64\), the total cost is similar to single-head attention with full dimensionality.

Encoder–decoder attention (queries from the previous decoder layer, keys and values from the encoder output, so every decoder position attends over the whole input); encoder self-attention (all from the previous encoder layer); and masked decoder self-attention (each position attends to positions up to and including itself).

Inside scaled dot-product attention, inputs to the softmax for illegal connections (from a position to later ones) are set to \(-\infty\), so their weights become 0 and the model stays auto-regressive.

Two linear maps with a ReLU, applied to each position separately and identically: \(\text{FFN}(x) = \max(0, x W_1 + b_1) W_2 + b_2\), with inner dimension \(d_{ff} = 2048\) and parameters differing between layers (equivalently, two convolutions with kernel size 1).

The same weight matrix is shared between the input and output embedding layers and the pre-softmax linear transformation; in the embedding layers the weights are multiplied by \(\sqrt{d_{\text{model}}}\).

Without recurrence or convolution the model has no notion of order, so encodings are added to the input embeddings: \(PE_{(pos, 2i)} = \sin(pos / 10000^{2i/d_{\text{model}}})\), \(PE_{(pos, 2i+1)} = \cos(pos / 10000^{2i/d_{\text{model}}})\), sinusoids with wavelengths from \(2\pi\) to \(10000 \cdot 2\pi\).

For any fixed offset \(k\), \(PE_{pos+k}\) is a linear function of \(PE_{pos}\), which may help attend by relative position, and sinusoids may extrapolate to sequences longer than those seen in training. Learned embeddings gave nearly identical results.

Per-layer cost, minimum sequential operations (seq.) and maximum path length between positions (\(n\) = sequence length, \(d\) = dimension, \(k\) = kernel width):

layer	cost	seq.	path
self-attn.	\(O(n^2 d)\)	\(O(1)\)	\(O(1)\)
recurrent	\(O(n d^2)\)	\(O(n)\)	\(O(n)\)
convol.	\(O(k n d^2)\)	\(O(1)\)	\(O(\log_k n)\)

When sequence length \(n\) is smaller than dimension \(d\), usually the case for sentences with word-piece or BPE tokens. For very long sequences, attention could be restricted to a neighbourhood of size \(r\), raising the maximum path length to \(O(n/r)\).

Learning long-range dependencies is a key challenge, and the shorter the paths forward and backward signals must traverse between any input and output positions, the easier such dependencies are to learn. Self-attention connects all positions in a constant number of steps.

WMT 2014 English–German (about 4.5M sentence pairs, shared BPE vocabulary of about 37,000 tokens) and English–French (36M sentences, 32,000 word-pieces). Batches held about 25,000 source and 25,000 target tokens, grouped by length.

Adam (\(\beta_1 = 0.9\), \(\beta_2 = 0.98\), \(\epsilon = 10^{-9}\)) with \(\text{lrate} = d_{\text{model}}^{-0.5} \cdot \min(\text{step}^{-0.5}, \text{step} \cdot \text{warmup}^{-1.5})\): linear warm-up for 4000 steps, then decay proportional to the inverse square root of the step number.

Residual dropout (on each sub-layer output before the residual addition, and on embeddings plus positional encodings; \(P_{\text{drop}} = 0.1\) for the base model) and label smoothing of 0.1, which hurts perplexity but improves accuracy and BLEU.

The base model, trained in 12 hours, already beat all previous models and ensembles on English–German; the big model reached 41.0 BLEU on English–French at less than 1/4 of the previous state of the art’s training cost. Checkpoint averaging and beam search (beam 4, length penalty 0.6) were used.

A single head was 0.9 BLEU worse than the best setting, but too many heads also hurt; smaller key size \(d_k\) hurt, suggesting compatibility is hard to compute; bigger models were better; dropout was very helpful against overfitting; learned positional embeddings matched sinusoids.

Applying attention-based models to other tasks and modalities (images, audio, video), local restricted attention for large inputs and outputs, and making generation less sequential.