Transformer Architecture Deep Dive

Overview

In our previous lesson on RNNs, LSTMs, and GRUs, we explored the sequential approach to modeling language. While these architectures revolutionized NLP, they still suffered from fundamental limitations in handling long-range dependencies and parallelization.

This lesson introduces the Transformer architecture, a paradigm shift that replaced recurrence with attention mechanisms. First introduced in the seminal 2017 paper "Attention Is All You Need" by Vaswani et al., Transformers have become the foundation of modern NLP models like BERT, GPT, and T5 that have dramatically advanced the state of the art.

Learning Objectives

After completing this lesson, you will be able to:

• Understand the key innovations and motivations behind the Transformer architecture
• Explain self-attention and multi-head attention mechanisms in detail
• Describe positional encoding and why it's necessary
• Compare encoder-only, decoder-only, and encoder-decoder transformer variants
• Implement basic transformer components
• Recognize how transformers enable modern language models

The Need for a New Architecture

The Limitations of RNNs Revisited

As we saw in the previous lesson, RNNs and their variants face several critical limitations:

1. Sequential Processing: Processing tokens one at a time creates a bottleneck for training and inference
2. Limited Context Window: Even LSTMs struggle with very long-range dependencies
3. Vanishing Gradients: Despite improvements, still an issue for very long sequences

Analogy: Information Highways vs. Relay Races

Think of an RNN as a relay race where information is passed from one runner (time step) to the next. If the race is long, messages can get distorted or lost along the way, and the entire race is only as fast as the slowest runner.

In contrast, a Transformer is like a highway system where every location has direct high-speed connections to every other location. Information doesn't have to flow sequentially but can take direct routes, and all routes can be traveled simultaneously.

The Transformer Architecture: A High-Level View

Architectural Overview

Key Innovations

The Transformer introduced several groundbreaking innovations:

1. Self-Attention: Allows each position to directly attend to all positions
2. Multi-Head Attention: Enables attention across different representation subspaces
3. Positional Encoding: Captures sequence order without recurrence
4. Residual Connections + Layer Normalization: Facilitates training of deep networks
5. Feed-Forward Networks: Adds non-linearity and transforms representations
6. Parallel Processing: Enables efficient training and inference

Self-Attention: The Core Mechanism

Understanding Attention

Attention allows a model to focus on relevant parts of the input sequence when making predictions. It computes a weighted sum of values, where weights reflect the relevance of each value to the current context.

The Intuition Behind Self-Attention

In the example above, to understand what "it" refers to, the model must determine which previous words are most relevant. Self-attention allows the model to learn these relevance patterns.

Query, Key, Value (QKV) Framework

Self-attention can be conceptualized using the Query-Key-Value framework:

1. Query (Q): What we're looking for
2. Key (K): What we match against
3. Value (V): What we retrieve if there's a match

Think of it as a sophisticated dictionary lookup:

• The Query is like your search term
• The Keys are like the dictionary entries
• The Values are the definitions you retrieve

Self-Attention Computation: Step-by-Step

1. Projection: Generate Query, Key, and Value vectors by multiplying input embeddings by weight matrices $\mathbf{Q} = \mathbf{X}\mathbf{W}^Q, \mathbf{K} = \mathbf{X}\mathbf{W}^K, \mathbf{V} = \mathbf{X}\mathbf{W}^V$

2. Score Calculation: Compute attention scores by multiplying Q and K matrices $\text{Score} = \mathbf{Q}\mathbf{K}^T$

3. Scaling: Divide by square root of dimension to prevent extremely small gradients $\text{Score}_{\text{scaled}} = \frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{d_k}}$

4. Masking (Decoder Only): Apply mask to prevent attending to future positions (for decoder) $\text{Score}_{\text{masked}} = \text{Score}_{\text{scaled}} + \text{Mask}$

5. Softmax: Apply softmax to get probability distribution across values $\text{Attention Weights} = \text{softmax}(\text{Score}_{\text{scaled}})$

6. Weighted Sum: Multiply attention weights by values $\text{Attention Output} = \text{Attention Weights} \times \mathbf{V}$

