« All Roles

LLM / Foundation-Model Engineer — Complete Learning Curriculum

Target Roles:

• Research Engineer, Pretraining (Anthropic, OpenAI, DeepMind, Meta, Mistral, xAI)
• LLM Infrastructure Engineer / ML Systems Engineer
• Foundation Model Engineer
• Post-training / Fine-tuning Engineer (RLHF, DPO, SFT)
• LLM Inference Engineer (vLLM/TGI/TensorRT-LLM class work)
• Model Evaluation Engineer
• Pretraining Data Engineer
• Applied AI / Production AI Engineer

Duration: 24 weeks core (6 months) — extendable to 12 months for deep specialization Goal: Reach interview-ready expertise with a portfolio competitive for senior LLM/foundation-model roles at frontier labs.

Why This Curriculum Exists

The hiring bar at frontier labs (Anthropic, OpenAI, DeepMind, Meta AI, Mistral, xAI, Cohere) is not "have you used ChatGPT" — it is "can you implement attention from scratch, debug a 64-GPU training run, profile a CUDA kernel, design a 100k-QPS inference gateway, and explain why DPO converges differently than PPO".

This curriculum is built backward from real job postings (referenced below) and is structured so that every lab maps to a real interview question or production system you would build on the job.

Reference Job Targets

• Anthropic — Research Engineer, Pretraining (JD) → Phases 4, 5, 10, Capstone 1
• Anthropic — Research Engineer, Production Model Post-Training → Phases 6, 8, Capstone 4
• OpenAI — Research Engineer, Applied AI (JD) → Phases 7, 9, Capstones 2 & 3
• Google DeepMind — Research Engineer, Gemini Latent Thinking → Phases 4, 5, 6, 8
• Meta AI — Research / Production roles (Careers) → Phases 5, 9, 10

What You Will Build

By the end of this curriculum you will have shipped:

• A working BPE tokenizer that matches GPT-2 output byte-for-byte
• Word2Vec, attention, and a transformer block — all from scratch in NumPy and PyTorch
• A nanoGPT-style model trained on a custom corpus (TinyStories or your own)
• A LoRA / QLoRA fine-tuning pipeline on an open 7B model
• A DPO preference-optimization run with reward analysis
• A production-grade RAG system with hybrid retrieval, re-ranking, and an eval harness
• An inference gateway with continuous batching, KV-cache, streaming, quantization, observability
• A pretraining data pipeline with deduplication (MinHash), quality filtering (FastText/heuristics), and tokenization at scale
• A multi-GPU training experiment using FSDP / DDP with mixed precision and gradient accumulation
• An evaluation harness comparing base, fine-tuned, and RAG-augmented models on MMLU/HellaSwag/HumanEval-style tasks
• A complete portfolio of 10+ GitHub repos with READMEs, benchmarks, diagrams, and ablations

Folder Structure

llm-inference-engineer/
├── README.md                              ← You are here (master roadmap)
├── phase-01-foundations-text/             ← Tokenization, BoW, TF-IDF, similarity, PyTorch
├── phase-02-classical-nlp-embeddings/     ← Word2Vec, GloVe, FastText, embedding eval
├── phase-03-rnns-language-modeling/       ← RNN/LSTM/GRU, char-LM, seq2seq, Bahdanau attention
├── phase-04-attention-transformers/       ← Self-attention, MHA, positional encodings, full transformer
├── phase-05-training-small-llms/          ← Mini-GPT, BPE, training loop, mixed precision, sampling
├── phase-06-finetuning-instruction/       ← SFT, LoRA/QLoRA, instruction data, RLHF/DPO
├── phase-07-rag-retrieval/                ← Vector DBs, hybrid search, re-ranking, agents/tool use
├── phase-08-evaluation-safety/            ← Eval harness, LLM-as-judge, red-teaming, benchmarks
├── phase-09-inference-optimization/       ← KV-cache, quantization, batching, vLLM/TGI, spec decoding
├── phase-10-distributed-production/       ← DDP/FSDP, pretraining data pipeline, observability
├── phase-11-capstone/                     ← 4 portfolio-grade end-to-end systems
├── system-design/                         ← LLM-specific system design walkthroughs
└── interview-prep/                        ← Concepts, coding, ML systems, behavioral

24-Week Schedule

Each Lab Structure

Every lab folder contains:

Project Specification Template

Every non-trivial project in this curriculum is described with the same template, so you can lift any lab into a portfolio-ready repo:

The phase READMEs (phase-XX/README.md) instantiate this template for every lab.

Prerequisites

• Python 3.10+
• Comfort with backend / distributed systems (you have this)
• Basic linear algebra (matrix multiply, eigenvectors) — Phase 1 has a refresher
• A Hugging Face account (free) for model + dataset access
• Optional: Weights & Biases / Comet ML account for experiment tracking

Hardware Recommendations

You do NOT need a GPU cluster. Every lab in this curriculum has a "small-model mode" that runs on Colab free tier. Capstones can be completed for under $50 of cloud GPU time.

System Design Philosophy

Every production-oriented lab (Phases 7, 9, 10) is evaluated on the same five axes that frontier-lab interviewers care about:

1. Throughput — tokens/sec at the system level (not just the model)
2. Latency — TTFT (time-to-first-token) and TPOT (time-per-output-token), P50/P99
3. Memory efficiency — KV-cache size, activation memory, parameter offloading
4. Cost — $/million-tokens served, $/training-run, GPU-hour utilization
5. Observability — request tracing, token-level metrics, drift detection, eval-in-production

Each capstone explicitly reports numbers on these axes.

Phase-by-Phase Overview

Each phase has its own README.md with full lab specs, concept list, deliverables, and interview questions. Below is the index — click into the phase folder for depth.

Phase 1 — Foundations: Text, Math, PyTorch

Concepts: Tokenization (whitespace → regex → byte-level), bag-of-words, TF-IDF, cosine similarity, PyTorch tensors/autograd, broadcasting, CPU/GPU dispatch. Difficulty: ⭐⭐☆☆☆ | Time: 1–2 weeks Deliverables: From-scratch TF-IDF search engine over a Wikipedia subset; PyTorch tensor playground notebook. Roles supported: All — this is non-negotiable foundation.

Phase 2 — Classical NLP & Static Embeddings

Concepts: Word2Vec (CBOW + skip-gram), negative sampling, GloVe, FastText subword, embedding evaluation (analogy, WordSim353), dimensionality reduction. Difficulty: ⭐⭐⭐☆☆ | Time: 1.5 weeks Deliverables: Skip-gram trained from scratch on text8; embedding visualization (t-SNE/UMAP); analogy benchmark report. Roles supported: Pretraining Data Engineer, Research Engineer.

Phase 3 — RNNs & Language Modeling

Concepts: Vanilla RNN forward/backward, vanishing gradients, LSTM gates, GRU, sequence-to-sequence, Bahdanau additive attention, teacher forcing, perplexity. Difficulty: ⭐⭐⭐☆☆ | Time: 1.5 weeks Deliverables: Char-RNN trained on Shakespeare; LSTM seq2seq translator (toy). Roles supported: Foundation Model Engineer (historical context); strong "explain attention" interview answer.

Phase 4 — Attention & Transformers (From Scratch)

Concepts: Scaled dot-product attention, masking (causal/padding), multi-head, sinusoidal/RoPE/ALiBi positional encodings, layer norm vs RMSNorm, residual streams, encoder/decoder/decoder-only. Difficulty: ⭐⭐⭐⭐☆ | Time: 2 weeks Deliverables: 200-line transformer that passes attention shape tests; visualized attention maps; ablation report (pre-norm vs post-norm). Roles supported: All research-engineer roles. The most-asked interview topic.

Phase 5 — Training Small LLMs

Concepts: BPE tokenization (matching GPT-2), nanoGPT architecture, AdamW, cosine LR schedule, mixed precision (BF16/FP16), gradient accumulation, gradient clipping, checkpointing, sampling (greedy/top-k/top-p/temperature/beam). Difficulty: ⭐⭐⭐⭐☆ | Time: 2.5 weeks Deliverables: BPE tokenizer matching tiktoken on test corpus; nanoGPT trained on TinyStories with W&B logs and loss curves. Roles supported: Research Engineer Pretraining, Foundation Model Engineer.

Phase 6 — Fine-tuning, Instruction Tuning, Preference Optimization

Concepts: SFT, chat templates, LoRA / QLoRA (NF4), reward modeling, RLHF (PPO conceptual), DPO / IPO / KTO, RLAIF, constitutional AI. Difficulty: ⭐⭐⭐⭐☆ | Time: 2.5 weeks Deliverables: QLoRA fine-tune of Llama-3-8B or Qwen2-7B on a domain dataset; DPO run with preference dataset; before/after eval table. Roles supported: Post-training Engineer, Production Model Post-Training (Anthropic-style).

Phase 7 — RAG, Retrieval, Agents

Concepts: Embedding models (sentence-transformers, E5, BGE), FAISS vs HNSW vs IVF, hybrid retrieval (BM25 + dense), re-ranking (cross-encoder, ColBERT), chunking strategies, query rewriting, agent loops, tool use, structured output (JSON schema, constrained decoding). Difficulty: ⭐⭐⭐⭐☆ | Time: 2 weeks Deliverables: Production-style RAG over a real corpus with eval (RAGAS); agent that uses 3+ tools. Roles supported: Applied AI Engineer (OpenAI-style), LLM Inference Engineer.

Phase 8 — Evaluation & Safety

Concepts: Benchmarks (MMLU, HellaSwag, GSM8K, HumanEval, IFEval, MT-Bench), perplexity vs downstream eval, LLM-as-judge bias, RAGAS, red-teaming, jailbreak taxonomy, safety classifiers. Difficulty: ⭐⭐⭐⭐☆ | Time: 1.5 weeks Deliverables: Forked lm-evaluation-harness task; LLM-as-judge harness with bias analysis; red-team report. Roles supported: Model Evaluation Engineer, Safety roles.

Phase 9 — Inference Optimization & Serving

Concepts: KV-cache mechanics + memory math, paged attention (vLLM), continuous batching, INT8/INT4 quantization (GPTQ, AWQ, bitsandbytes), speculative decoding, prefix caching, FlashAttention-2/3, CUDA graphs, TensorRT-LLM, streaming via SSE. Difficulty: ⭐⭐⭐⭐⭐ | Time: 2.5 weeks Deliverables: Custom inference server with KV-cache + continuous batching + INT4 quantization; benchmark report (TTFT/TPOT/throughput). Roles supported: LLM Inference Engineer, ML Systems Engineer. Highest-leverage phase for infrastructure roles.

Phase 10 — Distributed Training & Pretraining Data

Concepts: DDP, FSDP, ZeRO-1/2/3, tensor/pipeline parallelism (conceptual), mixed precision strategies, NCCL, gradient checkpointing, activation recomputation, MinHash dedup, quality filtering (perplexity, FastText, heuristics), tokenization at scale, Common Crawl pipeline. Difficulty: ⭐⭐⭐⭐⭐ | Time: 2 weeks Deliverables: 2-GPU FSDP training run (rentable for ~$5); pretraining data pipeline processing 10 GB → deduped + tokenized shards. Roles supported: Pretraining Data Engineer, ML Infrastructure Engineer, Research Engineer Pretraining.

Phase 11 — Capstone Projects

Four portfolio-grade systems. Pick at least 2 to ship publicly.

1. Mini-GPT pretrained on a custom corpus (your dataset, full pipeline, model card)
2. Production RAG with eval (hybrid retrieval, RAGAS, A/B harness)
3. LLM inference gateway (KV-cache, batching, quantization, streaming, observability)
4. Domain-assistant fine-tune (SFT + DPO + eval comparison vs base)

The Top 10 Projects to Prioritize (Resume-Critical)

These are the projects that, when present on a portfolio, change interview outcomes:

Top 20 Interview Questions (Curated)

Full answers in interview-prep/01-concepts-cheatsheet.md.

1. Derive scaled dot-product attention. Why divide by √d_k?
2. Explain causal masking. Implement it in 5 lines.
3. Compare sinusoidal, learned, RoPE, and ALiBi positional encodings.
4. Why does pre-norm train more stably than post-norm?
5. Walk through one forward + backward pass of a transformer block.
6. Explain KV-cache. What is its memory footprint? When does it become the bottleneck?
7. Compare LoRA, QLoRA, full fine-tuning. When would you use each?
8. Explain DPO derivation. Why does it not need a separate reward model?
9. Compare PPO and DPO. Pros and cons.
10. Explain ZeRO-1/2/3 and FSDP. What does each shard?
11. What is continuous batching? How does paged attention enable it?
12. Compare INT8 and INT4 quantization (GPTQ vs AWQ vs bitsandbytes NF4).
13. Speculative decoding — explain the algorithm and the speedup math.
14. Compare BM25, dense retrieval, ColBERT, and a cross-encoder re-ranker.
15. How would you build an LLM eval pipeline that catches regressions in prod?
16. Design a RAG system for 100M documents at 1k QPS.
17. Design an LLM inference gateway for 100k QPS with multi-model routing.
18. Walk through a pretraining data pipeline: filtering, dedup, tokenization, sharding.
19. Why is BPE the dominant tokenizer? What are its failure modes?
20. Explain mixed precision (BF16 vs FP16) and loss scaling.

A Recommended Learning Order

1 → 2 → 3 → 4  (theory + scratch builds — sequential, no skipping)
        ↓
        5  (training mechanics — sequential)
        ↓
        ├── 6  (fine-tuning) ──┐
        ├── 7  (RAG)        ──┼──> 8 (evaluation ties everything together)
        └── 9  (inference)  ──┘
                ↓
                10 (distributed) → 11 (capstones)

You can swap the order of 6 / 7 / 9 based on the role you're targeting.

Job Titles to Search For

Use these exact strings on LinkedIn / Greenhouse / Ashby / company career pages:

• "Research Engineer, Pretraining"
• "Research Engineer, Post-Training"
• "Research Engineer, Applied AI"
• "Foundation Model Engineer"
• "LLM Infrastructure Engineer"
• "ML Systems Engineer (LLM)"
• "LLM Inference Engineer"
• "ML Performance Engineer"
• "Machine Learning Engineer, Generative AI"
• "Model Evaluation Engineer"
• "AI Safety Engineer"
• "Pretraining Data Engineer"
• "Member of Technical Staff" (used by Anthropic, OpenAI, Mistral)

Skill Checklist — "Am I Ready to Apply?"

Apply when you can honestly check ✅ on at least 80% of these:

Theory

• Derive attention end-to-end on a whiteboard
• Implement multi-head attention from scratch in <50 lines
• Explain RoPE rotation math
• Compare LayerNorm vs RMSNorm and justify modern choice
• Explain KV-cache memory math
• Derive DPO loss from RLHF objective
• Explain LoRA's rank decomposition and why it works
• Compute the parameter count of a transformer given d_model, n_layers, n_heads, vocab_size

Engineering

• Train a transformer from scratch end-to-end
• Fine-tune a 7B+ model on a single 24GB GPU using QLoRA
• Run a multi-GPU FSDP training job
• Build a RAG system with hybrid retrieval and re-ranking
• Quantize a model to INT4 and measure quality regression
• Implement continuous batching for an inference server
• Build a pretraining data pipeline with MinHash dedup

Portfolio

• 8+ public GitHub repos with READMEs, benchmarks, diagrams
• At least 1 project with reproducible training run + W&B logs
• At least 1 project with profiling output (Nsight, PyTorch profiler)
• A blog post or technical writeup of one capstone
• A resume with quantified, LLM-specific bullets

6-Month Plan (Aggressive, ~15 hr/week)

12-Month Plan (Deeper, ~10 hr/week — recommended for career switchers)

Same as above but each month covers half the content; the extra months go to:

• Months 7–8: CUDA fundamentals + Triton kernels (write a fused softmax)
• Months 9–10: One frontier-paper reimplementation (FlashAttention, Mixture-of-Experts, Mamba)
• Months 11–12: Capstone polish, blog posts, open-source contributions to vLLM / TGI / Transformers / lm-eval-harness

GitHub Portfolio Structure (Recommended)

your-github/
├── llm-from-scratch/                   ← Phases 1–4 in one repo (educational)
│   ├── 01-tokenization/
│   ├── 02-word2vec/
│   ├── 03-rnn-lstm/
│   └── 04-transformer/
├── nanogpt-tinystories/                ← Phase 5 capstone (single repo, polished)
├── qlora-domain-assistant/             ← Phase 6 capstone with eval
├── rag-production/                     ← Phase 7 capstone, full README + diagrams
├── llm-inference-gateway/              ← Phase 9 capstone (the hire-magnet)
├── lm-eval-harness-extension/          ← Phase 8 — contribute to upstream
├── pretraining-data-pipeline/          ← Phase 10
└── blog/                               ← MDX or plain markdown — link from each repo

Each Repo's README Should Have

1. One-sentence pitch above the fold
2. Architecture diagram (Excalidraw, Mermaid, or draw.io PNG)
3. Benchmarks table (numbers > prose)
4. Reproduction steps (make train, make eval)
5. Tradeoffs section — why you chose X over Y
6. Limitations — shows engineering maturity
7. What I'd do next — shows extensibility thinking

Resume Bullet Patterns

Use the action → system → quantified outcome → technical depth pattern:

"Built an LLM inference gateway supporting continuous batching, paged KV-cache, and INT4 GPTQ quantization, achieving 3.2× throughput improvement (412 → 1,317 tok/s) and 41% lower P99 TTFT on Llama-3-8B at 32 concurrent requests."

"Implemented a MinHash-LSH deduplication and FastText quality-filtering pipeline processing 180 GB of CommonCrawl WET shards into 41 GB of training-ready tokens, with reproducible Snakemake DAG and per-shard quality histograms."

"Pre-trained a 42M-parameter decoder-only transformer from scratch on TinyStories using a custom BPE tokenizer matching GPT-2, mixed precision, gradient accumulation, and cosine LR schedule on a single A100; achieved train loss 1.42 / val 1.51 in 4.2 GPU-hours."

Tools & Technologies Covered

Languages:        Python 3.11+, shell, basic CUDA/Triton overview
Core ML:          PyTorch 2.x, NumPy
Models / Libs:    Hugging Face transformers, datasets, accelerate, peft, trl
Tokenizers:       tiktoken, sentencepiece, hf-tokenizers
Training:         Lightning / pure PyTorch, FSDP, DeepSpeed (overview), bitsandbytes
Fine-tuning:      LoRA, QLoRA, DPO/IPO/KTO via trl
Retrieval:        FAISS, Qdrant, pgvector, sentence-transformers, BM25 (rank_bm25)
Eval:             lm-evaluation-harness, RAGAS, MT-Bench, HELM concepts
Inference:        vLLM, TGI, llama.cpp, TensorRT-LLM (overview), ONNX Runtime
Serving:          FastAPI, Uvicorn, Triton Inference Server (overview)
Observability:    OpenTelemetry, Prometheus, Grafana, Langfuse, W&B
Data:             pyspark / dask / polars, datasketch (MinHash), fasttext
Hardware:         CUDA, NCCL, BF16/FP16, A100/H100/L4/T4, AWQ/GPTQ

Quick Start

# 1. Navigate to the curriculum root
cd /path/to/llm-inference-engineer

# 2. Create a virtual environment
python -m venv .venv && source .venv/bin/activate

# 3. Install Phase 1 deps and start
pip install -r phase-01-foundations-text/lab-01-tokenization-from-scratch/requirements.txt
code phase-01-foundations-text/README.md

Mindset: You are not learning LLMs as an end. You are learning them well enough to build, debug, and ship the systems that frontier labs hire for. Every lab in this curriculum was designed by working backward from a real interview loop or a real production system. Do the work, ship the repos, and apply.

Phase 1 — Foundations: Text, Math, PyTorch

Difficulty: ⭐⭐☆☆☆ | Estimated Time: 1–2 weeks Roles supported: All — non-negotiable foundation.

Why This Phase Exists

Every modern LLM stack — from FlashAttention to vLLM — is built on three things: (1) representing text as numbers, (2) doing linear algebra on those numbers efficiently, and (3) using PyTorch's autograd to learn the parameters. If you cannot tokenize a string, build a TF-IDF index, or write a clean PyTorch nn.Module, the rest of the curriculum will collapse under you.

This phase rebuilds the floor.

Concepts

• Text representation: characters → words → subwords
• Tokenization: whitespace, regex, byte-level (BPE preview)
• Vocabulary construction, OOV handling, special tokens
• Bag-of-words (BoW) and term-document matrices
• TF-IDF derivation and intuition
• Cosine similarity, Euclidean distance, dot-product retrieval
• Sparse vs dense vector representations
• PyTorch tensors, broadcasting, indexing
• Autograd: forward, backward, .grad, .detach(), .no_grad()
• CPU/GPU dispatch, .to(device), pinned memory basics
• Linear algebra refresher: matmul, transpose, einsum, eigendecomposition

Labs

Lab 01 — Tokenization From Scratch

Lab 02 — Bag-of-Words & TF-IDF From Scratch

Lab 03 — Cosine Similarity & Retrieval Playground

Lab 04 — PyTorch Essentials & Autograd

Deliverables Checklist

• Three tokenizers (whitespace / regex / byte) with round-trip tests
• TF-IDF search engine over 10k docs, validated against sklearn
• Pairwise-distance visualization notebook
• Linear regression in pure PyTorch with autograd

Interview Relevance

• "How does TF-IDF differ from a dense embedding retrieval?" (you can answer both)
• "Walk me through autograd."
• "What does .detach() do?"
• "Why is byte-level tokenization useful?"

Warmup Guide — Foundations: Text

Zero-to-expert primer for Phase 01. How raw bytes become the token IDs a language model consumes — Unicode, encodings, and the BPE algorithm you will implement from scratch — assuming only basic Python.

Table of Contents

• Chapter 1: Text Is Not Strings — Bytes, Code Points, Graphemes
• Chapter 2: Why Models Need Tokens at All
• Chapter 3: The Tokenization Design Space
• Chapter 4: Byte-Pair Encoding — The Algorithm
• Chapter 5: Byte-Level BPE — What GPT Actually Does
• Chapter 6: Vocabularies, Special Tokens, and the Embedding Contract
• Chapter 7: Tokenization Pathologies Every Inference Engineer Meets
• Lab Walkthrough Guidance
• Success Criteria
• Interview Q&A
• References

Chapter 1: Text Is Not Strings — Bytes, Code Points, Graphemes

From zero: three distinct layers, conflated at your peril:

• Bytes: what disks and networks carry. "é" is 0xC3 0xA9 in UTF-8 — two bytes.
• Code points: Unicode's atomic units, U+0000–U+10FFFF. é is U+00E9 — or, equally validly, e (U+0065) + combining acute (U+0301): two code points that render identically. Normalization forms (NFC composes, NFD decomposes) exist to canonicalize this — and whether a tokenizer normalizes changes its output.
• Graphemes: what a human calls "one character." 👩‍👩‍👧 is one grapheme, seven code points (three emoji + two zero-width joiners), eighteen UTF-8 bytes.

