Jie Zhu's Blog

The 2017 paper Attention Is All You Need by Vaswani et al. introduced the Transformer architecture, revolutionizing natural language processing and becoming the foundation for modern AI systems like GPT, BERT, and Claude. This article provides a comprehensive explanation of the attention mechanism with mathematical intuition and complete code implementation.

Part 1: Understanding the Attention Mechanism

The Intuition Behind Attention

When humans read a sentence, we don't process every word with equal focus. Our attention dynamically shifts based on context. The attention mechanism in Transformers mimics this behavior, allowing the model to focus on relevant parts of the input when producing each output.

Analogy: Information Retrieval

Imagine a library search system:

Query (查询): Your search question
Key (键): Book titles and index entries
Value (值): The actual book content

The attention mechanism computes similarity between your query and all keys, then retrieves values weighted by relevance.

Query, Key, and Value Explained

❓

Query (Q)

What am I looking for?

Represents the current token's request for information

🔑

Key (K)

What do I contain?

Represents each token's identifier for matching

📄

Value (V)

What information do I provide?

Represents the actual content to be retrieved

The Attention Formula

Scaled Dot-Product Attention

Attention(Q, K, V) = softmax(QK^T / √d_k) V

Q: Query matrix (sequence_length × d_k)

K: Key matrix (sequence_length × d_k)

V: Value matrix (sequence_length × d_v)

d_k: Dimension of keys (used for scaling)

Step-by-Step Breakdown

Compute Similarity Scores

QK^T: Matrix multiplication between Query and Key transposed. Each entry (i,j) represents how much token i should attend to token j.

Scale by √d_k

Prevents dot products from growing too large in magnitude, which would push softmax into regions with extremely small gradients.

Apply Softmax

Converts scores to probabilities that sum to 1. Higher scores get higher weight in the weighted sum.

Weighted Sum of Values

Multiply attention weights by Value matrix to get the final output representation.

Self-Attention Explained

Self-attention is a special case where Q, K, and V all come from the same input sequence. Each token attends to all other tokens in the same sequence, allowing the model to capture contextual relationships.

Why Self-Attention is Powerful

Long-range dependencies: Any token can directly attend to any other token
Parallel computation: All attention scores computed simultaneously
Interpretability: Attention weights show which tokens influence each other
Contextual representation: Each token's meaning depends on its context

Part 2: Complete Transformer Architecture

Transformer Structure Overview


Input Embedding
      ↓
Positional Encoding
      ↓
┌─────────────────────────────────────┐
│         Encoder (Nx layers)          │
│  ┌───────────────────────────────┐  │
│  │  Multi-Head Attention         │  │
│  │  + Add & Norm                 │  │
│  └───────────────────────────────┘  │
│              ↓                      │
│  ┌───────────────────────────────┐  │
│  │  Feed Forward                 │  │
│  │  + Add & Norm                 │  │
│  └───────────────────────────────┘  │
└─────────────────────────────────────┘
              ↓
┌─────────────────────────────────────┐
│         Decoder (Nx layers)          │
│  ┌───────────────────────────────┐  │
│  │  Masked Multi-Head Attention  │  │
│  │  + Add & Norm                 │  │
│  └───────────────────────────────┘  │
│              ↓                      │
│  ┌───────────────────────────────┐  │
│  │  Cross-Attention              │  │
│  │  + Add & Norm                 │  │
│  └───────────────────────────────┘  │
│              ↓                      │
│  ┌───────────────────────────────┐  │
│  │  Feed Forward                 │  │
│  │  + Add & Norm                 │  │
│  └───────────────────────────────┘  │
└─────────────────────────────────────┘
              ↓
        Linear + Softmax
              ↓
          Output

Multi-Head Attention

Instead of performing a single attention function, the Transformer uses multiple attention heads that learn different representation subspaces. This allows the model to jointly attend to information from different positions at different levels.

MultiHead(Q, K, V) = Concat(head_1, ..., head_h) W^O

where head_i = Attention(QW_i^Q, KW_i^K, VW_i^V)

Part 3: Code Implementation

Single Head Attention in PyTorch

import torch
import torch.nn as nn
import math

