Transformers: Attention Is All You Need

A translation network with no recurrence: each word gathers information from every other word directly, in one parallel step.

The Transformer, built up piece by piece: embeddings, scaled dot-product attention, multiple heads, sinusoidal positions, the encoder-decoder, and the training recipe, with the exact equations and worked numbers.

In 2017 it set translation records on WMT 2014 (28.4 BLEU English-to-German, 41.8 English-to-French) after 3.5 days on eight GPUs, and its layer became the building block of BERT and the GPT models.

The paper solves machine translation: English sentence in, German or French sentence out. Before it, the strong systems were recurrent networks that read a sentence one word at a time, often with an attention mechanism added so the decoder could look back at the source. The Transformer removes the recurrence and keeps only the attention. Every layer lets each word's vector take a weighted average of information from every other word's vector, and all words are processed at the same time.

This page takes the architecture apart in dependency order: why recurrence is slow; how words become vectors; what attention computes, with a worked example you can drag; why the scores are divided by $\sqrt{d_k}$; why there are eight attention heads instead of one; how word order gets back in; the feed-forward, residual and normalization parts of each layer; how the encoder and decoder connect; the training recipe; and what it all cost and scored.

Why recurrence is slow

A recurrent network keeps a hidden state vector $h_t$ and updates it once per word: $h_t$ is a function of the previous state $h_{t-1}$ and the word at position $t$. The state after the last word is the network's summary of the sentence. The sequence-to-sequence translators of 2014 worked this way, with one recurrent network reading the source and a second writing the translation.

The dependency of step $t$ on step $t-1$ has two costs. The first is time. Step $t$ cannot start until step $t-1$ is done, so a 50-word sentence takes 50 sequential steps however many processors the machine has. The paper puts it as “this inherently sequential nature precludes parallelization within training examples.” A GPU gets its speed from doing thousands of independent operations at once; a recurrent layer only offers it the work of one position at a time, across the batch.

The second cost is distance. For word 1 to affect the representation of word 50, its information has to pass through 49 state updates, and the training signal that adjusts how word 1 was used has to travel back through the same 49 updates. The paper's argument, citing earlier work on learning long-range dependencies, is that the shorter these forward and backward paths are, the easier such dependencies are to learn.

Self-attention connects every position to every other position directly. The longest path between two words is one step, and all positions are computed in parallel. Toggle between the two modes below and drag the sentence length: the recurrent path grows with $n$, the attention path stays at one hop.

Connecting every pair of $n$ words costs $n^2$ comparisons per layer, each a dot product of two $d$-number vectors, so about $n^2 \cdot d$ operations. A recurrent layer does only $n$ updates, but each multiplies a $d \times d$ matrix by a vector, so about $n \cdot d^2$. The paper's Table 1 makes the trade exact, and the last section works through it: for sentence-length inputs the $n^2$ term is smaller than the recurrent layer's own cost.

Tokens become vectors

The network works on vectors, so text is first cut into tokens and each token is replaced by a vector. For English-German the paper uses byte-pair encoding, which splits rare words into frequent pieces, with one vocabulary of about 37,000 tokens shared by both languages. A learned table holds one row per token, and each row has $d_{\text{model}} = 512$ numbers in the base model; embedding a token means looking up its row. The rows start random and are trained with everything else.

A sentence of $n$ tokens becomes a matrix of shape $n \times 512$, one row per token. Every layer of the Transformer takes a matrix of that shape and returns a new one of the same shape. The rows keep their positions; what changes is their content. After a few layers the row for “it” in “the animal crossed the road because it got tired” carries information taken from the row for “animal”.

The paper's embedding step has two further details. The table rows are multiplied by $\sqrt{d_{\text{model}}}$, about 22.6, before use. The paper states this multiply and gives no reason; a common reading is that it keeps the embeddings on a scale comparable to the positional encoding added next, whose entries lie between −1 and 1. And the same table is used three times: for the source tokens, for the target tokens, and (transposed) as the final layer that turns the decoder's output vectors back into scores over the vocabulary, following Press and Wolf's weight-tying result.

Attention as a soft lookup

