h_t = f(h_{t-1}, x_t)

This sequential nature created fundamental problems:

Vanishing and Exploding Gradients:

Gradients must flow through many time steps during backpropagation. They tend to shrink (vanish) or grow (explode), making it difficult to learn long-range dependencies.

Limited Context:

Despite theoretical infinite memory, practical RNNs struggle to maintain information over long distances. The hidden state must compress all history into a fixed-size vector.

Sequential Computation:

Processing must happen step-by-step. This prevents parallelization and makes training slow on modern hardware.

The Attention Revolution

Attention mechanisms offered a solution: instead of compressing history into a fixed state, allow the model to look directly at all previous positions and attend to relevant ones.

Key Insight:

If predicting a word requires context from 100 positions ago, attention can directly access that position rather than hoping information survives through 100 recurrent steps.

The Transformer Leap:

Rather than using attention to supplement recurrence, Transformers eliminated recurrence entirely—using "self-attention" as the primary mechanism for processing sequences.

Transformer Architecture Overview

High-Level Structure

A Transformer consists of:

Encoder (for models like BERT):

• Processes input sequence
• Produces contextualized representations
• Bidirectional attention (can look forward and backward)

Decoder (for models like GPT):

• Generates output sequence
• Autoregressive (each token depends on previous tokens)
• Causal attention (can only look backward)

Encoder-Decoder (for models like T5, original Transformer):

• Encoder processes input
• Decoder generates output
• Cross-attention connects them

Many modern language models are decoder-only (GPT, Claude, LLaMA) or encoder-only (BERT for understanding tasks).

Component Stack

Each Transformer block contains:

Input

↓

[Multi-Head Self-Attention]

↓

[Add & Normalize]

↓

[Feed-Forward Network]

↓

[Add & Normalize]

↓

Output

These blocks stack multiple times (GPT-3 has 96 layers).

Embeddings and Positional Encoding

Token Embeddings

Before processing, text must become numbers:

Tokenization:

Text is split into tokens (words, subwords, or characters):

"Transformers are amazing" → ["Transform", "ers", " are", " amazing"]

Modern tokenizers (BPE, WordPiece, SentencePiece) balance vocabulary size and representation efficiency.

Embedding Lookup:

Each token maps to a learned vector:

# Simplified

embedding_matrix = nn.Embedding(vocab_size, d_model)

token_embeddings = embedding_matrix(token_ids) # (batch, seq_len, d_model)

Typical dimensions: 768 (BERT-base), 1024 (BERT-large), 4096 (GPT-3), 12288 (GPT-4).

Positional Encoding

Self-attention is permutation invariant—it doesn't inherently know position. Positional encodings inject sequence order information.

Original Sinusoidal Encoding:

def positional_encoding(max_len, d_model):

position = np.arange(max_len)[:, np.newaxis]

div_term = np.exp(np.arange(0, d_model, 2) * -(np.log(10000.0) / d_model))

pe = np.zeros((max_len, d_model))

pe[:, 0::2] = np.sin(position * div_term)

pe[:, 1::2] = np.cos(position * div_term)

return pe

The sinusoidal pattern allows the model to learn relative positions and generalize to longer sequences than seen during training.

Learned Positional Embeddings:

Many modern models (GPT-2, BERT) learn positional embeddings directly:

position_embedding = nn.Embedding(max_seq_len, d_model)

Rotary Position Embeddings (RoPE):

Used in LLaMA and many modern models, RoPE encodes position through rotation in the embedding space, providing better length generalization.

Combined Input:

input_embedding = token_embedding + positional_encoding

The Attention Mechanism

Intuition

Attention answers: "When processing this position, how much should I focus on each other position?"

For each position:

1. Create a query representing "what am I looking for?"
2. Create keys at all positions representing "what do I contain?"
3. Compare query to all keys to get attention weights
4. Use weights to aggregate values from all positions

Mathematical Definition

