Multi-Head Attention: Theory and Implementation

Introduction

In our last session, we implemented the core engine of the Transformer: scaled dot-product attention. You now have a solid grasp of how this mechanism allows a model to weigh the importance of different tokens using the Query, Key, and Value paradigm. However, you can think of this as giving the model a single "lens" through which to view the relationships in a sequence.

Today, we will significantly upgrade that engine by implementing multi-head attention. This is how the Transformer learns to look at a sentence from multiple perspectives simultaneously—for instance, focusing on grammatical structure with one "head," semantic relationships with another, and pronoun references with a third. Our goal is to implement this as a PyTorch nn.Module and, more importantly, understand how this parallelization enables the model to learn and combine information from different representation subspaces.

The Limitation of a Single Head

Imagine the sentence: "The tired runner drank water from the river bank."

A single attention mechanism, operating on one set of Q, K, and V projections, might learn to associate "runner" with "drank" and "water". But what about the relationship between "river" and "bank"? This context is crucial to disambiguate "bank" from its financial meaning. A single head might struggle to capture both the "action" relationship (runner -> drank) and the "location" relationship (river -> bank) at the same time. The averaging process within a single attention mechanism can wash out these distinct signals.

This is the core motivation for multi-head attention: it allows the model to jointly attend to information from different perspectives at different positions.

To build a strong intuition for this, let's watch a segment from 3Blue1Brown.

The Multi-Head Attention Mechanism

As the video explained, the solution is to run multiple scaled dot-product attention operations in parallel. Each "head" gets its own set of learned linear projection matrices. This allows each head to project the input embeddings into a different, lower-dimensional representation subspace where it can focus on a specific type of relationship.

The process can be broken down into five steps:

1. Initial Projections: Instead of one set of projection matrices, we have h (the number of heads). In practice, this is often implemented efficiently using three large linear layers for Q, K, and V that project the input from dimension d_model to d_model.
2. Split into Heads: The resulting (batch, seq_len, d_model) tensors are reshaped into (batch, num_heads, seq_len, d_k), where d_k = d_model / num_heads. This effectively gives each head its own smaller Q, K, and V vectors.
3. Parallel Attention Calculation: The scaled dot-product attention function you implemented in the last lesson is applied to each head's Q, K, and V tensors simultaneously.
4. Concatenate Outputs: The outputs from all h heads are concatenated back together, resulting in a tensor of shape (batch, seq_len, d_model).
5. Final Linear Projection: This concatenated tensor is passed through one final linear layer (W_o) to produce the final output of the multi-head attention block. This layer allows the model to learn how to best combine the information learned by the different heads.

The "Annotated Transformer" blog post, which is the canonical implementation guide, formalizes this.

Implementing Multi-Head Attention in PyTorch

Now, let's translate this into code. Your experience building production ML systems and your comfort with PyTorch will be very helpful here, as the main challenge lies in correctly manipulating tensor dimensions.

We will build a MultiHeadAttention nn.Module. For a detailed guide and a clean codebase to follow, the "Transformers From Scratch" series is an excellent resource.

The MultiHeadAttention Class

Based on the article, here is the structure of our class. We'll need the scaled_dot_product_attention function we built in the last lesson.

import torch
import torch.nn as nn
import math

# Assume scaled_dot_product_attention is defined from the previous lesson
def scaled_dot_product_attention(query, key, value, mask=None, dropout=None):
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9) # A common way to apply mask
    
    attention_weights = torch.nn.functional.softmax(scores, dim=-1)
    
    if dropout is not None:
        attention_weights = dropout(attention_weights)
        
    output = torch.matmul(attention_weights, value)
    return output, attention_weights


class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1):
        super().__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  # Dimension of each head

        # Efficient implementation: single large linear layers for all heads
        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)
        
        # Final output projection layer
        self.W_o = nn.Linear(d_model, d_model)
        
        self.dropout = nn.Dropout(dropout)
        self.attention_weights = None

    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor = None):
        batch_size = query.shape[0]

        # 1. Apply linear projections: (batch, seq_len, d_model) -> (batch, seq_len, d_model)
        Q = self.W_q(query)
        K = self.W_k(key)
        V = self.W_v(value)

        # 2. Reshape for multi-head processing
        # (batch, seq_len, d_model) -> (batch, num_heads, seq_len, d_k)
        Q = Q.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        
        # 3. Apply scaled dot-product attention for each head in parallel
        # Q, K, V shape: (batch, num_heads, seq_len, d_k)
        output, self.attention_weights = scaled_dot_product_attention(Q, K, V, mask, self.dropout)
        # output shape: (batch, num_heads, seq_len_q, d_k)
        
        # 4. Concatenate heads and reshape back
        # (batch, num_heads, seq_len_q, d_k) -> (batch, seq_len_q, d_model)
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)

        # 5. Apply final linear projection
        # (batch, seq_len_q, d_model) -> (batch, seq_len_q, d_model)
        output = self.W_o(output)
        
        return output

The most critical part is Step 2, the reshaping operation.

• Q.view(batch_size, -1, self.num_heads, self.d_k) reshapes the d_model dimension into two separate dimensions: num_heads and d_k. The tensor shape becomes (batch_size, seq_len, num_heads, d_k).
• .transpose(1, 2) then swaps the seq_len and num_heads dimensions. The final shape is (batch_size, num_heads, seq_len, d_k). This groups all the data by head, allowing the scaled_dot_product_attention function to perform batched computations across both the batch_size and num_heads dimensions.

Step 4 simply reverses this process to combine the heads' outputs.

To see these tensor manipulations performed in a live coding session, the following video is an excellent supplement.

Conclusion

Congratulations! You have now implemented one of the most important components of the Transformer. By moving from a single attention head to multiple parallel heads, you've unlocked the model's ability to capture far more nuanced and varied relationships within the data.

Key Takeaways:

• Motivation: Multi-head attention overcomes the limitations of single-head attention by allowing the model to focus on different types of information (e.g., syntactic, semantic) in parallel.
• Representation Subspaces: Each head operates in a different learned "subspace," created by its unique set of projection matrices (W_q_i, W_k_i, W_v_i).
• Mechanism: The core idea is to split d_model into h smaller chunks, run scaled dot-product attention on each chunk in parallel, and then concatenate the results and pass them through a final mixing layer.
• Implementation: The key to an efficient implementation in PyTorch is the use of view() and transpose() to reshape tensors for parallel processing across heads, rather than using explicit loops or separate layers for each head.

What's Next?

So far, we have focused on self-attention, where Q, K, and V all come from the same input sequence. But what happens in an encoder-decoder model where the decoder needs to pay attention to the encoder's output?

In our next lesson, we will compare and contrast different applications of the attention mechanism: self-attention, cross-attention, and masked self-attention. We will also implement causal masking, which is essential for ensuring that a decoder can only attend to previous tokens when generating text, preserving the autoregressive property.