Multi-head attention is a core component of the Transformer (architecture) architecture in Deep learning. It extends the basic attention mechanism by computing several attention operations in parallel, each with its own learned projections of queries, keys, and values. This allows the model to capture different types of relationships and patterns within the input sequence, such as syntactic dependencies, coreference, or positional cues, across multiple representation subspaces. The outputs of these parallel attention heads are then concatenated and linearly projected to produce the final result, enabling the model to jointly attend to information from different representation subspaces at different positions.
The concept was introduced in the 2017 paper "Attention Is All You Need" by researchers at Google, including Jakob Uszkoreit, Lukasz Kaiser, and Niki Parmar. The paper proposed the Transformer, which relies entirely on attention mechanisms and dispenses with recurrence and convolutions, leading to significant improvements in training speed and performance on sequence-to-sequence tasks. Multi-head attention has since become a foundational element of Large language models such as GPT, BERT, and T5, and is widely used in Natural language processing, computer vision, and speech recognition.
Mechanism
In the attention mechanism, each input token is represented as a vector, often called a token embedding. For a sequence of tokens, the model computes three matrices: queries (Q), keys (K), and values (V). The attention score between a query and a key is typically computed as a dot product, scaled by the square root of the key dimension. These scores are normalized using a softmax function to produce attention weights, which are then used to compute a weighted sum of the values.
Multi-head attention runs this process multiple times in parallel, each with different learned linear projections of the original queries, keys, and values. For example, with 8 heads, the model computes 8 separate attention outputs, each focusing on different aspects of the sequence. The outputs are concatenated and passed through a final linear layer. This design allows the model to attend to information from different representation subspaces, which is particularly useful for capturing diverse linguistic patterns.
Unlike "hard" weights, which are set during training, attention weights are "soft" and computed during the forward pass, meaning they change with each input. This allows the model to dynamically focus on relevant parts of the input, as illustrated by the alignment example in machine translation: when translating "I love you" to French, the model assigns high attention weights to "I" when generating "je", to "you" when generating "t'", and to "love" when generating "aime".
History
The attention mechanism was developed to address weaknesses in Recurrent neural networks (RNNs), which tend to favor information at the end of a sentence and attenuate the significance of earlier words. Early attention mechanisms, such as Bahdanau-style additive attention (2014) and Luong-style multiplicative attention (2015), were grafted onto encoder-decoder RNN architectures for machine translation. These allowed the decoder to access all encoder hidden states directly, rather than relying solely on the final hidden state.
The major breakthrough came with self-attention, where each element in the input sequence attends to all others, enabling the model to capture global dependencies. This idea was central to the Transformer, which replaced recurrence with attention mechanisms entirely. The Transformer's parallelizable architecture allowed for significant speedups in training, leading to its adoption in models like BERT, T5, and generative pre-trained transformers (GPT).
Variants
Several variants of attention have been developed, many implementing soft weights. These include:
- Fast weight programmers (1992): A "slow" neural network outputs the "fast" weights of another neural network through outer products. This was later renamed "linearized self-attention".
- Bahdanau-style attention (additive attention): Uses a feed-forward network to compute alignment scores.
- Luong-style attention (multiplicative attention): Uses dot products between decoder and encoder states.
- Positional attention and factorized positional attention: Incorporate positional information into attention computations.
- Spatial and channel attention for convolutional neural networks: Operate on spatial dimensions or feature channels.
These variants recombine encoder-side inputs to redistribute effects to each target output, often using a correlation-style matrix of dot products for re-weighting coefficients.
Optimizations
The size of the attention matrix is proportional to the square of the number of input tokens, which can be memory-intensive for long sequences. Flash attention is an implementation that reduces memory needs and increases efficiency without sacrificing accuracy. It partitions the attention computation into smaller blocks that fit into the GPU's faster on-chip memory, reducing the need to store large intermediate matrices.
FlexAttention, developed by Meta, is an attention kernel that allows users to modify attention scores prior to softmax and dynamically chooses the optimal attention algorithm, providing flexibility for custom attention patterns.
Applications
Attention is widely used in natural language processing, computer vision, and speech recognition. In NLP, it improves context understanding in tasks like question answering and summarization. In vision, visual attention helps models focus on relevant image regions, enhancing object detection and image captioning.
In Computer vision, vision transformers (ViT) apply multi-head attention to image patches. Attention maps, which visualize attention scores as heat maps, have become a routine way to inspect the decision-making process of ViT models. Deeper layers tend to show more semantically meaningful visualizations. Attention rollout is a recursive algorithm that combines attention scores across all layers by computing the dot product of successive attention maps.
Because vision transformers are typically trained in a self-supervised manner, attention maps are generally not class-sensitive. When a classification head is attached, class-discriminative attention maps (CDAM) combine attention maps and gradients with respect to the class token to provide class-specific explanations.
Impact
Multi-head attention has been instrumental in the rise of Generative AI and Large language models. It enables models to process long sequences efficiently and capture complex dependencies, making it a key component in systems developed by organizations like OpenAI, Anthropic, and Google DeepMind. The architecture has also influenced hardware design, with companies like NVIDIA and Cerebras optimizing chips for attention-based computations.
Research continues on improving attention efficiency and interpretability. Techniques like sparse attention, linear attention, and kernel-based methods aim to reduce the quadratic complexity of attention, while interpretability tools help understand what models learn. Multi-head attention remains a vibrant area of study in Machine learning and Artificial intelligence.