Attention generalizes a dictionary lookup. A dictionary compares a query with every stored key, finds the one that matches, and returns that key's value. Attention compares the query with every key, gives each key a weight according to how well it matches, and returns the weighted average of all the values. With one very good match, the weights put nearly everything on that key and attention returns a vector close to that value; with several partial matches, it returns a blend of their values.

In a Transformer every token plays all three roles. Its vector is multiplied by three learned matrices to give a query (what this position is looking for), a key (what this position offers to be matched against) and a value (what this position contributes if it is selected). The match between a query and a key is their dot product: large and positive when the two vectors point the same way, near zero when they are unrelated, negative when they point apart. A softmax turns a row of scores into weights, by exponentiating each score and dividing by the sum, so the weights are positive and add up to 1 and the largest score gets the largest share.

In the sentence “the animal crossed the road because it got tired,” the query of “it” should match the key of “animal”, so the new vector for “it” is built mostly from the value of “animal”. Pick a query token below to see a weight on every token of the sentence.

Real queries and keys have 64 numbers each, which is hard to draw. With two numbers per vector every step of the computation fits in a picture, so here is a worked example at $d_k = 2$. The query from “it” is $q = (1.6,\ 0.6)$, and four tokens of the sentence have these keys and values (a real head would also score “it” against its own key and the remaining words):

The first score is the dot product of the query with the key of “animal”:

Each score is divided by $\sqrt{d_k} = \sqrt{2}$ (the next section explains why), and the softmax of the four scaled scores is 0.80, 0.11, 0.04, 0.04. The output is the weighted average of the values:

which lies close to the value of “animal”. That output becomes the new vector at the position of “it”. In the figure, drag the query on the left plane and watch three things: the weights under the planes, the output point on the right plane, and the shaded region the output never leaves. A weighted average with positive weights summing to 1 always lies inside the convex hull of the values. Then try the two presets with the same direction: the short query gives weights 0.39, 0.24, 0.19, 0.19, and the long one puts about 0.96 on “animal”. Scaling the query scales every score by the same factor, and larger scores make the softmax sharper.

Stack the queries of all positions as the rows of a matrix $Q$, and the keys and values as the rows of $K$ and $V$. Then every query-key score is one entry of the matrix product $QK^{\top}$, the softmax runs along each row, and multiplying by $V$ forms every weighted average at once. This is the paper's equation (1), which it calls scaled dot-product attention:

For $n$ queries and $m$ keys, $QK^{\top}$ is an $n \times m$ grid of scores, the softmax leaves the shape alone, and the product with the $m \times d_v$ value matrix gives one $d_v$-dimensional output per query. The operation is two matrix multiplications with a softmax between them, which is why it runs fast on GPUs:

# scaled dot-product attention, one head, equation (1)
# Q: [n, d_k]   K: [m, d_k]   V: [m, d_v]
scores  = Q @ K.T / sqrt(d_k)       # [n, m]: every query against every key
weights = softmax(scores, axis=-1)  # each row is positive and sums to 1
out     = weights @ V               # [n, d_v]: weighted average of value rows

When $Q$, $K$ and $V$ are all computed from the same sequence, the operation is called self-attention: the sentence attends to itself. The same equation also runs with queries from one sequence and keys and values from another, which is how the decoder reads the encoder's output in the assembled model.

Dividing by √d_k

Without the $\sqrt{d_k}$, equation (1) is the dot-product (“multiplicative”) attention that already existed; the paper says the two are identical except for that factor. The other common choice was additive attention from Bahdanau, Cho and Bengio, which scores a query-key pair with a small feed-forward network with one hidden layer. The paper picks the dot product because it runs as a matrix multiplication and is “much faster and more space-efficient in practice.” A comparison the paper cites found that at small $d_k$ the two perform alike, and at larger $d_k$ additive attention beats the unscaled dot product.

The paper's explanation is a suspicion, stated as one: that for large $d_k$ the dot products grow large in magnitude and push the softmax into regions with extremely small gradients. A footnote shows why they grow. Suppose the components of $q$ and $k$ are independent, with mean 0 and variance 1. Each product $q_i k_i$ then has mean 0 and variance 1, and the variances of independent terms add:

The typical size of a score is its standard deviation, $\sqrt{d_k}$. At the base model's $d_k = 64$ that is 8, so raw scores routinely differ by 10 or 20. A softmax over scores that far apart puts almost all the weight on the largest one. The gradient of a softmax weight $w_i$ with respect to its own score is $w_i(1 - w_i)$, which is at most 0.25 (at $w_i = 0.5$) and goes to zero as $w_i$ approaches 0 or 1. A saturated softmax therefore passes almost no gradient back to the queries and keys, and training cannot easily change which key a query picks. Dividing every score by $\sqrt{d_k}$ brings their variance back to 1 at any $d_k$.

At $d_k = 1$ the two panels match, since $\sqrt{1} = 1$. At $d_k = 64$ the unscaled peak weight is 0.997 and the largest gradient is 0.0035, about 70 times smaller than the scaled version's 0.245, and at $d_k = 256$ the unscaled gradient is about $10^{-5}$.

The divisor is the standard deviation $\sqrt{d_k}$ and not the variance $d_k$ because the goal is to bring the spread of the scores to 1. Dividing by $d_k = 64$ would shrink the spread to $8/64 = 0.125$, and a softmax over scores that close together gives nearly equal weights to every key, so attention could not select anything.

Multi-head attention

One attention computation gives each position one set of weights, and so one weighted average. A word often has several relations that matter at once: a verb to its subject and to its object, a pronoun to the noun it replaces and to its own verb. One set of weights has to split its mass between them, and the output mixes their values into one average. The paper names this cost as “reduced effective resolution due to averaging attention-weighted positions” and counters it with several attentions in parallel.

Multi-head attention runs $h = 8$ attention computations side by side. Head $i$ has its own three learned matrices $W_i^Q$, $W_i^K$, $W_i^V$, each of shape $512 \times 64$, which project the 512-dimensional token vectors down to 64-dimensional queries, keys and values. Each head is a separate learned projection of the whole vector, not a slice of 64 of its 512 numbers, so each head can learn its own notion of which tokens match. Pick a query below to see four heads spend their weights four ways, and the single-head average in the last row.

The eight head outputs, 64 numbers each, are concatenated back into 512 numbers per position and multiplied by one more learned $512 \times 512$ matrix $W^O$, which mixes the heads:

The projection matrices have these shapes:

The paper leaves these formulas unnumbered. The base model sets $d_k = d_v = d_{\text{model}}/h$, which is 64. Because each head works at one eighth of the width, eight heads cost about as much as one attention at full width: the projections together still hold $4 \cdot 512^2$, or 1,048,576, weights per attention sublayer, and the heads compute eight $n \times n$ score grids of 64-dimensional dot products instead of one grid of 512-dimensional ones.

# multi-head attention, h heads (base model: h = 8, d_k = d_v = 64)
def multi_head(Xq, Xkv, h):             # Xq: [n, 512]   Xkv: [m, 512]
    heads = []
    for i in range(h):                  # each head has its own projections
        Q = Xq  @ Wq[i]                 # [n, 64]   Wq[i]: [512, 64]
        K = Xkv @ Wk[i]                 # [m, 64]
        V = Xkv @ Wv[i]                 # [m, 64]
        heads.append(attention(Q, K, V))  # [n, 64]
    return concat(heads, axis=-1) @ Wo  # [n, 512] -> [n, 512]

The paper tests the head count at constant compute (Table 3, rows A, measured in BLEU on the English-German development set newstest2013; BLEU scores word overlap with human reference translations, higher is better). One head with $d_k = 512$ scores 24.9; four heads of 128 score 25.5; eight heads of 64 score 25.8; sixteen heads of 32 also score 25.8; thirty-two heads of 16 drop back to 25.4. One head is 0.9 BLEU worse than the best setting, and too many heads hurt as well. Rows B shrink only $d_k$, to 16 and 32, and lose 0.7 and 0.4 BLEU, which the authors read as a sign that deciding compatibility is not easy.

The paper also shows attention maps where individual heads appear to track sentence structure, including two heads in layer 5 that seem to resolve what “its” refers to. Treat head roles as a tendency. Michel, Levy and Neubig later found that a large fraction of heads can be removed from a trained model at test time without a significant loss.

Positional encoding

