Chapter 1: Attention Mechanism

Introduction

The attention mechanism is the core innovation that makes Transformers possible. It allows models to dynamically focus on relevant parts of the input when producing each output, overcoming the limitations of fixed-size context vectors.

The Intuition Behind Attention

A Human Analogy

Think about how you read a complex sentence:

"The animal didn't cross the street because it was too tired."

To understand what "it" refers to, you mentally attend to "the animal" rather than "the street". This is exactly what attention does computationally.

The Core Idea

Attention computes a weighted sum of values, where the weights are determined by how relevant each value is to a given query:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

Query, Key, Value: The Three Projections

Understanding Q, K, V

Component	Role	Analogy
Query (Q)	What we're looking for	A search query
Key (K)	What each position offers	Index entries
Value (V)	The actual content	The documents

Mathematical Definition

Given an input $X \in \mathbb{R}^{n \times d}$ (sequence of n tokens, each with d dimensions):

$$Q = XW_Q, \quad K = XW_K, \quad V = XW_V$$

Where:

• $W_Q \in \mathbb{R}^{d \times d_k}$
• $W_K \in \mathbb{R}^{d \times d_k}$
• $W_V \in \mathbb{R}^{d \times d_v}$

Implementation

import torch
import torch.nn as nn
import math

class AttentionProjections(nn.Module):
    def __init__(self, d_model, d_k, d_v):
        super().__init__()
        self.W_Q = nn.Linear(d_model, d_k, bias=False)
        self.W_K = nn.Linear(d_model, d_k, bias=False)
        self.W_V = nn.Linear(d_model, d_v, bias=False)
    
    def forward(self, x):
        """
        x: (batch, seq_len, d_model)
        Returns Q, K, V projections
        """
        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)
        return Q, K, V

Scaled Dot-Product Attention

Step-by-Step Computation

Step 1: Compute Attention Scores

$$\text{scores} = QK^T$$

This gives us a matrix where entry $(i, j)$ represents how much position $i$ should attend to position $j$.

# Q: (batch, seq_len, d_k)
# K: (batch, seq_len, d_k)
scores = torch.matmul(Q, K.transpose(-2, -1))
# scores: (batch, seq_len, seq_len)

Step 2: Scale the Scores

$$\text{scaled\_scores} = \frac{QK^T}{\sqrt{d_k}}$$

Why Scale?

Without scaling, for large $d_k$, the dot products can grow large in magnitude, pushing softmax into regions with extremely small gradients. Scaling by $\sqrt{d_k}$ keeps the variance stable.

d_k = Q.size(-1)
scaled_scores = scores / math.sqrt(d_k)

Step 3: Apply Softmax

$$\text{attention\_weights} = \text{softmax}(\text{scaled\_scores})$$

attention_weights = torch.softmax(scaled_scores, dim=-1)
# Each row sums to 1

Step 4: Weighted Sum of Values

$$\text{output} = \text{attention\_weights} \cdot V$$

output = torch.matmul(attention_weights, V)
# output: (batch, seq_len, d_v)

