The transformer block piece by piece

Ask an engineer where the compute goes in a transformer and the answer is usually attention. Ask where the parameters go and the answer is usually attention. Both are wrong for most models at most sequence lengths, and the correction takes about ten lines of multiplication.

The attention mechanism worked through in Attention is a lookup you can learn is one stage of six. This article assembles the other five, counts what each holds and what each costs, and identifies the sequence length at which the usual story becomes true.

The six stages

Work at the base configuration Vaswani and colleagues specify: a model width of 512, 8 attention heads with a key and value width of 64 each, and an inner width of 2,048 in the position-wise network. The block runs

1. Layer normalisation
2. Multi-head self-attention
3. Add the block's input back to the result
4. Layer normalisation
5. A position-wise multilayer perceptron
6. Add that stage's input back to the result

The pattern is the same twice: normalise, do something, add the input back. Everything a transformer does is that pair, repeated as many times as the budget allows.

The part that holds the parameters

Attention needs four square matrices: one each to produce queries, keys and values, and one to combine the heads afterwards. At a width of 512, each is 512 by 512, which is 262,144 numbers. Four of them is 1,048,576 parameters, or 1.05 million.

The position-wise network is two matrices. Vaswani and colleagues give it as a linear layer, a rectifier, and a second linear layer, with an inner width of 2,048 against the model width of 512. So one matrix is 512 by 2,048 and the other is 2,048 by 512, each holding 1,048,576 numbers, for 2,097,152 parameters, or 2.10 million.

The perceptron holds twice what attention holds. It is applied independently to every position, sharing its weights across positions the way a convolution shares weights across an image, so it is not a small appendix to attention. It is the larger half of the block.

The two normalisation layers hold a scale and a shift per dimension, which is 1,024 parameters each: about 0.03 percent of the block, and unavoidable.

Total, ignoring biases: about 3.15 million parameters per block. A twelve-block model is around 37.8 million before the embedding table, which is usually larger than any single block.

What is the perceptron for. Attention produces a weighted average of values, and an average is a linear operation, so a stack of attention layers with nothing between them collapses into one, exactly as the linear layers in A neuron is a weighted sum and a decision collapsed. The perceptron is the non-linearity of the whole architecture. Attention decides what each position gets to read; the perceptron decides what to do with it. Modern implementations often replace the rectifier in it with the GELU that Hendrycks and Gimpel define as the input times the standard Gaussian cumulative distribution function, on the grounds set out in Why activation functions matter.

Where the compute goes

Split the cost per token into two parts.

The part that does not care about sequence length is every matrix multiplication: the four attention projections at 4 times 512 squared, and the perceptron at 2 times 512 times 2,048. That is 1,048,576 plus 2,097,152, or 3,145,728 multiply-adds per token, and it equals the parameter count, as it always does for a matrix multiplication applied once per token.

The part that does care is the interaction itself: scoring one query against every key, and then mixing every value. Each of those is one pass over the sequence at full width, so it costs 2 times the sequence length times 512, which is 1,024 multiply-adds per token, per token in the sequence.

Set them equal. 1,024 times the sequence length equals 3,145,728 when the sequence length is 3,072.

Below about three thousand tokens the block is dominated by the matrix multiplications, which grow linearly with the number of tokens. Above it, the quadratic term takes over, and by 32,000 tokens attention is roughly ten times everything else. The received wisdom about attention being the expensive part is a statement about long contexts, and it is simply false for a classifier over sentences.

Memory tells the same story from the other side. Dao, Fu, Ermon, Rudra and Re point out that the time and memory complexity of self-attention are quadratic in sequence length, and their remedy is not a different mathematical operation but a different memory access pattern: tiling to reduce the number of reads and writes between GPU high bandwidth memory and on-chip SRAM. They report a 15 percent end-to-end speedup on BERT-large at length 512, 3 times on GPT-2 at length 1,000, and 2.4 times on long-range arena at 1,000 to 4,000. The arithmetic did not change; the traffic did.

