Multi-Head Attention Mechanism

Introduction

Welcome to " Deconstructing the Transformer Architecture "! You've just completed an incredible journey through Sequence Models & The Dawn of Attention, where you built the foundational understanding of attention mechanisms and created a standalone PyTorch module for scaled dot-product attention. Now you're ready to take the next major step in your exploration of modern NLP architectures. In this course, you'll systematically build the complete Transformer architecture from the ground up. You'll start by enhancing your attention mechanism with Multi-Head Attention, then explore positional encodings, layer normalization, and feed-forward networks. By the end, you'll have a fully functional Transformer that can handle real sequence-to-sequence tasks. Our first lesson focuses on Multi-Head Attention, the mechanism that allows Transformers to attend to different types of information simultaneously across multiple representation subspaces.

The Power of Multiple Attention Perspectives

h

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

The Architecture: Linear Projections and Head Division

Let's begin implementing Multi-Head Attention by establishing the core architecture. The first crucial component involves creating separate linear projections for queries, keys, and values, then splitting these projections across multiple attention heads:

Python

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads, dropout=0.1):
        """
        Multi-Head Attention module
        
        Args:
            d_model: Model dimension
            num_heads: Number of attention heads
            dropout: Dropout rate
        """
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        # Linear projections for Q, K, V and output
        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_o = nn.Linear(d_model, d_model)
        
        self.dropout = nn.Dropout(dropout)
        self._init_weights()

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads, dropout=0.1):
        """
        Multi-Head Attention module
        
        Args:
            d_model: Model dimension
            num_heads: Number of attention heads
            dropout: Dropout rate
        """
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        # Linear projections for Q, K, V and output
        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_o = nn.Linear(d_model, d_model)
        
        self.dropout = nn.Dropout(dropout)
        self._init_weights()

The constructor establishes several critical design decisions. We ensure d_model is divisible by num_heads because we'll split the model dimension evenly across heads. Each head operates on $d_k = d_{model} / num\_heads$ dimensions, meaning if you have a 512-dimensional model with 8 heads, each head works with 64-dimensional projections. The four linear layers create separate learned transformations for queries, keys, values, and the final output projection that combines results from all heads.

The key insight here is that while you maintain the same total model dimension, you're creating multiple parallel pathways for information flow. Each of the projection matrices $W^Q$ , $W^K$ , and $W^V$ has shape $(d_{model} \times d_{model})$ , but their outputs will be reshaped to create multiple heads. The output projection $W^O$ serves a crucial role in learning how to optimally combine information from all heads back into a unified representation.

Implementing Scaled Dot-Product for Multiple Heads

Now you implement the core attention computation that will be applied independently to each head. This method encapsulates the scaled dot-product attention you learned previously, but is designed to handle the multi-head tensor structure:

Python

def scaled_dot_product_attention(self, q, k, v, mask=None):
    """Scaled dot-product attention for multiple heads"""
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
    
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    
    attention_weights = F.softmax(scores, dim=-1)
    attention_weights = self.dropout(attention_weights)
    
    output = torch.matmul(attention_weights, v)
    return output, attention_weights

def scaled_dot_product_attention(self, q, k, v, mask=None):
    """Scaled dot-product attention for multiple heads"""
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
    
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    
    attention_weights = F.softmax(scores, dim=-1)
    attention_weights = self.dropout(attention_weights)
    
    output = torch.matmul(attention_weights, v)
    return output, attention_weights

This implementation mirrors your previous scaled dot-product attention but operates on tensors with an additional head dimension. The key insight is that the same mathematical operations apply regardless of whether you're computing attention for one head or multiple heads simultaneously. The transpose(-2, -1) operation and matrix multiplications work seamlessly across the head dimension through broadcasting, allowing you to compute attention for all heads in parallel rather than using expensive loops. The dropout applied to the attention weights serves as a regularization mechanism, randomly zeroing out attention connections during training to prevent the model from over-relying on specific attention patterns and improve generalization.

The attention computation for each head follows the formula:

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