Visualizing Self-Attention

Multi-Head Attention: Attending to Different Aspects

Why Multiple Attention Heads?

Self-attention with a single attention mechanism (or "head") can only capture one type of relationship between words. But language has many types of relationships (syntactic, semantic, referential, etc.).

Multiple attention heads allow the model to:

• Attend to different representation subspaces simultaneously
• Capture different types of dependencies (e.g., syntactic vs. semantic)
• Create a richer representation by combining these diverse perspectives

Multi-Head Attention Mechanism

💡 Tip: Use the attention patterns mode in the tool above to see how different heads capture different relationships.

Mathematical Formulation

For each head $i$: $\text{head}_i = \text{Attention}(\mathbf{X}\mathbf{W}_i^Q, \mathbf{X}\mathbf{W}_i^K, \mathbf{X}\mathbf{W}_i^V)$

The outputs from all heads are concatenated and linearly transformed: $\text{MultiHead}(\mathbf{X}) = \text{Concat}(\text{head}_1, \text{head}_2, ..., \text{head}_h)\mathbf{W}^O$

Analogy: Multiple Expert Consultants

Think of multi-head attention as consulting multiple experts who each focus on different aspects of a problem:

• One linguist focuses on grammar
• Another focuses on vocabulary
• A third focuses on cultural context
• A fourth focuses on tone

Each provides valuable insights from their perspective, and together they create a more comprehensive understanding than any single expert could provide.

Positional Encoding: Preserving Sequence Order

The Problem: Transformers Don't Know Position

Unlike RNNs, the self-attention mechanism is inherently permutation-invariant—it doesn't consider the order of tokens. This is a problem because word order is crucial in language understanding.

For example, these sentences have very different meanings despite using the same words:

• "The dog chased the cat"
• "The cat chased the dog"

Solution: Positional Encoding

Transformers add positional information to each word embedding using sinusoidal functions:

$PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right)$ $PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right)$

Where:

• $pos$ is the position
• $i$ is the dimension
• $d_{\text{model}}$ is the embedding dimension

Visualizing Positional Encoding

Key Properties of Sinusoidal Positional Encoding

1. Unique Pattern: Each position gets a unique encoding
2. Fixed Offset: The relative encoding between positions at a fixed offset is constant
3. Extrapolation: Can generalize to longer sequences than seen in training
4. No New Parameters: Unlike learned positional embeddings, requires no additional parameters

Embedding + Positional Encoding

The final input to the transformer is the sum of the word embeddings and the positional encodings:

$\text{Input} = \text{WordEmbedding} + \text{PositionalEncoding}$

Key Insight: The same positional encoding visualization above shows how these encodings combine with word embeddings to create the final input representation.

The Building Blocks: Encoder and Decoder

Transformer Encoder

The encoder processes the input sequence and consists of:

1. Multi-Head Self-Attention: Each position attends to all positions
2. Feed-Forward Neural Network: A two-layer network with ReLU activation
3. Residual Connections: Helps gradient flow and stabilizes training
4. Layer Normalization: Normalizes inputs to each sub-layer

Feed-Forward Network (FFN)

The FFN applies the same transformation to each position independently:

$\text{FFN}(x) = \max(0, x\mathbf{W}_1 + \mathbf{b}_1)\mathbf{W}_2 + \mathbf{b}_2$

This is equivalent to two dense layers with a ReLU activation in between. The FFN allows the model to transform its representations and introduces non-linearity.

Transformer Decoder

The decoder generates the output sequence and has three main components:

1. Masked Multi-Head Self-Attention: Each position attends only to previous positions
2. Cross-Attention: Attends to the encoder's output
3. Feed-Forward Neural Network: Same structure as in the encoder

Masking in the Decoder

The decoder must generate text autoregressively (one token at a time), so it can't "see" future tokens during training. This is achieved using a look-ahead mask.

💡 Tip: Use the self-attention mode in the tool above to see how masking prevents the decoder from attending to future positions.

Common pitfalls:

