Transformer Explorer

About

This is a complete, working transformer block, small enough to read: 32-dimensional vectors, two attention heads, a 64-unit feed-forward layer and a read-out over 78 words. Every number on screen is computed from the prompt you type, in your browser.

Start at stage 1 and press the right arrow to walk forward. Click any token row to follow it. At the end, press G to generate a word and watch the loop repeat. Then switch off the trained read-out to see the same machinery produce nonsense.

Background: the Transformers lesson for the ideas, the Attention Visualizer for the attention maths in detail, and Generative AI for how the same loop scales to modern chat models.

What is trained here

Honesty first. Of the 13,326 parameters:

• Token embeddings, W_V, W_O and the feed-forward weights are seeded random numbers. They are not trained. Random projections still carry information: different words get different vectors, and the random layers mix them in fixed ways.
• Head 1’s W_Q and W_K are built by hand to attend to the previous token, using the rotation property of sinusoidal encodings. Head 2 has zero W_Q and W_K, so it averages the context.
• The read-out W_U is trained: 300 steps of full-batch Adam on softmax cross-entropy over all 252 next-word examples in the 46-sentence corpus. That is a convex problem, so it trains reliably in about a second.

The result behaves like a small n-gram model with some context: after “the cat sat on the” it favours mat, after “she read a” it favours book. A real language model trains every weight by backpropagation on trillions of tokens, and its attention patterns are learned rather than built.

Positional encoding

Attention on its own is order-blind: shuffle the tokens and each one computes the same weighted average. Position has to be injected.

Each pair of dimensions is a clock hand turning at its own speed. Vaswani et al. (2017) chose this because PE(p + k) is a fixed linear function (a rotation) of PE(p), so a head can learn relative offsets. Head 1 here uses exactly that property.

Most current open models use rotary position embeddings (RoPE, Su et al. 2021), which rotate the query and key vectors themselves by position-dependent angles, so the attention score depends directly on relative distance.

At real scale

The shape of the computation is the same: embed, then repeat attention and feed-forward blocks with residual connections and normalisation, then read out a distribution over the vocabulary. Scale changes what the weights can learn, not the recipe.

References

• Vaswani, A. et al. (2017). Attention is all you need. NeurIPS.
• Ba, J. L., Kiros, J. R., Hinton, G. E. (2016). Layer normalization.
• Xiong, R. et al. (2020). On layer normalization in the transformer architecture. ICML.
• Brown, T. et al. (2020). Language models are few-shot learners. NeurIPS.
• Su, J. et al. (2021). RoFormer: enhanced transformer with rotary position embedding.
• Elhage, N. et al. (2021). A mathematical framework for transformer circuits. Anthropic.