Multi-head attention

Multi-head attention computes attention multiple times in parallel, each with different learned linear projections of the queries, keys, and values. The outputs of these heads are concatenated and linearly projected to produce the final result. This allows the model to jointly attend to information from different representation subspaces at different positions, capturing diverse relationships such as syntactic, semantic, and coreference patterns.

Formally, given input matrices Q, K, V, each head i computes attention as softmax(QW_i^Q (KW_i^K)^T / sqrt(d_k)) VW_i^V, where W_i are learned weight matrices. The outputs are concatenated and multiplied by a final weight matrix W^O. The scaling factor sqrt(d_k) prevents the dot products from growing too large, stabilizing gradients.

Multi-head attention is a key component of the Transformer architecture, used in both the encoder and decoder. In the encoder, it allows each token to attend to all other tokens in the input sequence, enabling the model to capture long-range dependencies. In the decoder, masked multi-head attention prevents attending to future tokens, preserving autoregressive generation.

The Transformer's success in natural language processing, computer vision, and other domains is largely attributed to multi-head attention's ability to parallelize computation and model complex interactions. It replaced recurrent and convolutional layers in many state-of-the-art models, leading to the development of BERT, GPT, and vision transformers.

Several variants of multi-head attention have been proposed to improve efficiency or performance. Multi-query attention and grouped-query attention share keys and values across heads, reducing memory bandwidth and speeding up inference. Sparse attention patterns, such as those in Longformer and BigBird, reduce computational complexity from O(n^2) to O(n log n) or O(n).

Other works have explored head pruning, showing that many heads are redundant and can be removed without significant performance loss. Some studies have analyzed the interpretability of heads, finding that certain heads specialize in syntactic or positional relationships. These insights have led to more efficient and interpretable attention mechanisms.

Multi-head attention was inspired by earlier work on multi-head neural networks and attention mechanisms in machine translation, but the specific formulation in the Transformer was novel. The authors of the original paper, Ashish Vaswani and colleagues at Google, initially experimented with a single attention head but found that multiple heads improved performance.

Research has shown that the number of heads can be reduced to as few as one or two with minimal impact on performance in some tasks, while other tasks benefit from many heads. The choice of head dimension and number of heads interacts with model width and depth, and optimal configurations vary across applications. Additionally, multi-head attention has been applied beyond text, including graph neural networks and protein structure prediction.

Attention

A mechanism that allows a model to weigh the importance of different parts of the input when producing an output.

Transformer

A neural network architecture based on self-attention, introduced in 2017, that has become dominant in NLP and other fields.

Query, Key, Value

Vectors used in attention: the query is matched against keys to compute attention weights, which are used to aggregate values.

Head

One of the parallel attention computations in multi-head attention, each with its own learned projections.