• Off-by-one masking: ensure strictly upper-triangular mask so position t cannot attend to ≥ t.
• Padding mask mixing: combine causal mask with key padding mask correctly to avoid leaking pads.

Cross-Attention

Cross-attention allows the decoder to focus on relevant parts of the input sequence:

The Full Architecture: Putting It All Together

Complete Transformer Architecture

Now that we've explored the individual components, let's see how they work together in the complete architecture. Use the overview mode in the tool above to see the full encoder-decoder stack, or explore other modes to dive deeper into specific mechanisms.

Training the Transformer

Transformers are typically trained with:

1. Teacher forcing: Using ground truth as decoder input during training
2. Label smoothing: Preventing overconfidence by softening the target distribution
3. Learning rate scheduling: Using warmup and decay for optimal convergence
4. Large batch sizes: Stabilizing training with more examples per update

Computational Complexity

The self-attention mechanism has quadratic complexity with respect to sequence length:

$\mathcal{O}(n^2 \cdot d)$

Where:

• $n$ is the sequence length
• $d$ is the representation dimension

This can be a limitation for very long sequences, leading to various efficient transformer variants that reduce this complexity.

Transformer Variants: Encoder-Only, Decoder-Only, and Encoder-Decoder

Transformer Variants: Architecture Comparison

The transformer architecture has evolved into three main variants, each optimized for different types of tasks:

Key Distinctions:

Encoder-Only Models (BERT, RoBERTa, DistilBERT)

• Bidirectional attention across all tokens
• Suitable for understanding tasks: classification, NER, sentiment analysis
• Cannot generate text autoregressively

Decoder-Only Models (GPT, GPT-2, GPT-3, GPT-4)

• Causal (masked) attention to prevent looking ahead
• Excellent for text generation and completion
• Can be adapted for understanding tasks with prompting

Encoder-Decoder Models (T5, BART, Pegasus)

• Best of both worlds: bidirectional encoding + autoregressive decoding
• Excel at sequence-to-sequence tasks: translation, summarization
• More complex but very versatile

Implementation: Building a Simple Transformer

Implementing Self-Attention in PyTorch

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

class SelfAttention(nn.Module):
    def __init__(self, embed_size, heads):
        super(SelfAttention, self).__init__()
        self.embed_size = embed_size
        self.heads = heads
Show all 45 lines

Implementing Positional Encoding

python
Copy
class PositionalEncoding(nn.Module):
    def __init__(self, embed_size, max_len=5000):
        super(PositionalEncoding, self).__init__()
        
        pe = torch.zeros(max_len, embed_size)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, embed_size, 2).float() * (-math.log(10000.0) / embed_size))
        
        # Apply sin to even indices
        pe[:, 0::2] = torch.sin(position * div_term)
Show all 19 lines

Transformer Encoder Layer

