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:

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:

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.

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

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:

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:

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!