Attention lets each word in a sentence draw on the other words most relevant to it. This part builds the mechanism and trains a small two-head layer in your browser.
In the sentence "The cat sat on the mat," understanding "sat" requires knowing who sat (the cat) and where (the mat). Which words matter depends on the word you are trying to understand.
Early neural translation models read a sentence in order and compressed it into one fixed-size vector. Bahdanau, Cho and Bengio (2014) suspected this vector was a bottleneck, since such models got worse on longer sentences, and added attention: at each output word the model takes a weighted average over all the input words, with more weight on the relevant ones. Transformers later dropped the sequential reading and kept only attention.
A dictionary lookup has a query (the word you want), keys (the entries) and values (the definitions). In a hard lookup the query matches one key exactly or nothing. Attention is a soft lookup: the query partly matches every key, and the output is a blend of all the values, weighted by how well each key matched.
Figure 1. Hard and soft lookup of "cat". A key's score is the number of letter pairs it shares with "cat", counting the word's start and end; the soft weights are the softmax of those scores. Hover, tap or tab to a key for its score.
Each word enters as a vector, a list of numbers, and the model multiplies it by three learned matrices, WQ, WK and WV. The results are the word's query (what it is looking for), its key (what it offers to be found by) and its value (what it passes along).
Most of the remaining figures come from a small model that your browser trains when the page loads: one attention layer with two heads, whose queries, keys and values have 6 numbers each. Each word's input marks which word it is and where it sits. Sentences come from a toy grammar, like "the old dog slept under a box", with an optional adjective in each half so that no word has a fixed position. From each word's output, a linear readout has to name two targets: the word it attaches to (for instance, a determiner names its noun and a verb its subject) and the word just before it. No attention pattern is written by hand: the matrices and the readout start random and are trained by backpropagation on the readout's cross-entropy loss for 400 steps.
Figure 2. The queries, keys and values of the trained model's first head, one row per word (blue positive, gray negative, each matrix on its own scale). The cyan bars are the softmax of the selected query's dot products with every key, divided by √6.
The raw score for a pair of words is the dot product of one's query with the other's key. If the dk entries of a query and a key were independent with mean 0 and variance 1, the dot product would have variance dk. Vaswani and colleagues (2017), who introduced the transformer, suspected that such large scores push softmax into regions with tiny gradients, and divided them by √dk. Figure 3 runs their thought experiment.
Figure 3. Simulated: random queries and keys with independent standard normal entries, 400 queries per dimension d. The panels average the largest softmax weight in a row and the size of the softmax's gradient (its Jacobian's norm), gray for raw scores, cyan for scores divided by √d.
With 8 keys at d = 64, the size of each head in the transformer paper, raw scores already put most of a row's weight on one word before any training, and the softmax passes back well under half the gradient. Divided by √d, both curves stay flat.
Figure 4. The first head's attention on "The cat sat on the mat" in three stages; click a row to select a query. Scores run from gray (negative) to cyan (positive), weights from pale to dark cyan. The slider replays training with the weights saved every 20 steps.
Q × KT gives a raw score for every pair of words.
Divide by √dk, so a random start does not begin with near-one-hot rows. Training can still sharpen rows when that lowers the loss.
Each row becomes weights that sum to 1. Multiplying a row by V gives that word's output.
At the slider's left end, with random WQ and WK each row spreads its weight over several words. As training proceeds most rows concentrate on one word, in this head usually the word that row has to name: "sat" ends on "cat", "on" and the second "the" on "mat". The readout can only name "cat" from the output at "sat" if the value of "cat" is mixed into it, so gradient descent pushes the weights there.
One head gives each word one set of weights, so it tends to fetch one kind of neighbor. The toy asks two questions, so the model has two heads, each with its own WQ, WK and WV. Both feed both readouts; nothing assigns a head to a question.
Figure 5. Both heads of the trained model. Green outlines mark each row's targets, solid for the attached word and dashed for the previous word. The text below counts rows whose largest weight lands on each and gives accuracy on 300 unseen sentences. Retraining draws new random weights and sentences.
In the first run head 1 mostly fetches the attached word and head 2 the previous word. The first "The" is the exception: head 1 looks at the word itself, which is that position's previous-word answer. The first retrain swaps the roles; later ones vary, some spreading a head across both questions while the readout still recovers both answers.
Attention itself does not know where a word sits. Every pair of words is scored the same way, so shuffling the words only moves each word's output with it. This model tells the two "the"s apart only because each input marks its position, which is why transformers add a position signal to every input. Figure 6 trains a second model without one.
Figure 6. Each word's output (2 heads × 6 numbers) and the two words the readout names for it, red where wrong in the original order. Without positions both "the"s get the same column and answers, so one must be wrong, and shuffling moves each column with its word.
The cost can be counted exactly. Over every sentence the grammar can make, weighted by its probability, no order-blind model can name the attached word more than 11/14 of the time (78.6%), or the previous word more than 17/24 (70.8%): repeated words, and determiners swapped between the two halves, look identical to it. Trained from 20 different random starts, the position-free model reached the first bound every time and came within 0.7 points of the second.
This model was told what to predict, not where to look, so its attention patterns are whichever ones made the predictions work. Real transformers learn the same way from a different objective, usually predicting the next or a masked-out word, with many layers of many heads. When Clark and colleagues (2019) inspected BERT's heads, they found some that attend to the previous or next token and some that line up with syntax, such as direct objects attending to their verbs and determiners to their nouns. Many others have no such clean reading.
Part 4 turns to principal component analysis, which has no target and looks for structure in the data alone.