python
Copy
class TransformerEncoderLayer(nn.Module):
    def __init__(self, embed_size, heads, dropout, forward_expansion):
        super(TransformerEncoderLayer, self).__init__()
        
        self.attention = SelfAttention(embed_size, heads)
        self.norm1 = nn.LayerNorm(embed_size)
        self.norm2 = nn.LayerNorm(embed_size)
        
        self.feed_forward = nn.Sequential(
            nn.Linear(embed_size, forward_expansion * embed_size),
Show all 26 lines

Transformer Decoder Layer

python
Copy
class TransformerDecoderLayer(nn.Module):
    def __init__(self, embed_size, heads, dropout, forward_expansion):
        super(TransformerDecoderLayer, self).__init__()
        
        self.attention = SelfAttention(embed_size, heads)
        self.cross_attention = SelfAttention(embed_size, heads)
        self.norm1 = nn.LayerNorm(embed_size)
        self.norm2 = nn.LayerNorm(embed_size)
        self.norm3 = nn.LayerNorm(embed_size)
        
Show all 34 lines

Applications: How Transformers Revolutionized NLP

Machine Translation

The original transformer model was designed for machine translation and significantly improved the state of the art on the WMT English-to-German and English-to-French translation tasks.

Qualitative Comparison

The machine translation improvements can be seen in the model comparison tool we used earlier. Switch to the "Model Comparison" mode in any of the TransformerExplorer tools above to compare translation quality across different architectures.

Language Modeling and Text Generation

Transformer-based language models like GPT can generate remarkably coherent and contextually appropriate text.

Code Example: Text Generation with a Simple Transformer

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

# Simple GPT-like model for text generation
class SimpleGPT(nn.Module):
    def __init__(self, vocab_size, embed_dim, num_heads, num_layers, max_seq_len, dropout=0.1):
        super().__init__()
        
Show all 99 lines

Bidirectional Understanding and Masked Language Modeling

BERT and its variants use transformer encoders with masked language modeling to develop bidirectional understanding of text.

Code Example: Masked Language Modeling with Transformer Encoder

python
Copy
import torch
import torch.nn as nn
import torch.nn.functional as F
import random

# BERT-like model for masked language modeling
class SimpleBERT(nn.Module):
    def __init__(self, vocab_size, embed_dim, num_heads, num_layers, max_seq_len, dropout=0.1):
        super().__init__()
        
Show all 141 lines

Limitations and Future Directions

Current Limitations

1. Quadratic Complexity: Self-attention scales poorly with sequence length
2. Context Window: Limited by training and architecture constraints
3. Interpretability: Understanding attention patterns isn't straightforward
4. Data Hunger: Requires massive amounts of data for best performance
5. Compute Resources: Training large models requires significant resources

Efficient Transformer Variants

Researchers are developing efficient transformers to address these limitations:

Note: Performance is relative to vanilla Transformer of similar size.

The Future of Transformers

Transformers continue to evolve in several exciting directions:

1. Multimodal Transformers: Processing text, images, audio, and video together
2. Domain-Specific Architectures: Specialized for specific fields (science, medicine)
3. Mixture of Experts: Using sparse activation to scale to trillions of parameters
4. Retrieval-Augmented Models: Enhancing LLMs with external knowledge access
5. More Efficient Attention: Continuing to reduce the quadratic complexity

Conclusion: The Foundation of Modern NLP

The transformer architecture represents one of the most significant breakthroughs in natural language processing. By introducing self-attention, positional encoding, and parallel processing, transformers solved the fundamental limitations of sequential models while enabling the creation of increasingly powerful language models.

Key innovations of the transformer:

• Self-attention: Direct modeling of relationships between all sequence positions
• Parallel processing: Elimination of sequential dependencies for faster training
• Scalability: Architecture that grows effectively with more data and compute
• Versatility: Success across numerous NLP tasks and domains

The transformer has become the foundation for modern language models like BERT, GPT, T5, and their successors. Understanding this architecture is essential for working with contemporary NLP systems.

In our next lessons, we'll explore how transformer architectures evolved into the powerful language models of today, including the deterministic and probabilistic methods for text generation, and then survey the modern landscape of language models from BERT and GPT to the latest innovations like Llama 3 and Claude 3.

Practice Exercises

1. Implement Self-Attention:

   • Write a simplified version of the self-attention mechanism
   • Visualize attention weights for a sample sentence
   • Experiment with different scaling factors

2. Positional Encoding Analysis:

   • Implement sinusoidal positional encoding
   • Analyze how different positions are represented
   • Visualize positional encoding vectors

3. Transformer Architecture Comparison:

   • Compare performance of RNN vs. Transformer on a simple task
   • Measure inference time for both architectures
   • Analyze computational complexity at different sequence lengths

4. Pre-trained Model Exploration:

   • Fine-tune a small pre-trained transformer for a classification task
   • Analyze attention patterns in different heads
   • Experiment with different layer freezing strategies

Additional Resources

• Attention Is All You Need - The original transformer paper
• The Illustrated Transformer by Jay Alammar
• Transfomer detailed visualization and explanation
• The Annotated Transformer - Implementation walkthrough
• Transformers from Scratch by Peter Bloem
• Hugging Face Transformers - Library with implementations of transformer models
• A Survey of Long-Term Context in Transformers - Overview of efficient transformer variants