1 Background

1.1 Origins of the attention mechanism

The concept of attention in neural networks was first formalized for sequence‑to‑sequence models, notably by Bahdanau et al. (2014) in the context of machine translation. This mechanism allowed a decoder to selectively focus on relevant parts of an input sequence, dynamically weighting the encoder’s hidden states. The basic idea—computing a weighted sum of values based on the compatibility between a query and a set of keys—soon became a cornerstone of modern deep learning.

1.2 Limitations of single‑head attention

A single attention head computes one weighted combination of values for each query, using a single set of learned projections. While effective, this design restricts the model to capturing only one type of relational pattern at a time. In tasks requiring attending to different syntactic, semantic, or positional features simultaneously, a single head may be insufficient, leading to a loss of information diversity.

1.3 Motivation for multiple heads

To overcome the constraints of single‑head attention, the Transformer architecture proposed splitting the attention computation into multiple parallel “heads.” Each head operates on different learned linear projections of the queries, keys, and values, enabling the model to jointly attend to information from distinct representation subspaces. This design allows the network to capture a richer variety of relationships within the same layer.

2 Core Mechanism

2.1 Linear projections to queries, keys, and values

Given an input sequence (often represented as a matrix of token embeddings), multi‑head attention first applies three separate learned linear transformations to produce multiple sets of queries, keys, and values. For a model with *h* heads, each head *i* receives its own projected query (Q_i), key (K_i), and value (V_i) matrices. These projections are obtained by multiplying the input with parameter matrices W_q, W_k, and W_v, each having output dimensions tailored for the per‑head subspace.

2.2 Parallel attention computation

Each head independently computes an attention function on its projected Q_i, K_i, V_i. Because the heads are independent, they can be processed in parallel, making the mechanism efficient on modern hardware. The most common attention function used is scaled dot‑product attention.

2.2.1 Scaled dot‑product attention

For a single head, the attention output is computed as:

\[ \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). The scaling factor \(\sqrt{d_k}\) prevents the dot products from growing too large, which could push the softmax into regions of extremely small gradients. The resulting attention weights are then used to compute a weighted sum of the values.

2.3 Concatenation and output projection

After each head produces its output vector (or sequence of vectors), the results from all *h* heads are concatenated along the feature dimension. This concatenated vector is then passed through a final linear transformation (with matrix W_o) to project it back to the desired output dimension. This step fuses the information from the different subspaces into a unified representation.

3 Mathematical Formulation

3.1 Single‑head attention function

Formally, let \(Q \in \mathbb{R}^{n \times d_k}\), \(K \in \mathbb{R}^{m \times d_k}\), \(V \in \mathbb{R}^{m \times d_v}\). The single‑head attention function is:

\[ \text{Attn}(Q, K, V) = \text{softmax}\left(\frac{Q K^\top}{\sqrt{d_k}}\right) V \]

The output is a matrix of shape \(n \times d_v\).

3.2 Multi‑head attention definition

Multi‑head attention computes:

\[ \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) \, W_o \]

where each head is defined as:

\[ \text{head}_i = \text{Attn}(Q W_{Q_i}, K W_{K_i}, V W_{V_i}) \]

Here \(W_{Q_i} \in \mathbb{R}^{d_{\text{model}} \times d_k}\), \(W_{K_i} \in \mathbb{R}^{d_{\text{model}} \times d_k}\), \(W_{V_i} \in \mathbb{R}^{d_{\text{model}} \times d_v}\), and \(W_o \in \mathbb{R}^{h d_v \times d_{\text{model}}}\).

3.2.1 Number of heads (h) and dimension per head

The number of heads *h* is a hyperparameter, typically chosen to divide the model dimension \(d_{\text{model}}\) evenly. Common values are 8, 12, or 16. The per‑head dimensions are usually \(d_k = d_v = d_{\text{model}} / h\). This ensures that the total computational cost is similar to that of a single full‑dimensional attention.

3.3 Parameter matrices (W_Q, W_K, W_V, W_O)

The parameters of multi‑head attention consist of four sets of learned weight matrices. The three projection matrices for queries, keys, and values are each of size \(d_{\text{model}} \times (h d_k)\) (or equivalently *h* separate matrices). The output projection matrix W_o has dimensions \((h d_v) \times d_{\text{model}}\). These matrices are updated during training via backpropagation.

4 Variants and Extensions

4.1 Multi‑query attention

