Scaled Dot-Product Attention and Masking in Transformers

Welcome back to our journey through "Sequence Models & The Dawn of Attention"! We've made excellent progress, and I'm excited to dive deeper with you into lesson 3. In our previous lessons, we witnessed the limitations of traditional RNNs and LSTMs when handling long sequences, then discovered how attention mechanisms like Luong and Bahdanau elegantly solve the fixed-context bottleneck by enabling selective focus on different parts of input sequences.

Today, we're taking a crucial step toward the Transformer architecture by exploring scaled dot-product attention, the foundational attention mechanism that powers modern language models like GPT and BERT. We'll understand why the scaling factor is mathematically essential, then dive deep into masking techniques that are absolutely critical for real-world applications. Masking isn't just a technical detail; it's what enables models to ignore padding tokens and maintain the autoregressive property that makes language generation possible. Let's build this understanding step by step!

The scaling factor of $\frac{1}{\sqrt{d_k}}$ is mathematically necessary to avoid crippling our attention mechanism. When we compute attention scores using dot products between queries and keys, the magnitude of these scores grows with the dimensionality $d_k$ of our vectors. If we have high-dimensional vectors (say, $d_k = 512$ or $d_k = 1024$ as in real Transformers), the dot products can become very large. This creates a problem when we apply the softmax function: large input values to softmax result in extremely sharp probability distributions, where one element gets almost all the weight (approaching 1) and others get nearly zero weight.

This sharp distribution has devastating effects on gradients during backpropagation. The softmax function's gradient becomes vanishingly small in regions where the output is close to 0 or 1, leading to the vanishing gradient problem. By scaling the dot products by $\frac{1}{\sqrt{d_k}}$, we keep the variance of the attention scores roughly constant regardless of the dimensionality, ensuring that softmax operates in a regime where gradients remain healthy and learning can proceed effectively.

Let's implement the scaled dot-product attention mechanism, which forms the core of the Transformer architecture. This mechanism combines the efficiency of dot-product similarity with the mathematical stability provided by the scaling factor:

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

def scaled_dot_product_attention(query, key, value, mask=None):
    """
    Implement scaled dot-product attention
    
    Args:
        query: (batch_size, seq_len_q, d_k)
        key: (batch_size, seq_len_k, d_k)
        value: (batch_size, seq_len_v, d_v)
        mask: Optional mask tensor
    
    Returns:
        output: (batch_size, seq_len_q, d_v)
        attention_weights: (batch_size, seq_len_q, seq_len_k)
    """
    d_k = query.size(-1)
    
    # Compute attention scores with scaling
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    
    # Apply mask if provided
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    
    # Apply softmax and compute output
    attention_weights = F.softmax(scores, dim=-1)
    output = torch.matmul(attention_weights, value)
    
    return output, attention_weights

This implementation follows the mathematical formula: $\text{Attention}(Q,K,V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$. We compute the dot products between queries and keys, apply the crucial scaling factor, optionally apply masking (we'll explore this next), then use softmax to convert scores into attention weights. Finally, we compute the weighted sum of values using these attention weights. Notice how the mask parameter allows us to selectively ignore certain positions, which is essential for handling padding tokens and implementing causal attention for language modeling.