where $i$ represents the $i$ -th attention head. The scaling factor $\sqrt{d_k}$ becomes even more critical in multi-head attention because each head operates on a smaller dimension, making the dot products more prone to extreme values without proper scaling.

Tensor Reshaping and Multi-Head Computation

The heart of Multi-Head Attention lies in how you reshape tensors to create multiple heads and then recombine their outputs. This section handles the complex tensor manipulations that enable parallel attention computation:

Python

def forward(self, query, key, value, mask=None):
    """Forward pass for multi-head attention"""
    batch_size, seq_len_q = query.size(0), query.size(1)
    seq_len_k = key.size(1)
    
    # Linear projections and reshape for multi-head attention
    Q = self.w_q(query).view(batch_size, seq_len_q, self.num_heads, self.d_k).transpose(1, 2)
    K = self.w_k(key).view(batch_size, seq_len_k, self.num_heads, self.d_k).transpose(1, 2)
    V = self.w_v(value).view(batch_size, seq_len_k, self.num_heads, self.d_k).transpose(1, 2)

def forward(self, query, key, value, mask=None):
    """Forward pass for multi-head attention"""
    batch_size, seq_len_q = query.size(0), query.size(1)
    seq_len_k = key.size(1)
    
    # Linear projections and reshape for multi-head attention
    Q = self.w_q(query).view(batch_size, seq_len_q, self.num_heads, self.d_k).transpose(1, 2)
    K = self.w_k(key).view(batch_size, seq_len_k, self.num_heads, self.d_k).transpose(1, 2)
    V = self.w_v(value).view(batch_size, seq_len_k, self.num_heads, self.d_k).transpose(1, 2)

The reshaping operation is crucial: you transform tensors from (batch_size, seq_len, d_model) to (batch_size, num_heads, seq_len, d_k). The view operation splits the model dimension into separate heads, while transpose(1, 2) moves the head dimension to the second position for efficient computation. This tensor manipulation is what enables the "multi-head" aspect — you're essentially creating multiple parallel attention computations from a single input.

Python

    # Adjust mask for multi-head attention
    if mask is not None:
        if mask.dim() == 2:  # [seq_len, seq_len]
            mask = mask.unsqueeze(0).unsqueeze(0)  # [1, 1, seq_len, seq_len]
        elif mask.dim() == 3:  # [batch_size, seq_len, seq_len]
            mask = mask.unsqueeze(1)  # [batch_size, 1, seq_len, seq_len]
    
    # Apply attention
    attention_output, attention_weights = self.scaled_dot_product_attention(Q, K, V, mask)
    
    # Concatenate heads and apply output projection
    attention_output = attention_output.transpose(1, 2).contiguous().view(
        batch_size, seq_len_q, self.d_model)
    output = self.w_o(attention_output)
    
    return output, attention_weights

    # Adjust mask for multi-head attention
    if mask is not None:
        if mask.dim() == 2:  # [seq_len, seq_len]
            mask = mask.unsqueeze(0).unsqueeze(0)  # [1, 1, seq_len, seq_len]
        elif mask.dim() == 3:  # [batch_size, seq_len, seq_len]
            mask = mask.unsqueeze(1)  # [batch_size, 1, seq_len, seq_len]
    
    # Apply attention
    attention_output, attention_weights = self.scaled_dot_product_attention(Q, K, V, mask)
    
    # Concatenate heads and apply output projection
    attention_output = attention_output.transpose(1, 2).contiguous().view(
        batch_size, seq_len_q, self.d_model)
    output = self.w_o(attention_output)
    
    return output, attention_weights

The mask handling ensures compatibility with the multi-head structure by adding necessary dimensions for proper broadcasting across all heads. After computing attention, you reverse the reshaping process: transposing back and using view to concatenate all head outputs into the original d_model dimension. The contiguous() call ensures the tensor is stored in a contiguous block of memory, which is required before reshaping. Finally, the output projection w_o learns how to best combine information from all heads, implementing the concatenation and linear transformation specified in the original Transformer paper.