Equation (1) has no notion of order. Shuffle the input tokens and the rows of $Q$, $K$ and $V$ shuffle with them; each token sees the same set of keys and values, gets the same weights, and produces the same output, now in the shuffled row. “Dog bites man” and “man bites dog” give each word the same vector. A recurrent network gets order from the sequence of its updates; the Transformer has to write position into the input.

The paper adds a positional encoding to each token's embedding: a vector of the same size, 512, that depends only on the position. Instead of learning these vectors, it uses sines and cosines at a range of frequencies:

The variable $pos$ is the token's position and $i$ runs over the 256 dimension pairs. Dimension $2i$ gets a sine and dimension $2i+1$ the cosine of the same angle, whose frequency is $\omega_i = 1/10000^{2i/512}$. Pair 0 has $\omega_0 = 1$ and repeats every $2\pi \approx 6.3$ positions. Each later pair is slower by a constant factor, so the wavelengths form a geometric progression from $2\pi$ toward $10000 \cdot 2\pi$; the last pair, $i = 255$, repeats every 60,611 positions. Reading all 512 values together gives each position a distinct pattern, the same way the hands of clocks running at very different speeds together pin down a moment. Drag the position in the heatmap to read one row.

The paper gives two reasons for sinusoids, both hedged. The first is a hypothesis: that the model can easily learn to attend by relative position, because for any fixed offset $k$, $PE_{pos+k}$ is a linear function of $PE_{pos}$. The mechanism is the angle-addition formula. Each sine-cosine pair of the encoding is a point on a unit circle, and moving forward $k$ positions turns it by the angle $k\,\omega_i$, whatever $pos$ is:

The matrix contains $k$ and not $pos$, so one fixed rotation per pair maps every position to the position $k$ later. A consequence is that the dot product between two encodings depends only on their distance. Each pair contributes $\sin a \sin b + \cos a \cos b$, which equals $\cos(a-b)$, and summing over the pairs gives

which equals 256 at $k = 0$ and falls as the offset grows. In the figure, set the offset and the pair, then drag the start position: the arc moves around the circle, and its angle and the dot product on the right stay fixed.

The second reason is also a “may”: a formula defined at every position might let the model handle sentences longer than any seen in training. The paper tested neither reason directly. It did try learned position embeddings, one trained vector per position, and got “nearly identical results”: 25.7 BLEU on the development set against 25.8 for the sinusoids (Table 3, row E). The rotation idea returned in 2021 in a different form: RoPE rotates the queries and keys themselves by position-dependent angles, so the query-key dot product depends on the offset by construction.

Feed-forward, residual, norm

Each layer has a second sublayer after attention: a position-wise feed-forward network, the same two-layer network applied to every position separately. Attention moves information between positions; the feed-forward network then transforms each position's vector on its own, with no view of the others. It is a linear layer, a ReLU ($\max(0, x)$, which zeroes negative entries), and a second linear layer:

$W_1$ is $512 \times 2048$ and $W_2$ is $2048 \times 512$, so the network widens each vector four times and narrows it back. Its parameter count is

twice the 1,048,576 of the attention sublayer's four projection matrices. Each layer uses different $W_1, W_2$.

Both sublayers are wrapped the same way. A residual connection, from He et al.'s ResNet, adds the sublayer's input to its output, so the sublayer only has to compute a change to its input, and the gradient has a direct additive path through all the layers. Then layer normalization rescales each position's 512 numbers to mean 0 and variance 1 and applies a learned gain and bias, which keeps the scale of the vectors from drifting as they pass through 6 layers. In the paper's notation each sublayer outputs LayerNorm(x + Sublayer(x)). Dropout at rate 0.1 is applied to the sublayer output before the addition.

# one encoder layer as the paper writes it (post-LN)
a = LayerNorm(x + dropout(multi_head(x, x, 8)))   # sublayer 1
z = LayerNorm(a + dropout(ffn(a)))                # sublayer 2
# pre-LN, used by GPT-2 and much later code, moves the norm inside:
#   a = x + dropout(multi_head(LayerNorm(x), LayerNorm(x), 8))