Scaled Dot-Product Attention:

Given queries Q, keys K, and values V:

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

Where:

• Q: Query matrix (seq_len × d_k)
• K: Key matrix (seq_len × d_k)
• V: Value matrix (seq_len × d_v)
• d_k: Key dimension
• √d_k: Scaling factor to prevent softmax saturation

Step-by-Step:

1. Compute attention scores:

scores = Q @ K.transpose(-2, -1) # (batch, seq_len, seq_len)

1. Scale:

scores = scores / math.sqrt(d_k)

1. Apply softmax:

attention_weights = softmax(scores, dim=-1)

1. Weight values:

output = attention_weights @ V # (batch, seq_len, d_v)

Why Scaling Matters

Without scaling, dot products grow with d_k. Large values push softmax into regions with extremely small gradients, harming training. The √d_k normalization keeps variance stable.

Masking

Padding Mask:

Prevent attention to padding tokens:

scores.masked_fill_(padding_mask == 0, -float('inf'))

Causal Mask (for decoders):

Prevent attending to future tokens:

causal_mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1)

scores.masked_fill_(causal_mask == 1, -float('inf'))

This ensures autoregressive generation—token i can only see tokens 1 to i.

Multi-Head Attention

Why Multiple Heads?

A single attention operation might learn a single type of relationship. Multiple heads can learn different relationship types:

• Head 1: Syntactic relationships
• Head 2: Semantic associations
• Head 3: Positional patterns
• Head 4: Entity co-references

Implementation

class MultiHeadAttention(nn.Module):

def __init__(self, d_model, num_heads):

super().__init__()

self.num_heads = num_heads

self.d_k = d_model // num_heads

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 forward(self, x, mask=None):

batch_size, seq_len, d_model = x.shape

# Linear projections

Q = self.W_q(x) # (batch, seq_len, d_model)

K = self.W_k(x)

V = self.W_v(x)

# Reshape to (batch, num_heads, seq_len, d_k)

Q = Q.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)

K = K.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)

V = V.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)

# Attention

scores = Q @ K.transpose(-2, -1) / math.sqrt(self.d_k)

if mask is not None:

scores.masked_fill_(mask == 0, -float('inf'))

attention = torch.softmax(scores, dim=-1)

# Apply attention to values

context = attention @ V # (batch, num_heads, seq_len, d_k)

# Concatenate heads

context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)

# Final projection

output = self.W_o(context)

return output

Typical configurations:

• BERT-base: 12 heads, d_model=768, d_k=64
• GPT-3: 96 heads, d_model=12288, d_k=128

Feed-Forward Networks

Purpose

After attention aggregates information across positions, the feed-forward network (FFN) processes each position independently with a deeper non-linear transformation.

Architecture

class FeedForward(nn.Module):

def __init__(self, d_model, d_ff, dropout=0.1):

super().__init__()

self.linear1 = nn.Linear(d_model, d_ff)

self.linear2 = nn.Linear(d_ff, d_model)

self.dropout = nn.Dropout(dropout)

self.activation = nn.GELU() # or ReLU in original

def forward(self, x):

x = self.linear1(x)

x = self.activation(x)

x = self.dropout(x)

x = self.linear2(x)

return x

Expansion Factor:

Typically d_ff = 4 × d_model. This expansion allows richer representations before projection back down.

Activation Function:

• Original: ReLU
• Modern: GELU (smoother, better gradient flow)
• Some models: SwiGLU (gated linear units)

Why FFN Matters

While attention handles inter-token relationships, FFN provides depth for intra-token processing. Studies show FFN layers often store factual knowledge—a form of key-value memory.

Normalization and Residual Connections

Residual Connections

Residual (skip) connections add the input to the output of each sublayer:

output = sublayer(x) + x

Benefits:

• Enables training very deep networks
• Provides gradient highways during backpropagation
• Allows layers to learn "refinements" rather than complete transformations

Layer Normalization

