Chapter 1 · Book pages 39–110

LLM Architecture and Optimization Methods

Main read ≈ 20 minutes · dotted words have hover/tap definitions · ▸ boxes go deeper · finish with the test (100% to complete)

The big picture

Everything an LLM does — chat, code, reasoning — is one loop: text is chopped into tokens, tokens become vectors of numbers, a stack of transformer layers processes them, the model outputs a probability for every possible next token, one gets picked, and the loop runs again with the new token appended. Watch it below.

Interactive · The autoregressive loop

Press play to watch a sentence being generated token by token: tokenize → process → predict → append → repeat.

This loop is the substrate for everything later in the book. RLHF (part II) trains the loop's choices; reasoning models (part III) run the loop longer to "think"; agents (part V) wrap the loop with tools and memory. If you understand this chapter, every later chapter has a foundation to attach to.

Stop 1 — Tokenization: how the model reads

The model never sees letters or words. It reads tokens: word pieces from a fixed menu (the vocabulary, 32K–128K entries). Common words are one piece; rare words get assembled — like LEGO for text. The menu is learned by BPE: start from characters, repeatedly merge the most frequent adjacent pair, stop when the menu is full.

BPE builds the token menu bottom-up by merging frequent character pairs (book p. 40).

Subword tokens are the compromise that works: characters make sequences too long (attention cost grows with the square of length), whole words can't handle new words. Special tokens like EOS carry structure — a model that can't emit EOS never stops generating.

Tokenization decides what the model can even perceive. Poor coverage of a language means more tokens per word ("fertility"), which means less effective context and higher cost for that language. Digit-by-digit tokenization is why some models suddenly got better at arithmetic. When agents later stream structured output, the special-token discipline from here is what keeps tool calls parseable.

Go deeper: other tokenizers, best practices

WordPiece (BERT) merges by likelihood instead of frequency; Unigram (SentencePiece/T5) starts big and prunes down; byte-level BPE (GPT-2 onward) works on raw bytes so no input is ever "unknown". Llama-3 uses a 128K vocabulary for better multilingual and code coverage, and tokenizes runs of spaces as single tokens so indented code stays cheap.

Practice rules from the book: never let a tokenizer split special tokens; mask loss on structural tokens during SFT; keep role markers as dedicated tokens rather than plain text. Book §1.2, p. 39–43.

Stop 2 — Embeddings: how pieces get meaning

Each token ID looks up a row in a big table: its embedding, a vector of ~4096 numbers. These coordinates place every token in a "meaning space" where similar concepts sit close together — and even arithmetic works: king − man + woman ≈ queen. Nobody designs this geometry; it emerges purely from next-token prediction, because words used in similar contexts end up with similar vectors.

The embedding table is a plain lookup matrix (Llama-3 8B: 128,256 rows × 4,096 numbers = 525M parameters, ~6.5% of the model). Many models reuse the same table at the output end ("weight tying"), so meaning-space geometry directly shapes output probabilities.

Embeddings are the bridge between symbols and math — and they come back everywhere: RAG (chapter 16) retrieves documents by comparing embedding vectors, and memory systems (chapter 17) store experiences as embeddings. The book's warning about anisotropy (raw embeddings clustering in a narrow cone, making every similarity score look high) is exactly why retrieval pipelines need contrastively-trained embedding models.

Go deeper: anisotropy and whitening

Pretrained embedding spaces are anisotropic: vectors occupy a narrow cone, so cosine similarity between any two texts comes out high regardless of content — which quietly breaks retrieval, clustering, and recommenders. A one-time linear fix called whitening (rotate and rescale so all directions have equal variance) restores meaningful similarity, and can shrink dimensionality at the same time. Alternatives: contrastive fine-tuning (SimCSE-style). Book §1.3.4, p. 46–48.

Stop 3 — Attention & the transformer: how it thinks

Attention lets each token pull in information from the tokens before it. Every token makes three vectors: a query ("what am I looking for?"), a key ("what do I contain?"), and a value ("what do I hand over if picked?"). Query·key scores go through softmax into weights, and the token's output is the weighted mix of values. The causal mask hides the future. Step through it:

Interactive · Attention, one query at a time

Step through how the word "it" finds what it refers to: query vs keys → scores → softmax weights → value mix.

A transformer layer is attention (tokens talk) followed by an FFN (each token thinks alone — research suggests the FFN acts as stored knowledge). Both steps are wrapped in residual connections and preceded by RMSNorm. Stack 32–126 of these layers and you have the model.