For generation there is a third cost. Serving one token at a time requires keeping the keys and values of every previous token. At a width of 512, twelve blocks, a 4,096 token context and two bytes per number, that cache is about 100 megabytes for a single sequence, and it scales linearly with both context and batch size. It is frequently the binding constraint on how many users one accelerator can serve.

The residual connections

Stage three adds the block's input to attention's output. Stage six does the same for the perceptron. Written out, the block computes input plus attention-of-input, rather than attention-of-input.

The training consequence follows from the arithmetic in Backpropagation worked by hand. Differentiating a sum gives a term for each branch, so the gradient arriving at a block splits: one copy goes through attention, and one copy passes straight through the addition, unmultiplied. Stack ninety-six blocks and there is still a path from the loss to the first block whose factor is exactly one. Without that path, the product of ninety-six factors would behave the way the products in Sequences, recurrence and its limits behave, and the model would be untrainable.

There is also a plain reading. The residual lets a block change its input a little rather than replace it, so a block that has nothing useful to add can learn to add nothing without breaking what came before.

The normalisation, and the one place order matters

Ba, Kiros and Hinton describe layer normalisation as computing the statistics from all of the summed inputs to the neurons in a layer on a single training case, rather than from the distribution of a neuron's input over a mini-batch of cases. That single-case property is why it suits sequences: it performs exactly the same computation at training and test time, with no dependence on what else is in the batch, and it applies directly to recurrent models by computing statistics separately at each step.

The placement is not a detail. The original design puts normalisation after each addition, which is called Post-LN. Putting it before each sub-layer, inside the residual branch, is called Pre-LN.

Xiong and colleagues proved the difference rather than asserting it. Using mean field theory they show that at initialisation, for the original Post-LN transformer, the expected gradients of the parameters near the output layer are large, which makes training with a large learning rate unstable, and that the warm-up stage is practically helpful for avoiding this problem. For Pre-LN the gradients at initialisation are well behaved, and they show that Pre-LN transformers can drop the warm-up stage entirely, reaching comparable results while requiring significantly less training time and hyper-parameter tuning.

That is why nearly every transformer written after 2020 is Pre-LN. It is also why an implementation copied from the original paper sometimes refuses to train while an apparently identical one from a modern library trains without complaint. The difference is the position of two lines.

What to check when a block is not behaving

• Count the parameters by hand before trusting a summary. Four times the width squared for attention, two times the width times the inner width for the perceptron. A count that disagrees with the code means the code is not the architecture in your head, which is the first question asked in Reading a model you did not build.
• If training diverges in the first few hundred steps, check whether the normalisation is Post-LN and whether there is a warm-up. That combination is a documented instability, not bad luck.
• Compute the crossover length for your width before optimising attention. Below it, the perceptron is the thing to make faster.
• Measure the key and value cache before promising a concurrency number. It grows with context and batch, and it is usually what runs out first.
• Change one placement at a time. Normalisation position, residual scaling and activation choice all interact, and a run that mixes three changes tells you nothing about any of them.

Every article so far has described what a network computes. The last one in this series is about the part that actually consumes the calendar: choosing a batch size, a learning rate and a stopping rule so that the run finishes before the quarter does.

References

1. Attention Is All You Need. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin, arXiv, 2017.
2. Layer Normalization. Jimmy Lei Ba, Jamie Ryan Kiros and Geoffrey E. Hinton, arXiv, 2016.
3. On Layer Normalization in the Transformer Architecture. Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang and Tie-Yan Liu, arXiv, 2020.
4. Gaussian Error Linear Units (GELUs). Dan Hendrycks and Kevin Gimpel, arXiv, 2016.
5. FlashAttention, Fast and Memory-Efficient Exact Attention with IO-Awareness. Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra and Christopher Re, arXiv, 2022.