This arrangement, with the norm after the residual addition, is called post-LN. Many later models use pre-LN instead, which normalizes the sublayer's input and computes x + Sublayer(LayerNorm(x)); GPT-2 (2019), for one, moved layer normalization to the input of each sublayer. Xiong et al. (2020) analyzed the two orders. They found that in post-LN the expected gradients of the parameters near the output are large at initialization, so a large learning rate early in training is unstable and the warmup described in the training section helps, and that pre-LN, already proposed by then, reaches comparable results without warmup. None of this is in the 2017 paper, which uses post-LN with a 4000-step warmup and gives no stability reason for either.

Encoder, decoder, mask

The translation model has two stacks of $N = 6$ layers. The encoder reads the source sentence: each of its layers is self-attention followed by the feed-forward network, each wrapped in a residual and a norm. Its output is one 512-dimensional vector per source token. The decoder produces the translation one token at a time, and each of its layers has three sublayers instead of two. Attention appears in three roles:

• Encoder self-attention. Queries, keys and values all come from the previous encoder layer. Every source position can attend to every source position.
• Masked decoder self-attention. Queries, keys and values come from the previous decoder layer, and each target position may attend only to itself and earlier target positions.
• Encoder-decoder attention, also called cross-attention. Queries come from the previous decoder sublayer; keys and values come from the encoder's final output, so each target position can read every token of the source sentence. It is the same equation (1) with a different source for each input, and it plays the role of the attention in Bahdanau-style recurrent translators.

The figure below is the base model as a clickable map, with the shapes for a 9-token source and a 10-token target. Click a block (or step with the buttons) to see what flows into it, the shapes, and its parameter count. It opens on the cross-attention block: its score grid is $10 \times 9$, one row per target position and one column per source token.

Adding up the blocks: an encoder layer holds 3.15 million parameters (one attention sublayer, one feed-forward network, two norms), a decoder layer 4.20 million (two attention sublayers), and the shared 37,000 × 512 table about 18.9 million. Six layers of each plus the table come to about 63 million, close to the 65 million the paper lists for the base model in Table 3; the vocabulary size is only given as “about 37000”.

The decoder's self-attention needs one restriction. At generation time the model writes word $i$ before word $i+1$ exists, so the prediction at position $i$ must not use anything later. During training the whole target sentence is available at once, and the decoder processes all positions in one parallel pass. Without a restriction, position $i$ could attend to the answer it is supposed to predict, learn to copy it, and fail at generation time. The causal mask sets every score from position $i$ to a later position $j > i$ to $-\infty$ before the softmax. Since $e^{-\infty} = 0$, those keys get weight exactly 0 and the allowed weights still sum to 1. The attention grid becomes lower-triangular.

The mask works together with a shift. The decoder's input is the target sequence moved one place to the right, with a start symbol in front, so the input at position $i$ is the true word $i-1$, and the output at position $i$ is trained to predict word $i$. With both, the prediction for position $i$ depends only on the words before it. Toggle the mask below: off is the encoder's full grid, on is the decoder's triangle.

At the top of the decoder, each 512-dimensional output is multiplied by the transposed embedding table to give a score for each of the ~37,000 tokens, and a softmax turns the scores into next-token probabilities. At inference the paper runs beam search with 4 candidate translations and a length penalty $\alpha = 0.6$, generating one token per step, with a maximum output length of the input length plus 50.

The training recipe

English-German training used the WMT 2014 set, about 4.5 million sentence pairs; English-French used 36 million sentences with a 32,000-token word-piece vocabulary. Sentence pairs of similar length were batched together, and each batch held about 25,000 source tokens and 25,000 target tokens. Training ran on one machine with 8 NVIDIA P100 GPUs: the base model took about 0.4 seconds per step and 100,000 steps (12 hours); the big model took 1.0 second per step and 300,000 steps (3.5 days).

The optimizer is Adam with $\beta_1 = 0.9$, $\beta_2 = 0.98$ and $\epsilon = 10^{-9}$, and the learning rate follows equation (3):

Both terms inside the minimum equal $\mathit{warmup}^{-0.5}$ at $\mathit{step} = \mathit{warmup}$. Before that the second term is smaller, so the rate grows linearly; after it the first is smaller, so the rate decays as $1/\sqrt{\mathit{step}}$. The paper uses $\mathit{warmup} = 4000$. For the base model the peak, at step 4000, is

