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). Some words matter more than others, and which ones matter depends on the word you are trying to understand.
Earlier translation models read the words in order and compressed the sentence into one fixed-size vector. Bahdanau, Cho and Bengio (2014) suspected that this vector was a bottleneck, since those models got worse as sentences got longer, and they added attention on top of the sequence model: 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, so every word can look at every other word at once.
Parts 1 and 2 trained small networks by gradient descent, with backpropagation supplying the gradients. This part builds an attention layer, trains one on the page the same way, and looks at which words its two heads learn to attend to.
Think about looking up a word in a dictionary. You have a query (the word you want), a set of keys (the dictionary entries), and the values (the definitions). In a hard lookup, your query either matches a key exactly or it does not.
Attention is a soft lookup. The query partially matches many keys at once, each match gets a weight, and the output is a blend of all the values mixed by those weights. The better a key matches the query, the more its value contributes.
Figure 1. Hard and soft lookup of "cat". A key's match score is the number of letter pairs it shares with "cat" (counting the word's start and end), and the weights are the softmax of those scores. Click a mode to see the transition, and hover, tap or tab to a key to see its score; all five scores are also listed under the figure. On a narrow screen the values column is squeezed to fit.
Each word enters as a vector, a list of numbers. To compute attention, the model turns each word's vector into three new ones. The query represents what a word is looking for, the key what it offers to be found by, and the value the information it contributes when attended to. Each comes from multiplying the input vector by a learned weight matrix: WQ, WK or WV.
In this series' colors, the words and the vectors computed from them are data, the weight matrices are the model, the attention weights and each word's output are the model's prediction, and the words the readout has to name are the target. Queries, keys and values all count as data here, so the figures tell them apart by label and position rather than by color.
Every word plays all three roles. "Cat" has a query (what is cat looking for?), a key (what makes cat findable?), and a value (what does cat pass along?).
The remaining figures come from a small model that your browser trains when the page loads. It has one attention layer with two heads, and its queries, keys and values have 6 numbers each. Each word's input vector marks which word it is and where it sits. The training 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 always sits at the same position. From each word's attention output, a linear readout has to name two other words, its targets: the word it attaches to (a determiner or adjective names its noun, the verb its subject, a preposition its noun, and each noun the verb) and the word just before it. No attention pattern is written by hand: WQ, WK, WV 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 first head of the trained model, one row per word of "The cat sat on the mat" (blue entries are positive, gray negative, and each matrix has its own scale). Click a word to select its query. The cyan bars below the matrices show the softmax of that query's dot product with every key, divided by √6.
The raw attention scores are the dot products of each query with each key, one score per pair of words. If the dk entries of a query and a key were independent with mean 0 and variance 1, their dot product would have variance dk, so its typical size grows like √dk. Large scores make softmax so peaked that one word gets nearly all the weight, and there its gradients are tiny, which slows learning. Vaswani and colleagues (2017), who introduced the transformer, give that gradient problem as their suspected reason for dividing the scores by √dk. Softmax then normalizes each row to sum to 1, and those weights mix the value vectors into each word's output.
Figure 3. The first head's attention computation on "The cat sat on the mat", in three stages. Step through with the buttons or by scrolling, and click a row to select a query. Raw and scaled scores run from gray (negative) to cyan (positive); the weights run from pale to dark cyan. The slider replays training: it shows the same computation with the weights saved every 20 steps, from the random start to the end.
Q × KT gives a raw score for every pair of words. Larger numbers mean the query and key point in more similar directions.
Divide by √dk. This keeps the scores of a random start from growing with the dimension, so softmax does not begin with near-one-hot rows and tiny gradients. Training can still sharpen rows later when that lowers the loss, as the next step shows.
Softmax normalizes each row to sum to 1, turning scores into attention weights. Each row is now a probability distribution over the sentence, and multiplying it by V gives that word's output.
Drag the training slider to the left end. With random WQ and WK, the weights in each row are spread over several words and follow no pattern. As training proceeds, most rows concentrate on one word, and in this head it is usually the word that row has to name: "sat" ends on "cat", "on" and the second "the" end on "mat". Gradient descent pushes the weights this way because the readout can only name "cat" from the output at "sat" if the value vector of "cat" is mixed into it.
A single head gives each word one set of weights, so it tends to fetch one kind of neighbor at a time. The toy task asks every word two questions, which is why the model has two heads, each with its own WQ, WK and WV. Both heads feed both readouts, and nothing assigns either head to either question. Full-size transformers have more heads, and their outputs are concatenated and projected back to the model's dimension.
Figure 4. Both heads of the trained model on the same sentence. The text under the grids counts, for each head, how many rows put their largest weight on the attached word and how many on the previous word, and reports the model's accuracy on 300 sentences it never saw in training. Green outlines mark each row's targets: solid for the attached word, dashed for the previous word. On a narrow screen head 2 sits below head 1. Retraining runs the same 400 steps from a different random start, which also draws different training sentences.
In the first run, head 1 mostly fetches the attached word and head 2 mostly fetches the previous word. The first "The" is an exception: head 1 looks at the word itself, which is that position's answer to the previous-word question, and head 2 fetches "cat". Retraining shows how much of the split depends on the random start. The first retrain swaps the roles, and later ones vary. In some runs the heads swap roles; in others a head spreads its rows across both questions, or looks at words that answer neither, and the readout still recovers both answers from the two heads together.
This model was told what to predict and nothing about where to look, so its attention patterns are whichever ones made the predictions work. Real transformers are trained the same way on a different objective, usually predicting the next word or a masked-out word, with many layers of many heads. When Clark and colleagues (2019) inspected the heads of BERT, they found some that attend to the previous or next token, and some that line up with syntax, such as heads where direct objects attend to their verbs and determiners to their nouns. Many other heads have no such clean reading.
Parts 1 to 3 trained models to match targets. Part 4 turns to principal component analysis, which has no target and looks for structure in the data alone.