Attention Mechanisms Deep Dive (Part 1: Foundations to Positional Encoding)
Attention Mechanisms in Large Language Models: A Deep Dive
From foundational concepts to deep understanding of how attention powers modern AI
Table of Contents
- Foundations: What Problem Does Attention Solve?
- The Mechanics of Self-Attention
- Multi-Head Attention
- The Residual Stream
- Positional Information
- Attention Patterns and What They Learn
- Putting It All Together
- Advanced Topics
Part 1: Foundations - What Problem Does Attention Solve?
The Limitation of Sequential Processing
Before transformers, the dominant approach for processing sequences (like text) was Recurrent Neural Networks (RNNs) and their variants like LSTMs. These models processed tokens one at a time, left to right, maintaining a hidden state that accumulated information.
Token 1 → [Hidden State] → Token 2 → [Hidden State] → Token 3 → [Hidden State] → ...
This approach had two critical problems:
Problem 1: The Bottleneck Effect
All information about earlier tokens had to be compressed into a fixed-size hidden state vector. Imagine trying to summarize an entire book into a single paragraph—information gets lost. For the sentence "The cat that lived in the house that Jack built slept," by the time the RNN reaches "slept," much of the information about "cat" has been diluted.
Problem 2: Sequential Computation
RNNs couldn't be parallelized. Each step depended on the previous step's output. This made training extremely slow on modern parallel hardware like GPUs.
The Core Idea: Direct Access
Attention solves both problems with a elegant idea: let every token directly "look at" every other token. Instead of passing information through a chain, create direct connections.
┌─────────────────────────────────────────┐
│ Token 1 ←→ Token 2 ←→ Token 3 │
│ ↕ ↕ ↕ │
│ └───────────┴───────────┘ │
│ (All tokens see all tokens) │
└─────────────────────────────────────────┘
Now when processing "slept," the model can directly examine "cat" without that information passing through every intermediate token.
Intuition: Attention as a Soft Dictionary Lookup
A helpful analogy is to think of attention as a soft dictionary lookup.
In a regular dictionary:
- You have a query (the word you're looking up)
- You match it against keys (dictionary entries)
- You retrieve the value (the definition)
In attention:
- Each token generates a query ("What am I looking for?")
- Each token also provides a key ("This is what I contain")
- Each token has a value ("This is the information to retrieve")
But unlike a hard dictionary lookup that returns one exact match, attention does a soft lookup—it returns a weighted combination of all values, where weights are determined by how well each key matches the query.
Think of it like searching a library: instead of finding the one perfect book, you find that 60% of what you need is in book A, 30% in book B, and 10% in book C, and you blend information from all three.
Part 2: The Mechanics of Self-Attention
Now let's get precise about how attention actually works. We'll build it up step by step.
Step 1: Tokens as Vectors (Embeddings)
Before any computation, each token is converted into a vector (a list of numbers). These are called embeddings.
For example, if our embedding dimension is 4 (real models use 768, 4096, or more), a sentence might look like:
"The" → [0.2, -0.5, 0.8, 0.1]
"cat" → [0.9, 0.3, -0.2, 0.7]
"sat" → [-0.1, 0.6, 0.4, 0.2]
These vectors are learned during training and capture semantic meaning—similar words have similar vectors.
Step 2: The Three Projections (Q, K, V)
Each token's embedding is transformed into three different vectors using learned weight matrices:
- Query (Q): "What am I looking for?"
- Key (K): "What do I contain?"
- Value (V): "What information should I provide if selected?"
Mathematically, for an input embedding x:
Q = x · W_Q (Query projection)
K = x · W_K (Key projection)
V = x · W_V (Value projection)
Where W_Q, W_K, and W_V are learned weight matrices.
Why three different projections?
Having separate Q, K, V allows the model to learn that what a token "is" might be different from what it's "looking for." The word "it" might be looking for (Q) a noun it refers to, while its content (K) might signal "I'm a pronoun."
Step 3: Computing Attention Scores
Now we calculate how much each token should attend to every other token. This is done by taking the dot product of the query with all keys.
For token i attending to token j:
score(i,j) = Q_i · K_j
The dot product measures similarity—higher values mean the query and key are more aligned.
Concrete Example:
Let's say we have three tokens with 4-dimensional Q and K vectors:
Q_0 = [1.0, 0.0, 0.5, 0.0] (for "The")
K_0 = [0.5, 0.2, 0.3, 0.1] (for "The")
K_1 = [0.8, 0.1, 0.4, 0.2] (for "cat")
K_2 = [0.1, 0.6, 0.2, 0.3] (for "sat")
Computing attention scores from token 0 to all tokens:
score(0,0) = (1.0×0.5) + (0.0×0.2) + (0.5×0.3) + (0.0×0.1) = 0.65
score(0,1) = (1.0×0.8) + (0.0×0.1) + (0.5×0.4) + (0.0×0.2) = 1.00
score(0,2) = (1.0×0.1) + (0.0×0.6) + (0.5×0.2) + (0.0×0.3) = 0.20
Token 0 ("The") attends most strongly to token 1 ("cat").
Step 4: Scaling
Before applying softmax, we divide scores by the square root of the key dimension (√d_k). This prevents the dot products from becoming too large, which would push softmax into regions with tiny gradients.
scaled_score(i,j) = score(i,j) / √d_k
If d_k = 4:
scaled_score(0,0) = 0.65 / 2 = 0.325
scaled_score(0,1) = 1.00 / 2 = 0.500
scaled_score(0,2) = 0.20 / 2 = 0.100
Step 5: Softmax - Turning Scores into Weights
The softmax function converts raw scores into a probability distribution (values between 0 and 1 that sum to 1).
attention_weight(i,j) = exp(scaled_score(i,j)) / Σ_k exp(scaled_score(i,k))
For our example:
exp(0.325) = 1.384
exp(0.500) = 1.649
exp(0.100) = 1.105
sum = 4.138
weight(0,0) = 1.384 / 4.138 = 0.334 (33.4%)
weight(0,1) = 1.649 / 4.138 = 0.399 (39.9%)
weight(0,2) = 1.105 / 4.138 = 0.267 (26.7%)
Now we have attention weights: token 0 pays 33.4% attention to itself, 39.9% to "cat", and 26.7% to "sat".
Step 6: Weighted Sum of Values
Finally, we compute the output by taking a weighted sum of all value vectors:
output_i = Σ_j (attention_weight(i,j) × V_j)
If our value vectors are:
V_0 = [0.1, 0.2, 0.3, 0.4]
V_1 = [0.5, 0.6, 0.7, 0.8]
V_2 = [0.2, 0.1, 0.4, 0.3]
Then:
output_0 = 0.334×[0.1, 0.2, 0.3, 0.4]
+ 0.399×[0.5, 0.6, 0.7, 0.8]
+ 0.267×[0.2, 0.1, 0.4, 0.3]
output_0 = [0.287, 0.333, 0.486, 0.533]
This output is a context-aware representation—it's not just about "The" anymore, but "The" in the context of this specific sentence.
The Complete Attention Formula
Putting it all together:
Attention(Q, K, V) = softmax(Q · K^T / √d_k) · V
This single formula captures all six steps above.
┌────────────────────────────────────────────────────────────┐
│ Self-Attention Flow │
│ │
│ Input Embeddings │
│ │ │
│ ├──→ W_Q ──→ Queries (Q)──────┐ │
│ ├──→ W_K ──→ Keys (K)─────────┼──→ Q·K^T/√d_k │
│ └──→ W_V ──→ Values (V) │ │ │
│ │ │ softmax │
│ │ │ │ │
│ └─────────────┴────→ weights·V │
│ │ │
│ Output │
└────────────────────────────────────────────────────────────┘
Complete Worked Example with Matrix Multiplication
Let's trace through a full example with concrete numbers and shapes at every step.
Setup:
- 3 tokens: "The", "cat", "sat"
- Embedding dimension: d_model = 4
- Key/Query/Value dimension: d_k = 4 (single head for simplicity)
Step 1: Token Embeddings
Look up each token in the embedding table:
X = Token Embeddings [3, 4]
d0 d1 d2 d3
┌─────────────────────────┐
The │ 0.1 0.2 0.3 0.4 │
cat │ 0.5 0.6 0.7 0.8 │
sat │ 0.2 0.1 0.4 0.3 │
└─────────────────────────┘
Shape: [3 tokens, 4 dimensions]
Step 2: Weight Matrices (Learned Parameters)
These are learned during training:
W_Q [4, 4] W_K [4, 4] W_V [4, 4]
┌─────────────────────┐ ┌─────────────────────┐ ┌─────────────────────┐
│ 0.1 0.2 0.0 0.1 │ │ 0.2 0.1 0.3 0.0 │ │ 0.1 0.0 0.2 0.1 │
│ 0.3 0.1 0.2 0.0 │ │ 0.0 0.2 0.1 0.3 │ │ 0.2 0.3 0.1 0.0 │
│ 0.0 0.3 0.1 0.2 │ │ 0.1 0.0 0.2 0.1 │ │ 0.0 0.1 0.3 0.2 │
│ 0.2 0.0 0.3 0.1 │ │ 0.3 0.1 0.0 0.2 │ │ 0.3 0.2 0.0 0.1 │
└─────────────────────┘ └─────────────────────┘ └─────────────────────┘
Step 3: Compute Q, K, V
Matrix multiply: Q = X @ W_Q, K = X @ W_K, V = X @ W_V
Q = X @ W_Q [3, 4]
┌─────────────────────────┐
│ 0.20 0.16 0.19 0.12 │ ← The
│ 0.52 0.40 0.47 0.28 │ ← cat
│ 0.23 0.20 0.17 0.13 │ ← sat
└─────────────────────────┘
K = X @ W_K [3, 4]
┌─────────────────────────┐
│ 0.17 0.11 0.13 0.15 │ ← The
│ 0.41 0.27 0.33 0.35 │ ← cat
│ 0.18 0.10 0.17 0.12 │ ← sat
└─────────────────────────┘
V = X @ W_V [3, 4]
┌─────────────────────────┐
│ 0.19 0.15 0.13 0.08 │ ← The
│ 0.47 0.39 0.33 0.20 │ ← cat
│ 0.15 0.12 0.18 0.11 │ ← sat
└─────────────────────────┘
Step 4: Attention Scores
Scores = Q @ K.T (K transposed)
Q [3,4] @ K.T [4,3] = Scores [3,3]
Scores [3,3]
┌─────────────────────┐
│ 0.10 0.25 0.10 │ The attends to [The, cat, sat]
│ 0.26 0.64 0.25 │ cat attends to [The, cat, sat]
│ 0.11 0.28 0.11 │ sat attends to [The, cat, sat]
└─────────────────────┘
Each row = one token's raw attention to all tokens
Step 5: Scale
Scaled = Scores / √d_k = Scores / √4 = Scores / 2
Scaled Scores [3,3]
┌─────────────────────┐
│ 0.05 0.125 0.05 │ The
│ 0.13 0.32 0.125│ cat
│ 0.055 0.14 0.055│ sat
└─────────────────────┘
Step 6: Softmax (row-wise)
Each row becomes a probability distribution (sums to 1):
Attention Weights [3,3]
The cat sat
┌────────────────────────┐
The → │ 0.31 0.38 0.31 │ (sums to 1.0)
cat → │ 0.29 0.42 0.29 │ (sums to 1.0)
sat → │ 0.30 0.40 0.30 │ (sums to 1.0)
└────────────────────────┘
Interpretation: Every token attends most to "cat" (~40%), less to others (~30% each).
Step 7: Weighted Sum of Values
Output = Attention_Weights @ V
Weights [3,3] @ V [3,4] = Output [3,4]
┌────────────────────┐ ┌─────────────────────────┐ ┌─────────────────────────┐
│ 0.31 0.38 0.31 │ │ 0.19 0.15 0.13 0.08 │ │ 0.29 0.24 0.21 0.13 │ The
│ 0.29 0.42 0.29 │ @ │ 0.47 0.39 0.33 0.20 │ = │ 0.30 0.25 0.22 0.13 │ cat
│ 0.30 0.40 0.30 │ │ 0.15 0.12 0.18 0.11 │ │ 0.29 0.24 0.21 0.13 │ sat
└────────────────────┘ └─────────────────────────┘ └─────────────────────────┘
Shape Flow Summary:
┌─────────────────────────────────────────────────────────────┐
│ Token IDs [3] │
│ ↓ (embedding lookup) │
│ X [3, 4] ← token embeddings │
│ ↓ │
│ ┌────┼────┬────┐ │
│ ↓ ↓ ↓ (multiply by W_Q, W_K, W_V) │
│ Q K V [3,4] each │
│ │ │ │
│ └──┬─┘ │
│ ↓ (Q @ K.T) │
│ Scores [3, 3] ← each token to each token │
│ ↓ (÷ √d_k) │
│ Scaled [3, 3] │
│ ↓ (softmax per row) │
│ Weights [3, 3] ← attention probabilities │
│ ↓ (@ V) │
│ Output [3, 4] ← context-aware embeddings │
└─────────────────────────────────────────────────────────────┘
What Changed?
Before attention: Each token knew only about itself.
After attention: Each token contains blended information from all tokens, weighted by relevance.
Shapes at Every Step (Reference Guide)
Understanding the exact tensor dimensions helps demystify attention. Here's a complete shape reference for a typical configuration:
Model Configuration Example:
- Vocabulary size: 50,257 (GPT-2)
- Sequence length (n): 1024 tokens
- Model dimension (d_model): 768
- Number of heads: 12
- Head dimension (d_head): 768 / 12 = 64
┌─────────────────────────────────────────────────────────────────────┐
│ SHAPE REFERENCE CARD │
├─────────────────────────────────────────────────────────────────────┤
│ EMBEDDING LOOKUP │
│ Token IDs: [batch, seq_len] e.g., [1, 1024]│
│ Embedding table: [vocab_size, d_model] [50257, 768] │
│ Token embeddings: [batch, seq_len, d_model] [1, 1024, 768] │
├─────────────────────────────────────────────────────────────────────┤
│ PROJECTION WEIGHTS (per head) │
│ W_Q, W_K, W_V: [d_model, d_head] [768, 64] │
│ W_O (output proj): [n_heads × d_head, d_model] [768, 768] │
├─────────────────────────────────────────────────────────────────────┤
│ Q, K, V MATRICES (per head, per batch) │
│ Q, K, V: [seq_len, d_head] [1024, 64] │
├─────────────────────────────────────────────────────────────────────┤
│ ATTENTION COMPUTATION │
│ Attention scores (Q @ K.T): [seq_len, seq_len] [1024, 1024] │
│ Attention weights (softmax): [seq_len, seq_len] [1024, 1024] │
│ Attention output: [seq_len, d_head] [1024, 64] │
├─────────────────────────────────────────────────────────────────────┤
│ MULTI-HEAD OUTPUT │
│ Concatenated heads: [seq_len, n_heads × d_head] [1024, 768] │
│ After W_O projection: [seq_len, d_model] [1024, 768] │
└─────────────────────────────────────────────────────────────────────┘
Why These Shapes Matter:
-
The attention matrix is quadratic: [seq_len × seq_len] means 1024² = ~1M elements per head. For 100K context, that's 10 billion elements—this is why long contexts are expensive.
-
Head dimension trades off: More heads with smaller d_head means more diverse attention patterns but less capacity per pattern.
-
Memory bottleneck: KV-cache stores [seq_len, d_head] per head per layer. With 32 layers × 32 heads × 128 d_head × 2 (K+V) × 4 bytes = 1MB per token of context.
FAQ: Where Do Token Embeddings Come From?
Q: Are token embeddings from Word2Vec or GloVe?
A: No. Token embeddings in transformers are learned from scratch during training—they're randomly initialized and updated via backpropagation.
Why not use Word2Vec/GloVe?
-
Tokenization mismatch: Modern LLMs use subword tokenization (BPE, SentencePiece). Word2Vec works with whole words. "unhappiness" might be tokenized as ["un", "happiness"] or ["un", "happ", "iness"]—Word2Vec has no embeddings for these pieces.
-
End-to-end learning works better: When the embedding layer trains alongside everything else, it learns representations optimized for the actual task.
-
Scale: With billions of training tokens, the model can learn better embeddings than pre-trained ones from smaller corpora.
How to find embeddings in code:
# In HuggingFace Transformers (GPT-2):
from transformers import GPT2Model
model = GPT2Model.from_pretrained('gpt2')
# Token embeddings
model.wte # "word token embeddings" - shape [50257, 768]
# Position embeddings (for absolute positional encoding)
model.wpe # "word position embeddings" - shape [1024, 768]
# PyTorch native:
# nn.Embedding(vocab_size, d_model)
Q: How much of the model is in attention vs MLPs?
This varies by model, but here's a typical breakdown:
┌────────────────────────────────────────────────────────────────┐
│ PARAMETER DISTRIBUTION (GPT-2 Small) │
├────────────────────────────────────────────────────────────────┤
│ │
│ Embeddings: ~40M params (32%) │
│ └─ Token: 50257 × 768 = 38.6M │
│ └─ Position: 1024 × 768 = 0.8M │
│ │
│ Per Layer (×12 layers): │
│ Attention: ~2.4M per layer (28% of layer) │
│ └─ W_Q, W_K, W_V, W_O: 4 × (768 × 768) = 2.36M │
│ MLP: ~4.7M per layer (72% of layer) │
│ └─ fc1: 768 × 3072 = 2.36M │
│ └─ fc2: 3072 × 768 = 2.36M │
│ │
│ Total Attention: ~29M params (23%) │
│ Total MLP: ~56M params (45%) │
│ │
│ TOTAL: ~124M parameters │
└────────────────────────────────────────────────────────────────┘
Key insight: MLPs actually have more parameters than attention (~2× in most architectures). This is because the MLP expands to 4× the model dimension (768 → 3072 → 768).
For GPT-3:
- ~33% of compute is in attention
- ~67% is in MLPs
This is why "attention is all you need" is somewhat misleading—MLPs are doing a lot of work too!
Part 3: Multi-Head Attention
Why One Attention Pattern Isn't Enough
A single attention mechanism can only capture one type of relationship at a time. But language has many simultaneous relationships:
- Syntactic: subject-verb agreement
- Semantic: pronoun-antecedent reference
- Positional: adjacent word relationships
- Long-range: topic coherence
In "The cat that I saw yesterday was sleeping," a single token might need to simultaneously:
- Attend to its syntactic subject
- Attend to semantically related words
- Attend to nearby words for local context
Multi-head attention solves this by running multiple attention operations in parallel.
How Heads Work in Parallel
Instead of one large attention operation, we split into multiple smaller "heads," each with its own Q, K, V projections.
If the model dimension is 512 and we use 8 heads:
- Each head works with 512/8 = 64 dimensions
- Each head has its own W_Q, W_K, W_V matrices (mapping 512 → 64)
- Each head computes attention independently
┌──────────────────────────────────────────────────────────┐
│ Multi-Head Attention │
│ │
│ Input (d=512) │
│ │ │
│ ├──→ Head 1 (d=64) ──→ [attention output 1] │
│ ├──→ Head 2 (d=64) ──→ [attention output 2] │
│ ├──→ Head 3 (d=64) ──→ [attention output 3] │
│ ├──→ Head 4 (d=64) ──→ [attention output 4] │
│ ├──→ Head 5 (d=64) ──→ [attention output 5] │
│ ├──→ Head 6 (d=64) ──→ [attention output 6] │
│ ├──→ Head 7 (d=64) ──→ [attention output 7] │
│ └──→ Head 8 (d=64) ──→ [attention output 8] │
│ │ │
│ Concatenate │
│ │ │
│ [Combined: d=512] │
│ │ │
│ W_O (final projection) │
│ │ │
│ Output (d=512) │
└──────────────────────────────────────────────────────────┘
What Different Heads Learn
Research has revealed that different heads specialize in different patterns:
Syntactic Heads: Some heads learn to attend based on grammatical structure—verbs attending to their subjects, adjectives to their nouns.
Positional Heads: Some heads attend primarily to specific relative positions—the previous token, the next token, or tokens at fixed distances.
Semantic Heads: Some heads group semantically related concepts, even across long distances.
Copy Heads: Some heads learn to identify when the model should "copy" or repeat information seen earlier.
This specialization emerges naturally from training—the model discovers that dividing labor this way helps it predict text better.
Concatenation and Final Projection
After all heads compute their outputs, we:
- Concatenate: Stack all head outputs together (8 heads × 64 dims = 512 dims)
- Project: Apply one more learned matrix W_O to mix information across heads
MultiHead(Q, K, V) = Concat(head_1, ..., head_h) · W_O
where head_i = Attention(Q·W_Q^i, K·W_K^i, V·W_V^i)
The final projection is crucial—it allows the model to learn how to combine the different perspectives from each head.
FAQ: What Prevents Heads from Learning the Same Thing?
Q: If all heads have the same architecture, what stops them from becoming redundant?
A: Nothing explicit—and some redundancy does occur. But several forces encourage diversity:
1. Random Initialization
Each head's W_Q, W_K, W_V weights start with different random values. This "symmetry breaking" puts heads on different trajectories from the start.
Head 1 W_Q initialized: [0.02, -0.01, 0.03, ...]
Head 2 W_Q initialized: [-0.01, 0.04, 0.00, ...]
Different starting points → different learned patterns
2. Gradient Pressure from the Loss
If two heads learned identical patterns, they'd contribute redundant information. The output projection (W_O) would effectively see duplicated inputs. Through backpropagation, the model has pressure to make each head useful—redundancy doesn't help minimize loss.
3. The Lottery Ticket Hypothesis
Some researchers believe that random initialization creates a "lottery" of potential circuits, and training finds the useful ones. Different random seeds give different winning "tickets," leading to diverse heads.
4. Emergent Specialization
Research shows heads do specialize (syntactic heads, positional heads, induction heads), but this emerges naturally—there's no explicit mechanism enforcing "you must learn syntax, you must learn position."
What We Observe:
┌───────────────────────────────────────────────────────────────┐
│ HEAD DIVERSITY IN PRACTICE │
├───────────────────────────────────────────────────────────────┤
│ │
│ ✓ Many heads DO specialize (subject-verb, previous token) │
│ ✓ Some heads are "backup" copies of important patterns │
│ ✗ Some heads appear nearly unused ("dead heads") │
│ ✗ Some heads have overlapping attention patterns │
│ │
│ Pruning experiments show: ~30-50% of heads can be removed │
│ with minimal performance loss (Voita et al., 2019) │
└───────────────────────────────────────────────────────────────┘
This partial redundancy may actually be beneficial—it provides robustness and allows the model to represent important patterns more reliably.
Part 4: The Residual Stream
What Is the Residual Stream?
The residual stream is the "main highway" of information flow through a transformer. It's the tensor that flows through the network, getting modified by each layer.
Think of it as a river: the initial input (token embeddings) starts the flow, and each layer adds to the river without replacing what's already there.
┌────────────────────────────────────────────────────────────┐
│ The Residual Stream │
│ │
│ Token Embeddings ──→ [Residual Stream: initial state] │
│ │ │
│ ↓ │
│ ┌──────┴──────┐ │
│ │ Layer 1 │ │
│ │ Attention │──→ writes output │
│ └──────┬──────┘ │ │
│ │←────────────────┘ │
│ ↓ (added, not replaced) │
│ ┌──────┴──────┐ │
│ │ Layer 1 │ │
│ │ MLP │──→ writes output │
│ └──────┬──────┘ │ │
│ │←────────────────┘ │
│ ↓ │
│ ┌──────┴──────┐ │
│ │ Layer 2 │ │
│ │ Attention │──→ writes output │
│ └──────┬──────┘ │ │
│ │←────────────────┘ │
│ ↓ │
│ ... │
│ ↓ │
│ [Final representation] │
└────────────────────────────────────────────────────────────┘
Skip Connections: Why They Matter
The key mechanism enabling the residual stream is the skip connection (or residual connection). Instead of:
output = Layer(input)
We have:
output = input + Layer(input)
This has profound implications:
1. Gradient Flow: During training, gradients can flow directly backward through skip connections, avoiding the "vanishing gradient" problem that plagued deep networks.
2. Additive Updates: Each layer contributes additively. The original information is preserved, and layers add refinements.
3. Easy Identity: If a layer has nothing useful to add, it can output near-zero, effectively passing input through unchanged. This makes deep networks easier to train.
How Attention Reads From and Writes To the Residual Stream
Attention layers interact with the residual stream in two ways:
Reading: When computing Q, K, V projections, the attention layer reads from the current state of the residual stream. This means later layers have access to everything earlier layers wrote.
Writing: The attention output is added back to the residual stream. This "writes" the attention layer's contribution.
residual_stream = token_embeddings + positional_embeddings
for layer in transformer_layers:
# Attention reads from residual stream
attn_output = layer.attention(residual_stream)
# Attention writes to residual stream
residual_stream = residual_stream + attn_output
# MLP reads from residual stream
mlp_output = layer.mlp(residual_stream)
# MLP writes to residual stream
residual_stream = residual_stream + mlp_output
Accumulating Information Across Layers
The residual stream accumulates information progressively:
- Early layers: Add basic syntactic and local features
- Middle layers: Build up compositional semantics and relationships
- Later layers: Refine task-specific representations and predictions
This is sometimes called the "residual stream as a communication channel"—components (attention heads, MLPs) read shared information and write back their contributions, building up a rich representation.
A key insight: attention heads in different layers can "communicate" through the residual stream. An early head might identify a subject, write a "marker" to the residual stream, and a later head can read that marker to correctly conjugate a verb.
Deep Dive: Training vs Inference
Understanding how transformers work during training versus inference clarifies many design decisions.
Training: Forward Pass with Concrete Numbers
Let's trace training on a mini-batch with a tiny model (2 layers, d_model=4, single head):
Input sequence: "The cat sat" → Predict "on"
Token IDs: [0, 1, 2] → Target: 3
Step 1: Embedding Lookup
─────────────────────────
Embedding table E [vocab_size=10, d_model=4]:
┌────────────────────┐
0 │ 0.1 0.2 0.3 0.4 │ ← "The"
1 │ 0.5 0.6 0.7 0.8 │ ← "cat"
2 │ 0.2 0.1 0.4 0.3 │ ← "sat"
3 │ 0.3 0.4 0.2 0.5 │ ← "on" (target)
└────────────────────┘
X = E[[0,1,2]] → shape [3, 4]
┌────────────────────┐
│ 0.1 0.2 0.3 0.4 │ The
│ 0.5 0.6 0.7 0.8 │ cat
│ 0.2 0.1 0.4 0.3 │ sat
└────────────────────┘
Step 2: Add Position Embeddings (learned)
─────────────────────────────────────────
X = X + pos_emb[[0,1,2]]
Step 3: Layer 1 Attention + MLP
───────────────────────────────
X = X + Attention(X) # Residual connection
X = X + MLP(X) # Residual connection
Step 4: Layer 2 Attention + MLP
───────────────────────────────
X = X + Attention(X)
X = X + MLP(X)
Step 5: Final LayerNorm + Projection to Vocabulary
──────────────────────────────────────────────────
logits = LayerNorm(X) @ W_vocab # Shape [3, vocab_size=10]
For position 2 (predicting what comes after "sat"):
logits[2] = [-1.2, 0.3, -0.5, 2.1, ...]
↑
"on" has highest score
Step 6: Loss Computation
────────────────────────
probs = softmax(logits[2]) # Convert to probabilities
loss = -log(probs[3]) # Cross-entropy: how wrong were we about "on"?
Training: Backward Pass (Gradient Flow)
┌──────────────────────────────────────────────────────────────────┐
│ GRADIENT FLOW (BACKWARD PASS) │
│ │
│ Loss = -log(softmax(logits)[target]) │
│ │ │
│ ├─→ ∂Loss/∂logits: Gradient through cross-entropy │
│ │ └─→ ∂Loss/∂W_vocab: Update vocabulary projection │
│ │ │
│ ├─→ ∂Loss/∂X_final: Gradient to residual stream │
│ │ │ │
│ │ ├─→ Through Layer 2 MLP: │
│ │ │ ∂Loss/∂W_mlp2, ∂Loss/∂X (via residual) │
│ │ │ │
│ │ ├─→ Through Layer 2 Attention: │
│ │ │ ∂Loss/∂W_Q, ∂Loss/∂W_K, ∂Loss/∂W_V, ∂Loss/∂W_O │
│ │ │ ∂Loss/∂X (via residual) │
│ │ │ │
│ │ ├─→ Through Layer 1 (same pattern) │
│ │ │ │
│ │ └─→ Through Embeddings: │
│ │ ∂Loss/∂E[[0,1,2]]: Update embeddings for │
│ │ "The", "cat", "sat" │
│ │
│ KEY: Skip connections let gradients flow directly back! │
│ Without them: gradient vanishes after many layers │
└──────────────────────────────────────────────────────────────────┘
Efficiency trick: During training, we process ALL positions in parallel. The causal mask ensures position 2 only sees positions 0, 1, 2—but we compute all positions simultaneously. This is called "teacher forcing."
Inference: Autoregressive Generation Step-by-Step
During generation, we produce tokens one at a time:
┌──────────────────────────────────────────────────────────────────┐
│ AUTOREGRESSIVE GENERATION WITH KV-CACHE │
├──────────────────────────────────────────────────────────────────┤
│ │
│ Prompt: "The cat" │
│ │
│ STEP 1: Process prompt (prefill) │
│ ───────────────────────────────── │
│ Input: [The, cat] │
│ │
│ For each layer: │
│ Q = X @ W_Q → shape [2, d_head] │
│ K = X @ W_K → shape [2, d_head] ← SAVE to KV-cache │
│ V = X @ W_V → shape [2, d_head] ← SAVE to KV-cache │
│ attn_out = attention(Q, K, V) │
│ │
│ KV-cache after step 1: │
│ K_cache = [K_The, K_cat] │
│ V_cache = [V_The, V_cat] │
│ │
│ Output: logits → sample → "sat" │
│ │
│ ───────────────────────────────────────────────────────────── │
│ │
│ STEP 2: Generate "sat" │
│ ────────────────────── │
│ Input: just [sat] (1 new token) │
│ │
│ For each layer: │
│ Q_new = X_sat @ W_Q → shape [1, d_head] │
│ K_new = X_sat @ W_K → shape [1, d_head] │
│ V_new = X_sat @ W_V → shape [1, d_head] │
│ │
│ K_full = concat(K_cache, K_new) → [3, d_head] │
│ V_full = concat(V_cache, V_new) → [3, d_head] │
│ │
│ # "sat" attends to ALL previous tokens │
│ attn_out = attention(Q_new, K_full, V_full) │
│ │
│ # Update cache │
│ K_cache = K_full │
│ V_cache = V_full │
│ │
│ Output: logits → sample → "on" │
│ │
│ ───────────────────────────────────────────────────────────── │
│ │
│ STEP 3: Generate "on" (same pattern) │
│ Q_new for "on", use cached K, V for [The, cat, sat] │
│ ... │
└──────────────────────────────────────────────────────────────────┘
Why "autoregressive"?
The term comes from statistics/time series:
- "Auto" = self (Greek: αὐτός)
- "Regressive" = predicting from previous values
An autoregressive model predicts the next value from its own previous outputs—creating a feedback loop where each prediction becomes input for the next.
Traditional regression: Y = f(X) # Output from input
Autoregressive: Y_t = f(Y_{t-1}, Y_{t-2}, ...) # Output from own past
FAQ: Why Can't We Cache Q? Why Only K and V?
Q: Why store K and V but recompute Q for each new token?
A: Because attention is asymmetric—old tokens never need to attend to new tokens.
The attention operation: Q_new @ K_old.T
This computes: "How much should the NEW token attend to OLD tokens?"
We DON'T need: Q_old @ K_new.T
Because: "How much should OLD tokens attend to NEW tokens?"
↑ This is forbidden by causal masking anyway!
┌───────────────────────────────────────────────────────┐
│ CAUSAL ATTENTION MATRIX │
│ │
│ K_The K_cat K_sat K_on │
│ ┌────────────────────────────┐ │
│ Q_The │ ✓ ✗ ✗ ✗ │ │
│ Q_cat │ ✓ ✓ ✗ ✗ │ │
│ Q_sat │ ✓ ✓ ✓ ✗ │ Masked out │
│ Q_on │ ✓ ✓ ✓ ✓ │ (future) │
│ └────────────────────────────┘ │
│ │
│ When generating "on", we ONLY need the bottom row! │
│ Q_on @ [K_The, K_cat, K_sat, K_on].T │
│ │
│ Old Q values (Q_The, Q_cat, Q_sat) are never used │
│ again because those positions are already processed. │
└───────────────────────────────────────────────────────┘
The new token needs K, V from all previous tokens (to attend to them), but previous tokens never need anything from the new token.
Deep Dive: LayerNorm Explained
LayerNorm appears before (or after) every attention and MLP block. What does it actually do?
Step-by-Step Calculation
Input vector x = [2.0, 4.0, 6.0, 8.0] # One token's representation
Step 1: Compute mean
────────────────────
μ = (2 + 4 + 6 + 8) / 4 = 5.0
Step 2: Compute variance
────────────────────────
variance = [(2-5)² + (4-5)² + (6-5)² + (8-5)²] / 4
= [9 + 1 + 1 + 9] / 4 = 5.0
Step 3: Normalize (zero mean, unit variance)
────────────────────────────────────────────
x_norm = (x - μ) / √(variance + ε)
= ([2,4,6,8] - 5) / √5.0001
= [-3, -1, 1, 3] / 2.236
= [-1.34, -0.45, 0.45, 1.34]
Step 4: Scale and shift (learned parameters γ, β)
─────────────────────────────────────────────────
output = γ * x_norm + β
If γ = [1,1,1,1] and β = [0,0,0,0]:
output = [-1.34, -0.45, 0.45, 1.34] # Same as x_norm
If γ = [2,2,2,2] and β = [1,1,1,1]:
output = [-1.68, 0.10, 1.90, 3.68] # Rescaled
Key insight: γ and β are learned parameters (shape [d_model]) that let the model undo the normalization if needed—but the normalization helps training stability.
LayerNorm vs BatchNorm
┌─────────────────────────────────────────────────────────────────┐
│ NORMALIZATION COMPARISON │
├─────────────────────────────────────────────────────────────────┤
│ │
│ INPUT TENSOR: [batch, seq_len, d_model] │
│ │
│ BatchNorm normalizes across: BATCH dimension │
│ ───────────────────────────────────────── │
│ For each feature d: │
│ mean/var computed over all (batch × seq_len) examples │
│ │
│ Problem: Requires large batches for stable statistics │
│ Problem: Different behavior at train vs inference │
│ Problem: Sequences have variable length │
│ │
│ LayerNorm normalizes across: FEATURE dimension (d_model) │
│ ────────────────────────────────────────────────────── │
│ For each token (batch, position): │
│ mean/var computed over the d_model features │
│ │
│ ✓ Works with any batch size (even batch=1) │
│ ✓ Same computation at train and inference │
│ ✓ Each token normalized independently │
│ │
│ VISUAL: │
│ │
│ Batch dimension → │
│ ┌─────────────────┐ │
│ │ ████████████████│ ← LayerNorm normalizes this row │
│ │ ████████████████│ (all d_model features for one token) │
│ │ ████████████████│ │
│ └─────────────────┘ │
│ ↑ │
│ BatchNorm normalizes this column │
│ (one feature across all examples) │
└─────────────────────────────────────────────────────────────────┘
Why LayerNorm for transformers?
- Inference simplicity: No running statistics to maintain
- Variable sequence lengths: Each token normalized independently
- Works with batch size 1: Critical for generation
- Training stability: Keeps activations in a reasonable range
Part 5: Positional Information
Why Transformers Need Position Encodings
Unlike RNNs, which process tokens sequentially (inherently knowing order), transformers process all tokens in parallel. Without explicit position information, the model couldn't distinguish:
- "The cat sat on the mat" from "The mat sat on the cat"
- "dog bites man" from "man bites dog"
Attention is permutation invariant—shuffling input order without position information would give the same output.
Position encodings solve this by adding information about where each token appears in the sequence.
Absolute Positional Encodings
The original transformer (Vaswani et al., 2017) used sinusoidal positional encodings—fixed mathematical functions based on position.
For position pos and dimension i:
PE(pos, 2i) = sin(pos / 10000^(2i/d))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d))
Each position gets a unique vector, added to the token embedding:
input = token_embedding + positional_embedding
Properties:
- Each position has a unique encoding
- The model can potentially extrapolate to longer sequences (the functions are defined for any position)
- Relative positions have consistent patterns (PE[pos+k] can be represented as a linear function of PE[pos])
Learned absolute positions are another variant—treat position embeddings as parameters learned during training, like word embeddings.
Relative Positional Encodings
Absolute encodings say "this token is at position 5." Relative encodings say "this token is 3 positions before that token."
Relative encodings are often more useful because:
- "The word before the verb" is more generalizable than "the word at position 7"
- They naturally handle varying sequence lengths
- They capture the intuition that relative distance matters more than absolute position
In relative position schemes, the attention score calculation is modified to include relative position information:
attention_score(i,j) = Q_i · K_j + position_bias(i-j)
Rotary Position Embeddings (RoPE)
RoPE (Rotary Position Embeddings) is used in most modern LLMs (LLaMA, GPT-NeoX, PaLM, and many others). It's elegant because it encodes absolute position in a way that naturally captures relative positions in attention.
The key idea: rotate the Q and K vectors based on their position. When computing Q·K, the angle between them encodes their relative position.
Mathematically, RoPE applies a rotation matrix based on position:
Q_rotated = R(θ_pos) · Q
K_rotated = R(θ_pos) · K
Where R(θ) is a rotation matrix and θ depends on position.
Why rotations work:
When you take the dot product of two rotated vectors:
Q_i_rotated · K_j_rotated
The rotation angles combine in a way that depends on (i - j), the relative position. The absolute positions "cancel out" leaving only relative information.
┌────────────────────────────────────────────────────────────┐
│ RoPE Intuition │
│ │
│ Token at position 3: Q₃ rotated by θ×3 │
│ Token at position 7: K₇ rotated by θ×7 │
│ │
│ Q₃ · K₇ depends on rotation difference = θ×(7-3) = θ×4 │
│ │
│ The relative distance (4) is encoded in the angle! │
└────────────────────────────────────────────────────────────┘
Benefits of RoPE:
- Naturally captures relative positions without explicit bias terms
- Can extrapolate to longer sequences than seen in training (with some degradation)
- Computationally efficient—rotations are cheap
- Works with linear attention variants
FAQ: Does RoPE Have Learned Parameters?
Q: Are the rotation angles in RoPE learned during training?
A: No. RoPE uses a fixed, deterministic formula. The rotation angles are computed from position indices using:
θ_i = 10000^(-2i/d) # Fixed frequencies for each dimension pair
angle(pos, i) = pos × θ_i
There are no trainable parameters in RoPE itself.
What IS learned:
- The W_Q and W_K matrices that produce Q and K before rotation
- These projections learn to work WITH the rotation geometry
Why fixed rotations work:
The model learns Q, K projections that exploit the rotation structure. It doesn't need to learn the rotations themselves—the mathematical properties of rotations naturally encode relative positions.
┌────────────────────────────────────────────────────────────────┐
│ WHERE PARAMETERS LIVE │
├────────────────────────────────────────────────────────────────┤
│ │
│ Learned (updated during training): │
│ ✓ W_Q, W_K, W_V, W_O projection matrices │
│ ✓ Token embeddings │
│ ✓ LayerNorm γ and β │
│ ✓ MLP weights │
│ │
│ Fixed (computed, not learned): │
│ ✗ RoPE rotation angles │
│ ✗ Sinusoidal position encodings (if used) │
│ ✗ Causal attention mask │
│ │
│ Note: Learned absolute position embeddings ARE trainable │
│ (GPT-2 style), but RoPE is not. │
└────────────────────────────────────────────────────────────────┘
Comparison: Position Encoding Methods
┌──────────────────────────────────────────────────────────────────────────┐
│ POSITION ENCODING COMPARISON │
├──────────────────────────────────────────────────────────────────────────┤
│ │
│ METHOD LEARNED? RELATIVE? EXTRAPOLATE? USED IN │
│ ───────────────── ──────── ───────── ──────────── ────────── │
│ Learned Absolute Yes No Poor GPT-2 │
│ Sinusoidal No No* Moderate Original │
│ Relative Bias Yes Yes Good T5, ALiBi │
│ RoPE No** Yes Good LLaMA, GPT-4 │
│ │
├──────────────────────────────────────────────────────────────────────────┤
│ * Sinusoidal can represent relative positions mathematically │
│ ** W_Q, W_K learn to work with fixed RoPE rotations │
├──────────────────────────────────────────────────────────────────────────┤
│ │
│ LEARNED ABSOLUTE (GPT-2) │
│ ──────────────────────── │
│ • Separate embedding table: pos_embed[pos] → vector │
│ • Added to token embedding before transformer │
│ • Simple but poor extrapolation: pos > max_len is undefined │
│ │
│ SINUSOIDAL (Original Transformer) │
│ ────────────────────────────────── │
│ • PE(pos,2i) = sin(pos / 10000^(2i/d)) │
│ • PE(pos,2i+1) = cos(pos / 10000^(2i/d)) │
│ • Defined for any position, but attention doesn't naturally use it │
│ │
│ RELATIVE POSITION BIAS (T5, ALiBi) │
│ ─────────────────────────────────── │
│ • Adds learned or fixed bias to attention scores │
│ • bias[i,j] depends on (i - j) relative distance │
│ • ALiBi: simple linear penalty for distance (no learning) │
│ │
│ ROPE (Rotary Position Embedding) │
│ ───────────────────────────────── │
│ • Rotates Q, K by position-dependent angle │
│ • Q_i · K_j naturally encodes (i - j) │
│ • Elegant: absolute encoding → relative attention │
│ │
└──────────────────────────────────────────────────────────────────────────┘
Modern preference: RoPE has become dominant because:
- Clean mathematical properties
- Good length extrapolation with extensions (YaRN, etc.)
- Efficient implementation (just element-wise multiplication after precomputing sin/cos)