Attention and self-attention

Where attention came from

Attention is a way for a neural network to focus on the parts of its input that matter most at each step. The mechanism was introduced for neural machine translation in September 2014 by Dzmitry Bahdanau, Kyunghyun Cho and Yoshua Bengio.^[1] Before it, translation models squeezed a whole source sentence into one fixed-length vector. Attention lets the decoder look back at every source word and weigh them differently for each word it produces.^[1]

From attention to self-attention

Self-attention applies the same idea inside one sequence. Instead of a decoder looking at a separate source sentence, every token computes how much to attend to every other token in the same text. The 2017 Transformer paper made self-attention the only sequence-mixing operation in its model.^[2]

How it works step by step

Each token is projected into three vectors: a query, a key and a value. To update one token, its query is compared with the key of every token, including its own, using a dot product. These scores are scaled by the square root of the key dimension and passed through a softmax, which turns them into weights that add up to one. The new representation is the weighted sum of all the value vectors.^[2]

A useful way to think about the three roles is this. The query says what a token is looking for, the key says what each token offers for matching, and the value is the content that actually gets passed along. Consider again the word "it" in a sentence about a cat and a mat. Its query matches the key of "cat" closely, so the dot product is large, the softmax gives "cat" most of the weight, and the output is built mostly from the value of "cat". The same operation runs in parallel for every token in the sequence.^[2]

Why the name

The word self refers to the fact that the sequence attends to itself rather than to a separate input. The same mechanism, with a decoder querying an encoder's output, is what the earlier translation models used. Self-attention keeps the mechanism and removes the second input.

Why it is useful

Self-attention connects any two positions in a single step. This helps when related words are far apart in a text, because the information does not have to pass through every word in between. The computation also parallelizes well on GPUs.^[2]

In practice, the output for each token is a blend of information from the whole sequence, and the blend is recomputed for every position. The same layer therefore applies one set of rules to short and long texts alike.

Limits

Every token is compared with every other token, so the cost grows with the square of the sequence length. That is why long context windows are expensive to run, and why many efficient-attention variants have been proposed.^[2]