Unlike batch normalization (which normalizes across batches), layer normalization normalizes across features:

class LayerNorm(nn.Module):

def __init__(self, d_model, eps=1e-6):

super().__init__()

self.gamma = nn.Parameter(torch.ones(d_model))

self.beta = nn.Parameter(torch.zeros(d_model))

self.eps = eps

def forward(self, x):

mean = x.mean(dim=-1, keepdim=True)

std = x.std(dim=-1, keepdim=True)

return self.gamma * (x - mean) / (std + self.eps) + self.beta

Placement:

• Post-LN (original): Normalize after each sublayer
• Pre-LN (common now): Normalize before each sublayer

Pre-LN typically trains more stably for very deep models.

Combined Transformer Block

class TransformerBlock(nn.Module):

def __init__(self, d_model, num_heads, d_ff, dropout=0.1):

super().__init__()

self.attention = MultiHeadAttention(d_model, num_heads)

self.ffn = FeedForward(d_model, d_ff, dropout)

self.norm1 = nn.LayerNorm(d_model)

self.norm2 = nn.LayerNorm(d_model)

self.dropout = nn.Dropout(dropout)

def forward(self, x, mask=None):

# Pre-LN variant

attn_output = self.attention(self.norm1(x), mask)

x = x + self.dropout(attn_output)

ffn_output = self.ffn(self.norm2(x))

x = x + self.dropout(ffn_output)

return x

Complete Transformer Implementation

Decoder-Only Model (GPT-style)

class GPTModel(nn.Module):

def __init__(self, vocab_size, d_model, num_heads, num_layers,

max_seq_len, d_ff, dropout=0.1):

super().__init__()

self.token_embedding = nn.Embedding(vocab_size, d_model)

self.position_embedding = nn.Embedding(max_seq_len, d_model)

self.dropout = nn.Dropout(dropout)

self.blocks = nn.ModuleList([

TransformerBlock(d_model, num_heads, d_ff, dropout)

for _ in range(num_layers)

])

self.ln_f = nn.LayerNorm(d_model)

self.head = nn.Linear(d_model, vocab_size, bias=False)

# Tie weights (optional but common)

self.head.weight = self.token_embedding.weight

def forward(self, token_ids):

batch_size, seq_len = token_ids.shape

# Create causal mask

mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()

mask = mask.to(token_ids.device)

# Embeddings

positions = torch.arange(seq_len, device=token_ids.device).unsqueeze(0)

x = self.token_embedding(token_ids) + self.position_embedding(positions)

x = self.dropout(x)

# Transformer blocks

for block in self.blocks:

x = block(x, mask)

# Output

x = self.ln_f(x)

logits = self.head(x) # (batch, seq_len, vocab_size)

return logits

Training

def train_step(model, batch, optimizer, criterion):

model.train()

# Inputs and targets offset by 1

inputs = batch[:, :-1]

targets = batch[:, 1:]

# Forward pass

logits = model(inputs)

# Compute loss

loss = criterion(logits.view(-1, logits.size(-1)), targets.view(-1))

# Backward pass

optimizer.zero_grad()

loss.backward()

optimizer.step()

return loss.item()

Generation

@torch.no_grad()

def generate(model, prompt_ids, max_new_tokens, temperature=1.0, top_k=50):

model.eval()

generated = prompt_ids.clone()

for _ in range(max_new_tokens):

# Get logits for next token

logits = model(generated)[:, -1, :] # Last position

# Apply temperature

logits = logits / temperature

# Top-k filtering

if top_k > 0:

values, _ = torch.topk(logits, top_k)

min_value = values[:, -1].unsqueeze(-1)

logits = torch.where(logits < min_value,

torch.full_like(logits, -float('inf')),

logits)

# Sample

probs = torch.softmax(logits, dim=-1)

next_token = torch.multinomial(probs, num_samples=1)

# Append

generated = torch.cat([generated, next_token], dim=-1)

