Self-Attention Mechanism

1. Historical Context & Motivation

The Problem with Sequential Models

Before transformers, the dominant architectures for sequence processing were Recurrent Neural Networks (RNNs) and their variants (LSTMs, GRUs). These models processed sequences one token at a time, maintaining a hidden state that accumulated information:

This sequential dependency created several critical problems:

• Vanishing/Exploding Gradients: Information had to pass through many time steps. Gradients either vanished (making early tokens hard to learn from) or exploded. LSTMs and GRUs mitigated but didn't solve this.
• Limited Long-Range Dependencies: Even LSTMs struggled to connect information across hundreds of tokens. The hidden state acted as a "bottleneck" - all context had to fit in a fixed-size vector.
• No Parallelization: Because h_t depends on h_t-1, you couldn't compute positions in parallel during training. This made RNNs extremely slow to train on modern GPUs.

The Attention Revolution

Attention was first introduced in 2014 for machine translation (Bahdanau et al.). Instead of compressing all source information into a single vector, the decoder could "look back" at all encoder states and decide which ones were relevant.

The 2017 paper "Attention Is All You Need" (Vaswani et al.) took this further: what if we used ONLY attention, with no recurrence at all? This gave birth to the Transformer architecture.

2. The Core Intuition

Attention as Dynamic, Content-Based Routing

Think of self-attention as a sophisticated information routing system. For each position in the sequence, we ask: "Given what I'm trying to compute here, which other positions in this sequence contain relevant information, and how should I combine that information?"

Unlike convolutions (which have fixed receptive fields) or RNNs (which have fixed routing through time), attention routing is:

• Dynamic: The routing pattern changes based on the actual content of the input, not just positions
• Learned: The model learns what constitutes "relevance" during training
• All-to-All: Every position can potentially attend to every other position

The Database Analogy

A useful mental model is to think of attention like a fuzzy database lookup:

3. Query, Key, Value: Deep Dive

The Linear Projections

Given an input sequence X ∈ ℝ^n×d (n tokens, each with d-dimensional embedding), we create Q, K, V through learned linear transformations:

Dimension Choices

In the original transformer and most modern LLMs:

• d_k = d_v = d / h, where h is the number of attention heads
• For example, with d=768 and h=12 heads, each head has d_k=d_v=64
• Some architectures (like Multi-Query Attention) use different dimensions for K and V to reduce memory

Interactive: QKV Vectors

What Do These Projections Learn?

During training, the model learns weight matrices that create useful query/key/value spaces:

• W_Q learns to project tokens into a "question space" - encoding what information each position needs
• W_K learns to project tokens into an "answer space" - encoding what information each position offers, in a way that matches relevant queries
• W_V learns to project tokens into a "content space" - encoding the actual information to be aggregated

Research has shown that attention heads often specialize. Some heads learn syntactic relationships (subject-verb), others learn positional patterns (attend to previous token), and others learn semantic relationships (entity coreference).

4. Mathematical Foundations

The Attention Formula

The complete scaled dot-product attention formula is:

Let's break down each component:

Step 1: Compute Attention Scores (Q K^T)

The matrix multiplication Q K^T computes the dot product between every query and every key:

The result is an n×n matrix where entry (i,j) represents how much position i should attend to position j. Higher values indicate higher relevance.

Step 2: Scale by √d_k

We divide by √d_k to counteract the effect of large dot products:

Step 3: Softmax Normalization

We apply softmax along the key dimension (each row):

Properties of the resulting attention weights α:

• Non-negative: α_ij ≥ 0 for all i, j
• Normalized: Σ_j α_ij = 1 for each row i
• Interpretable: α_ij can be viewed as "probability that position i attends to position j"

Step 4: Weighted Sum of Values

Finally, we compute a weighted combination of value vectors:

In matrix form: Output = α · V, where α is the n×n attention weight matrix and V is n×d_v.

5. Interactive Visualization

Understanding the Attention Pattern

The visualization shows a simulated attention pattern that a trained language model might produce. Notice:

• "sat" strongly attends to "cat" - the verb connects to its subject
• "the" (second) attends to "The" (first) - learning that these are the same word in different positions
• Each position attends to itself - often important for preserving its own information
• The lower triangle pattern - this is the causal mask, preventing positions from seeing the future