and by step 100,000 the rate is down to $1.4 \times 10^{-4}$. The $d_{\text{model}}^{-0.5}$ factor lowers the whole curve for wider models: the big model's peak is $4.9 \times 10^{-4}$. Drag the warmup below; the dashed curve keeps the paper's 4000 for comparison.

Two regularizers complete the recipe. Dropout at $P_{drop} = 0.1$ is applied to each sublayer output and to the sums of embeddings and positional encodings (the big English-German model used 0.3). Table 3 row D shows how much it matters: with no dropout the base model falls from 25.8 to 24.6 BLEU. Label smoothing with $\epsilon_{ls} = 0.1$ trains against a target that puts 0.9 of the probability on the correct token and spreads the remaining 0.1 over the vocabulary. The paper reports that this hurts perplexity, because the model learns to be less certain, but improves accuracy and BLEU. For evaluation, the base model is the average of the last 5 checkpoints (written every 10 minutes) and the big model the average of the last 20.

Cost and results

The paper's Table 1 compares layer types on three counts: the arithmetic per layer, the number of operations that must run one after another, and the longest path between two positions. In the table, $n$ is the sequence length, $d$ the representation size, $k$ a convolution's kernel width, and $r$ the neighborhood size of a restricted self-attention that only looks at nearby positions.

Self-attention needs $n^2 d$ operations (every pair of positions, one $d$-dimensional dot product each) and a recurrent layer $n d^2$ (a $d \times d$ matrix times a vector at each of $n$ steps). Self-attention is cheaper whenever $n < d$. For a 30-token sentence at $d = 512$, that is 0.46 million multiply-adds against 7.9 million:

The paper notes that sentence lengths in machine translation are usually below $d$. For long inputs the $n^2$ term dominates, in time and in the memory for the $n \times n$ score grids, and much later work (for example FlashAttention, which cuts the memory but not the quadratic arithmetic) targets it.

On WMT 2014 English-to-German, the big model scored 28.4 BLEU, more than 2.0 above the best previously reported result including ensembles (ConvS2S ensemble, 26.36). The base model scored 27.3, which already beats every earlier model and ensemble. On English-to-French, the big model scored 41.8, a new single-model record against 40.56 for the best earlier single model (a mixture-of-experts network); the base model scored 38.1.

The training-cost column is an estimate, not a count of operations: the paper multiplies training time, number of GPUs, and an assumed sustained throughput per GPU (9.5 TFLOPS for a P100). For the base model that is 12 hours × 8 GPUs × 9.5 TFLOPS $\approx 3.3 \times 10^{18}$; for the big model, 3.5 days gives $2.3 \times 10^{19}$. Against the English-French mixture-of-experts model at $1.2 \times 10^{20}$, the big Transformer used about a fifth of the compute, which is the paper's “less than 1/4 the training cost of the previous state-of-the-art model”. On English-German the base model used $3.3 \times 10^{18}$, against $9.6 \times 10^{18}$ for ConvS2S, the cheapest competitor with a listed cost.

The big model changes the sizes and the dropout rate: $d_{\text{model}} = 1024$, $d_{ff} = 4096$, $h = 16$, dropout 0.3 (0.1 for the English-French model), 213 million parameters, and 26.4 development BLEU against the base model's 25.8. To test the architecture outside translation, the authors also trained a 4-layer Transformer with $d_{\text{model}} = 1024$ for English constituency parsing (turning a sentence into its grammatical parse tree). Trained on the 40,000 Wall Street Journal sentences alone it reached 91.3 F1, and with about 17 million extra sentences 92.7 F1, with little task-specific tuning.

The paper describes the Transformer as “the first sequence transduction model based entirely on attention”, and its experiments are translation and parsing with the full encoder-decoder. The later uses came from other papers. BERT (2018) kept only the encoder stack and trained it to fill in masked words. The GPT line, through GPT-3, kept only the decoder with its causal mask and dropped the cross-attention. The Vision Transformer (2020) fed image patches to an encoder, with learned position embeddings in place of the sinusoids. The title is an overstatement on the paper's own terms, since the model still needs embeddings, positional encodings, feed-forward sublayers, residuals and normalization; the paper removed only recurrence and convolution.