One pre-norm transformer block, the shape used by all modern LLMs (book p. 43–44, 53–54).

Attention cost grows with the square of sequence length — at a 128K context window the score matrix alone would be 64 GB. This single fact drives half of modern architecture research: GQA and MLA shrink the KV cache, Flash Attention reorganizes the compute, linear-attention hybrids replace it.

Word order isn't free: attention alone can't tell "the cat sat" from "sat cat the". RoPE injects position by rotating queries and keys so scores depend on relative distance — and RoPE scaling tricks (NTK, YaRN) are how models trained at 8K serve 128K.

Attention quirks show up later as product behavior: the "attention sink" (models dump spare attention on the first token — evict it and quality collapses) shapes streaming inference, and "lost in the middle" (long-context retrieval favors beginning and end) is why RAG pipelines put the best evidence at the edges of the prompt.

Go deeper: multi-head, GQA/MLA, positional encodings, pathologies, interpretability

Multi-head attention runs many small attentions in parallel; heads specialize (syntax, coreference, "induction heads" that copy patterns — the mechanism behind in-context learning). GQA shares one K/V head across several query heads (Llama-3: 32 Q, 8 KV → 4× smaller cache); MLA (DeepSeek) compresses all K/V into one small latent vector per token (~93% cache reduction).

Positional encoding history: fixed sinusoids (2017) → learned positions (GPT-2, hard length cap) → RoPE (relative, extendable) and ALiBi (distance penalty, great extrapolation). Long contexts also need continued pretraining on long documents — supporting a length positionally is not the same as using it.

Interpretability ladder: attention maps (cheap, least faithful — "attention is not explanation") → probing → sparse autoencoders (decompose activations into single-concept features, scaled to 34M features on Claude 3 Sonnet) → causal tracing (most faithful, priciest). Book §1.3.5–1.3.12, p. 48–59.

Stop 4 — Output & decoding: how it speaks

After the last layer, a projection turns each position's vector into logits — raw scores for every vocabulary token — and softmax makes them probabilities. Picking the actual token is decoding, and the strategy changes everything: greedy is precise but repetitive, temperature dials randomness, top-p adapts the candidate pool to the model's confidence. Try it:

Interactive · Sampling playground

Drag the temperature and top-p sliders to reshape a real next-token distribution, then sample from it.

Same model, same prompt — decoding sets the character of the output. Defaults that matter: temperature ≈ 0 for code and math, ~0.7 for factual work, ~1.0+ for creative text; top-p 0.9 is the everyday workhorse.

The same backbone serves different jobs by swapping the final head: language modeling (pretraining), response-only loss (SFT), a single-number value head (RL, part II), a reward head (RLHF). And when a program consumes the output, constrained decoding masks illegal tokens so the result is guaranteed-valid JSON — the backbone of reliable agent tool calls in part V.

Go deeper: beam search, min-p, contrastive and constrained decoding

Beam search tracks several candidate sequences (great for translation, bland for open-ended prose). Min-p keeps tokens at least X% as likely as the best token — scales with confidence even better than top-p. Contrastive decoding subtracts a weak model's scores from a strong model's, stripping the generic signal both share. Repetition penalties push down already-used tokens.

Constrained decoding compiles a JSON schema or grammar into a state machine; at every step, tokens that would break the grammar get probability zero. Libraries: Outlines, Guidance, XGrammar (used inside vLLM). Book §1.4 p. 59–62, §1.12 p. 88–94.

Stop 5 — Training: how it learns

Training is hill-descending in the dark. The loss measures how wrong the next-token predictions are (its human-readable cousin is perplexity); the gradient tells each weight which way to nudge; the learning rate sets the step size. Repeat over trillions of tokens. Plain SGD fails at LLM scale because different layers have wildly different gradient sizes — AdamW gives every weight its own adaptive step plus clean weight decay.

The learning rate is the most important hyperparameter, full stop. Diverging loss → too high; early plateau → too low. Standard recipe: warmup ramp, then cosine or WSD decay, with gradient clipping at norm 1.0 as the seatbelt. Everything runs in BF16 mixed precision with FP32 master weights.

Pretraining is just next-token prediction over 1–15 trillion tokens (Llama-3: 15T; ~80% web, ~10% code). Reasoning, in-context learning, instruction-following — all emergent from that one objective. Chinchilla scaling laws balance model size and data; in practice models are over-trained because small-and-well-trained is cheaper to serve forever.

