Chapter 2: Self-Attention and Multi-Head Attention

Introduction

Building on our understanding of the attention mechanism, this chapter explores how attention is applied when a sequence attends to itself (self-attention) and why using multiple attention heads in parallel dramatically improves model capability.

Self-Attention: The Core of Transformers

What is Self-Attention?

In self-attention, a sequence attends to itself. Each position in the input can attend to every other position in the same input.

Input:  "The cat sat on the mat"

Self-Attention: Each word attends to ALL words
├─ "The" attends to ["The", "cat", "sat", "on", "the", "mat"]
├─ "cat" attends to ["The", "cat", "sat", "on", "the", "mat"]
├─ "sat" attends to ["The", "cat", "sat", "on", "the", "mat"]
└─ ... and so on

Why Self-Attention?

Self-attention enables:

1. Long-range dependencies: Direct connections between any two positions
2. Parallel processing: All positions computed simultaneously
3. Dynamic context: Different weights for different inputs

Mathematical Formulation

For self-attention, Q, K, and V all come from the same input:

$$\text{SelfAttention}(X) = \text{softmax}\left(\frac{XW_Q(XW_K)^T}{\sqrt{d_k}}\right)XW_V$$

Implementation

import torch
import torch.nn as nn
import math

class SelfAttention(nn.Module):
    """
    Self-attention layer where Q, K, V all come from the same input.
    """
    def __init__(self, d_model, d_k=None):
        super().__init__()
        d_k = d_k 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_k, bias=False)
        self.d_k = d_k
    
    def forward(self, x, mask=None):
        """
        x: Input tensor (batch, seq_len, d_model)
        mask: Optional attention mask
        """
        # All projections from the same input
        Q = self.W_Q(x)
        K = self.W_K(x)
        V = self.W_V(x)
        
        # Compute attention scores
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        attention_weights = torch.softmax(scores, dim=-1)
        output = torch.matmul(attention_weights, V)
        
        return output, attention_weights

Comparing Attention Types

Self-Attention vs. Cross-Attention

Type	Q Source	K, V Source	Use Case
Self-Attention	Input X	Input X	Encoder, Decoder
Cross-Attention	Decoder	Encoder	Encoder-Decoder models

Cross-Attention Example

class CrossAttention(nn.Module):
    """
    Cross-attention where Q comes from one source, K and V from another.
    Used in encoder-decoder architectures.
    """
    def __init__(self, d_model, d_k=None):
        super().__init__()
        d_k = d_k 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_k, bias=False)
        self.d_k = d_k
    
    def forward(self, query_input, key_value_input, mask=None):
        """
        query_input: Source for Q (e.g., decoder state)
        key_value_input: Source for K, V (e.g., encoder output)
        """
        Q = self.W_Q(query_input)
        K = self.W_K(key_value_input)
        V = self.W_V(key_value_input)
        
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        attention_weights = torch.softmax(scores, dim=-1)
        output = torch.matmul(attention_weights, V)
        
        return output, attention_weights

Multi-Head Attention

The Motivation

A single attention head can only capture one type of relationship. Multi-head attention runs multiple attention operations in parallel, allowing the model to jointly attend to information from different representation subspaces.

Intuition

Different heads can learn different patterns:

• Head 1: Syntactic relationships (subject-verb)
• Head 2: Positional relationships (nearby words)
• Head 3: Semantic relationships (synonyms)
• Head 4: Coreference (pronouns to nouns)

Mathematical Definition

$$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, ..., \text{head}_h)W^O$$

Where each head is:

$$\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$$

Parameter Dimensions

For $h$ heads with model dimension $d_{model}$:

• Each head has dimension $d_k = d_v = d_{model} / h$
• Total parameters stay roughly the same as single-head attention

# Example: d_model = 512, h = 8
d_model = 512
num_heads = 8
d_k = d_model // num_heads  # 64 per head

Complete Implementation