Multi‑query attention shares the same key and value projections across all heads, while keeping separate query projections. This reduces memory and improves inference speed, especially for autoregressive decoding, making it popular in large language models like PaLM and Falcon.

4.2 Grouped query attention

Grouped query attention strikes a balance between multi‑head and multi‑query attention by dividing heads into groups that share keys and values within each group. It provides a trade‑off between the expressivity of multi‑head attention and the efficiency of multi‑query attention, and is used in models such as LLaMA 2.

4.3 Flash attention (efficient implementation)

Flash attention is not a variant of the attention function itself but an IO‑aware algorithm that computes attention without materializing the full attention matrix. By tiling and exploiting the memory hierarchy of GPUs, it significantly reduces memory usage and speeds up training and inference, enabling longer context windows.

5 Applications

5.1 Transformer models (encoder and decoder)

Multi‑head attention is the core building block of the Transformer, used in both the encoder and decoder. In the encoder, it allows each position to attend to all positions in the input. In the decoder, masked multi‑head attention prevents attending to future tokens, while cross‑attention allows the decoder to attend to the encoder output.

5.2 Natural language processing (machine translation, BERT, GPT)

The mechanism powers nearly all modern NLP models. Machine translation (e.g., Google Neural Machine Translation), pre‑trained encoders like BERT, and autoregressive decoders like GPT all rely heavily on multi‑head attention to model long‑range dependencies and contextual relationships.

5.3 Computer vision (ViT, DETR)

Vision Transformers (ViT) apply multi‑head attention directly to image patches, achieving strong results on image classification. Detection Transformer (DETR) uses the same mechanism for object detection, replacing traditional region‑proposal and anchor‑based methods with an end‑to‑end set prediction approach.

6 Advantages

6.1 Capturing diverse contextual dependencies

By using multiple heads, the model can simultaneously learn different types of attention patterns—for example, one head might focus on syntactic relations while another captures semantic similarity. This diversity enriches the representation.

6.2 Parallelizable and computationally efficient

Multi‑head attention allows parallel computation across heads and across positions, unlike recurrent or convolutional alternatives. This makes it highly suitable for modern hardware (GPUs/TPUs), leading to faster training and inference.

6.3 Improved representation learning

The concatenation and projection of multiple subspace representations produces a richer, more expressive output. This contributes to the superior performance of Transformer‑based models across many tasks.

7 Limitations

7.1 Increased parameter count and memory usage

Compared to single‑head attention, multi‑head attention introduces additional weight matrices (projections plus output projection). This increases the number of parameters and the memory footprint, which can be a constraint for deployment on limited hardware.

7.2 Overhead in inference and training

The need to compute and concatenate multiple heads adds computational overhead. While parallelizable, the overall cost scales linearly with the number of heads, and the memory required for storing intermediate attention maps can be large for long sequences.

7.3 Sensitivity to hyperparameter (number of heads)

Choosing an inappropriate number of heads can degrade performance. Too few heads may reduce diversity, while too many can lead to redundancy or wasteful computation. The optimal value is task‑dependent and often requires empirical tuning.

8 Historical Development

8.1 Introduction in "Attention is All You Need" (2017)

Multi‑head attention was introduced by Vaswani et al. in the seminal paper “Attention is All You Need.” The authors proposed the Transformer architecture, which replaced recurrent layers entirely with multi‑head self‑attention, achieving state‑of‑the‑art results in machine translation with superior training speed.

8.2 Subsequent refinements and variants (2018–present)

Following the original paper, multi‑head attention became ubiquitous. Researchers introduced variants such as relative positional encodings, sparse attention patterns (e.g., Longformer, BigBird), and efficient implementations like FlashAttention. Multi‑query and grouped query attention emerged to reduce memory during decoding, and extensions like cross‑attention in multimodal models continued to expand the mechanism’s applicability.

9 See Also

  • Transformer (deep learning architecture)
  • Self‑attention
  • Scaled dot‑product attention
  • Positional encoding
  • BERT (language model)
  • GPT (model series)

10 References

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., & Polosukhin, I. (2017). Attention is all you need. *Advances in Neural Information Processing Systems*, 30.

Bahdanau, D., Cho, K., & Bengio, Y. (2014). Neural machine translation by jointly learning to align and translate. *arXiv preprint arXiv:1409.0473*.

Shazeer, N. (2019). Fast transformer decoding: One write‑head is all you need. *arXiv preprint arXiv:1911.02150*.

Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and memory‑efficient exact attention with IO‑awareness. *Advances in Neural Information Processing Systems*, 35.