"Mid-training" — a newer stage that up-weights math/code and stretches context before any RL — determines whether reinforcement learning works at all later: the same RL recipe boosted Qwen (reasoning-dense mid-training) and did nothing for Llama. Part II of the book stands on this stop.

Go deeper: Adam mechanics, Muon, precision frontier, failure modes

Adam keeps two running averages per weight: recent gradient direction (momentum) and magnitude (for the per-weight step). AdamW applies weight decay outside the adaptive scaling — always AdamW, never plain Adam. Muon (2025) orthogonalizes momentum for ~2× compute efficiency; adopted by GLM, DeepSeek-V4, and Kimi K2 (15.5T tokens, zero loss spikes) while AdamW stays the fine-tuning default.

Precision frontier: BF16 beats FP16 (same range as FP32, no loss-scaling gymnastics); DeepSeek-V3 trained 671B mostly in FP8 for ≈$5.6M; NVFP4 (4-bit) training already works with ~15% of layers kept in BF16. Pretraining pitfalls: loss spikes (roll back, skip the batch), memorization (deduplicate), short-context training that fails at long context. Book §1.5 p. 63–72, §1.7 p. 76–78.

Stop 6 — Fine-tuning & LoRA: becoming an assistant

A pretrained model completes text; it doesn't follow instructions. SFT fixes that with curated prompt→response pairs — mechanically identical to pretraining except loss is computed only on response tokens. That single masking change turns a text predictor into an assistant. Quality beats quantity: 1,000 excellent examples can match 50,000 noisy ones (the LIMA principle).

Full fine-tuning of a 70B model needs ~560 GB of GPU memory. LoRA sidesteps it: freeze the weights, learn a thin low-rank correction W′ = W + (α/r)·B·A. Watch what the rank buys:

Interactive · LoRA visualizer

Drag the rank slider and watch the adapter matrices grow while the trainable-parameter count stays a rounding error next to full fine-tuning.

Why LoRA works: fine-tuning changes live in a tiny subspace of parameter space ("intrinsic dimensionality"). Defaults: rank 16, α = 2r, all attention projections, LR ~10× higher than full fine-tuning. QLoRA quantizes the frozen base to 4-bit — a 70B fine-tune fits one 48 GB GPU.

SFT teaches format, not preference, honesty calibration, or deep reasoning — those need RL on top (chapters 4–13). The pipeline is always pretrain → SFT → RLHF/DPO. And LoRA's merge-back-with-zero-overhead property is why one base model can serve hundreds of customer-specific adapters in production.

Go deeper: LoRA variants and efficiency stacks

DoRA splits weight direction from magnitude and adapts only direction (+1–3% on reasoning). LoRA+ trains B with ~16× the LR of A (free ~2%). rsLoRA stabilizes high ranks (≥64). VeRA freezes shared random A/B and trains only tiny scaling vectors (~10× fewer parameters). Training stacks: Unsloth (fastest single-GPU QLoRA), Liger Kernel (fused Triton kernels, +20% throughput), torchtune (hackable pure PyTorch). Book §1.8–1.9, p. 78–84.

Stop 7 — Mixture of Experts: capacity without the bill

A MoE layer replaces one big FFN with many parallel "expert" FFNs plus a router that picks the top-K per token — a hospital with specialists instead of one generalist brain. Total knowledge is huge; per-token compute stays small. Mixtral 8x7B: 47B total, 13B active, matching a dense 70B.

Top-2 routing: only the chosen experts run; the rest are skipped entirely (book p. 85).

Two hard problems define MoE engineering: load balancing (left alone, the router overuses a few experts — "expert collapse") and trainability (a hard "pick top-2" has no gradient, so an expert never picked never learns; noise injection during training gives underdogs a chance).

MoE is how frontier models grew 10× in capacity without 10× serving cost — DeepSeek-V3: 671B total / 37B active (book p. 54); and, outside the book, Kimi K3 (2026): 2.8T total with 16-of-896 routed experts. When later chapters mention "active vs total parameters", this stop is the decoder ring.

Go deeper: balancing losses, noisy top-K, notable models

