Building Attention Modules

Welcome to the fourth and final lesson of "Sequence Models & The Dawn of Attention"! We've accomplished so much together in this course. We started by understanding the limitations of RNNs and LSTMs, then discovered how attention mechanisms revolutionize sequence modeling through the Query-Key-Value paradigm. In our most recent lesson, we mastered scaled dot-product attention and learned how masking enables real-world applications by handling padding tokens and maintaining causal ordering for autoregressive generation.

Today, we're taking a significant step toward building production-ready attention systems by creating a standalone PyTorch module for scaled dot-product attention. Rather than writing attention as a simple function, we'll encapsulate it in a reusable nn.Module that can be easily integrated into larger architectures. This modular approach mirrors how real Transformer implementations work, where attention is a fundamental building block that is reused across multiple layers. We'll design our module to handle various input configurations, support masking, and include essential features like dropout for robust training.

As we move toward building complex architectures like Transformers, treating attention as a standalone function becomes limiting. Real neural networks require components that can maintain state, participate in gradient computation graphs, and integrate seamlessly with PyTorch's training infrastructure. This is where nn.Module shines: it provides automatic parameter registration, gradient tracking, and clean interfaces for complex operations.

Our modular attention design will accept pre-projected Query, Key, and Value tensors as inputs. This separation of concerns is crucial because, in full Transformer architectures, the same attention module is reused across different layers and heads, while the projection matrices (that create Q, K, V from input embeddings) vary. By focusing our module purely on the attention computation itself, we create a flexible building block that can serve multiple purposes.

Additionally, a proper module allows us to incorporate training-specific features like dropout, which is essential for preventing overfitting in large attention-based models. The module will also handle the complex broadcasting requirements for masks, making it robust enough for various input configurations we'll encounter in real applications.

Let's begin building our attention module by establishing the class structure and constructor. Our module needs to inherit from nn.Module and handle the essential components for attention computation:

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class ScaledDotProductAttention(nn.Module):
    def __init__(self, dropout=0.1):
        """
        Scaled Dot-Product Attention module
        
        Args:
            dropout: Dropout rate for attention weights
        """
        super(ScaledDotProductAttention, self).__init__()
        self.dropout = nn.Dropout(dropout)

Our constructor is deliberately simple, focusing on the essential parameters needed for attention computation. The dropout parameter controls regularization applied to attention weights, which is a standard technique in Transformer training to prevent overfitting. We initialize a nn.Dropout layer that we'll apply to attention weights before computing the final output.

Notice that we don't store any learned parameters in this module; the attention computation itself is parameter-free. The actual learned parameters (projection matrices for creating Q, K, V) will live in higher-level modules that use our attention as a building block. This design keeps our attention module focused and reusable.

Now let's implement the complete forward pass that computes scaled dot-product attention. This method encapsulates all the mathematical operations we've studied, including attention score computation, masking, softmax normalization, and dropout:

def forward(self, query, key, value, mask=None):
    """
    Forward pass for scaled dot-product attention
    
    Args:
        query: (batch_size (B), num_heads (H), seq_len_q (L_q), d_k) or (batch_size, seq_len_q, d_k)
        key: (batch_size, num_heads, seq_len_k (L_k), d_k) or (batch_size, seq_len_k, d_k)
        value: (batch_size, num_heads, seq_len_v (L_v), d_v) or (batch_size, seq_len_v, d_v)
        mask: (batch_size, 1, seq_len_q, seq_len_k) or (batch_size, seq_len_q, seq_len_k)
            Mask should be broadcastable to the attention score shape.              
    Returns:
        output: (batch_size, num_heads, seq_len_q, d_v) or (batch_size, seq_len_q, d_v)
        attention_weights: (batch_size, num_heads, seq_len_q, seq_len_k) or (batch_size, seq_len_q, seq_len_k)
    """
    d_k = query.size(-1)
    
    # Compute attention scores with scaling
    # scores shape: (batch_size [B], num_heads [H], seq_len_q [L_q], seq_len_k [L_k])
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    
    # Apply mask if provided
    if mask is not None:
        # Mask shape should be broadcastable to scores shape
        # (batch_size [B], 1, 1, seq_len_k [L_k]) for padding mask on keys
        # (batch_size [B], 1, seq_len_q [L_q], seq_len_k [L_k]) for combined padding/look-ahead
        # (1, 1, seq_len_q [L_q], seq_len_k [L_k]) for look-ahead mask
        # Masked positions (where mask == 0) are set to a large negative value so their softmax becomes ~0
        scores = scores.masked_fill(mask == 0, -1e9)
    
    # Apply softmax to get attention weights
    # attention_weights shape: (batch_size [B], num_heads [H], seq_len_q [L_q], seq_len_k [L_k])
    attention_weights = F.softmax(scores, dim=-1)
    attention_weights = self.dropout(attention_weights) # Apply dropout to attention weights
    
    # Compute output by multiplying attention weights with values
    # output shape: (batch_size [B], num_heads [H], seq_len_q [L_q], d_v)
    output = torch.matmul(attention_weights, value)
    
    return output, attention_weights

The forward method signature is designed for maximum flexibility. Our tensors can have either 3 dimensions for single-head attention or 4 dimensions for multi-head scenarios. The transpose(-2, -1) operation swaps the last two dimensions of the key tensor, which works correctly regardless of tensor dimensionality. The scaling factor math.sqrt(d_k) ensures numerical stability by preventing the dot products from growing too large, as we learned in the previous lesson.

The masking logic uses masked_fill to set masked positions to very large negative values (-1e9), which become essentially zero after softmax. This approach is more numerically stable than setting them to negative infinity. The detailed comments about mask shapes help users understand the various broadcasting patterns they might encounter in practice.

