Attention mechanism

What it is

An attention mechanism determines the importance of each element of a sequence relative to the others, producing weights over the sequence^[1]. A recurrent network normally carries information forward one step at a time; attention lets each token access any part of the sentence directly instead of relying only on the previous hidden state^[1]. One consequence is worth stating plainly: without positional encoding, attention is permutation invariant, so the order of the tokens has to be supplied some other way^[1].

How it works

Attention is usually described with three roles: a query, which is what the model is currently looking for; a set of keys, which is what each position offers; and a set of values, which is the content that actually gets passed on. A scoring function compares the query with each key, the scores are normalised into weights, and the output is a weighted mixture of the values.

Two scoring functions are standard. Additive attention scores a query and key through a tanh non-linearity and a learned weight vector; multiplicative, or dot-product, attention uses a scaled dot product^[1]. The scaled dot-product form used in modern models is the softmax of Q times K transposed over the square root of the key dimension, multiplied by V^[1].

A small example makes this concrete. In machine translation, while predicting a particular target word, the model can soft-search for the parts of the source sentence that are relevant to that word, rather than having to carve the source into a hard segment first^[2]. Every source position receives a weight, and the weights are recomputed at each output step.

Where it came from

Bahdanau, Cho and Bengio introduced additive attention in 2014, in work aimed at removing the fixed-length vector bottleneck in sequence-to-sequence translation^[1]. Their paper, first submitted in September 2014, conjectures that the fixed-length vector used by encoder-decoder translation models is a bottleneck on performance^[2], and extends that architecture by allowing the model to automatically soft-search for parts of a source sentence relevant to predicting a target word, without forming those parts as an explicit hard segment^[2]. With this approach the authors achieved translation performance comparable to the existing state-of-the-art phrase-based system on English-to-French translation^[2].

Luong introduced dot-product, or multiplicative, attention for sequence-to-sequence machine translation in 2015^[1]. Self-attention, where queries, keys and values all come from the same sequence, grew out of differentiable neural computers, neural Turing machines and intra-attention^[1]. Vaswani and colleagues combined these lines of work into the transformer in 2017, in "Attention Is All You Need"^[1].

Where it is used today

Self-attention sits at the centre of the transformer, and the scaled dot-product formulation is the form used in modern models^[1]. Because every position can read every other position in a single step, the mechanism replaces the step-by-step handover of a recurrent network with a direct lookup over the whole sequence^[1].

Limits

The attention matrix grows with the square of the number of tokens, so long inputs need large amounts of GPU memory^[1]. FlashAttention reduces memory by processing the computation in blocks, and KV caching and FlexAttention are further optimisations^[1]. There is also a limit on interpretation rather than computation: attention scores are not necessarily equivalent to explanations of model behaviour, so a high weight is not by itself a justification of an output^[1].