The classic fix adds an auxiliary loss pushing routing toward uniform use; DeepSeek showed the loss fights specialization and replaced it with a per-expert bias adjusted outside the gradient. Noisy top-K gating adds learned Gaussian noise before selection (removed at inference); Gumbel-softmax is the differentiable alternative. Aggregating balance statistics over the global batch (not micro-batches) is what preserves expert specialization. Book §1.10, p. 85–87.

Stop 8 — Running fast & small

Serving is where architecture meets physics. A GPU has a big slow warehouse (HBM) and a tiny fast desk (SRAM) — and standard attention wastes its time walking between them. Flash Attention computes the exact same result in desk-sized tiles using a running-softmax trick, never materializing the giant score matrix: identical math, unchanged FLOP count, 2–4× faster, memory O(n²) → O(n).

Flash Attention: tile the computation to fit fast memory (book p. 73–74).

For the model itself: quantization stores weights in 4–8 bits (70B: 140 GB → 35 GB at ~2–3% quality cost — the book's rule: always quantize for serving, W4A16 is the sweet spot). Distillation trains a small student on a big teacher's full probability distribution — "dark knowledge": Paris 90%, Lyon 5%, banana ~0% teaches which mistakes are reasonable.

And speculative decoding attacks generation latency itself: a small drafter proposes ~5 tokens, the big model verifies them in one parallel pass, and a clever acceptance rule keeps the output distribution exactly the big model's. Pure speedup, zero quality change.

Speculative decoding: draft cheap, verify in parallel, keep the exact target distribution (book p. 104).

Attention at decode time is memory-bound, not compute-bound. Once you see serving as "move fewer bytes" — Flash Attention, KV-cache shrinking, quantization, speculative drafts — every inference optimization in the book clicks into one picture.

These are the levers that set the price-per-token you pay for any API — and chapter 2 (GPUs, vLLM, PagedAttention) continues exactly here. Flash Attention 4's hardware story traces back to this stop — as does the 4-bit quantization-aware training in recent frontier models (the book's NVFP4 example, or Kimi K3's MXFP4, 2026, outside the book).

Go deeper: FA versions, quantization methods, pruning, distillation modes

Flash Attention by GPU generation: FA2 (A100 — better parallelism, skips masked blocks), FA3 (H100 — async memory engine, warp specialization, 75% of peak), FA4 (B200 — the exponential became the bottleneck, so it's emulated with polynomials on tensor cores). Lesson: each hardware generation moves the bottleneck and demands new algorithms.

Quantization named methods: GPTQ, AWQ (protects the ~1% of weights that matter most), GGUF (llama.cpp local formats), SmoothQuant (W8A8), QAT (highest quality, expensive). Pruning: 2:4 structured sparsity gets exact 2× on NVIDIA tensor cores; SparseGPT/Wanda reach 50% one-shot. Distillation modes: offline (stored logits), online (co-training), black-box (API text only — degrades to imitation SFT), self-distillation. Speculative variants: Medusa (extra heads), Eagle (drafts from hidden states, 85–95% acceptance), n-gram lookup (free, great for code). Book §1.6 p. 73–76, §1.14–1.15 p. 101–107.

Stop 9 — Failures & safety

A hallucination is confident, fluent, wrong output — contradicting the given context, inventing facts, or ignoring instructions. Model-level detectors look for uncertainty: high entropy on output tokens, or self-disagreement (ask several times and compare meanings; correct knowledge is stable, confabulation varies). The honest caveat: these detect uncertainty, not incorrectness — a model can be confidently and consistently wrong. Reliable factuality needs retrieval (chapter 16).

Safety is layered through the whole pipeline, never bolted on: filtered pretraining data → refusal examples in SFT → a dedicated safety reward model in RLHF → red-teaming → serving-time guardrails. The central tension is helpfulness vs over-refusal, and it is never "done" — new attacks appear continuously.

For agents this stop is existential: an autonomous system that hallucinates tool results or gets jailbroken mid-task fails at machine speed. Everything in part V (guardrails, evaluation, human-in-the-loop UI) builds on the failure taxonomy defined here.

Go deeper: detection methods and the safety pipeline

Semantic entropy clusters sampled answers by meaning and measures disagreement; SelfCheckGPT verifies claims against the model's own alternative generations; DoLA contrasts late-layer vs early-layer logits (factual knowledge lives deeper) at decode time. Threat taxonomy: harmful content, bias, privacy leaks, jailbreaks, misinformation, dual-use. Helpfulness–safety balancing is a constrained optimization with separate reward models; over-refusal target <5%. Book §1.16–1.17, p. 108–110.