The film
About 6:19, hosted by Tallow. It asks you two questions and waits for your answer before it shows you. The narrated cut runs 5:08, voiced by Tallow in a synthesized voice made from the author's own voice; the author approved these cuts for publishing. Captions are on by default. The live version under it draws the same film in your browser, waits for your answer at each question, and runs a little longer; the full words are also below.
Nothing loads until you press play.
The lab: Tallow's lamp room
Train the same kind of model the film uses: 2 layers, 2 heads each, 20,864 numbers, starting from random. It learns from sentences about the coast, such as "rerun found a kettle . pylon put a cone on the", where the last word can only be guessed by looking back. Pick a head to light the sentence with, change the test sentence, or turn the lamp off and train again.
The film's own runs of 2,000 steps filled 99% of the copy holes with attention on, 17% when each word could see only itself, and 64% with even light and no choosing. Your run is a new one, so its numbers will differ a little.
The model trains on your computer, in this page. Nothing is sent anywhere.
A circuit, switched off
A second, short film: a different small model, trained on a computer's main processor to continue random strings with one repeated stretch. One head in its first layer moves each character one place forward; a head in its second layer finds the earlier copy and reads what came next. Switch the heads off and the model loses the repeat. That is a causal test, which a picture of attention weights alone is not.
About 73 seconds, with captions. The narration is a synthesized voice made from the author's own voice, and the author listened to this cut before it was published. Every number on screen comes from the model's data file in raw-native.
Everything the film says
Attention
Attention. How a language model decides what matters.
I keep the lamp at Night School. Tonight the lamp is the lesson.
A hole at the end
Here is a sentence with a hole at the end.
Rerun found a kettle. Pylon put a cone on the...
What word goes in the hole? (the film pauses here)
You said kettle. And you did not read the sentence again. You glanced back, past the full stop, to one word near the start, and took it.
The old way: a whisper down the line
A language model like this one does one job. It reads the words so far and guesses the next one.
The first neural translation models squeezed the whole sentence into one fixed-length list of numbers, and worked from that.
Think of a message whispered down a line of people. By the end of a long line, the message has worn thin.
In 2014, Bahdanau, Cho and Bengio let a translation model look at every word of the sentence while it wrote each new word, and weigh how much each one counts. They called it an alignment model. Today we call the idea attention.
Question, label, cargo
Here is how the looking works. Every word does three things.
It asks a question. It holds up a label. And it carries some cargo.
The word at the end shines its question back along the sentence. Each earlier word answers with its label.
A good match gets a high score. The scores are turned into shares that add up to one hundred percent.
Then the word takes a blend of the cargo of every word so far, itself included, in those shares. That blend is what it carries forward.
Inside the lamp
Let's go inside the lamp.
Every word here is a list of 32 numbers. One dot is one number.
A match is a dot product: multiply the two lists number by number, and add it all up.
Each head works on a shorter list of its own, 16 numbers here, and the score is divided by the square root of that length. Without it, long lists give huge scores, the shares go all or nothing, and the model learns slowly.
This is a real model, trained on sentences from our coast. It has 20,864 numbers in total, and nothing in it was set by hand.
Which earlier word will get most of the light? (the film pauses here)
On this sentence, the word at the end gives more than 99 percent of its light to kettle.
One beam, several heads
One beam asks one kind of question. A transformer runs several at once, and calls each one a head.
The base transformer in the 2017 paper had eight heads in every attention layer. The larger one had sixteen.
Ours has 2 heads in each of 2 layers. Nobody told them what to look for.
After training, a head in the second layer puts 71 percent of its light on the word to be copied.
In small attention-only models, researchers at Anthropic found pairs of heads, in different layers, that work together to copy. The first writes into each word a note about the word before it. The second uses that note to find an earlier copy of the current word, and reads the word that followed it. They called them induction heads.
Turn the lamp off
Here is the test that matters. We trained the same model again with the lamp turned off, so each word could see only itself.
With attention, it filled the copy holes correctly 99 percent of the time.
With each word seeing only itself, 17 percent.
Then one more run. Every word got an even share of light over everything before it, with no choosing. That run reached 64 percent.
Knowing which words came before helps. Choosing where to look is what finishes the job.
What the light costs
There's a bill. Every word compares itself with every word before it.
In each head, a thousand words means about half a million comparisons. Double the text, and the attention work goes up about four times.
That's why people build special machinery to do this fast. We take it apart in the cannery, in a later lesson.
What the light does not tell you
One warning before you go.
The light shows where a head is reading from. It does not, on its own, tell you why the model answered what it did.
In 2019, researchers showed that in some text classifiers, very different attention patterns can give almost the same answer. Others replied that it depends on what you mean by explain, and proposed tests. Treat the light as a clue, and check it.
Your turn
The lamp room is open below this film. Train the model yourself, and watch the heads learn where to look.
My voice is synthesized, made from the author's voice.
Sources
Every claim in the film, with the paper it comes from and what that paper does not prove. The numbers about our model come from its weights file, which you can load and check: model.json and facts.json.
Vaswani and others 2017, Attention Is All You Need. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., Polosukhin, I. (2017). Attention Is All You Need. Advances in Neural Information Processing Systems 30. https://arxiv.org/abs/1706.03762
Used for: Scaled dot-product attention computes softmax(Q K^T / sqrt(d_k)) V: queries are compared with keys by dot products, the scores are divided by the square root of the key dimension, turned into weights by a softmax, and used to average the values (section 3.2.1). The authors suspect that for large d_k the dot products grow large, pushing the softmax into regions with extremely small gradients, and scale by 1/sqrt(d_k) to counteract it (section 3.2.1). Multi-head attention runs several attention functions in parallel on learned projections; the base model uses h = 8 heads and the big model h = 16 (section 3.2.2, Table 3). The decoder's self-attention masks later positions, so a position can attend only to itself and earlier positions (section 3.2.3). Self-attention's cost per layer grows with the square of the sequence length (Table 1, O(n^2 d)).
Does not prove: Anything about how attention behaves inside a particular trained model; the paper defines the mechanism.
Bahdanau, Cho and Bengio 2015, neural machine translation by jointly learning to align and translate. Bahdanau, D., Cho, K., Bengio, Y. (2015). Neural Machine Translation by Jointly Learning to Align and Translate. ICLR 2015 (oral). arXiv:1409.0473, first posted 1 September 2014. https://arxiv.org/abs/1409.0473
Used for: Neural encoder-decoder translation models encoded a source sentence into a fixed-length vector, which the authors conjecture is a bottleneck for long sentences; the baseline's score drops sharply with sentence length (sections 1 and 5.1). Their model searches the source positions and weighs them for each output word; they call it an alignment model, and section 3.1 says it implements a mechanism of attention in the decoder.
Does not prove: That early words in particular fade; the paper shows degradation with sentence length.
Olsson and others 2022, In-context Learning and Induction Heads. Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., and others (2022). In-context Learning and Induction Heads. Transformer Circuits Thread, Anthropic. arXiv:2209.11895. https://arxiv.org/abs/2209.11895
Used for: An induction head attends to the token that followed an earlier occurrence of the current token and copies it ([A][B] ... [A] -> [B]). It works with a previous-token head in an earlier layer (K-composition), first described for two-layer attention-only models in Elhage and others 2021.
Does not prove: That the heads in our tiny model are induction heads; that is measured on our own model, separately.
Jain and Wallace 2019, Attention is not Explanation. Jain, S., Wallace, B. C. (2019). Attention is not Explanation. NAACL-HLT 2019. arXiv:1902.10186. https://arxiv.org/abs/1902.10186
Used for: In the recurrent text-classification and question-answering models studied, attention weights often correlate weakly with gradient-based importance, and very different attention distributions can give equivalent predictions; simpler averaging encoders behaved better.
Does not prove: That attention never explains a prediction, or anything about transformer language models in particular.
Wiegreffe and Pinter 2019, Attention is not not Explanation. Wiegreffe, S., Pinter, Y. (2019). Attention is not not Explanation. EMNLP-IJCNLP 2019. arXiv:1908.04626. https://arxiv.org/abs/1908.04626
Used for: Argues that whether attention is an explanation depends on the definition of explanation, and that the earlier experiments do not settle it; proposes tests for when attention can serve as one.
Does not prove: That attention is an explanation; it argues the earlier experiments do not settle the question.
Dao and others 2022, FlashAttention. Dao, T., Fu, D. Y., Ermon, S., Rudra, A., Re, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135. https://arxiv.org/abs/2205.14135
Used for: Standard attention's time and memory grow quadratically with sequence length; FlashAttention computes exact attention in tiles to cut reads and writes to GPU memory. FlashAttention keeps the exact result and the quadratic arithmetic; it cuts memory traffic and makes memory use linear in sequence length.
Does not prove: Anything about this lesson's model, which runs on a CPU.
Night School's tiny transformer, trained for this film. media/school/attention/model.json, written by tools/school/train_film_weights.mjs from system/school/tinyformer.mjs at a fixed seed; the same code the lab runs in the browser. https://github.com/HarperZ9/HarperZ9.github.io/tree/main/media/school/attention
Used for: Every number the film states about 'our model' (size, heads, the probe's attention share, copy accuracy with and without attention) is read from model.json by tools/school/fill_film.mjs.
Does not prove: Anything about large language models' attention; it is a two-layer model on a toy language.
Counting comparisons. A causal attention head compares each of n positions with itself and every earlier one: n (n + 1) / 2 scores, 500,500 for n = 1,000, or n (n - 1) / 2 = 499,500 counting only earlier words. Doubling n multiplies either by about four.
Used for: In each head, a thousand words means about half a million comparisons; doubling the text multiplies the attention work by about four (3.998).
Does not prove: Wall-clock time, which also depends on memory traffic (see dao2022).
Elhage and others 2021, A Mathematical Framework for Transformer Circuits. Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B., and others (2021). A Mathematical Framework for Transformer Circuits. Transformer Circuits Thread, Anthropic. https://transformer-circuits.pub/2021/framework/index.html
Used for: Induction heads were first described in two-layer attention-only models, working with a previous-token head through K-composition. An attention pattern can be very informative and also misleading about what a head does.
Does not prove: Anything about models with MLP layers, such as ours, without measuring them.
What this lesson leaves out
The model here is tiny: two layers on a made-up language of 46 words, so it can train in your browser in about a minute. Large models have dozens of layers and many heads, and their heads do jobs that are much harder to name. Where a head puts its light tells you where it reads from; on its own it does not tell you why the model answered as it did.
Next along the coast: the transformer block that attention sits in (L11), and attention rebuilt as one fast trip through the GPU (K23). Both are planned, not made yet.