Transformers
RNNs were the standard for sequences but had fundamental problems:
Transformers are a neural architecture built entirely on attention (no recurrence, no convolution). Introduced in “Attention Is All You Need” (Vaswani et al., 2017), they have become the dominant architecture in NLP and increasingly in vision, audio, and multimodal AI.
Table of Contents
- 1. Why Transformers?
- 2. The Original Architecture (Encoder-Decoder)
- 3. Components of a Transformer Layer
- 4. The Feedforward Network (FFN)
- 5. Positional Encoding
- 6. Multi-Head Self-Attention — Refresher
- 7. Putting It Together — Full Transformer Block
- 8. Parameter Count of a Transformer
- 9. The Three Flavors of Transformer
- 10. Training a Transformer
- 11. Inference — Autoregressive Generation
- 12. Variants and Optimizations
- 13. Tokenization
- 14. Scaling Laws
- 15. Strengths and Weaknesses
- 16. Transformer vs RNN vs CNN
- 17. Common Pitfalls
- 18. Production Considerations
- 19. WALKTHROUGH — Text Classification (Encoder-Only)
- 20. WALKTHROUGH — Translation (Encoder-Decoder)
- 21. The Three Attention Patterns
1. Why Transformers?
RNNs were the standard for sequences but had fundamental problems:
| Problem | RNN | Transformer |
|---|---|---|
| Sequential computation | Must process tokens one-by-one | All tokens in parallel |
| Long-range dependencies | Difficult (vanishing gradients) | Direct attention to any position |
| Training speed | Slow | Fast (GPU-friendly matmuls) |
| Scalability | Limited | Excellent — scales to billions of params |
Transformers solve all four by replacing recurrence with self-attention.
2. The Original Architecture (Encoder-Decoder)
The 2017 Transformer was an encoder-decoder model for machine translation:
Input sequence
│
▼
┌────────────────┐
│ ENCODER │ (N stacked layers)
│ (self-attn) │
└───────┬────────┘
│
▼ (encoder outputs)
┌────────────────┐
│ DECODER │ (N stacked layers)
│ (self-attn + │
│ cross-attn) │
└───────┬────────┘
│
▼
Output sequence
Modern variants:
- Encoder-only: BERT, RoBERTa — for understanding tasks (classification, NER)
- Decoder-only: GPT, LLaMA — for generation (LLMs)
- Encoder-decoder: T5, BART — for seq2seq tasks (translation, summarization)
3. Components of a Transformer Layer
A transformer encoder layer has two sublayers:
- Multi-Head Self-Attention
- Position-wise Feedforward Network (FFN)
Each sublayer is wrapped with residual connection + layer normalization:
\[\mathbf{x}' = \text{LayerNorm}(\mathbf{x} + \text{Sublayer}(\mathbf{x}))\]This is Post-LN (original). Modern Transformers usually use Pre-LN:
\[\mathbf{x}' = \mathbf{x} + \text{Sublayer}(\text{LayerNorm}(\mathbf{x}))\]Pre-LN is more stable and doesn’t require learning rate warmup.
Decoder layer adds a third sublayer
- Masked multi-head self-attention (causal)
- Cross-attention to encoder outputs
- Feedforward network
4. The Feedforward Network (FFN)
After attention mixes information across positions, the FFN processes each position independently:
\[\text{FFN}(\mathbf{x}) = \mathbf{W}_2 \cdot \text{ReLU}(\mathbf{W}_1 \mathbf{x} + \mathbf{b}_1) + \mathbf{b}_2\]- Inner dim is typically 4× the model dim ($d_{\text{ff}} = 4d$).
- Same FFN weights applied to every position (parameter sharing across positions, like CNNs across space).
- Modern variants use GELU or SwiGLU instead of ReLU.
Why FFN? Attention mixes information across positions but is linear (just weighted sums). The FFN adds nonlinear processing at each position — essential for expressive power. Roughly, attention = “communicate,” FFN = “think.”
Tip: Where the parameters live In a typical Transformer, about 2/3 of parameters are in FFN layers, not attention. Attention is the famous part, but FFN is doing a lot of heavy lifting.
5. Positional Encoding
Self-attention is permutation-equivariant — it doesn’t know token order. We add positional encodings to embeddings so the model knows position.
Sinusoidal (original)
\(PE_{(\text{pos}, 2i)} = \sin\left(\text{pos} / 10000^{2i/d}\right)\) \(PE_{(\text{pos}, 2i+1)} = \cos\left(\text{pos} / 10000^{2i/d}\right)\)
- Fixed (not learned)
- Allows extrapolation to longer sequences (in theory)
- Each dimension is a sinusoid of different frequency
Learned positional embeddings (BERT, GPT)
A learnable embedding vector for each position. Simple but doesn’t extrapolate beyond training length.
Rotary Position Embedding (RoPE) — modern (LLaMA, GPT-NeoX)
Rotates query/key vectors based on position. Naturally encodes relative position. State of the art.
ALiBi (Attention with Linear Biases)
Adds a position-dependent bias to attention scores. Simple, extrapolates well.
6. Multi-Head Self-Attention — Refresher
Recap from Attention:
\[\text{head}_i = \text{softmax}\left(\frac{\mathbf{Q}_i \mathbf{K}_i^\top}{\sqrt{d_k}}\right)\mathbf{V}_i\] \[\text{MultiHead}(\mathbf{X}) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)\mathbf{W}_O\]In a Transformer:
- Encoder self-attention: unmasked
- Decoder self-attention: masked (causal)
- Decoder cross-attention: unmasked, queries from decoder, keys/values from encoder
7. Putting It Together — Full Transformer Block
Encoder block (Pre-LN version):
x → LN → MHSA → + (residual) → x'
x' → LN → FFN → + (residual) → output
Decoder block:
x → LN → Masked-MHSA → + → x'
x' → LN → Cross-Attn(K,V from encoder) → + → x''
x'' → LN → FFN → + → output
Stack $N$ of these (typically $N = 6$–$96$).
8. Parameter Count of a Transformer
For model dim $d$, FFN dim $4d$, $h$ heads:
| Component | Params per layer |
|---|---|
| MHSA ($\mathbf{W}_Q, \mathbf{W}_K, \mathbf{W}_V, \mathbf{W}_O$) | $4 d^2$ |
| FFN ($\mathbf{W}_1, \mathbf{W}_2$) | $8 d^2$ |
| LayerNorm | $4d$ (small) |
| Total per layer | $\approx 12 d^2$ |
For $N$ layers + embeddings: $\approx 12 N d^2 + V \cdot d$ (where $V$ is vocab size).
Example: GPT-2 small ($d=768, N=12$): $\approx 12 \cdot 12 \cdot 768^2 \approx 85M$ params + embedding. Actual: 124M.
9. The Three Flavors of Transformer
Encoder-only (BERT, RoBERTa)
- Bidirectional self-attention (no mask)
- Trained with masked language modeling (predict masked tokens)
- For understanding tasks: classification, NER, QA
- Outputs: contextual embeddings for each token
Decoder-only (GPT, LLaMA, Claude, etc.)
- Causal (masked) self-attention only
- Trained with next-token prediction (autoregressive LM)
- For generation tasks
- Standard architecture for modern LLMs
Encoder-Decoder (T5, BART)
- Encoder reads input bidirectionally
- Decoder generates output autoregressively, cross-attending to encoder
- For seq2seq: translation, summarization, structured generation
10. Training a Transformer
Pre-training objectives
| Model type | Objective | Description |
|---|---|---|
| Decoder-only (GPT) | Causal LM | Predict next token given previous |
| Encoder-only (BERT) | Masked LM | Predict randomly masked tokens |
| Encoder-decoder (T5) | Span corruption | Predict masked spans |
Loss
Standard cross-entropy over vocabulary:
\[\mathcal{L} = -\frac{1}{T}\sum_{t=1}^{T} \log P(y_t \mid y_{<t}, \mathbf{x})\]Key training details
- Warmup + cosine decay for learning rate (critical for stability)
- AdamW optimizer (with $\beta_1 = 0.9, \beta_2 = 0.95$ for LLMs)
- Gradient clipping (typically max norm 1.0)
- Mixed precision (BF16) for speed
- Large batches (millions of tokens per batch for LLMs)
- Weight tying between input embedding and output projection (saves params)
11. Inference — Autoregressive Generation
For decoder-only models, generation is one token at a time:
tokens = [start_token]
for t in range(max_length):
logits = model(tokens)
next_token = sample(logits[-1]) # or argmax / top-k / top-p
tokens.append(next_token)
if next_token == eos:
break
Decoding strategies
| Method | How it picks next token | Property |
|---|---|---|
| Greedy | argmax | Deterministic, often repetitive |
| Beam search | Track top-k partial sequences | Better quality, slower |
| Sampling | Random per probability | Diverse but can be incoherent |
| Top-k sampling | Sample from top $k$ probable | Balanced |
| Top-p (nucleus) | Sample from smallest set with cumulative prob ≥ $p$ | Adaptive, popular |
| Temperature | Scale logits: $\text{logits}/T$ before softmax | $T<1$ sharpens, $T>1$ flattens |
KV Caching — Critical Optimization
At each generation step, we recompute attention over all previous tokens — but the keys and values for past tokens don’t change! Cache them:
- Without caching: $O(T^2)$ per step → $O(T^3)$ for full generation
- With caching: $O(T)$ per step → $O(T^2)$ total
KV caches consume significant memory in LLM serving — often the main inference bottleneck.
12. Variants and Optimizations
Long-context Attention
- Sparse attention (Longformer, BigBird): only attend to nearby + global tokens
- Linear attention (Linformer, Performer): kernel approximations, $O(T)$
- FlashAttention: IO-aware exact attention, much faster
- Sliding window (Mistral): local attention with limited range
Mixture of Experts (MoE)
Replace FFN with multiple “expert” FFNs, route each token to top-$k$ experts. More parameters with similar compute per token. Used in Mixtral, GPT-4 (rumored), DeepSeek.
Grouped-Query Attention (GQA) / Multi-Query Attention (MQA)
Share keys/values across multiple query heads. Reduces KV cache size significantly. Used in LLaMA-2, LLaMA-3.
Parameter-Efficient Fine-Tuning
- LoRA: low-rank updates to attention weights — fine-tune 0.1% of params
- Adapters: small inserted modules
- Prompt tuning: learn input prefixes
13. Tokenization
Transformers operate on tokens, not raw text. Methods:
| Tokenizer | Used by | How |
|---|---|---|
| Word-level | Old NLP | Each word → token. Huge vocab. |
| BPE (Byte-Pair Encoding) | GPT-2, RoBERTa | Merge frequent character pairs |
| WordPiece | BERT | Similar to BPE |
| SentencePiece | T5, LLaMA | Language-agnostic, subword |
| Tiktoken (BPE variant) | GPT-3/4 | Optimized BPE |
Subword tokenization handles rare words (split into known subwords) and keeps vocab manageable (typically 32k–100k).
14. Scaling Laws
Empirical finding: model performance follows power laws in:
- Parameters ($N$)
- Dataset size ($D$)
- Compute ($C$)
Larger models, more data, more compute → predictably better loss. Driving the LLM scaling race.
Chinchilla scaling law: For optimal compute budget, train smaller models on more data — many LLMs were “compute-inefficient” by being too large for their training data.
15. Strengths and Weaknesses
Strengths
- Parallel training over the whole sequence
- Long-range dependencies handled natively
- Scales beautifully with data and compute
- General-purpose — text, vision, audio, multimodal
- Transfer learning — pretrain once, fine-tune for many tasks
- Pretraining-friendly — billions of tokens
Weaknesses
- $O(T^2)$ attention — long context expensive (mitigated by FlashAttention, sparse variants)
- Memory-hungry at training and inference (KV cache)
- Data-hungry — needs huge corpora for good performance
- No inherent inductive bias for locality (must learn from data)
- Hallucinations in generation
- Inference latency for long generations
16. Transformer vs RNN vs CNN
| Aspect | RNN/LSTM | CNN | Transformer |
|---|---|---|---|
| Parallel training | No | Yes | Yes |
| Long-range deps | Hard | Limited | Easy |
| Inductive bias | Sequential | Locality | None |
| Compute per token | $O(d^2)$ | $O(k d^2)$ | $O(T d)$ |
| Memory | $O(T d)$ | $O(T d)$ | $O(T^2)$ |
| Streaming inference | Natural | Natural | Awkward (KV cache) |
| Default today | Niche | Vision | Most things |
17. Common Pitfalls
- No positional encoding → permutation-equivariant (treats inputs as a set)
- Wrong mask in decoder → model peeks at future tokens, gets impossibly good train loss, fails at inference
- No warmup → unstable training, loss diverges (especially with Post-LN)
- Forgetting to scale by $\sqrt{d_k}$ → vanishing gradients
- Wrong shape conventions —
(batch, seq, dim)vs(seq, batch, dim)(PyTorch nn.Transformer uses the latter by default!) - Tokenization mismatch between train and inference
- Not using KV cache during generation → very slow
- Massive KV cache running out of memory for long contexts
- Numerical instability in attention softmax with FP16 → use BF16 or stable softmax
18. Production Considerations
- Pretrained models are the starting point — fine-tune or prompt; rarely train from scratch
- Quantization (INT8, INT4, GGUF) for inference, especially LLMs
- Speculative decoding — small draft model proposes, big model verifies (2–3× speedup)
- FlashAttention for fast attention without quality loss
- Tensor parallelism / pipeline parallelism for very large models
- Distillation — train smaller student from larger teacher
- LoRA fine-tuning instead of full fine-tuning to save memory and compute
- Continuous batching in serving (vLLM, TensorRT-LLM) for throughput
19. WALKTHROUGH — Text Classification (Encoder-Only)
Task: Sentiment classification — given a sentence, predict positive/negative.
Example: "the movie was really good" → label: positive (1)
Model: Tiny encoder-only Transformer with 2 layers, $d=8$, 2 heads, FFN dim 32, vocab ~30k.
Step 1: Tokenization
"the movie was really good"
↓ (tokenizer)
[CLS] the movie was really good [SEP]
↓
[101, 1996, 3185, 2001, 2428, 2204, 102]
[CLS]at start — its final embedding will represent the whole sentence.[SEP]at end — sentence separator.- Sequence length: $T = 7$.
Step 2: Embedding Lookup
Each token ID is converted to a dense vector via an embedding table of shape (vocab_size, d) = (30000, 8).
Token IDs: [101, 1996, 3185, 2001, 2428, 2204, 102]
↓ embedding table
Embeddings: shape (7, 8)
Input to Transformer: $\mathbf{X} \in \mathbb{R}^{7 \times 8}$.
Step 3: Add Positional Encoding
Position embeddings (shape (7, 8)):
position 0: [0.1, -0.2, 0.0, ...]
position 1: [0.0, 0.3, 0.1, ...]
...
X = token_embedding + position_embedding
shape: (7, 8)
Now each row of $\mathbf{X}$ encodes both what the token is and where it is.
Step 4: First Encoder Layer
4a. LayerNorm: Normalize each row independently across 8 features.
4b. Multi-Head Self-Attention (2 heads, $d_k = 4$ each):
For each head:
Q = X @ W_Q → split into (7, 4) per head
K = X @ W_K → split
V = X @ W_V → split
Scores = Q @ K^T shape: (7, 7)
Scaled = Scores / sqrt(4)
Weights = softmax(Scaled) each row sums to 1
Output = Weights @ V shape: (7, 4)
Interpretation of the (7,7) weight matrix:
- The
[CLS]token (row 0) likely attends broadly to all tokens — it needs to summarize. "really"(row 4) might attend strongly to"good"(row 5) — modifier ↔ adjective.
Concatenate heads and project:
head_1: (7, 4)
head_2: (7, 4)
Concat: (7, 8)
@ W_O: (7, 8)
4c. Residual:
X' = X + attention_output shape: (7, 8)
4d. FFN (applied per row independently):
hidden = ReLU(x @ W_1 + b_1) shape (32,)
out = hidden @ W_2 + b_2 shape (8,)
Note: Same $\mathbf{W}_1, \mathbf{W}_2$ for every position. FFN does NOT mix info across tokens — only attention does that.
4e. Residual:
X'' = X' + FFN(X') shape: (7, 8)
Step 5: Stack Encoder Layer 2
Repeat with a second set of weights. After layer 2:
Final output: shape (7, 8)
Each row is now a contextualized embedding of that token — aggregated information from the whole sentence.
Step 6: Extract Classification Representation
cls_vector = final_output[0] # shape (8,)
The [CLS] token’s final embedding represents the whole sentence (attention let it aggregate info from all positions).
Step 7: Classification Head
logits = cls_vector @ W_cls + b_cls # shape (2,)
prob = softmax(logits)
For our example, hopefully prob ≈ [0.05, 0.95] → “positive.”
Step 8: Loss
true_label = 1 (positive)
loss = cross_entropy(prob, true_label)
Backprop through everything — token embeddings, position embeddings, attention weights, FFN weights, classifier.
Classification Pipeline Summary
text
→ tokenize
→ embedding lookup (T, d)
→ + positional encoding (T, d)
→ encoder layer × N (T, d)
→ take [CLS] embedding (d,)
→ classification head (num_classes,)
→ softmax → predict
20. WALKTHROUGH — Translation (Encoder-Decoder)
Task: English → French. Translate "the cat sat" → "le chat assis".
Model: 2 encoder layers, 2 decoder layers, $d=8$, 2 heads, FFN dim 32.
Step 1: Tokenize Both Sequences
Source (English):
"the cat sat" → [CLS] the cat sat [SEP]
→ [101, 1996, 4937, 2938, 102]
T_src = 5
Target (French) — Teacher Forcing:
Gold: "le chat assis"
Decoder input: [BOS] le chat assis (shifted right)
Decoder target: le chat assis [EOS] (what to predict)
T_tgt = 4
Important: Why “shifted right”? Position $t$ of the decoder’s output should predict target token $t$. So we feed the gold prefix as input and predict the next token at each position.
Step 2: Encoder — Process Source
Identical to classification’s encoder:
[CLS] the cat sat [SEP]
↓ embedding + positional
↓ encoder layer 1
↓ encoder layer 2
encoder_output: shape (5, 8)
Each row is a contextual embedding of the corresponding English token. The decoder will use these to know what to translate.
Step 3: Decoder Inputs
Decoder input: [BOS] le chat assis → [2, 105, 678, 423]
Embedding + positional → shape (4, 8)
Step 4: Decoder Layer (3 Sublayers)
4a. Masked Self-Attention
Compute Q, K, V from decoder input. Apply causal mask:
[1 0 0 0]
[1 1 0 0]
[1 1 1 0]
[1 1 1 1]
Set masked positions to $-\infty$ before softmax. This prevents position 2 ("chat") from peeking at position 3 ("assis").
Important: Why masking matters During training, the decoder sees
[BOS] le chat assisall at once (parallel for efficiency). But we want to simulate generation, where position 2 doesn’t know future tokens. The mask enforces this.
Output: (4, 8).
4b. Cross-Attention — Where Encoder Meets Decoder
This is the magic step.
Q = decoder_state @ W_Q shape: (4, 8)
K = encoder_output @ W_K shape: (5, 8)
V = encoder_output @ W_V shape: (5, 8)
Weights: Q @ K^T shape: (4, 5)
Softmax over English tokens.
Output: Weights @ V shape: (4, 8)
Interpretation:
- Row 0 (
[BOS]) attends to source — likely focuses on"the". - Row 1 (
"le") attends to"cat"to figure out next French word ("chat"). - Row 2 (
"chat") attends to"sat"to predict"assis". - Row 3 (
"assis") attends to[SEP]— translation winding down.
This is the soft alignment between source and target — solving the “context vector bottleneck” of old seq2seq.
4c. FFN
Position-wise, like in the encoder. Output: (4, 8).
(All three sublayers have residual + LayerNorm.)
Step 5: Stack Decoder Layer 2
Repeat. Final decoder output: shape (4, 8).
Step 6: Output Projection
logits = decoder_output @ W_vocab # shape (4, 30000)
probs = softmax(logits, dim=-1)
For each of the 4 positions, a probability distribution over the French vocab.
Step 7: Loss (Training Time)
Position 0 should predict "le" (token 105)
Position 1 should predict "chat" (token 678)
Position 2 should predict "assis" (token 423)
Position 3 should predict [EOS] (token 3)
loss = average cross-entropy over 4 positions
Backprop through encoder, decoder, all attention layers, embeddings, output projection.
Step 8: Inference (Generation)
At inference, we don’t have the gold target. We generate one token at a time:
Step 1:
Decoder input: [BOS]
→ forward pass → predict "le"
Step 2:
Decoder input: [BOS] le
→ forward pass → predict "chat"
Step 3:
Decoder input: [BOS] le chat
→ forward pass → predict "assis"
Step 4:
Decoder input: [BOS] le chat assis
→ forward pass → predict [EOS]
Stop. Final output: "le chat assis"
Important: Training vs Inference
- Training: all positions processed in parallel with teacher forcing.
- Inference: sequential, one token at a time. Each step does a full forward pass.
KV caching speeds this up: encoder runs once at start; decoder’s keys/values for past tokens are cached.
Translation Pipeline Summary
source text target text (training only)
→ tokenize → tokenize
→ embedding + pos → shift right + embedding + pos
→ encoder × N ↓
→ encoder_output ──────► cross-attn ◄── decoder × N (masked self-attn + cross-attn + FFN)
↓
→ output projection
→ softmax → token probs
21. The Three Attention Patterns
In a full encoder-decoder Transformer, there are three distinct attention layers, each with its own role:
| Where | Q from | K, V from | Masked? |
|---|---|---|---|
| Encoder self-attention | Source | Source | No |
| Decoder self-attention | Target (so far) | Target (so far) | Yes (causal) |
| Decoder cross-attention | Target | Encoder output | No |
What each one does
- Encoder self-attn: Source tokens talk to each other (resolves pronouns, captures syntax).
- Decoder self-attn: Target tokens talk to past target tokens (ensures fluent output language).
- Decoder cross-attn: Target tokens read from the source (the actual “translation” linking).
- FFN: Adds nonlinear processing at each position.
- Output projection: Converts representations to vocabulary probabilities.
Tip: What if you removed one?
- No cross-attn: Decoder becomes a language model that ignores the source.
- No decoder self-attn: Model generates each word without considering what it already wrote.
- No encoder self-attn: Source representations don’t contextualize each other.
Each piece plays a role.
Mental models
Classification: “Read the input, summarize it in the [CLS] token, classify.”
Translation: “Encoder builds a rich representation of the source. Decoder generates the target one token at a time, looking back at what it’s already written (self-attn, masked) and looking at the source (cross-attn).”
Related Notes
- Neural Networks
- CNNs
- RNNs
- Attention
- Vision Transformers (next)
- LLMs (upcoming)