6. PyTorch Implementation

Let's implement self-attention from scratch. We'll build up from a minimal version to a more complete implementation.

Minimal Implementation

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

def attention(Q, K, V, mask=None):
    """
    Minimal scaled dot-product attention.

    Args:
        Q: Queries (batch, seq_len, d_k)
        K: Keys (batch, seq_len, d_k)
        V: Values (batch, seq_len, d_v)
        mask: Optional mask (batch, seq_len, seq_len) or (1, seq_len, seq_len)

    Returns:
        output: Attended values (batch, seq_len, d_v)
        weights: Attention weights (batch, seq_len, seq_len)
    """
    d_k = Q.size(-1)

    # Compute attention scores
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)

    # Apply mask (if provided)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)

    # Softmax to get attention weights
    weights = F.softmax(scores, dim=-1)

    # Weighted sum of values
    output = torch.matmul(weights, V)

    return output, weights

Complete Self-Attention Module

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Optional, Tuple

class SelfAttention(nn.Module):
    """
    Scaled Dot-Product Self-Attention.

    This is single-head attention. For multi-head, see the next lesson.

    Args:
        d_model: Input embedding dimension
        d_k: Key/Query dimension (defaults to d_model)
        d_v: Value dimension (defaults to d_k)
        dropout: Dropout probability on attention weights
    """

    def __init__(
        self,
        d_model: int,
        d_k: Optional[int] = None,
        d_v: Optional[int] = None,
        dropout: float = 0.0
    ):
        super().__init__()

        self.d_k = d_k if d_k is not None else d_model
        self.d_v = d_v if d_v is not None else self.d_k
        self.scale = math.sqrt(self.d_k)

        # Linear projections
        # Note: bias=False is common in modern LLMs (LLaMA, etc.)
        self.W_q = nn.Linear(d_model, self.d_k, bias=False)
        self.W_k = nn.Linear(d_model, self.d_k, bias=False)
        self.W_v = nn.Linear(d_model, self.d_v, bias=False)

        # Output projection (projects back to d_model if d_v != d_model)
        self.W_o = nn.Linear(self.d_v, d_model, bias=False)

        # Dropout on attention weights (regularization)
        self.dropout = nn.Dropout(dropout)

    def forward(
        self,
        x: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
        return_weights: bool = False
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """
        Forward pass.

        Args:
            x: Input tensor (batch, seq_len, d_model)
            mask: Attention mask (batch, 1, seq_len) or (batch, seq_len, seq_len)
                  1 = attend, 0 = mask out
            return_weights: Whether to return attention weights

        Returns:
            output: (batch, seq_len, d_model)
            weights: (batch, seq_len, seq_len) if return_weights=True
        """
        batch_size, seq_len, _ = x.shape

        # Project to Q, K, V
        Q = self.W_q(x)  # (batch, seq_len, d_k)
        K = self.W_k(x)  # (batch, seq_len, d_k)
        V = self.W_v(x)  # (batch, seq_len, d_v)

        # Compute attention scores
        # (batch, seq_len, d_k) @ (batch, d_k, seq_len) -> (batch, seq_len, seq_len)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale

        # Apply mask
        if mask is not None:
            # Use -inf so softmax gives 0 probability
            scores = scores.masked_fill(mask == 0, float('-inf'))

        # Softmax over keys (last dimension)
        weights = F.softmax(scores, dim=-1)

        # Apply dropout to attention weights
        weights = self.dropout(weights)

        # Weighted sum of values
        # (batch, seq_len, seq_len) @ (batch, seq_len, d_v) -> (batch, seq_len, d_v)
        attended = torch.matmul(weights, V)

        # Project back to d_model
        output = self.W_o(attended)

        if return_weights:
            return output, weights
        return output, None


# ===== Example Usage =====
if __name__ == "__main__":
    # Hyperparameters
    batch_size = 2
    seq_len = 10
    d_model = 512

    # Create random input (simulating token embeddings)
    x = torch.randn(batch_size, seq_len, d_model)

    # Create attention layer
    attn = SelfAttention(d_model=d_model, dropout=0.1)

    # Forward pass (no mask = full attention)
    output, weights = attn(x, return_weights=True)

    print(f"Input shape:  {x.shape}")        # (2, 10, 512)
    print(f"Output shape: {output.shape}")   # (2, 10, 512)
    print(f"Weights shape: {weights.shape}") # (2, 10, 10)

    # Verify attention weights sum to 1
    print(f"Row sums: {weights[0].sum(dim=-1)}")  # All ~1.0

7. Causal Masking for Autoregressive Models

Why Causal Masking?

In autoregressive language models (GPT, LLaMA, Claude), we generate text left-to-right, predicting each token based only on previous tokens. During training, we need to enforce this constraint so the model can't "cheat" by looking at future tokens.

Causal mask (also called "look-ahead mask"): Position i can only attend to positions 0, 1, ..., i.

def create_causal_mask(seq_len: int, device: torch.device = None) -> torch.Tensor:
    """
    Creates a causal (lower triangular) attention mask.

    Args:
        seq_len: Sequence length
        device: Device to create tensor on

    Returns:
        mask: (1, seq_len, seq_len) tensor
              1 = can attend, 0 = cannot attend (will become -inf)
    """
    # torch.tril creates lower triangular matrix
    mask = torch.tril(torch.ones(seq_len, seq_len, device=device))
    # Add batch dimension for broadcasting
    return mask.unsqueeze(0)  # (1, seq_len, seq_len)


# Example: seq_len = 5
mask = create_causal_mask(5)
print(mask[0])
# tensor([[1., 0., 0., 0., 0.],   <- token 0 sees only itself
#         [1., 1., 0., 0., 0.],   <- token 1 sees tokens 0, 1
#         [1., 1., 1., 0., 0.],   <- token 2 sees tokens 0, 1, 2
#         [1., 1., 1., 1., 0.],   <- token 3 sees tokens 0, 1, 2, 3
#         [1., 1., 1., 1., 1.]])  <- token 4 sees all tokens


# Combining causal mask with padding mask
def create_combined_mask(
    seq_len: int,
    padding_mask: torch.Tensor = None  # (batch, seq_len), 1=real token, 0=padding
) -> torch.Tensor:
    """
    Creates mask that combines causal constraint with padding.

    Returns:
        mask: (batch, seq_len, seq_len)
    """
    # Start with causal mask
    causal = create_causal_mask(seq_len, device=padding_mask.device)

    if padding_mask is not None:
        # Expand padding mask: don't attend to padding tokens
        # (batch, seq_len) -> (batch, 1, seq_len)
        padding = padding_mask.unsqueeze(1)
        # Combine: must satisfy BOTH causal AND not-padding
        mask = causal * padding

    return mask

Bidirectional vs Causal Attention

8. Computational Complexity Analysis

Time Complexity

For a sequence of length n with dimension d:

Total: O(n² · d) - The quadratic dependence on n is the main bottleneck for long sequences.

Memory Complexity

The attention weight matrix α ∈ ℝ^n×n must be stored (at least temporarily):

• For n = 2048 tokens: 2048² × 4 bytes = 16 MB per attention layer per batch element
• For n = 100,000 tokens: 100,000² × 4 bytes = 40 GB per layer! (Infeasible)

Comparison with RNNs

For typical sequence lengths (hundreds to low thousands), self-attention's parallelizability makes it much faster to train on GPUs despite higher theoretical complexity.

9. Attention Variants (Preview)

The basic self-attention we've covered is the foundation, but modern LLMs use several important variants:

Multi-Head Attention

Instead of one attention operation, run h parallel attention "heads" with smaller dimensions (d_k = d/h), then concatenate results. This allows the model to jointly attend to information from different representation subspaces. → Next lesson

Grouped Query Attention (GQA)

Used in Llama 2, Mistral. Share K and V projections across groups of heads to reduce memory (KV cache). We'll cover this in Module 6.

Multi-Query Attention (MQA)

Extreme version of GQA: all heads share the same K and V. Faster inference but potentially lower quality.

Multi-head Latent Attention (MLA)

Used in DeepSeek-V2/V3. Uses low-rank projection to compress KV cache. Covered in Module 7.

Flash Attention

Not a different attention formula, but a memory-efficient implementation that avoids materializing the full attention matrix. Critical for long contexts. Covered in Module 6.