Attention is all you need, what the paper actually says

The paper that the current generation of language models descends from is a machine translation paper. It proposes an architecture, tests it on two translation benchmarks, reports better scores than the previous best systems, and stops. It does not mention language models trained on web text. It does not mention few-shot prompting, chat, code generation, or scale. Read today, the striking thing is how little of what followed is anticipated in it.

That gap is worth understanding, because the reason the architecture spread is not the reason the paper gives for it being good.

What it reports

Vaswani and colleagues describe a model "based solely on attention mechanisms, dispensing with recurrence and convolutions entirely". On the WMT 2014 English-to-German task it reaches 28.4 BLEU, "improving over the existing best results, including ensembles by over 2 BLEU". On English-to-French a single model reaches 41.8 BLEU, which was the best single-model score at the time, after training for 3.5 days on eight GPUs.

BLEU compares a machine translation against human reference translations by counting overlapping word sequences. Beating an ensemble by two BLEU points was a real result in 2017. It is also, now, the least interesting sentence in the paper.

The architecture it describes is an encoder and a decoder, six layers each. The base model uses a width of 512, an inner feed-forward width of 2,048, and 8 attention heads, so each head works in 512 divided by 8, which is 64 dimensions. The larger model widens to 1,024 and uses 16 heads. The mechanism inside is worked through arithmetic by arithmetic in Attention is a lookup you can learn, and the full block is assembled in The transformer block piece by piece, so neither is repeated here.

The table in the middle

The argument that mattered is not in the results section. It is in a comparison of three ways to connect positions in a sequence, measured on three axes: how much work one layer costs, how many of those steps must happen one after another, and how far information has to travel to get from one position to another.

Write n for the sequence length and d for the model width. A recurrent layer costs on the order of n times d squared, and it needs n sequential steps, because step 12 cannot start until step 11 has produced its state. A self-attention layer costs on the order of n squared times d, and needs one sequential step, because every position is computed from every other position in the same pass.

Put numbers in. Take the base model width of 512 and a sentence of 100 tokens.

• Recurrent: 100 times 512 squared. 512 squared is 262,144, so the layer costs about 26,214,400 units of work.
• Self-attention: 100 squared times 512. That is 10,000 times 512, or about 5,120,000 units.

Self-attention is roughly five times cheaper here, and the paper states the condition under which that holds: self-attention layers are faster than recurrent layers when the sequence length is smaller than the representation width. At a width of 512, the crossover is a sequence of 512 tokens.

Past it, attention loses on raw arithmetic. Take 4,096 tokens at the same width.

• Recurrent: 4,096 times 262,144, which is about 1.07 billion units.
• Self-attention: 4,096 squared times 512. 4,096 squared is 16,777,216, so about 8.59 billion units.

Eight times more work. On the arithmetic alone, the recurrent layer wins at long sequence lengths, and it keeps winning by a wider margin the longer the sequence gets.

The column that decided it

The architecture won anyway, and the reason is the middle axis: sequential steps.

The recurrent layer needs n of them. The attention layer needs one. That difference does not show up in a count of multiplications, because it is not about how much work there is. It is about whether the work can be done at the same time.

Six recurrent layers over 100 tokens is 600 steps that must happen in order. Six attention layers over the same 100 tokens is 6 steps that must happen in order, each of which internally is a large matrix multiplication that a GPU is built to do all at once. The 8.59 billion units of work in the long-sequence case are a lot of arithmetic, and arithmetic is the thing hardware has been getting cheaper at for decades. Waiting is not.

That is the trade the paper made: it converted a latency problem, which hardware was not fixing, into a throughput and memory problem, which hardware was fixing every year. A 4,096 token sequence needs 4,096 squared attention scores per head per layer, which is 16,777,216 numbers to compute and hold. That is expensive and it is parallel, and parallel and expensive is a much better position to be in than cheap and serial.