Testing Our Implementation

Now let's create tests to verify that your Multi-Head Attention implementation works correctly across different scenarios:

Python

def test_multi_head_attention():
    """Test Multi-Head Attention module"""
    print("Testing Multi-Head Attention Implementation...")
    
    # Configuration
    batch_size = 2
    seq_len = 8
    d_model = 64
    num_heads = 8
    
    # Create sample input
    torch.manual_seed(42)
    x = torch.randn(batch_size, seq_len, d_model)
    
    # Initialize Multi-Head Attention
    mha = MultiHeadAttention(d_model, num_heads)
    
    print(f"Input shape: {x.shape}")
    print(f"Model config: d_model={d_model}, num_heads={num_heads}, d_k={d_model//num_heads}")
    
    # Test self-attention (Q=K=V)
    output, attention_weights = mha(x, x, x)
    
    print(f"Output shape: {output.shape}")
    print(f"Attention weights shape: {attention_weights.shape}")
    
    # Verify dimensions are preserved
    assert output.shape == x.shape, f"Output shape {output.shape} != input shape {x.shape}"
    
    # Test with causal mask
    causal_mask = torch.tril(torch.ones(seq_len, seq_len))
    masked_output, masked_attn = mha(x, x, x, causal_mask)
    
    print(f"Masked output shape: {masked_output.shape}")
    
    # Test gradient flow
    loss = output.mean()
    loss.backward()
    
    print(f"Gradient verification:")
    for name, param in mha.named_parameters():
        if param.grad is not None:
            print(f"  {name}: grad_norm={param.grad.norm():.6f}")
    
    total_params = sum(p.numel() for p in mha.parameters())
    print(f"\nTotal parameters: {total_params:,}")

def test_multi_head_attention():
    """Test Multi-Head Attention module"""
    print("Testing Multi-Head Attention Implementation...")
    
    # Configuration
    batch_size = 2
    seq_len = 8
    d_model = 64
    num_heads = 8
    
    # Create sample input
    torch.manual_seed(42)
    x = torch.randn(batch_size, seq_len, d_model)
    
    # Initialize Multi-Head Attention
    mha = MultiHeadAttention(d_model, num_heads)
    
    print(f"Input shape: {x.shape}")
    print(f"Model config: d_model={d_model}, num_heads={num_heads}, d_k={d_model//num_heads}")
    
    # Test self-attention (Q=K=V)
    output, attention_weights = mha(x, x, x)
    
    print(f"Output shape: {output.shape}")
    print(f"Attention weights shape: {attention_weights.shape}")
    
    # Verify dimensions are preserved
    assert output.shape == x.shape, f"Output shape {output.shape} != input shape {x.shape}"
    
    # Test with causal mask
    causal_mask = torch.tril(torch.ones(seq_len, seq_len))
    masked_output, masked_attn = mha(x, x, x, causal_mask)
    
    print(f"Masked output shape: {masked_output.shape}")
    
    # Test gradient flow
    loss = output.mean()
    loss.backward()
    
    print(f"Gradient verification:")
    for name, param in mha.named_parameters():
        if param.grad is not None:
            print(f"  {name}: grad_norm={param.grad.norm():.6f}")
    
    total_params = sum(p.numel() for p in mha.parameters())
    print(f"\nTotal parameters: {total_params:,}")

This testing framework validates critical aspects of the implementation. You verify that input and output dimensions are preserved, test both unmasked and causal-masked scenarios, and examine gradient flow across all parameters. The self-attention setup (where query, key, and value are all the same input x) is fundamental to Transformer architectures and provides a clear test case for your implementation.

Introduction

The Power of Multiple Attention Perspectives

Join the 1M+ learners on CodeSignal

The Architecture: Linear Projections and Head Division

Implementing Scaled Dot-Product for Multiple Heads

Tensor Reshaping and Multi-Head Computation

Network Initialization

Testing Our Implementation

Verifying the Results

Computational Complexity and Optimization Considerations

Conclusion and Next Steps