class MultiHeadAttention(nn.Module):
    """
    Multi-head attention mechanism.
    """
    def __init__(self, d_model, num_heads, dropout=0.1):
        super().__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        # Linear projections for Q, K, V (all heads combined)
        self.W_Q = nn.Linear(d_model, d_model, bias=False)
        self.W_K = nn.Linear(d_model, d_model, bias=False)
        self.W_V = nn.Linear(d_model, d_model, bias=False)
        
        # Output projection
        self.W_O = nn.Linear(d_model, d_model, bias=False)
        
        self.dropout = nn.Dropout(dropout)
    
    def split_heads(self, x, batch_size):
        """
        Split the last dimension into (num_heads, d_k).
        Transpose to shape (batch, num_heads, seq_len, d_k)
        """
        x = x.view(batch_size, -1, self.num_heads, self.d_k)
        return x.transpose(1, 2)
    
    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)
        
        # Linear projections
        Q = self.W_Q(query)
        K = self.W_K(key)
        V = self.W_V(value)
        
        # Split into heads
        Q = self.split_heads(Q, batch_size)  # (batch, heads, seq, d_k)
        K = self.split_heads(K, batch_size)
        V = self.split_heads(V, batch_size)
        
        # Scaled dot-product attention for all heads in parallel
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        if mask is not None:
            # Expand mask for heads dimension
            if mask.dim() == 3:
                mask = mask.unsqueeze(1)
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        attention_weights = torch.softmax(scores, dim=-1)
        attention_weights = self.dropout(attention_weights)
        
        # Apply attention to values
        context = torch.matmul(attention_weights, V)
        
        # Concatenate heads
        context = context.transpose(1, 2).contiguous()
        context = context.view(batch_size, -1, self.d_model)
        
        # Final linear projection
        output = self.W_O(context)
        
        return output, attention_weights

Using Multi-Head Attention

# Example usage
d_model = 512
num_heads = 8
seq_len = 100
batch_size = 32

mha = MultiHeadAttention(d_model, num_heads)

# Input (same for self-attention)
x = torch.randn(batch_size, seq_len, d_model)

# Self-attention
output, attn_weights = mha(x, x, x)
print(f"Output shape: {output.shape}")  # (32, 100, 512)
print(f"Attention shape: {attn_weights.shape}")  # (32, 8, 100, 100)

Parallelization Benefits

Why Transformers Train Fast

Unlike RNNs, all positions in a Transformer can be computed simultaneously:

# RNN: Sequential (slow)
for t in range(seq_len):
    h_t = rnn_cell(x[t], h_{t-1})

# Transformer: Parallel (fast)
output = multi_head_attention(x, x, x)  # All positions at once

GPU Efficiency

Multi-head attention is particularly GPU-friendly:

# All heads computed in a single batched matrix multiplication
# Shape: (batch, heads, seq, d_k) @ (batch, heads, d_k, seq)
# → (batch, heads, seq, seq)
scores = torch.matmul(Q, K.transpose(-2, -1))  # Parallel across all heads!

Visualizing Multiple Heads

def visualize_multi_head_attention(attention_weights, tokens, num_heads_to_show=4):
    """
    Visualize attention patterns from multiple heads.
    attention_weights: (batch, heads, seq, seq)
    """
    import matplotlib.pyplot as plt
    
    fig, axes = plt.subplots(1, num_heads_to_show, figsize=(4*num_heads_to_show, 4))
    
    for head_idx in range(num_heads_to_show):
        ax = axes[head_idx]
        attn = attention_weights[0, head_idx].detach().numpy()  # First batch item
        
        im = ax.imshow(attn, cmap='Blues')
        ax.set_xticks(range(len(tokens)))
        ax.set_yticks(range(len(tokens)))
        ax.set_xticklabels(tokens, rotation=45, ha='right')
        ax.set_yticklabels(tokens)
        ax.set_title(f'Head {head_idx + 1}')
    
    plt.tight_layout()
    plt.show()

Attention Patterns in Practice

Research has shown that different heads learn specialized patterns:

Pattern Type	What It Captures	Example
Diagonal	Same position / identity	Token attends to itself
Previous token	Sequential relationships	Language model patterns
First token	Global aggregation	CLS token in BERT
Syntactic	Grammar structure	Subject-verb agreement
Semantic	Meaning relationships	Synonyms, antonyms

Key Takeaways

Summary

1. Self-attention allows each position to attend to all positions in the same sequence
2. Cross-attention enables attending to a different sequence (encoder-decoder)
3. Multi-head attention runs multiple attention operations in parallel
4. Different heads learn different types of relationships
5. Parallelization makes Transformers efficient on GPUs

Looking Ahead

In the next chapter, we'll explore how Transformers handle positional information through positional encoding, and how tokens are represented through embeddings.

---

← Chapter 1: Attention Mechanism

Chapter 3: Positional Encoding →