Complete Implementation

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Compute scaled dot-product attention.
    
    Args:
        Q: Query tensor (batch, seq_len, d_k)
        K: Key tensor (batch, seq_len, d_k)
        V: Value tensor (batch, seq_len, d_v)
        mask: Optional mask tensor
    
    Returns:
        output: Attention output (batch, seq_len, d_v)
        attention_weights: Attention weights (batch, seq_len, seq_len)
    """
    d_k = Q.size(-1)
    
    # Step 1 & 2: Compute scaled scores
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    
    # Optional: Apply mask (for causal attention)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    
    # Step 3: Softmax
    attention_weights = torch.softmax(scores, dim=-1)
    
    # Step 4: Weighted sum
    output = torch.matmul(attention_weights, V)
    
    return output, attention_weights

Visualizing Attention

Attention Matrix

For a sentence like "The cat sat":

         The   cat   sat
The    [ 0.7   0.2   0.1 ]
cat    [ 0.3   0.5   0.2 ]
sat    [ 0.2   0.4   0.4 ]

Each row shows how much that position attends to all other positions.

Code to Visualize

import matplotlib.pyplot as plt
import seaborn as sns

def visualize_attention(attention_weights, tokens):
    """
    Visualize attention weights as a heatmap.
    """
    plt.figure(figsize=(10, 8))
    sns.heatmap(
        attention_weights.detach().numpy(),
        xticklabels=tokens,
        yticklabels=tokens,
        cmap='Blues',
        annot=True,
        fmt='.2f'
    )
    plt.xlabel('Keys')
    plt.ylabel('Queries')
    plt.title('Attention Weights')
    plt.show()

Causal (Masked) Attention

For autoregressive language models, each position should only attend to previous positions.

The Causal Mask

def create_causal_mask(seq_len):
    """
    Create a causal mask for autoregressive attention.
    Position i can only attend to positions <= i.
    """
    mask = torch.tril(torch.ones(seq_len, seq_len))
    return mask

# Example for seq_len=4:
# [[1, 0, 0, 0],
#  [1, 1, 0, 0],
#  [1, 1, 1, 0],
#  [1, 1, 1, 1]]

Why Causal Masking?

During training, we want the model to predict the next token based only on previous tokens, not future ones (which it won't have at inference time).

def causal_attention(Q, K, V):
    """
    Attention with causal masking.
    """
    seq_len = Q.size(1)
    mask = create_causal_mask(seq_len).to(Q.device)
    return scaled_dot_product_attention(Q, K, V, mask)

Attention Complexity

Time Complexity

$$O(n^2 \cdot d)$$

Where:

• $n$ = sequence length
• $d$ = model dimension

This quadratic complexity in sequence length is why long-context models are challenging.

Memory Complexity

$$O(n^2)$$

For the attention matrix itself (storing all pairwise scores).

Optimization Techniques

Several approaches address attention's quadratic complexity:

1. Flash Attention: Memory-efficient implementation
2. Linear Attention: Approximate attention in O(n)
3. Sparse Attention: Only attend to subset of positions
4. Sliding Window: Local attention with global tokens

Putting It All Together

class Attention(nn.Module):
    """
    Complete attention module.
    """
    def __init__(self, d_model, d_k=None, d_v=None, causal=False):
        super().__init__()
        d_k = d_k or d_model
        d_v = d_v or d_model
        
        self.W_Q = nn.Linear(d_model, d_k, bias=False)
        self.W_K = nn.Linear(d_model, d_k, bias=False)
        self.W_V = nn.Linear(d_model, d_v, bias=False)
        self.W_O = nn.Linear(d_v, d_model, bias=False)
        
        self.causal = causal
        self.d_k = d_k
    
    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.size()
        
        # Project to Q, K, V
        Q = self.W_Q(x)
        K = self.W_K(x)
        V = self.W_V(x)
        
        # Compute attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        # Apply causal mask if needed
        if self.causal:
            causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=x.device))
            scores = scores.masked_fill(causal_mask == 0, float('-inf'))
        
        # Apply additional mask if provided
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        # Softmax and weighted sum
        attn_weights = torch.softmax(scores, dim=-1)
        output = torch.matmul(attn_weights, V)
        
        # Output projection
        output = self.W_O(output)
        
        return output, attn_weights

Key Takeaways

Summary

1. Attention computes weighted sums based on relevance
2. Q, K, V are three projections of the input with distinct roles
3. Scaling by $\sqrt{d_k}$ stabilizes training
4. Causal masking prevents looking at future tokens
5. Complexity is $O(n^2)$ in sequence length

Looking Ahead

In the next chapter, we'll explore how attention is used within a single sequence (self-attention) and how multiple attention mechanisms are combined (multi-head attention).

---

← Module Overview

Chapter 2: Self & Multi-Head Attention →