Attention, Step by Step

How a translation model decides which source words to look at while it writes each output word. The running example is the start of our sentence: Unha tía miña, que todo o solucionaba con ditos e refráns, dicía algo así como «un é escravo do que di e dono do que cala».

1The encoder: reading a sentence into one vector

A translation model starts by reading the Galician sentence. Each word has an embedding, a row in a table like p on the skip-gram page; here it has 2 dimensions, so every word is a point. The encoder reads the words in order and keeps a summary: one more vector, updated after every word. In this simplified encoder, each new word moves the summary halfway toward that word's point.

a word's embedding the encoder's summary

The update, with this step's numbers

How much of each word is in the summary now

The word vectors are made up for this example. Real encoders, recurrent neural networks, learn their update rule instead of always taking half and half, and they can learn to keep some things longer. But their summary is still one vector of a fixed size, whatever the length of the sentence.

2The decoder: writing one word at a time

The decoder writes the English sentence the way the n-gram page generated text: at every step it gives each word it knows a probability of coming next, writes the most likely one, and repeats. Two things are new. Its probabilities come from vectors, as on the skip-gram page. And it starts from the encoder's summary, so what it writes depends on the Galician sentence.

What it can write: its vocabulary

Where it starts

The loop, one step per English word (click a step)

This step, with its numbers

3Look back at every word

A translation model has an encoder, which reads the Galician words, and a decoder, which writes the English words one at a time. Each box below is the encoder's summary after one more word, as in section 1. Without attention, the decoder gets only the last box, so it uses the words in the proportions that box happens to hold: mostly the last ones. With attention, it can look back at every box, and it chooses new proportions for every word it writes.

Without attention: one vector

What the decoder uses for this word

With attention

What the decoder uses for this word

Weights for tía miña are the hand-set ones from the worked example; for the whole sentence they are the real model's (section 5, averaged over its heads). The English words are that model's translation. The top bar comes from the simplified encoder of section 1; the bottom bar from the attention weights. Before attention, a long sentence had to fit into the last box alone, and translations of long sentences got much worse.

4Query, key, score, softmax, value

Each source word offers a key: what it has. The word being written asks with a query: what it is looking for. The score is the dot product of query and key; softmax turns the scores into weights that sum to 1; the output is the weighted average of the source words' values. Drag the orange query.

Set the query to

The arithmetic

Attention weights

One query per output word: the attention grid

Values and the output

The two values are made up for this example, with dimensions we can name. The output always lies on the line between them, pulled toward the word with the larger weight.

Keys and queries as on the worked-example slide: ktía = (1, 0), kmiña = (0, 1), qaunt = (3, 0), qmy = (0, 3). A real model learns them, with hundreds of dimensions.

5A real model on our sentence

The same arithmetic inside a trained Galician–English model, Helsinki-NLP/opus-mt-gl-en. Its decoder has 6 layers with 8 attention heads each, and every head looks back at the source in its own way. Click an output word, or play the translation, to see where the model looked.

Source (Galician), shaded by attention weight

Output (English), as the model wrote it

The whole grid: one row per output word, one column per source word

Words are split into pieces the model knows, such as solucion·aba. ‹/s› marks the end of the sentence; models often park attention there when nothing in the source is needed.

6Self-attention, and many heads at once

The encoder uses the same mechanism inside one sentence: every Galician word attends to every Galician word, itself included. Each head has its own queries and keys, so each head learns its own pattern. The heads' outputs are then combined into one vector per word.

Pick a head, then click a word.

7Who decides what a head does?

Nobody. The people who build a model choose how many layers and heads it has, how big they are, and what it is trained to predict. They do not assign jobs to heads. Every head starts as random numbers, and training nudges all of them at once in whatever direction lowers the error. A head ends up with a job because, from its own random start, doing that job was the quickest way down. Researchers find the jobs afterwards, by looking, as we do here.

Below, a tiny model with 2 layers of 2 heads learns the classic sentence The animal didn't cross the street because it was tired: it has to say what “it” refers to. The adjective decides: tired means the animal, wide means the street. It trains on 96 such sentences. The three runs differ only in their random start.

Where “it” looks, in each head, for two sentences

Loss (lower is better) and the current step

The job each head ended up with, over all 96 sentences

Layer 1 lets every word look at every word; layer 2 looks from “it” and names the referent. Keys and queries have 2 dimensions. Adam, learning rate 0.03, 500 steps. “Switched off” means that head's output is set to zero after training. “Other words” are the, didn't, cross, because, was: models often use such positions to store or park information.