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
The Architecture: Linear Projections and Head Division
Implementing Scaled Dot-Product for Multiple Heads
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:
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.
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.
Network Initialization
Proper weight initialization and comprehensive testing are crucial for ensuring your Multi-Head Attention module functions correctly. Let's implement the initialization method and create a thorough testing framework:
def _init_weights(self): """Initialize weights using Xavier uniform initialization""" for linear in [self.w_q, self.w_k, self.w_v, self.w_o]: nn.init.xavier_uniform_(linear.weight) if linear.bias is not None: nn.init.constant_(linear.bias, 0)
Xavier uniform initialization helps ensure stable training by preventing vanishing or exploding gradients during the initial training phases. This initialization strategy considers the number of input and output connections to set appropriate initial weight magnitudes.
Testing Our Implementation
Now let's create tests to verify that your Multi-Head Attention implementation works correctly across different scenarios:
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.
Verifying the Results
Computational Complexity and Optimization Considerations
Conclusion and Next Steps
Congratulations! You've successfully implemented Multi-Head Attention, a cornerstone mechanism that enables Transformers to capture multiple types of relationships simultaneously. Your implementation demonstrates how parallel attention heads can attend to different representation subspaces, providing a much richer understanding than single-head mechanisms. The careful tensor reshaping and concatenation logic you've mastered form the foundation for more complex Transformer components.
In our next lesson, you'll explore positional encodings, the mechanism that gives Transformers their understanding of sequence order. Unlike RNNs that process sequences step by step, Transformers need explicit positional information, and you'll discover the elegant mathematical solutions that make this possible. Get ready to dive deeper into the architectural innovations that make Transformers so powerful!
Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal
While our single-head attention mechanism from the previous course works well, it has a fundamental limitation: it can only focus on one type of relationship at a time. Imagine reading a sentence where you need to simultaneously track grammatical dependencies, semantic relationships, and contextual references. A single attention head might excel at one of these tasks but struggle to capture all of them effectively.
Multi-Head Attention solves this by running multiple attention computations in parallel, each focusing on different aspects of the input relationships. Think of it as having multiple experts, each specializing in different types of patterns. One head might focus on local dependencies, another on long-range relationships, and yet another on specific semantic patterns. The key insight is that by running these computations in parallel and combining their results, you can capture much richer representations than any single attention mechanism could provide.
The mathematical foundation remains the same scaled dot-product attention you mastered previously, but now you apply it across multiple "heads" simultaneously. Each head operates on different learned projections of the input, allowing the model to attend to information from different representation subspaces at different positions. Mathematically, if you have h heads and model dimension dmodel, each head operates on dimension dk=dmodel/h, ensuring computational efficiency while maintaining expressiveness.
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:
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 dk=dmodel/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 WQ, WK, and WV has shape (dmodel×dmodel), but their outputs will be reshaped to create multiple heads. The output projection WO serves a crucial role in learning how to optimally combine information from all heads back into a unified representation.
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:
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:
Attention(Qi,Ki,Vi)=softmax(dkQiKiT)Vi
where i represents the i-th attention head. The scaling factor dk 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.
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 _init_weights(self): """Initialize weights using Xavier uniform initialization""" for linear in [self.w_q, self.w_k, self.w_v, self.w_o]: nn.init.xavier_uniform_(linear.weight) if linear.bias is not None: nn.init.constant_(linear.bias, 0)
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:,}")
When you run this complete test, it produces the following comprehensive output that confirms your implementation works correctly:
This output confirms several crucial aspects of your implementation: input and output shapes are preserved, attention weights have the expected multi-head structure (batch, heads, seq, seq), and gradient flow is verified across all parameters. The causal mask creates a lower triangular matrix that prevents positions from attending to future positions, essential for autoregressive generation tasks. The total parameter count of 16,640 reflects four projection matrices, each with 64×64=4,096 weights plus 64 bias terms, totaling 4×(4,096+64)=16,640 parameters.
Multi-Head Attention, while powerful, comes with significant computational costs that scale quadratically with sequence length. The time complexity is O(n2⋅d) where n is the sequence length and d is the model dimension, dominated by the computation of attention scores between all pairs of positions. Memory complexity follows a similar pattern, requiring O(n2) storage for attention matrices across all heads. For longer sequences (thousands of tokens), this quadratic scaling becomes prohibitive, making standard attention computationally expensive for tasks like document processing or long-form generation.
Several optimization techniques have emerged to address these limitations. Flash Attention restructures the computation to be more memory-efficient by avoiding the materialization of large attention matrices, achieving significant speedups especially on modern GPUs. Sparse attention patterns limit attention to local windows or specific patterns (like strided or dilated attention), reducing complexity to O(n⋅n) or O(n⋅logn). Other approaches include linear attention approximations that replace the softmax with more efficient operations, gradient checkpointing to trade computation for memory during training, and mixed precision training using FP16 or BF16 to reduce memory footprint while maintaining numerical stability. Understanding these trade-offs becomes crucial when deploying Transformer models in resource-constrained environments or scaling to very long sequences.
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()
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
