multi-head attention

1 Introduction

Multi-head attention is a mechanism in artificial neural networks that enables a model to focus on different parts of an input sequence simultaneously, capturing a variety of relational patterns. It is a core component of the transformer architecture, as introduced in the 2017 paper “Attention Is All You Need.” By computing multiple attention functions in parallel and combining their results, multi-head attention enhances the model’s ability to represent complex dependencies, such as syntactic and semantic relationships, across sequential data.

1.1 Motivation for Multiple Attention Heads

A single attention function computes a weighted sum of values based on queries and keys. However, a single set of weights may only capture one type of relationship—for example, focusing on nearby words or specific syntactic roles. Multi-head attention addresses this limitation by allowing the model to attend to information from different representation subspaces at different positions. Each head can learn a distinct attention pattern, and the combination of heads provides a richer, more flexible representation of the input.

1.2 Basic Single-Head Attention

Single-head attention takes a query, a set of keys, and a set of values, and computes an output as a weighted combination of the values. The weights are derived from the compatibility between the query and each key. The most common form used in transformers is scaled dot-product attention.

1.2.1 Scaled Dot-Product Attention

Scaled dot-product attention computes attention scores as the dot product between queries and keys, scaled by the inverse square root of the key dimension. The scores are then passed through a softmax function to obtain attention weights, which are multiplied by the values. The scaling factor prevents large dot products from pushing the softmax into regions with extremely small gradients.

2 Mechanism

The multi-head attention mechanism operates by projecting the input into multiple sets of queries, keys, and values, computing attention for each set independently, then concatenating and linearly projecting the results.

2.1 Linear Projections for Query, Key, Value

Given an input sequence represented as a matrix, the mechanism first applies separate learned linear transformations to produce queries (Q), keys (K), and values (V) for each head. The projections are defined by weight matrices: \(W_i^Q\), \(W_i^K\), \(W_i^V\) for head \(i\). Typically, the input dimension is reduced by a factor equal to the number of heads, so that each head works in a lower-dimensional subspace.

2.2 Parallel Attention Computation per Head

For each head \(i\), a scaled dot-product attention function is computed independently using its own \(Q_i\), \(K_i\), and \(V_i\). The output of head \(i\) is:

\[ \text{head}_i = \text{Attention}(Q_i, K_i, V_i) = \text{softmax}\left(\frac{Q_i K_i^\top}{\sqrt{d_k}}\right)V_i \]

where \(d_k\) is the dimension of the keys (and queries). Multiple heads are processed in parallel, often on modern hardware (e.g., GPUs) as batched matrix multiplications.

2.3 Concatenation and Final Output Projection

After computing all heads, their outputs are concatenated along the feature dimension. This concatenated vector is then passed through a final linear projection \(W^O\), which maps it back to the desired output dimension. The result is the multi-head attention output:

\[ \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) W^O \]

2.4 Mathematical Formulation

Let the number of heads be \(h\), and let the input dimensions be such that each head works in a subspace of dimension \(d_k = d_v = d_{\text{model}} / h\). The full computation can be written as:

\[ \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) W^O \]

where

\[ \text{head}_i = \text{Attention}(Q W_i^Q, K W_i^K, V W_i^V) \]

and the attention function is the scaled dot-product form above.

3 Advantages

Multi-head attention provides several key benefits over single-head attention, improving both representational power and training dynamics.

3.1 Capturing Different Relationship Types

Each attention head can specialize in different patterns, such as attending to nearby tokens, long-range dependencies, or specific grammatical structures. For example, in a language model, one head might focus on subject-verb agreement, while another captures semantic similarity.

3.2 Improved Gradient Flow and Training Stability

The parallel, lower-dimensional subspaces allow gradients to flow more easily through the network. The averaging effect across heads can also reduce variance during training, making optimization more stable compared to a single high-dimensional attention function.

3.3 Subspace Representation Learning

By projecting the input into multiple subspaces, the model learns to represent different facets of the data in distinct parts of the feature space. This decomposition enriches the overall representation without requiring a larger single attention matrix, which would be computationally expensive.

4 Variants and Extensions

Several modifications to multi-head attention have been proposed to improve efficiency or adapt to specific task requirements.

4.1 Multi-Query Attention

Multi-query attention reduces the number of key and value heads to one, while keeping multiple query heads. This cuts memory usage for key-value caches in autoregressive decoding, as only a single set of keys and values is stored per layer. It is commonly used in efficient decoder models.

4.2 Grouped-Query Attention

Grouped-query attention divides the query heads into groups, with each group sharing a single key-value head. This provides a middle ground between standard multi-head and multi-query attention, balancing quality and efficiency. It has been adopted in models like LLaMA 2.

4.3 Multi-Head Latent Attention

Multi-head latent attention introduces a compressed latent representation of keys and values, reducing memory and computation while preserving multi-head structure. It achieves this by projecting keys and values into a lower-dimensional latent space before caching or computation, making it suitable for long-context models.

5 Applications

Multi-head attention is a core mechanism in many state-of-the-art neural architectures across diverse domains.

5.1 Natural Language Processing

Multi-head attention is the foundation of transformer-based language models (e.g., GPT, BERT, T5) used for machine translation, text summarization, question answering, and sentiment analysis. It enables capturing long-range dependencies and contextual meanings in text.

5.2 Computer Vision

Vision transformers (ViT) apply multi-head attention to image patches, treating them as a sequence. This has achieved competitive results on image classification, object detection, and segmentation, often rivaling or surpassing convolutional neural networks.

5.3 Speech and Multimodal Models

In speech recognition (e.g., Whisper), multi-head attention aligns audio representations with text outputs. Multimodal models such as CLIP and Flamingo use cross-attention between visual and textual features, enabling tasks like image captioning and visual question answering.

6 Implementation Considerations

Efficient implementation of multi-head attention requires careful management of memory, computation, and hardware utilization.

6.1 Parallelization over Heads

All heads can be computed simultaneously by reshaping the queries, keys, and values into batched tensors (e.g., combining the head dimension with the batch dimension). Modern deep learning frameworks (PyTorch, TensorFlow) and hardware accelerators (GPUs, TPUs) perform these operations efficiently.

6.2 Memory and Compute Trade-offs

Increasing the number of heads improves representational capacity but adds proportional memory and compute costs for the linear projections and attention scores. For very long sequences, the attention matrix grows quadratically, prompting the use of sparse or approximate attention methods.

6.3 Head Dimension Choices

The head dimension is typically set to \(d_{\text{model}} / h\). Common choices for \(h\) range from 8 to 128, depending on model size. Smaller head dimensions reduce per-head capacity but allow more heads; empirical tuning is often required to balance performance and efficiency.

7 Relationship to Other Attention Mechanisms

Multi-head attention builds on foundational attention concepts and adapts them for different architectural roles.

7.1 Self-Attention vs. Cross-Attention

In self-attention, the queries, keys, and values all come from the same sequence, allowing each token to attend to every other token. In cross-attention, queries come from one sequence (e.g., decoder) and keys/values from another (e.g., encoder), enabling information transfer between sequences. Multi-head versions of both are standard.

7.2 Causal Masking in Decoders

During autoregressive decoding, causal masking prevents tokens from attending to future positions. This is applied per attention head by setting the attention scores of future positions to a large negative value before softmax. Multi-head attention in decoders uses the same mask for all heads, ensuring each head respects the causal ordering.