Attention mechanism

The attention mechanism is the part of an AI language model that works out which words in a text matter for each other. In "the customer returned the jacket because it was too small", attention helps the model link "it" to the jacket. It is the core of the transformer, the design behind ChatGPT and Gemini, and it explains why very long documents take so much computing power.

Where does it come from?

Attention is older than the transformer. In 2014 Bahdanau, Cho and Bengio added it to a translation model. Until then such a model squeezed the whole source sentence into one fixed-length vector, and its quality dropped quickly as sentences got longer. With attention, the model searches the source sentence, for every word it writes, for the parts that matter to that word.

In 2017 Vaswani et al. built the transformer "based solely on attention mechanisms". Because it no longer reads a text word by word, training runs largely in parallel: their translation model set a new best score after twelve hours of training on eight GPUs. GPT-4 and Gemini are built on this design.

How does it work?

Each token gets three vectors, computed with learned weights: a query, a key and a value. The score between two tokens is the dot product of the query of one and the key of the other, divided by √d_k, the square root of the size of the key vectors. Without that scaling, large scores push softmax into a range where it barely learns. Softmax turns the scores into weights that add up to 1, and each token's new representation is the weighted sum of the values:

Attention(Q, K, V) = softmax(QKᵀ / √d_k) V

In self-attention, queries, keys and values all come from the same text. Multi-head attention runs several of these calculations side by side, each with its own weights, so the model can attend to different kinds of information at once. The original transformer used 8 heads; the largest GPT-3 has 96 heads in each of its 96 layers. In a model that generates text, a mask stops each token from looking at the tokens that come after it.

Where does it go wrong?

Every token is compared with every other token, so the work per layer grows with the square of the text length. Twice as many tokens means four times as many pairs, and 100,000 tokens means 10 billion query-key pairs per head in each layer. Sparse attention (Child et al., 2019) brings this down to O(n√n) by comparing each token with a selection. FlashAttention (Dao et al., 2022) keeps exact attention and cuts the memory traffic on the GPU instead.

A long context window is not used evenly. Liu et al. (2023) found that models often answer best when the relevant information sits at the beginning or end of the input, and much worse when it sits in the middle, even models built for long contexts.

Attention weights are also a weak explanation of a model's answer. Jain and Wallace (2019) found them frequently uncorrelated with other measures of which input words mattered, and very different attention patterns can lead to the same prediction.

Example

The transformer paper shows two heads in layer 5 of 6 reading "The Law will never be perfect, but its application should be just". From "its", both heads point almost entirely at "Law" and "application": the model links the pronoun to the nouns it belongs to, which the authors call anaphora resolution. Clark et al. (2019) found similar specialists in BERT: heads that attend to the direct objects of verbs, the determiners of nouns and coreferent mentions "with remarkably high accuracy".