Attention Is All You Need

Ashish Vaswani; Noam Shazeer; Niki Parmar; Jakob Uszkoreit; Llion Jones; Aidan N. Gomez; Lukasz Kaiser; Illia Polosukhin

Papers · arXiv 1706.03762

Attention Is All You Need

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, Illia Polosukhin

Published 2017-06-12cs.CLcs.LG arXiv PDF

Read with translation and AI

Summary in three lines

🎯 Recurrent models execute sequentially along token positions, blocking parallelization within training sequences and struggling with long-range path lengths.
🔧 Introduces the Transformer, an encoder-decoder architecture based entirely on scaled dot-product multi-head self-attention and position-wise feed-forward networks.
📈 Achieves state-of-the-art translation with 28.4 BLEU on WMT 2014 English-to-German and 41.8 BLEU on English-to-French, training in 3.5 days on 8 GPUs.

Abstract

The dominant sequence transduction models are based on complex recurrent or convolutional neural networks in an encoder-decoder configuration. The best performing models also connect the encoder and decoder through an attention mechanism. We propose a new simple network architecture, the Transformer, based solely on attention mechanisms, dispensing with recurrence and convolutions entirely. Experiments on two machine translation tasks show these models to be superior in quality while being more parallelizable and requiring significantly less time to train. Our model achieves 28.4 BLEU on the WMT 2014 English-to-German translation task, improving over the existing best results, including ensembles by over 2 BLEU. On the WMT 2014 English-to-French translation task, our model establishes a new single-model state-of-the-art BLEU score of 41.8 after training for 3.5 days on eight GPUs, a small fraction of the training costs of the best models from the literature. We show that the Transformer generalizes well to other tasks by applying it successfully to English constituency parsing both with large and limited training data.

Reading guide

Problem and motivation

Dominant sequence transduction architectures relied on recurrent neural networks (such as LSTM or GRU) or convolutional neural networks arranged in encoder-decoder structures. Recurrent networks inherently compute hidden representations sequentially ( $h_{t}$ conditioned on $h_{t - 1}$ and input $t$ ), preventing parallelization across time steps within training examples. Furthermore, relating distant positions in recurrent networks requires $O (n)$ sequential operations, while convolutional models like ConvS2S and ByteNet require $O (n)$ or $O (lo g_{k} (n))$ operations to connect arbitrary positions. The authors aimed to build an architecture that eschews recurrence and convolutions entirely, enabling full parallelization across sequence lengths while reducing the maximum path length between any two tokens to $O (1)$ .

Key ideas

Pure Attention Architecture: Replace recurrent connections entirely with attention mechanisms that map dependencies directly across all positions in a sequence.
Scaled Dot-Product Attention: Scale multiplicative dot products by $\frac{1}{d _{k}}$ to prevent dot-product values from growing too large in high dimensions, which would push softmax functions into regions with vanishingly small gradients.
Multi-Head Attention: Linearly project queries, keys, and values into multiple lower-dimensional representation subspaces in parallel, allowing the network to simultaneously attend to information from different positions and representation spaces.
Positional Encoding: Because attention operations are permutation-invariant, fixed sinusoidal functions of varying frequencies are added to input and output embeddings to encode token ordering.

Method

The Transformer uses an encoder-decoder framework with $N = 6$ identical stacked layers in each.

Encoder and Decoder

Each encoder layer has two sub-layers: a multi-head self-attention mechanism and a position-wise fully connected feed-forward network (FFN). Residual connections and layer normalization are applied to each sub-layer: $LayerNorm (x + Sublayer (x))$ , operating over dimension $d_{model} = 512$ .

Each decoder layer contains three sub-layers: masked multi-head self-attention (which masks out future positions to preserve auto-regressive decoding), multi-head cross-attention over the encoder stack's outputs, and a position-wise FFN.

Mathematical Formulation

Scaled dot-product attention on packed matrices $Q, K, V$ is computed as: $Attention (Q, K, V) = softmax (\frac{Q K ^{T}}{d _{k}}) V$