The third axis says the same thing from another direction. In a recurrent network, information from position 1 reaches position 100 by passing through 99 intermediate states, and anything not preserved along the way is lost, which is the failure described in Sequences, recurrence and its limits. Under self-attention, position 100 reads position 1 directly, in one hop, at any distance.

Where the parameters actually go

The name has done some damage here, because it suggests the attention is the model. Count the weights in one base layer and it is not.

Attention needs four projection matrices: one each for queries, keys and values, and one to mix the heads back together. Each is the model width by the model width, so at a width of 512 that is 512 times 512, which is 262,144 weights apiece, and 4 times 262,144 is 1,048,576 for the whole attention sublayer.

The feed-forward network next to it is two matrices, one going from the width out to the inner width and one coming back. At a width of 512 and an inner width of 2,048 that is 512 times 2,048 in each direction, which is 1,048,576 each, or 2,097,152 together.

The part of the layer with no attention in it holds twice as many weights as the part the paper is named after. Both numbers are worth carrying: the attention sublayer is where positions talk to each other, and the feed-forward sublayer is where most of the capacity sits and where the non-linearity lives. A stack of attention with nothing beside it collapses, for the reason set out in A neuron is a weighted sum and a decision, because a weighted average of values is a linear operation and a stack of linear operations is one linear operation.

What had to be added back

Removing recurrence removed something useful along with the bottleneck. A recurrent network knows the order of its input because it consumes it in order. Attention does not consume anything in order. It scores a set.

The calculation in the linked article gives the same weights for "the river bank flooded" and "flooded bank river the". Order has to be supplied separately, which the paper does by adding a positional encoding to each token's embedding before the first layer. That is not a detail of implementation. It is the price of the thing that made the architecture fast, and every later model pays it in one form or another.

What the paper does not say

Almost everything the architecture is now used for.

The model in the paper is an encoder-decoder, which suits translation: a complete input sequence goes in, a different sequence comes out. Nearly nothing since has kept that shape for language. GPT and its descendants use the decoder alone and predict the next token from the tokens before it, as described in Language models predict the next token. BERT and its descendants use the encoder alone and read in both directions at once, which the next article in this series works through.

The paper does not propose either split. It does not suggest training on unlabelled text at scale, does not discuss what happens as parameters grow, and makes no claim that attention is the right primitive for anything beyond sequence transduction. The abstract's claim is bounded: this architecture, on these two translation tasks, is better and faster to train than what came before.

Reading it now, the sentence that aged best is the one about training time, and the reason is not that the authors were writing about the future. They were reporting that their model reached a state-of-the-art score in 3.5 days on eight GPUs, at a time when comparable systems took considerably longer. They meant it as evidence that the architecture was practical. What it turned out to be was evidence that the architecture was scalable, which is a different property, and one that only becomes visible once somebody tries to spend a hundred times more compute and finds that nothing structural stops them.

Why this is worth reading carefully

There is a habit of describing this paper as having invented attention. It did not. Bahdanau, Cho and Bengio introduced attention for translation three years earlier, to get past the fixed-length vector bottleneck in the encoder-decoder models of Sutskever, Vinyals and Le. Attention already existed and already worked. The 2017 contribution was to remove everything else and show that what remained was sufficient, and faster.

The lesson generalises past this one paper. The load-bearing claim in a piece of research is often not the headline number. Here the headline was a benchmark score on a task the field has largely moved past, and the durable contribution was a complexity argument about sequential dependency that occupies a few lines in the middle. A result gets cited for what it measured. It gets used for what it made possible, and those are rarely the same sentence.

The next article in this series takes the encoder half of this architecture and asks what changes when a model is allowed to read in both directions at once, and what that costs it.

References

1. Attention Is All You Need. Ashish Vaswani and colleagues, arXiv, 2017.
2. Neural Machine Translation by Jointly Learning to Align and Translate. Dzmitry Bahdanau, Kyunghyun Cho and Yoshua Bengio, arXiv, 2014.
3. Sequence to Sequence Learning with Neural Networks. Ilya Sutskever, Oriol Vinyals and Quoc V. Le, arXiv, 2014.