UTF-8, the encoding that won: variable-length (1–4 bytes per code point), ASCII is valid UTF-8 unchanged, no byte-order issues, self-synchronizing (continuation bytes are distinguishable from start bytes — you can find a character boundary from any offset). The practical law: all text at system boundaries is UTF-8 bytes; decode explicitly; never trust a default encoding (Phase 05 of the PMC track meets the same law as a SpotBugs pattern — it's universal).

Chapter 2: Why Models Need Tokens at All

A neural LM consumes a sequence of integers from a fixed vocabulary, each mapped to an embedding vector. The tokenizer is the bridge: text → [int]. The design tension:

• Character/byte level (vocab ~256): no out-of-vocabulary problem ever — but sequences are ~4× longer than word-level, and attention cost grows quadratically with length (the model-accuracy track's Phase 07 math). Long sequences also dilute what each position can learn.
• Word level (vocab ~1M): short sequences, but unbounded vocabulary — every typo, inflection, and new name is out-of-vocabulary (OOV); embedding tables explode; and morphology is invisible (run/running unrelated).
• Subword (vocab 32K–256K): the resolution — frequent strings become single tokens (the, ing, tion), rare words decompose into pieces (tokenization → token|ization). No OOV (worst case: decompose to bytes), bounded vocab, morphology partially visible. Every modern LLM lives here.

The compression view is the deepest framing: a tokenizer is a learned compression codec for its training distribution — typical English compresses to ~0.75 tokens/word in GPT-class vocabularies. Everything about cost (context windows, API pricing, latency) is denominated in this codec's output.

Chapter 3: The Tokenization Design Space

The three families (know all; implement one):

• BPE (Byte-Pair Encoding): bottom-up greedy merging of frequent pairs (Chapter 4). GPT-2/3/4, LLaMA, most modern models. Deterministic, simple, fast.
• WordPiece (BERT): like BPE but merges the pair maximizing likelihood gain (score = pair_count / (left_count × right_count)) rather than raw frequency — favors pairs that are informative, not just common. Continuation pieces marked ##.
• Unigram LM (SentencePiece's default mode): top-down — start with a huge candidate vocabulary, iteratively remove pieces whose removal least hurts corpus likelihood under a unigram model; tokenization is then the Viterbi-best segmentation. Probabilistic, supports sampling multiple segmentations (subword regularization).
• SentencePiece the library (often confused with the algorithm): treats input as a raw code-point stream (no pre-splitting on whitespace — language-agnostic; spaces become ▁), and implements both BPE and Unigram.

Chapter 4: Byte-Pair Encoding — The Algorithm

Training (what your lab implements):

1. Start with a base vocabulary (all bytes, or all characters in the corpus).
2. Count all adjacent symbol pairs in the corpus.
3. Merge the most frequent pair into a new symbol; add it to the vocab; record the merge rule (A, B) → AB in order.
4. Repeat until vocab reaches target size (the merge list — ordered — is the model).

Encoding new text: split to base symbols, then apply the merge rules in training order (not greedy-longest-match!) — at each step, find the present pair that was learned earliest and merge all its occurrences. This ordering rule is the #1 implementation bug: greedy longest-match produces different (wrong) tokenizations that silently disagree with the reference.

Complexity reality: naive training recounts all pairs each merge — O(merges × corpus); fine for the lab. Production trainers maintain pair counts incrementally with priority queues. Encoding is near-linear with the right data structure (linked-list of symbols + a heap of candidate merges).

The pre-tokenization detail that matters: GPT-class BPE first splits text with a regex (on whitespace/letter/number/punctuation boundaries, keeping the leading space attached to the word: " world" is one pre-token). Merges never cross pre-token boundaries. This is why "world" and " world" are different tokens with different IDs — and why prompts that end with a trailing space produce subtly worse completions (you've stranded the model off its learned distribution).

Chapter 5: Byte-Level BPE — What GPT Actually Does

Character-level base vocabularies still have an OOV problem (a new Unicode character). GPT-2's move: run BPE over bytes. Base vocab = 256, every possible input is representable, full stop. One wrinkle: raw bytes include unprintable values that make merge files unreadable, so GPT-2 maps each byte to a printable Unicode proxy character (the famous Ġ is the proxy for the space byte 0x20). When you see Ġworld in a vocab dump, you're reading "space + world" through that proxy alphabet.

Cost of byte-level: non-Latin scripts pay heavily — a Chinese character is 3 UTF-8 bytes, and with few learned merges for it, tokens-per-character ratios for non-English text run 2–4× English's. This tokenizer tax is why multilingual models train larger vocabularies (LLaMA-3: 128K; Qwen: 152K) — more merges to amortize the world's scripts — and it's directly an inference-cost and context-budget issue you will own (Chapter 7).

Chapter 6: Vocabularies, Special Tokens, and the Embedding Contract

• The tokenizer's output range [0, V) is the contract with the model's embedding table (V × d_model) and output head. Mismatch = garbage or crashes; this is why tokenizer files ship with checkpoints and why "I'll just use a different tokenizer" is never a thing.
• Special tokens are IDs reserved outside BPE: <|endoftext|> / <s> </s> (BOS/EOS), padding, and — in chat models — role markers (<|im_start|> etc.). Two properties matter operationally: they must never be producible by encoding user text (else prompt injection by literally typing the marker — tokenizers have an explicit "special tokens are not encoded from text" path for this), and chat templates (the exact arrangement of role markers) are part of the model contract — a wrong template silently degrades quality with no error anywhere.
• Vocab size trade-off: larger V = shorter sequences (cheaper attention) but a bigger embedding/output matrix (for small models, the embedding table can be >30% of all parameters) and rarer tokens train on fewer examples. 32K (LLaMA-2) → 128K (LLaMA-3) reflects multilingual pressure winning that argument at scale.

Chapter 7: Tokenization Pathologies Every Inference Engineer Meets

The debugging bestiary — each of these will be a production ticket someday:

• Trailing-space prompts: "The answer is " ends mid-pre-token; the model sees a distribution it rarely trained on. Symptom: oddly worse completions; fix: end prompts at token boundaries.
• Token-boundary string operations: truncating a prompt by characters can split a token (or a multi-byte character); always truncate in token space, decode, re-check.
• Streaming partial-token display: a generated token can be half a multi-byte character; decoders must buffer incomplete UTF-8 sequences (every streaming API bug report eventually traces here).
• Numbers: BPE chunks digits inconsistently (1234 may be 12|34, 7,000 three tokens) — part of why arithmetic is hard for LLMs; newer models force single-digit tokenization.
• The SolidGoldMagikarp class: tokens present in the vocab but nearly absent from training data have untrained embeddings — feeding them produces erratic behavior. Vocab and training data must be curated together.
• Counting: "how many tokens is this?" depends on the exact tokenizer+version; budget enforcement with the wrong tokenizer over/under-counts by 20%+ across languages. Use the model's own tokenizer, always.

Lab Walkthrough Guidance

Lab 01 — Tokenization from Scratch (build BPE end-to-end):

1. Implement byte-level base splitting + the pre-tokenization regex first; unit-test against known pre-token boundaries (" world" stays whole).
2. Training loop: pair counting → merge → repeat. Verify on a tiny corpus by hand (5 merges you can compute on paper) before scaling.
3. Encoding with merge-order priority (not longest-match) — test: your tokenizer's output must match tiktoken/HF reference token-for-token on a sample corpus; any divergence is the ordering bug until proven otherwise.
4. Decoding + round-trip property test: decode(encode(s)) == s for arbitrary Unicode (emoji, CJK, ZWJ sequences) — this catches byte-proxy and UTF-8 buffering mistakes.
5. Then the analytics: tokens-per-word across English/code/CJK samples — reproduce Chapter 5's tokenizer-tax observation with your own numbers.

Success Criteria

You are ready for Phase 02 when you can, from memory:

1. Distinguish bytes / code points / graphemes with the é and emoji examples, and state what NFC/NFD change.
2. Argue subword tokenization from both failure modes it resolves (char-level length, word-level OOV).
3. Write the BPE training loop in pseudocode and state the encode-time merge-ordering rule and why greedy-longest is wrong.
4. Explain Ġ, why world ≠ world, and the trailing-space pathology.
5. Name the special-token security property and why chat templates are model contract.
6. Quantify the tokenizer tax and its two production consequences (cost, context budget).

Interview Q&A

Q: Why do all modern LLMs use subword tokenization? It's the only point in the design space that bounds the vocabulary (so embedding tables are trainable), eliminates OOV (worst case decomposes to bytes), and keeps sequences ~4× shorter than character-level — which matters quadratically because of attention. Frequent strings get dedicated capacity; rare strings get compositional treatment.

Q: Your user reports the API "cut off their prompt mid-word." What happened? Truncation done in token space (correct) but displayed expectations in character space — or worse, truncation in character space splitting a multi-byte character/token. The fix is truncating by tokens with the model's own tokenizer, decoding the kept prefix, and surfacing the token count to the user. Bonus point: leading-space attachment means the visible "word" boundary and the token boundary genuinely differ.

Q: Why is the same text 3× more tokens in Thai than English on GPT-2's tokenizer? GPT-2's merges were learned on overwhelmingly English bytes; Thai gets few merges, so text decomposes to near-raw UTF-8 bytes — 3 bytes/char × ~1 token/byte. Consequences: 3× cost, 3× context consumption, and worse modeling (fewer chars per attention span). That's why multilingual models retrain larger vocabularies rather than reuse GPT-2's.

Q: What breaks if a user can type <|im_start|> and it encodes to the real special token? Role injection: the user's message can impersonate the system/assistant turn, overriding instructions — prompt injection at the tokenizer layer, below any application filtering. Correct tokenizers only emit special-token IDs from explicit API parameters, never from encoding user text; verifying that property is part of deploying any new tokenizer.

References

• Sennrich et al., Neural Machine Translation of Rare Words with Subword Units (2016) — arXiv:1508.07909 — BPE's introduction to NLP
• Radford et al., Language Models are Unsupervised Multitask Learners (GPT-2, 2019) — §2.2 for byte-level BPE
• Kudo & Richardson, SentencePiece (2018) — arXiv:1808.06226; Kudo, Subword Regularization (2018) — arXiv:1804.10959 for Unigram
• tiktoken — read the educational _educational.py BPE implementation
• Karpathy, Let's build the GPT Tokenizer — youtube.com/watch?v=zduSFxRajkE — the single best companion to this lab
• Unicode Standard Annex #15: Normalization Forms
• Rumbelow & Watkins, SolidGoldMagikarp (2023) — the glitch-token investigation

🛸 Hitchhiker's Guide — Phase 1: Foundations (Text, Math, PyTorch)

Read this if: You have never built an ML model, or you have but you're shaky on tokenization, autograd, or why cosine ≡ dot product on normalized vectors. By the end you should be able to explain every line of phase-01-foundations-text/lab-01-tokenization-from-scratch/solution.py to a stranger and know why every choice is made.

0. The 30-second mental model

A modern LLM is just three nested operations:

1. Tokenize: turn a string into a list of integers (token IDs).
2. Embed + Transform: look up an embedding vector per ID, then run a stack of matmul + nonlinearity layers.
3. Predict + Sample: produce a probability distribution over the next token, sample from it, append, repeat.

Phase 1 covers (1) and the linear-algebra + PyTorch substrate that (2) and (3) need. Everything else in the curriculum is built on top of these primitives.

1. Prerequisite knowledge

If any bullet looks unfamiliar, knock it out first.

1.1 Python (the floor)

You need fluency, not mastery. Specifically:

• Data structures: list, dict, set, tuple, collections.Counter, collections.defaultdict. Big-O of each.
• Iteration: for/while, comprehensions, generators (yield), itertools (chain, islice, groupby).
• Functions: positional vs keyword, *args/**kwargs, lambdas, decorators, functools.lru_cache.
• OOP: classes, __init__, __call__, __len__, __getitem__, dataclasses, @property.
• Files & I/O: context managers (with open(...)), pathlib, JSON/CSV.
• Typing: from __future__ import annotations, list[int], Optional, TypedDict, Protocol.
• Performance hygiene: vectorize with NumPy, avoid Python for over arrays, profile with cProfile/line_profiler.

References:

• Fluent Python, 2nd ed., Luciano Ramalho — the single best Python book for ML engineers.
• Python Speed/Performance Tips
• Real Python's generators tutorial

1.2 The shell, git, and an editor

• Bash basics (grep, find, xargs, redirection, pipes).
• git: branch / commit / rebase / cherry-pick / bisect / reflog.
• VS Code or Neovim with Python LSP, debugger configured, ability to set breakpoints.
• tmux or screen for long-running training jobs.

1.3 NumPy

NumPy is the lingua franca. If you can't think in arrays, you can't think in tensors.

• np.array, dtype, shape, stride.
• Broadcasting — read the rules until they are reflexes.
• Slicing, fancy indexing, boolean masks.
• np.einsum (your eventual best friend; covered in §3.4).
• Random: seeded Generator (not legacy np.random.*).

References:

• Python for Data Analysis, 3rd ed., Wes McKinney
• From Python to Numpy by Nicolas Rougier — free, deep
• 100 NumPy exercises

1.4 Math you actually need

You don't need a PhD. You need:

Books that won't waste your time:

• Math: Mathematics for Machine Learning (Deisenroth, Faisal, Ong) — free PDF, read Ch. 2–6.
• Probability: Introduction to Probability (Blitzstein, Hwang) — first 7 chapters.
• Linear algebra (geometric intuition): 3Blue1Brown's Essence of Linear Algebra YouTube series. Watch all 15 videos. Mandatory.
• Calculus (intuition): 3Blue1Brown's Essence of Calculus.

2. Concept 1 — Text → Numbers (Tokenization)

2.1 Why we tokenize at all

Neural networks ingest tensors of numbers. We must convert text (a sequence of Unicode codepoints) into a sequence of small integer IDs that index into an embedding table (a learned matrix E ∈ ℝ^{V × d}). The choice of what counts as a "token" is the most consequential design decision in NLP. It determines:

• Sequence length (and thus compute cost — attention is O(T²)).
• Vocabulary size (which scales the embedding table and the LM head).
• Out-of-vocabulary (OOV) behavior.
• Whether the model can spell, count letters, do arithmetic, code.

2.2 The four families of tokenization

Character: Simple, no OOV, but loses morphological structure. RNN char-LMs work but transformers struggle (sequences too long for O(T²) attention).

Word: Tokenizing on whitespace and punctuation. Big problem: "running", "ran", "runs" are unrelated tokens. And any new word is <UNK>. This is what Word2Vec used.

Subword (BPE — Byte Pair Encoding): The 2016 breakthrough (Sennrich et al.). Start with characters; iteratively merge the most frequent adjacent pair into a new symbol. Vocabulary becomes a mix of common whole words and morpheme-like fragments. "tokenization" might split into ["token", "ization"]. We'll implement this in Phase 4–5.

Byte-level BPE (GPT-2 onwards): start with 256 single bytes instead of Unicode characters, then BPE on top. Every UTF-8 string is encodable, period. No <UNK> exists. Lab 01 builds the byte-level skeleton.

2.3 GPT-2's pre-tokenization regex (worth understanding)

GPT-2 doesn't BPE-merge across word boundaries blindly. It first splits text using this regex:

GPT2_PAT = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""

Translated:

• '(?:[sdmt]|ll|ve|re) — contractions ('s, 't, 'll, …)
• ?\p{L}+ — letters with optional leading space
• ?\p{N}+ — digits
• ?[^\s\p{L}\p{N}]+ — punctuation
• \s+(?!\S)|\s+ — whitespace handling

This pre-split prevents merges like the_cat becoming a single token, while still letting BPE merge inside word groups. Lab 01 implements this.

2.4 Why "5 + 3 = 8" sometimes confuses LLMs

Tokenization shapes capabilities. If "857" tokenizes as [8, 57] but "858" as [85, 8], the model has to learn arithmetic across token boundaries that depend on the input. This is why models historically struggled with multi-digit math. Modern fixes: digit-by-digit tokenization (Llama-3), or trained with right-to-left number reversal.

Read: Karpathy's Let's build the GPT tokenizer video — 2 hours, the single best tokenization resource on the internet.

2.5 References

• Sennrich, Haddow, Birch (2016), Neural Machine Translation of Rare Words with Subword Units — the BPE paper.
• Kudo (2018), Subword Regularization — Unigram tokenizer.
• HuggingFace tokenizers documentation.
• OpenAI's tiktoken source — read it.

3. Concept 2 — Linear Algebra for Neural Networks

3.1 What you must know cold

A neural network is a sequence of affine transforms (y = W x + b) interleaved with elementwise nonlinearities (ReLU, GELU, …). All the action is in:

• Matrix multiplication (M, K) @ (K, N) → (M, N). Memorize the inner-dimension rule.
• Transpose Aᵀ. For batched tensors think of shape gymnastics.
• Outer product u vᵀ: a rank-1 matrix.
• Inner / dot product uᵀ v = Σ u_i v_i: a scalar.
• Norm ‖v‖₂ = √(vᵀ v). L2 norm.
• Cosine similarity cos(u, v) = (uᵀ v) / (‖u‖ ‖v‖).

Key identity for retrieval: if ‖u‖ = ‖v‖ = 1, then cos(u, v) = uᵀ v. That's why FAISS and Qdrant store normalized vectors and use inner-product search — it's the same math, but cheaper.

3.2 Eigenvalues, SVD — when do you actually need them?

You will not compute eigenvalues by hand. But you'll meet them in:

• PCA (Phase 2 dimensionality reduction).
• Spectral norm / weight-norm regularization.
• Initialization theory (kaiming/xavier are about preserving variance, related to spectra).
• Understanding why attention has a "rank collapse" problem (Dong et al. 2021).

For now: know that SVD decomposes any matrix as A = U Σ Vᵀ with U, V orthogonal and Σ diagonal of singular values. LoRA (Phase 6) is a low-rank approximation justified by the observation that fine-tuning updates have low effective rank.

3.3 Probability and information theory primer

• Random variable, distribution, density (continuous) vs mass (discrete).
• Bayes: P(A | B) = P(B | A) P(A) / P(B).
• Expectation 𝔼[X] = Σ x P(x).
• Variance Var(X) = 𝔼[X²] - 𝔼[X]².
• Entropy H(p) = -Σ p log p — the average "surprise". Maximum at the uniform distribution.
• Cross-entropy H(p, q) = -Σ p log q. The loss function for classification (and thus next-token prediction).
• KL divergence D_KL(p ‖ q) = Σ p log(p/q). Distance-like (asymmetric); shows up in RLHF (PPO's KL constraint), DPO, distillation.

When an LLM minimizes cross-entropy on next tokens, it is minimizing D_KL(data ‖ model) + H(data) — and H(data) is fixed, so it's equivalently doing maximum likelihood.

3.4 Einstein summation (einsum) — the universal hammer

Once you can read einsum, every transformer paper becomes 5× clearer.

# Standard matmul
torch.einsum("ik,kj->ij", A, B)   # == A @ B

# Batched matmul
torch.einsum("bik,bkj->bij", A, B)   # == torch.bmm(A, B)

# Attention scores
torch.einsum("bhid,bhjd->bhij", Q, K)   # == Q @ K.transpose(-1,-2)

# Multi-head value gather
torch.einsum("bhij,bhjd->bhid", attn, V)

Rules: indices that appear in inputs but not in output get summed over; indices that appear in both inputs and output are batched.

3.5 References

• 3Blue1Brown's Essence of Linear Algebra (mandatory).
• Strang, Introduction to Linear Algebra (book + MIT 18.06 lectures on YouTube).
• einsum is all you need by Tim Rocktäschel.
• Information Theory, Inference, and Learning Algorithms, MacKay — free PDF; read Ch. 2.

4. Concept 3 — Sparse Vector Retrieval (TF-IDF and BM25)

4.1 The problem

Given a query "best neural network books", rank N documents by relevance. The dense embedding approach (Phase 7) is overkill for many real workloads — a sparse keyword model gets you 80% of the way and is interpretable, fast, and updatable.

4.2 TF-IDF derivation

Term Frequency tf(t, d): how often term t appears in document d. Often log-scaled: tf' = 1 + log(tf) to dampen high counts.

Inverse Document Frequency idf(t) = log(N / df(t)), where df(t) is the number of documents containing t. Common terms ("the") get low weight; rare terms ("transformer") get high weight. The log is justified information-theoretically: if t appears in df of N documents, knowing it occurred carries -log(df/N) bits of information.

TF-IDF score of (t, d): tfidf(t, d) = tf'(t, d) · idf(t).

Document vector: a sparse vector indexed by vocabulary, with tfidf(t, d) at each position. L2-normalize so cosine similarity becomes a pure dot product.

Query: same transform, then score(d) = q · d. Top-k.

4.3 BM25 — the workhorse

BM25 is TF-IDF with two essential improvements: term-frequency saturation (the 50th occurrence of "neural" doesn't add 50× the signal) and length normalization (longer docs aren't unfairly favored). Formula:

$$ \text{BM25}(q, d) = \sum_{t \in q} \text{idf}(t) \cdot \frac{f(t, d)(k_1 + 1)}{f(t, d) + k_1 (1 - b + b \cdot |d|/\overline{|d|})} $$

with typical k_1 = 1.2, b = 0.75. This is what Elasticsearch / OpenSearch / Lucene actually use. Phase 7 covers hybrid search (BM25 + dense).

4.4 References

• Manning, Raghavan, Schütze, Introduction to Information Retrieval — free at nlp.stanford.edu/IR-book. Chapters 1, 6, 7.
• Robertson & Zaragoza (2009), The Probabilistic Relevance Framework: BM25 and Beyond.

5. Concept 4 — PyTorch and Autograd

5.1 Tensors

Think of a torch.Tensor as np.ndarray + (a) GPU dispatch and (b) automatic differentiation. The shape and dtype rules are nearly identical.

import torch
x = torch.zeros(3, 4)              # shape (3, 4), float32
x = x.to("cuda")                   # device move
x = x.to(torch.bfloat16)           # dtype cast
y = x[:, :2]                       # view; shares memory
z = x.contiguous().view(12)        # reshape; may copy

Strides matter: a transpose returns a view with non-contiguous strides; some ops require .contiguous() first.

5.2 Autograd — the only paragraph that matters

When you do y = f(x) with x.requires_grad=True, PyTorch builds a computational graph of all intermediate ops. When you call y.backward(), it walks the graph in reverse and fills .grad on every leaf tensor that participates. That's it. This is just the chain rule executed automatically:

If L = f(g(h(x))) then dL/dx = f'(g(h(x))) · g'(h(x)) · h'(x).

PyTorch records each f, g, h and replays the derivatives backward.

x = torch.tensor(2.0, requires_grad=True)
y = (x ** 3 + 2 * x).sin()         # y = sin(x³ + 2x)
y.backward()                        # populates x.grad
print(x.grad)                       # cos(12) * (3*4 + 2) = cos(12) * 14

Key APIs:

• loss.backward() — compute gradients.
• optimizer.zero_grad() — clear them before the next step (gradients accumulate by default).
• optimizer.step() — apply the update.
• with torch.no_grad(): — disable graph building (use in eval / inference). Saves memory.
• tensor.detach() — return a tensor that shares storage but is excluded from autograd.

5.3 The canonical training loop

model = MyModel().to(device)
opt = torch.optim.AdamW(model.parameters(), lr=3e-4)
for epoch in range(N):
    for batch in loader:
        opt.zero_grad()
        logits = model(batch.x.to(device))
        loss = F.cross_entropy(logits, batch.y.to(device))
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()

Memorize this. Every script in this curriculum is a variation on these 7 lines.

5.4 What nn.Module actually is

A class that registers parameters (nn.Parameter) and submodules so that .parameters() recursively collects everything. State dict (model.state_dict()) is just a flat OrderedDict of all params + buffers — that's how saving/loading works.

5.5 References

• PyTorch Tutorials — start with Deep Learning with PyTorch: A 60 Minute Blitz.
• Deep Learning with PyTorch, Stevens, Antiga, Viehmann (Manning).
• Andrej Karpathy's Neural Networks: Zero to Hero — the micrograd lecture builds autograd from scratch in 100 lines. Mandatory.
• PyTorch internals by Edward Yang.

6. How the labs in this phase exercise these ideas

For Lab 01 specifically, when you read the solution, ask yourself:

1. Why does RegexTokenizer need that exact regex (and not just \w+)?
2. What happens if I encode an emoji with WhitespaceTokenizer vs ByteLevelTokenizer?
3. Why is the byte-level vocab exactly 256 before adding merges?
4. How would I extend this to BPE (Phase 4 will)?

7. Common interview questions on Phase 1 material

1. Walk me through what .backward() does.
2. Why is byte-level tokenization useful? Are there downsides?
3. Why do BPE models sometimes count the letters in "strawberry" wrong?
4. What's the difference between cosine similarity and dot product? When are they equivalent?
5. Why is H(p, q) = -Σ p log q the right loss for classification?
6. What does with torch.no_grad(): do — and why does it matter for memory?
7. Implement TF-IDF on a whiteboard.
8. What's the time complexity of self-attention in sequence length T? Why does that matter?
9. Explain the difference between a view and a copy in PyTorch.
10. Given a (B, T, C) tensor, how do you compute per-batch row-wise softmax with einsum / broadcasting?

8. Going from solid → exceptional

After Phase 1, most candidates can do TF-IDF and write a training loop. To stand out:

• Implement BPE end-to-end (training + encoding) before anyone tells you to. Compare your output to tiktoken byte-for-byte.
• Read the GPT-2 tokenizer source in tiktoken and the original OpenAI repo. Understand the byte-to-unicode mapping (bytes_to_unicode() function — it's a clever hack to keep BPE in printable Unicode).
• Read micrograd by Karpathy (~150 lines). You should be able to reimplement it from scratch in 1 hour by the end of the phase.
• Profile a training step with torch.profiler and identify the kernel-launch overhead vs compute.
• Write a 1-page essay on "Why does tokenization shape model capability?" — using examples from the literature.

9. Recommended weekly cadence

Move on to Phase 2 only when you can write the BPE training loop, the TF-IDF formula, and a PyTorch training loop on a whiteboard with no reference.

Lab 01 — Tokenization From Scratch (Solution Walkthrough)

Phase: 1 — Foundations | Difficulty: ⭐⭐☆☆☆ | Time: 1–2 hours

Read ../HITCHHIKERS-GUIDE.md §Tokenization first. This document walks through the solution code line-by-line.

0. What you build and why

Three tokenizers, ranging from naïve to production:

You will see by direct measurement why naïve splitting fails on real text and why GPT-2 layers a regex on top of bytes/BPE.

Run

pip install -r requirements.txt
python lab.py        # TODO scaffold — fill in
python solution.py   # reference implementation

1. The GPT-2 pre-tokenization regex

GPT2_PAT = re.compile(
    r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
)

This is the single most important regex in modern NLP. Read each alternative left-to-right (the | operator):

1. '(?:[sdmt]|ll|ve|re) — English contractions. Captures 's, 'd, 'm, 't, 'll, 've, 're so Don't → Don, 't (two reusable tokens).
2. ' ?\p{L}+' — an optional leading space followed by 1+ Unicode letter characters (\p{L}). The leading space is part of the token — that's why GPT models render text by simple "".join(tokens) (no " ".join). " hello" is one token, distinct from "hello".
3. ' ?\p{N}+' — same for numbers. Splits "abc123" into ["abc", "123"].
4. ' ?[^\s\p{L}\p{N}]+' — runs of punctuation/symbols.
5. \s+(?!\S) — runs of whitespace not followed by non-whitespace (trailing whitespace).
6. \s+ — any other whitespace run (final fallback).

Why is this regex the right pre-tokenization? Because BPE merges look at adjacent characters within a pre-token; if you don't pre-split, "the dog" could merge across the space into a single token "the dog" — wasteful and brittle. By forcing " dog" as a self-contained pre-token, BPE can only learn merges within " dog", never across.

You don't run BPE in this lab (that's an extension), but you set up the pre-tokenization layer that BPE consumes.

2. WhitespaceTokenizer

class WhitespaceTokenizer:
    UNK = "<unk>"

    def __init__(self):
        self.token_to_id: dict[str, int] = {}
        self.id_to_token: dict[int, str] = {}

Two parallel dicts — cheaper than list.index() lookups.

    def train(self, corpus, min_freq=1, special_tokens=None):
        specials = [self.UNK] + (special_tokens or [])
        counts = Counter()
        for line in corpus:
            counts.update(line.split())
        vocab = list(specials) + [tok for tok, c in counts.most_common()
                                  if c >= min_freq and tok not in specials]
        self.token_to_id = {tok: i for i, tok in enumerate(vocab)}
        self.id_to_token = {i: tok for tok, i in self.token_to_id.items()}

Key choices:

• Special tokens reserve the lowest IDs (<unk> is always id 0). Convention; production code hard-codes pad_id=0 or unk_id=0 everywhere.
• most_common() ensures stable ordering: more frequent tokens get smaller ids → embedding tables are more cache-friendly.
• min_freq lets you drop hapaxes (words seen once). For natural text, ~50% of unique tokens appear once but contribute <2% of total tokens.

    def encode(self, text):
        unk = self.token_to_id[self.UNK]
        return [self.token_to_id.get(t, unk) for t in text.split()]

    def decode(self, ids):
        return " ".join(self.id_to_token.get(i, self.UNK) for i in ids)

encode is information-lossy on OOV → <unk>. decode reinserts spaces but cannot recover the original whitespace structure — multiple spaces, tabs, newlines all become single spaces.

3. RegexTokenizer

Only difference from WhitespaceTokenizer: replace line.split() with GPT2_PAT.findall(line). This single change fixes:

• Punctuation: "hello!" → ["hello", "!"].
• Contractions: "don't" → ["don", "'t"].
• Number boundaries: "abc123" → ["abc", "123"].

Decoding uses "".join(...) — the leading-space tokens already carry their spaces. This is the trick that makes round-trip preserve spacing for in-vocab tokens.

4. ByteLevelTokenizer

class ByteLevelTokenizer:
    def encode(self, text):
        return list(text.encode("utf-8"))

    def decode(self, ids):
        return bytes(ids).decode("utf-8", errors="replace")

Three lines. Yet this is what production LLMs fall back to.

• text.encode("utf-8") produces bytes in [0, 255]. UTF-8 is variable-length: ASCII is 1 byte, accented Latin is 2, CJK is 3, emoji is 4.
• The vocab is fixed at 256 — no training step.
• errors="replace" handles ill-formed byte sequences (truncated multi-byte chars from streaming generation) by inserting U+FFFD instead of crashing.

Round-trip safety: for any input, decode(encode(x)) == x. Try emoji, Arabic, code, unicode whitespace.

Trade-off: a token = a byte. English averages ~1 char ≈ 1 byte ≈ 1 token. After BPE we collapse common byte sequences into ~0.25 tokens/char. So pure byte-level inflates sequence length 4× vs BPE — slower but trivially robust.

5. The runner

sample = "Hello, world! Don't tokenize naively — GPT-2's regex is smart."
corpus = [sample] * 100

Corpus is the same sentence ×100. Enough to populate vocabularies; lab is about correctness not training a useful tokenizer.

The round_trip_ok flag will reveal:

• whitespace: True only because we trained on the exact string. Add an unseen word → False.
• regex: True for the same reason — but spacing/punctuation is preserved correctly.
• byte: Always True.

The sanity check against tiktoken confirms our regex matches GPT-2's. tiktoken then runs BPE merges on top — that's why its token count is even lower than our regex count.

6. Expected output

[whitespace] n_tokens= 10  round_trip_ok=True
[regex     ] n_tokens= 17  round_trip_ok=True
[byte      ] n_tokens= 64  round_trip_ok=True
[tiktoken  ] n_tokens= 17

Numbers may differ by ±1 depending on how the em-dash is encoded. Try a slightly modified input — change one word to something not in corpus. The whitespace tokenizer will produce <unk> and break round-trip; the others won't.

7. Common pitfalls

1. re vs regex — only the regex package supports \p{L} Unicode properties. Standard re will silently match nothing.
2. Decoding regex tokens with " ".join — adds extra spaces because the leading-space tokens already carry them. Always "".join.
3. Forgetting errors="replace" in byte-level decoding — production streaming generation will yield partial multi-byte chars, and bare .decode("utf-8") raises.
4. Reserving <unk> mid-vocab — always put specials at IDs 0..k-1. Many downstream libraries assume this.
5. text.split() vs text.split(" ") — split() (no arg) collapses runs of whitespace; split(" ") keeps empty strings.

8. Stretch exercises

• Implement BPE training (Sennrich 2016): start with byte vocab; repeatedly find the most frequent adjacent pair and merge it into a new token; stop at target vocab size. ~100 lines.
• Implement byte-level BPE like GPT-2: pre-tokenize with GPT2_PAT, then do BPE within each pre-token over its UTF-8 bytes. The bytes_to_unicode() helper is the trickiest piece.
• Compute compression ratio (chars/token) on 1 MB of English Wikipedia. Whitespace ~5.0, regex ~4.5, byte 1.0, tiktoken cl100k_base ~4.0.
• Run on Chinese / Arabic / code and compare. tiktoken gpt2 is famously bad on non-English (5–10× more tokens per char). cl100k_base (GPT-4) added more multilingual merges.
• Visualize token boundaries with color coding (Karpathy's video on tokenization shows this beautifully).

9. What this lab proves about you

You can answer tokenizer interview questions ("explain BPE", "why does GPT-4 sometimes count letters wrong in 'strawberry'", "what's the failure mode of pure-whitespace tokenization for LLM pretraining") with code-level confidence. That's the Phase-1 milestone.

Phase 2 — Classical NLP & Static Embeddings

Difficulty: ⭐⭐⭐☆☆ | Estimated Time: 1.5 weeks Roles supported: Pretraining Data Engineer, Research Engineer, Foundation Model Engineer.

Why This Phase Exists

Static embeddings (Word2Vec, GloVe, FastText) are the conceptual ancestors of every modern embedding model used in RAG, retrieval, and the input layer of every LLM. Implementing them from scratch teaches you negative sampling, contrastive objectives, and embedding evaluation — all of which reappear at scale in CLIP, sentence-transformers, and reward models.

You will leave this phase able to explain "what an embedding actually is" without hand-waving.

Concepts

• Distributional hypothesis
• CBOW vs Skip-gram
• Negative sampling derivation (and why it approximates softmax)
• Subsampling of frequent words
• Hierarchical softmax (overview)
• GloVe: co-occurrence matrix factorization
• FastText: subword n-grams, OOV handling
• Embedding evaluation: intrinsic (analogy, similarity) vs extrinsic (downstream task)
• Dimensionality reduction for visualization (t-SNE, UMAP)
• Anisotropy of embedding spaces

Labs

Lab 01 — Word2Vec Skip-Gram From Scratch (NumPy + PyTorch)

Lab 02 — GloVe & FastText (Hands-On)

Lab 03 — Embedding Evaluation & Visualization

Deliverables Checklist

• Skip-gram trained on text8 with intrinsic eval > 0.55
• GloVe + FastText comparison report
• Embedding eval harness reusable in later phases
• t-SNE / UMAP visualization

Interview Relevance

• "Explain negative sampling."
• "Why are static embeddings insufficient for modern NLP?"
• "How would you evaluate an embedding model for a RAG system?" (sets up Phase 7)

Warmup Guide — Classical NLP & Embeddings

Zero-to-expert primer for Phase 02: how meaning becomes geometry — from counting words to word2vec's learned vectors — the conceptual foundation every embedding system (including modern RAG) still stands on.

Table of Contents

• Chapter 1: The Distributional Hypothesis
• Chapter 2: Count-Based Representations — BoW and TF-IDF
• Chapter 3: Vector Similarity — Why Cosine
• Chapter 4: word2vec — Learning Vectors by Prediction
• Chapter 5: Negative Sampling — The Trick That Made It Tractable
• Chapter 6: What Embedding Spaces Encode (and Don't)
• Chapter 7: From Static to Contextual — Why This Still Matters
• Lab Walkthrough Guidance
• Success Criteria
• Interview Q&A
• References

Chapter 1: The Distributional Hypothesis

From zero: computers need numbers; words are symbols. The naive encoding — one-hot vectors (vocabulary-sized, single 1) — makes every pair of words equidistant: cat is as far from kitten as from carburetor. No similarity structure, vocabulary-sized dimensionality, and no generalization (what the model learns about cat says nothing about kitten).

The escape is a 1950s linguistic idea (Firth: "you shall know a word by the company it keeps"): words appearing in similar contexts have similar meanings. This turns semantics into a statistics problem — measure each word's context distribution, and words with similar distributions get similar representations. Every embedding method in this phase (and every contextual model after it) is an implementation of this single hypothesis; the methods differ only in how they compress context statistics into vectors.

Chapter 2: Count-Based Representations — BoW and TF-IDF

The pre-neural baselines — still in production everywhere (search engines, spam filters), and the honest baseline your lab compares against:

• Bag of Words: document → vector of word counts. Order discarded ("dog bites man" = "man bites dog"), but for topical tasks counts carry most of the signal.
• TF-IDF fixes BoW's flaw — common words (the) dominate counts while carrying no discriminative information:

$$\text{tfidf}(t, d) = \text{tf}(t, d) \cdot \log\frac{N}{\text{df}(t)}$$

Term frequency × inverse document frequency: a term scores high if frequent in this document and rare across documents. The log dampens; smoothing (+1) avoids division by zero. The IDF weighting idea — downweight the ubiquitous — recurs everywhere from BM25 (RAG's sparse-retrieval workhorse, Phase 07) to attention entropy analysis.

• Limits that motivate the rest of the phase: vectors are vocabulary-sized and sparse; synonyms remain orthogonal (car ⊥ automobile — exact-match only); word order and polysemy invisible.

Chapter 3: Vector Similarity — Why Cosine

$$\text{cosine}(u, v) = \frac{u \cdot v}{|u|,|v|} \in [-1, 1]$$

— the angle, ignoring magnitude. Why magnitude is ignored deliberately: in count spaces, document length inflates magnitude without changing topic; in learned embedding spaces, frequency-related magnitude artifacts similarly pollute. Direction carries the semantics.

Three operational facts you'll reuse for the rest of the curriculum:

1. On L2-normalized vectors, cosine, dot product, and Euclidean distance produce identical rankings ($|u-v|^2 = 2 - 2u\cdot v$) — which is why vector databases normalize once and use dot product (cheapest) internally.
2. Exact top-k search is O(N·d) per query — fine to ~1M vectors, then approximate indexes (HNSW, IVF — Phase 07's territory) trade recall for speed.
3. In high dimensions, random vectors concentrate near orthogonality — cosine similarities of unrelated items cluster around 0 with small variance, so small absolute differences (0.31 vs 0.35) can be meaningful. Calibrate thresholds empirically per embedding model; never copy a threshold across models.

Chapter 4: word2vec — Learning Vectors by Prediction

The 2013 leap: don't count contexts, predict them — and harvest the network's weights as the representation.

Skip-gram (the variant your lab implements): for each word, predict the words in a window around it. Architecture is almost embarrassingly simple — two matrices: $W_{in} \in \mathbb{R}^{V \times d}$ (center-word vectors) and $W_{out} \in \mathbb{R}^{V \times d}$ (context-word vectors). Score of context $o$ given center $c$: $u_o^\top v_c$, softmaxed over the vocabulary:

$$P(o \mid c) = \frac{\exp(u_o^\top v_c)}{\sum_{w \in V} \exp(u_w^\top v_c)}$$

Training maximizes this over all (center, context) pairs in the corpus. The product: $W_{in}$'s rows — vectors where similarity-of-context became geometric proximity, because words with the same neighbors received the same gradient pushes.

CBOW is the mirror (predict center from averaged context) — faster, slightly worse on rare words. The deeper unification (Levy & Goldberg): skip-gram with negative sampling implicitly factorizes the PMI (pointwise mutual information) co-occurrence matrix — the neural method and the count-based methods converge on the same statistics, which is the distributional hypothesis showing through.

Chapter 5: Negative Sampling — The Trick That Made It Tractable

The softmax's denominator sums over the whole vocabulary — at V = 100K, every training step costs a 100K-way normalization. Skip-Gram with Negative Sampling (SGNS) replaces the multiclass problem with binary discrimination: for each true (center, context) pair, sample $k$ random "negative" words and train the model to score real pairs high and fakes low:

$$\mathcal{L} = -\log \sigma(u_o^\top v_c) - \sum_{i=1}^{k} \log \sigma(-u_{n_i}^\top v_c)$$

Cost per step: $k+1$ dot products instead of $V$. Details that matter (and that the lab implements):

• Negatives are drawn from the unigram distribution raised to the 3/4 power — flattening it so rare words get sampled enough; the 0.75 is empirical and load-bearing.
• Frequent-word subsampling: drop training occurrences of very frequent words with probability $1 - \sqrt{t/f(w)}$ — the appears millions of times and teaches nothing new after the first thousand.
• $k$ = 5–20 for small corpora, 2–5 at scale.

This "turn an intractable softmax into sampled binary classification" move is a recurring pattern — you'll meet its cousins in contrastive learning (CLIP, Phase 06 of the CV track) and InfoNCE.

Chapter 6: What Embedding Spaces Encode (and Don't)

• The famous linear structure: king - man + woman ≈ queen — relations as directions (gender, tense, capital-of). Real but overstated: it works for frequent, clean relations and degrades elsewhere; evaluate with analogy suites, don't assume.
• Similarity vs relatedness: embeddings conflate them — coffee is close to cup (related) and to tea (similar). Tasks that need one but not the other (retrieval vs substitution) need different training objectives; this distinction returns with force in RAG retrieval quality (Phase 07).
• Antonyms are close: hot/cold share contexts almost perfectly. Distributional methods cannot see the polarity flip — a permanent limitation downstream tasks must handle.
• Bias is faithfully learned: occupation–gender associations and worse, present in the corpus statistics, become geometry. Debiasing-by-projection exists and is partial (Gonen & Goldberg's "lipstick on a pig"); the production answer is evaluation and task-level mitigation, not pretending the vectors are neutral.
• One vector per word: bank (river/finance) gets a frequency-weighted average of its senses — the polysemy failure that motivates Chapter 7.

Chapter 7: From Static to Contextual — Why This Still Matters

Static embeddings assign one vector per type; contextual models (ELMo → BERT → every LLM layer) produce a vector per occurrence, computed from the sentence — solving polysemy and word order at the cost of running a model per text. The modern landscape you're heading toward:

• Sentence/document embedding models (the RAG workhorses, Phase 07) are contextual encoders pooled to one vector and trained contrastively — conceptually: word2vec's objective, upgraded with transformers and curated positives/negatives.
• The mechanics transfer wholesale: cosine similarity, normalization, the similarity-vs-relatedness distinction, threshold calibration, nearest-neighbor search — everything in Chapters 3 and 6 is daily working knowledge for an inference engineer running a vector database.
• And inside every LLM, the embedding table (Phase 01's contract) is a static embedding layer — token vectors whose geometry you'll inspect when debugging glitch tokens and quantization damage to lm_head/embeddings (model-accuracy track).

Lab Walkthrough Guidance

Lab 01 — word2vec from Scratch:

1. Build the data pipeline first: tokenize (Phase 01's tools), build vocab with min frequency, generate (center, context) pairs with a window; verify pair counts by hand on a toy sentence.
2. Implement SGNS loss exactly as Chapter 5's equation — two embedding matrices, $k$ negatives from the unigram^0.75 table (build the sampling table once; verify its marginal distribution empirically).
3. Add frequent-word subsampling; measure its effect on training speed and on rare-word neighbor quality (it should help both — understand why before believing it).
4. Train on text8 or a Wikipedia slice; sanity-check during training with a fixed probe set (king, paris, coffee nearest neighbors every N steps — watching neighbors sharpen is the lab's payoff moment).
5. Evaluate: nearest neighbors, a handful of analogies, and a 2-D projection (PCA or t-SNE — with the standard caveat that t-SNE distorts global structure; don't over-read the picture).

Success Criteria

You are ready for Phase 03 when you can, from memory:

1. State the distributional hypothesis and show how both TF-IDF and word2vec implement it.
2. Write the TF-IDF formula and explain each factor's job.
3. Explain why cosine over raw distance, the normalization-makes-them-equivalent fact, and the high-dimensional concentration caveat.
4. Write the SGNS loss and justify the 3/4 power and subsampling.
5. Name four things embedding spaces get wrong (antonyms, polysemy, relatedness conflation, bias) with the mechanism behind each.
6. Trace the lineage: one-hot → counts → static learned → contextual, with the failure that drove each transition.

Interview Q&A

Q: Why did word2vec beat count-based methods if they capture the same statistics? Levy & Goldberg showed SGNS implicitly factorizes the PMI matrix — so the information is the same; the win was practical: dense low-dimensional vectors directly usable downstream, scalable online training over arbitrarily large corpora (no V×V matrix to materialize), and hyperparameters (subsampling, negative distribution, dynamic windows) that act as effective regularizers. With matched preprocessing, tuned SVD-over-PMI is competitive — knowing that is the difference between folklore and understanding.

Q: Your RAG system retrieves topically-related but unhelpful passages. Connect to this phase. Similarity-vs-relatedness conflation (Ch. 6): the embedding model scores coffee prices near coffee brewing because they share context distributions, but the task needed answer-bearing similarity. Fixes are objective-level — embedding models trained on (query, answer) pairs rather than symmetric similarity, hybrid BM25+dense scoring (exact terms carry intent), and rerankers (Phase 07). The diagnosis vocabulary comes from static-embedding days; the failure never went away.

Q: Why do hot and cold embed nearby, and when does it bite? They're distributionally near-identical ("the soup is ", " weather"). It bites in any polarity-sensitive downstream use: sentiment lexicon induction, contradiction detection, retrieval where negation flips the answer. Contextual models mitigate but don't eliminate (negation remains a known weak spot through modern LLMs); the mitigation is task-specific supervision, not better unsupervised geometry.

References

• Mikolov et al., Efficient Estimation of Word Representations in Vector Space (2013) — arXiv:1301.3781
• Mikolov et al., Distributed Representations of Words and Phrases (2013) — arXiv:1310.4546 — negative sampling, subsampling
• Levy & Goldberg, Neural Word Embedding as Implicit Matrix Factorization (NeurIPS 2014)
• Levy, Goldberg & Dagan, Improving Distributional Similarity with Lessons Learned from Word Embeddings (TACL 2015) — the hyperparameters-matter paper
• Pennington et al., GloVe (EMNLP 2014) — the count-based contemporary
• Gonen & Goldberg, Lipstick on a Pig (2019) — debiasing's limits
• Jurafsky & Martin, Speech and Language Processing (3rd ed. draft), ch. 6 — the best textbook treatment
• Illustrated word2vec — Jay Alammar

🛸 Hitchhiker's Guide — Phase 2: Classical NLP & Word Embeddings

Read this if: You can write a TF-IDF index but you don't yet feel in your bones why king − man + woman ≈ queen falls out of word2vec, or why "softmax over the whole vocabulary is too expensive" is the historical pivot that led to negative sampling.

0. The 30-second mental model

A word embedding is a learned dense vector that captures meaning by co-occurrence statistics. The training signal is: "predict context from word" (or vice versa). After enough data, vectors of similar words cluster together — and useful linear structure emerges (analogies). This is the conceptual ancestor of token embeddings inside transformers.

By the end of Phase 2 you should:

• Know the distributional hypothesis and why it makes sense.
• Be able to derive Skip-gram with negative sampling from scratch.
• Understand the difference between count-based (PPMI, GloVe) and predict-based (word2vec) embeddings, and the surprising 2014 result that they're closely related.
• Know how to evaluate an embedding (intrinsic vs extrinsic).
• Understand why static embeddings were superseded by contextual embeddings (ELMo → BERT) and where they are still used today (retrieval, recommender systems, cold start).

1. The Distributional Hypothesis

"You shall know a word by the company it keeps." — J. R. Firth, 1957

If two words appear in similar contexts, they probably mean similar things. That's the entire premise. Make a giant matrix M ∈ ℝ^{V × V} where M[i, j] = "how often word i appears near word j" — the rows are word representations. The rest of the field is "how do we make this matrix smaller and better".

1.1 Three classes of word representations

1. Count-based: build the co-occurrence matrix; reduce dimensionality (SVD on PPMI). Examples: LSA, HAL, PPMI+SVD.
2. Predict-based: train a neural model whose weights become the word vectors. Examples: word2vec, GloVe (hybrid), FastText.
3. Contextual: embeddings depend on the sentence around the word. Examples: ELMo, BERT, every modern LLM. Phase 4+.

1.2 PPMI — the bridge between counting and predicting

Pointwise Mutual Information: pmi(w, c) = log P(w, c) / (P(w) P(c)). Positive PMI: clip negatives to 0. The famous Levy & Goldberg (2014) result is that Skip-gram with negative sampling implicitly factorizes a shifted PMI matrix. So the seemingly different paradigms compute almost the same thing under the hood.

2. Word2Vec — Skip-Gram with Negative Sampling

This is the core of Lab 01. Internalize it.

2.1 The Skip-Gram task

Given a center word w_c (e.g., "neural"), predict its surrounding context words w_o within a window (e.g., the 5 words before and after). The model parameters are two embedding matrices:

• Input embeddings V ∈ ℝ^{|vocab| × d} — V[w] is the vector when w is the center.
• Output embeddings U ∈ ℝ^{|vocab| × d} — U[w] is the vector when w is a context.

Probability that w_o is a context for w_c:

$$ P(w_o \mid w_c) = \frac{\exp(U_{w_o}^\top V_{w_c})}{\sum_{w} \exp(U_w^\top V_{w_c})} $$

This denominator sums over the entire vocabulary at every step. That's prohibitive (vocab can be millions). Two historical fixes:

1. Hierarchical softmax — replace the flat softmax with a binary tree (Huffman code) so each prediction is a sequence of log V binary choices. O(log V).
2. Negative sampling — don't compute the partition function at all; turn it into a binary classification task.

2.2 Negative sampling derivation

For each true (center, context) pair (w_c, w_o), sample K "negative" context words w_neg ~ P_n(w). Train a logistic regression: predict 1 for the true pair, 0 for each negative.

Loss for a single positive example with K negatives:

$$ \mathcal{L} = -\log \sigma(U_{w_o}^\top V_{w_c}) - \sum_{k=1}^K \mathbb{E}{w_k \sim P_n}\left[\log \sigma(-U{w_k}^\top V_{w_c})\right] $$

where σ is the sigmoid. Notice: each gradient step touches only K + 1 rows of U instead of all |vocab|. That's the speedup.

The negative distribution is the unigram raised to 0.75:

$$ P_n(w) \propto f(w)^{0.75} $$

This empirical heuristic dampens very frequent words and boosts rare ones. The 0.75 is mostly folklore (Mikolov et al. tried a few values and it worked).

2.3 Subsampling frequent words

Mikolov also discards each occurrence of word w with probability:

$$ P_\text{discard}(w) = 1 - \sqrt{\frac{t}{f(w)}} $$

with t ≈ 1e-5. This removes noise from "the", "and", etc., yielding both faster training and higher-quality vectors.

2.4 Why analogies work (linear structure)

The famous king − man + woman ≈ queen is a consequence of how the model encodes multiple, additive semantic axes. If "royal-ness" and "gender" are roughly orthogonal directions in the embedding space, then subtracting "man" from "king" removes the gender component and adding "woman" restores it as female. Levy & Goldberg (2014) and Arora et al. (2016) explain this rigorously; the short version: log-bilinear models produce vectors whose inner products approximate PMI, and PMI has linear-additive structure for many semantic features.

2.5 Two main objectives — Skip-Gram vs CBOW

• Skip-Gram: predict context given center. Better for rare words.
• CBOW (Continuous Bag of Words): predict center given averaged context. Faster to train.

Skip-gram with negative sampling won historically.

2.6 References

• Mikolov, Sutskever, Chen, Corrado, Dean (2013), Distributed Representations of Words and Phrases and their Compositionality — the SGNS paper.
• Mikolov, Chen, Corrado, Dean (2013), Efficient Estimation of Word Representations in Vector Space.
• Levy & Goldberg (2014), Neural Word Embedding as Implicit Matrix Factorization.
• Goldberg's chapter word2vec Explained (free).
• Goldberg, Neural Network Methods for Natural Language Processing (Morgan & Claypool) — the best textbook for this era.

3. GloVe and FastText (the cousins)

3.1 GloVe

Pennington, Socher, Manning (2014) at Stanford. Trains on the logarithm of the co-occurrence matrix directly with a weighted least-squares loss:

$$ \mathcal{L} = \sum_{i, j} f(X_{ij}) \left(w_i^\top \tilde{w}_j + b_i + \tilde{b}j - \log X{ij}\right)^2 $$

The weighting f(X_{ij}) damps very frequent pairs. GloVe sits philosophically between count-based and predict-based methods.

3.2 FastText

Bojanowski, Grave, Joulin, Mikolov (2017). Each word is represented as the sum of its character n-gram vectors. So "where" = <wh + whe + her + ere + re> + <where>. Two huge wins:

1. OOV handling: any new word can be embedded by summing its n-grams.
2. Morphology: Inflected forms (run/runs/running/ran) share n-grams and thus geometry.

FastText is still a great default for non-English languages (Arabic, Finnish, Turkish) and tasks where a small model + cold-start handling matters.

3.3 References

• Pennington, Socher, Manning (2014), GloVe: Global Vectors for Word Representation.
• Bojanowski et al. (2017), Enriching Word Vectors with Subword Information.

4. Sentence and Document Embeddings

A single vector per word doesn't help if you want to retrieve passages. Three eras:

1. Average / weighted-average of word vectors (Arora SIF). Dirt simple, surprisingly effective baseline.
2. InferSent / Universal Sentence Encoder — supervised on NLI / multitask data.
3. Contrastive sentence transformers (Phase 7 will use these): SBERT, E5, BGE, Cohere embed-v3, OpenAI text-embedding-3. Trained with triplet loss or InfoNCE on (query, positive, negative) pairs.

The key idea connecting Phases 2 → 7: a contrastive loss is just negative sampling on sentence pairs. The math is the same; the unit is bigger.

Read: Reimers & Gurevych (2019), Sentence-BERT. Wang et al. (2022), Text Embeddings by Weakly-Supervised Contrastive Pre-training (E5).

5. Evaluating Embeddings

5.1 Intrinsic

• Word similarity: human-rated pairs (WordSim-353, SimLex-999). Score: Spearman correlation between cosine and human ratings.
• Analogies: Google analogy set (king:man :: queen:?). Score: top-1 accuracy on arg max_w cos(w, b - a + c).
• Clustering coherence: do nearest neighbors of "Java" all relate to programming or to coffee?

5.2 Extrinsic

• Plug embeddings into downstream tasks (sentiment classification, NER, retrieval) and measure end-task metric.
• Extrinsic almost always wins as the truth — intrinsic benchmarks can be gamed.

5.3 Bias auditing

Bolukbasi et al. (2016), Man is to Computer Programmer as Woman is to Homemaker? showed word2vec encodes gender bias along measurable axes. This is a recurring topic in safety interviews.

6. Dimensionality Reduction & Visualization

To inspect a 300-dim embedding space, project to 2D:

• PCA: linear, fast. Use as a first glance.
• t-SNE (van der Maaten & Hinton, 2008): nonlinear, preserves local neighborhoods. Notoriously misleading at the global scale (clusters far apart in t-SNE may be close in reality).
• UMAP (McInnes, Healy, Melville, 2018): nonlinear, faster than t-SNE, preserves more global structure. Default in 2024+.

Always plot a known-labeled subset (countries, animals, programming languages) to sanity-check that semantic clusters appear.

7. Lab 01 walkthrough (lab-01-word2vec-from-scratch)

The lab implements Skip-Gram + negative sampling end to end. Things to internalize while reading the solution:

1. Vocab construction — frequency cutoff, then index assignment. Output a dict {word: id} and its inverse.
2. Subsampling — apply Mikolov's discard rule per occurrence.
3. Iterable dataset — yields (center, positive_context, negatives[K]) tuples. The IterableDataset pattern is essential for streaming corpora that don't fit in memory.
4. Negative sampling distribution — pre-compute a frequency-table ^0.75 once; sample by binary search on cumulative.
5. Forward: score = sigmoid(U[w_o] · V[w_c]). Loss = BCE on the binary labels.
6. Two embedding matrices: input and output. The "word vector" you keep at inference is the input matrix V (or sum of both — Mikolov used V only).
7. Nearest-neighbor demo — cosine over a (V, d) matrix → top-k.

Things you should be able to explain afterwards:

• Why two embedding matrices and not one?
• What happens if K = 0? (Degenerates; can't learn discrimination.)
• Why is ^0.75 a "good enough" hack?
• How would you make this run in a few minutes on a single GPU? (Bigger batches, fewer epochs, smaller window.)

8. Common interview questions on Phase 2 material

1. Derive Skip-gram with negative sampling on the whiteboard.
2. Why does king − man + woman ≈ queen work?
3. What's the difference between word2vec, GloVe, and FastText?
4. Why is the unigram raised to 0.75 in negative sampling?
5. How would you handle a new word that wasn't in your vocab? (FastText answer.)
6. What's PMI, and how does it relate to word2vec?
7. Compare static embeddings (word2vec) to contextual ones (BERT). When would you still use the former?
8. How do you evaluate an embedding model?
9. What's the complexity of a softmax over a 1M-word vocab? How do hierarchical softmax and negative sampling help?
10. Explain the connection between negative sampling and the modern InfoNCE / contrastive loss.

9. From solid → exceptional

• Reimplement word2vec with NEG-K loss in pure NumPy (no PyTorch). It's ~150 lines.
• Train on a 1 GB Wikipedia dump; evaluate on Google analogies; report both accuracy and training time.
• Re-derive the gradient updates for SGNS by hand. Confirm against autograd.
• Read the Levy & Goldberg paper and explain in your own words why SGNS factorizes shifted PMI.
• Train fastText on Arabic Wikipedia and show it handles morphology via n-gram averaging.
• Compare nearest neighbors in your word2vec space with those from BGE — note the differences (BGE captures factual relatedness; word2vec captures distributional similarity).

10. Recommended cadence

Lab 01 — Word2Vec Skip-Gram with Negative Sampling (Solution Walkthrough)

Phase: 2 — Classical NLP & Embeddings | Difficulty: ⭐⭐⭐☆☆ | Time: 3–5 hours

Concept primer: ../HITCHHIKERS-GUIDE.md §Word2Vec. This document walks through the code in solution.py and explains every non-obvious choice.

Run

pip install -r requirements.txt
wget http://mattmahoney.net/dc/text8.zip && unzip text8.zip -d data/
python solution.py --data ./data/text8 --epochs 3

0. The mission

Train a 100-dim word embedding on text8 (a 100 MB cleaned slice of English Wikipedia) using Skip-Gram with Negative Sampling (SGNS). At the end:

nearest("king")  → ["queen", "prince", "throne", "kings", "monarch", ...]
nearest("paris") → ["france", "london", "berlin", "vienna", "rome", ...]

…with no labels, just from co-occurrence. The experiment that launched modern NLP (Mikolov et al. 2013).

1. The math

For each (center $c$, context $o$) pair:

$$ \mathcal{L} = -\log \sigma(v_c \cdot v_o) - \sum_{k=1}^{K} \log \sigma(-v_c \cdot v_{n_k}), \quad n_k \sim P_n $$

where $P_n(w) \propto \text{freq}(w)^{0.75}$ is the negative-sampling distribution, and $K$ (5–20) is the number of negatives per positive.

Two embedding tables: an input matrix $V$ for centers, an output matrix $U$ for context/negatives. By convention we keep $V$ as the final embeddings.

2. build_vocab — three things in one

def build_vocab(words, min_count=5):
    counts = Counter(words)
    vocab = [w for w, c in counts.items() if c >= min_count]
    w2i = {w: i for i, w in enumerate(vocab)}
    freqs = np.array([counts[w] for w in vocab], dtype=np.float64)
    neg_dist = freqs ** 0.75
    neg_dist /= neg_dist.sum()
    return w2i, vocab, neg_dist

The exponent 0.75 is Mikolov's empirical choice: smaller than 1 down-weights very common words (you don't want every negative to be "the"); larger than 0 doesn't make rare words too likely (which would be uninformative).

min_count=5: drop any word seen <5 times. For text8 (~17M tokens) this prunes ~250k unique words to ~70k. Removes most typos and proper-noun fluff.

3. SkipGramDataset — subsampling and pair generation

3.1 Frequent-word subsampling

self.keep = np.minimum(1.0, np.sqrt(subsample_t / f) + subsample_t / f)

For each center occurrence, probabilistically drop with probability 1 - keep[w]. Why?

• "the" appears with frequency ~5%. Without subsampling, half your training pairs would have "the" as the center — useless because "the" co-occurs with everything.
• For very common words f >> t (with t=1e-4), so keep ≈ sqrt(t/f) ≪ 1.
• For rare words f ≪ t, so keep saturates at 1 → never dropped.

This trick gives ~2× quality improvement (Mikolov 2013).

3.2 Dynamic window with random shrinking

for i, center in enumerate(self.ids):
    if rng.random() > self.keep[center]:
        continue
    w = rng.randint(1, self.window)   # 👈 random window per sample
    for j in range(max(0, i - w), min(len(self.ids), i + w + 1)):
        if j == i: continue
        yield center, self.ids[j]

The window size is resampled per center word. This implicitly weights nearer context words more (they're sampled in every window size; far words only at large window sizes). Mathematically equivalent to a triangular weighting kernel — for free.

IterableDataset (vs Dataset) means we stream pairs instead of materializing all ~100M of them.

4. The model — SkipGramNS

class SkipGramNS(nn.Module):
    def __init__(self, vocab_size, dim=100):
        super().__init__()
        self.in_emb = nn.Embedding(vocab_size, dim)
        self.out_emb = nn.Embedding(vocab_size, dim)
        nn.init.uniform_(self.in_emb.weight, -0.5/dim, 0.5/dim)
        nn.init.zeros_(self.out_emb.weight)

• Two tables, not one. The math fundamentally needs both.
• Init scale 0.5/dim — keeps dot products $v_c \cdot v_o$ in a sensible range early.
• Output init 0 — at step 0, $v_c \cdot v_o = 0$ → $\sigma(0) = 0.5$ → loss = $\log 2 \approx 0.69$. Clean baseline.

def forward(self, center, pos, neg):
    v_c = self.in_emb(center)              # (B, D)
    v_p = self.out_emb(pos)                # (B, D)
    v_n = self.out_emb(neg)                # (B, K, D)
    pos_score = (v_c * v_p).sum(-1)
    neg_score = torch.bmm(v_n, v_c.unsqueeze(-1)).squeeze(-1)
    loss = -F.logsigmoid(pos_score).mean() - F.logsigmoid(-neg_score).mean()
    return loss

• (v_c * v_p).sum(-1) is elementwise multiply + sum — the per-row dot product (cheaper than bmm).
• bmm(v_n, v_c.unsqueeze(-1)) is a batched matrix-vector product: K-many dot products of v_n against v_c.
• F.logsigmoid not log(sigmoid(x)) — numerically stable. Naive composition produces nan for very negative x.

5. collate — batching pairs and sampling negatives

negatives = torch.multinomial(neg_dist_t, len(batch) * n_neg, replacement=True).view(-1, n_neg)

• replacement=True is essential — without it you'd be sampling without replacement from a 70k-element distribution, hitting the rare tail too often.
• We don't filter cases where the negative equals the positive — probability is ≤ 1/|V| ≈ 1/70000, dominated by other negatives.

6. The training loop

opt = torch.optim.Adam(model.parameters(), lr=2.5e-3)

LR is high (~10× a typical transformer LR) because (a) embeddings are linear → no exploding-gradient risk, (b) each parameter is touched rarely (sparse access pattern), so per-update steps must be larger.

batch_size=512, n_neg=5 → each step processes 512 positives + 2560 negatives = 3072 dot products per layer.

7. nearest

W = F.normalize(model.in_emb.weight.detach(), dim=1)
q = W[w2i[word]]
sims = (W @ q).cpu().numpy()

F.normalize(..., dim=1) makes each row unit-norm. Then W @ q is cosine similarity (since cos(a,b) = a_unit · b_unit).

We use the input embedding (in_emb) for query and key. Convention; out_emb works similarly.

8. Expected output

After 3 epochs (~10 min on a 4090, ~30 min on CPU):

chars=70123  tokens=17,005,207
  ep 0 step    1000  loss=4.2143
  ep 2 step  100000  loss=1.6234

Nearest neighbors:
  king        [('prince', 0.71), ('queen', 0.69), ('throne', 0.62), ...]
  paris       [('france', 0.73), ('london', 0.66), ('berlin', 0.62), ...]
  computer    [('computers', 0.78), ('software', 0.71), ('hardware', 0.66), ...]

Sanity bar: if king's top-5 doesn't include queen, something is wrong — most likely (a) min_count too high, (b) too few epochs, (c) you accidentally averaged input+output before training.

9. The famous analogy test

v = W[w2i["king"]] - W[w2i["man"]] + W[w2i["woman"]]
# nearest to v, excluding king/man/woman → should produce "queen"

This is the demo that made Word2Vec famous. Works because the embedding space encodes gender as a roughly linear direction.

It also fails in revealing ways: try nurse - woman + man and you may get doctor. The bias that motivated debias and counterfactual-augmentation research.

10. Common pitfalls

1. Forgetting subsampling → 2× slower convergence, worse quality.
2. Same random seed across DataLoader workers → all workers yield the same pair sequence. Use num_workers=0 here.
3. log(sigmoid(x)) instead of F.logsigmoid(x) → NaN losses at high negatives.
4. Sampling without replacement for negatives → biases toward rare words.
5. Only positive pairs (no negatives) → embeddings collapse to one vector.
6. Computing similarity without normalizing → returns dot products, correlated with vector norms.

11. Stretch exercises

• Add CBOW (Continuous Bag-of-Words): predict center from average of context. Compare quality.
• Implement GloVe (Pennington 2014): factorize the global co-occurrence matrix's log-counts.
• Visualize with t-SNE/UMAP. Plot 5000 most-frequent words. Observe clusters: countries, days, professions.
• Replicate Levy & Goldberg: SGNS implicitly factorizes the shifted PPMI matrix. Compute SVD of PPMI and compare cosine sims.
• Plug into a downstream task (e.g., SST-2 sentiment). Compare to randomly-initialized embeddings.
• FastText extension: hash character n-grams; sum subword vectors. Handles OOV.

12. What this lab proves about you

You can implement the foundational embedding model without scaffolding, derive the SGNS loss from cross-entropy, explain every hyperparameter, and link it forward to attention (which generalizes "context = nearby tokens" to "context = all tokens with learned weights"). Phase-2 milestone.

Phase 3 — RNNs & Language Modeling

Difficulty: ⭐⭐⭐☆☆ | Estimated Time: 1.5 weeks Roles supported: Foundation Model Engineer (historical literacy), all research-engineer roles (interview "explain attention" answer requires you to know what came before).

Why This Phase Exists

You will not deploy an RNN to production in 2026. But you will be asked in interviews:

• "Why did transformers replace RNNs?"
• "Explain LSTM gating mathematically."
• "What is teacher forcing?"
• "Where did attention come from?"

Building a char-RNN and a seq2seq model with Bahdanau attention is the cheapest way to internalize these answers — and it makes the leap to transformers in Phase 4 trivial.

Concepts

• Sequence modeling: P(x_t | x_<t)
• Vanilla RNN: hidden-state recurrence h_t = tanh(W_x x_t + W_h h_{t-1})
• Backpropagation through time (BPTT)
• Vanishing/exploding gradients (and the math behind why)
• LSTM: forget / input / output gates, cell state
• GRU: reset / update gates (simpler, often comparable)
• Sequence-to-sequence: encoder-decoder, fixed-context-vector bottleneck
• Bahdanau (additive) attention — the precursor to transformer attention
• Teacher forcing, scheduled sampling
• Perplexity = exp(cross-entropy loss)

Labs

Lab 01 — Vanilla RNN Char-Language-Model From Scratch

Lab 02 — LSTM & GRU (And Why They Help)

Lab 03 — Seq2Seq + Bahdanau Attention (Toy Translation)

Deliverables Checklist

• Char-RNN trained on Shakespeare with temperature sampling
• LSTM vs GRU vs RNN comparison + gradient-norm plot
• Seq2seq with attention + alignment heatmap

Interview Relevance

• "Why did transformers replace RNNs?" — parallelism + long-range dependencies
• "Walk me through LSTM gates"
• "Where does scaled-dot-product attention come from historically?"

Warmup Guide — RNNs & Language Modeling

Zero-to-expert primer for Phase 03: what a language model is (the probability framing that survives every architecture change), and recurrent networks — the architecture whose failures explain why transformers look the way they do.

Table of Contents

• Chapter 1: Language Modeling — The Definition That Never Changes
• Chapter 2: Cross-Entropy and Perplexity
• Chapter 3: The Recurrent Idea
• Chapter 4: Backpropagation Through Time and the Vanishing Gradient
• Chapter 5: LSTM and GRU — Gating as Gradient Plumbing
• Chapter 6: Sampling from a Language Model
• Chapter 7: Why RNNs Lost — and Where They Won
• Lab Walkthrough Guidance
• Success Criteria
• Interview Q&A
• References

Chapter 1: Language Modeling — The Definition That Never Changes

A language model assigns probability to sequences. By the chain rule, exactly:

$$P(w_1, \ldots, w_n) = \prod_{i=1}^{n} P(w_i \mid w_1, \ldots, w_{i-1})$$

So the entire field reduces to one learnable function: given a prefix, output a probability distribution over the next token. Every architecture in this curriculum — n-grams, the RNN you build here, the transformer in Phase 04, GPT-4 — is a different parameterization of $P(w_i \mid w_{<i})$. Generation is just repeated sampling from it. Internalizing this framing is the phase's real product: when later phases discuss KV caches or speculative decoding, they're discussing engineering of this conditional distribution's evaluation, nothing more.

The n-gram baseline (know it; it calibrates everything): approximate the condition by the last $k{-}1$ words and estimate by counting. Fails by sparsity — most 5-grams never occur even in huge corpora (smoothing/backoff is the classic patchwork) — and by having no notion of similarity: counts for cat sat teach nothing about kitten sat. Neural LMs fix both at once: embeddings give similarity (Phase 02), and a parametric function generalizes across contexts.

Chapter 2: Cross-Entropy and Perplexity

Training minimizes cross-entropy: average negative log-probability the model assigns to the actual next token:

$$\mathcal{L} = -\frac{1}{N}\sum_{i} \log P_\theta(w_i \mid w_{<i})$$

— equivalently, maximum likelihood. Perplexity is its exponential, $\text{PPL} = e^{\mathcal{L}}$: the effective branching factor ("as uncertain as a fair choice among PPL options"). Calibration numbers worth carrying: a uniform model over vocab V has PPL = V; a character-level model on English text reaching PPL ~3–4 (≈1.6–2.0 bits/char) is learning real structure; word/subword PPL of strong LLMs on WikiText runs single digits. (The full measurement discipline — tokenizer dependence, sliding windows — is in the model-accuracy track's Phase 09; here you just need the loss-curve intuition: watch bits-per-character fall during training and know what a given level "feels like" in samples. The lab makes you do exactly that.)

Chapter 3: The Recurrent Idea

Feed-forward nets take fixed-size inputs; text is variable-length. The recurrent answer: process one token at a time, carrying a fixed-size hidden state $h_t$ as a running summary of everything seen:

$$h_t = \tanh(W_{xh} x_t + W_{hh} h_{t-1} + b), \qquad y_t = W_{hy} h_t$$

Three properties define the design:

• Parameter sharing across time — the same $W$ at every step (like convolution shares across space): generalization across positions, and any-length sequences.
• O(1) state: memory of the whole past is compressed into $h$ — exactly the property that makes RNN inference cheap (foreshadowing: this is what Mamba resurrects, model-accuracy Phase 02 Ch. 8).
• Inherently sequential: $h_t$ needs $h_{t-1}$ — training cannot parallelize across time. Hold that thought for Chapter 7.

Chapter 4: Backpropagation Through Time and the Vanishing Gradient

Training unrolls the recurrence into a deep computation graph (one "layer" per time step) and backpropagates — BPTT. The gradient from a loss at step $t$ to the state at step $k$ passes through a product of Jacobians:

$$\frac{\partial h_t}{\partial h_k} = \prod_{i=k+1}^{t} \frac{\partial h_i}{\partial h_{i-1}} = \prod_{i} ,\text{diag}(\tanh') , W_{hh}^\top$$

A product of $t-k$ matrices: if their spectral norms sit below 1, the product decays exponentially — gradients from distant errors vanish, and the network simply cannot learn long-range dependencies (it could represent them; it can't be taught them). Norms above 1 explode instead — sudden loss spikes, NaNs.

The standard mitigations (all in your lab): gradient clipping for explosion (cap the global norm — crude, universal, still used in every LLM training run today), truncated BPTT (backprop only K steps — bounds cost, also bounds learnable dependency length), careful init (orthogonal $W_{hh}$), and — the real fix — architectural: Chapter 5.

Chapter 5: LSTM and GRU — Gating as Gradient Plumbing

The LSTM's move: add a cell state $c_t$ — a conveyor belt modified only by element-wise, gated operations:

$$c_t = f_t \odot c_{t-1} + i_t \odot \tilde{c}_t$$

with forget gate $f_t$, input gate $i_t$, candidate $\tilde{c}t$, and output gate $o_t$ producing $h_t = o_t \odot \tanh(c_t)$ — each gate a small sigmoid network of $(x_t, h{t-1})$.

Why this fixes vanishing: the gradient path along $c$ is multiplication by $f_t$ per step — no repeated $W_{hh}$ matrix product, no tanh-derivative shrinkage. With $f_t \approx 1$, gradients flow back nearly unattenuated for as long as the forget gate chooses. The gates make memory learned and content-dependent: keep this, overwrite that. (Squint and you can see both the residual stream of transformers and Mamba's selective state — gating-as-gradient-highway is one of deep learning's most recycled ideas.) Practical detail the lab uses: initialize forget-gate bias positive (~1.0) so training starts in "remember" mode.

GRU: the 2014 simplification — merges cell and hidden state, two gates (update/reset), ~25% fewer parameters, usually within noise of LSTM quality. The default when you want recurrence cheap.

Chapter 6: Sampling from a Language Model

The model outputs logits → softmax → a distribution. How you pick from it shapes everything users see (this section is permanent knowledge — identical for your char-RNN and for GPT-4):

• Greedy (argmax): deterministic, repetitive, gets trapped in loops ("the the the") — fine for short factual continuations only.
• Temperature $T$: divide logits by $T$ before softmax. $T \to 0$ approaches greedy; $T = 1$ is the model's honest distribution; $T > 1$ flattens toward chaos. The lab's most instructive experiment: the same checkpoint at T = 0.3 / 0.8 / 1.5 — coherence vs creativity as one knob.
• Top-k: zero out all but the k highest logits, renormalize — a hard cutoff on the tail where degenerate tokens live.
• Top-p (nucleus): keep the smallest set whose cumulative probability ≥ p — adaptive cutoff: narrow when the model is confident, wide when uncertain; generally preferred over fixed k.
• Combinations apply in order (temperature → top-k → top-p), and every serving stack (Phase 09) implements precisely this pipeline per token.

Chapter 7: Why RNNs Lost — and Where They Won

The scorecard against what's coming in Phase 04:

Transformers won on the training column — parallelism let them eat the whole internet, and the O(1) gradient path made long-range learning easy rather than heroic. But read the inference column: RNNs are the better inference shape, and that trade resurfaces constantly in your career — streaming/edge workloads, the KV-cache memory wall (Phase 09), and the state-space-model renaissance (Mamba) which is explicitly "RNN inference with parallelizable training." Phase 03 isn't history class; it's the other pole of a tradeoff you'll navigate professionally.

Lab Walkthrough Guidance

Lab 01 — Char-RNN (Karpathy's classic, built honestly):

1. Data first: character vocab over your corpus (Shakespeare is traditional), contiguous batching with correct (input, target=input-shifted-by-1) alignment — off-by-one here trains a copy machine; test the alignment on a tiny string.
2. Implement the vanilla RNN cell from Chapter 3's equations yourself (then optionally swap nn.LSTM); train with truncated BPTT — carry the hidden state across chunks (detached!) so the model sees long context without unbounded graphs.
3. Add gradient-norm clipping; log the pre-clip norm — watching it spike is the vanishing/exploding chapter made visible.
4. Track bits-per-character; sample at fixed prompts every N steps at several temperatures — the qualitative arc (noise → words → grammar → style) is the most instructive training curve in the curriculum; save the samples.
5. Compare vanilla vs LSTM on the same budget: loss curves and long-range behavior (does a quote opened 200 chars ago get closed?).

Success Criteria

You are ready for Phase 04 when you can, from memory:

1. Write the chain-rule factorization and explain why every LM reduces to next-token distribution modeling.
2. Connect cross-entropy ↔ likelihood ↔ perplexity and calibrate a bits-per-char number.
3. Derive (sketch) the Jacobian-product argument for vanishing/exploding gradients and name the four mitigations.
4. Explain the LSTM cell-state gradient path and why $f_t$ replaces $W_{hh}$-products.
5. Implement temperature/top-k/top-p from logits on a whiteboard, with what each fixes.
6. Reproduce the RNN-vs-transformer scorecard and argue the inference column's modern relevance.

Interview Q&A

Q: Why did transformers replace LSTMs — and what did we give up? Two structural wins: training parallelism across sequence positions (the GPU-era scaling unlock) and O(1) gradient paths between any two tokens (long-range learning by construction instead of through a gated bottleneck). We gave up O(1) inference state — transformers pay a KV cache that grows with context and dominates serving memory. Naming the loss is the senior half of the answer; SSMs/Mamba exist precisely to claw it back.

Q: Your generation loops endlessly repeating a phrase. Diagnose across the stack. Decoding first: greedy or near-zero temperature makes repetition self-reinforcing (the repeated phrase becomes ever more likely in context) — raise temperature, add top-p, or apply a repetition penalty. If it persists: degenerate model (undertrained, or trained on duplicated data — Phase 10's dedup matters here). The mechanism to articulate: argmax decoding + a model that locally overweights recent n-grams = a fixed point; sampling breaks the loop stochastically.

Q: What does gradient clipping actually do, and what does it not do? It rescales the gradient norm when it exceeds a threshold — direction preserved, magnitude capped — turning explosion (rare, catastrophic steps) into survivable ones. It does nothing for vanishing (you can't rescale a signal that's already ~0), which needs architectural fixes (gating, residuals) — the two pathologies are opposite and people conflate them; the distinction is the question's point. Still standard in LLM training (spikes from data/loss anomalies), so it's not retro knowledge.

References

• Karpathy, The Unreasonable Effectiveness of Recurrent Neural Networks (2015) — the lab's spiritual source; read it with samples open
• Hochreiter & Schmidhuber, Long Short-Term Memory (1997)
• Pascanu et al., On the difficulty of training recurrent neural networks (2013) — arXiv:1211.5063 — the vanishing/exploding analysis + clipping
• Cho et al., Learning Phrase Representations using RNN Encoder–Decoder (2014) — GRU
• Holtzman et al., The Curious Case of Neural Text Degeneration (2020) — arXiv:1904.09751 — nucleus sampling and why greedy degenerates
• Olah, Understanding LSTM Networks — colah.github.io — the canonical visual walkthrough
• Jurafsky & Martin, ch. 3 (n-grams) and ch. 8–9 (RNNs) for the textbook depth

🛸 Hitchhiker's Guide — Phase 3: RNNs and Language Modeling

Read this if: You want to internalize why every modern LLM is a "language model", what perplexity means, and where the conceptual bridges are between an RNN and a Transformer. RNNs are not in production for new LLMs in 2026 (transformers and SSMs replaced them) — but their failure modes are exactly what attention was invented to fix, so understanding them sharpens transformer intuition immensely.

0. The 30-second mental model

A language model is a probability distribution over sequences:

$$ P(w_1, w_2, \ldots, w_T) = \prod_{t=1}^T P(w_t \mid w_1, \ldots, w_{t-1}) $$

A neural language model parameterizes that conditional with a network. An RNN maintains a recurrent hidden state h_t = f(h_{t-1}, x_t) that's supposed to summarize all prior tokens; an LSTM does the same with gates that protect against vanishing gradients; a transformer throws the recurrence away and lets every token attend to every other token in parallel. Same task, three architectures.

By the end of Phase 3 you should:

• Know what an n-gram baseline gives you and why it's the floor for any LM evaluation.
• Be able to derive Backpropagation Through Time on the whiteboard.
• Explain vanishing/exploding gradients and how LSTM gates fix them.
• Compute and interpret perplexity, bits-per-character, and bits-per-byte.
• Implement a character-level RNN from raw cells (no nn.RNN) and use it to generate Shakespearean text.

1. Language modeling as a discipline

1.1 Why predict the next token?

Because everything is the next token. Translation, summarization, code generation, chat — they're all "given some prefix, what comes next?" If a model assigns high probability to true continuations across a vast and diverse corpus, it has implicitly learned grammar, facts, reasoning patterns, style, and code structure. This is the core hypothesis on which every LLM stands.

1.2 The chain rule and the autoregressive factorization

Any joint distribution over a sequence factorizes as a product of conditionals (chain rule of probability). A model that computes P(w_t | w_<t) for every t is sufficient to:

• Score any sequence (just multiply).
• Sample from the model (sample one token at a time, append, repeat).

That's the autoregressive style. There are non-AR alternatives (BERT-style masked LM, diffusion LMs, SSMs) but AR has won for generation.

1.3 Cross-entropy = next-token loss

Training a language model means minimizing the negative log-likelihood of the true next token at every position:

$$ \mathcal{L} = -\sum_t \log P(w_t \mid w_{<t}; \theta) $$

Equivalently: cross-entropy between the model's distribution and the one-hot distribution at the true token. This is the loss. Pretraining, fine-tuning, distillation all start from this.

2. n-gram models — your baseline

Before deep learning, language models were tables of conditional probabilities:

$$ P(w_t \mid w_{t-n+1}, \ldots, w_{t-1}) = \frac{\text{count}(w_{t-n+1}, \ldots, w_t)}{\text{count}(w_{t-n+1}, \ldots, w_{t-1})} $$

For unseen n-grams: smoothing (add-1, Kneser-Ney). Kneser-Ney is the gold-standard pre-deep-learning smoothing. Read Jurafsky & Martin Ch. 3.

A 5-gram Kneser-Ney model on 1B words gets ~80 perplexity on PTB. A modern transformer LM gets ~10–20. Always include the n-gram baseline before claiming your model is good.

Reference: Jurafsky & Martin, Speech and Language Processing, 3rd ed., Ch. 3 (free draft).

3. The Recurrent Neural Network

3.1 The vanilla RNN cell

$$ h_t = \tanh(W_{xh} x_t + W_{hh} h_{t-1} + b_h) $$

That's it. The hidden state h_t is a fixed-size vector (e.g., 256 dims) that's supposed to summarize all prior tokens. The output prediction is a softmax over W_{hy} h_t + b_y.

Because the same W_{hh} is applied at every step, the network has a fixed parameter count regardless of sequence length. That's beautiful — and dooms it.

3.2 Backpropagation Through Time (BPTT)

To train, "unroll" the recurrence into a deep feed-forward network of length T. Apply standard backprop. The gradient of the loss with respect to h_0 involves a product:

$$ \frac{\partial \mathcal{L}}{\partial h_0} \propto \prod_{t=1}^T \frac{\partial h_t}{\partial h_{t-1}} = \prod_{t=1}^T W_{hh}^\top , \text{diag}(\tanh'(\cdot)) $$

This is a long product of matrices.

• If the spectral radius of W_{hh} < 1 (and tanh' ≤ 1), the product vanishes. The model can't learn long-range dependencies.
• If > 1, the product explodes.

Both are catastrophic. Vanishing is the more common problem. Exploding is mitigated cheaply by gradient clipping (torch.nn.utils.clip_grad_norm_).

For long sequences: truncated BPTT — backprop only through the last K steps; detach the hidden state across boundaries.

3.3 LSTM — gating to the rescue

Hochreiter & Schmidhuber (1997). Add a cell state c_t that flows through with mostly identity-like updates, controlled by three gates (forget f, input i, output o):

$$ \begin{aligned} f_t &= \sigma(W_f [x_t, h_{t-1}] + b_f) \ i_t &= \sigma(W_i [x_t, h_{t-1}] + b_i) \ o_t &= \sigma(W_o [x_t, h_{t-1}] + b_o) \ g_t &= \tanh(W_g [x_t, h_{t-1}] + b_g) \ c_t &= f_t \odot c_{t-1} + i_t \odot g_t \ h_t &= o_t \odot \tanh(c_t) \end{aligned} $$

Why it works: the cell state c_t is updated additively (c_{t-1} + ...), so the gradient through c is roughly the identity matrix times the forget gate. If the forget gate is near 1, gradients flow through hundreds of steps without vanishing.

3.4 GRU — fewer gates

Cho et al. (2014). Merges forget+input into a single gate. Slightly fewer params; usually comparable to LSTM in practice.

3.5 Stacking and bidirectionality

• Stacked: feed h_t^{(1)} of layer 1 as input to layer 2. Each layer learns higher-level features. Beyond ~3 layers, returns diminish.
• Bidirectional: a forward RNN + a backward RNN; concatenate. Useful for tagging/classification but not for autoregressive generation (you can't see the future at inference).

3.6 Why RNNs lost to Transformers

The last row is interesting — RNNs have constant memory at inference, which is why State Space Models (Mamba, S5, Hyena) are mounting a comeback for very long contexts. A modern RNN literacy still matters.

4. Perplexity and friends

4.1 Perplexity

$$ \text{PPL} = \exp\left(\frac{1}{N} \sum_{i=1}^N -\log P(w_i \mid w_{<i})\right) = \exp(\bar{\mathcal{L}}) $$

Intuition: "if the model treated every step as a uniform choice over PPL options, it would have the same loss." Lower is better. PPL = vocab_size means random; PPL = 1 means perfect.

PPL is not comparable across tokenizers — a model with a 50k subword vocab cannot be PPL-compared to a model with a 30k vocab. To compare across tokenizers use:

4.2 Bits-per-character (BPC) / Bits-per-byte (BPB)

$$ \text{BPB} = \frac{\text{loss in nats} \cdot \log_2 e}{\text{number of bytes in the original text}} $$

Because bytes are tokenizer-agnostic, BPB lets you fairly compare any LM. State-of-the-art LMs on enwik8 reach ~0.94 BPB.

4.3 What "good" perplexity looks like

• 5-gram Kneser-Ney on PTB: ~80 PPL.
• Char-RNN on Tiny Shakespeare (small): ~5–10 PPL (chars are easier per-step).
• GPT-2 small on WikiText-103: ~30 PPL.
• GPT-3 175B on PTB: ~20 PPL.
• Frontier LLMs on web text: ~6–10 PPL on held-out web.

5. Sampling from a language model

You'll meet these again in Phase 9. Preview:

• Greedy (argmax): deterministic; can repeat.
• Beam search: keep top-k partial sequences. Better for translation; rare in chat (boring outputs).
• Temperature: divide logits by T. T < 1 sharpens, T > 1 flattens.
• Top-k: sample only from the k most-likely tokens.
• Top-p (nucleus): sample from the smallest set whose cumulative prob ≥ p. Adapts to entropy.
• Repetition penalty / no-repeat n-gram: hacks to prevent loops.

6. Lab 01 walkthrough (lab-01-char-rnn)

6.1 What you'll build

• A VanillaRNNCell — implemented as the raw tanh(Wxh x + Whh h) math, not nn.RNN. The point is to see autograd handle BPTT.
• A CharRNN module — embedding → stacked RNN cells → linear projection to vocab.
• A train() loop that processes Tiny Shakespeare in fixed-length sequences, with TBPTT (detach() the hidden state between batches).
• A sample() method that generates new text given a seed string.

6.2 Things to internalize while reading the solution

• Why detach() between batches? Without it, autograd builds an infinitely long graph and OOMs. Detaching pretends the prior hidden state is a constant input.
• Why is the loss reshaped to (B*T, V) for cross-entropy? Because F.cross_entropy expects a 2D logits tensor and a 1D target tensor. The (B, T) structure is irrelevant to the per-position loss.
• Why no causal mask? Because RNNs are causal by construction — h_t only depends on h_{<t}.
• Why stack the cells but not parallelize them? Each layer must wait for the previous layer's output at the same time step. Sequence dimension is sequential; layer dimension can be batched in a single for loop with shared compute pattern.

6.3 Watch the loss curve

Early training: loss drops fast as the model learns the unigram distribution. After a few hundred steps it learns bigram statistics, then short word fragments, then real words, then word ordering. By 5k steps, it should produce something that looks like Shakespeare-flavored gibberish. By 20k+, full pseudo-grammatical lines. (Famous Karpathy 2015 blog post.)

7. References

• Karpathy, The Unreasonable Effectiveness of Recurrent Neural Networks (2015) — required reading.
• Olah, Understanding LSTM Networks (2015) — required reading; the diagrams.
• Hochreiter & Schmidhuber (1997), Long Short-Term Memory.
• Cho et al. (2014), Learning Phrase Representations using RNN Encoder–Decoder for Statistical Machine Translation — GRU.
• Sutskever, Vinyals, Le (2014), Sequence to Sequence Learning with Neural Networks — the seq2seq paper.
• Bahdanau, Cho, Bengio (2015), Neural Machine Translation by Jointly Learning to Align and Translate — the attention paper that started everything Phase 4 covers.
• Jurafsky & Martin, Speech and Language Processing, 3rd ed., Ch. 9 (RNNs and LSTMs).
• Deep Learning (Goodfellow, Bengio, Courville), Ch. 10.
• Pascanu, Mikolov, Bengio (2013), On the difficulty of training recurrent neural networks — the vanishing/exploding gradient analysis.

8. Common interview questions on Phase 3 material

1. Walk me through BPTT on a 3-step RNN.
2. What causes vanishing gradients in vanilla RNNs and how do LSTMs help?
3. Compute perplexity from a cross-entropy loss of 2.3 nats per token.
4. Why is BPB more honest than PPL across tokenizers?
5. What's a Kneser-Ney 5-gram baseline and when is it competitive?
6. Why didn't RNNs scale to GPT-3 sizes?
7. What's truncated BPTT and why do we need it?
8. Compare LSTM vs GRU.
9. Why are state-space models (Mamba) suddenly interesting again?
10. Implement an LSTM cell on a whiteboard.

9. From solid → exceptional

• Reimplement an LSTM cell from scratch (no nn.LSTMCell) and train on Tiny Shakespeare. Compare loss curves and sample quality vs vanilla RNN.
• Reproduce Karpathy's char-RNN results on Linux source code; show the model learns to balance braces and indent.
• Implement a GRU alongside; benchmark perplexity at equal parameter count.
• Train a 1-layer LSTM on enwik8; compute BPB; compare to the famous IndyLSTM / mLSTM numbers (~1.0 BPB).
• Read the original attention paper (Bahdanau 2015) and implement attention as an add-on to a seq2seq RNN encoder-decoder. This gives you the conceptual bridge to Phase 4.
• Skim the Mamba paper (Gu & Dao, 2023) and write a one-page comparison: how is Mamba different from an LSTM?

10. Recommended cadence

Lab 01 — Char-Level RNN (Solution Walkthrough)

Phase: 3 — RNNs & Language Modeling | Difficulty: ⭐⭐⭐☆☆ | Time: 2–4 hours

Concept primer: ../HITCHHIKERS-GUIDE.md §RNNs and §BPTT. This document walks through solution.py.

Run

pip install -r requirements.txt
curl -O https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
python solution.py --data input.txt --steps 2000

0. The mission

Train a vanilla RNN (no nn.RNN, no LSTM — by hand) to predict the next character on Tiny Shakespeare (~1 MB). At the end you'll sample text that looks Elizabethan even if it's nonsense. Karpathy's 2015 demo that proved RNNs could do generative modeling.

You are deliberately implementing the worst sequence model — the simplest possible RNN — to feel why Transformers were invented:

• Sequential decoding (no parallel forward over T positions).
• Vanishing gradients past ~50 timesteps.
• Information bottleneck through a single hidden state.

1. The math

A vanilla recurrent cell at each step:

$$ h_t = \tanh(W_{ih} x_t + W_{hh} h_{t-1} + b) $$

For LM, project to logits: $\hat{p}t = \mathrm{softmax}(W{ho} h_t)$. Train with cross-entropy on next-character.

2. VanillaRNNCell

class VanillaRNNCell(nn.Module):
    def __init__(self, in_dim, hidden_dim):
        super().__init__()
        self.W_ih = nn.Linear(in_dim, hidden_dim, bias=False)
        self.W_hh = nn.Linear(hidden_dim, hidden_dim, bias=True)

    def forward(self, x, h):
        return torch.tanh(self.W_ih(x) + self.W_hh(h))

• Two linear layers, one bias. Convention: bias on the recurrent path only (input path's bias is redundant after summing).
• tanh not ReLU. ReLU + recurrent multiplication is unstable: positive activations grow without bound across time-steps. tanh ∈ [-1, 1] keeps state bounded.

3. CharRNN

class CharRNN(nn.Module):
    def __init__(self, vocab_size, hidden_dim=256):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, hidden_dim)
        self.cell = VanillaRNNCell(hidden_dim, hidden_dim)
        self.head = nn.Linear(hidden_dim, vocab_size)
        self.hidden_dim = hidden_dim

    def forward(self, x):
        B, T = x.shape
        h = x.new_zeros(B, self.hidden_dim, dtype=torch.float)
        e = self.embed(x)
        outs = []
        for t in range(T):
            h = self.cell(e[:, t], h)
            outs.append(h)
        out = torch.stack(outs, dim=1)
        return self.head(out)

The for loop over time is the heart of the inefficiency. Every forward call serializes T cell evaluations. For T=128, that's 128 sequential CUDA launches → tiny kernels → GPU idle most of the time. Compare a Transformer where the whole sequence is processed in one matmul.

x.new_zeros(...) creates a zero tensor on the same device + dtype family as x — avoids manual .to(device).

4. sample — autoregressive generation

@torch.no_grad()
def sample(self, ctx, n, temperature=1.0):
    h = ctx.new_zeros(1, self.hidden_dim, dtype=torch.float)
    # Warm up on context
    for t in range(ctx.size(1)):
        e = self.embed(ctx[:, t])
        h = self.cell(e, h)
    out = [ctx]
    last = ctx[:, -1]
    for _ in range(n):
        e = self.embed(last)
        h = self.cell(e, h)
        logits = self.head(h) / max(1e-6, temperature)
        probs = F.softmax(logits, dim=-1)
        last = torch.multinomial(probs, 1).squeeze(-1)
        out.append(last.unsqueeze(1))
    return torch.cat(out, dim=1)

Two phases — exactly the same pattern you'll see in Phase 9's KV-cache lab:

1. Warm-up / prefill: run the cell across the prompt to populate the hidden state.
2. Decode: feed back the previously sampled token, take one cell step, sample.

Temperature: divides logits before softmax.

• T = 1: model's natural distribution.
• T < 1: sharper → more confident → more repetition.
• T > 1: flatter → more diverse → more nonsense.

5. The training loop

data = torch.tensor([stoi[c] for c in text], dtype=torch.long)

def get_batch():
    ix = torch.randint(0, len(data) - args.seq_len - 1, (args.batch,))
    x = torch.stack([data[i:i + args.seq_len] for i in ix])
    y = torch.stack([data[i + 1:i + 1 + args.seq_len] for i in ix])
    return x.to(device), y.to(device)

For Tiny Shakespeare: vocab ~65 chars. Whole dataset fits as a single 1.1M-element tensor.

This is truncated BPTT: gradients only flow within each seq_len-long chunk; we never connect chunks across batch boundaries. For seq_len=128 we backprop through 128 timesteps. Beyond that, vanishing gradients would erase the signal anyway.

torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)

Non-optional for RNNs. Without it, exploding gradients (which RNNs do regularly) cause loss spikes and NaN. The clip threshold of 5.0 is empirically standard.

6. Expected output

Tiny Shakespeare, 2000 steps, hidden=256, seq_len=128, on a 4090 (~3 minutes):

step     0  loss=4.1742  ppl=64.97
step  1000  loss=1.7634  ppl=5.83
step  2000  loss=1.5897  ppl=4.91

----- T=0.8 -----
ROMEO: I will not be a man, and the king shall be the world,
And the world is the world's the world to thee...

Sanity numbers:

• Initial loss ≈ log(65) ≈ 4.17. ✅
• After training: loss ≈ 1.5–1.7, perplexity 4.5–5.5. Vanilla RNN can't go much lower; LSTMs reach ~1.4, Transformers ~1.2.
• BPC (bits per char) = loss / log(2) ≈ 2.3.

7. Why is this so much worse than a Transformer?

Compared to Phase 4's MiniGPT on the same data:

Two reasons:

1. Vanishing gradient — info from 100 chars ago contributes ~0 to the current hidden state. The Transformer attends directly with no decay.
2. Single bottleneck — entire history compressed into one 256-vector. Attention has 128×256 effective state.

The lab is about feeling this gap, not closing it.

8. Common pitfalls

1. Forgetting clip_grad_norm_ → loss explodes around step ~500 with nan output.
2. Don't carry h across batches in this lab — that's stateful RNN training, more complex.
3. reshape vs view — view requires contiguous memory; reshape doesn't. The model output from torch.stack is contiguous.
4. Forgetting @torch.no_grad() on sample — slowdown 2–3× and OOM on long generations.

9. Stretch exercises

• Replace VanillaRNNCell with LSTMCell (still by hand). LSTM has 4 gates: input, forget, cell, output. Train for 2000 steps; expect loss → ~1.4 (vs 1.6 for vanilla).
• Implement GRU (3 gates). Compare to LSTM.
• Add layer normalization inside the cell — stabilizes longer-context training.
• Statefulness: carry h across batches within an epoch; reset at epoch boundary.
• Bigger context: train with seq_len=256 or 512. Watch loss saturate earlier than the Transformer would.
• Sampling tricks: implement top-k and top-p (nucleus) sampling. Compare quality at the same temperature.
• Time it: profile and confirm the for-loop over T dominates wall-time. The exact reason Transformers won.

10. What this lab proves about you

You can implement an autoregressive sequence model from scratch, train it stably, sample from it, and articulate exactly why this architecture lost to attention. Phase-3 milestone.

Phase 4 — Attention & Transformers (From Scratch)

Difficulty: ⭐⭐⭐⭐☆ | Estimated Time: 2 weeks Roles supported: ALL research-engineer roles. The single most-asked LLM interview topic.

Why This Phase Exists

If you can derive scaled dot-product attention on a whiteboard, implement multi-head attention in <50 lines, explain RoPE, and walk through one forward pass of a transformer block — you pass the technical bar of nearly every LLM-engineering interview I have seen.

This is the most important phase. Do not rush it.

Concepts

• Self-attention as content-based addressable memory
• Scaled dot-product attention: softmax(QK^T / √d_k) V
• Why divide by √d_k (variance argument)
• Causal masking (decoder) vs padding masking (encoder)
• Multi-head attention: parallel subspace projections
• Positional encoding flavors:
  • Sinusoidal (original Transformer)
  • Learned absolute
  • RoPE (rotary, used in Llama / Qwen / most modern decoders)
  • ALiBi (used in MPT / BLOOM)
• Layer normalization vs RMSNorm
• Pre-norm vs post-norm (training stability)
• Residual stream view (Anthropic's interpretability framing)
• Feed-forward block (MLP) — usually 4× hidden dim, GELU/SwiGLU
• Encoder vs decoder vs encoder-decoder topology
• Parameter counting

Labs

Lab 01 — Scaled Dot-Product Attention From Scratch

Lab 02 — Multi-Head Attention

Lab 03 — Positional Encodings: Sinusoidal, RoPE, ALiBi

Lab 04 — Mini Transformer Block + Full Decoder

Deliverables Checklist

• Attention implementation (3 ways) with tests
• Multi-head attention benchmarked against nn.MultiheadAttention
• Positional-encoding ablation report
• 200-line transformer that overfits a single batch

Interview Relevance

This phase is the technical heart of LLM interviews. Expect:

• Whiteboard derivation of attention
• "Implement multi-head attention in 30 minutes"
• "Compare RoPE and ALiBi"
• "Walk through a transformer block"
• Parameter-count math problems

Warmup Guide — Attention & Transformers

Zero-to-expert primer for Phase 04: the transformer derived piece by piece — why each component exists, what breaks without it — culminating in the decoder-only GPT architecture you implement in the mini-transformer lab.

Table of Contents

• Chapter 1: The Problem Attention Solves
• Chapter 2: Scaled Dot-Product Attention, Derived
• Chapter 3: Causal Masking — Training on Every Position at Once
• Chapter 4: Multi-Head Attention
• Chapter 5: The Rest of the Block — FFN, Residuals, LayerNorm
• Chapter 6: Position Information
• Chapter 7: The Full GPT — Assembly and Parameter Accounting
• Chapter 8: Encoder vs Decoder vs Encoder-Decoder
• Lab Walkthrough Guidance
• Success Criteria
• Interview Q&A
• References

Chapter 1: The Problem Attention Solves

Phase 03 ended with the RNN's twin constraints: information from token 5 reaches token 500 only through 495 sequential state updates (a lossy game of telephone), and training can't parallelize over time. Attention dissolves both with one move: let every position directly query every other position, with content-dependent weights, all positions computed simultaneously. Path length between any two tokens: 1. Training parallelism: total. The price — O(n²) pairs — was accepted in 2017 and has been the field's central engineering battle ever since (this track's Phase 09; the model-accuracy track's FlashAttention). Phase 04 is where you internalize what that price buys.

Chapter 2: Scaled Dot-Product Attention, Derived

Build it from the retrieval metaphor: each token asks a question and offers an answer. From each input $x_i$, three learned projections:

• query $q_i = W_Q x_i$ — what am I looking for?
• key $k_i = W_K x_i$ — what can I be found by?
• value $v_i = W_V x_i$ — what do I contribute if selected?

Relevance of $j$ to $i$: the dot product $q_i \cdot k_j$. Softmax the scores into weights; output the weighted sum of values:

$$\text{Attention}(Q, K, V) = \text{softmax}!\left(\frac{QK^\top}{\sqrt{d_k}}\right) V$$

Why each piece:

• Separate Q and K (not $x_i \cdot x_j$ directly): asymmetry — "looking for" and "findable as" are different roles ("it" queries for nouns; it doesn't offer itself as one).
• Separate V: what a token contributes differs from what it matches on.
• $\sqrt{d_k}$: components i.i.d. with unit variance make $q \cdot k$ have variance $d_k$; at $d_k = 64$, unscaled logits of magnitude ~8 saturate softmax — winner-take- all weights and vanishing gradients through the softmax. Dividing by $\sqrt{d_k}$ restores unit variance. (Verify empirically in the lab — it's a two-line experiment.)
• Softmax: turns arbitrary scores into a convex combination — outputs stay in the values' span, bounded, differentiable.

Chapter 3: Causal Masking — Training on Every Position at Once

A language model must not see the future (Phase 03 Ch. 1's conditional). In attention, enforcement is brutally simple: before softmax, set $S_{ij} = -\infty$ for $j > i$ (an upper-triangular mask). After softmax those weights are exactly 0.

The consequence that makes LLM training affordable: one forward pass over a sequence trains all n positions simultaneously — position i's output predicts token i+1, every position is a training example, and the mask guarantees no leakage. This "teacher forcing in parallel" is the transformer's economic engine; the lab's correctness test (perturb a future token, assert outputs at earlier positions are bit-identical) is the single most important test you'll write in this phase.

Chapter 4: Multi-Head Attention

One attention computes one weighted average per position — one "relation type" at a time. Multi-head splits $d_{model}$ into $h$ subspaces of $d_k = d_{model}/h$, runs attention independently in each, concatenates, and projects ($W_O$):

$$\text{MHA}(X) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h),W_O$$

Same total FLOPs as one full-width head — but $h$ different relation patterns (syntactic heads, positional heads, rare-token heads emerge under analysis). Two implementation truths the lab teaches: it's all done as one batched tensor op (reshape to (B, h, n, d_k), never a Python loop over heads), and the (B, h, n, n) attention-weights tensor is the memory hog — materializing it is what FlashAttention later avoids (model-accuracy Phase 08, when you get there).

Chapter 5: The Rest of the Block — FFN, Residuals, LayerNorm

A transformer block is attention plus three pieces, each load-bearing:

• FFN: per-position MLP, $W_2,\text{GELU}(W_1 x)$ with hidden dim $4\times d_{model}$. Attention moves information between positions; the FFN processes it pointwise — ~⅔ of the model's parameters live here, and mechanistic work reads it as key-value memory over features. No interaction across positions (that's attention's job alone).
• Residual connections: $x + \text{Sublayer}(x)$ — the gradient highway (the LSTM cell-state idea, Phase 03 Ch. 5, reborn as architecture). Depth-100 stacks train because identity paths exist around every sublayer.
• LayerNorm: re-center/re-scale each token's vector (learned $\gamma, \beta$) — keeps activation scales stable across depth. Pre-norm placement (norm before each sublayer, used by GPT-2 onward) leaves the residual path untouched and trains stably without warmup heroics; post-norm (the 2017 original) normalizes the sum and destabilizes deep stacks. Your lab uses pre-norm; know why.

Block: x = x + Attn(LN(x)); x = x + FFN(LN(x)). Stack N times. That's the architecture.

Chapter 6: Position Information

Attention is permutation-equivariant — shuffle inputs, get shuffled outputs. Order must be injected:

• Learned absolute embeddings (GPT-2, and your lab): a trainable vector per position index, added to token embeddings. Simple, effective, but no extrapolation beyond trained length and no explicit relative structure.
• Sinusoidal (the 2017 original): fixed sin/cos at geometric frequencies — no parameters, theoretical extrapolation (rarely realized in practice).
• The modern answers — RoPE and ALiBi — rotate Q/K by position-dependent angles (relative position emerges in the dot product) or bias logits by distance. They're covered in depth in the model-accuracy track's Phase 02 warmup (Ch. 5); for this phase, know that learned-absolute is the pedagogical baseline and why its extrapolation fails (untrained rows in the position table).

Chapter 7: The Full GPT — Assembly and Parameter Accounting

tokens → token_emb (V×d) + pos_emb (n_ctx×d)
       → N × [pre-norm attention block + pre-norm FFN block]
       → final LayerNorm → lm_head (d×V) → logits

Two assembly details with outsized importance:

• Weight tying: lm_head shares the token-embedding matrix (transpose) — saves V×d parameters (significant at small scale) and improves quality; both GPT-2 and your lab do it.
• Parameter accounting (do this once by hand — the lab asks for it): per block ≈ $12 d^2$ ($4d^2$ attention QKVO + $8d^2$ FFN); total ≈ $12 N d^2$ + embeddings $Vd$. GPT-2-small: N=12, d=768, V=50257 → ~85M block + ~39M embedding ≈ 124M. ✓ Being able to do this arithmetic is how you sanity-check any config — and it's the FLOPs-per-token estimate ($\approx 2 \times$ params) from the model-accuracy track's Phase 07 in embryo.

Chapter 8: Encoder vs Decoder vs Encoder-Decoder

The same blocks, three wirings — know which and why:

• Decoder-only (GPT, LLaMA, ~everything generative): causal mask everywhere; one stack does both understanding and generation. Won because of training simplicity, the in-context-learning emergent bonus, and KV-cache-friendly inference.
• Encoder-only (BERT): bidirectional attention (no mask) + masked-token training — better representations per parameter for understanding tasks (classification, retrieval embeddings — your Phase 07 RAG encoders are these). Cannot generate autoregressively.
• Encoder-decoder (T5, translation, Whisper): bidirectional encoder over input, causal decoder with cross-attention (decoder queries, encoder keys/values) — still the right shape when input and output are genuinely different sequences/ modalities.

Lab Walkthrough Guidance

Lab 04 — Mini-Transformer (decoder-only GPT, trained on the Phase 03 corpus):

1. Single-head causal attention first; write the no-future-leakage test (perturb token j, assert positions < j unchanged) before training anything.
2. Verify the $\sqrt{d_k}$ claim empirically: log attention-logit variance with and without scaling at d_k = 64.
3. Multi-head via the batched reshape — test: output matches a loop-over-heads reference within 1e-6.
4. Assemble pre-norm blocks + weight tying; do the Chapter 7 parameter count and assert it matches sum(p.numel()) exactly — off-by-anything means a wiring bug.
5. Train on the same data as your char-RNN with the same budget; compare loss curves and samples — the gap is the phase's thesis. Then overfit a tiny batch to near-zero loss (the standard can-it-learn sanity check) before any long run.

Success Criteria

You are ready for Phase 05 when you can, from memory:

1. Derive attention from the Q/K/V retrieval story, justifying all three projections and the $\sqrt{d_k}$.
2. Explain how the causal mask enables all-positions-at-once training and write its correctness test.
3. State what FFN does that attention doesn't (and vice versa), and why pre-norm.
4. Count GPT-2-small's parameters on paper to within a few percent.
5. Choose decoder-only vs encoder-only vs enc-dec for: chat model, embedding model, speech-to-text — with reasons.
6. Name the O(n²) cost's two manifestations (training compute, inference KV/attention) — the bridge to Phases 05 and 09.

Interview Q&A

Q: Why three separate projections Q, K, V instead of using the embeddings directly? Roles differ: what a token searches for (Q), what it's discoverable by (K), and what it contributes when found (V) are three different functions of its content. Collapsing them forces symmetric attention ($x_i \cdot x_j$) and couples routing to content transport. Empirically and mechanistically, the asymmetry is where attention's expressiveness lives — heads implement little programs like "pronouns query for recent nouns," which need Q ≠ K.

Q: Walk me through why training a transformer LM is one parallel pass but generating is sequential. Training: all tokens are known, so all n next-token predictions compute simultaneously under the causal mask — the mask, not time, enforces order. Generation: token t+1's identity depends on sampling from position t's output — an inherent data dependency no parallelism removes. That asymmetry creates the prefill/decode split and the KV cache (Phase 09): cache K/V of the fixed prefix; each new token costs one query against cached history instead of recomputing the past.

Q: Remove the FFN entirely — what happens? You lose all per-position nonlinear processing; the network becomes (normed, gated) weighted averaging of value vectors — outputs confined near the span of linear transforms of inputs, and stacking attention alone is known to collapse toward rank-deficient mixing. Practically: most of the parameter budget and the "memory/ feature computation" capacity is in FFNs; models degrade catastrophically. The clean division — attention mixes across positions, FFN computes within them — is the architecture's actual design principle.

References

• Vaswani et al., Attention Is All You Need (2017) — arXiv:1706.03762
• Radford et al., GPT-2 report (2019) — the decoder-only recipe your lab follows
• Karpathy, Let's build GPT from scratch — youtube.com/watch?v=kCc8FmEb1nY — pairs exactly with this lab
• The Illustrated Transformer and The Annotated Transformer
• Xiong et al., On Layer Normalization in the Transformer Architecture (2020) — pre-norm vs post-norm, formally
• Elhage et al., A Mathematical Framework for Transformer Circuits (2021) — transformer-circuits.pub — heads-as-programs, for the curious
• Press & Wolf, Using the Output Embedding to Improve Language Models (2017) — weight tying

🛸 Hitchhiker's Guide — Phase 4: Attention and Transformers

Read this if: You want to be able to implement a transformer from scratch on a whiteboard, defend every design choice, and answer every variant of "explain attention" you'll get in an interview. This is the most important phase of the curriculum. Spend twice as long here as anywhere else.

0. The 30-second mental model

Attention is a content-based, weighted average. Given a query vector q and a set of key-value pairs {(k_i, v_i)}, compute similarities s_i = q · k_i, normalize them with softmax to get weights α_i, and return Σ α_i v_i. That's it. Everything else — multi-head, causal masking, RoPE, KV cache, FlashAttention — is a refinement of that one operation.

A transformer is a stack of "blocks", where each block applies (a) self-attention so every token can pull information from every other token, and (b) a position-wise MLP that processes each token's representation independently. Repeat 12, 32, 80, 96 times. Add a softmax head to predict the next token. Done.

By the end of Phase 4 you should:

• Derive scaled dot-product attention from first principles.
• Know exactly why we divide by √d_k, why we use multi-head, why we use causal masking.
• Implement RoPE (and explain why it's "relative" without an explicit (i-j)).
• Compare LayerNorm vs RMSNorm, GELU vs SwiGLU, post-norm vs pre-norm.
• Reason about KV-cache memory and its scaling.
• Implement a MiniGPT from blank file in 30 minutes (the lab does ~150 lines).

1. The road to attention

1.1 Why RNNs needed help

In a seq2seq translation model, the encoder RNN summarizes the source sentence into a single fixed vector — and the decoder must squeeze the entire meaning of "The agreement on the European Economic Area was signed in August 1992" through this bottleneck. Disaster on long sentences.

1.2 Bahdanau attention (2015)

Bahdanau, Cho, Bengio added an "alignment" mechanism: at each decoder step, look at all encoder hidden states and softmax over their similarities to the current decoder state. Now the decoder gets a weighted average focused on the source tokens that matter for the current target token. Translation quality jumped immediately.

This is the seed crystal. Everything after is "attention but more so".

1.3 Attention Is All You Need (Vaswani et al., 2017)

The Google Brain team noticed: if attention is so good, why have the RNN at all? Replace the recurrence with attention layers. Add positional encodings (so the model knows token order without recurrence). Stack. Train.

The result was the Transformer. Every modern foundation model — GPT-4, Claude 4, Gemini 2.5, Llama-3, Mistral, DeepSeek — is a descendent of this paper.

2. Scaled Dot-Product Attention (the unit)

2.1 The math

Inputs: queries Q ∈ ℝ^{T×d_k}, keys K ∈ ℝ^{T×d_k}, values V ∈ ℝ^{T×d_v}. Output:

$$ \text{Attention}(Q, K, V) = \text{softmax}!\left(\frac{Q K^\top}{\sqrt{d_k}}\right) V $$

Step-by-step:

1. S = Q K^⊤ / √d_k — pairwise scores. Shape (T, T). Each row S_i says how much token i cares about every other token.
2. P = softmax(S, dim=-1) — row-wise normalize.
3. O = P V — output is a weighted sum of value vectors.

2.2 Why divide by √d_k?

If Q and K entries have unit variance and zero mean, then the dot product q · k (a sum of d_k independent products) has variance d_k. For d_k = 64, that's stddev 8. Pushing such large values into softmax saturates it: most weight goes to one element, gradients vanish.

Dividing by √d_k keeps the score variance ≈ 1 regardless of d_k. This is purely a numerical-stability trick at initialization, not a "more correct" formulation.

2.3 Causal masking (for decoder-only LMs)

For autoregressive generation, token t must not attend to tokens > t. Implement by setting the upper triangular entries of S to -∞ before softmax:

mask = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
scores = scores.masked_fill(mask, float("-inf"))

After softmax those entries become 0. This is what makes the transformer a language model in the autoregressive sense.

2.4 Why is attention O(T²)?

The score matrix is T × T. For long context (32k, 128k, 1M), this is the bottleneck. FlashAttention (Dao 2022) doesn't reduce the FLOPs but eliminates the materialization of the matrix in HBM, dramatically improving wall-clock and memory. Sparse / linear attention (Reformer, Linformer, Performer, Longformer) trades quality for sub-quadratic compute. Phase 9 covers all of these.

3. Multi-Head Attention

3.1 The intuition

Different "heads" can specialize in different relationships: one head tracks subject–verb agreement, another co-references pronouns, another keeps positional adjacency. A single attention does one weighted average; h heads do h of them in parallel and concatenate.

3.2 The math (and the parameter count)

Pick n_heads and d_head such that n_heads × d_head = d_model. Project the input three times with shape-d_model × d_model matrices W_Q, W_K, W_V, then reshape the result into (B, n_heads, T, d_head). Run scaled dot-product attention per head, concatenate, project with W_O.

# (B, T, C)  —->  (B, n_heads, T, d_head)
q = self.W_q(x).view(B, T, n_heads, d_head).transpose(1, 2)

Total parameters in attention: 4 d_model² (Q, K, V, O). The MLP block is 8 d_model² (typically 4× expansion factor up and back). Each transformer block is ~12 d_model² parameters; total ≈ 12 d_model² × n_layers.

3.3 MHA → MQA → GQA

• MHA (vanilla): each head has its own K and V projections. Best quality, biggest KV cache.
• MQA (Shazeer 2019): all heads share one K and V. KV cache shrinks by n_heads×. Slight quality drop on hard tasks.
• GQA (Ainslie 2023): heads grouped; one K/V per group. Tunable middle ground (Llama-3 8B: 32 query heads, 8 KV groups). Now standard.

The motivation for MQA/GQA is inference: at long context, the KV cache dominates GPU memory, so reducing KV size directly increases batch-size headroom and throughput.

4. Position information

A transformer is permutation-equivariant without positional information — shuffle the input tokens and the output set is the same. We must inject positional signal somehow.

4.1 Sinusoidal positional encoding (Vaswani 2017)

Hand-designed sin/cos features added to the token embeddings. Each dimension oscillates at a different wavelength. Conceptually elegant; rarely used in modern LLMs.

4.2 Learned absolute positional embedding (BERT, GPT-2)

A learned (max_pos, d_model) matrix added to token embeddings. Simple but doesn't extrapolate beyond max_pos.

4.3 ALiBi (Press et al., 2022)

Adds a position-dependent bias to attention scores: s_{ij} ← s_{ij} - m · |i - j| for a per-head slope m. Linear penalty on distance. No vector positional encoding at all. Extrapolates to longer contexts than seen at train time.

4.4 RoPE (Su et al., 2021) — the modern winner

Rotary Positional Embedding rotates Q and K vectors by an angle that depends on position. Pair adjacent dimensions (x_{2i}, x_{2i+1}) into a 2D point, rotate by θ_i = pos · base^{-2i/d}. Critically, after rotation, the dot product q_i · k_j becomes a function purely of (i - j):

$$ q'_i \cdot k'_j = q_i \cdot k_j \cdot \cos((i-j)\theta) + (\text{cross terms involving } i-j) $$

So RoPE is relative without an explicit (i-j) term. Used by Llama, Mistral, Qwen, Gemma, and most open models.

Length extension tricks: NTK-aware scaling, YaRN, position interpolation. These adjust base or θ to extend a 4k-trained model to 32k or beyond at inference.

4.5 References

• Su et al. (2021), RoFormer.
• Press et al. (2022), Train Short, Test Long: Attention with Linear Biases (ALiBi).
• bloc97's NTK-aware RoPE blog post and YaRN (Peng et al. 2023).

5. The Transformer Block

5.1 The standard recipe (pre-norm, modern)

input x
  ┌─→ LayerNorm → CausalSelfAttention ─→ + (residual)
  │                                       │
  └───────────────────────────────────────┘
                                          │
  ┌─→ LayerNorm → MLP ───────────────────→ + (residual)
  │                                        │
  └────────────────────────────────────────┘
output

That is: x = x + Attn(LN(x)) then x = x + MLP(LN(x)). Repeat N times.

5.2 Pre-norm vs Post-norm

• Post-norm (original 2017): x = LN(x + Sublayer(x)). Gradients flow through the LayerNorm — vanish for deep stacks. Required learning-rate warmup gymnastics.
• Pre-norm: x = x + Sublayer(LN(x)). Gradient has a clean residual highway. Stable past 100+ layers.

Every modern LLM is pre-norm.

5.3 LayerNorm vs RMSNorm

LayerNorm: y = γ · (x - μ) / σ + β — subtract mean, divide by std, scale, shift.

RMSNorm: y = γ · x / RMS(x) — drop the mean subtraction, drop the bias. ~10% faster, no quality loss in practice. Used by Llama, Mistral, Qwen.

Why does dropping the mean work? Empirical observation backed by some analysis: the centering operation is largely redundant once activations are well-conditioned at depth.

5.4 The MLP block

mlp_out = down_proj(activation(up_proj(x)))

For most transformers, up_proj expands by 4× (so a d_model = 4096 model has a 16384-wide hidden layer in the MLP). Activation choices:

• ReLU: original; rarely used now.
• GELU: smooth ReLU; used by GPT-2, BERT.
• SwiGLU (Shazeer 2020): (W_up x) ⊙ silu(W_gate x) — gated linear unit with Swish gating. Costs 50% more params but better quality at fixed FLOPs. Used by Llama, Qwen, Mistral.

5.5 Weight tying

The token embedding matrix E ∈ ℝ^{V × d} and the LM head matrix W_lm ∈ ℝ^{d × V} are often shared (W_lm = E^⊤). Saves V × d parameters (significant: 50k × 4096 = 200M). Justified theoretically by symmetry and empirically by similar or better perplexity. The MiniGPT lab implements this.

5.6 Initialization

You can't init transformer weights from a uniform [-1, 1]. Standard recipe (GPT-style):

• Token embeddings: N(0, 0.02)
• Linear layers: N(0, 0.02)
• Residual-stream output projections (W_O, W_down): N(0, 0.02 / √(2 N)) where N is the number of layers — counteracts variance growth through the residual stream.

A correctly initialized model should have an initial loss of ≈ log(vocab_size) (uniform-distribution prediction). The lab's sanity_init_loss test checks exactly this.

6. Putting it together — the GPT-style architecture

input: token IDs (B, T)
   │
   ▼
[Token Embedding] (V, d) → (B, T, d)
   +
[Positional encoding (or RoPE applied inside attention)]
   │
   ▼
[Block 1] = pre-norm + causal MHA + residual + pre-norm + MLP + residual
[Block 2]
   ...
[Block N]
   │
   ▼
[Final LayerNorm]
   │
   ▼
[LM Head] (d, V) — weight-tied to embedding
   │
   ▼
logits (B, T, V)
   │
   ▼
softmax → probabilities → loss (cross-entropy vs next-token target)

That's a complete decoder-only LLM. Llama, GPT-3, Claude, Gemini — same skeleton, different sizes and tweaks (RoPE flavor, GQA group count, SwiGLU, RMSNorm, attention bias removal).

6.1 Encoder vs decoder vs encoder-decoder

• Encoder (BERT): bidirectional attention; trained with masked LM. Used for classification, embeddings.
• Decoder (GPT, Claude, Llama): causal attention; autoregressive. Used for generation.
• Encoder-Decoder (T5, BART, original transformer): encoder reads input bidirectionally, decoder generates output causally with cross-attention to encoder. Used for translation, summarization (legacy).

In 2024+, decoder-only dominates. Why? Empirically, decoder-only with prompt-based learning matches encoder-decoder quality and is simpler to scale.

7. Lab walkthrough (lab-04-mini-transformer)

7.1 Architecture

The lab builds MiniGPT:

• GPTConfig dataclass — vocab_size, n_layer, n_head, d_model, block_size, dropout.
• CausalSelfAttention — fused QKV projection (one matmul producing all three), reshape to heads, scaled dot-product, mask, softmax, weighted sum, output projection.
• MLP — Linear → GELU → Linear with 4× expansion.
• Block — pre-norm + attn + residual + pre-norm + MLP + residual.
• MiniGPT — embedding + position embedding + N blocks + final LN + tied LM head.

7.2 The two sanity tests

sanity_init_loss(): a freshly-initialized model on random tokens should produce a loss ≈ log(vocab_size). If yours is much higher, your init is broken; if much lower, you have a target leak.

sanity_overfit_one_batch(): take 1 batch, train for ~100 steps; loss should go to near zero. If it doesn't, you have a bug — gradient not flowing, wrong target alignment, frozen parameters. This is the single most useful debugging test.

7.3 Things to read in the solution

• The fused QKV projection: qkv = self.c_attn(x) produces (B, T, 3*d_model) in one matmul; split into Q/K/V. Faster than three separate matmuls (better tensor-core utilization).
• Causal mask is registered as a buffer — not a parameter, but moves with .to(device).
• The view → transpose → matmul → transpose → contiguous → view dance for multi-head — make sure you trace shapes by hand.
• Weight tying: self.lm_head.weight = self.token_emb.weight.

8. References

Required:

• Vaswani et al. (2017), Attention Is All You Need — read it twice.
• Karpathy, Let's build GPT: from scratch, in code, spelled out — the YouTube lecture (~2 hours). Mandatory.
• Karpathy's nanoGPT — read every line.
• Lilian Weng, The Transformer Family — comprehensive blog overview.
• Jay Alammar, The Illustrated Transformer — best diagrams.

Important:

• Radford et al. (2018), Improving Language Understanding by Generative Pre-Training — GPT-1.
• Radford et al. (2019), Language Models are Unsupervised Multitask Learners — GPT-2.
• Brown et al. (2020), Language Models are Few-Shot Learners — GPT-3.
• Touvron et al. (2023), LLaMA: Open and Efficient Foundation Language Models; Llama-2 and Llama-3 papers.
• Devlin et al. (2018), BERT.

Architecture variants:

• Su et al. (2021), RoFormer (RoPE).
• Shazeer (2019), Fast Transformer Decoding: One Write-Head Is All You Need (MQA).
• Ainslie et al. (2023), GQA: Training Generalized Multi-Query Transformer Models.
• Shazeer (2020), GLU Variants Improve Transformer.
• Zhang & Sennrich (2019), Root Mean Square Layer Normalization (RMSNorm).

Theoretical:

• Elhage et al. (2021), A Mathematical Framework for Transformer Circuits (Anthropic) — circuits-level interpretability of attention.
• Olsson et al. (2022), In-Context Learning and Induction Heads (Anthropic).
• Phuong & Hutter (2022), Formal Algorithms for Transformers — pseudocode for everything.

9. Common interview questions on Phase 4 material

1. Implement scaled dot-product attention on a whiteboard.
2. Why divide by √d_k?
3. Why multi-head and not single-head with bigger d?
4. Compare MHA, MQA, GQA. When would you pick each?
5. Compare absolute positional, ALiBi, and RoPE.
6. Walk me through what happens during one forward pass of a 12-layer GPT.
7. Why pre-norm and not post-norm?
8. Why RMSNorm and not LayerNorm?
9. What's weight tying and why does it help?
10. What's the parameter count of a 32-layer, 4096-dim transformer with vocab 50k?
11. Why is the time complexity of attention O(T²) and what can you do about it?
12. Sketch how you'd add a KV cache to your MiniGPT. (Bridges to Phase 9.)
13. Explain SwiGLU vs GELU.
14. What's a residual stream? Why is it useful for analysis?
15. What fails first as you scale a transformer to 70B and 1024 GPUs? (Bridges to Phase 10.)

10. From solid → exceptional

• Implement MiniGPT from a blank file in 30 minutes without consulting solution.py. Time yourself.
• Add RoPE to your MiniGPT (replace the additive position embedding). Compare loss curves.
• Add MQA, then GQA. Measure throughput at long context.
• Replace GELU with SwiGLU. Compare equal-FLOP runs.
• Implement attention three ways (einsum, manual bmm, F.scaled_dot_product_attention). Benchmark each.
• Read Anthropic's A Mathematical Framework for Transformer Circuits and write a one-page summary of "induction heads".
• Pick a real released model (Llama-3 8B, Mistral 7B, Qwen2 7B). Read its config; identify every architectural choice and explain why it was made.
• Do a line-by-line annotation of nanoGPT's model.py in a markdown file. This is the most valuable single hour you can spend.

11. Recommended cadence

Lab 04 — Mini Transformer (Solution Walkthrough)

Phase: 4 — Attention & Transformers | Difficulty: ⭐⭐⭐⭐☆ | Time: 4–6 hours

Concept primer: ../HITCHHIKERS-GUIDE.md §Attention and §Transformer architecture. This is the most important lab in the curriculum — every later phase reuses or extends this code.

Run

pip install -r requirements.txt
python solution.py   # runs init-loss + single-batch overfit sanity checks

0. The mission

Build a decoder-only Transformer from scratch in ~200 lines that:

• Implements scaled dot-product attention with causal masking.
• Uses pre-norm with residual connections.
• Passes the two universal sanity tests: init-loss matches the entropy of the uniform vocab distribution, and the model can overfit a single batch to ~zero loss in 200 steps.

This is the kernel that Phase 5 trains on TinyStories, Phase 6 fine-tunes via LoRA, Phase 9 retro-fits with a KV cache. Get this right and the rest of the curriculum compiles.

1. The math

For each token position $t$:

$$ \mathrm{Attn}(Q, K, V) = \mathrm{softmax}!\left(\frac{Q K^\top}{\sqrt{d_\text{head}}} + M\right) V $$

where $M$ is the causal mask: $M_{ij} = 0$ if $i \ge j$ else $-\infty$. Multi-head attention runs n_head of these in parallel on slices of dim d_head = d_model / n_head, then concatenates.

The full block (pre-norm):

$$ \begin{aligned} x &\leftarrow x + \mathrm{Attn}(\mathrm{LN}(x)) \ x &\leftarrow x + \mathrm{MLP}(\mathrm{LN}(x)) \end{aligned} $$

A model is n_layer blocks stacked plus token + position embeddings at the input and a linear head at the output.

2. GPTConfig

@dataclass
class GPTConfig:
    vocab_size: int = 50257
    n_layer: int = 6
    n_head: int = 8
    d_model: int = 512
    d_ff: int = 2048           # typically 4 * d_model
    block_size: int = 1024     # max context length
    dropout: float = 0.0
    tie_weights: bool = True

• vocab_size = 50257 matches GPT-2 BPE.
• d_ff = 4 * d_model is the universal heuristic from "Attention Is All You Need" — gives the MLP enough capacity to act as the model's "memory" (Geva et al. 2021 showed MLP weights store factual knowledge).
• block_size is the maximum sequence the position-embedding table supports.
• tie_weights=True shares the vocab × d_model matrix between the input embedding and the output head — saves ~50 MB on a small model, ~1 GB on 7B. Quality identical or slightly better.

3. CausalSelfAttention — the centerpiece

3.1 The fused QKV projection

self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False)

One large matmul is faster than three smaller ones (better GPU utilization). Mathematically identical to three separate linears. bias=False is the modern default — biases add parameters without measurable quality benefit at scale.

3.2 The causal mask buffer

self.register_buffer(
    "mask",
    torch.tril(torch.ones(cfg.block_size, cfg.block_size, dtype=torch.bool))
         .view(1, 1, cfg.block_size, cfg.block_size),
    persistent=False,
)

• torch.tril(...) gives a lower-triangular boolean matrix: True on and below the diagonal. Position i can attend to position j iff i ≥ j.
• Shape (1, 1, T, T) so it broadcasts over batch and head dims.
• register_buffer so the mask moves to GPU with .to(device). persistent=False keeps it out of state_dict (deterministically reconstructable).

3.3 The forward — six lines that contain the whole transformer

def forward(self, x):
    B, T, C = x.shape
    qkv = self.qkv(x)                                              # (B, T, 3C)
    q, k, v = qkv.split(C, dim=-1)
    q = q.view(B, T, self.n_head, self.d_head).transpose(1, 2)
    k = k.view(B, T, self.n_head, self.d_head).transpose(1, 2)
    v = v.view(B, T, self.n_head, self.d_head).transpose(1, 2)
    att = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_head)
    att = att.masked_fill(~self.mask[:, :, :T, :T], float("-inf"))
    att = F.softmax(att, dim=-1)
    att = self.attn_drop(att)
    y = att @ v
    y = y.transpose(1, 2).contiguous().view(B, T, C)
    return self.resid_drop(self.proj(y))

Decoded:

1. qkv.split(C, -1) — split the fused projection into Q, K, V each of shape (B, T, C).
2. view + transpose(1, 2) — reshape to (B, n_head, T, d_head). The transpose is the canonical position for the head dim; what cuBLAS expects for batched matmul efficiency.
3. q @ k.transpose(-2, -1) — batched matmul → attention scores (B, n_head, T, T).
4. / math.sqrt(self.d_head) — the most important divisor in deep learning. Without it, scores have variance d_head, push softmax into saturation, gradients vanish.
5. masked_fill(~mask, -inf) — -inf not -1e9 because -1e9 plus a moderately positive score can still produce >1e-30 after softmax, polluting attention.
6. softmax(dim=-1) — normalize across the key dimension. Each row sums to 1.
7. att @ v → (B, n_head, T, d_head) — weighted sum of values.
8. transpose(1, 2).contiguous().view(B, T, C) — un-do the head split. contiguous() is required before view because transpose only changes strides.
9. self.proj(y) — output projection (per-block recombination of head info).

3.4 Why two dropouts?

attn_drop masks attention weights (random tokens become "ignored"); resid_drop masks the output before adding to residual stream. Both at 0 in this skeleton — turn on for fine-tuning small datasets.

4. MLP

class MLP(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.fc = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
        self.proj = nn.Linear(cfg.d_ff, cfg.d_model, bias=False)
        self.drop = nn.Dropout(cfg.dropout)

    def forward(self, x):
        return self.drop(self.proj(F.gelu(self.fc(x))))

GELU = x * Φ(x) (smooth ReLU); empirically better than ReLU for transformers.

Modern variants use SwiGLU (Llama, Qwen): (SiLU(W_g x)) * (W_u x) then W_d. Three matrices instead of two — adds 50% MLP params, gives ~2% perplexity improvement.

5. Block — the pre-norm layout

class Block(nn.Module):
    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x

Pre-norm vs post-norm matters more than any other architecture choice:

• Post-norm (original 2017 paper): x = LN(x + sublayer(x)). Trains poorly without warmup; gradients pass through LN on every residual.
• Pre-norm (GPT-2 onwards): x = x + sublayer(LN(x)). Residual stream is "clean" — gradients flow unimpeded through every layer. Trains stably without warmup at any depth.

Modern alternative: RMSNorm (Llama) — drops mean-subtraction; ~10% faster, identical quality.

6. MiniGPT

self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
self.pos_emb = nn.Embedding(cfg.block_size, cfg.d_model)
self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)])
self.ln_f = nn.LayerNorm(cfg.d_model)
self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
if cfg.tie_weights:
    self.head.weight = self.tok_emb.weight

• Learned absolute position embeddings (GPT-2 style). Modern models use RoPE (rotary, applied in attention itself) — handles longer contexts and extrapolates better.
• Final LayerNorm before head (ln_f) — important for training stability.
• Weight tying by direct assignment. Both tok_emb.weight and head.weight point to the same tensor → only one tensor in the optimizer.

6.1 Init

nn.init.normal_(m.weight, mean=0.0, std=0.02)

std=0.02 is GPT-2's choice. Theoretically 0.02 / sqrt(2 * n_layer) is better for residual-path projections (keeps activation variance constant across layers), but 0.02 everywhere works fine for small models.

6.2 generate

for _ in range(max_new_tokens):
    ctx = idx[:, -self.cfg.block_size:]
    logits, _ = self(ctx)
    logits = logits[:, -1, :] / max(1e-6, temperature)
    if top_k is not None:
        v, _ = torch.topk(logits, top_k)
        logits[logits < v[:, [-1]]] = float("-inf")
    probs = F.softmax(logits, dim=-1)
    next_id = torch.multinomial(probs, 1)
    idx = torch.cat([idx, next_id], dim=1)

• ctx = idx[:, -block_size:] truncates context to the model's max — naive but correct. The KV-cache lab in Phase 9 makes this efficient.
• This is O(T²) per generated token because we re-process the entire context. Phase 9 fixes this with KV cache → O(T).

7. The two sanity tests

These should be the first things you run on any from-scratch transformer.

7.1 Init loss

A randomly-initialized transformer should output approximately uniform logits. Cross-entropy of uniform over V classes is -log(1/V) = log(V). For V=1000, that's 6.91.

If init loss is way off:

• Way higher → bad init scale; logits not centered around 0; softmax saturating.
• Way lower → you accidentally have a constant-output bias somewhere.

7.2 Single-batch overfit

A correctly-wired transformer must memorize a single batch (loss → 0). If it can't:

• Bug in the causal mask (try removing it — does it then overfit? If yes, your mask is upside-down).
• Bug in residual connections (forgetting x = x + ...).
• Bug in positional embeddings (model can't tell positions apart).
• LR way too high (loss explodes) or too low (no progress).

Hitting final_loss < 0.5 in 200 steps confirms forward + backward + optimizer all wire correctly.

8. Expected output

params = 526,464
[init-loss]  got=6.9085  expected≈6.9078  ok=True
[overfit]   step  200  loss=0.0264  ok=True

If init-loss matches log(vocab_size) to two decimals and single-batch overfit drives loss < 0.5, your transformer is wired correctly.

9. Common pitfalls

1. Forgetting / math.sqrt(d_head) — softmax saturates → gradients vanish.
2. Mask shape mismatch when T < block_size → must slice with [:, :, :T, :T].
3. Forgetting contiguous() before view after transpose → runtime error.
4. Missing residuals — x = self.attn(self.ln1(x)) (forgot the x +) — model trains but quality is terrible. Sanity tests catch this.
5. Wrong mask direction — triu instead of tril → tokens attend only to the future. Loss might still go down but generation produces garbage.
6. Tied weights only on init — must assign self.head.weight = self.tok_emb.weight not copy values.
7. F.cross_entropy expects raw logits, not log-softmax. Don't double-softmax.

10. Stretch exercises

• Implement RoPE (rotary positional embeddings). Apply rotation to Q, K inside attention. Drop the pos_emb table.
• Implement RMSNorm. Replace LayerNorm. ~10 lines, ~10% faster.
• Implement SwiGLU MLP.
• Implement GQA (grouped-query attention). Set n_kv_head < n_head; broadcast K, V across query heads. Halves the KV cache.
• Use torch.nn.functional.scaled_dot_product_attention to dispatch FlashAttention. Compare wall-clock — should be 2-3× faster at long contexts.
• Profile with torch.profiler: where is time spent? (~60% matmuls, ~20% softmax, ~10% everything else.)
• Reproduce the GPT-2 124M architecture exactly: 12 layers, 12 heads, d=768.

11. Connecting to later phases

You'll come back to this file 5+ times across the curriculum. Internalize it.

12. What this lab proves about you

You can implement causal multi-head attention from raw matmuls, articulate every design decision, verify correctness via init-loss + overfit, and modify it for new architectures (RoPE, SwiGLU, GQA) without breaking it. The bar for a Phase-4 milestone — and the single most-asked area of LLM interviews.

Phase 5 — Training Small LLMs

Difficulty: ⭐⭐⭐⭐☆ | Estimated Time: 2.5 weeks Roles supported: Research Engineer Pretraining, Foundation Model Engineer.

Why This Phase Exists

The Anthropic / OpenAI / DeepMind pretraining job descriptions all say variations of: "experience training transformer models end-to-end". Reading about it is not the same as having stared at a loss curve at 3 AM, debugged a NaN, and explained to yourself why your gradients exploded. This phase produces that experience cheaply.

By the end you will have trained a real (small) language model from scratch with a tokenizer you wrote, on data you cleaned, with a training loop you understand line-by-line.

Concepts

• Byte-Pair Encoding (BPE) algorithm + GPT-2 / Llama tokenizer details
• Tokenizer training: word frequencies → merges → vocab
• nanoGPT-style architecture (Andrej Karpathy)
• Dataset packing & sequence packing
• Optimizers: AdamW, Lion, Sophia (overview)
• Learning-rate schedules: warmup + cosine decay
• Mixed precision: BF16 vs FP16, loss scaling
• Gradient accumulation (simulating larger batch sizes)
• Gradient clipping
• Checkpointing strategy (save best, save last, save every N)
• Sampling: greedy, multinomial, temperature, top-k, top-p (nucleus), beam, contrastive
• Chinchilla scaling laws (intuition)
• W&B / Tensorboard logging hygiene

Labs

Lab 01 — BPE Tokenizer From Scratch (Matching GPT-2)

Lab 02 — nanoGPT From Scratch on TinyStories

Lab 03 — Training Loop Mechanics (Mixed Precision, Grad Accumulation, Checkpointing)

Lab 04 — Sampling Strategies & Generation

Deliverables Checklist

• BPE tokenizer matching tiktoken on test data
• nanoGPT trained on TinyStories with W&B logs and generated samples
• Resumable training loop with grad accumulation + clipping
• Sampling library + comparison report

Interview Relevance

• "Walk me through your training loop"
• "How would you debug a NaN loss?"
• "Why BF16 over FP16?"
• "Explain top-p sampling"
• "How would you scale this to 1B parameters?" (sets up Phase 10)

Warmup Guide — Training Small LLMs

Zero-to-expert primer for Phase 05: everything between "I have a transformer" and "it trained well" — optimization (AdamW, schedules), stability, mixed precision, scaling laws, and the experimental discipline of the nanoGPT lab.

Table of Contents

• Chapter 1: The Training Loop, Honestly Stated
• Chapter 2: Adam and AdamW — Why They Won
• Chapter 3: Learning-Rate Schedules — Warmup and Cosine
• Chapter 4: Batch Size, Gradient Accumulation, and Tokens-Per-Step
• Chapter 5: Mixed Precision — BF16 and the Loss-Scale Question
• Chapter 6: Initialization and Stability
• Chapter 7: Scaling Laws — How Big, How Much Data
• Chapter 8: The Experimental Method
• Lab Walkthrough Guidance
• Success Criteria
• Interview Q&A
• References

Chapter 1: The Training Loop, Honestly Stated

The loop is five lines; every line hides a chapter:

for batch in data:                  # Ch. 4: what's a batch, really
    logits = model(x)               # Ch. 5: in what precision
    loss = cross_entropy(logits, y) # Phase 03 Ch. 2: the objective
    loss.backward()                 # autograd (model-accuracy Phase 01)
    clip_grad_norm_(params, 1.0)    # Ch. 6: stability
    optimizer.step(); scheduler.step(); optimizer.zero_grad()   # Ch. 2–3

The phase's real subject is the gap between this loop running and this loop working: the silent failure modes (wrong LR, bad init, precision bugs) that produce a loss curve that goes down — just to a worse place than it should. The defense is the calibration mindset: know what loss value, what curve shape, and what sample quality to expect at each stage, so deviation is informative (Chapter 8).

Chapter 2: Adam and AdamW — Why They Won

SGD's problem for transformers: one global learning rate across parameters whose gradient scales differ by orders of magnitude (embedding rows for rare tokens vs LayerNorm gains). Adam maintains per-parameter statistics — first moment (EMA of gradients, $m$) and second moment (EMA of squared gradients, $v$) — and updates:

$$\theta \mathrel{-}= \eta \cdot \frac{\hat{m}}{\sqrt{\hat{v}} + \epsilon}$$

The $\sqrt{\hat{v}}$ denominator is an automatic per-parameter LR: rarely-updated parameters (small $v$) take big steps; noisy ones get damped. Bias correction ($\hat{m} = m/(1-\beta_1^t)$) fixes the zero-initialized EMAs' early-step underestimate — without it the first steps are silently tiny.

AdamW's fix to Adam: in Adam, L2 "weight decay" added to the gradient gets divided by $\sqrt{\hat{v}}$ too — parameters with large gradient history get less regularization, which is backwards. AdamW applies decay decoupled, directly to the weights ($\theta \mathrel{-}= \eta \lambda \theta$), restoring decay's meaning. All modern LLMs: AdamW, $\beta_1{=}0.9$, $\beta_2{=}0.95$ (lower than the 0.999 default — faster-moving curvature estimates suit LLM gradient noise), decay ~0.1 applied to weight matrices but not to biases, LayerNorm parameters, or (usually) embeddings.

The cost nobody mentions until it hurts: two FP32 states per parameter — optimizer memory = 2× model size in FP32, i.e., more than the model itself. This number drives everything in distributed training (Phase 10: ZeRO exists for this) and fine-tuning (Phase 06: LoRA exists substantially for this).

Chapter 3: Learning-Rate Schedules — Warmup and Cosine

The standard LLM schedule is linear warmup → cosine decay:

• Warmup (a few hundred–few thousand steps, ramping 0 → peak): Adam's $v$ estimate is garbage for the first steps (built from a handful of noisy gradients), and transformer training is most fragile at init — a full-size step at step 1 can spike the loss into a divergence it never recovers from. Warmup lets the optimizer's statistics and the loss landscape's early descent stabilize first.
• Cosine decay to ~10% of peak: large steps explore early; small steps settle into a minimum late. Empirically robust; the exact shape (cosine vs linear) matters less than (a) peak LR (the single most important hyperparameter — too high diverges or plateaus noisily; too low underfits the budget) and (b) decaying to a low floor by the end of the token budget.
• The lab's required experiment: an LR range sweep (3e-3 / 6e-4 / 1e-4 on the same config) — the canonical "too hot / right / too cold" triptych of curves is something you should have personally produced once, because you'll be diagnosing it from others' curves forever.

Chapter 4: Batch Size, Gradient Accumulation, and Tokens-Per-Step

• The meaningful unit is tokens per optimizer step = micro_batch × seq_len × accumulation × (data-parallel ranks). nanoGPT-class runs use ~0.5M tokens/step; report and reason in this unit, not "batch size."
• Gradient accumulation: run N micro-batch forward/backwards, summing grads, then one optimizer step — mathematically identical to the big batch (within FP non-associativity), trading wall-clock for memory. The classic bug: forgetting to scale the loss by 1/N (or zero_grad placement), silently multiplying the effective LR.
• Critical batch size intuition (McCandlish et al.): below it, bigger batches are nearly free speedups (gradient noise dominates); above it, diminishing returns — you're averaging already-clean gradients. Small models on small data have small critical batch sizes — another reason the lab's modest config is correct, not just convenient.
• Batch size and LR interact (bigger batch → cleaner gradient → supports higher LR); change them together or not at all (Chapter 8's one-variable rule has this one sanctioned exception, with the linear-scaling heuristic as the starting point).

Chapter 5: Mixed Precision — BF16 and the Loss-Scale Question

Train in 16-bit where it's safe, 32-bit where it isn't:

• What stays FP32: master weights (inside AdamW), optimizer states, and the big reductions (loss, norms). What runs 16-bit: matmuls, activations — the bulk of compute and memory.
• FP16 vs BF16 (the exponent/mantissa trade from the model-accuracy track's Phase 03 Ch. 1): FP16's 5-bit exponent overflows at 65504 and underflows small gradients to zero — hence loss scaling (multiply loss by ~2¹⁵ so gradients shift into representable range, unscale before the step, skip steps on inf/nan). BF16 has FP32's 8-bit exponent — no overflow/underflow drama, no loss scaling, at the cost of ~3 decimal digits of mantissa, which training tolerates. On any hardware that has it (A100+), BF16 is the answer and the lab uses torch.autocast(dtype=bfloat16).
• The debugging signature to memorize: FP16 run with occasional inf → loss-scaler skips (normal if rare); BF16 run with loss spikes → it's not precision overflow, look at data or LR (this elimination step saves real days).

Chapter 6: Initialization and Stability

Why the lab's init code looks the way it does:

• Linear layers: normal(0, 0.02); embeddings likewise — small enough that pre-norm residual streams start near-identity.
• The GPT-2 residual-projection trick: scale the output projections of attention and FFN by $1/\sqrt{2N}$ (N = layer count) — each block adds to the residual stream, and without the scaling the stream's variance grows linearly with depth; the scaling keeps the sum's variance constant. This single line is the difference between 12 layers training smoothly and not.
• Stability kit, in escalation order: gradient clipping at norm 1.0 (always on; watch the pre-clip norm — its trend is an early-warning instrument), warmup (Ch. 3), then if spikes persist: check data (a pathological document), lower peak LR, raise $\epsilon$, suspect precision last (Ch. 5's elimination).
• The overfit-one-batch test before any long run: a healthy model+loop drives one small batch to ~zero loss in a few hundred steps. Failure means a wiring bug (shifted targets, mask, LR) — never start a multi-hour run without this two-minute test.

Chapter 7: Scaling Laws — How Big, How Much Data

The empirical regularities that turn "how big a model?" from taste into arithmetic:

• Kaplan et al. (2020): loss falls as a power law in parameters N, data D, and compute C, over many orders of magnitude — smooth, predictable returns.
• Chinchilla (2022): for a fixed compute budget $C \approx 6ND$, the optimum is roughly D ≈ 20 N — tokens ≈ 20× parameters. GPT-3 (175B params, 300B tokens) was badly under-trained by this rule; Chinchilla (70B, 1.4T) beat it at the same compute.
• The inference-cost asterisk (the practitioner's correction): Chinchilla optimizes training compute only. If a model will serve billions of requests, over-training a smaller model far past 20:1 (LLaMA-class models: 100–2000 tokens/param) buys lower inference cost forever at modest training premium — which is the actual logic behind every production "small" model you'll serve in Phase 09.
• For the lab's scale: a ~10M-param model wants ~200M+ tokens by the 20:1 rule — Shakespeare (~1M) is hopelessly small, which is why the lab also runs OpenWebText- class data and why your Shakespeare model memorizes (watch val loss diverge from train — the overfitting lesson live).

Chapter 8: The Experimental Method

The discipline that separates training-as-science from GPU-warming, the same ledger ethics as the model-accuracy track's capstone (Phase 11 Ch. 3):

1. One variable per run; config serialized with every checkpoint; seeds fixed (and the run-to-run σ at fixed seed measured once, so you know what differences are real).
2. Watch validation, not train, loss — and watch samples: fixed prompts generated every N steps; the qualitative arc is a debugging instrument no scalar replaces.
3. Baseline before improvement: the stock config trained to completion is the reference every change is measured against. "It felt better" without the baseline delta is how teams lose months.
4. Log the boring numbers: tokens/sec (throughput regressions are bugs), grad norm, LR actually applied. When a run dies at 3am, the logs are the post-mortem.

Lab Walkthrough Guidance

Lab 02 — nanoGPT:

1. Run the overfit-one-batch test on your Phase 04 model wired into this loop. Only then start real runs.
2. Shakespeare char-level first (minutes/run): produce the LR triptych (Ch. 3), the batch-size/accumulation equivalence check (same tokens/step two ways → same curve), and the warmup ablation (watch the no-warmup run spike).
3. Add mixed precision; verify the loss curve matches FP32 within noise and measure the throughput gain — both numbers go in your notes.
4. Scale to the larger dataset/config; apply Chapter 7's arithmetic to predict where val loss should land relative to the small run before looking. Calibration is the skill; the prediction-vs-actual delta is the lesson.
5. Keep the ledger (Ch. 8) from run #1 — the lab's deliverable is the table of runs, not the final checkpoint.

Success Criteria

You are ready for Phase 06 when you can, from memory:

1. Write Adam's update and explain $\sqrt{\hat{v}}$, bias correction, and AdamW's decoupling — plus the 2×-FP32 memory fact and what it later justifies (ZeRO, LoRA).
2. Justify warmup and cosine decay; sketch the too-hot/right/too-cold triptych.
3. Compute tokens-per-step for any config and state the accumulation-scaling bug.
4. Choose BF16 vs FP16 with the exponent argument; explain loss scaling and when it's unnecessary.
5. Explain the $1/\sqrt{2N}$ residual init and the overfit-one-batch test's purpose.
6. State Chinchilla's 20:1, its derivation context, and the inference-cost correction that explains over-trained small models.

Interview Q&A

Q: Your loss spikes at step 40K and recovers but lands at a worse plateau. Walk through it. First the data: spikes are most often a pathological batch (corrupt/duplicated document) — find what was sampled at that step (this is why ledgers log data order/ seed). Then optimizer state: a spike corrupts Adam's $v$ (huge squared grads), damping subsequent steps — recovery-to-worse is consistent with poisoned state; mitigations are tighter clipping, lower $\beta_2$, or restart from pre-spike checkpoint with the data fixed. Precision last: BF16 makes overflow unlikely; FP16 scaler logs would show it. The answer's structure — data, optimizer state, precision, in that order — is what's being graded.

Q: Why does everyone use AdamW for transformers when SGD trains ResNets fine? Transformer gradient scales are wildly heterogeneous across parameter types (embedding rows, LayerNorm gains, attention vs FFN matrices) and the loss landscape has sharper curvature anisotropy; per-parameter adaptive steps absorb this where one global LR can't (SGD on transformers needs delicate per-layer LR surgery to compete). Plus decoupled decay gives clean regularization semantics. Cost: 2× FP32 optimizer state — which is a real systems consequence, not a footnote.

Q: With a fixed training budget, would you train a 7B for 1 epoch or a 1.4B for 5× the tokens — and what changes your answer? Chinchilla says match D ≈ 20N for best training-compute loss — compute both options' N:D against that. But the deployment profile dominates: if the model serves at scale, the smaller over-trained model wins on lifetime cost (inference forever beats a small pretraining delta); if it's a research artifact or will be distilled anyway, the bigger model's better loss matters. Saying "Chinchilla, then correct for inference" is the complete answer.

References

• Kingma & Ba, Adam (2014) — arXiv:1412.6980; Loshchilov & Hutter, Decoupled Weight Decay (AdamW) (2017) — arXiv:1711.05101
• Kaplan et al., Scaling Laws for Neural Language Models (2020) — arXiv:2001.08361
• Hoffmann et al., Training Compute-Optimal LLMs (Chinchilla) (2022) — arXiv:2203.15556
• McCandlish et al., An Empirical Model of Large-Batch Training (2018) — arXiv:1812.06162 — critical batch size
• Micikevicius et al., Mixed Precision Training (2017) — arXiv:1710.03740 — loss scaling's origin
• Karpathy, nanoGPT — read train.py line by line; it is this chapter as code
• Karpathy, A Recipe for Training Neural Networks — karpathy.github.io/2019/04/25/recipe — Chapter 8, canonically

🛸 Hitchhiker's Guide — Phase 5: Training Small LLMs

Read this if: You can build a MiniGPT but you've never trained one to convergence on real data, or you don't yet have a feel for "this loss curve looks healthy", "this is the LR I should use for a 124M model", "this is what 50 GPU-hours of pretraining buys you".

0. The 30-second mental model

Pretraining = run AdamW on a MiniGPT-style architecture for billions of next-token prediction steps over a giant deduplicated text corpus, using mixed precision, with a warmup-then-decay learning rate schedule, gradient accumulation to reach a large effective batch, and frequent checkpointing. Watch the loss go down. Sample. Cry tears of joy. That's pretraining.

By the end of Phase 5 you should:

• Train nanoGPT on TinyStories and produce coherent toy text.
• Understand and tune: batch size, learning rate, warmup, weight decay, gradient clipping, gradient accumulation, mixed precision (bf16/fp16/fp8).
• Read and apply scaling laws (Kaplan, Chinchilla, MoE corrections).
• Diagnose loss spikes, NaN, slow convergence, and undertraining.
• Know the data preparation pipeline: tokenize → shard → memory-map → uint16 .bin.
• Be ready to discuss real pretraining at the 1B–70B scale (Phase 10 will go deeper).

1. The pretraining objective

Same as Phase 3: minimize cross-entropy of next-token prediction. For a sequence of token IDs x_0, x_1, …, x_{T-1}, the model produces logits (T, V) and the loss is:

loss = F.cross_entropy(logits[:-1].reshape(-1, V), x[1:].reshape(-1))

Note the shift by 1: position t predicts position t+1. A common bug is forgetting this shift; the model then learns identity (loss → 0 instantly). The lab's sanity_overfit_one_batch catches it.

2. Optimizers — what's actually happening

2.1 SGD — the conceptual baseline

θ ← θ - η · ∇_θ L. Simple, but for transformers it's terrible without momentum and tuning.

2.2 Momentum / Nesterov

Track a running average of gradients; update with that. Smooths out noisy gradients.

2.3 Adam (Kingma & Ba, 2014)

For each parameter, maintain two moving averages:

• m_t = β₁ m_{t-1} + (1 - β₁) g_t — first moment (mean of gradient).
• v_t = β₂ v_{t-1} + (1 - β₂) g_t² — second moment (uncentered variance).

Bias-correct (m̂ = m / (1 - β₁ᵗ), etc.), then update:

$$ θ ← θ - η · \hat{m} / (\sqrt{\hat{v}} + ε) $$

Intuition: Adam is per-parameter learning-rate adaptation. Parameters with consistently large gradients get smaller effective updates; sparse-gradient parameters get larger ones.

2.4 AdamW (Loshchilov & Hutter, 2019)

Vanilla Adam with L2 regularization couples decay with the adaptive lr — wrong. AdamW decouples: θ ← θ - η (m̂/√v̂ + ε + λ θ). Same intuition, decay applied directly to weights. Always use AdamW, never Adam, for transformers.

Hyperparameters (sane defaults for transformers):

• β = (0.9, 0.95) (note: β₂ = 0.95, not 0.999 — empirically better for LLMs)
• weight_decay = 0.1
• eps = 1e-8

2.5 Lion, Sophia, etc.

Recent alternatives. Lion (Chen et al. 2023) uses sign-of-momentum updates; smaller memory footprint. Sophia (Liu et al. 2023) uses Hessian estimates. Neither has displaced AdamW universally yet.

2.6 Memory cost

AdamW stores 2 floats per parameter (m, v). At fp32 that's 8 × params bytes. A 7B model = 56 GB just for optimizer states — more than the weights themselves. This is why we shard them in FSDP (Phase 10).

3. Learning rate schedules

The single biggest training-stability lever after batch size.

3.1 Warmup → Cosine decay (the workhorse)

• Warmup (first 1–2% of steps): linearly increase from 0 to peak_lr. Without it, early steps with random weights produce huge gradients that destabilize training.
• Cosine decay (remaining steps): lr = min_lr + 0.5 (peak_lr - min_lr) (1 + cos(π t/T_max)). Smooth descent to ~10% of peak.

3.2 Warmup-Stable-Decay (WSD)

• Warmup → constant peak_lr for ~80% of training → fast cosine decay over last 10–20%.
• Lets you take any intermediate checkpoint and "finalize" it with a short decay run. No need to commit to a token budget upfront.
• Used in MiniCPM, DeepSeek and increasingly elsewhere.

3.3 What peak_lr to pick?

Empirical rule: peak_lr ≈ 6e-4 × (124M / params)^0.5 for GPT-style. For nanoGPT (124M): 6e-4. For 1B: ~2e-4. For 7B: ~1e-4. For 70B: ~3e-5.

You can also do a lr range test (Smith 2017): train for a few hundred steps with linearly-increasing lr; pick the lr where loss starts diverging, divide by 4–10. Lab 02 uses fixed sane defaults rather than tuning.

4. Batch size and gradient accumulation

4.1 Effective batch and tokens-per-step

Modern LLMs train at 0.5M–4M tokens per step (effective batch). You rarely fit that in one micro-batch on one GPU, so:

effective_batch_size = micro_batch × n_gpus × grad_accum_steps

grad_accum_steps accumulates gradients across forward/backward passes before the optimizer step:

opt.zero_grad()
for k in range(grad_accum_steps):
    micro = next_batch()
    loss = model(micro) / grad_accum_steps   # divide so loss is averaged
    loss.backward()                          # accumulates into .grad
opt.step()

This is mathematically equivalent to a single bigger batch (assuming no batch-norm — which transformers don't use).

4.2 The batch-size–LR coupling

When you increase the batch by k, you can usually increase the LR by k (linear scaling) or √k (sqrt scaling) without instability. For transformers the sqrt scaling is more conservative.

4.3 Critical batch size

McCandlish et al. (2018) showed each task has a critical batch size beyond which throughput improvements diminish. For LLMs the critical batch grows with model size — so you can use larger batches as you scale up.

5. Mixed precision

Goal: use lower-precision math to get more throughput per GPU and fit bigger models.

5.1 The four datatypes

BF16 is the default for pretraining in 2024+. Same exponent range as FP32 means you don't need loss scaling. Mantissa precision is enough for most ops if you keep certain reductions in FP32.

5.2 The recipe (PyTorch AMP)

scaler = torch.cuda.amp.GradScaler()              # FP16 path
with torch.amp.autocast("cuda", dtype=torch.bfloat16):
    logits = model(x)
    loss = F.cross_entropy(logits, y)
loss.backward()
opt.step()
opt.zero_grad()

For BF16 you don't need GradScaler. For FP16 you do, because FP16's tiny range (~6e-5 minimum normal) underflows easily; the scaler multiplies the loss by a large number to keep gradients in range, then unscales before the optimizer step.

5.3 FP8 on H100

Hopper TensorCores natively run FP8 matmul at 2× the rate of BF16. Used with per-tensor delayed scaling (or per-block scaling for finer granularity). Library: NVIDIA's transformer_engine. Phase 10 covers it more deeply.

6. Scaling laws — the most important paper of the era

6.1 Kaplan et al. (2020) — Scaling Laws for Neural Language Models

Loss as a function of compute, parameters, and data follows a clean power law:

$$ L(N) \approx (N_c / N)^{α_N} $$

(Same for D and C.) The bombshell was: at fixed compute C ≈ 6 N D, the optimal allocation favored bigger models. GPT-3 was sized accordingly: 175B params, ~300B tokens.

6.2 Chinchilla (Hoffmann et al., 2022) — Training Compute-Optimal Large Language Models

DeepMind redid the analysis carefully and found N and D should scale equally at fixed compute — i.e., ~20 tokens per parameter is optimal. Implication: GPT-3 was massively undertrained. The 70B Chinchilla model trained on 1.4T tokens beat the 280B Gopher trained on 300B tokens.

This single finding reshaped the field. Llama models train at 200×+ tokens per param (LLama-3 8B trained on 15T tokens — far beyond Chinchilla optimal but yields better inference economics).

6.3 The compute equation

For a dense transformer:

$$ C ≈ 6 N D \text{ FLOPs} $$

where N = non-embedding parameters, D = training tokens. The 6 comes from 2 (multiply-add) × 3 (forward + backward + optimizer-related). Useful for back-of-envelope cost estimates.

6.4 References

• Kaplan et al. (2020), Scaling Laws for Neural Language Models.
• Hoffmann et al. (2022), Training Compute-Optimal Large Language Models (Chinchilla).
• Henighan et al. (2020), Scaling Laws for Autoregressive Generative Modeling (multimodal).
• Hoffmann's Chinchilla follow-ups. Replications: Pearce et al. (2024).

7. Data preparation for pretraining

Phase 10 covers this in depth. Quick preview:

1. Source: CommonCrawl (web), GitHub (code), arXiv (science), books, Wikipedia.
2. Filter: language ID, quality classifier, Gopher rules, perplexity filter.
3. Dedup: URL → exact → MinHash near-dup.
4. PII scrub: regex + Presidio.
5. Tokenize: with your tokenizer; output uint16 (vocab ≤ 65535) or uint32 .bin shards.
6. Mix and shuffle: weighted source mixing, deterministic shuffle.

For Lab 02 (nanoGPT on TinyStories): step 5 only. The dataset is small and pre-cleaned.

8. The lab walkthrough (lab-02-nano-gpt)

8.1 What you'll build

A working prepare → train → sample CLI that:

1. Prepare: downloads TinyStories (Eldan & Li, 2023; ~500MB of GPT-3.5-generated 4-year-old-level stories with vocabulary ~1500 words), tokenizes with GPT-2's tokenizer, dumps to train.bin / val.bin (uint16 memory-mapped arrays).
2. Train: imports MiniGPT and GPTConfig from Phase 4; trains for max_iters (default 5000) with bf16 AMP, gradient accumulation, cosine schedule.
3. Sample: loads checkpoint, runs autoregressive generation with top-k + temperature.

8.2 What "healthy" looks like

• Initial loss ≈ log(50257) ≈ 10.8.
• After 100 steps: ~6 (model has learned unigram distribution).
• After 1000 steps: ~3 (basic word patterns).
• After 5000 steps on TinyStories with a 6-layer 384-dim model: ~1.5–2.0 (coherent simple stories).

8.3 Why memory-mapped uint16 .bin?

A 5GB tokenized corpus loaded into RAM = 5GB. As np.memmap, it costs ~0 — only the active page is in memory. Cheap random access for batch sampling. uint16 (2 bytes/token) halves disk vs uint32.

8.4 Things to read carefully

• get_batch() — random offsets within the .bin, slice block_size + 1 tokens, split into (x, y) with the +1 shift.
• The training loop's grad_accum arithmetic.
• The cosine schedule with warmup function.
• torch.amp.autocast placement (only the forward; backward and optim step run in original precision).
• The @torch.no_grad() eval block — saves memory.

8.5 Cost expectation

On a single A100 40GB, the default config (~10M params, 5k steps, batch 64 × 256 tokens) trains in ~15–30 minutes. On consumer GPU (4090): ~30–60 minutes. Generates believable toddler stories.

9. Diagnosing training problems

10. References

Core:

• Karpathy's nanoGPT repo and video lecture.
• Kaplan et al. (2020) and Hoffmann et al. (2022) — scaling laws.
• Loshchilov & Hutter (2019), Decoupled Weight Decay Regularization (AdamW).
• Smith (2017), Cyclical Learning Rates for Training Neural Networks — LR range test.
• Eldan & Li (2023), TinyStories: How Small Can Language Models Be and Still Speak Coherent English?

Production-scale recipes (read once you finish the lab):

• OPT (Zhang et al. 2022) — has a release log of every restart and bug for a 175B model. Eye-opening.
• Llama-3 tech report.
• DeepSeek-V2 and DeepSeek-V3 tech reports.
• Qwen-2 tech report.
• Pythia (Biderman et al. 2023) — releases all checkpoints; great for studying training dynamics.

11. Common interview questions on Phase 5 material

1. Why AdamW and not Adam?
2. Why do we need LR warmup?
3. What's the Chinchilla finding in one sentence? Why did it overturn Kaplan?
4. How do you decide effective batch size?
5. Walk me through gradient accumulation.
6. Why BF16 over FP16 for pretraining?
7. What does the AdamW optimizer cost in memory per parameter?
8. Loss is NaN at step 200. How do you debug?
9. You have $50k of compute. What size model and how many tokens?
10. What's WSD and why is it interesting?
11. Sketch the training loop on a whiteboard.
12. How would you know if your model is undertrained?

12. From solid → exceptional

• Train nanoGPT on TinyStories. Then train on all of Wikipedia (~30GB tokenized). Document loss curves and final perplexity.
• Implement gradient checkpointing by hand (re-compute forward activations during backward instead of storing them). Measure the memory ↔ throughput tradeoff.
• Implement torch.compile wrapping; benchmark step time before/after.
• Add bf16 mixed precision with FP32 reductions explicitly (not via autocast); confirm equivalent loss.
• Read the OPT log book end-to-end; pick three failures and write what you would have done differently.
• Implement a scaling-law ablation: train models at sizes 6M, 12M, 25M, 50M for matched compute budgets; fit the power law; predict the loss at 100M; train and verify.
• Write a one-page cost model: $/M-tokens-trained for various model sizes on H100 spot.

13. Recommended cadence

Lab 02 — nanoGPT on TinyStories (Solution Walkthrough)

Phase: 5 — Training Small LLMs | Difficulty: ⭐⭐⭐⭐☆ | Time: 4–8 hours (incl. training)

Reuses the model from ../../phase-04-attention-transformers/lab-04-mini-transformer/solution.py. Concept primer: ../HITCHHIKERS-GUIDE.md §Pretraining mechanics.

Run

pip install -r requirements.txt
python solution.py --prepare        # tokenizes → ./data/train.bin, val.bin
python solution.py --train --steps 2000
python solution.py --sample --prompt "Once upon a time"

0. The mission

Go from raw text to a generating model in one script. End-to-end:

1. --prepare — download TinyStories, tokenize with tiktoken GPT-2 BPE, write packed uint16 shards.
2. --train — mixed-precision (BF16) training with gradient accumulation, cosine LR, AdamW, periodic eval + checkpoints.
3. --sample — load checkpoint, generate text from a prompt.

Default config trains a ~10M-param model in ~30 minutes on a T4 (Colab free) and produces grammatical English. Scale up to d=512, 8 layers and you have a real (if tiny) language model.

1. --prepare — the data pipeline

import tiktoken
enc = tiktoken.get_encoding("gpt2")
ids = enc.encode_ordinary(text)        # ~100M tokens for TinyStories
ids.append(enc.eot_token)              # "<|endoftext|>" id 50256 between docs
arr = np.array(ids, dtype=np.uint16)   # 50257 < 65536 → fits in uint16
arr.tofile(out_dir / "train.bin")

• encode_ordinary strips special tokens — we don't want stray <|endoftext|> tokens accidentally appearing inside docs.
• uint16 halves disk footprint vs int32. Required because GPT-2 vocab is 50257 < 65536.
• EOT between docs so the model learns where stories end. During training we randomly slice across boundaries — the EOT token is the only signal.
• We write train.bin and val.bin (90/10 split). Loading is np.memmap(...) so a 100 MB file uses zero RAM.

def get_batch(split, block_size, batch_size):
    data = np.memmap(out_dir / f"{split}.bin", dtype=np.uint16, mode="r")
    ix = np.random.randint(0, len(data) - block_size - 1, (batch_size,))
    x = np.stack([data[i:i+block_size].astype(np.int64) for i in ix])
    y = np.stack([data[i+1:i+1+block_size].astype(np.int64) for i in ix])
    return torch.from_numpy(x).to(device), torch.from_numpy(y).to(device)

Random-offset slicing is the standard trick: every batch is a fresh random crop. No shuffling overhead. The model sees ~steps * batch * block_size tokens total; for 2000 steps × batch 64 × block 256 ≈ 33M tokens (1/3 epoch over TinyStories).

2. --train — the training loop

2.1 Optimizer setup

def configure_optimizer(model, lr, weight_decay):
    decay, no_decay = [], []
    for n, p in model.named_parameters():
        if p.dim() >= 2:
            decay.append(p)            # weight matrices, embeddings
        else:
            no_decay.append(p)         # biases, LayerNorm gain/beta
    groups = [
        {"params": decay, "weight_decay": weight_decay},
        {"params": no_decay, "weight_decay": 0.0},
    ]
    return torch.optim.AdamW(groups, lr=lr, betas=(0.9, 0.95), fused=True)

Three non-obvious choices:

1. No weight decay on 1D parameters. Decaying LayerNorm gains pulls them toward 0, distorting the normalization. Decaying biases is similarly harmful and pointless. Standard since GPT-2.
2. betas=(0.9, 0.95) — Llama/GPT-3's choice. Default is (0.9, 0.999). The lower β₂ makes the second-moment estimate more responsive to recent gradients — crucial when LR is high and gradient stats change quickly.
3. fused=True — PyTorch 2.x fused AdamW kernel. ~30% faster on GPU. Only works on CUDA.

2.2 Cosine LR schedule with warmup

def get_lr(step, warmup, max_steps, lr_max, lr_min):
    if step < warmup:
        return lr_max * step / warmup
    if step > max_steps:
        return lr_min
    decay_ratio = (step - warmup) / (max_steps - warmup)
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return lr_min + coeff * (lr_max - lr_min)

• Warmup — 100–2000 steps. Without it, the first big update from random init explodes activations; AdamW's second-moment estimate is also unreliable until enough gradients accumulate. Skipping warmup is the #1 cause of NaN losses.
• Cosine decay to lr_min = 0.1 * lr_max. Empirically beats linear, exponential, or step decay.
• LR is set per-step via for g in opt.param_groups: g["lr"] = lr.

2.3 Mixed precision + gradient accumulation

scaler = torch.cuda.amp.GradScaler(enabled=(dtype == torch.float16))
ctx = torch.amp.autocast("cuda", dtype=dtype)

for micro in range(grad_accum_steps):
    x, y = get_batch("train", block_size, batch_size)
    with ctx:
        _, loss = model(x, y)
        loss = loss / grad_accum_steps
    scaler.scale(loss).backward()
scaler.unscale_(opt)
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
scaler.step(opt)
scaler.update()
opt.zero_grad(set_to_none=True)

• BF16 preferred over FP16 when your GPU supports it (Ampere+). Same dynamic range as FP32; no GradScaler needed (enabled=False).
• Grad accumulation simulates a larger batch: with grad_accum=8, the effective batch is batch_size * 8 * world_size. Loss divided by grad_accum_steps so the gradient magnitude matches the full batch.
• clip_grad_norm_(..., 1.0) prevents occasional spikes from corrupting the running optimizer state.
• zero_grad(set_to_none=True) is faster than zero_grad() (avoids touching every param).

2.4 Periodic eval + checkpoint

if step % eval_interval == 0:
    model.eval()
    with torch.no_grad():
        losses = []
        for _ in range(eval_iters):
            xv, yv = get_batch("val", block_size, batch_size)
            with ctx:
                _, loss = model(xv, yv)
            losses.append(loss.item())
        val_loss = sum(losses) / len(losses)
    model.train()
    if val_loss < best_val:
        best_val = val_loss
        torch.save({"model": model.state_dict(), "opt": opt.state_dict(),
                    "step": step, "val_loss": val_loss}, ckpt_path)

Saving optimizer state allows resume. Saving only the best-val checkpoint avoids disk bloat. For real runs, also save a last checkpoint every N steps for crash recovery.

3. --sample — generation

Load checkpoint, tokenize prompt with tiktoken, call model.generate(...) from Phase 4. Use top_k=200, temperature=0.8 for stories (slightly conservative).

4. Expected output

Default config (d=128, 6 layers, 8 heads, block=256), 2000 steps, T4 GPU:

step    0  loss=10.4321  lr=0.000e+00  ms/step=  N/A
step  100  loss= 5.1234  lr=2.99e-04  ms/step= 280
step  500  loss= 3.4567  lr=5.95e-04  ms/step= 282
step 1000  loss= 2.8902  lr=4.50e-04  ms/step= 281
step 2000  loss= 2.4521  lr=6.00e-05  ms/step= 280
val_loss=2.41 (best, saved)

[sample] Once upon a time, there was a little girl named Lily.
She loved to play with her toys. One day, she found a big box.

Sanity numbers:

• Initial loss ≈ log(50257) ≈ 10.83. ✅
• Final val loss for 10M params on TinyStories: ~2.3–2.5 (scales like Chinchilla predicts).
• 280 ms/step on T4 is normal; 90 ms/step on a 4090.

5. Diagnosing training pathologies

6. Common pitfalls

1. Running --prepare every time — cache the .bin files; tokenization is slow.
2. Forgetting device_type in autocast on CPU — BF16 autocast on CPU only works in PyTorch 2.0+.
3. memmap on a remote/Network file — random access is brutal on NFS. Copy to local SSD.
4. torch.compile(model) can help but breaks eager debugging — enable last.
5. Checkpoint with model.state_dict() only — lose optimizer state → can't resume cleanly.

7. Stretch exercises

• Scale up to d=512, 8 layers, block=512. ~30M params, ~4 hours on a single A100. Val loss should reach ~1.9.
• Replace LayerNorm with RMSNorm — ~10% speedup, no quality loss.
• Add RoPE (rotary position embeddings) — better long-context generalization.
• Use SwiGLU MLP — ~2% perplexity improvement for ~50% more MLP params.
• Compute Chinchilla compute-optimal for your params: tokens ≈ 20 × params. For 10M params, train on 200M tokens.
• Run on FineWeb-Edu sample instead of TinyStories — better quality data, harder to learn from.
• Visualize attention at a checkpoint: pick a position, plot attention weights across all layers. Identify induction heads.

8. What this lab proves about you

You can run a complete pretraining loop end-to-end, choose every hyperparameter with justification, debug loss-curve pathologies, and ship a generating model from raw text. This is the bar Anthropic/OpenAI use for applied research engineers — the difference between someone who knows transformers and someone who can train them.

Phase 6 — Fine-tuning, Instruction Tuning, Preference Optimization

Difficulty: ⭐⭐⭐⭐☆ | Estimated Time: 2.5 weeks Roles supported: Post-training Engineer, Production Model Post-Training (Anthropic-style), Applied AI Engineer.

Why This Phase Exists

The frontier-lab post-training stack — SFT → reward model → preference optimization — is what turns a base LM into Claude / ChatGPT / Gemini. Anthropic's "Production Model Post-Training" role explicitly asks for hands-on experience with this exact pipeline.

You will fine-tune a real 7B model on a single 24 GB GPU using QLoRA, then run DPO with a preference dataset, and produce a quantitative before/after eval.

Concepts

• Pretraining vs SFT vs preference optimization
• Chat templates (ChatML, Llama-3, Mistral) — and why they matter
• Loss masking on prompt tokens
• LoRA: low-rank adapters, math, parameter savings (A ∈ R^{d×r}, B ∈ R^{r×d})
• QLoRA: 4-bit base + LoRA on top, NF4 quantization, double quantization
• PEFT library mechanics
• Reward modeling: pairwise loss, Bradley-Terry assumption
• RLHF / PPO conceptual flow (without implementing PPO end-to-end)
• DPO derivation from RLHF objective
• IPO, KTO, ORPO — the DPO family
• RLAIF (AI feedback) and Constitutional AI overview
• Catastrophic forgetting & mitigation

Labs

Lab 01 — Supervised Fine-Tuning (SFT) on Instruction Data

Lab 02 — LoRA & QLoRA on a 7B Model (Single GPU)

Lab 03 — Building an Instruction Dataset (Synthetic + Curated)

Lab 04 — Reward Modeling + DPO Preference Optimization

Deliverables Checklist

• SFT-trained small model with eval comparison
• QLoRA fine-tune of 7B on 24 GB GPU
• 5k-example synthetic instruction dataset
• DPO-trained model with preference win-rate report

Interview Relevance

• "Compare SFT, RLHF, DPO"
• "Walk through LoRA math"
• "Why does QLoRA work? What's NF4?"
• "Derive the DPO loss"
• "How would you build a preference dataset?"

Warmup Guide — Fine-tuning & Instruction Following

Zero-to-expert primer for Phase 06: how a next-token predictor becomes an assistant — SFT and chat templates, LoRA/QLoRA mechanics, and the preference-tuning landscape (RLHF → DPO) — assuming Phase 05's training fluency.

Table of Contents

• Chapter 1: The Gap Between a Language Model and an Assistant
• Chapter 2: Supervised Fine-Tuning — Mechanics That Matter
• Chapter 3: Chat Templates — The Underrated Contract
• Chapter 4: LoRA — Fine-Tuning as a Low-Rank Update
• Chapter 5: QLoRA — Training on a Quantized Base
• Chapter 6: Preference Tuning — RLHF and DPO
• Chapter 7: What Fine-Tuning Can and Cannot Do
• Lab Walkthrough Guidance
• Success Criteria
• Interview Q&A
• References

Chapter 1: The Gap Between a Language Model and an Assistant

A pretrained LM completes text — ask it "How do I bake bread?" and a plausible continuation is more questions ("How long does it take? What flour should I use?"), because the internet contains question lists. Nothing in pretraining says "answer."

Alignment closes the gap in (classically) three stages:

1. SFT (this phase's core): train on (instruction → good response) pairs — teaches the format and role of being an assistant.
2. Preference tuning (RLHF/DPO, Ch. 6): teach which of two responses is better — quality, harmlessness, tone — things hard to demonstrate but easy to compare.
3. The base model's knowledge comes along for free: alignment is widely understood as mostly eliciting and formatting capabilities pretraining built, not adding knowledge — the "superficial alignment hypothesis" (LIMA: 1K excellent examples sufficed for format). This framing predicts both fine-tuning's cheapness and its limits (Ch. 7).

Chapter 2: Supervised Fine-Tuning — Mechanics That Matter

SFT is Phase 05's training loop with three changes that carry all the difficulty:

• Loss masking: the example is prompt + response, but loss is computed on response tokens only (prompt positions set to ignore_index=-100). Training on prompt tokens teaches the model to generate prompts — a real and common bug whose symptom is the model continuing user-style text. (You met the identical masking in the model-accuracy track's multimodal lab — same mechanism, same off-by-one risk at the boundary.)
• Hyperparameters shrink: LR ~1e-5–2e-5 (vs 1e-3-ish pretraining), 1–3 epochs, small batches. You are adjusting a finished model, not training one — overcooked SFT shows as repetitive, sycophantic, distribution-collapsed outputs.
• Data quality dominates data quantity: 1K–10K excellent, diverse examples beat 100K scraped mediocre ones (LIMA's lesson, repeatedly replicated). Curation — dedup, decontamination (your eval sets! — Phase 08), length/format balance — is most of the actual work. Catastrophic forgetting is the tax: narrow SFT data measurably degrades broad capabilities; mitigations are mixing in general data and short training.

Chapter 3: Chat Templates — The Underrated Contract

Multi-turn conversations serialize to one token stream via a template — special tokens marking roles and turns (Phase 01 Ch. 6's special-token machinery, now load-bearing):

<|im_start|>system ... <|im_end|> <|im_start|>user ... <|im_end|> <|im_start|>assistant ...

The facts that cause production incidents:

• The template is part of the model's weights-contract: it must match training exactly — wrong template degrades quality silently (no error, just a worse model). This is the #1 cause of "this model is bad actually" reports; check tokenizer.apply_chat_template against the model card before any other debugging.
• EOS/EOT discipline: training must teach the model to emit the end-of-turn token (loss on that token!), else generation never stops — the "model rambles forever" bug is usually a masking bug at the turn boundary.
• In multi-turn training data, mask everything except assistant turns; whether to train on all assistant turns or only the last is a real design choice (all-turns is standard; it reuses context efficiently).

Chapter 4: LoRA — Fine-Tuning as a Low-Rank Update

Full fine-tuning of a 7B model needs weights + grads + Adam states ≈ 70+ GB (Phase 05 Ch. 2's 2×-FP32 fact) — and produces a full copy per task. LoRA: freeze $W$; learn a low-rank delta:

$$W' = W + \frac{\alpha}{r} BA, \qquad B \in \mathbb{R}^{d \times r} \text{ (init 0)}, \quad A \in \mathbb{R}^{r \times k} \text{ (init gaussian)}, \quad r \in [4, 64]$$

• Why it works: task adaptation empirically has low intrinsic rank — the change needed is far simpler than the model. B=0 init means training starts exactly at the base model (no cold-start damage); $\alpha/r$ decouples update magnitude from rank.
• What to adapt: classically attention's $W_q, W_v$; modern practice (QLoRA paper's finding) adapts all linear layers — with rank low, parameters stay <1%.
• Trainable fraction: r=16 on a 7B → ~0.1–0.5% of parameters; optimizer state shrinks proportionally — this, not gradient compute, is the memory win.
• Serving: merge ($W + \frac{\alpha}{r}BA$, zero inference overhead, loses swappability) vs keep-separate (tiny extra matmul; hot-swappable per-task adapters over one base — multi-tenant LoRA serving (S-LoRA-style) is a Phase 09-adjacent production pattern worth naming in interviews).

Chapter 5: QLoRA — Training on a Quantized Base

QLoRA composes LoRA with a 4-bit frozen base — backprop flows through the quantized weights into FP16 adapters:

• NF4: 4-bit grid placed at the quantiles of a standard normal — equal probability mass per level for normal-ish weights — information-theoretically optimal storage when you'll dequantize to float for compute (vs uniform INT4, which is what integer arithmetic hardware wants — the storage-vs-compute distinction).
• Double quantization: the per-block (64) absmax constants are themselves quantized — ~0.4 bits/param saved (~370 MB at 7B).
• Paged optimizer states handle memory spikes.
• Net: 7B fine-tune in ~10 GB, 65B in 48 GB — the democratization moment of 2023, and the lab's actual configuration. Quality: QLoRA ≈ LoRA ≈ full-FT on instruction tasks at matched data (the paper's Guanaco result) — with the caveat that aggressive new domains stress the frozen-base assumption.

(The deep quantization math — and using adapters to recover quantization accuracy — lives in the model-accuracy track's Phase 03/10 warmups; cross-reference rather than re-derive.)

Chapter 6: Preference Tuning — RLHF and DPO

SFT teaches format; preferences teach better-vs-worse — judgments easy to collect pairwise, hard to demonstrate.

• RLHF (the classic pipeline): (1) train a reward model on human preference pairs (Bradley-Terry loss on chosen-vs-rejected); (2) optimize the policy against it with PPO, with a KL penalty to the SFT model as the leash. Why the leash: unconstrained reward maximization finds the RM's blind spots — reward hacking (length inflation, sycophancy, weird token exploits). RLHF works (it built ChatGPT) but is operationally heavy: four models in memory (policy, reference, RM, value), RL instability, and the RM's quality caps everything.
• DPO (the 2023 simplification): the KL-constrained RLHF objective has a closed-form optimal policy; inverting it turns preference optimization into a supervised loss on the policy directly:

$$\mathcal{L} = -\log \sigma!\left(\beta\left[\log\tfrac{\pi(y_w|x)}{\pi_{ref}(y_w|x)} - \log\tfrac{\pi(y_l|x)}{\pi_{ref}(y_l|x)}\right]\right)$$

— raise the chosen response's likelihood relative to the reference, lower the rejected's, with $\beta$ playing the KL knob. No reward model, no RL loop, two models in memory. Tradeoffs to know: DPO is bounded by its offline preference data (no exploration), can over-optimize toward verbosity too, and variants (IPO, KTO, ORPO) patch specific failure modes. Industry reality: DPO-family for most open work; RLHF/ online methods at the frontier labs.

Chapter 7: What Fine-Tuning Can and Cannot Do

The judgment chapter — when an inference engineer should say "don't fine-tune":

• Great for: format/style/persona, task specialization (SQL, extraction schemas), tool-calling patterns, domain vocabulary, latency wins (a tuned 8B replacing a prompted 70B).
• Poor for: adding knowledge — facts injected by fine-tuning are brittle and hallucination-prone (the model learns to assert in-domain rather than to know); retrieval (RAG, Phase 07) is the right tool for knowledge, fine-tuning for behavior. The slogan that survives scrutiny: RAG for what the model should know, fine-tuning for how it should act.
• Dangerous for: safety properties (fine-tuning on even benign data measurably erodes refusal training — a real deployment consideration), and anything where eval contamination from your SFT data corrupts your metrics (Phase 08's discipline).
• The decision ladder to recite: prompt engineering → few-shot → RAG → fine-tune → pretrain-from-scratch, escalating only when the cheaper rung measurably fails.

Lab Walkthrough Guidance

Lab 02 — LoRA & QLoRA:

1. Implement the LoRA module yourself before importing PEFT: wrap a frozen nn.Linear, add $BA$ with B-zero init; test that step 0 output equals the base model exactly, and that only adapter params receive grads (the same two tests as the model-accuracy multimodal lab — the freezing discipline generalizes).
2. Prepare the instruction dataset with explicit loss masking; print one fully rendered example (template + mask visualized) and check the EOT token carries loss — five minutes here prevents the two classic failure modes (Ch. 2–3).
3. SFT a small base (1–3B class) with LoRA; track val loss and generations on fixed prompts; verify the assistant-formatting behavior emerges.
4. Switch to QLoRA (4-bit NF4 base): match LoRA's quality within noise while logging peak VRAM for both — the memory table is the deliverable.
5. Ablate: r ∈ {4, 16, 64} and attention-only vs all-linear targets — quality vs adapter size; then merge the best adapter and verify merged == unmerged outputs.

Success Criteria

You are ready for Phase 07 when you can, from memory:

1. Explain the three alignment stages and the superficial-alignment hypothesis with its evidence and limits.
2. State the SFT loss-masking rule, the two bugs it prevents, and the EOS-discipline bug's symptom.
3. Write the LoRA update with both inits and the $\alpha/r$ role; give the trainable-% and the merge-vs-swap tradeoff.
4. Explain NF4-vs-INT4 as storage-vs-compute optimality and double quantization's saving.
5. Sketch RLHF's pipeline with the KL leash's purpose, then derive DPO's pitch (closed- form inversion → supervised loss) and its offline limitation.
6. Recite the RAG-vs-fine-tuning division and the escalation ladder.

Interview Q&A

Q: A team fine-tuned on their product docs and the model now confidently invents product features. What happened and what should they do? Fine-tuning taught behavior (assert fluently about the product domain) not knowledge (the docs' facts aren't reliably stored). In-domain confidence without in-domain grounding = amplified hallucination. Fix: RAG over the docs for facts — fine-tune only for format/tone if needed; keep an eval set of answerable+unanswerable product questions, and measure refusal-when-ungrounded. The conceptual error — fine-tuning as knowledge injection — is the thing to name.

Q: Why does DPO work without a reward model — what's the trick? The RLHF objective (maximize reward minus β·KL-to-reference) has a closed-form optimum: $\pi^(y|x) \propto \pi_{ref}(y|x)\exp(r(y,x)/\beta)$. Solve for $r$ in terms of $\pi^/\pi_{ref}$ and substitute into the Bradley-Terry preference likelihood: the reward cancels into log-likelihood-ratios of the policy itself. The "reward model" is implicit in the policy — you optimize it directly by supervised learning on preference pairs. The cost: you inherit the offline data's coverage; no exploration.

Q: When would you choose full fine-tuning over LoRA today? When the adaptation is large-rank by nature: substantial new domains/languages, changing base behaviors deeply, or continued-pretraining-scale token budgets — measured by LoRA-vs-full ablation gaps, not vibes. Also when serving exactly one task at scale (merge erases LoRA's deployment advantage) and you have the memory anyway. For the standard instruct/persona/task tune, LoRA-family matches quality at ~1% of the optimizer footprint, so it's the default; the burden of proof sits on full FT.

References

• Ouyang et al., Training language models to follow instructions (InstructGPT) (2022) — arXiv:2203.02155 — the SFT+RLHF blueprint
• Zhou et al., LIMA: Less Is More for Alignment (2023) — arXiv:2305.11206
• Hu et al., LoRA (2021) — arXiv:2106.09685
• Dettmers et al., QLoRA (2023) — arXiv:2305.14314
• Rafailov et al., Direct Preference Optimization (2023) — arXiv:2305.18290
• Qi et al., Fine-tuning Aligned Language Models Compromises Safety (2023) — arXiv:2310.03693
• HuggingFace: chat templating docs and TRL library — read DPOTrainer's loss
• Sheng et al., S-LoRA: Serving Thousands of Concurrent LoRA Adapters (2023) — arXiv:2311.03285

🛸 Hitchhiker's Guide — Phase 6: Fine-Tuning & Instruction Tuning

Read this if: You can pretrain a small LM, but you don't yet know the difference between SFT, RLHF, DPO, ORPO; you've heard "LoRA" but can't write its math; or you can't explain why QLoRA lets you fine-tune 70B on a single A100.

0. The 30-second mental model

A pretrained "base" model is a calculator that loves to complete the most likely text. To turn it into a useful assistant, you do post-training in 1–3 stages:

1. SFT (Supervised Fine-Tuning): train on (prompt, ideal_response) pairs to teach the format and behavior. ~10k–1M examples.
2. Preference learning (RLHF, DPO, ORPO): align outputs with human preferences using (prompt, chosen, rejected) triplets. The model learns subtle quality, helpfulness, and refusal behaviors that are easier to prefer than to write.
3. (Optional) Constitutional AI / RLAIF — use an LLM to generate the preference labels at scale.

Plus a separate axis: how you fine-tune.

• Full fine-tune: update every parameter. Highest quality, biggest cost (memory + storage).
• LoRA (Low-Rank Adaptation): add tiny rank-r adapters; freeze base. ~100× less memory, near-equal quality.
• QLoRA: LoRA on top of a 4-bit quantized base. Lets you fine-tune 70B on one A100 80GB.

By the end of Phase 6 you should:

• Build an SFT dataset and run a real SFT job with HuggingFace trl's SFTTrainer.
• Derive LoRA's math; explain r and α.
• Configure QLoRA correctly (NF4, double-quant, paged optimizers).
• Explain DPO's loss derivation from PPO's optimum.
• Know when to fine-tune vs RAG vs prompt-engineer.

1. The post-training pipeline at a glance

Base model  ──SFT on demos──►  SFT model  ──preference learning──►  Aligned model
   (lossy completer)              (instruction follower)             (helpful + harmless)

Real production stacks (OpenAI, Anthropic, Llama-3): SFT on millions of demos → DPO (or RLHF) on hundreds of thousands of preferences → optional rejection sampling, constitutional AI, red-teaming, eval gates.

2. Stage 1 — Supervised Fine-Tuning (SFT)

2.1 The data

Each example is (prompt, response). Crucially, loss is computed only on the response tokens, not the prompt. The prompt is conditioning context.

Common templates:

• ChatML / OpenAI format:

  <|im_start|>system
  You are a helpful assistant.
  <|im_end|>
  <|im_start|>user
  Explain attention.
  <|im_end|>
  <|im_start|>assistant
  Sure! Attention is a mechanism that...
  <|im_end|>
  

• Alpaca format:

  Below is an instruction...
  ### Instruction:
  Explain attention.
  ### Response:
  Sure! Attention is a mechanism that...
  

• Llama-3 format has its own special tokens.

The exact template MUST be consistent between training and inference. A common bug: training with one template, serving with another → garbled outputs.

2.2 Loss masking

Compute loss only on the assistant's tokens. Implementation: build a labels tensor identical to input_ids, then set labels[i] = -100 for every token that's part of the prompt. PyTorch's cross_entropy ignores -100.

trl's SFTTrainer does this automatically when you pass formatting_func and a response_template.

2.3 The classic SFT datasets

• Alpaca (52k, GPT-3.5 generated) — historical baseline, low quality but shows the format.
• Dolly-15k (Databricks, 2023) — 15k human-written; permissively licensed. Used in Lab 02.
• OpenAssistant Conversations — 161k human conversations.
• UltraChat — 1.5M GPT-3.5 conversations.
• ShareGPT — real ChatGPT conversations.

A common pattern at frontier labs: ~100k–1M examples, with ~70% LLM-generated and ~30% human-curated/filtered.

2.4 SFT hyperparameters that matter

• LR: small. ~1e-5 to 5e-5 for full fine-tune; ~1e-4 to 3e-4 for LoRA.
• Epochs: 1–3. SFT overfits fast. More epochs ≠ better.
• Batch size: large effective batch (64–256) via gradient accumulation.
• Cosine decay with short warmup (3% of steps).

3. Parameter-Efficient Fine-Tuning (PEFT)

3.1 Why PEFT exists

A 70B model needs ~140GB for weights, ~280GB for fp32 AdamW state, ~10–50GB for activations. That's ~500GB peak — eight A100 80GBs. Most practitioners cannot afford this.

PEFT methods freeze the base and train tiny additions. The full base + adapter at inference is identical in size to the base; only the adapter (~few hundred MB) needs to be stored per fine-tune.

3.2 LoRA — Low-Rank Adaptation (Hu et al., 2021)

Key observation: empirically, fine-tuning updates ΔW to weight matrices have low intrinsic rank. So decompose ΔW as the product of two thin matrices:

$$ W_{\text{eff}} = W_0 + \Delta W = W_0 + B A $$

where A ∈ ℝ^{r×k}, B ∈ ℝ^{d×r}, r ≪ \min(d, k). Only A and B train; W_0 is frozen.

Forward pass:

$$ y = W_0 x + (α/r) \cdot B (A x) $$

The α/r is the LoRA scaling. Convention: α = 2r (so the scaling is 2), but it's tunable — it controls how strongly the adapter influences the output.

Parameter savings

For a d × k = 4096 × 4096 weight: full update = 16M params. LoRA r = 16: 16 × (4096 + 4096) = 131k params. 122× fewer. Apply LoRA to all attention QKV+O and MLP up/down/gate: ~7 matrices/layer × 32 layers = ~225 matrices, total adapter ≈ 30M params for a 7B model. Optimizer states for those 30M params fit in <1GB.

Initialization

A initialized with kaiming_uniform, B initialized to zero. So BA = 0 at start, the adapter is initially the identity perturbation, and the model behaves exactly like the base. Loss starts at the base model's loss; training improves from there.

Where to apply LoRA

The Hu paper applied only to W_q and W_v. Modern practice: apply to all attention and MLP projections (q_proj, k_proj, v_proj, o_proj, up_proj, gate_proj, down_proj). More targets = more adapter params = better quality. Lab 02 uses this set.

Choosing r

Typical: 8, 16, 32, 64. Bigger r = more capacity to fit the new task. r = 16 is a great default. For very different downstream tasks (e.g., teaching a new language), r = 64 may help.

3.3 QLoRA (Dettmers et al., 2023)

QLoRA = LoRA on top of a 4-bit quantized base model. Three innovations:

1. NF4 (NormalFloat-4): a 4-bit datatype whose quantization levels are chosen to be information-theoretically optimal for normal-distributed data. Pretrained weights are approximately N(0, σ), so NF4 minimizes quantization error in the relevant range. (Standard 4-bit integer quantization wastes bits on values that rarely occur.)
2. Double quantization: the per-block quantization scales themselves are quantized, saving another ~0.4 bits/param on average.
3. Paged optimizers: optimizer state pages move between GPU and CPU memory via NVIDIA's Unified Memory, avoiding OOM spikes during gradient checkpointing.

End result: fine-tune 70B on a single A100 80GB at near-equal quality to full fp16 fine-tuning. Bombshell paper.

In Lab 02 you'll set up QLoRA via:

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
)

3.4 Other PEFT methods

• Prefix tuning / Prompt tuning — train soft "virtual tokens" prepended to inputs. Older, less popular than LoRA now.
• (IA)³ — scale activations by learned vectors. Tiny but limited capacity.
• DoRA (Liu et al. 2024) — decomposes weight updates into magnitude + direction; small quality bump over LoRA at same r.

4. Stage 2 — Preference Learning

4.1 The data

Triplets (prompt, chosen_response, rejected_response). Sources:

• Human annotators ranking pairs (most expensive, highest signal).
• AI judges (RLAIF) — cheap; quality bounded by judge.
• Self-rejection sampling — generate multiple, score with a reward model, keep best/worst.

4.2 RLHF (PPO) — the original recipe

Three steps:

1. Train a reward model r_φ(x, y): a small head on top of the SFT model that outputs a scalar. Trained on preferences with the Bradley-Terry loss: $$\mathcal{L}_{RM} = -\log \sigma(r(x, y_w) - r(x, y_l))$$ where y_w is chosen, y_l is rejected.
2. Use PPO (Proximal Policy Optimization, Schulman et al. 2017) to optimize the SFT model against the reward, with a KL penalty to a frozen reference (the SFT model itself): $$\mathcal{L}{RLHF} = \mathbb{E}y[r(x, y)] - β , D{KL}(\pi(\cdot|x) | \pi{\text{ref}}(\cdot|x))$$ The KL prevents the policy from drifting too far and reward-hacking.
3. Generate rollouts, score with reward model, run PPO updates. Repeat.

PPO is complex: ~7 hyperparams; unstable; needs distributed rollout infrastructure; reward-hacking is real (model finds adversarial paths to high reward). Cost is huge — 4× SFT cost easily.

4.3 DPO — Direct Preference Optimization (Rafailov et al., 2023)

Insight: you can derive PPO's optimal policy in closed form (assuming the KL-constrained reward objective), and inverting that derivation gives a contrastive loss directly on (chosen, rejected) pairs. No reward model. No rollouts. Just SFT-like training.

Loss:

$$ \mathcal{L}{DPO} = -\log \sigma!\left(β \log \frac{\pi(y_w | x)}{\pi{\text{ref}}(y_w | x)} - β \log \frac{\pi(y_l | x)}{\pi_{\text{ref}}(y_l | x)}\right) $$

Where π is the trainable policy and π_ref is the frozen SFT model. Intuitively: increase π's probability of chosen relative to ref, decrease for rejected.

DPO has effectively replaced PPO as the default for new projects in 2024+. Simpler, more stable, often matches or beats PPO. Llama-3 instruct uses DPO.

4.4 ORPO (Hong et al., 2024)

Combines SFT and preference learning in a single stage. Loss = standard cross-entropy on chosen + odds-ratio penalty against rejected. Skips the SFT-then-DPO sequence; one-shot post-training.

4.5 Constitutional AI / RLAIF (Bai et al., 2022 — Anthropic)

Use an LLM to critique and revise outputs against a written "constitution" of principles, generating preference pairs at scale without humans. Anthropic's main alignment recipe.

4.6 References

• Christiano et al. (2017), Deep RL from Human Preferences — the first RLHF paper.
• Stiennon et al. (2020), Learning to Summarize with Human Feedback — first compelling RLHF for LLMs.
• Ouyang et al. (2022), Training Language Models to Follow Instructions with Human Feedback — InstructGPT, the basis of ChatGPT.
• Schulman et al. (2017), Proximal Policy Optimization Algorithms.
• Rafailov et al. (2023), Direct Preference Optimization.
• Hong et al. (2024), ORPO: Monolithic Preference Optimization without Reference Model.
• Bai et al. (2022), Constitutional AI: Harmlessness from AI Feedback.

5. The lab walkthrough (lab-02-lora-qlora)

5.1 What you'll build

Fine-tune Mistral-7B (or similar 7B base) on Dolly-15k with QLoRA:

• Load 4-bit quantized base via BitsAndBytesConfig.
• Configure LoRA with r=16, α=32, applied to attention QKV+O and MLP up/down/gate.
• Use paged_adamw_8bit optimizer.
• Train with SFTTrainer for 1 epoch, ~30 minutes on a single A100 40GB.
• Save the LoRA adapter (~100MB).
• Inference: load the base + adapter, generate.

5.2 Things to read carefully

• The exact target_modules list — this depends on the model architecture. For Llama/Mistral: ["q_proj", "k_proj", "v_proj", "o_proj", "up_proj", "gate_proj", "down_proj"].
• The prepare_model_for_kbit_training() call — disables some incompatible features and casts the LM head to fp32 for numerical stability.
• The formatting_func and response_template — these tell SFTTrainer how to mask labels.
• The merge step (model.merge_and_unload()) — fuses adapter weights into the base for deployment.

5.3 Sanity checks

• Initial loss should be the base model's loss on the format (~2–3).
• Loss should drop to ~1.0–1.5 by epoch end on Dolly.
• Generated responses should be grammatical and follow Dolly's tone.

6. When to fine-tune (vs RAG vs prompt)

A common mistake: trying to fine-tune in facts that change weekly. Use RAG.

7. Common interview questions on Phase 6 material

1. Walk through SFT, RLHF, and DPO. When would you use each?
2. Derive LoRA's math. Explain r and α.
3. What's NF4 and why is it better than INT4?
4. Why does QLoRA let you fine-tune 70B on one GPU?
5. Compare PPO's KL penalty and DPO's reference model — they're related, how?
6. What's reward hacking and how do you mitigate it?
7. When would you fine-tune instead of RAG?
8. How do you mask labels for SFT so the model doesn't train on the prompt?
9. Sketch the Bradley-Terry reward model loss.
10. What's the role of α / r scaling in LoRA at inference?
11. How would you serve 100 different LoRA adapters in production? (Bridges to Phase 9.)
12. Why is constitutional AI scalable in a way RLHF isn't?

8. From solid → exceptional

• Implement LoRA from scratch (no peft): wrap an nn.Linear so its forward adds a low-rank update. Confirm gradient flow only into the adapter.
• Implement the DPO loss in pure PyTorch (no trl); compute against a tiny preference dataset. Verify against trl reference.
• Run a side-by-side SFT vs SFT+DPO vs SFT+ORPO on the same base, evaluate with MT-Bench. Report numbers.
• Implement rejection sampling with reward model: generate 16 responses per prompt, score with a separate RM, keep top-1. Compare to base sampling.
• Read the Constitutional AI paper and write a one-page summary; sketch how you'd build a small CAI loop on a 1B model.
• Train multiple LoRA adapters for different tasks; demonstrate hot-swapping at inference (e.g., via peft's set_adapter).

9. Recommended cadence

Lab 02 — QLoRA Fine-Tune of a 7B Model (Solution Walkthrough)

Phase: 6 — Fine-tuning & Instruction Following | Difficulty: ⭐⭐⭐⭐☆ | Time: 3–6 hours (incl. training)

Concept primer: ../HITCHHIKERS-GUIDE.md §LoRA, §QLoRA, §SFT.

Run

pip install -r requirements.txt
huggingface-cli login   # for gated Llama-3
python solution.py

Hardware: 24 GB GPU (RTX 3090/4090/A5000/A10). For Colab T4 (16 GB), use Qwen/Qwen2-1.5B.

0. The mission

Fine-tune Llama-3-8B (or Qwen2-7B) on a single 24 GB consumer GPU using QLoRA: 4-bit base + LoRA adapters in BF16. The fully-merged model would need ~32 GB just for weights in BF16; QLoRA reduces this to ~6 GB and trainable parameters to ~50 MB.

This is the technique that democratized LLM fine-tuning. Every "I fine-tuned a 7B model on my gaming GPU" project uses it.

1. The math

1.1 LoRA decomposition

For any linear layer $y = Wx$ with $W \in \mathbb{R}^{d \times k}$, freeze $W$ and add a low-rank update:

$$ y = Wx + BAx, \quad B \in \mathbb{R}^{d \times r}, ; A \in \mathbb{R}^{r \times k}, ; r \ll \min(d, k) $$

$A$ is initialized to random Gaussian, $B$ to zero — so $BA = 0$ at step 0 (model output is unchanged). With $r = 16$ and $d = k = 4096$, trainable params per layer drop from $16{,}777{,}216$ to $131{,}072$ (a 128× reduction).

A scalar $\alpha / r$ scales the update: $y = Wx + (\alpha / r) BAx$. Convention: $\alpha = 2r$ so the scale is 2.0.

1.2 QLoRA's three tricks

1. NF4 quantization — a 4-bit data type optimized for normally-distributed weights (which neural-net weights approximately are). Quantization levels are placed at the quantiles of $\mathcal{N}(0, 1)$. Less quantization error than uniform INT4.
2. Double quantization — quantize the per-block quantization constants themselves. Saves ~0.4 bits/param on top of NF4. Free.
3. Paged optimizer — use NVIDIA unified memory to swap optimizer states to CPU when GPU memory pressure spikes. Lets you fine-tune without OOM crashes during memory peaks.

Backward pass dequantizes 4-bit weights to BF16 on the fly — no quality loss. Forward + backward in BF16. Optimizer (AdamW) only updates LoRA params, so optimizer state is tiny.

2. Loading the model in 4-bit

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

bnb = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
)
model = AutoModelForCausalLM.from_pretrained(
    base_model,
    quantization_config=bnb,
    device_map="auto",
    attn_implementation="flash_attention_2",
)

• bnb_4bit_compute_dtype=torch.bfloat16 — dequantize to BF16 for the matmul. (FP16 also works but BF16 is more stable.)
• device_map="auto" — transformers' accelerate-based dispatcher places layers on available GPUs.
• flash_attention_2 — ~2× faster + much lower memory. Required for long-context fine-tuning.

model.config.use_cache = False                  # incompatible with grad checkpointing
model.gradient_checkpointing_enable()
model = prepare_model_for_kbit_training(model)

• Gradient checkpointing — trade compute for memory. Recompute activations during backward instead of storing them. Cost: ~30% slower; benefit: ~5× less activation memory — essential for 8B at 24 GB.
• prepare_model_for_kbit_training — casts LayerNorm/embedding outputs to FP32 for stability, enables requires_grad on input embeddings (so gradients flow back through the frozen base).

3. Attaching LoRA adapters

from peft import LoraConfig, get_peft_model

lora = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora)
model.print_trainable_parameters()
# trainable params: 41,943,040 || all params: 8,071,016,448 || trainable%: 0.52

Key choices:

• r=16, alpha=32 — the modal QLoRA settings. alpha = 2r is convention; some prefer alpha = r (scale 1.0). Both work; 2r is slightly more aggressive.
• All linear layers — attention (q/k/v/o) and MLP (gate/up/down). The QLoRA paper showed that targeting all linears gives ~2 perplexity points improvement over attention-only.
• lora_dropout=0.05 — small dropout on the LoRA path only (frozen base unaffected). Helps when fine-tuning on small datasets.
• bias="none" — don't train biases. Could try "lora_only" or "all" but rarely worth it.

4. Dataset & chat template

ds = load_dataset("tatsu-lab/alpaca", split="train").select(range(2000))

def format_example(ex):
    msgs = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": ex["instruction"] + ("\n\n" + ex["input"] if ex["input"] else "")},
        {"role": "assistant", "content": ex["output"]},
    ]
    return {"text": tokenizer.apply_chat_template(msgs, tokenize=False)}

ds = ds.map(format_example)

• Use the model's own chat template (apply_chat_template). Llama-3 uses <|begin_of_text|><|start_header_id|>system<|end_header_id|>.... Qwen uses <|im_start|>system\n...<|im_end|>. Hardcoding the wrong template silently destroys quality.
• 2000 examples is enough to teach instruction-following style on a base model. For domain knowledge, you need 10k+.

5. SFTTrainer setup

from trl import SFTTrainer, SFTConfig

cfg = SFTConfig(
    output_dir="./qlora-out",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,                # effective batch = 16
    num_train_epochs=2,
    learning_rate=2e-4,
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    bf16=True,
    optim="paged_adamw_8bit",                     # 👈 QLoRA's paged optimizer
    max_seq_length=1024,
    packing=True,                                  # concat short examples → fill seq
    logging_steps=20,
    save_steps=200,
    report_to="none",
)
trainer = SFTTrainer(model=model, args=cfg, train_dataset=ds, dataset_text_field="text")
trainer.train()

Key choices:

• learning_rate=2e-4 — ~10× higher than full fine-tuning. LoRA params are randomly initialized and need bigger steps to learn.
• optim="paged_adamw_8bit" — the 8-bit AdamW from bitsandbytes with paging. Keeps optimizer state at ~25% of FP32 size and survives memory spikes.
• packing=True — concatenates short examples to fill max_seq_length. Eliminates padding waste. Critical for instruction datasets where most examples are <500 tokens.
• bf16=True — BF16 forward/backward. (FP16 with QLoRA is unstable.)
• warmup_ratio=0.03 — first 3% of steps are linear warmup. Smaller than pretraining warmup because we're fine-tuning, not training from scratch.

6. Saving and merging

trainer.model.save_pretrained("./qlora-out/adapter")

This saves only the LoRA adapter (~50 MB). For deployment, you typically merge:

from peft import PeftModel
base = AutoModelForCausalLM.from_pretrained(base_model, torch_dtype=torch.bfloat16)
merged = PeftModel.from_pretrained(base, "./qlora-out/adapter").merge_and_unload()
merged.save_pretrained("./merged-bf16")

• Merging requires a non-quantized base — you can't merge a LoRA adapter into a 4-bit base while preserving quality. Load the base in BF16, merge, save.
• After merging, the model has the same architecture as the base (no adapter overhead at inference).

7. Expected output

trainable params: 41,943,040 || all params: 8,071,016,448 || trainable%: 0.52

{'loss': 1.4521, 'learning_rate': 6e-05, 'epoch': 0.04}
{'loss': 1.1234, 'learning_rate': 1.99e-04, 'epoch': 0.20}
...
{'loss': 0.8021, 'learning_rate': 2e-06, 'epoch': 1.99}
{'train_runtime': 5400.0, 'train_samples_per_second': 0.74}

Sanity checks:

• Loss starts near 2.0, ends near 0.8–1.0 for typical SFT data.
• VRAM usage during training: ~14–18 GB on a 24 GB card. If you OOM, lower per_device_train_batch_size to 1 or max_seq_length to 512.
• Sample from the merged model afterward and compare to the base — the fine-tune should follow instructions in the assistant turn instead of continuing the prompt.

8. Common pitfalls

1. Wrong chat template — silent quality killer. Always use tokenizer.apply_chat_template, never hand-format.
2. Forgetting model.config.use_cache = False with grad checkpointing → silent slowdown + warning.
3. load_in_8bit instead of 4-bit — 8-bit doesn't fit 8B in 24 GB during training (only inference).
4. flash_attention_2 not installed — fall back to eager attention, doubles VRAM, halves throughput.
5. Training a chat model on raw text (no chat template) — you wreck the model's existing instruction-following.
6. Saving the full model instead of the adapter — wastes 16 GB of disk per checkpoint.
7. Merging at FP16 precision — quality loss vs BF16. Always merge in BF16.

9. Stretch exercises

• DPO on top of SFT: take your SFT'd model + a preference dataset (e.g., argilla/distilabel-intel-orca-dpo-pairs) and run trl.DPOTrainer. Measure win-rate vs the SFT-only model.
• Multi-LoRA serving: train two adapters on different domains; load both into one base; route at inference time.
• Compare ranks: train at r=4, 16, 64. Plot loss vs trainable params. The 4↓16 jump should be large; 16↓64 small.
• Compare full FT vs LoRA at same compute: full fine-tune a 1.5B model vs LoRA on a 7B — which is better at the same wall-clock?
• Eval with lm-eval-harness on MMLU/GSM8K before and after — by how much does instruction tuning hurt raw-knowledge benchmarks (the alignment tax)?
• Try GaLore or DoRA as alternatives to LoRA — newer parameter-efficient methods with slightly different tradeoffs.

10. What this lab proves about you

You can stand up a production fine-tuning pipeline for a 7B+ model on consumer hardware, justify every hyperparameter (rank, alpha, target modules, optimizer choice, packing), and articulate the QLoRA tricks that make it possible. This is the bar for Phase-6 — and it's the most-demanded skill in current LLM engineering job postings.

Phase 7 — RAG, Retrieval, Agents

Difficulty: ⭐⭐⭐⭐☆ | Estimated Time: 2 weeks Roles supported: Applied AI Engineer (OpenAI-style), LLM Inference Engineer, ML Systems Engineer.

Why This Phase Exists

RAG is the most-deployed LLM pattern in industry. The OpenAI Applied AI Engineering JD is essentially "build production RAG and agentic systems". The interview bar is no longer "did you call a vector DB" — it is "did you compare BM25 vs dense vs hybrid vs ColBERT, did you re-rank, did you measure with RAGAS, did you handle long-context tradeoffs, did you build observability".

Concepts

• Embedding models for retrieval: sentence-transformers, E5, BGE, Cohere embed, OpenAI text-embedding-3
• Vector index types: flat, IVF, HNSW, PQ, IVF-PQ tradeoffs
• Vector DBs: FAISS (library), Qdrant, Weaviate, pgvector, Milvus
• Sparse retrieval: BM25, TF-IDF
• Hybrid retrieval: RRF (reciprocal rank fusion), weighted sum
• Re-ranking: cross-encoders (BGE-reranker), ColBERT (late interaction)
• Chunking: fixed-size, sentence, recursive, semantic, late-chunking
• Query rewriting / HyDE / multi-query
• RAG evaluation: RAGAS (faithfulness, answer relevance, context precision/recall)
• Agents: ReAct loop, tool use, function calling
• Structured outputs: JSON schema, constrained decoding (Outlines, lm-format-enforcer, OpenAI structured outputs)
• Long-context vs RAG tradeoff

Labs

Lab 01 — Embeddings & Vector Search Fundamentals

Lab 02 — Production RAG Pipeline (End-to-End)

Lab 03 — Hybrid Retrieval + Re-Ranking

Lab 04 — Agents, Tool Use, Structured Output

Deliverables Checklist

• FAISS embedding-model benchmark
• Production RAG service with citations + streaming
• Hybrid retrieval + re-ranking with quality lift report
• Tool-using agent with constrained outputs

Interview Relevance

• "Design a RAG system for 100M docs at 1k QPS" (system design — see system-design/)
• "How do you evaluate RAG quality?"
• "Compare BM25, dense, hybrid"
• "How would you build an agent reliably?"

Warmup Guide — Retrieval, RAG & Agents

Zero-to-expert primer for Phase 07: grounding LLMs in external knowledge — sparse and dense retrieval, chunking, vector indexes, the RAG pipeline's failure points, and the agent loop — assuming Phase 02's embedding foundations.

Table of Contents

• Chapter 1: Why Retrieval — The Knowledge Problem
• Chapter 2: Sparse Retrieval — BM25, Still Undefeated
• Chapter 3: Dense Retrieval — Bi-Encoders and Their Training
• Chapter 4: Approximate Nearest Neighbors — HNSW and IVF
• Chapter 5: The RAG Pipeline — Chunking to Generation
• Chapter 6: Where RAG Fails — A Debugging Taxonomy
• Chapter 7: Rerankers and Hybrid Search
• Chapter 8: Agents — The Loop, Tools, and Failure Containment
• Lab Walkthrough Guidance
• Success Criteria
• Interview Q&A
• References

Chapter 1: Why Retrieval — The Knowledge Problem

An LLM's knowledge is frozen at training time, stored diffusely in weights, unattributable, and expensive to update (Phase 06 Ch. 7: fine-tuning injects behavior, not reliable facts). Retrieval-Augmented Generation splits the job: a retriever finds relevant text from a controllable corpus; the LLM reads it in-context and synthesizes. What this buys, precisely: updatable knowledge (re-index, done), attribution (citations to real passages), access control (retrieve only what this user may see), and a hallucination reduction (not elimination — Ch. 6). The price: you now operate a search engine in front of your LLM, and the search engine's quality bounds the system — garbage retrieved, garbage generated. This phase is mostly about the search engine.

Chapter 2: Sparse Retrieval — BM25, Still Undefeated

BM25 is TF-IDF (Phase 02 Ch. 2) with two refinements, and it remains the baseline that embarrasses fancy systems:

$$\text{BM25}(q, d) = \sum_{t \in q} \text{IDF}(t) \cdot \frac{tf(t,d),(k_1 + 1)}{tf(t,d) + k_1\left(1 - b + b,\frac{|d|}{\text{avgdl}}\right)}$$

• Term-frequency saturation ($k_1 \approx 1.2$): the 10th occurrence of a term adds far less than the 2nd — relevance isn't linear in counts.
• Length normalization ($b \approx 0.75$): long documents accumulate matches by bulk; penalize proportionally.

Why it endures: exact lexical match carries intent — IDs, error codes, function names, rare entities ("XK-2447 timeout") are precisely what dense embeddings blur and exactly what users search for. Zero training, interpretable scores, mature infrastructure (Lucene/Elasticsearch). The professional default: BM25 is the baseline every dense system must beat on your corpus, and hybrid (Ch. 7) usually beats both.

Chapter 3: Dense Retrieval — Bi-Encoders and Their Training

Dense retrieval embeds queries and documents into one vector space (Phase 02's machinery, industrialized): a bi-encoder encodes each independently — documents offline into an index, the query at request time — similarity = dot product.

What makes a retrieval embedding model (vs generic sentence similarity):

• Trained contrastively on (query, relevant-doc) pairs: InfoNCE loss pulls true pairs together against in-batch + hard negatives (BM25-retrieved near-misses — the ingredient that most improves quality; DPR's central lesson).
• Asymmetry: queries are short questions, documents are long prose — models use instruction prefixes ("query: …" / "passage: …", or task instructions in modern e5/bge/gte models) to condition the encoder per side. Omitting the prefix at inference silently costs recall — a top-3 production bug.
• Pooling: CLS or mean over tokens → one vector; L2-normalized so dot = cosine (Phase 02 Ch. 3's equivalence, now operational).
• The bi-encoder's structural weakness: query and document never interact during encoding — one vector must anticipate every question a document answers. That ceiling is what rerankers fix (Ch. 7).

Chapter 4: Approximate Nearest Neighbors — HNSW and IVF

Exact top-k over N vectors is O(N·d) per query — fine to ~1M, then you buy speed with recall:

• HNSW (the default): a multi-layer skip-list-like proximity graph — search greedily descends from sparse upper layers to the dense bottom layer. Sub-millisecond at ~95–99% recall@10. Knobs: M (edges/node — memory vs recall), efConstruction (build quality), efSearch (query-time recall vs latency — the one you tune live). Costs: RAM-resident (graph + vectors), slow builds, deletes are awkward (tombstones + rebuild).
• IVF (inverted file): k-means the corpus into nlist cells; search probes the nprobe nearest cells only. Cheaper memory, faster builds, natural for disk/segmented systems; recall cliff if the true neighbor sits in an unprobed cell. Usually paired with PQ (product quantization — compress vectors ~16–64×; the quantization worldview of the model-accuracy track, applied to the index itself) for billion-scale.
• The two facts to keep: recall@k of the ANN layer upper-bounds end-to-end RAG quality (measure it against brute force on a sample — it's one of the lab's checks), and filtering interacts badly with ANN (metadata-filtered search degrades graph traversal; engines differ wildly here — ask this question of any vector DB you evaluate).

Chapter 5: The RAG Pipeline — Chunking to Generation

The full assembly (Lab 02 builds it end-to-end):

1. Ingest & chunk: split documents into retrieval units. The tension: small chunks (200–400 tokens) embed precisely but lose context; large chunks (1–2K) keep context but blur the embedding and waste prompt budget. Structure-aware splitting (headings, paragraphs, code blocks) with ~10–15% overlap is the sane default; chunk-expansion (retrieve small, hand the LLM the surrounding section) decouples the two needs and is the single highest-leverage upgrade.
2. Embed & index: batch-embed chunks (with the passage prefix!), store vectors + text + metadata (source, position — you need provenance for citations).
3. Retrieve: embed query (query prefix!), ANN top-k (k≈20–50, over-retrieve for the reranker).
4. Rerank (Ch. 7) → top 3–8 into the prompt.
5. Generate: a prompt that instructs grounding ("answer from the passages; say 'not found' if absent; cite [1][2]") — the instruction measurably matters.
6. Evaluate (the part everyone skips): retrieval metrics (recall@k, MRR against a labeled set — even 50 hand-labeled queries transform your ability to iterate) and generation metrics (faithfulness: are claims supported by the retrieved text? LLM-as-judge with spot-checking is the pragmatic tool — Phase 08 formalizes).

Chapter 6: Where RAG Fails — A Debugging Taxonomy

Diagnose by stage — each failure has a distinct signature (the lab plants several):

The meta-rule (same as every debugging chapter in this curriculum): bisect by stage with fixed inputs — log query → retrieved chunks → prompt → answer for every request, and you can localize any complaint in minutes; skip the logging and every incident is archaeology.

Chapter 7: Rerankers and Hybrid Search

• Hybrid search: run BM25 and dense in parallel, fuse with Reciprocal Rank Fusion: $\text{RRF}(d) = \sum_r 1/(60 + \text{rank}_r(d))$ — rank-based, no score calibration needed between incommensurable scorers, robust default. Hybrid covers Ch. 6's lexical gap and the dense model's paraphrase strength simultaneously; it's the standard production answer.
• Cross-encoder rerankers: feed (query, document) jointly through a transformer → relevance score. Full token-level interaction — far more accurate than bi-encoders — but O(pairs) inference, so it only reranks the top 20–100. The two-stage pattern (cheap recall → expensive precision) is the same architecture as every ranking system from web search to ads; recognizing it as such is the senior framing. Late-interaction models (ColBERT) sit between: token-level vectors with cheap MaxSim scoring.

Chapter 8: Agents — The Loop, Tools, and Failure Containment

An agent is an LLM in a loop with tools:

while not done: thought → tool call → observation → (repeat) → answer

(ReAct's pattern, now native via structured tool-calling APIs — the model emits JSON conforming to tool schemas; the runtime executes and returns observations.)

What the engineer actually owns:

• Tool design dominates agent quality: few, orthogonal tools with crisp descriptions, typed arguments, and informative error returns (the model reads errors and retries — a good error message is agent UX). RAG-as-a-tool ("search the KB") composing with this loop is the standard enterprise assistant shape.
• Failure containment: max-iteration caps (loops happen), timeouts, idempotent or confirmation-gated side-effecting tools, sandboxed execution, and observability — full traces of every thought/call/observation, because debugging an agent without traces is impossible.
• Honest economics: each loop iteration is an LLM call with a growing context — latency and cost compound; multi-step agents amplify single-step error rates (0.95⁵ ≈ 0.77). Use the dumbest sufficient pattern: a fixed pipeline beats an agent when the workflow is known; agents earn their cost when the path is genuinely dynamic.

Lab Walkthrough Guidance

Lab 02 — RAG Pipeline:

1. Corpus + labeled eval queries first (~30–50 questions with known source passages). Building eval before the pipeline is the discipline that makes every later step measurable.
2. BM25 baseline; record recall@5/MRR. This number is the bar.
3. Dense path: embed with prefixes, exact search first (no ANN) — measure; then add the ANN index and measure the recall delta at several efSearch values (Ch. 4's upper-bound check, made personal).
4. Hybrid RRF; then a cross-encoder reranker on top-30 — by now you have a five-row ablation table (BM25 / dense / hybrid / +rerank / oracle), which is the lab's deliverable.
5. Wire generation with grounding instructions + citations; run the faithfulness check on 20 answers; deliberately break one stage (wrong prefix, tiny efSearch) and confirm your logging localizes it (Ch. 6's bisection drill).

Success Criteria

You are ready for Phase 08 when you can, from memory:

1. Write BM25's formula and explain saturation and length normalization; state why lexical match still matters.
2. Describe bi-encoder training (InfoNCE, hard negatives, prefixes) and its structural ceiling vs cross-encoders.
3. Compare HNSW and IVF-PQ on memory/build/recall/deletes, and name the filter-interaction caveat.
4. Defend a chunking strategy including the small-vs-large tension and chunk expansion.
5. Recite six rows of the failure taxonomy with signatures.
6. Explain RRF (why rank-based) and the two-stage retrieve-rerank economics.
7. State the agent loop, two containment mechanisms, and the error-compounding arithmetic.

Interview Q&A

Q: Your RAG system answers correctly for verbatim queries but fails on paraphrases — yet your embedding benchmark scores are great. Diagnose. Suspects in order: (1) missing instruction prefixes at inference (benchmarks used them, your service doesn't — silent recall loss); (2) domain shift — the benchmark isn't your corpus; build the 50-query labeled set and measure your recall@k; (3) chunking that splits answers (the retrieval is "right" but partial); (4) ANN recall loss masking as model failure (compare against brute force). The structure — measure per stage against your own labels before swapping models — is the answer being graded.

Q: When is RAG the wrong tool? When the knowledge is behavioral (format, style, procedures → fine-tune, Phase 06); when the corpus fits comfortably in context (long-context stuffing beats a retrieval pipeline's complexity below ~50–100K tokens of stable docs); when queries need aggregation over many documents ("how many customers complained about X" — that's analytics/SQL, not top-k retrieval); and when latency budgets can't fit embed+search+rerank+generate. Naming retrieval's aggregation blindspot is the differentiator.

Q: How would you evaluate a RAG system end to end? Layered: ANN recall vs brute force (index health); retrieval recall@k/MRR on labeled query→passage pairs (the bottleneck metric); reranker NDCG on the same; generation faithfulness (claims supported by context — LLM-judge with human spot-checks) and answer correctness on a QA set; plus abstention quality (does it say "not found" when the corpus lacks the answer — measured with deliberately unanswerable queries). One labeled set of ~100 queries powers all of it; refusing to operate without one is the senior move.

References

• Robertson & Zaragoza, The Probabilistic Relevance Framework: BM25 and Beyond (2009)
• Karpukhin et al., Dense Passage Retrieval (2020) — arXiv:2004.04906
• Malkov & Yashunin, HNSW (2016) — arXiv:1603.09320
• Jégou et al., Product Quantization for Nearest Neighbor Search (2011)
• Lewis et al., Retrieval-Augmented Generation (2020) — arXiv:2005.11401
• Liu et al., Lost in the Middle (2023) — arXiv:2307.03172
• Cormack et al., Reciprocal Rank Fusion (SIGIR 2009)
• Khattab & Zaharia, ColBERT (2020) — arXiv:2004.12832
• Yao et al., ReAct: Synergizing Reasoning and Acting (2022) — arXiv:2210.03629
• Anthropic: Building effective agents — the use-the-dumbest-sufficient-pattern argument, from production experience

🛸 Hitchhiker's Guide — Phase 7: Retrieval, RAG & Agents

Read this if: You can fine-tune a model but you're hazy on dense vs sparse retrieval, why hybrid search wins, what "RAG faithfulness" measures, or how a tool-use agent loop actually works under the hood.

Folder note: this curriculum has both phase-07-rag-retrieval/ (older spec) and phase-07-retrieval-rag-agents/ (current). Prefer the latter for labs.

0. The 30-second mental model

A pretrained LLM is great at language, weak at facts (especially fresh or private ones). RAG (Retrieval-Augmented Generation) fixes this by retrieving relevant text at query time and stuffing it into the model's context window. Agents extend this further: the LLM can choose to call tools (search, code-execute, query a DB, send an email) and iterate.

The full stack:

query → embed → vector + keyword search → rerank → top-k passages
       ↓
       prompt = [system, retrieved passages, query] → LLM → answer with citations

For an agent:

loop:
  thought ← LLM(history)
  action ← LLM(thought)        # e.g., {tool: "search", args: ...}
  observation ← tool(action.args)
  history.append([thought, action, observation])
  if action.tool == "final_answer": break

By the end of Phase 7 you should:

• Know how dense embeddings work (Phase 2 → contrastive loss → SBERT/E5/BGE).
• Implement HNSW conceptually and know when to use which vector DB.
• Build a token-aware chunker, embed with a real model, index in Qdrant, retrieve, and stream answers from an LLM via Server-Sent Events.
• Combine BM25 + dense + reranker → understand why hybrid wins.
• Reason about RAG quality (RAGAS metrics: faithfulness, answer relevance, context precision/recall).
• Be able to design a tool-use agent loop and discuss its failure modes (loops, halting, cost).

1. Sentence and document embeddings

1.1 The journey from word2vec to E5

Phase 2 covered static word embeddings. For RAG we need sentence/passage embeddings — a single vector per chunk that captures meaning at the passage level.

Eras:

1. Average word vectors (or Arora SIF) — a 2017 baseline that's surprisingly hard to beat with naive pooling.
2. InferSent / Universal Sentence Encoder — supervised on NLI.
3. SBERT (Reimers & Gurevych, 2019) — fine-tune BERT with siamese networks on NLI/STS, take pooled output. The breakthrough that made dense retrieval practical at scale.
4. Contrastive sentence encoders (E5, BGE, GTE, Cohere embed-v3, OpenAI text-embedding-3): trained at scale with InfoNCE loss on (query, positive_passage, hard_negatives). Current SOTA.

1.2 The InfoNCE / contrastive loss

For a batch of B (query, positive) pairs, treat the other queries' positives as negatives within the same batch. Loss for query i:

$$ \mathcal{L}_i = -\log \frac{\exp(\text{sim}(q_i, p_i)/\tau)}{\sum_j \exp(\text{sim}(q_i, p_j)/\tau)} $$

τ is a temperature (typically 0.05). This is the same idea as word2vec's negative sampling but at the sentence level. Hard negatives (semantically close but irrelevant) are critical for high-quality retrievers.

1.3 Picking an embedding model

Always check the MTEB leaderboard (huggingface.co/spaces/mteb/leaderboard) for current SOTA in your domain.

2. Approximate Nearest Neighbor (ANN) Search

2.1 Why we need approximation

Exact NN: argmax_d cos(q, d) requires O(N) time. For 100M vectors at 1024 dims, that's ~400 GB of FLOPs per query. Unworkable.

Approximate methods trade a tiny recall@k drop for orders-of-magnitude speedup.

2.2 IVF — Inverted File

K-means cluster the corpus into nlist centroids; each vector belongs to one cluster. At query: find the nprobe nearest centroids, search only their members. Easy, fast, decent recall. Used in older FAISS.

2.3 HNSW — Hierarchical Navigable Small World (Malkov & Yashunin, 2018)

The dominant graph-based ANN. Build a multi-layer "small-world" graph; search starts at the top (sparse) layer and greedily descends. O(log N) query, very high recall. Widely used: FAISS, Qdrant, Vespa, Milvus, Pinecone.

Key parameters:

• M (typically 16–32): number of edges per node.
• ef_construction (200): candidates considered during build.
• ef_search (50–200): candidates considered during query. Bigger = higher recall, slower.

2.4 Product Quantization (PQ)

Compress vectors to ~8–16 bytes by splitting into subvectors and quantizing each independently with a small codebook. Combine with IVF for IVFPQ — billion-scale ANN on a single machine. Cost: small accuracy loss.

2.5 ScaNN, DiskANN, RaBitQ