Attention Mechanisms Deep Dive (Part 2: Patterns, Implementation & Modern Optimizations)
Attention Mechanisms in Large Language Models: A Deep Dive (Part 2)
Continuation: Patterns, Implementation, and Modern Optimizations
Part 6: Attention Patterns and What They Learn
Induction Heads and In-Context Learning
One of the most important discoveries in mechanistic interpretability is the induction head—a circuit that enables in-context learning.
An induction head performs this operation:
- "If I've seen pattern [A][B] before, and I now see [A] again, predict [B]"
For example, if the context contains "Harry Potter" and later "Harry" appears, the induction head helps predict "Potter."
Induction heads typically require two attention heads working together:
-
Previous token head (in an earlier layer): Copies information from each token to the next position. When processing "Harry," it looks at "Harry" and writes information about what followed (Potter).
-
Induction head (in a later layer): When seeing "Harry" again, it matches this against keys that contain "what followed Harry before" and retrieves "Potter."
┌─────────────────────────────────────────────────────────────┐
│ Induction Head Circuit │
│ │
│ Context: "Harry Potter went to... Harry [???]" │
│ │
│ Step 1 (Early layer - Previous Token Head): │
│ At "Potter", attend to "Harry", write info about │
│ "Potter follows Harry" │
│ │
│ Step 2 (Later layer - Induction Head): │
│ At second "Harry", look for "what came after Harry" │
│ Find: "Potter" │
│ Predict: "Potter" │
└─────────────────────────────────────────────────────────────┘
Previous Token Heads
Some attention heads learn to simply attend to the immediately preceding token. This seems simple but is crucial for:
- Bigram statistics (common word pairs)
- Providing the induction head circuit with shifted information
- Local syntactic patterns
Positional Attention Patterns
Some heads learn to attend based primarily on position:
- Beginning-of-sequence heads: Attend to the first token (often special tokens like BOS)
- Fixed-offset heads: Attend to position i-k for some fixed k
- Diagonal heads: Attend primarily to the same position (self-attention in the literal sense)
These provide the model with consistent positional structure.
How Layers Specialize
Research has shown a rough progression of what layers focus on:
Early layers (1-4):
- Basic token identity
- Local patterns and syntax
- Previous token relationships
Middle layers (5-20):
- Semantic grouping
- Entity tracking
- Syntactic tree structure
- Induction patterns
Late layers (20+):
- Task-specific processing
- Output formatting
- Final prediction refinement
This isn't rigid—there's significant overlap and variation across models—but it provides intuition for how transformers build up understanding.
Specific Syntactic Heads: What Research Has Found
Researchers have identified specific attention heads that encode grammatical relationships. Here are documented examples:
┌────────────────────────────────────────────────────────────────────┐
│ SYNTACTIC HEAD CATALOG │
│ (Examples from BERT/GPT research) │
├────────────────────────────────────────────────────────────────────┤
│ │
│ SUBJECT-VERB HEADS │
│ ────────────────── │
│ "The cat [that I saw] sleeps" │
│ ↑ attends strongly to "cat" (not "I") │
│ Function: Track subject across long distances for agreement │
│ │
│ DIRECT OBJECT HEADS │
│ ─────────────────── │
│ "She gave [him] [the book]" │
│ ↑─────┘ │
│ Function: Link verbs to their objects │
│ │
│ DETERMINER-NOUN HEADS │
│ ───────────────────── │
│ "the big red [ball]" │
│ ↑───────────────┘ │
│ Function: Connect articles/adjectives to their nouns │
│ │
│ COREFERENCE HEADS │
│ ───────────────── │
│ "John said [he] would come" │
│ ↑──────┘ │
│ Function: Resolve pronouns to their antecedents │
│ │
│ RARE/REFLEXIVE HEADS │
│ ──────────────────── │
│ "The cat licked [itself]" │
│ ↑────────────┘ │
│ Function: Connect reflexive pronouns to subjects │
│ │
└────────────────────────────────────────────────────────────────────┘
Key Research Papers:
- Clark et al. (2019): "What Does BERT Look At?" — Analyzed attention patterns across BERT layers
- Voita et al. (2019): "Analyzing Multi-Head Self-Attention" — Identified syntactic, positional, and rare-token heads
- Htut et al. (2019): "Do Attention Heads in BERT Track Syntactic Dependencies?" — Systematic evaluation of syntactic encoding
Important caveat: These patterns are tendencies, not guarantees. A "subject-verb" head might attend to subjects 60-70% of the time for relevant constructions—not 100%.
How Syntax Gets Encoded in Weights
Q: How do weight matrices learn to encode syntax?
A: Through geometric structure in Q/K space. The projections create regions where syntactic roles cluster.
┌────────────────────────────────────────────────────────────────────┐
│ GEOMETRIC ENCODING OF SYNTAX │
├────────────────────────────────────────────────────────────────────┤
│ │
│ INTUITION: W_Q and W_K learn complementary projections │
│ │
│ For a subject-verb head: │
│ │
│ W_K might project: │
│ • Nouns → region A (potential subjects) │
│ • Verbs → region B │
│ • Others → scattered │
│ │
│ W_Q might project: │
│ • Verbs → queries that match region A │
│ • (Verbs "look for" nouns) │
│ │
│ K-space: Q-space: │
│ │
│ ┌──────────┐ ┌──────────┐ │
│ │ •N │ │ │ │
│ │ •N •N │ │ •V→ │ ← Verb queries │
│ │ •N │ │ •V→ │ point toward │
│ │ │ │ •V→ │ noun region │
│ │ •V •V │ │ │ │
│ │ •V │ │ │ │
│ └──────────┘ └──────────┘ │
│ │
│ Q_verb · K_noun = HIGH (verbs attend to nouns) │
│ Q_verb · K_verb = LOW (verbs don't self-attend as much) │
│ │
├────────────────────────────────────────────────────────────────────┤
│ │
│ THE LEARNING PROCESS: │
│ │
│ 1. Random init: Q, K projections point in random directions │
│ │
│ 2. Training signal: "sleeps" should predict based on "cat" │
│ not "I" → gradient pushes verbs to attend to subjects │
│ │
│ 3. Convergence: W_Q, W_K settle into complementary structure │
│ where syntactic relationships produce high dot products │
│ │
│ 4. Result: The geometry of Q/K space implicitly encodes grammar │
│ │
└────────────────────────────────────────────────────────────────────┘
Analogy: It's like learning that questions (Q) and answers (K) should match. The model learns that certain types of questions (verb queries) match certain types of answers (noun keys).
Why this works without explicit grammar labels:
- The training objective is next-token prediction
- Getting "sleeps" right requires tracking "cat"
- Gradient descent finds that organizing Q/K by syntactic role helps
- Grammar emerges as a useful structure for prediction
Part 7: Putting It All Together
Full Forward Pass Through Attention
Let's trace a complete forward pass for the sentence "The cat sat" through one transformer layer.
Step 0: Input Preparation
Token IDs: [1, 42, 87] (for "The", "cat", "sat")
Embeddings: Look up in embedding table → 3 vectors of dim 512
+ Positions: Add RoPE rotation (or add positional embeddings)
→ Input X: Shape [3, 512]
Step 1: Layer Normalization
X_norm = LayerNorm(X)
Normalizes each position's vector to have zero mean and unit variance.
Step 2: Q, K, V Projections (for all 8 heads)
For each head h (in parallel):
Q_h = X_norm · W_Q_h # Shape [3, 64]
K_h = X_norm · W_K_h # Shape [3, 64]
V_h = X_norm · W_V_h # Shape [3, 64]
Step 3: Apply RoPE to Q and K
Q_h_rotated = apply_rope(Q_h, positions)
K_h_rotated = apply_rope(K_h, positions)
Step 4: Attention Scores
For each head h:
scores_h = Q_h_rotated · K_h_rotated.T / √64
# Shape [3, 3] - each position attends to each position
# Apply causal mask (for autoregressive models)
scores_h = scores_h + causal_mask # -inf for future positions
Step 5: Softmax → Attention Weights
For each head h:
weights_h = softmax(scores_h, dim=-1)
# Shape [3, 3], rows sum to 1
# Example for position 2 ("sat"):
# weights might be [0.2, 0.5, 0.3]
# meaning: 20% attention to "The", 50% to "cat", 30% to "sat"
Step 6: Weighted Sum of Values
For each head h:
attn_output_h = weights_h · V_h
# Shape [3, 64]
Step 7: Concatenate Heads and Project
attn_output = Concat(attn_output_1, ..., attn_output_8) · W_O
# Shape [3, 512]
Step 8: Residual Connection
X = X + attn_output # Add attention output to residual stream
Step 9: MLP Block (simplified)
X_norm2 = LayerNorm(X)
mlp_output = MLP(X_norm2) # Typically: Linear → GELU → Linear
X = X + mlp_output # Add MLP output to residual stream
The residual stream X now contains a richer representation, ready for the next layer.
How Information Flows from Input to Output
┌────────────────────────────────────────────────────────────────┐
│ Information Flow │
│ │
│ "The cat sat [PREDICT NEXT]" │
│ │
│ Layer 0: Embeddings + Positions │
│ └→ Basic token identity established │
│ │
│ Layers 1-3: Local patterns │
│ └→ "The" modifies "cat", "cat" is subject │
│ └→ Previous token info propagated │
│ │
│ Layers 4-10: Semantic building │
│ └→ "cat" is an animal, "sat" is past tense action │
│ └→ Subject-verb relationship encoded │
│ │
│ Layers 11-20: Compositional meaning │
│ └→ "A cat performed the action of sitting" │
│ └→ Likely completions narrowed based on meaning │
│ │
│ Layers 21-32: Task completion │
│ └→ Output probability distribution refined │
│ └→ "on", "down", "there" likely next tokens │
│ │
│ Final: Predict "on" as most likely next token │
└────────────────────────────────────────────────────────────────┘
The Role of Attention in "Moving Information"
A crucial insight: attention moves information between positions.
In "The cat that I saw yesterday was sleeping":
- Position of "was" needs to know the subject is "cat" (singular)
- Attention heads can attend from "was" → "cat"
- They copy/transform relevant information (singular noun, animal, etc.)
- This information influences the prediction that "was" (not "were") is correct
Each attention layer gives tokens another chance to gather information from other positions. With 32 layers, there are 32 opportunities for each token to look at and incorporate information from elsewhere.
Information doesn't just flow left-to-right—through attention, even early tokens get updated based on later tokens (in bidirectional models), and even in causal models, the representation of "The" at layer 20 contains information about "cat" and "sat" because those tokens have written to the residual stream.
Complete Layer-by-Layer Shape Trace
Here's every shape transformation through a full transformer layer:
┌──────────────────────────────────────────────────────────────────────┐
│ COMPLETE SHAPE TRACE: ONE TRANSFORMER LAYER │
│ (GPT-2 Small: d_model=768, n_heads=12, d_head=64) │
├──────────────────────────────────────────────────────────────────────┤
│ │
│ INPUT: X [batch=1, seq=512, d_model=768] │
│ │
│ ═══════════════════════════════════════════════════════════════════ │
│ ATTENTION BLOCK │
│ ═══════════════════════════════════════════════════════════════════ │
│ │
│ 1. Pre-LayerNorm │
│ ────────────── │
│ X_norm = LayerNorm(X) │
│ Shape: [1, 512, 768] → [1, 512, 768] │
│ Params: γ [768] + β [768] = 1,536 │
│ │
│ 2. Q, K, V Projections (all heads computed together) │
│ ───────────────────────────────────────────────── │
│ Q_all = X_norm @ W_Q │
│ Shape: [1, 512, 768] @ [768, 768] → [1, 512, 768] │
│ Params: W_Q [768, 768] = 589,824 │
│ │
│ K_all = X_norm @ W_K │
│ Shape: [1, 512, 768] @ [768, 768] → [1, 512, 768] │
│ Params: W_K [768, 768] = 589,824 │
│ │
│ V_all = X_norm @ W_V │
│ Shape: [1, 512, 768] @ [768, 768] → [1, 512, 768] │
│ Params: W_V [768, 768] = 589,824 │
│ │
│ 3. Reshape for Multi-Head │
│ ────────────────────── │
│ Q = Q_all.reshape(1, 512, 12, 64).transpose(1,2) │
│ Shape: [1, 512, 768] → [1, 12, 512, 64] │
│ (batch, heads, seq, d_head) │
│ │
│ K, V same reshape → [1, 12, 512, 64] each │
│ │
│ 4. Attention Scores │
│ ───────────────── │
│ scores = Q @ K.transpose(-2, -1) / √64 │
│ Shape: [1, 12, 512, 64] @ [1, 12, 64, 512] → [1, 12, 512, 512] │
│ │
│ 5. Causal Mask + Softmax │
│ ───────────────────── │
│ scores = scores + causal_mask # -inf for future positions │
│ weights = softmax(scores, dim=-1) │
│ Shape: [1, 12, 512, 512] │
│ │
│ 6. Weighted Sum │
│ ──────────── │
│ attn_out = weights @ V │
│ Shape: [1, 12, 512, 512] @ [1, 12, 512, 64] → [1, 12, 512, 64] │
│ │
│ 7. Reshape + Output Projection │
│ ─────────────────────────── │
│ attn_out = attn_out.transpose(1,2).reshape(1, 512, 768) │
│ attn_out = attn_out @ W_O │
│ Shape: [1, 512, 768] @ [768, 768] → [1, 512, 768] │
│ Params: W_O [768, 768] = 589,824 │
│ │
│ 8. Residual Connection │
│ ──────────────────── │
│ X = X + attn_out │
│ Shape: [1, 512, 768] │
│ │
│ ═══════════════════════════════════════════════════════════════════ │
│ MLP BLOCK │
│ ═══════════════════════════════════════════════════════════════════ │
│ │
│ 9. Pre-LayerNorm │
│ ────────────── │
│ X_norm = LayerNorm(X) │
│ Shape: [1, 512, 768] │
│ Params: γ [768] + β [768] = 1,536 │
│ │
│ 10. MLP Expansion (fc1) │
│ ──────────────────── │
│ hidden = X_norm @ W_fc1 + b_fc1 │
│ Shape: [1, 512, 768] @ [768, 3072] → [1, 512, 3072] │
│ Params: W_fc1 [768, 3072] + b_fc1 [3072] = 2,362,368 │
│ │
│ 11. Activation (GELU) │
│ ───────────────── │
│ hidden = GELU(hidden) │
│ Shape: [1, 512, 3072] (unchanged) │
│ │
│ 12. MLP Contraction (fc2) │
│ ────────────────────── │
│ mlp_out = hidden @ W_fc2 + b_fc2 │
│ Shape: [1, 512, 3072] @ [3072, 768] → [1, 512, 768] │
│ Params: W_fc2 [3072, 768] + b_fc2 [768] = 2,360,064 │
│ │
│ 13. Residual Connection │
│ ──────────────────── │
│ X = X + mlp_out │
│ Shape: [1, 512, 768] │
│ │
│ OUTPUT: X [1, 512, 768] │
│ │
└──────────────────────────────────────────────────────────────────────┘
Detailed Parameter Count: GPT-2 Small
┌────────────────────────────────────────────────────────────────────┐
│ GPT-2 SMALL PARAMETER BREAKDOWN │
│ (124M parameters total) │
├────────────────────────────────────────────────────────────────────┤
│ │
│ EMBEDDINGS │
│ ─────────── │
│ Token embeddings: 50257 × 768 = 38,597,376 (31.1%) │
│ Position embeddings: 1024 × 768 = 786,432 ( 0.6%) │
│ ─────────────────────────────────────────────────────── │
│ Embedding total: 39,383,808 (31.8%) │
│ │
│ PER LAYER (×12 layers) │
│ ───────────────────── │
│ Attention: │
│ W_Q: 768 × 768 = 589,824 │
│ W_K: 768 × 768 = 589,824 │
│ W_V: 768 × 768 = 589,824 │
│ W_O: 768 × 768 = 589,824 │
│ LayerNorm (pre): 768 × 2 = 1,536 │
│ ──────────────────────────────────────────────── │
│ Attention subtotal: 2,360,832 (1.9%) │
│ │
│ MLP: │
│ W_fc1: 768 × 3072 = 2,359,296 │
│ b_fc1: 3072 = 3,072 │
│ W_fc2: 3072 × 768 = 2,359,296 │
│ b_fc2: 768 = 768 │
│ LayerNorm (pre): 768 × 2 = 1,536 │
│ ──────────────────────────────────────────────── │
│ MLP subtotal: 4,723,968 (3.8%) │
│ │
│ Per layer total: 7,084,800 (5.7%) │
│ ×12 layers: 85,017,600 (68.5%) │
│ │
│ FINAL LAYER │
│ ─────────── │
│ Final LayerNorm: 768 × 2 = 1,536 ( 0.0%) │
│ LM head (tied): (shares token embeddings) ( 0 ) │
│ │
│ ═══════════════════════════════════════════════════════════════ │
│ GRAND TOTAL: 124,402,944 │
│ ═══════════════════════════════════════════════════════════════ │
│ │
│ BREAKDOWN BY TYPE: │
│ Embeddings: 39.4M (31.7%) │
│ Attention (all): 28.3M (22.8%) │
│ MLP (all): 56.7M (45.5%) │
│ LayerNorms: 0.04M ( 0.0%) │
│ │
└────────────────────────────────────────────────────────────────────┘
Simplified Forward Pass in Python
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
class SimplifiedAttention(nn.Module):
def __init__(self, d_model=768, n_heads=12):
super().__init__()
self.d_model = d_model
self.n_heads = n_heads
self.d_head = d_model // n_heads
# Projections
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)
self.W_O = nn.Linear(d_model, d_model, bias=False)
def forward(self, x):
batch, seq_len, _ = x.shape
# Project to Q, K, V
Q = self.W_Q(x) # [batch, seq, d_model]
K = self.W_K(x)
V = self.W_V(x)
# Reshape for multi-head: [batch, n_heads, seq, d_head]
Q = Q.view(batch, seq_len, self.n_heads, self.d_head).transpose(1, 2)
K = K.view(batch, seq_len, self.n_heads, self.d_head).transpose(1, 2)
V = V.view(batch, seq_len, self.n_heads, self.d_head).transpose(1, 2)
# Attention scores: [batch, n_heads, seq, seq]
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_head)
# Causal mask
mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
scores = scores.masked_fill(mask, float('-inf'))
# Softmax → weights
weights = F.softmax(scores, dim=-1)
# Weighted sum of values
attn_out = torch.matmul(weights, V) # [batch, n_heads, seq, d_head]
# Reshape back: [batch, seq, d_model]
attn_out = attn_out.transpose(1, 2).contiguous().view(batch, seq_len, self.d_model)
# Output projection
return self.W_O(attn_out)
class SimplifiedTransformerLayer(nn.Module):
def __init__(self, d_model=768, n_heads=12, d_ff=3072):
super().__init__()
self.ln1 = nn.LayerNorm(d_model)
self.attn = SimplifiedAttention(d_model, n_heads)
self.ln2 = nn.LayerNorm(d_model)
self.mlp = nn.Sequential(
nn.Linear(d_model, d_ff),
nn.GELU(),
nn.Linear(d_ff, d_model)
)
def forward(self, x):
# Pre-norm attention with residual
x = x + self.attn(self.ln1(x))
# Pre-norm MLP with residual
x = x + self.mlp(self.ln2(x))
return x
# Usage example:
# x = torch.randn(1, 512, 768) # [batch, seq, d_model]
# layer = SimplifiedTransformerLayer()
# output = layer(x) # [1, 512, 768]
This code demonstrates the core mechanics without optimizations like Flash Attention or fused kernels.
Part 8: Advanced Topics
Causal Masking in Decoder-Only Models
Modern LLMs like GPT and LLaMA are decoder-only models that generate text left-to-right. During training, they learn to predict the next token given only previous tokens.
Causal masking enforces this constraint in attention. Position i can only attend to positions ≤ i:
Attention Mask (for sequence length 4):
pos 0 pos 1 pos 2 pos 3
pos 0 [ 1 0 0 0 ]
pos 1 [ 1 1 0 0 ]
pos 2 [ 1 1 1 0 ]
pos 3 [ 1 1 1 1 ]
1 = can attend, 0 = masked (set to -∞ before softmax)
This ensures the model can't "cheat" by looking at future tokens during training, and enables efficient parallel training while maintaining the autoregressive property.
KV-Cache for Efficient Inference
During generation, we produce one token at a time. Naively, each new token would require recomputing attention over the entire sequence—wasteful since previous tokens' K and V don't change.
KV-cache stores the K and V vectors from previous positions:
Step 1: Generate token 1
- Compute K₁, V₁
- Store in cache: K_cache = [K₁], V_cache = [V₁]
Step 2: Generate token 2
- Compute K₂, V₂
- Retrieve from cache: K_cache = [K₁], V_cache = [V₁]
- Append: K_cache = [K₁, K₂], V_cache = [V₁, V₂]
- Attention only needs new Q₂
Step 3: Generate token 3
- Compute K₃, V₃
- Retrieve: K_cache = [K₁, K₂], V_cache = [V₁, V₂]
- Append: K_cache = [K₁, K₂, K₃], etc.
This reduces generation from O(n²) to O(n) in computation per token, at the cost of O(n) memory for the cache.
Memory challenge: For long contexts (100K+ tokens) with many layers and heads, KV-cache can require tens of gigabytes of GPU memory.
Flash Attention and Memory Optimization
Standard attention has a memory problem: the attention matrix is [sequence_length × sequence_length]. For 100K tokens, that's 10 billion elements per head—too large to fit in fast GPU memory.
Flash Attention (Dao et al., 2022) is an algorithm that:
- Never materializes the full attention matrix
- Processes attention in blocks, keeping data in fast SRAM
- Is mathematically equivalent to standard attention, just computed differently
Standard Attention Memory: O(n²)
Flash Attention Memory: O(n) [linear in sequence length]
Flash Attention achieves 2-4x speedup and enables much longer context lengths without running out of memory.
Sparse Attention Variants
For very long sequences, even O(n²) computation (not just memory) becomes prohibitive. Sparse attention reduces this by having each token attend to only a subset of other tokens.
Common patterns:
-
Local attention: Each token attends only to a window of nearby tokens
Token 50 attends to tokens 40-60 only -
Strided attention: Attend to every k-th token
Token 100 attends to tokens 0, 10, 20, 30, ... -
Global attention: Certain tokens (like [CLS] or special markers) attend to everything
[CLS] attends to all tokens; all tokens attend to [CLS] -
Learned sparse patterns: Let the model learn which positions to attend to (Routing Transformers, etc.)
Models like Longformer, BigBird, and Sparse Transformers combine these patterns to achieve O(n) or O(n log n) complexity while maintaining most of the model's capability.
Deep Dive: KV-Cache Memory Costs
Understanding KV-cache memory requirements is crucial for deployment:
┌────────────────────────────────────────────────────────────────────┐
│ KV-CACHE MEMORY CALCULATION │
├────────────────────────────────────────────────────────────────────┤
│ │
│ For each token in context, we store: │
│ K vector: [d_head] per head per layer │
│ V vector: [d_head] per head per layer │
│ │
│ FORMULA: │
│ Memory per token = 2 × n_layers × n_heads × d_head × bytes │
│ = 2 × n_layers × d_model × bytes │
│ │
│ EXAMPLE: LLaMA 7B │
│ ───────────────── │
│ n_layers = 32, n_heads = 32, d_head = 128, d_model = 4096 │
│ Using FP16 (2 bytes per float): │
│ │
│ Per token: 2 × 32 × 4096 × 2 bytes = 524,288 bytes = 0.5 MB │
│ │
│ For different context lengths: │
│ 2K tokens: 0.5 × 2048 = 1 GB │
│ 8K tokens: 0.5 × 8192 = 4 GB │
│ 32K tokens: 0.5 × 32768 = 16 GB │
│ 128K tokens: 0.5 × 131072 = 64 GB │
│ │
│ EXAMPLE: LLaMA 70B │
│ ────────────────── │
│ n_layers = 80, d_model = 8192 │
│ │
│ Per token: 2 × 80 × 8192 × 2 bytes = 2.6 MB │
│ │
│ For 32K context: 2.6 × 32768 = 85 GB (just for KV-cache!) │
│ │
└────────────────────────────────────────────────────────────────────┘
This is why KV-cache optimization (GQA, quantization) is critical for long contexts.
Flash Attention: How It Works
Flash Attention achieves O(n) memory by never materializing the full attention matrix.
┌────────────────────────────────────────────────────────────────────┐
│ FLASH ATTENTION ALGORITHM │
├────────────────────────────────────────────────────────────────────┤
│ │
│ STANDARD ATTENTION: │
│ ─────────────────── │
│ 1. Compute full S = Q @ K.T [seq × seq] ← HUGE! │
│ 2. Apply mask to S │
│ 3. Compute P = softmax(S) [seq × seq] ← HUGE! │
│ 4. Compute O = P @ V │
│ Memory: O(n²) │
│ │
│ FLASH ATTENTION (block-wise): │
│ ───────────────────────────── │
│ Process Q in blocks, K/V in blocks: │
│ │
│ for q_block in Q_blocks: # Iterate over Q │
│ # Accumulate output for this Q block │
│ O_block = 0 │
│ running_max = -inf │
│ running_sum = 0 │
│ │
│ for kv_block in KV_blocks: # Iterate over K, V │
│ # Compute attention for this block pair │
│ S_block = q_block @ k_block.T # Small matrix! │
│ S_block = S_block / √d_k │
│ │
│ # Online softmax: update running statistics │
│ block_max = max(S_block) │
│ new_max = max(running_max, block_max) │
│ │
│ # Rescale previous accumulator │
│ scale = exp(running_max - new_max) │
│ O_block = O_block * scale │
│ running_sum = running_sum * scale │
│ │
│ # Add contribution from this block │
│ P_block = exp(S_block - new_max) │
│ O_block += P_block @ v_block │
│ running_sum += sum(P_block) │
│ │
│ running_max = new_max │
│ │
│ # Final normalization │
│ O_block = O_block / running_sum │
│ │
│ KEY INSIGHT: Never store full [seq × seq] matrix! │
│ Memory: O(n) - linear in sequence length │
│ │
│ MEMORY HIERARCHY: │
│ ───────────────── │
│ HBM (GPU main memory): High capacity, slow │
│ SRAM (on-chip): Low capacity, fast │
│ │
│ Flash Attention keeps hot data in SRAM, minimizes HBM transfers. │
│ │
└────────────────────────────────────────────────────────────────────┘
The online softmax trick is crucial: normally softmax needs to see all values to compute the normalizer. Flash Attention uses a mathematical reformulation that allows incremental computation.
Grouped Query Attention (GQA)
GQA reduces KV-cache memory by sharing K, V across groups of heads.
┌────────────────────────────────────────────────────────────────────┐
│ ATTENTION VARIANTS COMPARISON │
├────────────────────────────────────────────────────────────────────┤
│ │
│ MULTI-HEAD ATTENTION (MHA) - Original │
│ ───────────────────────────────────── │
│ Each head has its own Q, K, V projections │
│ │
│ Q heads: [Q₁] [Q₂] [Q₃] [Q₄] [Q₅] [Q₆] [Q₇] [Q₈] │
│ K heads: [K₁] [K₂] [K₃] [K₄] [K₅] [K₆] [K₇] [K₈] │
│ V heads: [V₁] [V₂] [V₃] [V₄] [V₅] [V₆] [V₇] [V₈] │
│ │
│ KV-cache: 8 K heads + 8 V heads = 16 vectors per layer │
│ │
│ ──────────────────────────────────────────────────────────────── │
│ │
│ MULTI-QUERY ATTENTION (MQA) - Extreme sharing │
│ ───────────────────────────────────────────── │
│ All Q heads share a single K and V │
│ │
│ Q heads: [Q₁] [Q₂] [Q₃] [Q₄] [Q₅] [Q₆] [Q₇] [Q₈] │
│ K heads: [────────────── K ──────────────────] │
│ V heads: [────────────── V ──────────────────] │
│ │
│ KV-cache: 1 K + 1 V = 2 vectors per layer (8× reduction!) │
│ Downside: Some quality loss due to limited K/V diversity │
│ │
│ ──────────────────────────────────────────────────────────────── │
│ │
│ GROUPED QUERY ATTENTION (GQA) - Balanced │
│ ──────────────────────────────────────── │
│ Q heads grouped, each group shares K/V │
│ │
│ Q heads: [Q₁ Q₂] [Q₃ Q₄] [Q₅ Q₆] [Q₇ Q₈] │
│ K heads: [ K₁ ] [ K₂ ] [ K₃ ] [ K₄ ] │
│ V heads: [ V₁ ] [ V₂ ] [ V₃ ] [ V₄ ] │
│ │
│ KV-cache: 4 K + 4 V = 8 vectors per layer (2× reduction) │
│ Better quality than MQA, still significant memory savings │
│ │
│ USED IN: LLaMA 2 70B, LLaMA 3, Mistral, and most modern LLMs │
│ │
└────────────────────────────────────────────────────────────────────┘
Why GQA works: Research found that K/V don't need as much diversity as Q. The "what to look for" (Q) benefits from variety, but "what to find" (K) and "what to return" (V) can be shared across heads without much loss.
RMSNorm: Simpler Than LayerNorm
Modern LLMs often use RMSNorm instead of LayerNorm:
┌────────────────────────────────────────────────────────────────────┐
│ RMSNORM vs LAYERNORM │
├────────────────────────────────────────────────────────────────────┤
│ │
│ LAYERNORM: │
│ ────────── │
│ 1. Compute mean: μ = mean(x) │
│ 2. Compute variance: σ² = var(x) │
│ 3. Normalize: x_norm = (x - μ) / √(σ² + ε) │
│ 4. Scale & shift: output = γ * x_norm + β │
│ │
│ Parameters: γ (scale), β (shift) - both [d_model] │
│ │
│ ──────────────────────────────────────────────────────────────── │
│ │
│ RMSNORM (Root Mean Square Normalization): │
│ ───────────────────────────────────────── │
│ 1. Compute RMS: rms = √(mean(x²)) │
│ 2. Normalize: x_norm = x / rms │
│ 3. Scale only: output = γ * x_norm │
│ │
│ Parameters: γ (scale) only - [d_model] │
│ No mean subtraction, no shift parameter β │
│ │
│ ──────────────────────────────────────────────────────────────── │
│ │
│ WHY RMSNORM? │
│ ──────────── │
│ ✓ Faster: fewer operations (no mean computation) │
│ ✓ Fewer params: no β vector │
│ ✓ Empirically: works just as well for LLMs │
│ ✓ The mean subtraction in LayerNorm isn't necessary │
│ │
│ COMPUTATION: │
│ │
│ LayerNorm: 2 reductions (mean, var) + 2 param vectors │
│ RMSNorm: 1 reduction (sum of squares) + 1 param vector │
│ │
└────────────────────────────────────────────────────────────────────┘
SwiGLU: Gated MLP Activation
Modern LLMs replace GELU with SwiGLU in the MLP block:
┌────────────────────────────────────────────────────────────────────┐
│ MLP ACTIVATION COMPARISON │
├────────────────────────────────────────────────────────────────────┤
│ │
│ ORIGINAL MLP (GPT-2): │
│ ───────────────────── │
│ hidden = x @ W_fc1 # Expand: [d] → [4d] │
│ hidden = GELU(hidden) # Activation │
│ output = hidden @ W_fc2 # Contract: [4d] → [d] │
│ │
│ Parameters: W_fc1 [d, 4d] + W_fc2 [4d, d] │
│ │
│ ──────────────────────────────────────────────────────────────── │
│ │
│ SWIGLU MLP (LLaMA, etc.): │
│ ───────────────────────── │
│ gate = x @ W_gate # Gate projection │
│ hidden = x @ W_up # Up projection │
│ hidden = SiLU(gate) * hidden # Gated activation │
│ output = hidden @ W_down # Down projection │
│ │
│ SiLU(x) = x * sigmoid(x) (also called "Swish") │
│ │
│ Parameters: W_gate [d, 4d] + W_up [d, 4d] + W_down [4d, d] │
│ (actually uses 8/3 × d to match param count) │
│ │
│ ──────────────────────────────────────────────────────────────── │
│ │
│ WHY SWIGLU? │
│ ─────────── │
│ The gating mechanism (SiLU(gate) * hidden) allows the model │
│ to control information flow more precisely: │
│ │
│ • gate ≈ 0 → block this information │
│ • gate ≈ 1 → pass this information through │
│ │
│ Empirically: better performance at same parameter count │
│ │
└────────────────────────────────────────────────────────────────────┘
Architecture Comparison: GPT-2 vs Modern LLMs
┌────────────────────────────────────────────────────────────────────────────┐
│ ARCHITECTURE EVOLUTION: GPT-2 → LLaMA 3 │
├────────────────────────────────────────────────────────────────────────────┤
│ │
│ COMPONENT GPT-2 (2019) LLaMA 3 (2024) │
│ ───────────────── ─────────────────── ───────────────────────── │
│ │
│ Position Encoding Learned absolute RoPE (rotary) │
│ [max_seq, d_model] No params, better extrap. │
│ │
│ Normalization LayerNorm RMSNorm │
│ mean + var + β RMS only, faster │
│ │
│ Norm Placement Post-norm Pre-norm │
│ x + Norm(Attn(x)) x + Attn(Norm(x)) │
│ │
│ Attention Multi-Head (MHA) Grouped Query (GQA) │
│ n_kv = n_heads n_kv < n_heads │
│ │
│ MLP Activation GELU SwiGLU │
│ GELU(xW₁)W₂ SiLU(xW_g) * (xW_up) W_down │
│ │
│ MLP Ratio 4× ~2.7× (8/3) │
│ d_ff = 4 × d_model Adjusted for SwiGLU params │
│ │
│ Attention Impl. Standard Flash Attention 2 │
│ O(n²) memory O(n) memory │
│ │
│ Biases Yes (in attention, No (removed from most) │
│ MLP, LayerNorm) │
│ │
│ Vocab Size 50,257 128,000+ │
│ BPE BPE with larger vocab │
│ │
│ Context Length 1,024 8,192 - 128,000+ │
│ │
├────────────────────────────────────────────────────────────────────────────┤
│ │
│ EFFICIENCY GAINS FROM MODERN CHANGES: │
│ ──────────────────────────────────── │
│ • RoPE: No position embedding table, better length generalization │
│ • RMSNorm: ~15% faster than LayerNorm │
│ • GQA: 2-8× reduction in KV-cache memory │
│ • Flash Attention: 2-4× speedup, enables long contexts │
│ • SwiGLU: ~1-2% quality improvement at same params │
│ • No biases: Fewer params, marginally faster │
│ │
└────────────────────────────────────────────────────────────────────────────┘
Key insight: Modern LLMs aren't just "bigger GPT-2"—they incorporate numerous architectural improvements discovered over 5 years of research. Each change is small, but together they add up to significantly better efficiency and capability.
Summary: Key Takeaways
-
Attention solves the bottleneck problem by letting every token directly access every other token, replacing sequential processing with parallel, direct connections.
-
Q, K, V projections transform each token into three roles: what it's looking for (Query), what it contains (Key), and what it provides (Value).
-
Softmax converts raw attention scores into a probability distribution, creating a soft weighted average rather than a hard lookup.
-
Multi-head attention runs multiple attention patterns in parallel, allowing the model to capture different types of relationships simultaneously.
-
The residual stream is the central highway of information, with each layer reading from and writing to it additively.
-
Position encodings (especially RoPE) inject positional information that attention otherwise lacks.
-
Attention patterns specialize: induction heads enable in-context learning, previous token heads provide local context, and different layers focus on different levels of abstraction.
-
Information flows through attention: the mechanism "moves" relevant information between positions in the sequence, building up context-aware representations.
-
Practical optimizations like KV-cache, Flash Attention, and sparse attention make transformers efficient enough to run at scale.
Further Reading
- Vaswani et al. (2017): "Attention Is All You Need" - The original transformer paper
- Elhage et al. (2021): "A Mathematical Framework for Transformer Circuits" - Deep dive into mechanistic interpretability
- Dao et al. (2022): "FlashAttention" - Memory-efficient attention
- Su et al. (2021): "RoFormer: Enhanced Transformer with Rotary Position Embedding" - RoPE paper
- Olsson et al. (2022): "In-context Learning and Induction Heads" - Induction heads discovery