Applying dropout to attention weights (rather than the final output) is a subtle but important design choice from the original Transformer paper. This technique randomly zeroes some attention connections during training, forcing the model to maintain diverse attention strategies and preventing overfitting to specific patterns.

Before testing our module, let's create utility functions that generate realistic test data and demonstrate our module's flexibility:

def create_sample_qkv(batch_size, seq_len_q, seq_len_k, d_k, d_v):
    """Create sample Q, K, V tensors"""
    torch.manual_seed(42)
    q = torch.randn(batch_size, seq_len_q, d_k, requires_grad=True)
    k = torch.randn(batch_size, seq_len_k, d_k, requires_grad=True)
    v = torch.randn(batch_size, seq_len_k, d_v, requires_grad=True) # seq_len_k == seq_len_v
    return q, k, v

This utility function creates sample Query, Key, and Value tensors with different sequence lengths and dimensions. The requires_grad=True parameter ensures our tensors participate in gradient computation, validating that our module integrates properly with PyTorch's automatic differentiation system. Notice how the value tensor shares its sequence length with the key tensor (seq_len_k). This is not a coincidence: in attention mechanisms, the set of positions over which we attend (the keys) must exactly match the set of positions from which we gather information (the values); this means that for every key position, there is a corresponding value position, and the attention weights computed for each query are used to take a weighted sum over these value positions. As a result, the sequence length of the value tensor (seq_len_v) must always be equal to the sequence length of the key tensor (seq_len_k). This alignment is a fundamental property of all attention mechanisms.

Now let's create comprehensive tests to validate that our attention module works correctly across different scenarios. We'll test both masked and unmasked attention to ensure our implementation handles various real-world conditions:

def main():
    print("Testing ScaledDotProductAttention Module")
    
    # Configuration
    batch_size = 2
    seq_len_q = 5
    seq_len_k = 7 # seq_len_v will be same as seq_len_k
    d_k = 16 # Dimension of keys/queries
    d_v = 20 # Dimension of values
    
    # Create sample data
    query, key, value = create_sample_qkv(batch_size, seq_len_q, seq_len_k, d_k, d_v)
    
    print(f"Input shapes - Q: {query.shape}, K: {key.shape}, V: {value.shape}")
    
    # Initialize attention module
    attention_module = ScaledDotProductAttention(dropout=0.1)
    
    # Test forward pass without mask
    output_no_mask, attn_weights_no_mask = attention_module(query, key, value)
    print(f"Output shape (no mask): {output_no_mask.shape}")
    print(f"Attention weights shape (no mask): {attn_weights_no_mask.shape}")
    assert output_no_mask.shape == (batch_size, seq_len_q, d_v)
    assert attn_weights_no_mask.shape == (batch_size, seq_len_q, seq_len_k)
    
    # Test with a mask
    mask = torch.ones(batch_size, seq_len_q, seq_len_k)
    mask[:, :, -2:] = 0 # Mask out last two key positions for all queries
    
    output_masked, attn_weights_masked = attention_module(query, key, value, mask=mask)
    print(f"Output shape (masked): {output_masked.shape}")
    print(f"Attention weights shape (masked): {attn_weights_masked.shape}")
    assert output_masked.shape == (batch_size, seq_len_q, d_v)

Our test setup creates realistic scenarios where queries and keys have different sequence lengths (5 vs. 7), simulating cross-attention scenarios common in encoder-decoder architectures. We use different dimensions for keys/queries (d_k=16) and values (d_v=20) to test that our module correctly handles scenarios where the output dimension differs from the input dimension. This flexibility is essential in real Transformer architectures, where different components may use different dimensional representations.

The masking test creates a simple scenario where we block attention to the last two key positions. This simulates a common real-world case where certain positions (like future tokens or padding) should be ignored. The assertions ensure our module produces outputs with the expected shapes, confirming dimensional correctness.

When we run our complete test, it produces the following output:

Testing ScaledDotProductAttention Module
Input shapes - Q: torch.Size([2, 5, 16]), K: torch.Size([2, 7, 16]), V: torch.Size([2, 7, 20])
Output shape (no mask): torch.Size([2, 5, 20])
Attention weights shape (no mask): torch.Size([2, 5, 7])
Output shape (masked): torch.Size([2, 5, 20])
Attention weights shape (masked): torch.Size([2, 5, 7])

This output confirms our module correctly transforms the input dimensions: we start with queries of shape (2, 5, 16) and values of shape (2, 7, 20), and produce outputs of shape (2, 5, 20). The attention mechanism properly aggregates information from all 7 key-value positions to produce representations for each of the 5 query positions, with the final dimension matching the value dimension.

Congratulations! You've reached the end of our course "Sequence Models & The Dawn of Attention", the first course in the "Exploring the Transformer Architecture" course path! What an incredible journey we've taken together. We started by understanding the fundamental limitations of RNNs and LSTMs, discovered how attention mechanisms revolutionize sequence modeling, mastered scaled dot-product attention with comprehensive masking strategies, and today built a production-ready attention module that serves as a cornerstone of modern NLP architectures. The ScaledDotProductAttention module we've created isn't just an academic exercise: it's the exact type of building block used in real Transformer implementations.

Your journey into the world of Transformers is just beginning! In our next course, "Deconstructing the Transformer Architecture", we'll take these attention fundamentals and build the complete Transformer architecture from the ground up. We'll explore multi-head attention, positional encoding, layer normalization, and the encoder-decoder structure that powers models like GPT and BERT. Get ready to see how all these pieces come together to create the most influential architecture in modern AI!