101 puffs

Chapter 4 · The transformerPuff 18 of 23

Stacking it all up

You’ve met the parts: tokens, embeddings, attention with many heads, position codes. A chatbot’s model is these parts arranged into one block, and that block copied dozens of times. GPT-2’s smallest model has 12 copies; GPT-3 has 96. Follow one word through, then find out what makes a model big.

Guess first

GPT-3 is over a thousand times bigger than GPT-2’s smallest model. What makes it bigger?

Play with it

Fig. 1 · The transformer One word’s journey, and the size of the stack

Press “Next” to follow one word through the model, one stage at a time.

How many knobs?

Estimate = 12 × blocks × width² + vocabulary × width. It leaves out the small stuff: biases, the normalising steps and the position codes. Published sizes are from the GPT-2 and GPT-3 papers and OpenAI’s GPT-2 release.

How it works

Tap a step to see it in the figure.

Check yourself

Why the “add” shortcuts around attention and the feed-forward network?

Go deeper: the mathThe block’s equations and where the knob count comes from

One block

Write the word’s running list as x (it has d numbers). The original Transformer’s block, as drawn in the figure:

h = LayerNorm(x + Attention(x))
y = LayerNorm(h + FFN(h))

The “+ x” and “+ h” are the shortcuts, also called residual connections. GPT-2 moved the normalising to the start of each part, h = x + Attention(LayerNorm(x)), and GPT-3 kept that order. The feed-forward network widens each word’s list four times, applies a bend (ReLU in the original, GELU in GPT-2), and narrows it back:

FFN(h) = W2 · act(W1 h + b1) + b2

Layer normalisation rescales the d numbers so their average is 0 and their spread is 1, then applies two learned knobs per number:

LayerNorm(x) = γ · x − μσ + β

Counting the knobs

  • Attention: four tables of d × d (queries, keys, values, and the mixer WO), so 4d2, however many heads.
  • Feed-forward: W1 is 4d × d and W2 is d × 4d, so 8d2.
  • So each block has about 12d2, and L blocks have 12Ld2.
  • The word table: V × d for a vocabulary of V tokens. GPT-2 reuses it in the final layer, so it is counted once.

knobs ≈ 12 L d2 + V d

With your settings: ….

The estimate leaves out each block’s biases and normalising knobs (13d per block), the final normalisation (2d) and the position table. Adding them for GPT-2’s smallest model, with its 1,024 positions, gives … knobs, which is OpenAI’s 124 million. (The GPT-2 paper first listed its smallest and largest models as 117 and 1,542 million; OpenAI later corrected the counts to 124 and 1,558 million.)

Why stacking works

Shortcuts were made popular by ResNet (He et al., 2015), which used them to train image networks over 100 layers deep. Because each block adds to the running list, the blame for an error flows straight back to early blocks through the additions. As for what the blocks store, Geva and colleagues (2021) found that feed-forward layers behave like stores of patterns and the words they tend to predict.

Sources

  1. A. Vaswani et al., “Attention Is All You Need”, 2017. The block, figure 1.
  2. A. Radford et al., “Language Models are Unsupervised Multitask Learners”, OpenAI, 2019 (GPT-2 sizes, table 2), with the corrected counts in OpenAI’s gpt-2 code release.
  3. T. Brown et al., “Language Models are Few-Shot Learners”, 2020. GPT-3 sizes, table 2.1.
  4. K. He et al., “Deep Residual Learning for Image Recognition”, 2015; J. L. Ba, J. R. Kiros and G. E. Hinton, “Layer Normalization”, 2016; M. Geva et al., “Transformer Feed-Forward Layers Are Key-Value Memories”, 2021.