Multi-Head AttentionMHA

Multi-Head Attention (MHA) is a core component of the Transformer architecture. Instead of relying on a single attention

Multi-Head Attention (MHA) is a core component of the Transformer architecture. Instead of relying on a single attention operation, it executes multiple attention mechanisms in parallel, enabling the model to attend to different parts of an input sequence across several representation subspaces. While standard self-attention captures global dependencies between all positions, a single perspective may fail to encode rich local patterns or long-range relationships simultaneously. MHA addresses this limitation by projecting queries, keys, and values into multiple lower-dimensional subspaces, where each head can specialize in a distinct type of relationship.

Mathematical Formulation

Given an input sequence $X \in \mathbb{R}^{N \times d_{model}}$, where $N$ denotes the sequence length and $d_{model}$ the model dimension, MHA first projects $X$ into $h$ sets of queries, keys, and values:

$$
Q_i = X W_i^Q, \quad K_i = X W_i^K, \quad V_i = X W_i^V
$$

Here, $W_i^Q, W_i^K, W_i^V \in \mathbb{R}^{d_{model} \times d_k}$ are learnable weight matrices, $d_k$ is the dimensionality of each head, and typically $h \cdot d_k = d_{model}$. For the $i$-th head, the attention output is computed as:

$$
\text{Attention}(Q_i, K_i, V_i) = \text{softmax}\left(\frac{Q_i K_i^T}{\sqrt{d_k}}\right) V_i
$$

The outputs from all heads are concatenated and passed through a final linear projection $W^O$:

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

This design allows the model to jointly attend to syntactic structure, semantic similarity, and positional proximity, resulting in richer and more expressive representations.

Key Advantages

Compared with single-head attention, MHA offers several notable benefits:

  • Multiple representation subspaces: Different heads can learn different attention patterns. Some heads may focus on local dependencies between neighboring tokens, while others capture long-range semantic associations.
  • Computational efficiency: Although the total parameter count increases, each head operates in a smaller dimensional space and all heads are computed in parallel, making the layer well suited for GPU acceleration.
  • Improved generalization: By aggregating attention over multiple subspaces, the model becomes more robust to noise and irrelevant information, which often translates into better performance on complex downstream tasks.

Practical Implementation in PyTorch

PyTorch provides the nn.MultiheadAttention module, which simplifies the construction of MHA layers. The following example demonstrates a basic setup:

import torch
import torch.nn as nn

class MultiHeadAttentionDemo(nn.Module):
    def __init__(self, d_model: int = 512, num_heads: int = 8):
        super().__init__()
        self.attention = nn.MultiheadAttention(
            embed_dim=d_model,
            num_heads=num_heads,
            batch_first=True
        )

    def forward(self, x: torch.Tensor):
        # x shape: (batch_size, seq_len, d_model)
        output, weights = self.attention(query=x, key=x, value=x)
        return output, weights


model = MultiHeadAttentionDemo(d_model=512, num_heads=8)
dummy_input = torch.randn(2, 10, 512)
output, attn_weights = model(dummy_input)

print(f"Output shape: {output.shape}")           # torch.Size([2, 10, 512])
print(f"Attention weights shape: {attn_weights.shape}")  # torch.Size([2, 8, 10, 10])

Setting batch_first=True ensures that input and output tensors follow the (batch, seq, feature) convention commonly used in NLP pipelines. The returned attn_weights reveal the attention distribution across time steps for each head, which is valuable for model interpretability analysis.

Usage Considerations and Optimization

When applying MHA, several practical factors should be kept in mind:

  • Number of heads: The head count $h$ is usually chosen as a factor of $d_{model}$ (for example, 8, 16, or 32). Too many heads can reduce the dimensionality of each head and weaken expressiveness, while too few heads may limit the diversity of captured patterns.
  • Scaling factor: Dividing the dot-product scores by $\sqrt{d_k}$ prevents excessively large logits from pushing the softmax into regions with vanishing gradients. This scaling detail is critical for stable training.
  • Masking: In autoregressive tasks such as language modeling, a causal mask must be applied to prevent the model from attending to future positions. PyTorch’s nn.MultiheadAttention supports custom masks through the attn_mask argument.
  • Memory complexity: The attention matrix scales as $O(N^2)$ with sequence length, which can become a bottleneck for long inputs. Sparse attention or linear attention variants are common alternatives in such scenarios.

Applications

MHA is widely used across multiple domains. In natural language processing, it underpins machine translation, text summarization, and sentiment analysis. In computer vision, Vision Transformers rely on MHA to model relationships between image patches. The mechanism also proves effective in time-series forecasting, where it can capture complex temporal dependencies across different time scales.

Conclusion

Multi-Head Attention remains a foundational module for building high-performance Transformer models. A solid understanding of its internal mechanics, combined with careful configuration of hyperparameters such as head count and masking strategy, is essential for achieving strong results across a wide range of sequence modeling tasks.