return generated

Advanced Concepts

Key-Value Caching

During generation, computing attention over the full sequence for each new token is wasteful. Key-Value caching stores previous K and V computations:

def forward_with_cache(self, x, past_kv=None):

# Compute Q, K, V for current input

Q = self.W_q(x)

K = self.W_k(x)

V = self.W_v(x)

if past_kv is not None:

past_K, past_V = past_kv

K = torch.cat([past_K, K], dim=1)

V = torch.cat([past_V, V], dim=1)

# Standard attention computation

# ...

return output, (K, V) # Return new cache

This reduces generation from O(n²) to O(n) per token.

Grouped Query Attention (GQA)

Standard multi-head attention has separate K and V for each head. GQA groups multiple Q heads to share fewer K/V heads:

• MHA: 32 Q heads, 32 K heads, 32 V heads
• GQA: 32 Q heads, 8 K heads, 8 V heads
• MQA: 32 Q heads, 1 K head, 1 V head

This dramatically reduces KV cache size with minimal quality loss.

Flash Attention

Standard attention materializes the full attention matrix, using O(n²) memory. Flash Attention fuses operations and uses tiling to keep memory O(n):

# Pseudocode for concept

def flash_attention(Q, K, V, block_size):

# Process in blocks, accumulating output

# Never materialize full n×n matrix

for q_block in partition(Q, block_size):

for kv_block in partition(K, V, block_size):

# Compute block attention with numerically stable accumulation

pass

Flash Attention enables much longer context lengths with less memory.

Mixture of Experts (MoE)

Instead of one large FFN, MoE routes each token to a subset of "expert" FFNs:

class MoELayer(nn.Module):

def __init__(self, d_model, num_experts, top_k):

self.gate = nn.Linear(d_model, num_experts)

self.experts = nn.ModuleList([FFN(d_model) for _ in range(num_experts)])

self.top_k = top_k

def forward(self, x):

# Route to top-k experts

gate_logits = self.gate(x)

weights, indices = torch.topk(gate_logits, self.top_k)

weights = torch.softmax(weights, dim=-1)

# Compute expert outputs (simplified)

output = sum(w * self.expertsi for w, i in zip(weights, indices))

return output

MoE allows massive parameter counts with manageable compute costs (e.g., Mixtral, GPT-4 rumored).

Transformer Variants

Encoder-Only (BERT-style)

For understanding tasks:

• Bidirectional attention (no causal mask)
• [CLS] token for classification
• Masked language modeling during training
• Use cases: Classification, NER, similarity

Decoder-Only (GPT-style)

For generation tasks:

• Causal attention mask
• Next-token prediction
• Autoregressive generation
• Use cases: Text generation, chat, code

Encoder-Decoder (T5-style)

For sequence-to-sequence:

• Encoder processes input
• Decoder attends to encoder (cross-attention) and previous outputs
• Use cases: Translation, summarization, question-answering

Vision Transformers (ViT)

Adapting Transformers for images:

• Split image into patches
• Treat patches as "tokens"
• Apply standard Transformer architecture
• Use cases: Image classification, object detection

Multimodal Transformers

Combining modalities:

• Shared or separate encoders for different modalities
• Cross-modal attention
• Examples: CLIP, GPT-4V, Gemini

Computational Considerations

Complexity Analysis

Attention:

• Time: O(n² · d)
• Memory: O(n² + n·d)

Feed-Forward:

• Time: O(n · d · d_ff)
• Memory: O(n · d_ff)

For very long sequences, attention dominates. This motivates:

• Sparse attention patterns
• Linear attention approximations
• Hierarchical approaches

Scaling Laws

Research has identified consistent relationships:

Compute-Optimal Training (Chinchilla):

N ∝ C^0.5 (parameters)

D ∝ C^0.5 (data tokens)

Equal scaling of model size and data is optimal.

Loss Scaling:

L ≈ A/N^α + B/D^β + L_∞