Multi-head attention combines $h$ parallel projections: $MultiHead (Q, K, V) = Concat (head_{1}, \dots, head_{h}) W^{O}$ $where head_{i} = Attention (Q W_{i}^{Q}, K W_{i}^{K}, V W_{i}^{V})$ For the base model, $h = 8$ and $d_{k} = d_{v} = d_{model} / h = 64$ .

The position-wise feed-forward sub-layer is applied identically to each position: $FFN (x) = max (0, x W_{1} + b_{1}) W_{2} + b_{2}$ where the inner-layer dimension is $d_{ff} = 2048$ .

Positional encodings are defined by sine and cosine functions: $P E_{(p os, 2 i)} = sin (\frac{p os}{1000 0 ^{2 i / d_{model}}})$ $P E_{(p os, 2 i + 1)} = cos (\frac{p os}{1000 0 ^{2 i / d_{model}}})$

Results

WMT 2014 English-to-German: The Transformer (big) achieved a state-of-the-art score of 28.4 BLEU, outperforming prior models and ensembles by over 2.0 BLEU. The base model achieved 27.3 BLEU using $3.3 \times 1 0^{18}$ FLOPs, a fraction of the cost of previous architectures.
WMT 2014 English-to-French: The big model achieved 41.8 BLEU (single model), setting a new state of the art after 3.5 days of training on 8 NVIDIA P100 GPUs ( $2.3 \times 1 0^{19}$ FLOPs).
English Constituency Parsing: Evaluated on WSJ Section 23 without task-specific tuning, a 4-layer Transformer achieved 91.3 $F_{1}$ in the discriminative setting and 92.7 $F_{1}$ in the semi-supervised setting, outperforming RNN baselines.

Limitations and open questions

Standard self-attention has a per-layer computational complexity of $O (n^{2} \cdot d)$ , which becomes computationally expensive for very long sequences $n$ .
Generation in the decoder remains auto-regressive (generating one token at a time sequentially), which bounds inference latency.
The paper notes the need to investigate local or restricted attention neighborhoods ( $r$ ) for handling massive input and output contexts (e.g., images, audio, video) efficiently.

Key terms

Transformer: A sequence transduction model architecture based entirely on self-attention mechanisms without recurrence or convolutions.
Scaled Dot-Product Attention: An attention function where dot products of queries and keys are divided by the square root of the key dimension $d_{k}$ before computing a softmax over values.
Multi-Head Attention: An attention mechanism that projects queries, keys, and values into multiple representation subspaces in parallel before concatenating and projecting the outputs.
Self-Attention: An intra-sequence attention mechanism that relates different positions of a single sequence to compute a representation of that sequence.
Positional Encoding: Fixed or learned vector representations added to token embeddings to inject token sequence order information into the permutation-invariant attention layers.
Residual Dropout: Dropout applied to the output of each sub-layer before adding it to the sub-layer input and applying layer normalization, as well as to token embeddings.
Label Smoothing: A regularization technique during training with value $ϵ_{l s} = 0.1$ that penalizes overconfident predictions, improving BLEU scores despite hurting perplexity.
Constituency Parsing: A syntactic NLP task requiring structurally constrained and long target outputs, used to demonstrate the generalizability of the Transformer architecture.

Questions

Why does Scaled Dot-Product Attention divide by $\sqrt{d_k}$?

For large values of $d_{k}$ , the dot products grow large in magnitude (under assumption of zero mean and unit variance, the dot product has variance $d_{k}$ ), pushing the softmax function into regions with extremely small gradients. Scaling by $\frac{1}{d _{k}}$ counteracts this effect.

How does the Transformer retain token order without recurrence?

It adds sinusoidal positional encodings directly to the input and output embeddings at the bottom of the encoder and decoder stacks, assigning a unique frequency-based pattern to each token position.

What prevents decoder positions from attending to future tokens during training?

The decoder's self-attention sub-layer uses masking that sets values in the softmax input corresponding to illegal future connections to $- \infty$ .

What were the training hardware and compute requirements?

The base model trained on a single machine with 8 NVIDIA P100 GPUs for 100,000 steps (12 hours at 0.4 seconds/step). The big model trained for 300,000 steps (3.5 days at 1.0 second/step).

The overview on this page was written by AI from the paper and may contain mistakes. Read the paper for the details.

Read with translation and AI