class SingleHeadAttention(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.d_model = d_model
        
        # Linear projections for Q, K, V
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        
    def forward(self, x, mask=None):
        """
        Args:
            x: Input tensor of shape (batch_size, seq_len, d_model)
            mask: Optional mask tensor
        Returns:
            output: Attention output of shape (batch_size, seq_len, d_model)
            attention_weights: Attention weights for visualization
        """
        batch_size, seq_len, d_model = x.shape
        
        # Step 1: Compute Q, K, V
        Q = self.W_q(x)  # (batch_size, seq_len, d_model)
        K = self.W_k(x)  # (batch_size, seq_len, d_model)
        V = self.W_v(x)  # (batch_size, seq_len, d_model)
        
        # Step 2: Compute attention scores
        # QK^T: (batch_size, seq_len, seq_len)
        scores = torch.matmul(Q, K.transpose(-2, -1))
        
        # Step 3: Scale by sqrt(d_model)
        scores = scores / math.sqrt(d_model)
        
        # Step 4: Apply mask (if provided)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        # Step 5: Apply softmax
        attention_weights = torch.softmax(scores, dim=-1)
        
        # Step 6: Weighted sum of values
        output = torch.matmul(attention_weights, V)
        
        return output, attention_weights


# Example usage
if __name__ == "__main__":
    # Parameters
    batch_size = 2
    seq_len = 4
    d_model = 8
    
    # Create random input
    x = torch.randn(batch_size, seq_len, d_model)
    
    # Initialize attention module
    attention = SingleHeadAttention(d_model)
    
    # Forward pass
    output, weights = attention(x)
    
    print(f"Input shape: {x.shape}")
    print(f"Output shape: {output.shape}")
    print(f"Attention weights shape: {weights.shape}")
    print(f"\nAttention weights (first sample):")
    print(weights[0].round(decimals=3))

Multi-Head Attention Implementation

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        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  # Dimension per head
        
        # Linear projections
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        
    def scaled_dot_product_attention(self, Q, K, V, mask=None):
        """
        Compute attention for all heads in parallel
        Q, K, V: (batch_size, num_heads, seq_len, d_k)
        """
        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
    
    def forward(self, x, mask=None):
        batch_size, seq_len, d_model = x.shape
        
        # Linear projections and reshape for multi-head
        # (batch_size, seq_len, d_model) -> (batch_size, num_heads, seq_len, d_k)
        Q = self.W_q(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        
        # Apply attention to all heads
        attn_output, weights = self.scaled_dot_product_attention(Q, K, V, mask)
        
        # Concatenate heads and apply final linear
        # (batch_size, num_heads, seq_len, d_k) -> (batch_size, seq_len, d_model)
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)
        output = self.W_o(attn_output)
        
        return output, weights


# Example usage
mha = MultiHeadAttention(d_model=512, num_heads=8)
x = torch.randn(2, 10, 512)  # (batch_size, seq_len, d_model)
output, weights = mha(x)
print(f"Multi-head output shape: {output.shape}")  # (2, 10, 512)
print(f"Attention weights shape: {weights.shape}")  # (2, 8, 10, 10)

Complete Transformer Encoder Layer

class TransformerEncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        
        # Multi-head attention
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        
        # Feed-forward network
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model)
        )
        
        # Layer normalization
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        
        # Dropout
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x, mask=None):
        # Self-attention with residual connection
        attn_output, _ = self.self_attn(x, mask)
        x = self.norm1(x + self.dropout(attn_output))
        
        # Feed-forward with residual connection
        ffn_output = self.ffn(x)
        x = self.norm2(x + self.dropout(ffn_output))
        
        return x

Part 4: Visualizing Attention

Attention Pattern Examples

Example: The cat sat on the mat

	The	cat	sat	on	the	mat
The	0.4	0.2	0.1	0.1	0.1	0.1
cat	0.2	0.3	0.3	0.1	0.05	0.05
sat	0.1	0.3	0.3	0.2	0.05	0.05
mat	0.05	0.05	0.1	0.2	0.3	0.3

Yellow cells show strong attention. Notice how 'sat' attends to both 'cat' (subject) and 'mat' (location), demonstrating contextual understanding.

Conclusion

The Transformer architecture, powered by the attention mechanism, has fundamentally changed how we approach sequence modeling. By understanding Query, Key, and Value, and implementing the attention formula Attention(Q,K,V) = softmax(QK^T/√d_k)V, we can build models that capture complex relationships in data.

Key takeaways from this article:

Attention is a weighted retrieval mechanism: Query searches for relevant Keys to retrieve Values
Scaling is crucial: Dividing by √d_k prevents softmax saturation
Multi-head attention enables diverse representations: Different heads learn different patterns
Self-attention captures context: Each token's representation depends on all other tokens

The code implementations provided here form the building blocks for understanding modern NLP systems. From this foundation, you can explore more advanced topics like positional encoding, masked attention for decoding, and the full encoder-decoder architecture used in machine translation and text generation.