About This Edition

This mdBook edition is a derivative work adapted from Haggai Roitman's The Hitchhiker's Guide to Agentic AI: From Foundations to Systems (arXiv:2606.24937). Changes to the original work include conversion to mdBook format and adjustments to typography and layout.

The original work is licensed under the Creative Commons Attribution-ShareAlike 4.0 International License (CC BY-SA 4.0). This mdBook edition is made available under the same CC BY-SA 4.0 license.

Dedication from the Original Work

To my beloved wife Janna and daughters Inbar and Einav.

Disclaimer

This document is an independent survey and educational resource prepared solely by the author. The views, opinions, and conclusions expressed herein are those of the author alone and do not necessarily reflect the views of any employer, organization, or institution with which the author is or has been affiliated.

This work does not contain any proprietary, confidential, or trade-secret information. All referenced material is drawn from publicly available sources, including peer-reviewed publications, open-access preprints, official documentation, and open-source repositories.

The content is provided “as is” without warranty of any kind, express or implied. The author makes no representations regarding the accuracy, completeness, or suitability of the information for any particular purpose. Readers should independently verify any claims, formulas, or implementation details before applying them in production systems.

Feedback welcome. If you find any mistakes, inaccuracies, or have suggestions for improvement, you are encouraged to send feedback to [first_name]r@gmail.com.

AI Disclosure. Large language models were used as a research and drafting aid. All content was edited and verified by the author on a best-effort basis.

License. This work is licensed under the Creative Commons Attribution-ShareAlike 4.0 International License (CC BY-SA 4.0). You are free to share (copy and redistribute) and adapt (remix, transform, and build upon) this material for any purpose, including commercially, provided you give appropriate credit, provide a link to the license, indicate if changes were made, and distribute any derivative works under the same license. Full license text: https://creativecommons.org/licenses/by-sa/4.0/.

About the Author

Haggai Roitman has spent over two decades at the intersection of AI research and large-scale production systems. His work bridges theory and practice—from publishing foundational research to shipping systems that serve millions of users.

His research interests span information retrieval, recommender systems, natural language processing, large language models, reinforcement learning for LLMs, and agentic AI. He has authored more than 100 peer-reviewed publications and holds approximately 100 patents. He earned his BSc (Cum Laude) and PhD from the Technion — Israel Institute of Technology.

When not thinking about gradient flows and KV caches, he can be found behind a set of turntables, mixing progressive trance and deep house.

Preface

Why This Guide Exists

Building intelligent AI systems in 2026 requires mastering an extraordinary breadth of knowledge — from how transformers process language internally, through the hardware and systems that make training possible, the optimization techniques that make it efficient, the reinforcement learning algorithms that teach models to reason and align with human intent, all the way to multi-agent architectures that coordinate autonomous systems at scale.

This knowledge is scattered across hundreds of papers, blog posts, GitHub repositories, and tribal knowledge within a handful of labs. This guide exists because practitioners need a single, unified reference that covers the entire stack — not just the theory, but the implementation details that make things actually work.

A Personal Journey to Agentic AI

My fascination with intelligent agents began two decades ago, when I still studied for my first degree in information systems engineering. I took courses on Agent-Oriented Software Engineering (AOSE) (Wooldridge et al. 2000) and learned to build multi-agent systems using JADE (Bellifemine et al. 2007) (Java Agent DEvelopment Framework)—a FIPA-compliant (Intelligent Physical Agents 2002) platform where agents communicated via structured protocols, negotiated over shared resources, and coordinated autonomously. Around the same time, Berners-Lee, Hendler, and Lassila’s seminal paper “The Semantic Web” (Berners-Lee et al. 2001) painted a vision of machine-readable knowledge that agents could reason over. These two threads—autonomous agent architectures and semantic knowledge representation—planted a seed that has guided my career ever since. One early project that crystallized this vision was an attempt to build a shopping agent—developed with OntoBuilder (Gal et al. 2005) under the guidance of my respected future academic advisor Prof. Avigdor Gal—a system that could automatically fill product search queries and orders across different heterogeneous websites, understanding their varied schemas through ontology matching and mapping. The Semantic Web promised that such agents would thrive in a world of structured, machine-readable data. In practice, the brittleness of hand-crafted ontologies, the messiness of real-world product data, and the lack of robust natural language understanding made the vision perpetually “five years away.”

Over the following years, I tracked each wave of AI progress as it arrived: neural networks and heuristic search for combinatorial optimization; deep learning and representation learning; information retrieval and personalization at scale; and most recently, the revolution of large language models. Each wave brought powerful new tools, but the dream remained the same: systems that understand, reason, and act autonomously in complex environments.

What makes 2024–2026 remarkable is that these threads have finally converged. LLMs provide the language understanding and generation; reinforcement learning teaches them to reason and align with human intent; tool-use protocols (MCP) give them hands to act in the world; and agent orchestration frameworks provide the coordination layer that JADE envisioned twenty years ago—now powered by foundation models instead of hand-coded ontologies. This guide is, in many ways, the reference I wish I had at each step of that journey.

The Landscape in 2026

The journey to today’s agentic AI systems spans three decades of compounding breakthroughs across architecture, training, and deployment:

1. Architectural foundations (2017–2020): The Transformer (Vaswani et al. 2017) introduced self-attention as a universal sequence-processing primitive. Scaling laws revealed that larger models trained on more data reliably improve. GPT-2 and GPT-3 demonstrated that decoder-only transformers, scaled sufficiently, become capable few-shot learners.

2. Systems and efficiency (2020–2023): Flash Attention (Dao et al. 2022) made training 2–4\(\times\) faster by eliminating memory bottlenecks. LoRA (Hu et al. 2022a) enabled fine-tuning 70B+ models on a single node. Mixture-of-Experts (MoE) decoupled model capacity from compute cost. Inference engines like vLLM brought throughput within reach of real-time applications.

3. Alignment via RL (2022–2024): RLHF (Ouyang et al. 2022) transformed capable-but-unhelpful base models into useful assistants — the recipe behind ChatGPT. DPO (Rafailov et al. 2023) collapsed the reward model and RL loop into a single supervised loss, democratizing alignment. Variants proliferated: KTO (Ethayarajh et al. 2024), IPO (Azar et al. 2024), ORPO (J. Hong et al. 2024), GRPO (Shao et al. 2024).

4. Reasoning and autonomy (2024–2026): DeepSeek-R1 (DeepSeek-AI et al. 2025) and OpenAI’s o1/o3 demonstrated that RL could teach reasoning itself — models spontaneously discover chain-of-thought, backtracking, and self-verification. Simultaneously, the Model Context Protocol (MCP) standardized tool access, Agent-to-Agent (A2A) enabled inter-agent communication, and production-grade orchestration frameworks matured.

Who This Guide Is For

This document is written for practitioners who build things:

• ML engineers who need to understand transformer internals, training infrastructure, optimization, and why things diverge.

• Applied researchers evaluating architectures, fine-tuning strategies, and RL methods for their specific domains.

• Agent developers building production systems who need orchestration patterns, memory architectures, tool integration (MCP), and multi-agent coordination (A2A).

• Systems engineers responsible for training infrastructure, GPU clusters, distributed training, and inference deployment.

• Technical leaders making architectural and resourcing decisions across the full stack.

We assume familiarity with neural networks and basic probability. No prior LLM, RL, or systems knowledge is required — the guide builds from first principles.

What You Will Gain

After reading this guide, you will be able to:

• Understand LLM internals — attention mechanisms, positional encodings, MoE routing, Flash Attention, and why architectural choices matter for downstream capability.

• Reason about systems — GPU memory budgets, distributed training strategies (FSDP, tensor/pipeline parallelism), inference optimization, and production deployment with vLLM.

• Train and fine-tune efficiently — LoRA/QLoRA, quantization, knowledge distillation, optimizer selection, and learning rate scheduling.

• Align models with human preferences — implement RLHF/DPO/GRPO/KTO pipelines, debug reward hacking and mode collapse, choose the right algorithm among 20+ methods.

• Build reasoning models — understand how DeepSeek-R1, o1/o3, and QwQ discover chain-of-thought through RL without explicit demonstrations.

• Architect agentic systems — select orchestration patterns, design memory, integrate tools via MCP, coordinate agents via A2A, evaluate with production benchmarks.

• Evaluate rigorously — apply appropriate metrics, benchmarks, and LLM-as-Judge patterns for both model quality and agent capability.

How This Guide Is Organized

The guide spans 29 chapters organized in five parts:

1. Part I — Foundations (Chapters 1–3): LLM architecture and optimization (transformers, attention, positional encodings, Flash Attention, LoRA, MoE), systems foundations (GPU hierarchies, distributed training, vLLM), and classical RL theory (MDPs, policy gradients, actor-critic).

2. Part II — RL Methods for LLMs (Chapters 4–12): The complete RL-for-LLMs toolkit. RL foundations for language models, then full mathematical treatment of PPO, DPO, GRPO, and preference optimization variants (Online DPO, KTO, IPO, ORPO, SimPO), reward model training, SFT best practices, system architecture at scale, and agentic training with trajectory-level RL.

3. Part III — Reasoning (Chapter 13): Large reasoning models — DeepSeek-R1, OpenAI o1/o3/o4-mini, QwQ — how RL discovers chain-of-thought, MCTS, process reward models, and test-time compute scaling.

4. Part IV — Evaluation (Chapter 14): Comprehensive LLM evaluation methodology — metrics, LLM-as-Judge, human annotation, benchmark suites, contamination detection, and agentic evaluation.

5. Part V — Agentic AI (Chapters 15–26): The complete agentic stack — introduction to agentic AI, RAG and retrieval, memory systems, orchestration and context management, design patterns, agentic environments and benchmarks, Model Context Protocol (MCP), agent skills, Agent-to-Agent communication (A2A), multi-agent systems, development frameworks, and agentic UI.

6. Part VI — Assessment & Reference (Chapters 27–29): 108 detailed quiz questions with comprehensive answers spanning all topics, a quick-reference chapter consolidating key equations, API references, and failure mode diagnostics, and a conclusion with future directions.

The guide includes over 100 detailed quiz questions with comprehensive answers spanning all topics, plus a quick-reference chapter consolidating key equations, API references, and failure mode diagnostics.

Design Philosophy

Three principles guide this document:

1. Intuition first, formalism second. Every equation is preceded by a plain-English explanation of what it means and why it matters.

2. Implementation-aware. Theory is useless without knowing how to make it work. We include code, hyperparameter tables, memory budgets, architecture diagrams, and debugging strategies throughout.

3. Honest about what works. We clearly state which approaches are production-tested and which are research explorations.

Scope and Deliberate Omissions

This guide focuses on text-in, text-out language models and the RL, systems, and agentic infrastructure built around them. Several important areas are intentionally excluded:

• Multimodal models (vision–language, audio, video). Multimodal architectures introduce distinct training pipelines (contrastive pre-training, cross-modal alignment, modality-specific encoders), data curation challenges, and evaluation protocols that each merit book-length treatment. Including them would double the scope without deepening the RL and agentic core that unifies this guide.

• Domain-specific deployments (healthcare, legal, finance, scientific discovery). Domain adaptation introduces regulatory constraints, specialized evaluation, and data-access issues that are orthogonal to the general methods presented here. The algorithms and architectures we cover are the building blocks practitioners adapt to these domains, but the adaptation details are better served by dedicated references.

• Personalization and recommendation systems. Personalization relies on user modeling, collaborative filtering, and interaction-history architectures that form a parallel research tradition. While LLMs are increasingly used within recommender systems, the core techniques (sequential models, bandit-based exploration, cold-start handling) are sufficiently distinct to warrant separate coverage.

By maintaining this boundary, we keep a single coherent thread—from architectural foundations and systems infrastructure, through the training algorithms that produce aligned and reasoning models, to the orchestration and deployment of autonomous agents—without fragmenting the narrative across modalities and verticals.

— Haggai Roitman, 2026

Introduction

The Big Picture

This guide takes you from first principles to production systems. It is written for practitioners — researchers, engineers, and applied scientists — who want to understand and build the full stack of modern AI: from transformer architectures and the hardware that runs them, through the training algorithms that align models with human intent and teach them to reason, to the agentic architectures that deploy them as autonomous systems.

The core thesis is simple: building great AI systems requires understanding the entire pipeline, not just one layer. An engineer debugging a training run needs to understand GPU memory hierarchies and optimizer dynamics. A fine-tuning practitioner needs to know when LoRA suffices and when full-parameter training is worth the cost. An agent developer needs to understand how the underlying model was trained. A technical leader evaluating frameworks needs to understand what trade-offs each one makes. This guide provides that complete picture.

The Road to Agentic AI: A Brief History

Today’s agentic AI systems did not emerge in a vacuum. They stand on decades of milestone systems — each solving a narrower problem but collectively building the techniques, hardware, and ambition that made autonomous agents possible.

1. Deep Blue (1997) (Campbell et al. 2002) — IBM’s chess engine defeated world champion Garry Kasparov using brute-force search (200 million positions/second) with handcrafted evaluation functions. It proved machines could exceed human performance in well-defined adversarial domains, but generalized to nothing else.

2. IBM Watson — Jeopardy! (2011) (Ferrucci et al. 2010) — Watson combined information retrieval, NLP, and massive parallelism to defeat human champions at open-domain question answering. It demonstrated that AI could process unstructured text at scale, but required years of domain-specific engineering and couldn’t learn new domains without substantial human effort.

3. AlexNet and the Deep Learning Revolution (2012) (Krizhevsky et al. 2012) — Krizhevsky et al.’s CNN won ImageNet by a stunning margin, proving that deep neural networks trained on GPUs could learn representations from raw data. This single result triggered the modern deep learning era and the hardware investment that eventually made LLMs possible.

4. AlphaGo (2016) (Silver et al. 2016) — DeepMind’s system defeated Go world champion Lee Sedol using deep RL (policy networks + value networks + Monte Carlo Tree Search). Unlike Deep Blue’s brute force, AlphaGo learned to play — demonstrating that RL could master domains where search alone was intractable (\(10^{170}\) board states). AlphaGo Zero (2017) (Silver et al. 2017) later learned entirely from self-play, needing no human games at all.

5. GPT-2/GPT-3 (2019–2020) (Brown et al. 2020) — OpenAI showed that scaling decoder-only transformers to billions of parameters produced emergent few-shot learning. GPT-3 (175B parameters) could perform tasks it was never explicitly trained for — translation, arithmetic, code generation — simply from in-context examples. The era of foundation models began.

6. AlphaFold (2020) (Jumper et al. 2021) — DeepMind solved the 50-year protein folding problem, predicting 3D protein structures with atomic accuracy. AlphaFold demonstrated that deep learning could crack fundamental scientific problems previously considered decades away. It also showcased the power of architecture innovation (attention over residue pairs) combined with massive compute.

7. ChatGPT and RLHF (2022) (Ouyang et al. 2022) — InstructGPT/ChatGPT proved that a capable base model, when aligned via RLHF, becomes a genuinely useful assistant. This was the inflection point: AI went from a research tool to a consumer product used by hundreds of millions. The alignment techniques (reward models, PPO) became the template for all subsequent LLM post-training.

8. GPT-4 and Multimodal Models (2023) (OpenAI 2023) — Multimodal capabilities (vision + language), longer contexts, and improved reasoning pushed LLMs toward general-purpose cognition. Tool use (code interpreter, web browsing) hinted at agentic capabilities.

9. Reasoning Models (2024) (DeepSeek-AI et al. 2025) — OpenAI’s o1 and DeepSeek-R1 showed that RL could teach models to reason: chain-of-thought, backtracking, self-verification emerged spontaneously from reward signals alone. Models began solving competition-level mathematics and complex coding tasks.

10. Agentic AI (2025–present) — The convergence point: LLMs with reasoning capabilities, equipped with standardized tool access (MCP), inter-agent communication (A2A), persistent memory, and sophisticated orchestration frameworks. Agents now autonomously write code, conduct research, manage workflows, and coordinate with other agents — the subject of this guide.

Tip

Each milestone shares a common arc: architecture innovation \(+\) scale \(+\) learning signal \(=\) breakthrough. Deep Blue used handcrafted search. AlphaGo learned from self-play. GPT-3 learned from internet text. Today’s agentic systems learn from human feedback, verifiable rewards, and environment interaction. The learning signal has expanded from game outcomes to open-ended human preferences — and the architectures have grown to match.

This guide picks up the story at the foundation model era and carries it forward through alignment, reasoning, and autonomous agency.

What You Should Expect

Part I: Foundations (Chapters 1–3) builds the base knowledge the rest of the guide depends on. We start with how LLMs work internally — the architecture decisions that determine capability — then cover the hardware and systems that make training and inference possible, and finally introduce reinforcement learning from first principles.

• Chapter 1 — LLM Architecture and Optimization: Transformer internals (self-attention, multi-head attention, RoPE, GQA), Flash Attention, optimization methods (AdamW, learning rate schedules, gradient clipping), mixed precision, LoRA/QLoRA, quantization, knowledge distillation, and Mixture of Experts.

• Chapter 2 — Systems Foundations: GPU architecture (A100/H100/B200), memory hierarchies, NVLink/NVSwitch, distributed training (FSDP, DeepSpeed ZeRO, tensor/pipeline parallelism), and vLLM for high-throughput inference.

• Chapter 3 — Introduction to RL: MDPs, Bellman equations, TD learning, Q-learning, policy gradients (REINFORCE), actor-critic methods, GAE — the algorithmic toolkit that underpins Part II.

Part II: RL Methods for LLMs (Chapters 4–12) is the training and alignment core. Here you learn how to align, improve, and fine-tune language models — from full mathematical derivations to working code.

• Chapters 4–8: Every major RL/preference algorithm with math, intuition, and TRL code — PPO, DPO, GRPO, and preference optimization variants (Online DPO, KTO, IPO, ORPO, SimPO, Best-of-N).

• Chapters 9–10: Reward model training (Bradley–Terry, scaling laws, reward hacking) and SFT best practices (data quality, curriculum, formatting).

• Chapters 11–12: System architecture at scale (decoupled training, fault tolerance, GPU allocation) and LLM agentic training — how to train agents end-to-end with trajectory-level RL.

Part III: Reasoning (Chapter 13) covers the frontier of model capability — teaching LLMs to reason through multi-step problems.

• Chapter 13 — RL for Large Reasoning Models: DeepSeek-R1, OpenAI o1/o3/o4-mini, QwQ — how RL discovers chain-of-thought, MCTS, process reward models, and test-time compute scaling.

Part IV: Evaluation (Chapter 14) provides the methodology for measuring whether any of this actually works.

• Chapter 14 — LLM Evaluation: Metrics (perplexity, pass@k, ELO), LLM-as-Judge patterns, contamination detection, benchmark suites, and agentic evaluation methodology.

Part V: Agentic AI (Chapters 15–26) takes you from a trained model to a deployed autonomous system. This is the largest part, covering everything an agent needs to operate in the real world.

• Chapter 15 — Introduction to Agentic AI: What makes a system agentic, the spectrum from chatbots to autonomous agents, and the foundational concepts for the rest of Part V.

• Chapter 16 — RAG: Retrieval methods, chunking, embedding models, hybrid search, reranking, and production architectures.

• Chapter 17 — Memory Systems: Working, episodic, semantic, and procedural memory for persistent agent knowledge.

• Chapter 18 — Orchestration: ReAct, Plan-and-Execute, LLM Compiler, reflexion patterns, context management, and harness design.

• Chapter 19 — Design Patterns: Prompt chaining, routing, parallelization, evaluation-driven orchestration, and the simplicity principle.

• Chapter 20 — Environments and Benchmarks: WebArena, SWE-bench, OSWorld, GAIA — evaluation environments for agentic capability.

• Chapter 21 — Model Context Protocol (MCP): Architecture, transport layers, tool/resource/prompt primitives, security, and deployment.

• Chapter 22 — Agent Skills: Skill libraries, tool composition, and capability abstraction.

• Chapter 23 — A2A Communication: Google’s Agent-to-Agent protocol — Agent Cards, task lifecycle, streaming, enterprise patterns.

• Chapter 24 — Multi-Agent Systems: Hierarchical, debate, marketplace, and swarm architectures — coordination at scale.

• Chapter 25 — Development Frameworks: LangGraph, CrewAI, AutoGen, OpenAI Agents SDK, Google ADK — comparative analysis with code.

• Chapter 26 — Agentic UI: Streaming interfaces, generative UI, canvas paradigms, tool visualization, human-in-the-loop patterns.

The Modern AI Pipeline

The full pipeline from base model to deployed agent:

LLM Architecture and Optimization Methods

This section covers the foundational architecture of large language models and the key optimization techniques that make training and inference efficient. Topics are ordered as a curriculum: we begin with the transformer itself, then cover how to train it efficiently, how to adapt it cheaply, how to compress it, how to scale it, and how to accelerate its inference.

How LLMs Work: An Intuitive Overview

Before diving into architectural details, let us build intuition for how a large language model transforms text into text. The entire process follows a simple pipeline: text \(\to\) tokens \(\to\) representations \(\to\) tokens \(\to\) text.

Important

The Four Key Stages

1. Tokenization: Raw text is split into subword pieces (not characters, not full words) using a learned vocabulary. “unhappiness” might become [“un”, “happiness”] or [“unhapp”, “iness”].

2. Embedding: Each token ID indexes into a learned embedding table, producing a dense vector in \(\mathbb{R}^d\) (typically \(d = 4096\)). These vectors capture semantic meaning—similar words get similar vectors.

3. Contextual Processing: The transformer stack processes all embeddings in parallel, using self-attention to let each position “read” from all other positions. After \(L\) layers, each position’s hidden state encodes rich contextual information.

4. Prediction: The final hidden state is projected to a probability distribution over the full vocabulary, and a decoding strategy selects the next token.

Tokenization

Tokenization is the critical first step that converts raw text into the discrete symbols a language model operates on. The choice of tokenizer directly affects model quality, multilingual capability, and computational efficiency.

Tip

Why Subwords?

Character-level models need very long sequences (expensive attention). Word-level models cannot handle rare or novel words. Subword tokenization strikes the ideal balance: common words are single tokens (“the” \(\to\) [the]), rare words decompose into known pieces (“cryptocurrency” \(\to\) [“crypt”, “ocur”, “rency”]), and the vocabulary stays manageable (32K–128K tokens).

Why Not Characters or Words?

Trade-offs of different tokenization granularities.

Byte-Pair Encoding (BPE)

BPE (Sennrich et al. 2016) is the dominant tokenization algorithm used by GPT, Llama, Mistral, and most modern LLMs.

Important

BPE Algorithm

1. Start with a vocabulary of individual characters (bytes)

2. Count all adjacent symbol pairs in the training corpus

3. Merge the most frequent pair into a new symbol

4. Repeat steps 2–3 for \(k\) iterations (until desired vocabulary size)

Other Tokenization Methods

Comparison of subword tokenization algorithms.

Tokenization Best Practices

1. Vocabulary size matters: 32K is minimal; 128K enables better multilingual coverage and code handling. Llama-3 uses 128K tokens.

2. Special tokens: Always include <bos>, <eos>, <pad>, <unk>. For instruction-tuned models, add role markers (<|user|>, <|assistant|>).

3. Fertility: Measure tokens-per-word across languages. High fertility (many tokens per word) indicates poor coverage for that language.

4. Never tokenize across boundaries: Spaces, punctuation, and digits should be handled consistently. Most modern tokenizers prepend a space marker (“the”) to distinguish word-initial vs. continuation tokens.

5. Numbers: Consider digit-level tokenization for arithmetic tasks. “2024” as [“2”,“0”,“2”,“4”] enables digit-by-digit reasoning.

6. Code: Ensure whitespace (indentation) is tokenized efficiently. Llama-3 tokenizes runs of spaces as single tokens.

Tokenization in Practice: HuggingFace Example

The transformers library provides a unified interface for all tokenizers. The following demonstrates encoding and decoding with a modern LLM tokenizer:

from transformers import AutoTokenizer

# Load Llama-3 tokenizer (128K vocabulary, byte-level BPE)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B")

text = "Reinforcement learning optimizes long-term rewards."

# Encode: text -> token IDs
token_ids = tokenizer.encode(text)
print(token_ids)
# [128000, 29934, 262, 11008, 4815, 6900, 1317, 9860, 21845, 13]

# Decode individual tokens to see subword splits
tokens = tokenizer.convert_ids_to_tokens(token_ids)
print(tokens)
# ['<|begin_of_text|>', 'Re', 'inforce', 'ment', ' learning',
#  ' optimizes', ' long', '-term', ' rewards', '.']

# Decode back to text (round-trip)
reconstructed = tokenizer.decode(token_ids, skip_special_tokens=True)
assert reconstructed == text  # Perfect reconstruction

# Tokenize with attention mask (for batched inputs with padding)
batch = tokenizer(
    ["Short text.", "A much longer input sentence for comparison."],
    padding=True, return_tensors="pt"
)
print(batch.keys())  # dict_keys(['input_ids', 'attention_mask'])

Special Tokens and Structured Prompts

Special tokens are reserved vocabulary entries that carry structural meaning rather than linguistic content. They are critical for controlling model behavior.

Common special tokens across LLM families.

Role Markers for Instruction-Tuned Models.

Modern chat models use special tokens to delineate conversational structure. These are not trained to carry semantic meaning—they are structural delimiters that the model learns to parse:

# Llama-3 chat template
messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Explain PPO in one sentence."},
]

# apply_chat_template handles all special token insertion
prompt = tokenizer.apply_chat_template(messages, tokenize=False)
print(prompt)
# <|begin_of_text|><|start_header_id|>system<|end_header_id|>
#
# You are a helpful assistant.<|eot_id|><|start_header_id|>user<|end_header_id|>
#
# Explain PPO in one sentence.<|eot_id|><|start_header_id|>assistant<|end_header_id|>
#
#

Important

Special Token Best Practices

• Never split special tokens: They must be atomic—ensure your tokenizer treats them as single units, not character sequences.

• Mask loss on special tokens: During SFT, do not compute loss on structural tokens (role markers, separators). The model should not “learn” to predict formatting.

• Use templates for structure: Encode task semantics via special tokens rather than natural language instructions. E.g., <|tool_call|> is more reliable than “Now I will call a tool:”.

• Tool/function calling: Define dedicated tokens like <|function|>, <|result|> to create unambiguous boundaries between reasoning and action.

• Consistent handling in RL: During PPO/GRPO, ensure the reference model and policy model use identical tokenization and special token handling—mismatches corrupt KL computation.

• EOS handling: During generation, ensure EOS is included in the action space. If the model cannot emit EOS, responses grow unbounded (common RL failure mode).

The Transformer Architecture

The Transformer (Vaswani et al. 2017) is the foundation of all modern LLMs. Understanding its components is essential for grasping every optimization and training method in this guide.

High-Level Structure

A decoder-only transformer processes tokens sequentially through embedding, repeated attention+FFN blocks, and a final projection to vocabulary logits. Figure 1.1 shows the complete architecture.

The Original Encoder-Decoder Transformer

The Transformer was originally introduced (Vaswani et al. 2017) as an encoder-decoder architecture for sequence-to-sequence tasks (machine translation, summarization). While modern LLMs predominantly use decoder-only variants (GPT-style), understanding the full architecture is essential because cross-attention and masked self-attention — both originating here — remain fundamental building blocks.

Encoder.

The encoder processes the entire input sequence bidirectionally — each token attends to all other tokens (no causal mask). This produces a rich contextual representation \(\mathbf{H}^{\text{enc}} \in \mathbb{R}^{n \times d}\) where each position encodes information about the full input:

• Input: Token embeddings + sinusoidal positional encodings

• Each layer: Multi-Head Self-Attention \(\to\) Add & Norm \(\to\) FFN \(\to\) Add & Norm

• No causal mask: Position \(i\) attends to all positions \(1, \ldots, n\)

• Output: Contextual representations of the full input sequence

Decoder — Masked Multi-Head Self-Attention.

The decoder generates output tokens one at a time (autoregressively). To prevent the model from “seeing the future,” the self-attention in the decoder uses a causal mask:

\[ \text{MaskedAttn}(Q, K, V) = \text{softmax}!\left(\frac{QK^T}{\sqrt{d_k}} + M\right) V \]

where the mask \(M\) is:

\[ M_{ij} = \begin{cases} 0 & \text{if } i \geq j \text{ (can attend)} \\ -\infty & \text{if } i < j \text{ (future token --- blocked)} \end{cases} \]

Tip

Why Masking Matters

During training, the decoder processes the entire target sequence in parallel (teacher forcing), but each position must only attend to previous positions to maintain the autoregressive property. The mask ensures that generating token \(t\) uses only information from tokens \(1, \ldots, t{-}1\). At inference, tokens are generated one-by-one so the mask is implicit — but during training it enables parallel computation while preserving causality.

Decoder — Cross-Attention.

After masked self-attention, each decoder layer applies cross-attention where the decoder attends to the encoder’s output representations. This is the mechanism by which the decoder “reads” the input:

\[ \text{CrossAttn}(Q_{\text{dec}}, K_{\text{enc}}, V_{\text{enc}}) = \text{softmax}!\left(\frac{Q_{\text{dec}} K_{\text{enc}}^T}{\sqrt{d_k}}\right) V_{\text{enc}} \]

• Queries come from the decoder’s previous sublayer (the masked self-attention output)

• Keys and Values come from the encoder’s final output \(\mathbf{H}^{\text{enc}}\)

• No mask is applied — every decoder position can attend to every encoder position

• This allows the decoder to dynamically focus on different parts of the input at each generation step (e.g., attending to “cat” when translating to “gato” in English\(\to\)Spanish)

Full Decoder Layer.

Each decoder layer contains three sublayers (vs. two in the encoder):

1. Masked Multi-Head Self-Attention + Residual + LayerNorm

2. Multi-Head Cross-Attention (to encoder output) + Residual + LayerNorm

3. Feed-Forward Network + Residual + LayerNorm

From Encoder-Decoder to Decoder-Only.

Modern LLMs (GPT, Llama, Qwen) use only the decoder, removing both the encoder and cross-attention layers entirely. The key insight: for generative language modeling, a single causal (masked) self-attention stack is sufficient — the model learns to encode context and generate continuations in a single pass. This simplifies architecture, training, and inference while scaling more effectively. Encoder-decoder models (T5, BART) remain relevant for tasks with distinct input/output structure (translation, summarization), and cross-attention reappears in multimodal models where vision encoders provide keys/values to language decoders.

Decoder-Only vs Encoder-Decoder

Modern LLMs almost exclusively use decoder-only architectures, but understanding the trade-offs with encoder-decoder designs clarifies why.

Warning

Why Decoder-Only Won

Decoder-only models are simpler (one model, one loss), scale better (all parameters contribute to generation), and support unified training (pretraining = next-token prediction = fine-tuning objective). Encoder-decoder models waste capacity on the encoder for pure generation tasks.

Embeddings: From Discrete Tokens to Continuous Space

Before any attention or computation happens, the transformer must convert discrete token IDs into continuous vectors that neural networks can process. This is the role of the embedding layer.

What is an Embedding?

An embedding is a learned dense vector representation of a discrete symbol. Instead of representing the word “king” as a one-hot vector of size \(\vert \mathcal{V}\vert = 128{,}000\) (mostly zeros), we represent it as a compact vector in \(\mathbb{R}^d\) (e.g., \(d = 4096\)) that captures its meaning.

The key insight: similar concepts get nearby vectors. In a well-trained embedding space:

• “king” and “queen” are close (both royalty)

• “king” and “bicycle” are far apart (unrelated)

• Vector arithmetic captures relationships: \(\vec{\text{king}} - \vec{\text{man}} + \vec{\text{woman}} \approx \vec{\text{queen}}\)

The Embedding Table.

In practice, the embedding layer is simply a matrix \(\mathbf{E} \in \mathbb{R}^{\vert \mathcal{V}\vert \times d}\) where row \(i\) stores the embedding vector for token \(i\):

\[ \text{embed}(x_t) = \mathbf{E}[x_t] \in \mathbb{R}^d \]

For a sequence of token IDs \([x_1, x_2, \ldots, x_n]\), embedding is a simple table lookup (indexing operation):

\[ \mathbf{H}_0 = [\mathbf{E}[x_1];; \mathbf{E}[x_2];; \ldots;; \mathbf{E}[x_n]] \in \mathbb{R}^{n \times d} \]

Important

Embedding Table in Transformers

• Size: \(\vert \mathcal{V}\vert \times d\). For Llama-3: \(128{,}256 \times 4{,}096 = 525\)M parameters (6.5% of 8B model).

• Initialization: Random (Xavier/normal), then learned via backpropagation.

• Weight tying: Many models share the embedding matrix with the output projection head: \(W_{\text{head}} = \mathbf{E}^T\). This saves parameters and creates a symmetric encode-decode structure.

• Input: Token ID (integer) \(\to\) Output: Dense vector in \(\mathbb{R}^d\).

• Gradient flow: During training, only the rows corresponding to tokens in the current batch receive gradient updates (sparse update).

Tip

Why Embeddings Work

The embedding table is learned end-to-end with the rest of the model. Because the model is trained to predict the next token, it must learn representations where tokens that appear in similar contexts get similar vectors. This is the distributional hypothesis: “you shall know a word by the company it keeps” (Firth 1957). The embedding layer compresses this statistical structure into dense geometry.

The Anisotropy Problem.

A critical issue arises when using pretrained embeddings (e.g., from BERT or GPT-2) for downstream tasks like retrieval (RAG) or bootstrapping recommender systems: the learned representations are highly anisotropic—they occupy a narrow cone in the embedding space rather than being uniformly distributed across all directions (Ethayarajh 2019).

Why this matters for applications:

• RAG / Retrieval: If all embeddings have cosine similarity \(>0.7\) regardless of content, retrieval rankings become nearly random—the system cannot distinguish relevant from irrelevant passages.

• Recommender systems: Using pretrained LLM embeddings to represent items/users only works if the geometry preserves meaningful similarity structure.

• Clustering: Anisotropic embeddings collapse clusters, making it impossible to discover natural groupings.

Resolution: Whitening.

A simple and effective fix is whitening (Su et al. 2021)—a linear transformation that makes the embedding distribution isotropic (zero mean, identity covariance):

\[ \tilde{\mathbf{h}} = \mathbf{D}^{-1/2} \mathbf{U}^T (\mathbf{h} - \boldsymbol{\mu}) \]

where \(\boldsymbol{\mu}\) is the mean embedding, and \(\mathbf{U}\mathbf{D}\mathbf{U}^T\) is the eigendecomposition of the covariance matrix \(\Sigma = \frac{1}{N}\sum_i (\mathbf{h}_i - \boldsymbol{\mu})(\mathbf{h}_i - \boldsymbol{\mu})^T\).

Important

Whitening in Practice

• What it does: Rotates and scales the embedding space so all directions have equal variance (unit covariance).

• Effect: Cosine similarity becomes meaningful—semantically similar pairs score high, dissimilar pairs score low.

• Bonus: Can simultaneously reduce dimensionality by keeping only the top-\(k\) eigenvectors (similar to PCA), making retrieval faster.

• Cost: Requires computing the covariance matrix over a representative corpus (one-time, \(O(N \cdot d^2)\)). The transform itself is a simple matrix multiply at inference.

• Alternative approaches: Contrastive fine-tuning (SimCSE), flow-based normalization, or training with isotropy-promoting regularizers.

Self-Attention Mechanism

Self-attention is the core operation that allows each token to attend to every other token in the sequence, computing a weighted combination based on relevance.

Important

Scaled Dot-Product Attention

Given input sequence \(X \in \mathbb{R}^{n \times d}\), we compute: \[ Q = XW_Q, \quad K = XW_K, \quad V = XW_V \quad (W_Q, W_K, W_V \in \mathbb{R}^{d \times d_k}) \]

\[ \text{Attention}(Q, K, V) = \text{softmax}!\left(\frac{QK^T}{\sqrt{d_k}} + M\right) V \] where \(M\) is the causal mask (for autoregressive models): \(M_{ij} = 0\) if \(i \geq j\), else \(-\infty\).

Intuition: Each token “attends” to all previous tokens, computing a weighted average of their values based on query-key similarity.

Computational Complexity.

The naive attention computation has quadratic cost in sequence length:

• Time: \(O(n^2 \cdot d)\) — computing \(QK^T\) requires \(n^2\) dot products, each of dimension \(d_k\).

• Memory: \(O(n^2)\) — the full attention matrix must be materialized to apply softmax.

For a 128K-token context with \(d = 4096\), the attention matrix alone is \(128\text{K} \times 128\text{K} = 16.4\) billion entries (64 GB in FP32). This quadratic scaling is the fundamental bottleneck for long-context LLMs.

Attention cost scaling: why naive implementation is prohibitive for long sequences.

Approaches to Taming Attention Cost.

Several families of solutions address this quadratic bottleneck:

1. Exact attention with IO-awareness (FlashAttention (Dao et al. 2022)): Does not reduce computational complexity but eliminates the need to materialize the \(n \times n\) matrix in HBM by computing attention in tiles that fit in SRAM. Crucially, FlashAttention is orthogonal to the sparse patterns below—it is an execution engine, not an attention pattern. Production systems routinely combine FlashAttention with sliding windows or block-sparse masks, getting both IO efficiency and reduced FLOPs. We cover the algorithm in detail in Section 1.6.

2. Sliding window / local attention: Each token only attends to the \(w\) nearest tokens (e.g., \(w = 4096\)). Cost becomes \(O(n \cdot w)\)—linear in \(n\). Used by Mistral (Jiang et al. 2023) (window \(= 4096\)) and Longformer (Beltagy et al. 2020). Trades global context for efficiency; works well because most attention is local in practice. In modern stacks, the sliding-window mask is executed inside a FlashAttention kernel.

3. Sparse attention patterns: Combine local windows with periodic global tokens (e.g., every 512th token attends to all). BigBird (Zaheer et al. 2020) and LongT5 (Guo et al. 2022) use this. Preserves some long-range connectivity at \(O(n\sqrt{n})\) cost. Again, FlashAttention serves as the underlying kernel for the non-zero attention blocks.

4. Linear attention / state-space models: Replace \(\text{softmax}(QK^T)V\) with \(\phi(Q)(\phi(K)^T V)\) using associativity, or reformulate as a recurrence (Mamba (Gu and Dao 2023), RWKV (Bo Peng et al. 2023)). Theoretically \(O(n \cdot d^2)\) total. Unlike approaches 2–3 above, these are architectural replacements that alter model expressiveness—softmax-free attention is fundamentally less expressive, and empirically these models still lag behind transformers on tasks requiring precise long-range retrieval or complex reasoning.

5. KV cache compression: At inference, compress or evict old KV pairs to bound memory. Techniques include: H\(_2\)O (Z. Zhang et al. 2023) (heavy-hitter oracle—keep only high-attention keys), StreamingLLM (Xiao et al. 2024a) (keep initial “attention sink” tokens + recent window), and quantized KV caches (Z. Liu et al. 2024).

Tip

FlashAttention + Sparse Patterns = Best of Both Worlds

A common misconception is that FlashAttention is an alternative to sparse attention. It is not—it is an IO optimization for the attention kernel that composes freely with any attention mask. Modern production systems (e.g., Mistral, DeepSeek) use FlashAttention as the execution engine underneath a sliding-window or block-sparse mask. This gives you both reduced FLOPs (from sparsity) and optimal memory access patterns (from tiling). RingAttention (H. Liu et al. 2023) extends this further to multi-device settings, distributing the tiled computation across GPUs along the sequence dimension.

Linear attention and state-space models (Mamba, RWKV) are a genuinely different architectural choice—they sacrifice the full pairwise interaction for \(O(n)\) compute. While theoretically elegant, they have not matched transformer quality on knowledge-intensive or long-range reasoning tasks, and frontier labs continue to use exact attention (with FlashAttention + sparsity) as the backbone.

Multi-Head Attention

Rather than computing a single attention function, multi-head attention runs several attention operations in parallel, each learning to focus on different aspects of the input (syntax, semantics, position, etc.).

Important

Multi-Head Attention

Instead of one attention function with \(d\)-dimensional keys/values, use \(H\) parallel heads with dimension \(d_k = d/H\): \[ \text{MultiHead}(X) = \text{Concat}(\text{head}_1, \ldots, \text{head}_H) W_O \] Each head can learn different attention patterns (e.g., one head for syntax, another for semantics, another for positional proximity).

Grouped Query Attention (GQA): Llama-3 (Grattafiori et al. 2024) uses fewer K,V heads than Q heads (e.g., 8 KV heads shared across 32 Q heads). This reduces KV cache size by \(4\times\) with minimal quality loss.

Positional Encodings

Transformers are permutation-equivariant by construction — without positional information, the model cannot distinguish “the cat sat on the mat” from “mat the on sat cat the”. Positional encodings inject sequence-order signal so that attention can reason about token distance and direction.

Positional encoding methods in modern LLMs.

Sinusoidal (Fixed) Positional Encoding.

Introduced in the original Transformer (Vaswani et al. 2017), this method uses fixed sinusoidal functions at geometrically-spaced frequencies:

\[ \text{PE}(pos, 2i) = \sin!\Bigl(\frac{pos}{10000^{2i/d}}\Bigr), \qquad \text{PE}(pos, 2i{+}1) = \cos!\Bigl(\frac{pos}{10000^{2i/d}}\Bigr) \]

where \(pos\) is the token position, \(i\) is the dimension index, and \(d\) is the model dimension.

Motivation: Each frequency encodes position at a different scale (analogous to binary counting). The authors hypothesised that the model could learn to attend to relative positions because \(\text{PE}(pos+k)\) can be expressed as a linear function of \(\text{PE}(pos)\).

Pros: Zero learned parameters; deterministic; theoretically supports arbitrary lengths.
Cons: In practice, does not extrapolate well beyond training lengths; the model must learn to decode relative position from absolute signals indirectly; largely superseded.

Learned Absolute Positional Embedding.

Used by GPT-2 (Radford et al. 2019) and BERT (Devlin et al. 2019): a learnable embedding matrix \(\mathbf{E}_{\text{pos}} \in \mathbb{R}^{L_{\max} \times d}\) is added to token embeddings:

\[ h_0^{(pos)} = \text{TokenEmbed}(x_{pos}) + \mathbf{E}_{\text{pos}}[pos] \]

Motivation: Let the model learn whatever positional representation is optimal for the task, rather than imposing a fixed structure.

Pros: Maximum flexibility; simple implementation; often outperforms sinusoidal for short sequences.
Cons: Hard-coded maximum length \(L_{\max}\); no generalisation beyond it; embeddings near the end of \(L_{\max}\) are under-trained; adds \(L_{\max} \times d\) parameters.

Rotary Position Embedding (RoPE).

RoPE (Su et al. 2024) encodes position by rotating query and key vectors in 2D subspaces:

\[ \text{RoPE}(x_m, m) = \begin{pmatrix} x_m^{(1)} \\ x_m^{(2)} \\ \vdots \\ x_m^{(d-1)} \\ x_m^{(d)} \end{pmatrix} \odot \begin{pmatrix} \cos m\theta_1 \\ \cos m\theta_1 \\ \vdots \\ \cos m\theta_{d/2} \\ \cos m\theta_{d/2} \end{pmatrix} + \begin{pmatrix} -x_m^{(2)} \\ x_m^{(1)} \\ \vdots \\ -x_m^{(d)} \\ x_m^{(d-1)} \end{pmatrix} \odot \begin{pmatrix} \sin m\theta_1 \\ \sin m\theta_1 \\ \vdots \\ \sin m\theta_{d/2} \\ \sin m\theta_{d/2} \end{pmatrix} \]

where \(\theta_i = 10000^{-2i/d}\) and \(m\) is the position index. The key property is that the dot product between rotated queries and keys depends only on relative position:

\[ \langle \text{RoPE}(q_m, m),; \text{RoPE}(k_n, n) \rangle = f(q_m, k_n, m-n) \]

Motivation: Achieve relative position encoding without explicit bias terms, while maintaining compatibility with linear attention and KV-caching.

Pros: Naturally relative; no extra parameters; compatible with efficient inference; can be extended to longer contexts via NTK-aware scaling (Bowen Peng et al. 2023) or YaRN (adjusting \(\theta\) base or interpolating frequencies).
Cons: Slightly more compute per attention operation (rotation + interleaving); extrapolation requires explicit scaling strategies; rotation in 2D subspaces imposes structure that may not be optimal for all tasks.

Tip

RoPE Length Extension

To extend a RoPE model trained at \(L\) to context length \(L' > L\):

• Position interpolation: Scale positions by \(L/L'\) so all positions fit in \([0, L]\). Simple but compresses resolution.

• NTK-aware scaling: Increase the \(\theta\) base (e.g. \(10000 \to 10000 \cdot (L'/L)^{d/(d-2)}\)), effectively stretching high-frequency components while preserving low-frequency ones.

• YaRN (Bowen Peng et al. 2023): Combines NTK scaling with an attention temperature correction \(t = 0.1 \ln(s) + 1\) to compensate for increased entropy at longer distances.

ALiBi (Attention with Linear Biases).

ALiBi (Press et al. 2022) takes a radically different approach: no positional embedding at all. Instead, a static linear penalty is subtracted from attention scores:

\[ \text{Attention}(Q, K, V) = \text{softmax}!\left(\frac{QK^T}{\sqrt{d_k}} - m \cdot \bigl[|i-j|\bigr]_{i,j}\right) V \]

where \(m\) is a head-specific slope (set geometrically: \(m_h = 2^{-8h/H}\) for head \(h\) of \(H\) total). The bias \(-m\vert i-j\vert\) creates a soft local attention window whose width varies by head.

Motivation: Position should bias attention toward nearby tokens (recency prior) without interfering with the embedding space. By operating purely in attention-score space, ALiBi avoids polluting token representations with positional signal.

Pros: Excellent length extrapolation (trained at 1k, works at 8k+); zero parameters; trivial to implement; head-specific slopes give multi-scale locality.
Cons: Less expressive for tasks requiring precise long-range positional reasoning (e.g. “what was the 5th word?”); the linear decay is a strong inductive bias that may not suit all domains; largely overtaken by RoPE in recent models due to RoPE’s better short-context performance.

Positional encoding comparison: practical trade-offs.

Scaling to Extremely Long Contexts (100K–1M+ Tokens).

Modern frontier models (Claude (Anthropic 2024c) with 200K–1M context, Gemini 1.5 (Gemini Team 2024) at 1M+, GPT-4 (OpenAI 2023) at 128K) require positional encodings that remain faithful far beyond training lengths. The dominant solutions today:

1. RoPE with frequency scaling: The standard approach for extending RoPE beyond training length. Rather than retraining, the base frequency \(\theta\) is rescaled:

   \[ \theta'_i = \theta_i \cdot \left(\frac{L_{\text{target}}}{L_{\text{train}}}\right)^{2i/d} \]

   Variants include:

   • Linear scaling (Position Interpolation) (S. Chen et al. 2023): Simply divide position indices by a factor \(s\). Cheap but degrades quality at high extension ratios.

   • NTK-aware scaling (Bowen Peng et al. 2023): Scale the base frequency \(\theta = 10000 \to 10000 \cdot s^{d/(d-2)}\). Preserves high-frequency (local) information while extending low-frequency (global) range.

   • YaRN (Bowen Peng et al. 2023) (Yet another RoPE extensioN): Combines NTK scaling with an attention temperature correction and fine-tuning on a small long-context corpus. Used by Llama-3 to extend from 8K training to 128K deployment.

   • Dynamic NTK (Bowen Peng et al. 2023): Adjusts the scaling factor on-the-fly based on actual sequence length at inference. No fixed extension ratio needed—the model adapts as context grows.

2. Continued pretraining on long data: Even with RoPE scaling, models benefit from a short continued pretraining phase (1–5B tokens) on long documents. This teaches the model to actually use distant context, not just tolerate it positionally. Llama-3.1 used a progressive schedule: 8K \(\to\) 64K \(\to\) 128K.

3. Ring Attention / Blockwise Parallel (H. Liu et al. 2023): For sequences exceeding single-GPU memory (1M+ tokens), Ring Attention distributes the sequence across GPUs in a ring topology. Each GPU holds a block and passes KV blocks around the ring, computing local attention tiles. This enables linear memory scaling with GPU count while preserving exact attention.

4. Hybrid architectures: Some systems combine a local sliding window (e.g., 4K) for most layers with full attention at select layers (e.g., every 4th layer). This provides \(O(n \cdot w)\) cost for most computation while maintaining global information flow.

Warning

Long Context $\neq$ Long Context Usage

A model with 1M context length does not necessarily use all 1M tokens effectively. The “lost in the middle” phenomenon (N. F. Liu et al. 2024a) shows that models tend to focus on the beginning and end of long contexts, underutilizing information in the middle. Effective long-context utilization requires both positional encoding support and training on tasks that reward long-range retrieval.

Feed-Forward Network (MLP)

Each transformer block contains an MLP applied independently to each position:

\[ \text{FFN}(x) = W_2 \cdot \sigma(W_1 x + b_1) + b_2 \]

where \(W_1 \in \mathbb{R}^{d \times 4d}\), \(W_2 \in \mathbb{R}^{4d \times d}\). Modern LLMs use:

• SwiGLU activation: \(\text{FFN}(x) = W_2 (\text{Swish}(W_1 x) \odot W_3 x)\) — used by Llama (Grattafiori et al. 2024), Mistral (Jiang et al. 2023). Requires 3 weight matrices but gives better performance.

• Hidden dimension is typically \(8/3 \times d\) (rounded to multiples of 256 for Tensor Core efficiency).

Tip

FFN as Memory

Recent work (Geva et al. 2021) suggests the FFN layers act as a key-value memory: \(W_1\) rows are keys (patterns to match), \(W_2\) columns are values (information to output). The FFN “retrieves” stored knowledge based on the current hidden state.

Layer Normalization

Layer normalization stabilizes training by normalizing activations across the feature dimension. Its placement relative to the attention/FFN sublayers significantly affects training dynamics.

How LayerNorm Works.

Given a hidden state vector \(\mathbf{x} \in \mathbb{R}^d\) (a single token’s representation), LayerNorm (Ba et al. 2016) computes:

\[ \text{LayerNorm}(\mathbf{x}) = \gamma \odot \frac{\mathbf{x} - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta \]

where:

• \(\mu = \frac{1}{d}\sum_{i=1}^{d} x_i\) (mean across the \(d\) feature dimensions)

• \(\sigma^2 = \frac{1}{d}\sum_{i=1}^{d} (x_i - \mu)^2\) (variance across features)

• \(\gamma, \beta \in \mathbb{R}^d\) are learned scale and shift parameters (per-dimension)

• \(\epsilon \approx 10^{-5}\) prevents division by zero

Key distinction from BatchNorm: LayerNorm normalizes across the feature dimension of a single example, not across the batch. This makes it independent of batch size and works identically at training and inference.

RMSNorm — The Modern Simplification.

RMSNorm (Zhang and Sennrich 2019) drops the mean-centering step, normalizing only by the root-mean-square:

\[ \text{RMSNorm}(\mathbf{x}) = \gamma \odot \frac{\mathbf{x}}{\text{RMS}(\mathbf{x})}, \qquad \text{RMS}(\mathbf{x}) = \sqrt{\frac{1}{d}\sum_{i=1}^{d} x_i^2} \]

No \(\beta\) (shift) parameter and no mean subtraction — just scale. This saves one reduction operation per token and is \(\sim\)5–10% faster on GPUs while achieving equivalent model quality. All modern LLMs (Llama, Mistral, Qwen) use RMSNorm.

Important

Pre-LN vs Post-LN

• Post-LN (original Transformer): \(h + \text{LayerNorm}(\text{Attn}(h))\). Requires careful warmup; training can be unstable.

• Pre-LN (GPT-2+, all modern LLMs): \(h + \text{Attn}(\text{LayerNorm}(h))\). Stabilizes training; enables higher learning rates.

• RMSNorm (Llama (Grattafiori et al. 2024), Mistral (Jiang et al. 2023)): Simplified LayerNorm without mean-centering: \(\text{RMSNorm}(x) = x / \text{RMS}(x) \cdot \gamma\). Slightly faster, same quality.

Tip

Why Normalization Matters for Deep Networks

Without normalization, activations tend to grow or shrink exponentially through layers (exploding/vanishing activations). A 128-layer transformer without LayerNorm would see magnitudes vary by \(10^{30}\times\) between the first and last layer. Normalization constrains each layer’s output to a predictable range, enabling stable gradient flow and allowing the optimizer to use consistent learning rates throughout the network.

Model Size Reference

The following table summarizes key architectural parameters for widely-used open-weight models (latest versions as of 2025), providing a quick reference for understanding scale and design choices.

Architecture parameters for popular open-weight LLMs (2024–2025 generation).

Note: Models marked with “active” parameters use Mixture-of-Experts (MoE) architecture—total parameters indicate model capacity, while active parameters reflect per-token compute cost. DeepSeek-V3 uses Multi-head Latent Attention (MLA) instead of standard GQA, compressing KV into a low-rank latent space.

Attention Pathologies

While the attention mechanism is powerful, it exhibits systematic failure modes that practitioners must understand—especially when scaling to long contexts or interpreting model behaviour.

Attention Sink

The phenomenon.

Xiao et al. (Xiao et al. 2024b) discovered that transformer models allocate disproportionately high attention scores to the first token in the sequence—regardless of its semantic content. Even when the first token is a meaningless <BOS> marker, attention heads across all layers consistently attend to it, sometimes with 20–50% of total attention mass.

Why it happens.

Softmax attention must produce a valid probability distribution (\(\sum_j \alpha_j = 1\)). When no key is particularly relevant to a query, the model needs a “dump” location for unused attention mass. During training, the first token becomes this default sink because it is always present and positionally predictable. It functions as a no-op attention target—the model has learned to route irrelevant attention there rather than distributing it unpredictably.

\[ \alpha_{\text{sink}} = \frac{\exp(q^\top k_0 / \sqrt{d})}{\sum_{j} \exp(q^\top k_j / \sqrt{d})} \gg \frac{1}{n} \quad \text{(even when } k_0 \text{ is semantically irrelevant)} \]

Consequences.

• Streaming inference failure: When using sliding-window KV caches, evicting the first token causes perplexity to spike catastrophically—the model loses its attention sink.

• Misleading interpretability: Naive attention visualizations suggest the first token is “important” when it is merely a mathematical artefact.

• Context window waste: The sink token occupies a KV cache slot without carrying useful information.

Solutions.

• StreamingLLM (Xiao et al. 2024b): Always keep the first \(k\) tokens (“attention sinks”) in the KV cache alongside the recent sliding window. Enables infinite-length generation with bounded memory.

• Sink tokens by design: Some models (e.g., Mistral) prepend dedicated sink tokens during training that are explicitly meant to absorb residual attention.

• Softmax alternatives: Replace softmax with ReLU attention or sigmoid gating, where zero attention is representable without requiring a dump target.

Attention Dilution

The phenomenon.

As sequence length \(n\) grows, each query must distribute its attention budget across more keys. The average attention weight per token decreases as \(O(1/n)\), making it progressively harder for the model to concentrate on the few truly relevant positions—a problem known as attention dilution or attention diffusion (N. F. Liu et al. 2024a).

The “Lost in the Middle” effect.

Liu et al. (N. F. Liu et al. 2024a) showed that LLMs exhibit a U-shaped retrieval curve: information placed at the beginning or end of long contexts is retrieved reliably, but information in the middle is often ignored. This is a direct consequence of attention dilution compounded with positional biases from RoPE/ALiBi:

Why it happens.

• Softmax saturation: With many keys, the softmax temperature effectively decreases, making the distribution more uniform (entropic).

• Positional decay: RoPE’s relative positional encoding introduces a natural decay with distance, suppressing attention to middle positions that are far from both start and end.

• Training distribution: Models trained on shorter sequences develop attention patterns biased toward recent context.

Mitigation strategies.

• Explicit retrieval: Place relevant context at the beginning or end of the prompt; use RAG to avoid relying on middle positions.

• Long-context training: Train on long documents with varied placement of key information (Yao Fu et al. 2024).

• Hierarchical attention: Architectures like Mamba (Gu and Dao 2024) or RWKV that avoid the \(O(n^2)\) attention bottleneck entirely.

• Landmark tokens: Insert retrievable markers in the context that act as “signposts” for attention.

• Temperature scaling: Some implementations scale the attention logits by \(\log n\) to counteract dilution in long sequences.

Other Attention Phenomena

Additional attention patterns observed in large transformers.

Visualizing Attention for Explainability

Attention weights provide a window into model reasoning—but must be interpreted carefully.

Attention Visualization Methods

Raw attention maps.

The simplest approach: plot the \(n \times n\) attention matrix \(A = \text{softmax}(QK^\top/\sqrt{d})\) as a heatmap for each head and layer. Tools like BertViz (Vig 2019) render interactive multi-head visualizations.

Attention rollout.

Raw attention at a single layer is misleading because information flows through residual connections across all layers. Abnar and Zuidema (Abnar and Zuidema 2020) propose attention rollout: multiply attention matrices across layers to approximate the total information flow from input to output:

\[ R^{(l)} = A^{(l)} \cdot R^{(l-1)}, \quad R^{(0)} = I \]

where \(A^{(l)}\) is the (averaged across heads) attention matrix at layer \(l\), adjusted to include the residual connection: \(A^{(l)} = 0.5 \cdot A^{(l)}_{\text{raw}} + 0.5 \cdot I\).

Gradient-weighted attention.

Combine attention weights with gradient information to identify which attended tokens actually influence the output (Barkan et al. 2021):

\[ \text{Relevance}(i) = \alpha_i \cdot \left|\frac{\partial y}{\partial h_i}\right| \]

This addresses the criticism that high attention \(\neq\) high influence (a token can receive high attention but be processed through a near-zero-weight path).

Warning

Attention Is Not Explanation

Jain and Wallace (Jain and Wallace 2019) showed that attention weights often do not correlate with gradient-based feature importance and that adversarial attention distributions can produce identical outputs. Use attention visualization as a hypothesis generator, not as a faithful explanation. For causal attribution, prefer gradient-based methods, probing, or mechanistic interpretability.

Mechanistic Interpretability with Sparse Autoencoders (SAEs)

The interpretability problem.

Individual neurons in transformer MLPs and residual streams are typically polysemantic—a single neuron activates for multiple unrelated concepts (e.g., “the colour blue AND academic citations AND the word ‘the’”). This makes direct neuron-level interpretation unreliable.

Sparse Autoencoders.

Cunningham et al. (Cunningham et al. 2024) and Bricken et al. (Bricken et al. 2023) demonstrated that training a sparse autoencoder (SAE) on model activations can decompose polysemantic representations into monosemantic features—interpretable directions that each correspond to a single concept:

\[ h = W_{\text{dec}} \cdot \text{ReLU}(W_{\text{enc}} \cdot x + b_{\text{enc}}) + b_{\text{dec}} \]

where \(W_{\text{enc}} \in \mathbb{R}^{m \times d}\) with \(m \gg d\) (overcomplete basis), and the ReLU + sparsity penalty ensures only a few features activate per input.

Key findings from SAE interpretability:

• Features are monosemantic: each encodes a single human-interpretable concept (“code in Python,” “mentions of the Golden Gate Bridge,” “first-person narrative”) (Bricken et al. 2023).

• Features are steerable: clamping a feature’s activation high/low directly controls model behaviour (e.g., forcing the “Golden Gate Bridge” feature on makes the model mention it in every response) (Templeton et al. 2024).

• Features compose: complex behaviours emerge from combinations of simple features.

• SAEs scale: Templeton et al. (Templeton et al. 2024) trained SAEs with up to 34M features on Claude 3 Sonnet, finding interpretable features for safety-relevant concepts (deception, sycophancy, dangerous requests).

Important

SAE Training Recipe

1. Collect activations from a specific model layer across a large corpus.

2. Train a sparse autoencoder with \(L_1\) penalty on the hidden layer: \(\mathcal{L} = \\vert x - \hat{x}\\vert _2^2 + \lambda \\vert z\\vert _1\).

3. The learned encoder directions (\(W_{\text{enc}}\) rows) are candidate features.

4. Validate: for each feature, find max-activating examples and check semantic coherence.

5. Optionally: measure feature absorption and dead features to assess SAE quality.

Natural Language Autoencoders (Anthropic, 2026)

While SAEs decompose activations into interpretable vectors, their features still require human inspection of max-activating examples to understand. Anthropic’s Natural Language Autoencoders (NLAEs) (Anthropic 2026) take a fundamentally different approach: they replace the sparse bottleneck with natural language descriptions, making interpretability automatic.

How NLAEs work.

1. Encoder: A language model reads the hidden activations (or the input text) and produces a natural language description of the active concepts: e.g., “The text discusses French cuisine and uses formal academic tone.”

2. Decoder: A second language model reads the natural language description and reconstructs the original activations (or predicts the next token).

3. Training: Both encoder and decoder are trained end-to-end to minimize reconstruction loss, with the bottleneck being a variable-length natural language string rather than a sparse vector.

Advantages over SAEs.

• Self-interpreting: Features are literally natural language—no manual labelling needed.

• Compositional: Can express complex, relational concepts (“a sarcastic response to a factual claim”) that SAE features cannot represent as single directions.

• Hierarchical: Descriptions can capture both fine-grained (word-level) and coarse (document-level) properties in the same representation.

• Auditable: The bottleneck description is human-readable, enabling direct inspection of what information the model “thinks” is present.

Limitations.

NLAEs introduce a language-model-in-the-loop, making them computationally expensive and potentially subject to the same faithfulness concerns as any model-generated explanation. They also cannot easily represent sub-symbolic features (geometric patterns, exact numerical values) that SAEs handle naturally as activation magnitudes.

Tip

The Interpretability Stack

Think of interpretability tools as a hierarchy:

1. Attention maps: “What is the model looking at?” (cheapest, least faithful)

2. Probing classifiers: “What information is encoded at this layer?”

3. Sparse Autoencoders: “What monosemantic features are active?” (scalable, requires human labelling)

4. Natural Language Autoencoders: “What does the model think is happening?” (self-interpreting, expensive)

5. Causal tracing / patching: “Which components actually cause this output?” (most faithful, most expensive)

Each level trades off between cost, scalability, and faithfulness of explanation.

Prediction Heads: What Transformers Output

The transformer body produces contextual hidden states \(\mathbf{h}_t \in \mathbb{R}^d\) for each position. What we do with these hidden states—the prediction head—defines the task. The same transformer backbone can serve radically different purposes simply by swapping the head.

Language Modeling Head (Pretraining)

The standard LM head projects the final hidden state to vocabulary logits and trains with cross-entropy loss over the next token:

\[ P(x_{t+1} | x_{\leq t}) = \text{softmax}(\mathbf{W}_{\text{head}} \cdot \mathbf{h}_t + \mathbf{b}) \]

where \(\mathbf{W}_{\text{head}} \in \mathbb{R}^{\vert \mathcal{V}\vert \times d}\) (often tied with the embedding matrix: \(\mathbf{W}_{\text{head}} = \mathbf{E}^T\)).

Important

LM Head Properties

• Training objective: Causal language modeling (predict next token for every position)

• Loss: \(\mathcal{L}_{\text{LM}} = -\frac{1}{T}\sum_{t=1}^{T} \log P(x_t \vert x_{<t})\)

• Label: Every token is both input (shifted right) and target (shifted left)

• Used during: Pretraining on large corpora (trillions of tokens)

• Key insight: The model learns general language understanding as a byproduct of next-token prediction

Conditional Generation Head (SFT / Instruction Following)

For supervised fine-tuning (SFT), the architecture is identical to the LM head—the same linear projection to vocabulary logits. The difference is purely in what we compute loss on:

\[ \mathcal{L}_{\text{SFT}} = -\frac{1}{|y|}\sum_{t=1}^{|y|} \log P(y_t | x_{\text{prompt}}, y_{<t}) \]

Important

Conditional Head – Key Differences from LM Head

• Loss masking: Only compute loss on the response tokens, not the prompt/instruction. The prompt provides context but no gradient signal.

• Conditioning: The model learns to generate responses conditioned on specific instruction formats (system prompts, user queries, tool calls).

• Format tokens: Special tokens (<|user|>, <|assistant|>) guide the model to produce structured outputs.

• Used during: SFT on curated instruction-response pairs; also during RL policy generation (the policy head that produces actions/responses).

Tip

Same Head – Different Training Signal

The LM head and SFT head are architecturally identical (same \(\mathbf{W}_{\text{head}}\)). The only difference is that during SFT, we mask the loss on prompt tokens. This subtle change transforms a general text predictor into a instruction-following assistant. The head learns to “activate” different generation modes based on the conditioning context.

Value Head (Regression for RL)

In reinforcement learning (PPO, GRPO), we need to estimate how good a state is—this requires a scalar output, not vocabulary logits. The value head replaces the LM projection with a simple regression layer:

\[ V(s_t) = \mathbf{w}_{\text{value}}^T \cdot \mathbf{h}_t + b \in \mathbb{R} \]

where \(\mathbf{w}_{\text{value}} \in \mathbb{R}^d\) and \(b \in \mathbb{R}\).

Important

Value Head Properties

• Output: Single scalar (expected cumulative reward from this state)

• Loss: MSE between predicted and actual returns: \(\mathcal{L}_V = \frac{1}{T}\sum_t (V(s_t) - R_t)^2\)

• Architecture: Linear layer \(\mathbb{R}^d \to \mathbb{R}^1\) (sometimes with a small MLP: \(d \to 256 \to 1\))

• Backbone sharing: Often shares the transformer body with the policy (with a separate value head), or uses a completely separate critic network

• Used during: PPO advantage estimation (GAE), reward model scoring

Head Selection Summary

Prediction heads used throughout this paper and their training contexts.

Warning

Head Initialization Matters

When adding a value head to a pretrained LM, initialize it near zero (small random weights). If initialized with large values, the initial value estimates will be wildly wrong, causing huge advantages and unstable PPO updates. Common practice: initialize the final linear layer with \(\mathcal{N}(0, 1/\sqrt{d})\) or simply zeros.

HuggingFace Implementation

from transformers import (
    AutoModelForCausalLM,          # LM head (pretraining + SFT)
    AutoModelForSequenceClassification,  # Reward head
    AutoTokenizer,
)
from trl import AutoModelForCausalLMWithValueHead  # Value head (PPO)
import torch

model_name = "meta-llama/Llama-3.1-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# === 1. LM Head (Pretraining / SFT) ===
# The default CausalLM model -- projects hidden states to vocab logits
lm_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
# lm_model.lm_head: Linear(hidden_size -> vocab_size)
# Output: logits of shape (batch, seq_len, vocab_size)

inputs = tokenizer("The capital of France is", return_tensors="pt")
outputs = lm_model(**inputs)
next_token_logits = outputs.logits[:, -1, :]  # (batch, vocab_size)
probs = torch.softmax(next_token_logits, dim=-1)

# === 2. Conditional Head (SFT) ===
# Architecturally identical to LM head -- difference is in loss masking
# During SFT, we only compute loss on response tokens:
messages = [
    {"role": "user", "content": "What is 2+2?"},
    {"role": "assistant", "content": "4"},
]
formatted = tokenizer.apply_chat_template(messages, return_tensors="pt")
labels = formatted.clone()
# Mask prompt tokens (set to -100 so cross-entropy ignores them)
prompt_len = len(tokenizer.apply_chat_template(messages[:1]))
labels[:, :prompt_len] = -100
loss = lm_model(input_ids=formatted, labels=labels).loss

# === 3. Value Head (PPO Critic) ===
# Adds a Linear(hidden_size -> 1) on top of the LM backbone
value_model = AutoModelForCausalLMWithValueHead.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
# value_model.v_head: Linear(hidden_size -> 1)
# Returns both LM logits AND per-token value estimates

inputs = tokenizer("Explain quantum computing", return_tensors="pt")
lm_logits, loss, values = value_model(
    **inputs, return_dict=False
)
# values shape: (batch, seq_len, 1) -- scalar estimate per token

# === 4. Reward Head (Reward Model) ===
# Classification head: Linear(hidden_size -> 1) on last token
reward_model = AutoModelForSequenceClassification.from_pretrained(
    model_name,
    num_labels=1,              # single scalar output
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
# Scores entire sequence by pooling the last token's hidden state
inputs = tokenizer("Good response here", return_tensors="pt")
reward_score = reward_model(**inputs).logits  # shape: (batch, 1)

Tip

Weight Tying: LM Head = Embedding Matrix Transposed

Most modern LLMs tie the LM head weights with the input embedding matrix: lm_head.weight = model.embed_tokens.weight. This means the LM head is not a separately learned layer—it reuses the embedding table. Benefits: fewer parameters (\(\vert \mathcal{V}\vert \times d\) saved), better generalization, and the geometric structure of the embedding space directly determines token probabilities. You can verify this in HuggingFace: model.lm_head.weight is model.model.embed_tokens.weight returns True for most models.

Optimization Theory for LLM Training

Training a large language model means finding the set of parameters \(\theta\) (billions of weights) that minimizes the loss function \(\mathcal{L}(\theta)\) — typically the negative log-likelihood of the next token. This is an optimization problem in extraordinarily high-dimensional space, and the algorithm used to navigate this space determines whether training succeeds, diverges, or stalls.

Gradient Descent: The Foundation

What is a Gradient?

The gradient \(\nabla_\theta \mathcal{L}\) is a vector that points in the direction of steepest increase of the loss. Each component \(\frac{\partial \mathcal{L}}{\partial \theta_i}\) tells us how much the loss would change if we slightly increased parameter \(\theta_i\). To decrease the loss, we move in the opposite direction:

\[ \theta_{t+1} = \theta_t - \eta \nabla_\theta \mathcal{L}(\theta_t) \]

where \(\eta > 0\) is the learning rate — the step size. This is gradient descent (Rumelhart et al. 1986).

Why Full Gradient Descent is Impractical.

Computing the exact gradient requires evaluating the loss over the entire training dataset (trillions of tokens for LLMs). This is computationally prohibitive — a single gradient step would require a full pass over all data.

Stochastic Gradient Descent (SGD).

The solution: estimate the gradient from a small random subset (mini-batch) of the data (Robbins and Monro 1951):

\[ \nabla_\theta \mathcal{L}(\theta) \approx \frac{1}{B}\sum_{i=1}^{B} \nabla_\theta \ell(\theta; x_i) \]

where \(B\) is the batch size (typically 1K–4M tokens for LLMs). The mini-batch gradient is a noisy but unbiased estimate of the true gradient.

Important

Why Mini-Batch SGD Works

• Computational efficiency: Each step costs \(O(B)\) instead of \(O(N_{\text{total}})\). With \(B = 4096\) tokens and 15T total tokens, each step is \(\sim\)4 billion\(\times\) cheaper.

• Noise as regularization: The stochastic noise helps escape sharp local minima, finding flatter regions that generalize better.

• GPU utilization: Mini-batches are large enough to saturate GPU parallelism (matrix multiplications become compute-bound rather than memory-bound).

• Convergence: Theoretically converges to a local minimum at rate \(O(1/\sqrt{T})\) (slower than exact GD’s \(O(1/T)\), but each step is millions of times cheaper).

From SGD to Adaptive Methods.

While SGD with momentum works well for vision models (CNNs), LLM training requires adaptive optimizers — algorithms that maintain a per-parameter learning rate.

Why Vanilla SGD Fails for LLMs

Stochastic Gradient Descent updates weights as:

\[ \theta_{t+1} = \theta_t - \eta \nabla_\theta \mathcal{L}(\theta_t) \]

Warning

SGD Problems for LLMs

• Different gradient scales per layer: Early layers in a transformer have much smaller gradients than later layers (vanishing gradients). A single learning rate \(\eta\) is too large for some parameters and too small for others.

• Sparse gradients: Embedding layers receive gradients only for tokens in the current batch. Most embedding rows have zero gradient. SGD with momentum wastes momentum on zero-gradient rows.

• Saddle points: High-dimensional loss landscapes have many saddle points. SGD can stall; adaptive methods escape faster.

• Sensitivity to learning rate: SGD requires careful tuning; a 2\(\times\) change in \(\eta\) can cause divergence.

Adam – Adaptive Moment Estimation

Adam (Kingma and Ba 2015) maintains per-parameter estimates of the first moment (mean of gradients) and second moment (uncentered variance of gradients).

Important

Adam Update Equations

Given gradient \(g_t = \nabla_\theta \mathcal{L}(\theta_t)\), hyperparameters \(\beta_1, \beta_2, \epsilon, \eta\):

Step 1 – Update biased first moment estimate: \[ m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t \]

Step 2 – Update biased second moment estimate: \[ v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2 \]

Step 3 – Bias correction: \[ \hat{m}_t = \frac{m_t}{1 - \beta_1^t}, \qquad \hat{v}_t = \frac{v_t}{1 - \beta_2^t} \]

Step 4 – Parameter update: \[ \theta_{t+1} = \theta_t - \eta \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} \]

Typical values: \(\beta_1 = 0.9\), \(\beta_2 = 0.95\) or \(0.999\), \(\epsilon = 10^{-8}\), \(\eta = 10^{-4}\) to \(10^{-5}\).

Tip

What Each Term Does

• \(m_t\) (momentum): Exponential moving average of gradients. Smooths out noisy gradient estimates. \(\beta_1 = 0.9\) means the current gradient contributes 10% and the history contributes 90%.

• \(v_t\) (adaptive LR): EMA of squared gradients. Parameters with consistently large gradients get a smaller effective learning rate (\(\eta / \sqrt{v_t}\)). Parameters with small gradients get a larger effective LR. This is the key to handling different gradient scales per layer.

• \(\hat{m}_t, \hat{v}_t\) (bias correction): At \(t=1\), \(m_1 = (1-\beta_1)g_1\) is much smaller than the true mean. Dividing by \((1-\beta_1^t)\) corrects this initialization bias. Without it, early steps are too small.

• \(\epsilon\) (numerical stability): Prevents division by zero. Also acts as a floor on the effective learning rate.

AdamW – Decoupled Weight Decay

AdamW (Loshchilov and Hutter 2019) fixes a subtle but important issue with how weight decay interacts with adaptive optimizers.

Important

Why L2 Regularization $\neq$ Weight Decay in Adam

With L2 regularization, the loss becomes \(\mathcal{L} + \frac{\lambda}{2}\\vert \theta\\vert ^2\), so the gradient is \(g_t + \lambda \theta_t\). In Adam, this regularization gradient is scaled by the adaptive factor \(1/\sqrt{\hat{v}_t}\): \[ \theta_{t+1} = \theta_t - \eta \cdot \frac{\hat{m}_t + \lambda \theta_t}{\sqrt{\hat{v}_t} + \epsilon} \] Parameters with large \(v_t\) (large gradient variance) get less regularization. This is not what we want – weight decay should be uniform.

Important

AdamW – Decoupled Weight Decay

AdamW (Loshchilov & Hutter, 2017) applies weight decay directly to the parameters, outside the adaptive scaling: \[ \theta_{t+1} = \theta_t - \eta \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}

• \eta \lambda \theta_t \] The weight decay term \(\eta \lambda \theta_t\) is not divided by \(\sqrt{\hat{v}_t}\). This gives uniform regularization across all parameters regardless of their gradient history.

Typical value: \(\lambda = 0.1\) for LLM training.

Warning

Always Use AdamW – Never Plain Adam – for LLMs

The difference between Adam and AdamW is subtle but matters. With Adam + L2, the effective weight decay is stronger for parameters with small gradient variance (e.g., biases, LayerNorm parameters) and weaker for parameters with large gradient variance (e.g., attention weights). AdamW gives the intended uniform regularization. Most frameworks default to AdamW; double-check your optimizer class.

Learning Rate – The Most Important Hyperparameter

Important

Typical Learning Rates by Training Phase

Phase	Typical LR	Notes
Pretraining (from scratch)	\(1\text{e-}4\) to \(3\text{e-}4\)	Large model, large batch
Continued pretraining	\(1\text{e-}5\) to \(1\text{e-}4\)	Smaller LR to preserve knowledge
SFT (supervised fine-tune)	\(1\text{e-}5\) to \(2\text{e-}5\)	Standard range
LoRA fine-tuning	\(1\text{e-}4\) to \(3\text{e-}4\)	Higher LR for adapter weights

For RL learning rates (PPO, DPO, GRPO) see §8.15.

Learning Rate Warmup

Important

Why Warmup is Necessary

At the start of training, \(v_t\) (the second moment estimate) is initialized to zero. After bias correction: \(\hat{v}_t = v_t / (1 - \beta_2^t)\). At \(t=1\) with \(\beta_2 = 0.999\): \(\hat{v}_1 = v_1 / (1 - 0.999) = 1000 v_1\). This means the effective learning rate is \(\eta / \sqrt{1000 v_1}\) – much smaller than intended.

However, if the first gradient is unusually large (common at initialization), the second moment estimate can be dominated by this outlier, causing erratic early steps. Warmup mitigates this by starting with a very small LR and gradually increasing it, giving \(v_t\) time to accumulate a reliable estimate.

• Linear warmup: \(\eta_t = \eta_{\max} \times t / T_{\text{warmup}}\)

• Typical warmup duration: 1–5% of total steps for pretraining; 3–10% for fine-tuning (shorter runs need proportionally more warmup)

• For SFT: 50–200 warmup steps is typical

Learning Rate Schedules

(a) Constant.

Simplest schedule. Good for short fine-tuning runs where you want to avoid over-decaying the LR. Risk: no annealing means the model may not converge to the sharpest minimum.

(b) Cosine Decay.

\[ \eta_t = \eta_{\min} + \frac{1}{2}(\eta_{\max} - \eta_{\min}) \left(1 + \cos!\left(\frac{t - T_{\text{warmup}}}{T - T_{\text{warmup}}} \pi\right)\right) \]

Standard for pretraining and SFT. Smooth decay avoids abrupt LR changes. \(\eta_{\min}\) is typically \(\eta_{\max} / 10\).

(c) Linear Decay.

Simpler than cosine, similar empirical results. Preferred when you want predictable LR at any step.

(d) WSD – Warmup-Stable-Decay.

The new standard for large-scale pretraining (S. Hu et al. 2024; Grattafiori et al. 2024). Three phases:

1. Warmup: Linear ramp to \(\eta_{\max}\) (1–5% of steps)

2. Stable: Constant \(\eta_{\max}\) for the majority of training

3. Decay: Fast cosine or linear decay to \(\eta_{\min}\) (last 10–20% of steps)

Key advantage: the stable phase allows checkpointing at any point and continuing training. The decay phase can be applied at the end of any run.

(e) Cosine with Restarts (SGDR).

Periodic restarts reset the LR to \(\eta_{\max}\). Can help escape local minima. Less common for LLMs; more useful for smaller models.

Gradient Clipping

Important

Gradient Clipping

Gradient clipping rescales the gradient if its global norm exceeds a threshold: \[ g_t \leftarrow g_t \cdot \min!\left(1,; \frac{\tau}{|g_t|_2}\right) \] where \(\tau\) is max_grad_norm (typically 1.0).

Tip

Gradient Clipping vs. LR Reduction

Gradient clipping and reducing the learning rate both limit the size of parameter updates. The difference: clipping preserves the direction of the gradient (just scales the magnitude), while a smaller LR scales all updates uniformly. Clipping is better for handling occasional large gradients without slowing down normal training steps.

Putting It Together: HuggingFace Optimizer Configuration

The following snippet shows how the concepts from this section—AdamW with decoupled weight decay (§1.6.6), cosine learning-rate scheduling with linear warmup (§1.6.7), and gradient clipping (§1.6.8)—come together in practice using the HuggingFace transformers library.

from transformers import TrainingArguments, Trainer
from transformers import get_cosine_schedule_with_warmup
import torch

# --- Option 1: Using TrainingArguments (recommended) ---
training_args = TrainingArguments(
    output_dir="./checkpoints",

    # AdamW optimizer (decoupled weight decay, S1.6.6)
    optim="adamw_torch",
    learning_rate=2e-5,           # peak LR after warmup
    adam_beta1=0.9,               # first moment decay
    adam_beta2=0.999,             # second moment decay
    adam_epsilon=1e-8,            # numerical stability
    weight_decay=0.01,           # decoupled L2 penalty

    # Learning rate schedule (S1.6.7)
    lr_scheduler_type="cosine",  # cosine decay to 0
    warmup_ratio=0.1,            # 10% of steps = linear warmup

    # Gradient clipping (S1.6.8)
    max_grad_norm=1.0,           # clip by global L2 norm

    # Mixed precision (S1.6.9)
    bf16=True,                   # use BFloat16 on Ampere+ GPUs

    # Training duration
    num_train_epochs=3,
    per_device_train_batch_size=8,
    gradient_accumulation_steps=4,  # effective batch = 8*4 = 32
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
)
trainer.train()

# --- Option 2: Manual control (for custom training loops) ---
from torch.optim import AdamW

# Separate weight-decay groups (don't regularize biases/norms)
no_decay = ["bias", "LayerNorm.weight", "layernorm.weight"]
param_groups = [
    {
        "params": [p for n, p in model.named_parameters()
                   if not any(nd in n for nd in no_decay)],
        "weight_decay": 0.01,
    },
    {
        "params": [p for n, p in model.named_parameters()
                   if any(nd in n for nd in no_decay)],
        "weight_decay": 0.0,
    },
]

optimizer = AdamW(param_groups, lr=2e-5, betas=(0.9, 0.999))

# Cosine schedule with linear warmup
total_steps = len(train_dataloader) * num_epochs
warmup_steps = int(0.1 * total_steps)
scheduler = get_cosine_schedule_with_warmup(
    optimizer,
    num_warmup_steps=warmup_steps,
    num_training_steps=total_steps,
)

# Training loop with gradient clipping
for batch in train_dataloader:
    outputs = model(**batch)
    loss = outputs.loss
    loss.backward()

    # Clip gradients before optimizer step
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

    optimizer.step()
    scheduler.step()
    optimizer.zero_grad()

Important

Practical Tips

• Weight decay exclusion: bias terms and layer-norm weights should not be regularized—they have few parameters and regularizing them hurts performance (Loshchilov and Hutter 2019).

• Warmup ratio: 5–10% of total steps is standard; too little warmup with a high LR can destabilize early training.

• Gradient accumulation: simulates larger batches on limited GPU memory; clipping applies to the accumulated gradient.

• BF16 vs. FP16: prefer bf16=True on Ampere+ GPUs (wider dynamic range avoids loss scaling); fall back to fp16=True on older hardware.

Mixed Precision Training

Important

BF16 vs. FP16

Format	Exponent bits	Mantissa bits	Dynamic range
FP32	8	23	\(\sim 10^{-38}\) to \(10^{38}\)
BF16	8	7	Same as FP32 (same exponent)
FP16	5	10	\(\sim 6 \times 10^{-5}\) to \(65504\)

Tip

BF16 Over FP16: Why Range Beats Precision in LLM Training

BF16 has the same exponent range as FP32, so it can represent the same range of values (just with less precision in the mantissa). FP16 has a much smaller dynamic range – gradients or activations that exceed 65504 cause overflow (NaN/Inf). This is why FP16 training requires loss scaling (multiplying the loss by a large constant to keep gradients in FP16 range), while BF16 training typically does not. A100 and H100 support BF16 natively; use BF16 unless you have a specific reason for FP16.

Loss Scaling (FP16 only).

1. Multiply loss by scale factor \(S\) (e.g., \(S = 2^{15}\))

2. Compute gradients in FP16 (scaled by \(S\))

3. Before optimizer step, divide gradients by \(S\)

4. Check for overflow (NaN/Inf); if found, skip step and reduce \(S\)

5. If no overflow for \(N\) consecutive steps, increase \(S\)

FP32 Master Weights.

In mixed precision training, weights are stored in FP32 (master copy) and cast to BF16/FP16 for the forward/backward pass. The optimizer step is done in FP32. This is important because:

• Small gradient updates (\(\Delta\theta \ll \theta\)) would be lost in BF16 precision (7 mantissa bits \(\approx\) 0.8% relative precision)

• FP32 master weights ensure accurate accumulation of small updates over many steps

• Memory cost: 2\(\times\) weight storage (FP32 + BF16 copy)

Warning

When FP32 Master Weights Are Critical

FP32 master weights are most important for:

• Long training runs (many small gradient steps accumulate)

• Small learning rates (updates are tiny relative to weight magnitude)

For short SFT runs with large LR, BF16-only training (no FP32 master weights) often works fine and saves memory. For RL training, FP32 master weights are essential—see §8.15.

Mixed Precision in Practice: HuggingFace

# === HuggingFace TrainingArguments (simplest approach) ===
from transformers import TrainingArguments

# BF16 on Ampere+ GPUs (A100, H100, RTX 30xx/40xx)
args_bf16 = TrainingArguments(
    output_dir="./out",
    bf16=True,               # BF16 forward/backward; FP32 master weights
    bf16_full_eval=True,     # also use BF16 during evaluation
    # No loss scaling needed -- BF16 has FP32-equivalent range
)

# FP16 on older GPUs (V100, T4, RTX 20xx)
args_fp16 = TrainingArguments(
    output_dir="./out",
    fp16=True,               # FP16 forward/backward
    fp16_full_eval=False,    # keep eval in FP32 for accuracy
    # Loss scaling is automatic via PyTorch GradScaler
)

# === Manual PyTorch AMP (for custom training loops) ===
import torch

# Setup (PyTorch 2.x API)
use_fp16 = not torch.cuda.is_bf16_supported()
scaler = torch.amp.GradScaler("cuda", enabled=use_fp16)  # only needed for FP16
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
dtype = torch.float16 if use_fp16 else torch.bfloat16

for batch in train_dataloader:
    optimizer.zero_grad()

    # Autocast: run forward pass in reduced precision
    with torch.autocast("cuda", dtype=dtype):
        outputs = model(**batch)
        loss = outputs.loss

    if use_fp16:
        # FP16 path: scale loss to prevent gradient underflow
        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)          # unscale before clipping
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(optimizer)              # skips step on overflow
        scaler.update()                     # adjust scale factor
    else:
        # BF16 path: no scaling needed
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()

    scheduler.step()

Tip

Key Differences: BF16 vs. FP16 in Code

• BF16: just wrap with autocast(dtype=torch.bfloat16)—no scaler needed. Simpler code and more numerically stable.

• FP16: requires GradScaler to prevent gradient underflow. The scaler dynamically adjusts a multiplier; if overflow is detected (NaN), the optimizer step is skipped and the scale is reduced.

• Gradient clipping + FP16: you must call scaler.unscale_(optimizer) before clip_grad_norm_, otherwise you’re clipping scaled gradients (wrong threshold).

• Memory savings: % reduction in activation memory (activations stored in 16-bit); weight memory depends on whether you keep FP32 master copies.

Practical Optimizer Settings by Training Phase

Important

Optimizer Hyperparameter Reference Table

Phase	Optimizer	LR	WD	Warmup	Schedule
Pretraining	AdamW	\(3\text{e-}4\)	0.1	2000 steps	WSD or Cosine
SFT	AdamW	\(2\text{e-}5\)	0.01	100 steps	Cosine
LoRA SFT	AdamW	\(2\text{e-}4\)	0.01	100 steps	Cosine

All use: \(\beta_1{=}0.9\), \(\beta_2{=}0.95\), \(\epsilon{=}10^{-8}\), max_grad_norm=1.0, BF16. For RL settings see §8.15.

Note

Diagnosing Training Instability

# Monitor these metrics to diagnose optimizer issues:
# 1. Gradient norm -- should be < max_grad_norm most of the time
# 2. Loss scale (FP16) -- should be stable, not constantly decreasing
# 3. Parameter update norm -- should be << parameter norm

import torch

def log_optimizer_stats(model, optimizer, step):
    # Gradient norm (before clipping)
    total_norm = 0.0
    for p in model.parameters():
        if p.grad is not None:
            total_norm += p.grad.data.norm(2).item() ** 2
    total_norm = total_norm ** 0.5

    # Adam second moment stats (proxy for adaptive LR)
    v_norms = []
    for group in optimizer.param_groups:
        for p in group['params']:
            state = optimizer.state[p]
            if 'exp_avg_sq' in state:
                v_norms.append(state['exp_avg_sq'].mean().item())

    print(f"Step {step}: grad_norm={total_norm:.3f}, "
          f"mean_v={sum(v_norms)/len(v_norms):.6f}")

# Red flags:
# grad_norm >> 1.0 repeatedly -> reduce LR or increase warmup
# grad_norm == 0.0 -> gradient vanishing or wrong loss
# loss_scale decreasing -> FP16 overflow, switch to BF16
# v very small -> Adam not warmed up yet, extend warmup

Tip

The Learning Rate is the Most Important Hyperparameter

In practice, getting the learning rate right matters more than any other hyperparameter. A rule of thumb for LLM fine-tuning:

• Start with the values in the table above

• If loss diverges (increases after initial decrease): LR is too high

• If loss decreases very slowly and plateaus early: LR is too low

• If loss is unstable (oscillates): LR is too high or warmup is too short

The second most important hyperparameter is batch size (affects gradient noise and effective LR via the linear scaling rule). Everything else is secondary.

Flash Attention – Algorithm and Hardware Awareness

Flash Attention (Dao et al. 2022; Dao 2024) is one of the most impactful algorithmic innovations in deep learning since the transformer itself. It does not change the mathematical result of attention – it computes exactly the same output – but it restructures the memory access pattern so that the GPU’s limited fast SRAM does all the heavy lifting, cutting HBM footprint from \(O(n^2)\) to \(O(n)\) and delivering 2–4\(\times\) end-to-end wall-clock gains on typical workloads.

The Standard Attention Memory Problem

Standard scaled dot-product attention is:

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

Important

Standard Attention Memory Complexity

For sequence length \(n\) and head dimension \(d\):

• \(Q, K, V \in \mathbb{R}^{n \times d}\): \(O(nd)\) memory

• \(S = QK^T \in \mathbb{R}^{n \times n}\): \(O(n^2)\) memory – the bottleneck

• \(P = \text{softmax}(S) \in \mathbb{R}^{n \times n}\): another \(O(n^2)\)

• \(O = PV \in \mathbb{R}^{n \times d}\): \(O(nd)\)

At \(n=8192\), \(d=128\), BF16: the attention matrix alone is \(8192^2 \times 2 \approx 134\) MB per head. With 32 heads, that is 4.3 GB just for one layer’s attention scores.

Tip

Why $O(n^2)$ is Catastrophic

The attention matrix must be written to HBM (it does not fit in SRAM for long sequences), then read back for the softmax, then read again for the \(PV\) product. Each of these HBM round-trips is slow. For \(n=32768\) (32K context), the attention matrix is \(32768^2 \times 2 \approx 2\) GB per head – completely infeasible to store.

The Flash Attention Key Insight – Tiling and Online Softmax

The core insight is: we never need the full \(n \times n\) matrix in memory at once. We can compute the output \(O\) block-by-block if we use the online softmax trick.

Online Softmax.

Recall that softmax requires a global maximum for numerical stability:

\[ \text{softmax}(x_i) = \frac{e^{x_i - m}}{\sum_j e^{x_j - m}}, \quad m = \max_j x_j \]

The trick: we can update the running maximum and normalization factor as we process new blocks, without ever materializing the full row.

Important

Online Softmax Update Rule

Given a running state \((m_{\text{old}}, \ell_{\text{old}}, O_{\text{old}})\) and a new block of scores \(s_{\text{new}}\):

1. \(m_{\text{new}} = \max(m_{\text{old}},; \max(s_{\text{new}}))\)

2. [ \ell_{\text{new}} = e^{m_{\text{old}} - m_{\text{new}}} \cdot \ell_{\text{old}} + \sum_j e^{s_{\text{new},j} - m_{\text{new}}} ]

3. [ O_{\text{new}} = \frac{1}{\ell_{\text{new}}} \left( e^{m_{\text{old}} - m_{\text{new}}} \cdot \ell_{\text{old}} \cdot O_{\text{old}} + e^{s_{\text{new}} - m_{\text{new}}} \cdot V_{\text{new}} \right) ]

This is mathematically equivalent to computing softmax over all blocks at once.

The Flash Attention Algorithm

Note

Flash Attention Forward Pass – Block Tiling

Setup: SRAM size \(M\). Block sizes \(B_r = \lceil M / (4d) \rceil\), \(B_c = \min(\lceil M / (4d) \rceil, d)\).

1. Divide \(Q\) into \(T_r = \lceil n / B_r \rceil\) blocks \(Q_1, \ldots, Q_{T_r}\)

2. Divide \(K, V\) into \(T_c = \lceil n / B_c \rceil\) blocks \(K_1, \ldots, K_{T_c}\)

3. Initialize output \(O \in \mathbb{R}^{n \times d}\), running max \(m \in \mathbb{R}^n\), running sum \(\ell \in \mathbb{R}^n\) (all in HBM)

4. Outer loop over \(j = 1, \ldots, T_c\):

   1. Load \(K_j, V_j\) from HBM to SRAM

   2. Inner loop over \(i = 1, \ldots, T_r\):

      1. Load \(Q_i, O_i, m_i, \ell_i\) from HBM to SRAM

      2. Compute \(S_{ij} = Q_i K_j^T / \sqrt{d}\) (stays in SRAM)

      3. Apply online softmax update to get new \(m_i, \ell_i, O_i\)

      4. Write \(O_i, m_i, \ell_i\) back to HBM

5. Return \(O\)

Key: \(S_{ij}\) (the attention tile) is computed and discarded in SRAM. It is never written to HBM.

Important

Flash Attention Complexity

	Standard Attention	Flash Attention
Memory (HBM)	\(O(n^2)\)	\(O(n)\)
HBM reads/writes	\(O(n^2 d)\)	\(O(n^2 d / M)\)
FLOPs	\(O(n^2 d)\)	\(O(n^2 d)\) (same)
Speedup	1\(\times\)	2–4\(\times\)

In the forward pass, the total FLOPs remain \(O(n^2 d)\) – identical to standard attention. Flash Attention gains speed entirely by slashing slow HBM traffic, not by reducing arithmetic. (The backward pass actually performs more FLOPs due to recomputation, but the wall-clock time is still lower because the saved memory bandwidth dominates.)

Flash Attention 2 – Better Parallelism

Flash Attention 2 (Dao 2024) made three key improvements:

1. Reduced non-matmul FLOPs: The original FA had unnecessary rescaling operations in the inner loop. FA2 restructures the loop to minimize these. On A100, Tensor Core matrix multiplications outpace scalar operations by roughly 16\(\times\), so even a small fraction of non-matmul work in the inner loop becomes the latency bottleneck.

2. Better parallelism across sequence dimension: FA1 parallelized over batch and heads only. FA2 also parallelizes over the query sequence dimension, enabling better GPU utilization for long sequences with small batch sizes.

3. Causal masking optimization: For autoregressive (causal) attention, roughly half the blocks are fully masked. FA2 skips these blocks entirely, giving \(\sim\)2\(\times\) speedup for causal attention vs. bidirectional.

Flash Attention 3 – Hopper Architecture

Flash Attention 3 (Shah et al. 2024) is designed specifically for H100 and exploits three Hopper-specific features:

• TMA (Tensor Memory Accelerator): H100 has a dedicated hardware unit for asynchronous bulk data movement between HBM and SRAM. FA3 uses TMA to overlap data loading with computation, hiding memory latency.

• Warp-specialization: FA3 assigns different warps to different roles (producer warps load data via TMA; consumer warps compute MMA). This is a software pipelining technique that keeps both the memory system and Tensor Cores busy simultaneously.

• FP8 support: H100 supports FP8 (E4M3/E5M2) Tensor Core operations at 2\(\times\) the throughput of BF16. FA3 supports FP8 attention with per-block quantization to maintain accuracy.

FA3 achieves up to 75% of H100 theoretical peak for FP16 attention, compared to \(\sim\)35% for FA2.

Flash Attention 4 – Blackwell Architecture

Flash Attention 4 (Zadouri et al. 2026) targets NVIDIA’s Blackwell GPUs (B200/GB200), which double Tensor Core throughput to 2.25 PFLOP/s (BF16) while non-matmul units (exponential, shared memory bandwidth) scale at a slower rate. This asymmetric hardware scaling means that the bottleneck shifts: on Blackwell, attention is limited not by matmul but by the softmax exponentials and shared memory traffic surrounding them.

FA4 addresses this with four key techniques:

• Fully asynchronous MMA pipelines: Blackwell’s MMA instructions are fully asynchronous (unlike Hopper’s wgmma which still blocked on completion). FA4 redesigns the pipeline to overlap MMA, TMA loads, and softmax rescaling across larger tile sizes, keeping all hardware units saturated.

• Software-emulated exponential: Instead of calling the hardware ex2 unit (which is the throughput bottleneck), FA4 emulates \(e^x\) using polynomial approximations executed on the much faster Tensor Cores themselves. This trades extra matmul instructions for exponential-unit stalls.

• Conditional softmax rescaling: Standard FlashAttention rescales the running \(\max\) every tile. FA4 skips the rescaling when the new tile’s max does not exceed the running max (common in practice), saving both register shuffles and synchronization barriers.

• Tensor Memory + 2-CTA MMA mode (backward pass): The backward pass uses Blackwell’s Tensor Memory (a per-SM scratchpad larger than shared memory) and a 2-CTA cooperative mode that fuses \(dQ\) accumulation across two thread-block clusters, halving shared memory round-trips.

Important

FA4 Implementation: CuTe-DSL

FA4 is the first FlashAttention version written in CuTe-DSL, a Python-embedded domain-specific language for GPU kernels (part of CUTLASS 4.x). CuTe-DSL compiles 20–30\(\times\) faster than C++ CUTLASS templates while retaining full control over register allocation and pipeline scheduling. This dramatically lowers the iteration time for kernel development.

Results.

On B200 with BF16 head-dim 128 (causal, seq-len 8K):

• 1613 TFLOP/s – 71% of Blackwell peak utilization

• 1.3\(\times\) faster than cuDNN 9.13 (NVIDIA’s proprietary fused kernel)

• 2.7\(\times\) faster than Triton on the same hardware

Tip

Hardware–Software Co-evolution

The FlashAttention series illustrates a key principle: each GPU generation shifts the bottleneck, demanding new algorithmic ideas rather than just re-compilation. A80 \(\to\) memory bandwidth limited (FA1/FA2: tiling + recomputation). H100 \(\to\) data movement limited (FA3: TMA + warp-specialization). B200 \(\to\) non-matmul compute limited (FA4: software-emulated exp + conditional rescaling). Understanding where the hardware bottleneck lies is the prerequisite for writing efficient kernels.

Pretraining: Best Practices

Pretraining is the most expensive phase of LLM development—consuming millions of GPU-hours and requiring careful orchestration of data, compute, and hyperparameters. This section distills key lessons from Llama-3 (Grattafiori et al. 2024), Chinchilla (Hoffmann et al. 2022), and GPT-4 (OpenAI 2023).

Training Objective

All modern decoder-only LLMs use causal language modeling (CLM):

\[ \mathcal{L}_\text{CLM} = -\frac{1}{T}\sum_{t=1}^T \log P_\theta(x_t \mid x_{<t}) \]

This simple objective—with enough data and scale—produces emergent capabilities (in-context learning, reasoning, instruction following) without explicit supervision (Brown et al. 2020).

Data Pipeline

Important

Pretraining Data Recipe

• Scale: 1–15 trillion tokens for frontier models (Llama-3: 15T tokens)

• Sources: Web crawl (80%), code (10%), books/papers (5%), curated (5%)

• Deduplication: MinHash + exact substring dedup reduces memorization (Lee et al. 2022)

• Quality filtering: Perplexity-based classifier, heuristic filters (length, language ID, toxicity)

• Data mixing: Temperature-weighted sampling across domains; upweight code and math for reasoning

Scaling Laws

Hoffmann et al. (Hoffmann et al. 2022) showed that compute-optimal training requires balancing model size \(N\) and data size \(D\): \(N_\text{opt} \propto C^{0.50}\), \(D_\text{opt} \propto C^{0.50}\). A 70B model is compute-optimal at \(\sim\)1.4T tokens. In practice, models are over-trained (more tokens than Chinchilla-optimal) because inference cost scales with model size, not training tokens—smaller over-trained models are cheaper to deploy.

Key Hyperparameters

Pretraining hyperparameters from published models.

Common Failure Modes

Warning

Pretraining Pitfalls

• Loss spikes: Sudden loss increases from bad data batches or numerical instability. Llama-3 reports rolling back to checkpoints and skipping offending batches.

• Memorization: Model regurgitates training data verbatim. Fix: deduplicate aggressively; monitor extraction attacks.

• Context length: Training on short sequences then deploying at long context fails. Use continued pretraining on long documents + RoPE scaling.

Supervised Fine-Tuning (SFT)

SFT transforms a pretrained language model into an instruction-following assistant by training on curated prompt–response pairs. This is the bridge between raw language modeling and RLHF.

SFT Objective

The loss is identical to CLM, but computed only on response tokens:

\[ \mathcal{L}_\text{SFT} = -\frac{1}{|y|}\sum_{t=1}^{|y|} \log P_\theta(y_t \mid x_\text{prompt}, y_{<t}) \]

Prompt tokens provide context but receive no gradient (labels set to \(-100\)).

Data Quality: The LIMA Principle

Zhou et al. (C. Zhou et al. 2023) demonstrated that 1,000 carefully curated examples can match models trained on 50K+ noisy examples. Key requirements:

• Diversity: Cover QA, summarization, code, math, creative writing, multi-turn dialogue

• Correctness: Every response must be factually accurate and well-formatted

• Length balance: Mix short (1-sentence) and long (multi-paragraph) responses

• Decontamination: Remove overlap with evaluation benchmarks

Training Configuration

from trl import SFTTrainer, SFTConfig

sft_config = SFTConfig(
    output_dir="./sft_output",
    max_seq_length=4096,
    packing=True,              # Pack short examples into full sequences
    learning_rate=2e-5,
    lr_scheduler_type="cosine",
    warmup_ratio=0.1,
    weight_decay=0.01,
    max_grad_norm=1.0,
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,
    bf16=True,
    gradient_checkpointing=True,
)
trainer = SFTTrainer(model=model, args=sft_config,
                     train_dataset=dataset, processing_class=tokenizer)
trainer.train()

Efficient Training Solutions

Standard HuggingFace training leaves significant performance on the table. Several libraries provide drop-in efficiency gains for SFT workloads:

Liger Kernel (Hsu et al. 2024).

An open-source set of Triton-fused kernels from LinkedIn that replace standard PyTorch operators during training. Key fusions include:

• Fused Cross-Entropy: Merges the final linear projection, softmax, and loss computation into a single kernel—avoids materializing the full \((\text{batch} \times \text{seq} \times \text{vocab})\) logit tensor.

• Fused RMSNorm / SwiGLU / RoPE: Eliminates intermediate memory allocations for common LLM building blocks.

• Chunked operations: Processes large tensors in tiles to keep peak memory bounded.

Result: 20% higher throughput and up to 60% memory reduction with a one-line integration (apply_liger_kernel_to_llama()). Compatible with FSDP, DeepSpeed, and LoRA.

Unsloth (Han and Han 2024).

A specialized fine-tuning library that combines custom CUDA/Triton kernels with aggressive memory optimization:

• Manual backpropagation through LoRA layers (avoids autograd overhead).

• 4-bit QLoRA with fused dequantization—trains 70B models on a single 48 GB GPU.

• Intelligent RoPE and attention kernel fusion specific to each architecture (Llama, Mistral, Qwen, Gemma).

Result: 2–5\(\times\) faster than vanilla HuggingFace + PEFT, with 60–70% less VRAM. Particularly impactful for single-GPU and consumer-hardware workflows.

torchtune (P. Team 2024).

Meta’s native PyTorch fine-tuning library (development wound down in 2025), designed around composability rather than monolithic abstractions:

• Pure PyTorch—no trainer class; recipes are readable single-file scripts.

• Native integration with torch.compile, FSDP2, and activation checkpointing.

• First-class support for QLoRA, full fine-tuning, and knowledge distillation.

• Built-in quantization-aware training (QAT) for post-training compression.

Result: Comparable speed to custom solutions but with full debuggability and no framework lock-in.

Important

Choosing an Efficiency Stack

• Quick LoRA/QLoRA on \(\leq\)1 GPU: Unsloth (fastest time-to-train, minimal setup)

• Multi-GPU full fine-tune: TRL/DeepSpeed + Liger Kernel (best throughput at scale)

• Research / custom training loops: torchtune (transparent, hackable, native PyTorch)

These are complementary: Liger kernels can be used inside both TRL and torchtune workflows.

Best Practices

SFT training guidelines.

Tip

SFT Is Not Enough

SFT teaches format and basic instruction following, but cannot reliably teach: preference (which response is better—needs RLHF/DPO), refusal (when not to answer—needs safety training), calibration (saying “I don’t know”—needs RL with truthfulness rewards), or complex reasoning (multi-step chains—needs RL with verifiable rewards). The full pipeline is: Pretrain \(\to\) SFT \(\to\) RLHF/DPO.

LoRA and Parameter-Efficient Fine-Tuning

Full fine-tuning of a 70B model requires storing 70B trainable parameters plus their optimizer states (560+ GB of memory). LoRA (Hu et al. 2022b) (Low-Rank Adaptation) provides a way to fine-tune with \(<\)1% of the parameters while achieving comparable quality.

The LoRA Insight

Important

LoRA Core Idea

Instead of updating a full weight matrix \(W \in \mathbb{R}^{d \times d}\), learn a low-rank perturbation: \[ W' = W + \frac{\alpha}{r} \cdot BA, \quad B \in \mathbb{R}^{d \times r}, ; A \in \mathbb{R}^{r \times d} \]

• \(W\) is frozen (no gradients, no optimizer states)

• Only \(B\) and \(A\) are trained: \(2 \times d \times r\) parameters instead of \(d^2\)

• At rank \(r=16\), \(d=4096\): LoRA adds \(2 \times 4096 \times 16 = 131K\) params per layer vs. \(16.8M\) for full matrix

• \(\alpha/r\) scaling controls the magnitude of the update

Tip

Why Low-Rank Works

Aghajanyan et al. (Aghajanyan et al. 2021) showed that fine-tuning operates in a very low-dimensional subspace — the “intrinsic dimensionality” of the fine-tuning task is much smaller than the model’s parameter count. A 175B model’s fine-tuning task may have intrinsic dimensionality \(<\)10,000. LoRA exploits this directly: rank \(r\) constrains the update to an \(r\)-dimensional subspace per weight matrix.

Tip

Why the $\alpha/r$ Scaling Matters

Without scaling, doubling the rank \(r\) would roughly double the magnitude of \(\Delta W = BA\) (more columns in \(B\) contribute to the sum). This means changing rank would also change how much the model is perturbed—you’d need to re-tune the learning rate every time you adjust \(r\).

The \(\alpha/r\) factor normalizes the update magnitude so that it stays approximately constant regardless of rank: \[ W' = W + \frac{\alpha}{r} \cdot BA \]

• Fix \(\alpha\), sweep \(r\): The effective update magnitude stays \(\sim\alpha\) regardless of rank. You can try \(r \in \{8, 16, 32, 64\}\) without re-tuning LR.

• Common practice: Set \(\alpha = r\) (so \(\alpha/r = 1\)) or \(\alpha = 2r\) (so \(\alpha/r = 2\)). This is a convenient default where the scaling factor is a small integer.

• Why not just tune LR? You could, but \(\alpha/r\) provides a rank-independent knob. Teams can share LR recipes across experiments with different ranks.

• rsLoRA insight (Kalajdzievski 2023): At high ranks (\(r \geq 64\)), empirical evidence shows \(\alpha/\sqrt{r}\) is more stable than \(\alpha/r\), because the variance of \(BA\) scales with \(\sqrt{r}\), not \(r\).

LoRA Hyperparameters

Choosing LoRA hyperparameters correctly is critical — the wrong rank or alpha can either under-fit (too constrained) or waste memory (too expressive).

LoRA hyperparameter guide.

Warning

Rank Selection Rules of Thumb

• r=8: Simple tasks (single-domain chat, classification). Very memory-efficient.

• r=16: General-purpose fine-tuning. Good default.

• r=32–64: Complex tasks (math, code, multi-turn reasoning). Approaches full fine-tune quality.

• r=128+: Diminishing returns; consider full fine-tuning or QLoRA with higher rank.

• Key indicator: If training loss plateaus well above full fine-tune loss, increase rank.

LoRA Variants

LoRA variants and their innovations.

Key Extensions Explained

DoRA – Weight-Decomposed Low-Rank Adaptation.

DoRA (S.-Y. Liu et al. 2024) observes that full fine-tuning tends to change the direction of weight vectors more than their magnitude. Standard LoRA conflates both. DoRA decomposes each weight column into magnitude \(m = \\vert W\\vert _\text{col}\) and direction \(\hat{V} = W / \\vert W\\vert _\text{col}\), then applies LoRA only to the direction:

\[ W' = m \odot \hat{V}', \quad \hat{V}' = \frac{W + BA}{|W + BA|_\text{col}} \]

Magnitude \(m\) is a separate learnable vector (one scalar per column). This consistently outperforms LoRA by 1–3% on reasoning and instruction-following benchmarks with no additional inference cost (merged at deployment).

LoRA+ – Asymmetric Learning Rates.

Hayou et al. (Hayou et al. 2024) show that matrices \(A\) and \(B\) in LoRA have different optimal learning rates. Since \(B\) is initialized to zero, it starts in a very different regime than \(A\) (initialized from \(\mathcal{N}(0, \sigma^2)\)). Setting \(\eta_B \approx 16 \times \eta_A\) improves convergence speed and final quality by \(\sim\)2% — a free gain requiring only a one-line config change:

# LoRA+ in PEFT: set different LRs per matrix
optimizer_grouped_parameters = [
    {"params": [p for n, p in model.named_parameters() if "lora_B" in n],
     "lr": 2e-4 * 16},   # B matrix: higher LR
    {"params": [p for n, p in model.named_parameters() if "lora_A" in n],
     "lr": 2e-4},         # A matrix: base LR
]

VeRA – Vector-based Random Matrix Adaptation.

VeRA (Kopiczko et al. 2024) takes parameter efficiency to the extreme: instead of learning \(A\) and \(B\), it freezes them as shared random matrices across all layers and only trains two diagonal scaling vectors \(d_b \in \mathbb{R}^r\) and \(d_a \in \mathbb{R}^d\):

\[ \Delta W = B \cdot \text{diag}(d_b) \cdot A \cdot \text{diag}(d_a) \]

This reduces trainable parameters by \(\sim\)10\(\times\) vs. LoRA (only \(r + d\) params per layer) while achieving 90–95% of LoRA quality. Best for scenarios where you need hundreds of task-specific adapters with minimal storage.

Note

QLoRA Memory Savings

70B model full fine-tune: 140 GB (weights) + 280 GB (optimizer) + 140 GB (gradients) = 560 GB (7\(\times\) A100-80GB).

70B QLoRA (r=16, all linear layers):

• Base model in NF4: \(70\text{B} \times 0.5 = 35\) GB

• LoRA adapters in BF16: \(\sim\)160 MB

• Optimizer states (only for adapters): \(\sim\)320 MB

• Activations (gradient checkpointing): \(\sim\)8 GB

• Total: \(\sim\)44 GB — fits in a single 48GB GPU!

# QLoRA configuration with PEFT
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from transformers import BitsAndBytesConfig
import torch

# 4-bit quantization config
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",           # NormalFloat4 - optimal for weights
    bnb_4bit_compute_dtype=torch.bfloat16, # Compute in BF16
    bnb_4bit_use_double_quant=True,       # Quantize the quantization constants
)

# LoRA config
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,                        # alpha/r = 2x scaling
    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 = prepare_model_for_kbit_training(model)  # Prepare for QLoRA
model = get_peft_model(model, lora_config)       # Add LoRA adapters
model.print_trainable_parameters()
# Output: trainable params: 83,886,080 || all params: 70,553,706,496 || 0.12%

Other PEFT Approaches

LoRA dominates modern practice, but it is not the only parameter-efficient method. For completeness, the main alternatives:

PEFT method families. LoRA is the de facto standard for LLM fine-tuning; the others are included for historical context and niche use cases.

Tip

Why LoRA Won

LoRA became the standard because it uniquely combines: (1) zero inference overhead — adapters merge into base weights, unlike Adapters or Prefix Tuning which add latency or consume context; (2) composability — multiple LoRA adapters can be swapped at serving time for multi-tenant deployments; (3) ecosystem support — HuggingFace PEFT, TRL, vLLM, and all major frameworks have first-class LoRA support; (4) proven at scale — used in production by Meta, Google, and most open-source fine-tunes on HuggingFace. Unless you have a specific constraint that LoRA cannot satisfy, it should be your default choice.

Mixture of Experts (MoE)

Mixture of Experts models (Shazeer et al. 2017; Jiang et al. 2024) scale model capacity without proportionally scaling compute cost by activating only a subset of parameters for each token.

Architecture

Important

MoE Layer

In a MoE transformer, the FFN layer in each block is replaced by \(N\) parallel “expert” FFNs plus a router that selects which experts to use: \[ \text{MoE}(x) = \sum_{i=1}^{N} g_i(x) \cdot E_i(x), \quad g(x) = \text{TopK}(\text{softmax}(W_r x)) \]

• \(E_i\) are expert networks (standard FFN layers)

• \(g_i(x)\) are gating weights from the router (only top-\(K\) are non-zero)

• Typically \(K=2\) out of \(N=8\)–64 experts are active per token

• Total params scale with \(N\); active params scale with \(K/N\) of FFN size

Load Balancing

Warning

The Load Balancing Problem

Without constraints, the router may send most tokens to the same 1–2 experts (“expert collapse”). This wastes capacity and creates compute imbalance across GPUs (each expert typically lives on a different GPU).

Solution: Add an auxiliary load-balancing loss: \[ \mathcal{L}_{\text{bal}} = \alpha \cdot N \sum_{i=1}^{N} f_i \cdot p_i \] where \(f_i\) = fraction of tokens routed to expert \(i\), \(p_i\) = mean router probability for expert \(i\). This encourages uniform expert utilization.

Noisy Top-K Gating: Making Discrete Routing Trainable

The core challenge in MoE is that top-\(k\) selection is not differentiable — you can’t backpropagate through a hard “pick the top 2” operation. The field has developed two key tricks to solve this:

Tip

The Routing Differentiability Problem

The router computes logits \(h(x) = W_r \cdot x\) for each expert, then selects the top-\(k\). But:

• The selected experts get gradients through their gate weights (softmax over selected)

• The selection decision itself (which \(k\) to pick) has zero gradient

• Without a trick, the router can get stuck: an expert never selected \(\rightarrow\) never gets a gradient signal \(\rightarrow\) never gets selected

Approach 1: Noisy Top-K Gating (Shazeer et al. 2017).

Add learnable Gaussian noise to the router logits before the top-\(k\) selection:

\[ \begin{aligned} h(x) &= W_g \cdot x && \text{(clean logits)} \\ H(x) &= h(x) + \epsilon \cdot \text{Softplus}(W_{\text{noise}} \cdot x), \quad \epsilon \sim \mathcal{N}(0, 1) && \text{(noisy logits)} \\ \text{TopK}(v, k)_i &= \begin{cases} v_i & \text{if } v_i \text{ is in the top } k \\ -\infty & \text{otherwise} \end{cases} \\ g(x) &= \text{softmax}\big(\text{TopK}(H(x),, k)\big) && \text{(sparse gates)} \end{aligned} \]

• \(W_{\text{noise}}\) is a learned noise magnitude — the model learns how much exploration each expert needs

• During training, noise occasionally promotes “underdog” experts into the top-\(k\), giving them gradient signal

• At inference, noise is removed: use clean logits \(h(x)\) for deterministic routing

• The Softplus ensures noise scale is always positive

Approach 2: Gumbel-Softmax Trick (for differentiable discrete sampling).

An alternative from the variational inference literature (Jang et al. 2017). The Gumbel-Max trick provides exact sampling from a categorical distribution:

\[ z = \arg\max_i \left[ \log \pi_i + G_i \right], \quad G_i \sim \text{Gumbel}(0,1) \]

where Gumbel noise is generated as \(G_i = -\log(-\log(U_i)),; U_i \sim \text{Uniform}(0,1)\).

For top-\(k\) routing: taking the top-\(k\) of \((\log \pi_i + G_i)\) gives \(k\) samples without replacement from the categorical distribution defined by \(\pi\).

Since \(\arg\max\) is non-differentiable, the Gumbel-Softmax relaxation replaces it with a temperature-controlled softmax:

\[ \hat{g}_i = \frac{\exp\left((\log \pi_i + G_i) / \tau\right)}{\sum_j \exp\left((\log \pi_j + G_j) / \tau\right)} \]

• \(\tau \to 0\): approaches a hard one-hot (exact but non-differentiable)

• \(\tau \to \infty\): approaches uniform (differentiable but uninformative)

• In practice, anneal \(\tau\) from 1.0 down to 0.1–0.5 during training

• Straight-through estimator: use hard top-\(k\) in the forward pass, but Gumbel-Softmax gradients in the backward pass — best of both worlds

Important

Which Approach Is Used in Practice?

• Sparsely-Gated MoE (Shazeer et al. 2017), Mixtral (Jiang et al. 2024), DeepSeek-V2 (DeepSeek-AI 2024a): Use Noisy Top-K with Gaussian noise. Simple, effective, well-proven at scale.

• Switch Transformer (Fedus et al. 2022): Simplified to Top-1 with no noise (relies on load-balancing loss alone).

• Research / smaller-scale MoE: Some use Gumbel-Softmax for fully differentiable routing, especially when learning the routing itself is the research objective.

• Key insight: Both approaches solve the same problem (making discrete selection trainable) via noise injection. Gaussian noise is simpler; Gumbel noise has stronger theoretical guarantees for categorical sampling.

Notable MoE Models

Diversity in LLM Training

Diversity — in training data, model outputs, and optimization trajectories — is critical for preventing mode collapse and ensuring robust, general-purpose LLMs. This section covers the key diversity mechanisms applicable to all LLM training phases.

Sampling Diversity

Important

Sampling Strategies for Diverse Generation

• Temperature \(\tau\): \(P(x_i) \propto \exp(\text{logit}_i / \tau)\). Higher \(\tau\) = more uniform distribution = more diverse. Typical: \(\tau=0.7\)–\(1.0\) for RLHF generation.

• Top-\(k\): Only sample from the \(k\) highest-probability tokens. Prevents degenerate low-probability tokens.

• Top-\(p\) (nucleus): Sample from the smallest set of tokens whose cumulative probability \(\geq p\). Adaptive: more diverse when the model is uncertain.

• Min-\(p\): Only keep tokens with \(P \geq p_{\min} \times P_{\max}\). More principled than top-\(k\).

• Frequency/presence penalty: Penalize tokens that appeared in the response. Encourages lexical diversity.

Training Data Diversity

• Prompt diversity: Cover different domains, difficulty levels, and formats. The Goldilocks principle: prompts should have 20–80% success rate.

• Deduplication: Remove near-duplicate training examples (MinHash, n-gram overlap). Duplicates cause overfitting to specific patterns.

• Data mixing: Balance across tasks/domains using temperature-weighted sampling or curriculum strategies.

Diversity-Promoting Methods

Text Generation: Decoding Methods

A trained language model outputs a probability distribution over the vocabulary at each step: \(P(x_t \vert x_{<t})\). The decoding strategy determines how we select the next token from this distribution. This choice profoundly affects output quality, diversity, and coherence.

Greedy Decoding

The simplest strategy: always pick the highest-probability token.

\[ x_t = \arg\max_{v \in \mathcal{V}} P(v | x_{<t}) \]

Intuition: Like always taking the most obvious next word in a sentence. “The capital of France is...” \(\to\) “Paris” (probability 0.92).

Pros: Deterministic, fast, no hyperparameters.
Cons: Produces repetitive, generic text. Misses high-quality sequences where an early low-probability token leads to a globally better output. No diversity.

Beam Search

Maintain \(B\) (beam width) partial hypotheses in parallel, expanding each by the top-\(k\) tokens and keeping the \(B\) highest-scoring complete sequences:

\[ \text{score}(y_{1:t}) = \sum_{i=1}^{t} \log P(y_i | y_{<i}) \]

With length normalization to avoid favoring short sequences:

\[ \text{score}_{\text{norm}}(y) = \frac{1}{|y|^\alpha} \sum_{i=1}^{|y|} \log P(y_i | y_{<i}), \quad \alpha \in [0.6, 1.0] \]

Intuition: Like exploring multiple paths in a maze simultaneously, keeping only the \(B\) most promising ones at each junction.

Pros: Finds higher-likelihood sequences than greedy; good for translation and summarization where there’s a single “correct” output.
Cons: Still tends toward generic, repetitive text for open-ended generation; \(B \times\) more compute; all beams often converge to similar outputs.

Diverse Beam Search

Standard beam search produces near-duplicate beams. Diverse beam search (Vijayakumar et al. 2018) partitions beams into \(G\) groups and adds a dissimilarity penalty between groups:

\[ \text{score}_g(y_t) = \log P(y_t | y_{<t}) - \lambda \sum_{g'<g} \Delta(y_t, Y_{g'}) \]

where \(\Delta\) measures overlap (e.g., Hamming diversity) with tokens already selected by earlier groups, and \(\lambda\) controls diversity strength.

Intuition: Like forcing a brainstorming group to generate different ideas — each subgroup is penalized for repeating what earlier subgroups said.

Pros: Produces genuinely different candidate sequences; useful for reranking pipelines.
Cons: Diversity penalty can degrade individual beam quality; more hyperparameters (\(G\), \(\lambda\)).

Top-\(k\) Sampling

Sample from only the \(k\) most probable tokens, redistributing probability mass:

\[ P'(v | x_{<t}) = \begin{cases} \dfrac{P(v | x_{<t})}{\sum_{v' \in \text{Top-}k} P(v' | x_{<t})} & \text{if } v \in \text{Top-}k \\[6pt] 0 & \text{otherwise} \end{cases} \]

Intuition: After “The cat sat on the...”, only consider the top \(k\) plausible continuations (“mat”, “floor”, “couch”, ...) and ignore extremely unlikely ones (“quantum”, “archipelago”).

Pros: Removes tail noise; simple to implement.
Cons: Fixed \(k\) is too restrictive for peaked distributions (wastes probability mass) and too permissive for flat distributions (lets in garbage tokens).

Top-\(p\) (Nucleus) Sampling

Sample from the smallest set of tokens whose cumulative probability exceeds \(p\):

\[ \text{Top-}p = \min \left\{ S \subseteq \mathcal{V} : \sum_{v \in S} P(v | x_{<t}) \geq p \right\} \]

where tokens are sorted by descending probability and added until the threshold \(p\) is reached.

Intuition: Adaptively resize the candidate pool. If the model is confident (“Paris” at 95%), the nucleus is tiny. If uncertain (“The movie was...”), the nucleus expands to include many plausible adjectives.

Pros: Adapts to distribution shape; widely used default (\(p=0.9\)–\(0.95\)).
Cons: Still includes some low-quality tokens at the tail of the nucleus; the threshold is a single global hyperparameter.

Tip

Top-kk vs. Top-pp

Consider predicting the next word:

• After “2 + 2 =”: distribution is peaked — top-1 token (“4”) has 99% mass. Top-\(k\)=50 wastefully considers 49 wrong answers. Top-\(p\)=0.9 correctly picks just “4”.

• After “I enjoy eating”: distribution is flat — many foods are plausible. Top-\(k\)=5 is too restrictive. Top-\(p\)=0.9 might include 50+ tokens, matching the actual uncertainty.

Top-\(p\) adapts; top-\(k\) doesn’t. In practice, both are often combined: sample from top-\(p\) intersected with top-\(k\).

Min-\(p\) Sampling

A recent alternative that sets a relative probability floor (Nguyen 2024):

\[ \text{Min-}p = \left\{ v \in \mathcal{V} : P(v | x_{<t}) \geq p_{\min} \cdot \max_{v'} P(v' | x_{<t}) \right\} \]

Only tokens with probability at least \(p_{\min}\) times the top token’s probability are kept.

Intuition: “Only consider tokens that are at least 10% as likely as the best token.” If the top token has probability 0.8, only tokens above 0.08 survive. If the top token has probability 0.05 (very uncertain), tokens above 0.005 survive — naturally expanding the pool.

Pros: Scales naturally with model confidence; fewer degenerate samples than top-\(p\) on peaked distributions; single intuitive parameter.
Cons: Newer, less battle-tested; not yet standard in all inference frameworks.

Temperature Scaling

Before applying any sampling strategy, logits are divided by temperature \(T\):

\[ P_T(v | x_{<t}) = \frac{\exp(z_v / T)}{\sum_{v'} \exp(z_{v'} / T)} \]

• \(T < 1\): Sharpens distribution \(\to\) more deterministic, focused outputs.

• \(T = 1\): Unmodified model distribution.

• \(T > 1\): Flattens distribution \(\to\) more random, creative outputs.

• \(T \to 0\): Becomes greedy decoding. \(T \to \infty\): Becomes uniform sampling.

Common settings: \(T=0.7\) for factual tasks, \(T=1.0\)–\(1.2\) for creative writing, \(T=0.0\) (greedy) for code/math.

Contrastive Decoding

Contrastive decoding (X. Li et al. 2023) exploits the difference between a strong model (expert) and a weak model (amateur) to amplify the expert’s unique knowledge:

\[ x_t = \arg\max_{v \in \mathcal{V}(x_{<t})} \left[ \log P_{\text{expert}}(v | x_{<t}) - \log P_{\text{amateur}}(v | x_{<t}) \right] \]

where \(\mathcal{V}(x_{<t}) = \{v : P_{\text{expert}}(v \vert x_{<t}) \geq \alpha \cdot \max_{v'} P_{\text{expert}}(v' \vert x_{<t})\}\) is an adaptive plausibility constraint.

Intuition: The amateur model captures generic, obvious patterns (common words, repetition). Subtracting its log-probabilities removes this “generic signal,” leaving the expert’s distinctive knowledge and reasoning. Like removing background noise from a recording to hear the signal.

Pros: Reduces repetition and generic phrasing; improves factuality and coherence without additional training; works with any model pair.
Cons: Requires running two models (2\(\times\) compute); sensitive to amateur model choice; the plausibility threshold \(\alpha\) needs tuning.

Repetition Penalties

Orthogonal to the sampling strategy, repetition penalties discourage the model from repeating tokens. Given the raw logit \(z_v\) for token \(v\) (i.e., the unnormalized score output by the LM head before softmax), the penalized logit is:

\[ z_v' = \begin{cases} z_v / \theta & \text{if } v \in \text{generated tokens and } z_v > 0 \\ z_v \cdot \theta & \text{if } v \in \text{generated tokens and } z_v < 0 \end{cases} \]

where \(\theta > 1\) is the penalty factor (typically 1.1–1.3). In both cases, the effect is to push the logit toward zero—reducing the probability of previously generated tokens. Frequency and presence penalties are simpler additive variants used by OpenAI APIs:

\[ z_v' = z_v - \alpha \cdot \text{count}(v) - \beta \cdot \mathbf{1}[v \in \text{generated}] \]

where \(\alpha\) is the frequency penalty (proportional to how many times \(v\) appeared) and \(\beta\) is the presence penalty (flat penalty for any prior occurrence).

Practical Comparison

Decoding method comparison for LLM text generation.

Note

Decoding in Practice: “Once upon a time”

Given the prompt “Once upon a time,”:

• Greedy: “there was a young girl who lived in a small village...” (generic fairy tale)

• Top-\(p\)=0.9, \(T\)=1.0: “the rivers ran backwards and the fish learned to fly...” (creative, surprising)

• Top-\(p\)=0.9, \(T\)=0.3: “there was a kingdom ruled by a wise and just king...” (coherent, conventional)

• Contrastive: “in the amber-lit corridors of a collapsing star, two minds argued about the nature of time...” (distinctive, avoids clichés)

Same model, same prompt — decoding strategy determines the character of the output.

Constrained Decoding (Structured Generation)

All methods above sample from the full vocabulary at each step. Constrained decoding restricts the set of allowed tokens so that the output is guaranteed to conform to a formal grammar—typically a JSON schema, regex, or context-free grammar (CFG).

Core mechanism.

At each decoding step \(t\), a token mask \(M_t \subseteq \mathcal{V}\) is computed from the current parser state. Only tokens in \(M_t\) receive their original logits; all others are set to \(-\infty\) before softmax:

\[ P'(v | x_{<t}) = \begin{cases} P(v | x_{<t}) / Z & \text{if } v \in M_t \\ 0 & \text{otherwise} \end{cases} \]

where \(Z = \sum_{v \in M_t} P(v \vert x_{<t})\) renormalizes. Because the mask changes every step (it depends on what has been generated so far), the constraint is enforced incrementally—the model cannot produce an invalid prefix at any point.

From schema to mask.

The compilation pipeline is:

\[ \text{JSON Schema} ;\xrightarrow{\text{compile}}; \text{Regex} ;\xrightarrow{\text{compile}}; \text{FSM (DFA)} ;\xrightarrow{\text{index}}; \text{Token Mask per State} \]

The FSM states correspond to positions in the regex. For each state, all vocabulary tokens that would keep the string in the language are precomputed into an index (a one-time cost per schema). At runtime, looking up the mask is an \(O(1)\) table access—adding negligible latency to each decoding step.

Key libraries.

• Outlines (Willard and Louf 2023): Compiles JSON schemas and regexes into interleaved FSM-guided generation. Supports any model with a logits interface.

• lm-format-enforcer¹: Similar FSM approach with a focus on integration with serving frameworks (vLLM, TGI).

• Guidance² (Microsoft): Interleaves constrained generation with control flow (loops, conditions), enabling complex structured outputs beyond flat schemas.

• XGrammar (Dong et al. 2024): Pushdown-automaton-based engine supporting full context-free grammars (not just regular languages), used in MLC-LLM and vLLM for grammar-mode decoding.

Trade-offs.

Constrained decoding guarantees syntactic validity—no post-hoc parsing failures, no retries. However:

• Semantic quality: Forcing structure can degrade content quality if the model’s probability mass for the “correct” answer lies outside the grammar. In practice this is rare for well-trained models on well-designed schemas.

• Compilation cost: The FSM index must be built per schema. For complex schemas this can take 1–5 s, but it is amortized over all requests using that schema.

• Grammar coverage: Regex/FSM handles JSON, YAML, SQL fragments, and most structured formats. Full CFGs (via XGrammar or LALR parsers) cover languages like Python or XML.

Important

When to Use Constrained Decoding

Use constrained decoding whenever the consumer of the model’s output is a program rather than a human. Tool-calling agents, API backends, and data-extraction pipelines all benefit from guaranteed valid structure. For free-form prose or creative text, unconstrained sampling remains appropriate.

Prompt Engineering

Prompt engineering is the discipline of designing inputs to LLMs that reliably elicit desired behaviour—without changing model weights. While fine-tuning modifies the model, prompt engineering exploits the model’s existing capabilities through careful framing, examples, and structure. It is the fastest, cheapest, and most accessible lever for improving LLM outputs, and remains essential even when using fine-tuned models.

In-Context Learning (ICL)

In-context learning (Brown et al. 2020) is the remarkable ability of large language models to learn tasks at inference time purely from examples provided in the prompt—with no gradient updates. The model implicitly infers the task from the pattern of input–output pairs and generalizes to new inputs.

Important

Why In-Context Learning Works

• Implicit Bayesian inference: The model has seen millions of tasks during pretraining. The prompt examples locate the relevant task in the model’s learned distribution (Xie et al. 2022).

• Induction heads: Specific attention heads learn to copy patterns (“A is to B as C is to ”), enabling in-context generalization (Olsson et al. 2022).

• Task vectors: ICL creates implicit task representations in the residual stream that steer generation toward the demonstrated format and content (Todd et al. 2024).

Scaling behaviour.

ICL emerges primarily in models above \(\sim\)1B parameters and improves log-linearly with model scale (Brown et al. 2020). Smaller models can memorize examples but struggle to generalize to novel inputs within the same context window.

Zero-Shot Prompting

Zero-shot prompting provides no examples—only a task description or instruction. The model must rely entirely on its pretrained knowledge and instruction-tuning to produce the correct format and content.

Note

Zero-Shot Classification

Classify the following movie review as POSITIVE or NEGATIVE.

Review: "The cinematography was breathtaking but the plot 
felt rushed and predictable."

Sentiment:

When zero-shot works well:

• Tasks the model has seen extensively during pretraining/SFT (translation, summarization, sentiment)

• Well-specified instructions with unambiguous output format

• Instruction-tuned models (e.g., ChatGPT, Claude, Llama-3-Instruct) significantly outperform base models at zero-shot tasks (Ouyang et al. 2022)

When zero-shot fails:

Novel formats, domain-specific labeling schemes, or ambiguous tasks where the model cannot infer your exact requirements from the instruction alone.

Few-Shot Prompting

Few-shot prompting (Brown et al. 2020) provides \(k\) input–output examples (“shots”) before the actual query. This is the most common form of in-context learning and remains one of the most effective prompting strategies.

Note

Few-Shot Named Entity Recognition

Extract named entities from the text. Format: ENTITY

Text: "Apple released the iPhone 15 in Cupertino."
Entities: [Apple](ORG), [iPhone 15](PRODUCT), [Cupertino](LOC)

Text: "Elon Musk announced Tesla's new factory in Berlin."
Entities: [Elon Musk](PER), [Tesla](ORG), [Berlin](LOC)

Text: "OpenAI partnered with Microsoft to deploy GPT-4."
Entities:

Key design principles for few-shot examples:

1. Diversity: Cover the range of expected inputs (different lengths, edge cases, categories).

2. Ordering: Place harder or more representative examples last (recency bias) (Lu et al. 2022).

3. Label balance: If classifying, include examples from all classes to avoid majority-class bias.

4. Format consistency: Every example must follow the exact same structure. The model mimics the pattern.

5. Relevance: Use examples semantically similar to the target query for best results (J. Liu et al. 2022).

How many shots?

Performance typically improves from 0 to 4–8 examples, then plateaus. Beyond \(\sim\)20 examples, gains are marginal and you risk filling the context window. Min et al. (Min et al. 2022) showed that the format and label space of examples matter more than label correctness—even random labels help (though correct labels help more).

Instruction-Following Prompts

Instruction-tuned models respond best to clear, structured instructions. The key insight: treat the prompt as a specification, not a suggestion.

Important

Anatomy of an Effective Instruction Prompt

1. Role/Persona: Define who the model is (“You are a senior data scientist...”)

2. Task: What to do, stated clearly and unambiguously

3. Context: Background information the model needs

4. Constraints: Length limits, tone, what to avoid, output format

5. Examples (optional): Show the desired output format

6. Input: The actual data to process

Note

Instruction Prompt with Constraints

Role: You are a medical literature reviewer.

Task: Summarize the following research abstract for a 
general audience.

Constraints:
- Maximum 3 sentences
- No jargon (explain any technical terms)
- Include the key finding and its clinical implication
- Do NOT speculate beyond what the abstract states

Abstract: [...]

System prompts vs. user prompts.

Modern chat APIs separate the system prompt (persistent instructions, role definition) from the user message (per-turn input). System prompts are processed with higher attention priority in most models and provide a natural place for role definitions, constraints, and output format specifications (OpenAI 2023).

Structured Output Prompts (JSON/XML)

For programmatic use, the most critical prompting technique is enforcing structured output—particularly JSON.

Note

JSON Output Prompt

Extract the following information from the customer email. Respond ONLY with valid JSON, no other text.

Schema:
{
  "intent": "refund|complaint|question|praise",
  "urgency": "low|medium|high",
  "product_mentioned": "string or null",
  "summary": "one sentence summary"
}

Email: [...]

Techniques for reliable structured output:

• Schema-first: Show the exact JSON schema before the input. The model treats it as a template.

• Constrained decoding: Use grammar-based sampling (e.g., Outlines (Willard and Louf 2023), Guidance) to guarantee syntactically valid JSON at the token level.

• XML tags: For nested or multi-part outputs, XML tags (e.g., <thinking>...</thinking>) provide unambiguous delimiters that models follow reliably.

• Pydantic/TypeScript types: Providing type definitions helps models understand field constraints (OpenAI’s function calling uses JSON Schema internally).

Warning

JSON in Prompts — Common Pitfalls

• Models may add markdown fences (‘‘‘json ... ‘‘‘) — instruct explicitly to output raw JSON.

• Nested objects and arrays increase hallucination risk — flatten schemas where possible.

• Enum fields (fixed choices) are much more reliable than free-text fields.

• Always validate outputs programmatically; no prompt guarantees 100% compliance without constrained decoding.

JSON Prompting: Structuring the Input.

A distinct but complementary technique is JSON prompting—formatting the prompt itself as JSON rather than natural language. This exploits the model’s extensive pre-training on structured data (APIs, configs, code) to improve instruction adherence, reduce ambiguity, and enable deterministic parsing of multi-field requests.

Note

JSON Prompting with System Prompt

=== SYSTEM === You are a senior code reviewer. Analyze code for bugs, security issues, and style violations. Always respond in the JSON schema provided.

=== USER (JSON prompt) ===
{
  "task": "code_review",
  "language": "python",
  "severity_filter": "high",
  "code": "def login(user, pw):\n    query = ...",
  "output_schema": {
    "issues": [{
      "line": "int",
      "severity": "critical|high|medium|low",
      "category": "security|bug|style|performance",
      "description": "string",
      "fix": "string"
    }],
    "overall_risk": "critical|high|medium|low"
  }
}

Why JSON prompting works:

• Unambiguous field boundaries: No confusion about where one instruction ends and another begins.

• Typed constraints: Fields like "severity_filter": "high" are clearer than “only show high severity issues.”

• Schema-as-contract: Including output_schema in the input mirrors API design patterns the model has seen extensively during pre-training.

• System prompt still essential: The system message provides role, tone, and behavioral constraints that don’t fit naturally in a JSON payload.

Chain-of-Thought (CoT) Prompting

Chain-of-thought prompting (Wei et al. 2022) asks the model to produce intermediate reasoning steps before giving a final answer. This simple technique dramatically improves performance on tasks requiring multi-step reasoning: arithmetic, logic, commonsense inference, and code generation.

Why CoT works:

• Serializes computation: Transformers have fixed depth but variable-length generation. CoT converts parallel (hard) problems into sequential (easy) steps, effectively increasing the model’s computational budget.

• Reduces compounding errors: Each step is a simpler sub-problem with lower per-step error rate.

• Exposes intermediate state: Makes reasoning auditable and debuggable.

Important

Chain-of-Thought Variants

Method	Description
Zero-shot CoT (Kojima et al. 2022)	Append “Let’s think step by step” to any prompt
Few-shot CoT (Wei et al. 2022)	Provide examples with explicit reasoning chains
Self-Consistency (X. Wang, Wei, Schuurmans, Q. Le, et al. 2023)	Sample \(N\) CoT paths; majority-vote the final answer
Tree of Thoughts (S. Yao, Yu, et al. 2023)	Explore multiple reasoning branches with backtracking
Plan-and-Solve (L. Wang et al. 2023)	First plan the steps; then execute each step
ReAct (S. Yao, Zhao, et al. 2023)	Interleave Reasoning and Acting (tool use)

Note

Zero-Shot Chain-of-Thought

Q: A store has 45 apples. They sell 3/5 of them in the morning and 1/3 of the remaining in the afternoon. How many apples are left?

Let's think step by step.

A: Morning sales: 45 * 3/5 = 27 apples sold.
Remaining after morning: 45 - 27 = 18.
Afternoon sales: 18 * 1/3 = 6 apples sold.
Remaining: 18 - 6 = 12 apples.

Self-Consistency.

Wang et al. (X. Wang, Wei, Schuurmans, Q. Le, et al. 2023) showed that sampling multiple chain-of-thought reasoning paths and taking a majority vote over final answers significantly outperforms single-path CoT. The intuition: correct reasoning paths tend to converge on the same answer, while errors are typically idiosyncratic. This trades compute (generating \(N\) samples) for accuracy—practical when latency is less important than correctness.

When CoT hurts.

CoT is not universally beneficial. For simple tasks (single-step classification, retrieval, formatting), CoT adds unnecessary tokens, increases latency, and can even introduce errors through overthinking. Use CoT selectively for tasks where you expect multi-step reasoning to be required.

Advanced Prompting Techniques

Retrieval-Augmented Generation (RAG).

Rather than relying solely on the model’s parametric memory, RAG (P. Lewis et al. 2020) retrieves relevant documents and includes them in the prompt:

Context (retrieved): [document chunks]
Question: [user query]
Answer based ONLY on the provided context.

This grounds the model’s responses in verifiable sources and dramatically reduces hallucinations for knowledge-intensive tasks.

Prompt Chaining and Decomposition.

Complex tasks benefit from being broken into a pipeline of simpler prompts, where the output of one becomes the input to the next:

1. Extract key facts from document

2. Reason over extracted facts

3. Format final answer

Each step can use a different prompt template, model, or temperature setting. This is more controllable than a single monolithic prompt and enables targeted debugging.

Constitutional AI / Self-Critique.

Bai et al. (Bai et al. 2022) introduce prompts that ask the model to critique and revise its own output against a set of principles:

[Generate initial response]
Critique: Does this response violate any of the following 
principles? [list principles]
Revision: Rewrite the response addressing the critique.

Meta-Prompting and Prompt Optimization.

Rather than hand-crafting prompts, recent work automates prompt design:

• APE (Y. Zhou et al. 2023): Uses an LLM to generate and score candidate prompts automatically.

• DSPy (Khattab et al. 2024): Compiles declarative task descriptions into optimized prompt pipelines with learned few-shot examples.

• OPRO (C. Yang et al. 2024): Treats prompt optimization as an optimization problem, using an LLM as the optimizer.

Attentive Reasoning Queries (ARQ).

ARQ (Yang et al. 2025) addresses a fundamental weakness of standard prompting: as context length grows, models increasingly “lose” critical information in the middle of the prompt (the lost-in-the-middle effect). ARQ mitigates this by decomposing a complex query into multiple focused sub-queries, each designed to direct the model’s attention to a specific part of the context:

1. Query decomposition: Break the user question into atomic sub-questions that each target a narrow aspect.

2. Attentive retrieval: For each sub-query, retrieve or highlight only the relevant context slice—forcing the model to attend to it.

3. Aggregation: Combine sub-answers into a coherent final response.

This is particularly effective for long-document QA, multi-hop reasoning over large retrieval sets, and agentic tasks where the context window contains many tool outputs. ARQ can be seen as a structured form of chain-of-thought that explicitly manages where the model looks, not just how it reasons.

Best Practices: Crafting Effective Prompts

Based on empirical findings across the literature and practitioner experience, the following principles reliably improve prompt quality:

Important

The Prompt Engineering Checklist

1. Be specific and unambiguous: Replace “summarize this” with “summarize in 2–3 bullet points, each under 20 words, focusing on actionable findings.”

2. Show, don’t tell: One good example is worth 100 words of instruction. When in doubt, add a few-shot example.

3. Define the output format explicitly: Specify JSON schema, bullet points, table format, or exact delimiters. Never leave format to interpretation.

4. Use delimiters for input data: Wrap user inputs in clear delimiters (""", <input>...</input>, ---) to separate instructions from data.

5. Assign a role: “You are a [domain expert] who [specific behaviour]” primes relevant knowledge and tone.

6. Specify what NOT to do: Negative constraints (“do not explain your reasoning”, “never output more than 5 items”) are often more effective than positive ones.

7. Add chain-of-thought for reasoning tasks: Append “Think step by step” or provide worked examples for math, logic, or multi-hop questions.

8. Control temperature appropriately: Use \(T \approx 0\) for factual/deterministic tasks; \(T \approx 0.7\)–\(1.0\) for creative/diverse outputs.

9. Iterate empirically: Treat prompts as code—version them, A/B test them, and measure performance on representative eval sets.

10. Leverage recency bias: Place the most critical instructions and examples at the end of the prompt (closest to the generation point).

Common prompting failure modes and solutions.

Tip

The Prompt Engineering Mindset

Think of prompt engineering as programming in natural language. The model is a powerful but literal interpreter—it will do exactly what you ask, interpreted in the most likely way given its training distribution. Common principles from software engineering apply:

• DRY (Don’t Repeat Yourself): Unless fighting attention decay in long contexts

• Separation of concerns: Different prompt sections for role, constraints, examples, and input

• Test-driven development: Define expected outputs before writing the prompt

• Version control: Track prompt iterations and their eval scores

• Modularity: Build reusable prompt templates; parameterize variable parts

When prompting fails to achieve the desired quality after systematic iteration, that is the signal to move to fine-tuning (SFT) or reinforcement learning (RLHF/DPO).

Model Compression Methods

Model compression reduces model size and inference cost while preserving quality. Three main approaches: quantization (reduce precision), pruning (remove parameters), and distillation (train a smaller model to mimic a larger one).

Quantization

Quantization reduces model size and inference cost by representing weights (and optionally activations) in lower-precision formats. The core trade-off is compression ratio versus quality degradation.

Important

Quantization Overview

Quantization reduces the numerical precision of model weights (and optionally activations) from FP32/BF16 to lower-bit formats: \[ x_q = \text{round}!\left(\frac{x - z}{s}\right), \quad x_{\text{dequant}} = s \cdot x_q + z \] where \(s\) is the scale factor and \(z\) is the zero-point.

Quantization methods for LLMs.

Tip

When to Quantize

• Inference serving: Always quantize. W4A16 (4-bit weights, BF16 activations) is the sweet spot — 2\(\times\) memory savings, \(<\)1% quality loss for 70B+ models.

• Training: FP8 on H100 gives 2\(\times\) throughput with minimal quality loss. BF16 is still the default for smaller models.

• Edge deployment: GGUF Q4_K_M for local inference on consumer hardware.

• RLHF: Quantize the frozen models (reference, reward model) to INT8/FP8. Keep the policy in BF16 for training precision.

Pruning

Why Prune?

Modern LLMs contain billions of parameters, yet empirical studies consistently show that a large fraction of these weights contribute minimally to model outputs. Pruning exploits this over-parameterization: by removing redundant weights, we reduce memory footprint (enabling deployment on smaller GPUs or edge devices), inference latency (fewer multiply-accumulate operations per forward pass), and serving cost (higher throughput per dollar). Unlike quantization, which reduces the precision of all weights uniformly, pruning selectively eliminates weights—enabling multiplicative savings when combined with quantization (e.g., a 50% sparse, 4-bit model uses \(4\times\) less memory than the dense BF16 baseline). The challenge is achieving high sparsity without degrading generation quality, which has driven the development of principled one-shot methods that require no retraining.

Important

Pruning Methods

• Unstructured pruning: Zero out individual weights below a threshold. High sparsity (50–90%) possible. Requires sparse GEMM kernels (2:4 on A100/H100).

• Structured pruning: Remove entire attention heads, layers, or FFN neurons. Directly reduces FLOPS without specialized kernels.

• SparseGPT (Frantar and Alistarh 2023): One-shot pruning using approximate inverse Hessian. 50% unstructured sparsity with minimal quality loss on 175B models.

• Wanda (Sun et al. 2024): Prune by \(\vert w\vert \times \\vert x\\vert\) (weight magnitude times input activation norm). No calibration data needed. Competitive with SparseGPT.

Warning

NVIDIA 2:4 Structured Sparsity

A100/H100 Tensor Cores natively support 2:4 sparsity: out of every 4 elements, at most 2 are non-zero. This gives exactly 2\(\times\) speedup on supported operations with hardware acceleration. The constraint: you must achieve exactly 50% sparsity in this specific pattern, which limits flexibility compared to arbitrary sparsity.

Knowledge Distillation

Knowledge distillation (Hinton et al. 2015) transfers the learned behaviour of a large teacher model into a smaller, cheaper student model. The core idea is that the teacher’s output distribution over tokens carries far richer signal than ground-truth hard labels alone — revealing inter-class similarities, calibration, and uncertainty that the student can exploit.

Temperature-Scaled Softmax.

To expose the “dark knowledge” in the teacher’s logits we soften the distribution with a temperature \(T > 1\):

\[ p_i^{(T)} = \frac{\exp(z_i / T)}{\sum_j \exp(z_j / T)} \]

At high temperature the probability mass spreads across more tokens, making near-miss alternatives visible. During training the same temperature is applied to the student; at inference the student uses \(T=1\).

General Distillation Loss.

\[ \mathcal{L}_{\text{distill}} = \alpha , T^2 \cdot \text{KL}!\bigl(P_{\text{teacher}}^{(T)} ;|; P_{\text{student}}^{(T)}\bigr) ;+; (1-\alpha) \cdot \mathcal{L}_{\text{CE}}(y, P_{\text{student}}^{(1)}) \]

The \(T^2\) factor compensates for the reduced gradient magnitude of softened distributions. Typical values: \(T \in [2, 20]\), \(\alpha \in [0.5, 0.9]\) (more weight on KL when teacher quality is high).

Knowledge distillation paradigms for LLMs.

Offline (White-Box) Distillation.

The teacher’s full logit vector (or top-\(k\) logits for storage efficiency) is recorded for each training token. The student minimises the KL divergence against these stored distributions. This is the most data-efficient paradigm when teacher access is unrestricted.

Motivation: Decouple teacher inference from student training — run the teacher once on high-end hardware, then train many students cheaply.

Pros: Deterministic, reproducible; teacher cost is amortised; full distributional signal.
Cons: Requires storing \(\vert V\vert\)-dimensional vectors per token (mitigated by top-\(k\) pruning); teacher cannot adapt to student failures.

Online (Co-Training) Distillation.

Teacher and student are run jointly: the teacher generates logits for the student’s current training batch.

Motivation: Let the teacher focus on inputs where the student currently struggles (curriculum-like).

Pros: Freshness; can use student-generated inputs for on-policy distillation.
Cons: Double the GPU cost; synchronisation complexity; harder to scale.

Black-Box (API) Distillation.

When only text outputs are available (e.g. distilling from a proprietary API), the student is trained via SFT on the teacher’s generations, optionally augmented with chain-of-thought traces.

Motivation: Practical reality — most frontier models do not expose logits.

Pros: Simple pipeline; works with any model behind an API.
Cons: No soft-label signal; prone to hallucination amplification; effectively supervised fine-tuning.

Self-Distillation.

A model distils from a larger version within the same architecture family (e.g. Llama-3 70B \(\to\) 8B) or from its own checkpoints during training.

Motivation: Avoid training a separate teacher; leverage the model’s own capacity at different scales.

Pros: Architecture compatibility; no external dependency.
Cons: Teacher ceiling equals model ceiling; cannot introduce genuinely new knowledge.

Tip

Dark Knowledge

Consider a language model predicting the next word after “The capital of France is”. Hard labels say only “Paris” is correct. But the teacher’s soft distribution might assign 5% to “Lyon”, 2% to “Marseille”, and near-zero to “banana” — telling the student which errors are reasonable, which dramatically improves calibration and generalisation.

Practical Considerations for LLM Distillation.

• Sequence-level vs. token-level: Token-level KL is standard; sequence-level distillation (minimising KL over full sequences) better captures long-range coherence but is harder to optimise.

• Layer-wise hints: Matching intermediate representations (attention maps, hidden states) provides additional learning signal — especially useful when student architecture differs.

• Data selection: Distillation data quality matters; curating diverse, hard examples yields better students than random sampling.

• Student capacity: Diminishing returns below \(\sim\)10% of teacher parameters; at extreme compression, architecture changes (e.g. MoE \(\to\) dense) may be needed.

• Combining with quantization: Distillation + 4-bit quantization (e.g. QLoRA-distilled models) achieves near-teacher quality at \(20\times\) compression.

Note

Compression Method Comparison – 70B Model

Method	Size	Speed	Quality	Use Case
BF16 (baseline)	140 GB	1\(\times\)	100%	Training, reference
FP8 (E4M3)	70 GB	2\(\times\)	99.5%	H100 inference
INT8 (SmoothQuant)	70 GB	1.8\(\times\)	99%	A100 inference
4-bit (AWQ)	35 GB	2.5\(\times\)	97–98%	Serving at scale
2-bit (AQLM)	17.5 GB	3\(\times\)	90–93%	Edge, experimental
Pruned 50% (2:4)	70 GB	1.8\(\times\)	97%	Structured speedup
Distilled 8B	16 GB	10\(\times\)	80–85%	Mobile, edge

Speculative Decoding Methods

Speculative decoding (Leviathan et al. 2023) accelerates autoregressive generation by predicting multiple tokens simultaneously, then verifying them in a single forward pass of the target model. It produces identical output distribution to standard decoding (no quality loss) while achieving 2–3\(\times\) speedup.

Core Principle

Important

Speculative Decoding Framework

1. A fast draft mechanism proposes \(k\) candidate tokens: \(\hat{x}_1, \ldots, \hat{x}_k\)

2. The large target model runs a single forward pass on all \(k\) tokens (batched)

3. Verification: Accept tokens left-to-right while \(P_{\text{target}}(\hat{x}_i) \geq r_i \cdot P_{\text{draft}}(\hat{x}_i)\) (where \(r_i \sim U[0,1]\))

4. On first rejection at position \(j\): resample \(x_j\) from an adjusted distribution, discard \(\hat{x}_{j+1}, \ldots, \hat{x}_k\)

Key property: This acceptance/rejection scheme guarantees the final distribution equals \(P_{\text{target}}\) exactly.

Speedup: If acceptance rate is \(\alpha\), expected tokens per step = \(\frac{1 - \alpha^{k+1}}{1 - \alpha}\). At \(\alpha=0.8\), \(k=5\): expected 3.4 tokens/step vs. 1 for standard decoding.

Methods Comparison

Speculative decoding methods supported by modern inference engines.

Medusa: Multi-Head Speculative Decoding

Important

How Medusa Works

Medusa adds \(k\) additional “prediction heads” to the LLM (sharing the same backbone):

• Head 0 (original): predicts token at position \(t+1\) (standard next-token)

• Head 1: predicts token at position \(t+2\) (skipping one)

• Head \(i\): predicts token at position \(t+i+1\)

• All heads run in parallel during a single forward pass

• A tree-structured verification validates multiple candidate sequences simultaneously

Training: Fine-tune only the Medusa heads (backbone frozen). Cost: \(\sim\)1 epoch on representative data.

Advantage: No separate draft model; heads are tiny (one linear layer each). Memory overhead: \(<\)1%.

Eagle: Feature-Level Drafting

Tip

Why Eagle Outperforms Medusa

Medusa’s heads predict independently at each position — they cannot condition on their own previous predictions (token at \(t+2\) doesn’t know what was predicted at \(t+1\)). Eagle fixes this with a lightweight autoregressive decoder that operates on the target model’s hidden states:

1. Extract hidden states from the target model’s last layer

2. Feed into a small (1-layer) decoder that autoregressively generates draft tokens conditioned on previous hidden states

3. This captures inter-token dependencies that Medusa misses

Result: Eagle achieves 85–95% acceptance rate vs. Medusa’s 60–80%.

N-gram Speculative Decoding

Important

N-gram Lookup Method

The simplest speculative decoding — requires no additional model or training:

1. Maintain a cache of n-grams from the prompt and previously generated text

2. At each step, check if the current context’s last \(n-1\) tokens match any cached n-gram

3. If yes: propose the continuation as draft tokens

4. Verify against target model as usual

Best for: Code generation (repetitive patterns), structured outputs (JSON/XML), and prompts with repeated elements. Cost: Essentially zero.

Integration with vLLM

from vllm import LLM, SamplingParams

# Standard speculative decoding (separate draft model)
llm = LLM(
    model="meta-llama/Llama-3-70B",
    tensor_parallel_size=4,
    speculative_config={
        "model": "meta-llama/Llama-3-8B",
        "num_speculative_tokens": 5,
    },
)

# N-gram speculation (zero-cost, no draft model needed)
llm = LLM(
    model="meta-llama/Llama-3-70B",
    speculative_config={
        "method": "ngram",
        "num_speculative_tokens": 5,
        "prompt_lookup_max": 4,  # Match up to 4-grams from prompt
    },
)

# EAGLE-style (feature-level draft, high acceptance rate)
llm = LLM(
    model="meta-llama/Meta-Llama-3-8B-Instruct",
    tensor_parallel_size=4,
    speculative_config={
        "model": "yuhuili/EAGLE-LLaMA3-Instruct-8B",
        "num_speculative_tokens": 2,
        "method": "eagle",
        "draft_tensor_parallel_size": 1,
    },
)

# MLP speculator (IBM-style, lightweight head)
llm = LLM(
    model="meta-llama/Meta-Llama-3.1-70B-Instruct",
    tensor_parallel_size=4,
    speculative_config={
        "model": "ibm-ai-platform/llama3-70b-accelerator",
        "draft_tensor_parallel_size": 1,
    },
)

Warning

When NOT to Use Speculative Decoding

• High batch sizes: At batch \(\geq 64\), generation is already compute-efficient. Speculation adds overhead (draft generation + verification) that doesn’t pay off.

• Very different distributions: If draft model is too dissimilar to target, acceptance rate drops below 50% and speculation is slower than standard decoding.

• Short outputs: For \(<\)20 token outputs, the setup cost of speculation exceeds savings.

• Rule of thumb: Speculation helps most for latency-sensitive, single-stream generation (chatbots, interactive code completion).

Hallucination Detection

LLMs generate fluent text that may be factually incorrect—a phenomenon called hallucination (Ji et al. 2023). This section covers basic detection methods at the model level (without external retrieval or multi-agent verification).

Types of Hallucination

Important

Hallucination Taxonomy

• Intrinsic: Contradicts the provided input/context (e.g., summary says the opposite of the source)

• Extrinsic: Generates claims that cannot be verified from the input and are factually wrong

• Faithfulness: Output diverges from the instruction or specified constraints

Detection Methods (Model-Level)

Basic hallucination detection methods that operate at the model level.

Semantic Entropy.

Kuhn et al. (Kuhn et al. 2023) observe that token-level entropy is unreliable (paraphrases have different tokens but same meaning). Instead, they generate multiple responses, cluster them by semantic equivalence (via NLI), and compute entropy over meaning clusters:

\[ SE = -\sum_{c \in \text{clusters}} P(c) \log P(c) \]

High SE means the model produces semantically different answers—a strong hallucination signal.

SelfCheckGPT.

Manakul et al. (Manakul et al. 2023) detect hallucinations by checking self-consistency: generate multiple responses and verify whether claims in the main response are supported by the alternatives. If the model “disagrees with itself,” the claim is likely hallucinated. No external knowledge needed.

DoLA (Decoding by Contrasting Layers).

Chuang et al. (Chuang et al. 2024) observe that factual knowledge emerges in later transformer layers while earlier layers retain more generic/uncertain representations. DoLA contrasts the logit distributions between a later (“mature”) layer and an earlier (“premature”) layer at each decoding step:

\[ \text{DoLA}(x_t) = \text{softmax}!\bigl(\log P_{\text{late}}(x_t) - \log P_{\text{early}}(x_t)\bigr) \]

By amplifying the signal from factual knowledge encoded in deeper layers, DoLA reduces hallucinations at inference time without any retraining—requiring only a single additional forward pass through the contrasted layer. It is complementary to sampling-based methods and can be combined with them.

Warning

Limitations of Model-Level Detection

These methods detect uncertainty, not incorrectness. A model can be confidently wrong (low entropy, consistent responses—but factually false). For reliable detection, combine with retrieval-based verification (RAG) or external fact-checking tools.

LLM Safety and Responsible AI

Safety is not an afterthought—it is an integral part of the LLM training pipeline. This section covers the key dimensions of LLM safety and the mechanisms used to enforce responsible behavior.

Threat Taxonomy

LLM safety threat categories.

Safety Training Pipeline

Key Safety Mechanisms

Important

Safety Techniques

• Data filtering: Remove toxic, biased, and PII-containing text from pretraining corpora

• Safety SFT: Train on examples of appropriate refusals (“I can’t help with that because…”)

• Constitutional AI (Bai et al. 2022): Self-critique using principles; model revises its own outputs against a constitution of rules

• Safety reward model: Separate RM trained on safety-annotated pairs; combined with helpfulness RM during RLHF via weighted sum

• Guardrails: Input/output classifiers that block harmful requests/responses at serving time

• Red teaming (Perez et al. 2022): Systematic adversarial evaluation to find failure modes before deployment

The Helpfulness–Safety Tradeoff

Tip

Balancing Helpfulness and Safety

Over-optimizing for safety creates an over-refusal problem: the model declines benign requests (e.g., refusing to discuss historical violence in an educational context). The goal is a Pareto-optimal policy that is maximally helpful within safety constraints: \[ \max_\theta ; \mathbb{E}[R_\text{helpful}] \quad \text{subject to} \quad \mathbb{E}[R_\text{safety}] \geq \tau \] In practice, this is implemented as a weighted reward: \(R = \alpha R_\text{helpful} + (1-\alpha) R_\text{safety}\) with careful tuning of \(\alpha\) (typically 0.6–0.8). Meta’s Llama-3 reports using distinct safety and helpfulness reward models with margin-based weighting (Grattafiori et al. 2024).

Evaluation

• Safety benchmarks: ToxiGen, RealToxicityPrompts, BBQ (bias), CrowS-Pairs

• Jailbreak robustness: GCG attacks (Zou et al. 2023), multi-turn jailbreaks, encoded prompts

• Over-refusal rate: Measure false-positive refusals on benign prompts (target \(<\)5%)

• Red team evaluations: Human adversarial testing with domain experts (biosecurity, cybersecurity)

Warning

Safety Is Never Complete

No combination of techniques provides absolute safety. New attack vectors are discovered continuously (multi-modal jailbreaks, fine-tuning attacks that remove safety training, many-shot prompting). Safety requires ongoing monitoring, rapid response to new threats, and defense-in-depth (multiple independent layers).

1. https://github.com/noamgat/lm-format-enforcer ↩

2. https://github.com/guidance-ai/guidance ↩

Systems Foundations for LLMs

GPU Architecture – From Silicon to LLM Training

Modern large language models are trained and served almost exclusively on GPUs (Graphics Processing Units). Understanding GPU architecture is essential for making informed decisions about parallelism strategies, memory management, kernel optimization, and infrastructure sizing. This section provides a comprehensive introduction to GPU hardware as it relates to LLM workloads.

Why GPUs for Deep Learning?

GPUs and CPUs represent fundamentally different hardware philosophies. Understanding this difference explains why LLM training is 100–1000\(\times\) faster on GPUs.

Tip

CPUs vs. GPUs – Fundamental Design Philosophy

• CPUs are optimized for latency – they execute a few threads as fast as possible, with large caches, branch predictors, and out-of-order execution. A modern CPU has 8–96 cores.

• GPUs are optimized for throughput – they execute thousands of threads in parallel, each doing simple work. A modern GPU has thousands of “cores” (execution units) grouped into Streaming Multiprocessors (SMs).

Deep learning workloads are dominated by matrix multiplications (\(O(n^3)\) operations on \(O(n^2)\) data), which are embarrassingly parallel. A single transformer forward pass for a 70B model requires \(\sim\)140 TFLOP of compute per token – perfect for GPU throughput.

NVIDIA GPU Microarchitecture Generations

NVIDIA has released a series of GPU architectures, each bringing key innovations for deep learning:

NVIDIA GPU microarchitecture timeline for deep learning.

Common GPUs for LLM Training and Inference

Warning

Which GPU to Choose?

• Training 70B+ models: H100/B200 nodes with NVLink (need fast interconnect for tensor parallelism). Minimum 8\(\times\)H100 per instance.

• Inference (latency-sensitive): H100/H200 for high BW; MI300X for memory-bound (huge KV caches).

• Fine-tuning 7B–13B: A100-80GB is cost-effective. Single GPU with LoRA.

• Budget: A100-40GB or even A10 (24GB) for LoRA on 7B models.

GPU Internal Architecture – The Streaming Multiprocessor (SM)

A GPU is organized as an array of Streaming Multiprocessors (SMs), each of which is an independent processor with its own register file, shared memory, and execution units. Understanding SMs is key to understanding GPU performance.

Important

Key SM Components

• CUDA Cores: Scalar ALUs for FP32/INT32 operations. 64 per SM on A100. Used for element-wise ops, reductions, and non-matrix operations.

• Tensor Cores: Specialized matrix-multiply-accumulate (MMA) units. Each performs a \(4{\times}4{\times}4\) fused multiply-add per cycle. 4 per SM on A100, delivering \(16\times\) throughput over CUDA cores for supported precisions.

• Register File: Fastest storage (1 cycle latency). Shared among all active threads. Spilling to L1 causes significant slowdown.

• Shared Memory / L1: On-chip SRAM explicitly managed by the programmer. The key to Flash Attention’s performance (tiles fit entirely in shared memory).

• Warp Schedulers: Each SM has 4 warp schedulers (A100). A warp = 32 threads executing in lockstep (SIMT model). Schedulers hide memory latency by switching between warps.

Tip

The SIMT Execution Model

GPUs use Single Instruction, Multiple Threads (SIMT) execution. Within a warp (32 threads), all threads execute the same instruction but on different data. When threads diverge (e.g., if/else), both paths are serialized – called warp divergence. This is why GPU kernels must minimize branching.

For LLM workloads, the main operations (GEMM, attention, softmax) have uniform control flow across threads, making them ideal for SIMT execution.

GPU Chip Scaling Across Generations

The evolution of NVIDIA’s GPU architectures shows consistent scaling of compute density, on-chip memory, and specialized units for deep learning:

SM-level scaling across NVIDIA architectures.

GPU Memory Hierarchy and Bandwidth

Modern GPU training and inference performance is almost entirely determined by how well you manage data movement across the memory hierarchy. Understanding the hierarchy is not optional – it is the foundation for every optimization technique discussed in later sections.

Important

GPU Memory Hierarchy -- A100 80GB Reference Numbers

Level	Capacity	Bandwidth	Latency	Location
Registers	\(\sim\)256 KB/SM	\(>\)100 TB/s	1 cycle	On-chip, per-thread
SRAM (shared)	164 KB/SM	\(\sim\)19 TB/s	\(\sim\)20 cy	On-chip, per-SM
L2 Cache	40 MB total	\(\sim\)5 TB/s	\(\sim\)200 cy	On-chip, shared
HBM2e (VRAM)	80 GB	2 TB/s	\(\sim\)200 ns	On-package (5 stacks)
CPU DRAM	512 GB+	\(\sim\)25 GB/s	\(\sim\)10 \(\mu\)s	Host (PCIe 4)
NVMe SSD	TBs	7 GB/s	\(\sim\)100 \(\mu\)s	Host storage

Tip

Why the Gaps Are So Large

Each level of the hierarchy is roughly 10\(\times\) slower and 100–1000\(\times\) larger than the one above it. The A100 has 312 TFLOP/s of BF16 tensor-core throughput but only 2 TB/s of HBM bandwidth. That means for every byte loaded from HBM you can do \(312 \times 10^{12} / (2 \times 10^{12}) \approx 156\) floating-point operations before the next byte arrives. If your kernel does fewer than 156 FLOPs per byte, it is memory-bound – the compute units are idle waiting for data.

Registers.

Each CUDA thread has access to a private register file. Registers are the fastest storage on the chip – reads and writes happen in a single clock cycle with no arbitration. The A100 has 65,536 32-bit registers per SM. Spilling registers to local memory (L1/L2) is a major performance hazard.

SRAM – Shared Memory / L1.

Each SM has a combined L1/shared memory pool of 192 KB on A100 (256 KB on H100), with up to 164 KB configurable as shared memory on A100. Shared memory is explicitly managed by the programmer (or by the compiler in newer CUDA versions). Flash Attention, for example, is entirely built around the insight that the attention tile computation fits in SRAM.

L2 Cache.

The 40 MB L2 on A100 is shared across all 108 SMs. It acts as a staging area between SRAM and HBM. For workloads with good spatial locality (e.g., weight matrices accessed repeatedly across a batch), L2 hit rates can dramatically reduce effective HBM traffic.

HBM – High Bandwidth Memory.

HBM is stacked DRAM mounted directly on the GPU package, connected via a wide interposer. The A100 SXM has 80 GB of HBM2e at 2 TB/s. The H100 SXM5 has 80 GB of HBM3 at 3.35 TB/s. This is the primary working memory for model weights, KV caches, activations, and optimizer states.

CPU DRAM via PCIe.

Data transfer between GPU HBM and CPU DRAM traverses the PCIe bus. PCIe Gen4 \(\times\)16 provides \(\sim\)32 GB/s per direction (64 GB/s bidirectional); Gen5 doubles this. This is a \(\sim\)60\(\times\) bandwidth reduction compared to HBM (per-direction). CPU offloading (ZeRO-Infinity, DeepSpeed) exploits this link but must be used carefully to avoid becoming the bottleneck.

NVMe.

NVMe SSDs (e.g., Samsung 990 Pro) reach \(\sim\)7 GB/s sequential read. ZeRO-Infinity can offload optimizer states to NVMe, but this is only viable when the compute-to-IO ratio is very high (large batch sizes, slow training steps).

Arithmetic Intensity and the Roofline Model

Important

Arithmetic Intensity

\[ I = \frac{\text{FLOPs}}{\text{Bytes accessed from HBM}} \quad \text{(FLOPs / Byte)} \] A kernel is memory-bound when \(I < I_{\text{ridge}}\) and compute-bound when \(I > I_{\text{ridge}}\), where \[ I_{\text{ridge}} = \frac{\text{Peak FLOP/s}}{\text{Peak Bandwidth}} = \frac{312 \times 10^{12}}{2 \times 10^{12}} = 156 \text{ FLOP/Byte (A100 BF16)} \]

Note

Attention Arithmetic Intensity

For a single attention head with sequence length \(n=4096\), head dim \(d=128\):

FLOPs: \(QK^T\) costs \(2n^2d\), softmax is \(O(n^2)\), \(\text{Attn} \times V\) costs \(2n^2d\). Total: \(\approx 4n^2 d = 4 \times 4096^2 \times 128 \approx 8.6\) GFLOP.

Memory traffic (standard, non-Flash implementation):

• Read \(Q, K\): \(2 \times n \times d \times 2 = 2\) MB

• Write attention scores \(S = QK^T\): \(n^2 \times 2 = 33.5\) MB

• Read \(S\) for softmax: \(n^2 \times 2 = 33.5\) MB

• Write softmax output \(P\): \(n^2 \times 2 = 33.5\) MB

• Read \(P\) and \(V\) for final matmul: \(n^2 \times 2 + n \times d \times 2 = 34.5\) MB

• Write output \(O\): \(n \times d \times 2 = 1\) MB

Total memory: \(\approx 138\) MB (dominated by 4 passes over the \(n^2\) attention matrix).

Arithmetic intensity: \[ I = \frac{8.6 \times 10^9}{138 \times 10^6} \approx 62 \text{ FLOP/Byte} \]

This is \(62/156 = 40%\) of the A100 ridge point — firmly memory-bound. The GPU is 60% idle waiting for memory.

Flash Attention fix: By never materializing the \(n \times n\) matrix (tiling \(Q, K, V\) in SRAM), Flash Attention reduces HBM traffic to just reading \(Q, K, V\) and writing \(O\): \(4 \times n \times d \times 2 = 4\) MB. Each byte loaded is reused in \(O(n)\) computations (every query attends to every key), so: \[ I = \frac{4n^2 d}{4 \cdot n \cdot d \cdot 2} = \frac{n}{2} = \frac{4096}{2} = 2048 \text{ FLOP/Byte} \]

This is \(13\times\) above the ridge point (156) — deeply compute-bound. The GPU hits its peak 312 TFLOPS, needing only \(312\text{T}/2048 \approx 152\) GB/s of bandwidth (7.6% of HBM capacity). Memory is no longer the bottleneck.

Attention is Memory-Bound; FFN is Compute-Bound

Tip

Two Regimes in a Transformer

A transformer block has two main components with very different arithmetic intensities:

• Attention: Operates on \(n \times d\) tensors. The \(QK^T\) product is \(O(n^2 d)\) FLOPs but requires \(O(n^2)\) memory for the attention scores. At long sequences, memory traffic dominates – attention is memory-bound.

• FFN (MLP): Two large linear layers with weight matrices of shape \([d_{\text{model}}, 4d_{\text{model}}]\). These are large GEMMs with high arithmetic intensity – FFN is compute-bound.

This is why Flash Attention (memory optimization) helps attention but not FFN, while quantization (reducing weight size) helps FFN more than attention.

Tensor Cores

Important

What Are Tensor Cores?

Tensor Cores are specialized matrix-multiply-accumulate (MMA) units introduced in Volta (2017). Each Tensor Core performs a \(4\times4\times4\) matrix multiply in a single clock cycle: \[ D = A \times B + C \quad (4\times4 \text{ matrices}) \] The A100 has 432 Tensor Cores across 108 SMs (4 per SM, one per sub-partition). At BF16 precision, they deliver 312 TFLOP/s – roughly \(16\times\) the throughput of FP32 CUDA cores.

• Supported precisions: FP64, TF32, BF16, FP16, INT8, FP8 (H100+).

• Accumulation: Always in FP32 internally, even for BF16 inputs. This prevents catastrophic cancellation during the dot product.

• Requirement: Tensor Cores are most efficient when matrix dimensions are multiples of 8 (BF16) or 16 (FP8). Padding to these multiples is often worthwhile.

• WGMMA (H100): Hopper introduces warpgroup-level MMA instructions that operate on larger tiles (64\(\times\)256\(\times\)16) and can be pipelined with TMA (Tensor Memory Accelerator) data movement.

Warning

The Tensor Core Trap

Tensor Cores only help if your kernel is compute-bound. If you are running a small batch (batch size 1, inference), the GEMM tiles are tiny, Tensor Core utilization is low, and you are back in the memory-bound regime. This is why inference engines batch requests aggressively.

Communication Architecture – NVLink, InfiniBand, and PCIe

Distributed LLM training and inference require moving enormous amounts of data between GPUs, nodes, and storage. The communication fabric is often the bottleneck for large-scale training.

PCIe – The Host-Device Link.

Important

PCIe Generations

Generation	x16 BW (each dir.)	Bidirectional	Notes
PCIe Gen3	16 GB/s	32 GB/s	Common in older servers
PCIe Gen4	32 GB/s	64 GB/s	A100 PCIe, most current servers
PCIe Gen5	64 GB/s	128 GB/s	H100 PCIe, emerging

PCIe is used for:

• CPU \(\leftrightarrow\) GPU data transfers (model loading, CPU offloading)

• Cross-node GPU communication when NVLink is unavailable (rare, very slow)

• NVMe storage access (via CPU)

Warning

PCIe is Not for GPU-GPU Communication

Never route GPU-GPU communication through PCIe if NVLink is available. PCIe bandwidth (32 GB/s) is 28\(\times\) lower than NVLink 4 (900 GB/s). In a multi-GPU server without NVLink (e.g., consumer GPUs), inter-GPU bandwidth is limited to PCIe, making tensor parallelism extremely slow.

NVLink – Intra-Node High-Speed Interconnect.

Important

NVLink Generations

Generation	Links	Total BW	GPU
NVLink 2	6	300 GB/s	V100
NVLink 3	12	600 GB/s	A100
NVLink 4	18	900 GB/s	H100
NVLink 5	18	1800 GB/s	B200 (Blackwell)

NVLink is a point-to-point interconnect between GPUs on the same node. Each link is bidirectional. The H100 SXM5 has 18 NVLink 4 links, each providing 50 GB/s bidirectional, for a total of 900 GB/s.

NVSwitch.

In DGX H100 systems, all 8 GPUs are connected via NVSwitch – a dedicated switching chip that provides full bisection bandwidth. This means any GPU can communicate with any other GPU at full NVLink speed simultaneously, not just neighbors in a ring.

Tip

Ring vs. Full Bisection

In a ring topology (8 GPUs), an AllReduce requires data to travel around the ring. Each link must carry \(\frac{2(N-1)}{N}\) of the total data, so the algorithm bandwidth is \(B_{\text{link}} \times \frac{N}{2(N-1)}\) (about \(0.57 \times B_{\text{link}}\) for \(N=8\)). With NVSwitch full bisection, AllReduce can use all links simultaneously with tree-based algorithms, achieving near-peak bandwidth. In practice on DGX H100: ring achieves \(\sim\)700 GB/s bus bandwidth, NVSwitch achieves \(\sim\)900 GB/s.

InfiniBand – Inter-Node Communication.

For communication between nodes (servers), InfiniBand provides high-bandwidth, low-latency networking with direct GPU memory access.

Important

InfiniBand NDR

• NDR 400Gb/s = 50 GB/s per port (unidirectional)

• HDR 200Gb/s = 25 GB/s per port (previous generation)

• RDMA: Remote Direct Memory Access – GPU can read/write remote GPU memory without involving the remote CPU

• GPUDirect RDMA: Data goes directly HBM \(\to\) NIC \(\to\) network \(\to\) NIC \(\to\) HBM, bypassing CPU and system DRAM entirely

• Latency: \(\sim\)1–2 \(\mu\)s for small messages (vs. \(\sim\)100 \(\mu\)s for TCP/IP)

Fat-Tree Topology.

Large GPU clusters use fat-tree network topologies. A 3-level fat-tree with \(k\)-port switches supports \(k^3/4\) nodes with full bisection bandwidth. For 400Gb/s NDR switches with \(k=64\) ports: \(64^3/4 = 65{,}536\) nodes.

Rail-Optimized Topology.

In practice, clusters use rail-optimized topologies where each GPU in a node connects to a different top-of-rack switch. This ensures that AllReduce operations (which involve all GPUs) use all network links simultaneously, maximizing bandwidth.

Communication Patterns in Distributed LLM Training.

Distributed training relies on collective communication primitives. The choice of primitive determines bandwidth requirements and scaling behavior.

Important

Communication Primitives

Primitive	Use Case	Volume
AllReduce	Gradient sync (DDP, FSDP)	\(2(N-1)/N \times\) param size
AllGather	Collect sharded weights (FSDP)	\((N-1)/N \times\) param size
ReduceScatter	Scatter gradients (FSDP)	\((N-1)/N \times\) param size
AllGather	Tensor parallel activation	activation size
Point-to-Point	Pipeline parallel (send/recv)	micro-batch activation
Broadcast	Weight sync (new workers)	full model size

Note

Bandwidth Calculation – Gradient AllReduce for 70B Model

Setup: 70B parameter model, BF16 gradients, 8 nodes \(\times\) 8 GPUs = 64 GPUs. Data parallel degree = 64.

Gradient size: \(70 \times 10^9 \times 2\) bytes \(= 140\) GB.

AllReduce volume per GPU (ring): \(2 \times (64-1)/64 \times 140 \approx 275\) GB.

Available inter-node bandwidth: 8 GPUs/node \(\times\) 50 GB/s/GPU \(= 400\) GB/s (with rail-optimized topology, all 8 NICs active).

AllReduce time: \(275 / 400 \approx 0.69\) seconds per step.

Implication: For a 1-second compute step, communication adds 0.69 seconds (41% of total step time). This is why gradient compression, mixed precision, and FSDP (which overlaps communication with computation) are critical.

Network Topology Diagram.

The following diagram illustrates a typical two-node GPU cluster topology showing both intra-node (NVLink) and inter-node (InfiniBand) communication paths.

Tip

Choosing Parallelism Based on Bandwidth

• Tensor Parallelism (TP): Requires all-reduce every layer – use only within a node over NVLink. TP=8 is standard for H100 DGX nodes.

• Pipeline Parallelism (PP): Point-to-point between stages – can cross nodes, but adds pipeline bubble overhead. Use when model is too large for TP alone.

• Data Parallelism (DP): AllReduce of gradients – can cross nodes via IB. Scales well with fast IB.

• FSDP/ZeRO: AllGather + ReduceScatter – similar to DP but shards optimizer states. Preferred over DP for large models.

vLLM – PagedAttention and High-Throughput Inference

vLLM (Kwon et al. 2023) introduced PagedAttention, which borrows the paging abstraction that operating systems use for RAM and applies it to the GPU’s KV cache. During LLM inference, the KV cache – the stored key and value tensors for all previous tokens – is the dominant memory consumer. Managing it efficiently is the central challenge of high-throughput inference.

The KV Cache Fragmentation Problem

Important

KV Cache Memory Formula

For a model with \(L\) layers, \(H\) heads, head dimension \(d\), and a sequence of \(n\) tokens: \[ \text{KV cache size} = 2 \times L \times H \times d \times n \times \text{bytes_per_element} \] For Llama-3 70B (BF16): \(L=80\), \(H=8\) (GQA), \(d=128\): \[ = 2 \times 80 \times 8 \times 128 \times n \times 2 = 327{,}680 \times n \text{ bytes} \] At \(n=4096\) tokens: \(\approx 1.3\) GB per sequence.

Tip

Internal and External Fragmentation

Traditional inference systems pre-allocate a contiguous memory block for each sequence’s KV cache, sized to the maximum possible sequence length. This causes two types of waste:

Internal fragmentation: A sequence that generates only 500 tokens still holds a block reserved for 4096 tokens. The unused 3596 token slots are wasted.

External fragmentation: After many sequences complete, the free memory consists of many small non-contiguous gaps. A new long sequence cannot be allocated even if total free memory is sufficient, because no single contiguous block is large enough.

In practice, GPU memory utilization with naive allocation is often only 20–40%.

PagedAttention – Virtual Memory for KV Caches

PagedAttention (Kwon et al., 2023) borrows the paging abstraction from operating systems. Instead of one contiguous block per sequence, the KV cache is carved into fixed-size pages (blocks), and an indirection table—analogous to a CPU page table—translates each sequence’s logical token positions into scattered physical GPU memory addresses.

Important

PagedAttention Core Concepts

• Block size: Typically 16 tokens per block (tunable). Each block stores \(16 \times 2 \times L \times H \times d\) elements.

• Block table: A per-sequence mapping from logical block index to physical block index in the GPU memory pool.

• Physical block pool: A pre-allocated pool of fixed-size blocks. Allocation is \(O(1)\) – just pop from a free list.

• Attention kernel: Modified to gather KV blocks from non-contiguous physical locations using the block table during attention computation.

Note

Block Table Example

Suppose block size = 4 tokens, and we have two sequences:

• Sequence A (7 tokens): logical blocks [0,1] \(\to\) physical blocks [3, 7]

• Sequence B (5 tokens): logical blocks [0,1] \(\to\) physical blocks [1, 5]

Physical block 3 holds tokens 0–3 of sequence A. Physical block 7 holds tokens 4–6 of sequence A (partially filled). The attention kernel for sequence A reads from physical blocks 3 and 7 in order, using the block table as an indirection layer.

Benefits of PagedAttention

Near-zero waste.

Internal fragmentation is bounded by at most one partially-filled block per sequence (the last block). With block size 16, worst-case waste is 15 tokens per sequence – negligible. External fragmentation is eliminated because blocks are fixed-size and interchangeable.

Dynamic allocation.

Blocks are allocated on demand as the sequence grows. No need to know the final sequence length in advance. This is critical for generation, where output length is unknown.

Prefix sharing (copy-on-write).

Multiple sequences sharing a common prefix (e.g., a system prompt) can share the same physical blocks for that prefix. The block table simply points multiple sequences to the same physical blocks. When a sequence needs to write to a shared block (diverging from the prefix), a copy-on-write is triggered.

Tip

Prefix Sharing Savings

In a chatbot with a 1000-token system prompt serving 128 concurrent users:

• Without prefix sharing: \(128 \times 1000 \times 327{,}680 / 10^9 \approx 42\) GB just for system prompt KV cache

• With prefix sharing: \(1 \times 1000 \times 327{,}680 / 10^9 \approx 0.33\) GB

• Savings: \(\sim\)128\(\times\) for the shared prefix portion

Preemption via swap.

When GPU memory is exhausted, vLLM can preempt a sequence by swapping its KV blocks to CPU DRAM (or simply discarding them and recomputing later). This is only feasible because blocks are self-contained and non-contiguous – swapping a contiguous allocation would require copying the entire buffer.

Continuous Batching

Traditional batching (“static batching“) waits until all sequences in a batch finish before starting new ones. If one sequence generates 500 tokens and another generates 10, the GPU is idle for 490 steps on the short sequence. This is extremely wasteful.

Important

Continuous Batching

Continuous batching (also called iteration-level scheduling) processes one decode step at a time. After each step:

1. Check which sequences have finished (generated EOS token)

2. Remove finished sequences from the batch, freeing their KV blocks

3. Add new waiting sequences to fill the freed slots

4. Run the next decode step with the updated batch

The batch composition changes every step — sequences join and leave dynamically. This keeps GPU utilization near 100% and dramatically improves throughput (1.5–3\(\times\) over static batching). PagedAttention is essential here: adding/removing sequences mid-batch requires dynamic KV block allocation/deallocation, which is only efficient with paged memory.

Speculative Decoding in vLLM

Speculative decoding uses a small draft model (e.g., 1B parameters) to propose \(k\) candidate tokens quickly, which the large target model verifies in a single forward pass. All tokens up to the first rejection are accepted (expected acceptance: 3–5 tokens per verification step). This yields 2–3\(\times\) speedup for latency-sensitive single-sequence generation without any quality loss.

vLLM integrates speculative decoding with PagedAttention:

• Draft tokens are allocated speculative KV blocks

• On rejection, speculative blocks are freed (cheap with paged allocation)

• On acceptance, speculative blocks are promoted to the main sequence

• The block table update is \(O(k)\) – just updating a few table entries

Concrete Memory Savings – 70B Model at Scale

Note

Memory Budget – 70B BF16 Inference

Setup: Llama-3 70B, BF16, single A100 80GB node (8 GPUs, tensor parallel).

Model weights: \(70 \times 10^9 \times 2\) bytes \(= 140\) GB \(\div\) 8 GPUs \(= 17.5\) GB/GPU.

Remaining for KV cache: \(80 - 17.5 - 3\) (overhead) \(= 59.5\) GB/GPU.

KV cache per token per GPU (with TP=8, each GPU holds \(1/8\) of heads): \(2 \times 80 \times 1 \times 128 \times 2 = 40{,}960\) bytes \(\approx 40\) KB/token.

Max tokens in KV cache: \(59.5 \times 10^9 / 40{,}960 \approx 1.45\) million tokens.

With 128 concurrent sequences of 4096 tokens each: \(128 \times 4096 = 524{,}288\) tokens – well within budget.

Without PagedAttention (pre-allocating max length 4096 for each): Same math, but fragmentation wastes \(\sim\)50% on average \(\to\) only 64 sequences fit.

Warning

Block Size Tradeoff

Larger block sizes reduce the overhead of the block table and improve memory access locality (fewer scattered reads). Smaller block sizes reduce internal fragmentation and enable finer-grained prefix sharing. vLLM defaults to 16 tokens/block, which is a good balance. For very long sequences (100K+ tokens), larger blocks (32–64) may be preferable.

vLLM: End-to-End System

vLLM wraps PagedAttention inside a full serving stack: continuous batching, prefix caching, speculative decoding, and tensor-parallel model sharding all work together to maximize throughput per GPU dollar.

Architecture Overview

Core Components

• API Server: Accepts OpenAI-compatible requests (completions, chat). Tokenizes inputs and creates “sequence groups” (for beam search or multiple samples).

• Scheduler: The brain of vLLM. Maintains three queues:

  • waiting: New requests not yet started (prefill pending)

  • running: Actively generating tokens (decode phase)

  • swapped: Preempted requests whose KV cache was offloaded to CPU

  Each iteration, the scheduler decides which requests to run based on available GPU memory blocks.

• Block Manager: Implements the virtual memory abstraction for KV caches. Maps logical blocks (per-sequence) to physical blocks (in GPU memory pool). Handles:

  • Allocation (new tokens generated \(\rightarrow\) new blocks needed)

  • Copy-on-write (for beam search: multiple beams share prefix blocks, copy only on divergence)

  • Swap (GPU \(\leftrightarrow\) CPU migration when preempting/resuming)

  • Prefix caching (reuse cached blocks when prompts share common prefixes)

• Model Executor: Runs the actual LLM forward pass. Manages tensor parallelism across GPUs, dispatches attention kernels that read from paged KV cache blocks.

• KV Cache Pool: Pre-allocated GPU memory divided into fixed-size blocks (default: 16 tokens \(\times\) num_heads \(\times\) head_dim \(\times\) 2 bytes per block). No dynamic allocation at runtime \(\rightarrow\) zero fragmentation.

Request Lifecycle (End-to-End Flow)

1. Arrival: Client sends prompt. API server tokenizes it, creates a SequenceGroup, places it in the waiting queue.

2. Scheduling: At each step, the scheduler runs:

   1. Check if any swapped sequences can be resumed (enough free blocks).

   2. Check if any waiting sequences can start prefill (enough blocks for the full prompt).

   3. Budget remaining blocks across running sequences (need 1 new block per sequence per step if current block is full).

   4. If over budget: preempt lowest-priority running sequences (swap KV to CPU or recompute later).

3. Prefill (first iteration for a request): The entire prompt is processed in one forward pass. KV cache is computed for all prompt tokens and stored in allocated blocks. This is compute-bound (large batch of tokens).

4. Decode (subsequent iterations): One new token generated per sequence per step. All running sequences are batched together (continuous batching). This is memory-bound (reads full KV cache, generates 1 token).

5. Block Allocation: After each decode step, if the last block for a sequence is full, the Block Manager allocates a new physical block and maps it to the next logical block.

6. Completion: When a sequence hits EOS or max length, it’s removed from running. Its physical blocks are freed immediately \(\rightarrow\) available for other sequences. Response is streamed back to client.

Prefix Caching (Automatic Prompt Caching)

When multiple requests share a common prefix (system prompt, few-shot examples):

1. Hash the token content of each logical block.

2. On new request arrival, check if any prefix blocks are already in the cache.

3. If hit: skip prefill for those tokens, directly reuse physical KV blocks. Time-to-first-token drops dramatically.

4. Eviction: LRU policy. Cached blocks are freed only when memory pressure requires it.

Impact: For chat applications with long system prompts (2K+ tokens shared across all users), prefix caching reduces TTFT by 60–80%.

Guided (Constrained) Decoding in vLLM

vLLM natively supports constrained decoding (Section 1.12.11) through pluggable backends, enabling guaranteed structured output at serving time with minimal performance overhead.

Supported constraint types.

The OpenAI-compatible API accepts constraints via the guided_* parameters or the response_format field:

from openai import OpenAI
client = OpenAI(base_url="http://localhost:8000/v1")

# --- JSON Schema constraint ---
response = client.chat.completions.create(
    model="meta-llama/Llama-3-70B-Instruct",
    messages=[{"role": "user",
               "content": "Extract: name, age, city from: "
                          "'John is 30 and lives in NYC'"}],
    extra_body={
        "guided_json": {
            "type": "object",
            "properties": {
                "name": {"type": "string"},
                "age": {"type": "integer"},
                "city": {"type": "string"}
            },
            "required": ["name", "age", "city"]
        }
    }
)
# Output is guaranteed valid JSON matching the schema

# --- Regex constraint ---
response = client.completions.create(
    model="meta-llama/Llama-3-70B-Instruct",
    prompt="Generate an IPv4 address: ",
    extra_body={
        "guided_regex": r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}"
    }
)

# --- Choice constraint ---
response = client.completions.create(
    model="meta-llama/Llama-3-70B-Instruct",
    prompt="Sentiment: ",
    extra_body={"guided_choice": ["positive", "negative", "neutral"]}
)

Backend architecture.

vLLM delegates mask computation to a backend engine:

• XGrammar (default since v0.7): Pushdown-automaton engine supporting JSON schemas, regexes, and arbitrary EBNF grammars. Fastest for complex schemas due to efficient C++ core.

• Outlines (Willard and Louf 2023): FSM-based; supports JSON and regex. Used as fallback when XGrammar is unavailable.

The mask is applied after the model’s forward pass produces logits and before sampling—adding \(<\)1 ms per step in practice, since the FSM/PDA state transition and precomputed index lookup are \(O(1)\).

Performance impact.

Because the constraint only masks logits (no recomputation of attention or FFN), throughput loss is negligible (\(<\)2% in benchmarks). The main cost is compilation of the schema into an FSM/PDA index, which takes 0.5–5 s depending on schema complexity. vLLM caches compiled schemas across requests, so this cost is paid once per unique schema.

Warning

Structured Output $\neq$ Correct Output

Constrained decoding guarantees the output is syntactically valid (parses as JSON, matches the schema types). It does not guarantee semantic correctness—the model may still hallucinate values that parse correctly but are factually wrong. Always validate business logic downstream.

vLLM performance vs. alternatives (70B model, A100 \(\times\) 4, TP=4).

Introduction to Reinforcement Learning

Reinforcement Learning (RL) is a paradigm where an agent learns to make sequential decisions by interacting with an environment, receiving rewards as feedback, and optimizing its policy to maximize cumulative reward over time (Sutton and Barto 2018). Unlike supervised learning (which requires labeled input-output pairs), RL discovers optimal behavior through trial and error.

The Markov Decision Process (MDP)

An MDP is a 5-tuple \((S, A, P, R, \gamma)\):

• \(S\): State space — all possible configurations of the environment

• \(A\): Action space — all actions available to the agent

• \(P(s'\vert s, a)\): Transition function — probability of reaching state \(s'\) from state \(s\) after taking action \(a\)

• \(R(s, a, s')\): Reward function — immediate scalar feedback for a transition

• \(\gamma \in [0, 1]\): Discount factor — how much future rewards are valued relative to immediate ones

The Markov Property: The future depends only on the current state, not the history:

\[ P(s_{t+1} | s_t, a_t, s_{t-1}, a_{t-1}, \ldots) = P(s_{t+1} | s_t, a_t) \]

This makes the problem tractable.

Tip

Agent-Environment Interaction Loop

At each time step \(t\):

1. Agent observes state \(s_t\)

2. Agent selects action \(a_t\) according to policy \(\pi(a\vert s)\)

3. Environment transitions to \(s_{t+1} \sim P(\cdot\vert s_t, a_t)\)

4. Agent receives reward \(r_t = R(s_t, a_t, s_{t+1})\)

5. Repeat until terminal state or horizon \(T\)

Core Concepts and Definitions

Policy \(\pi(a\vert s)\): A mapping from states to action probabilities. Deterministic: \(a = \pi(s)\). Stochastic: \(a \sim \pi(\cdot\vert s)\).

Return (cumulative discounted reward):

\[ G_t = \sum_{k=0}^{\infty} \gamma^k r_{t+k} = r_t + \gamma r_{t+1} + \gamma^2 r_{t+2} + \cdots \]

Value Function (expected return from state \(s\) under policy \(\pi\)):

\[ V^\pi(s) = \mathbb{E}_\pi\left[G_t \mid s_t = s\right] = \mathbb{E}_\pi\left[\sum_{k=0}^{\infty} \gamma^k r_{t+k} \mid s_t = s\right] \]

Action-Value Function (expected return from state \(s\), taking action \(a\), then following \(\pi\)):

\[ Q^\pi(s, a) = \mathbb{E}_\pi\left[G_t \mid s_t = s, a_t = a\right] \]

Advantage Function (how much better action \(a\) is compared to average):

\[ A^\pi(s, a) = Q^\pi(s, a) - V^\pi(s) \]

Bellman Equations (recursive relationship):

\[ \begin{aligned} V^\pi(s) &= \sum_a \pi(a|s) \sum_{s'} P(s'|s,a)\left[R(s,a,s') + \gamma V^\pi(s')\right] \\ Q^\pi(s,a) &= \sum_{s'} P(s'|s,a)\left[R(s,a,s') + \gamma \sum_{a'} \pi(a'|s') Q^\pi(s', a')\right] \end{aligned} \]

Important

Optimal Policy and Bellman Optimality

The optimal policy \(\pi^*\) satisfies: \[ V^*(s) = \max_a \sum_{s'} P(s'|s,a)\left[R(s,a,s') + \gamma V^*(s')\right] \]

\[ Q^*(s,a) = \sum_{s'} P(s'|s,a)\left[R(s,a,s') + \gamma \max_{a'} Q^*(s', a')\right] \] Once \(Q^*\) is known, the optimal policy is simply: \(\pi^*(s) = \arg\max_a Q^*(s,a)\).

Taxonomy of RL Methods

Reinforcement learning algorithms can be classified along several axes. Understanding this taxonomy helps select the right approach for a given problem.

Important

Key Taxonomy Distinctions

Model-Free vs Model-Based:

• Model-Free: Learn policy or value function directly from experience. No knowledge of environment dynamics. Most practical for LLMs (language dynamics are intractable to model).

• Model-Based: Learn or use a model of environment transitions \(P(s'\vert s,a)\). Can plan ahead. More sample-efficient but requires accurate model.

Value-Based vs Policy-Based:

• Value-Based: Learn \(Q(s,a)\) or \(V(s)\), derive policy as \(\arg\max_a Q(s,a)\). Works well for discrete, small action spaces (e.g., Atari). Struggles with continuous/large action spaces.

• Policy-Based: Directly parameterize and optimize \(\pi_\theta(a\vert s)\). Natural for continuous/high-dimensional action spaces. Essential for LLMs (vocabulary = 32K–128K actions).

• Actor-Critic: Combine both — policy (actor) proposes actions, value function (critic) evaluates them. PPO for LLMs is actor-critic.

On-Policy vs Off-Policy:

• On-Policy: Learn from data generated by the current policy. Must regenerate data after each update. Examples: REINFORCE, PPO, A2C. More stable but less sample-efficient.

• Off-Policy: Learn from data generated by any policy (including old versions or other agents). Can reuse past experience. Examples: Q-Learning, DQN, SAC. More sample-efficient but harder to stabilize.

Temporal Difference (TD) Learning

TD learning (Sutton 1988) bootstraps — it updates value estimates using other value estimates, without waiting for the full episode to end.

Understanding TD Error: “Surprise” as a Learning Signal

TD error measures the discrepancy between an agent’s current estimate of future reward and a newly updated estimate after taking one step. Put simply, it is the difference between what the agent thought would happen and what actually happened plus what it expects next. It represents the agent’s “surprise.”

Note

The Driving Analogy

Imagine you are driving home and expect the drive to take 30 minutes.

• The prediction: 30 minutes total.

• The reality shift: After 10 minutes, you hit unexpected road construction. Your GPS updates, saying you now have 35 minutes left.

• The TD Error: Total expected time is now 45 minutes (10 elapsed + 35 remaining). The difference between new estimate (45 min) and old estimate (30 min) is a +15 minute TD error. You use this “surprise” to change your route next time.

A positive TD error means the outcome was better than expected \(\rightarrow\) boost this state’s value.
A negative TD error means it was worse than expected \(\rightarrow\) lower this state’s value.

The TD Error Formula

\[ \boxed{\delta_t = R_{t+1} + \gamma V(S_{t+1}) - V(S_t)} \]

• \(R_{t+1}\): The immediate reward received after taking an action.

• \(\gamma V(S_{t+1})\): The estimated discounted value of the next state (what the agent expects to get from the next state onward, scaled by discount factor \(\gamma\)).

• \(V(S_t)\): The original estimate of the current state’s value.

The combined term \((R_{t+1} + \gamma V(S_{t+1}))\) is called the TD Target. Therefore:

\[ \text{TD Error} = \text{TD Target} - \text{Old Estimate} \]

How the Agent Uses TD Error

The agent adjusts its value function to drive TD error toward zero:

\[ \boxed{V(S_t) \leftarrow V(S_t) + \alpha \cdot \delta_t} \]

• If \(\delta_t > 0\): outcome was better than predicted \(\rightarrow\) increase \(V(S_t)\) so the agent seeks this state.

• If \(\delta_t < 0\): outcome was worse than predicted \(\rightarrow\) decrease \(V(S_t)\) so the agent avoids this state.

• If \(\delta_t = 0\): prediction was perfect \(\rightarrow\) no update needed (convergence).

Tip

TD vs Monte Carlo

Monte Carlo: Wait until episode ends, use actual return \(G_t\). Unbiased but high variance (one full trajectory may be unrepresentative).

TD: Update after every step using estimated future value \(\gamma V(s_{t+1})\). Biased (depends on \(V\) accuracy) but much lower variance (one-step updates, doesn’t compound noise).

TD(\(\lambda\)): Interpolate between TD(0) and Monte Carlo. \(\lambda=0\): pure TD. \(\lambda=1\): pure MC. This is exactly what GAE does for PPO (with \(\lambda=0.95\)).

TD Target: \(y_t = r_t + \gamma V(s_{t+1})\) — the “better estimate” we move toward.

Multi-step TD (n-step returns):

\[ G_t^{(n)} = r_t + \gamma r_{t+1} + \cdots + \gamma^{n-1} r_{t+n-1} + \gamma^n V(s_{t+n}) \]

Q-Learning

Q-Learning (Watkins 1989) is the foundational off-policy, value-based algorithm. It learns the optimal \(Q^*\) directly, regardless of the policy being followed.

Update rule:

\[ \boxed{Q(s_t, a_t) \leftarrow Q(s_t, a_t) + \alpha\left[r_t + \gamma \max_{a'} Q(s_{t+1}, a') - Q(s_t, a_t)\right]} \]

Tip

Why Q-Learning is Off-Policy

The update uses \(\max_{a'} Q(s_{t+1}, a')\) — the value of the best action at the next state, regardless of which action the agent actually took. This means the target is always computed under the optimal policy, even if the behavior policy explores randomly (\(\epsilon\)-greedy).

This is why Q-Learning can learn from replay buffers, demonstrations, or any source of experience. The data doesn’t need to come from the current policy.

SARSA (Rummery and Niranjan 1994) (on-policy alternative): Uses the action actually taken instead of the max:

\[ Q(s_t, a_t) \leftarrow Q(s_t, a_t) + \alpha\left[r_t + \gamma Q(s_{t+1}, a_{t+1}) - Q(s_t, a_t)\right] \]

Deep Q-Networks (DQN) (Mnih et al. 2015): Replace tabular \(Q(s,a)\) with a neural network \(Q_\theta(s,a)\). Key innovations: experience replay buffer (off-policy data reuse), target network (stability), \(\epsilon\)-greedy exploration.

DQN Loss Function: The network is trained to minimize the mean squared TD error over mini-batches sampled from the replay buffer:

\[ \boxed{\mathcal{L}(\theta) = \mathbb{E}_{(s,a,r,s') \sim \mathcal{B}}!\left[\left(r + \gamma \max_{a'} Q_{\bar{\theta}}(s', a') - Q_\theta(s, a)\right)^2\right]} \]

where \(Q_{\bar{\theta}}\) is the target network — a frozen copy of \(Q_\theta\) updated only every \(C\) steps (e.g., \(C = 10{,}000\)). This prevents the moving target problem: without it, both the prediction and the target shift simultaneously, causing divergence.

Gradient update: Taking the gradient of the loss w.r.t. \(\theta\) (note: the target \(y\) is treated as a constant — no gradient flows through \(\bar{\theta}\)):

\[ \nabla_\theta \mathcal{L} = -\mathbb{E}!\left[\underbrace{\left(r + \gamma \max_{a'} Q_{\bar{\theta}}(s', a') - Q_\theta(s, a)\right)}_{\text{TD error } \delta}; \nabla_\theta Q_\theta(s, a)\right] \]

\[ \theta \leftarrow \theta - \alpha \cdot \delta \cdot \nabla_\theta Q_\theta(s, a) \]

Learning scheme (per training step):

1. Act: Select action via \(\epsilon\)-greedy: with probability \(\epsilon\) take random action, otherwise \(a = \arg\max_a Q_\theta(s, a)\). Anneal \(\epsilon\) from 1.0 \(\rightarrow\) 0.01 over first 1M steps.

2. Store: Save transition \((s, a, r, s', d)\) in replay buffer \(\mathcal{B}\) (capacity \(\sim\)1M).

3. Sample: Draw mini-batch of 32 transitions uniformly from \(\mathcal{B}\).

4. Compute target: \(y = r + \gamma(1 - d)\max_{a'} Q_{\bar{\theta}}(s', a')\) (zero future value if terminal).

5. Update: Gradient descent on \((y - Q_\theta(s,a))^2\). Clip gradients to \([-1, 1]\) (Huber loss variant).

6. Sync target: Every \(C\) steps, copy \(\bar{\theta} \leftarrow \theta\).

Understanding Replay Buffers

A replay buffer (Lin 1992) (experience replay) is a data storage mechanism that saves past experiences so an agent can relearn from them later. Instead of discarding data immediately after an action, the agent stores transitions in a memory bank and samples random mini-batches for training.

What’s stored: Each transition is a tuple:

\[ e_t = (s_t, a_t, r_t, s_{t+1}, d_t) \]

where \(d_t\) is a boolean flag indicating if the episode ended.

Important

Why Replay Buffers Are Essential

• Breaks data correlation: Consecutive steps are highly correlated. Neural networks generalize poorly on sequential data. Random sampling from a buffer makes training samples approximately i.i.d.

• Prevents catastrophic forgetting: Without a buffer, an agent that passes a difficult level might forget how to clear it while spending the next 10K steps failing on a later level. The buffer ensures it continues to practice old scenarios.

• Improves sample efficiency: Running environments can be slow. A replay buffer allows multiple weight updates from the same transition, extracting more value from every step.

import random
from collections import deque

class ReplayBuffer:
    def __init__(self, capacity):
        self.buffer = deque(maxlen=capacity)  # Bounded queue
    
    def push(self, state, action, reward, next_state, done):
        self.buffer.append((state, action, reward, next_state, done))
        
    def sample(self, batch_size):
        # Break correlation by selecting random experiences
        return random.sample(self.buffer, batch_size)
        
    def __len__(self):
        return len(self.buffer)

Tip

Prioritized Experience Replay (PER)

In a standard buffer, all experiences have equal sampling probability. But some are much more educational. PER (Schaul et al. 2016) scales sampling probability by TD error magnitude — if a transition caused massive “surprise” (high \(\vert \delta_t\vert\)), the agent samples it more frequently to correct its model faster. This accelerates learning by 2–3\(\times\) on Atari benchmarks.

Warning

Why Q-Learning Fails for LLMs

The action space in language generation is the full vocabulary (\(\vert A\vert = 32\text{K}\)–\(128\text{K}\)), and the state space is all possible token sequences (infinite). Computing \(\max_a Q(s,a)\) over 128K actions at every token position is intractable. This is why LLM RL uses policy-based methods (PPO, GRPO) instead.

Policy Gradient Methods — REINFORCE

Instead of learning a value function and deriving a policy, directly optimize the policy parameters \(\theta\) to maximize expected return (Williams 1992).

Objective: \(J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}[R(\tau)] = \mathbb{E}_{\pi_\theta}\left[\sum_{t=0}^T r_t\right]\)

Policy Gradient Theorem:

\[ \boxed{\nabla_\theta J(\theta) = \mathbb{E}_{\pi_\theta}\left[\sum_{t=0}^T \nabla_\theta \log \pi_\theta(a_t|s_t) \cdot G_t\right]} \]

Important

Policy Gradient Theorem — Formal Derivation (5 Steps)

Step 1: Define the objective. We want to maximize expected return: \[ J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}!\left[\sum_{t=0}^T r_t\right] = \sum_\tau P(\tau|\theta) R(\tau) \] where \(P(\tau\vert \theta) = p(s_0)\prod_{t=0}^T \pi_\theta(a_t\vert s_t), p(s_{t+1}\vert s_t, a_t)\) is the trajectory probability.

Step 2: Take the gradient. Only \(\pi_\theta\) terms depend on \(\theta\) (dynamics \(p\) doesn’t): \[ \nabla_\theta J = \sum_\tau \nabla_\theta P(\tau|\theta), R(\tau) \]

Step 3: Apply the log-derivative trick: \(\nabla_\theta P(\tau\vert \theta) = P(\tau\vert \theta), \nabla_\theta \log P(\tau\vert \theta)\): \[ \nabla_\theta J = \mathbb{E}_{\tau \sim \pi_\theta}!\left[\nabla_\theta \log P(\tau|\theta), R(\tau)\right] \]

Step 4: Expand \(\log P(\tau\vert \theta)\). The \(\log p(s_0)\) and \(\log p(s_{t+1}\vert s_t,a_t)\) terms vanish under \(\nabla_\theta\): \[ \nabla_\theta \log P(\tau|\theta) = \sum_{t=0}^T \nabla_\theta \log \pi_\theta(a_t|s_t) \]

Step 5: Combine. Future rewards don’t depend on past actions (causality), so each \(\nabla\log\pi\) pairs only with future return \(G_t = \sum_{t'=t}^T r_{t'}\): \[ \boxed{\nabla_\theta J = \mathbb{E}_{\pi_\theta}!\left[\sum_{t=0}^T \nabla_\theta \log \pi_\theta(a_t|s_t) \cdot G_t\right]} \]

Tip

Why This Is Beautiful

The gradient does not require differentiating through the environment dynamics \(p(s'\vert s,a)\). The log-derivative trick converts it into an expectation we can estimate by simply running the policy and observing rewards. Replacing \(G_t\) with advantage \(\hat{A}_t = G_t - V(s_t)\) reduces variance without bias (since \(\mathbb{E}[\nabla\log\pi \cdot b(s)] = 0\) for any state-dependent baseline).

REINFORCE Algorithm (Williams 1992) (Williams, 1992):

1. Sample complete trajectory \(\tau = (s_0, a_0, r_0, s_1, a_1, r_1, \ldots)\) under \(\pi_\theta\)

2. Compute return \(G_t = \sum_{k=0}^{T-t} \gamma^k r_{t+k}\) for each time step

3. Update: \(\theta \leftarrow \theta + \alpha \sum_t \nabla_\theta \log \pi_\theta(a_t\vert s_t) \cdot G_t\)

Tip

REINFORCE Intuition — “Reward-Weighted Maximum Likelihood”

\(\nabla_\theta \log \pi_\theta(a_t\vert s_t)\) is the direction that increases the probability of action \(a_t\). Multiplying by \(G_t\) means:

• High-reward trajectories: increase probability of all actions taken (positive \(G_t\))

• Low-reward trajectories: decrease probability of actions taken (negative \(G_t\) after baseline)

It’s supervised learning where the “labels” are the actions you took, weighted by how good they turned out to be.

Variance Reduction with Baseline:

\[ \nabla_\theta J(\theta) = \mathbb{E}_{\pi_\theta}\left[\sum_{t=0}^T \nabla_\theta \log \pi_\theta(a_t|s_t) \cdot (G_t - b(s_t))\right] \]

Any baseline \(b(s_t)\) that doesn’t depend on \(a_t\) keeps the gradient unbiased but reduces variance. Best choice: \(b(s_t) = V^\pi(s_t)\). Then \(G_t - V(s_t) \approx A^\pi(s_t, a_t)\) = advantage.

Warning

REINFORCE Limitations

• High variance: Each gradient uses one trajectory. Thousands of samples needed for stable updates.

• No bootstrapping: Must wait for full episode (no partial credit).

• Sample inefficient: Data is used once then discarded (on-policy).

• No step-size control: Can take catastrophically large policy steps.

These limitations motivate the progression: REINFORCE \(\to\) Actor-Critic \(\to\) TRPO \(\to\) PPO.

Actor-Critic Methods

Combine policy gradient (actor) with learned value function (critic) to reduce variance while maintaining the flexibility of policy optimization.

Architecture:

• Actor \(\pi_\theta(a\vert s)\): The policy. Proposes actions.

• Critic \(V_\phi(s)\) or \(Q_\phi(s,a)\): Evaluates how good a state/action is. Provides low-variance baseline.

Actor update (using advantage from critic):

\[ \nabla_\theta J = \mathbb{E}\left[\nabla_\theta \log \pi_\theta(a_t|s_t) \cdot \hat{A}_t\right], \quad \hat{A}_t = r_t + \gamma V_\phi(s_{t+1}) - V_\phi(s_t) \]

Critic update (minimize TD error):

\[ \mathcal{L}_\text{critic} = \mathbb{E}\left[(r_t + \gamma V_\phi(s_{t+1}) - V_\phi(s_t))^2\right] \]

Important

Evolution to PPO for LLMs

1. REINFORCE (Williams 1992): High variance, no bootstrapping \(\rightarrow\) impractical for LLMs

2. A2C/A3C (Mnih et al. 2016) (Advantage Actor-Critic): Uses TD-based advantage. Lower variance. But unbounded step sizes.

3. TRPO (Schulman et al. 2015): Constrains KL divergence between policy updates. Stable but expensive (second-order).

4. PPO (Schulman et al. 2017): Clips the policy ratio to achieve similar stability as TRPO with first-order optimization only. The standard for LLM RL training.

5. GRPO: Removes the critic entirely. Uses group statistics as baseline. Simpler and effective for verifiable rewards.

Generalized Advantage Estimation (GAE)

Motivation: The Actor-Critic framework needs a good estimate of the advantage \(A(s,a) = Q(s,a) - V(s)\) — how much better was this action than average? But there’s a fundamental tension:

• 1-step TD advantage (\(r_t + \gamma V(s_{t+1}) - V(s_t)\)): Low variance (only one random step), but biased — if the value function \(V\) is wrong, the advantage estimate is systematically off.

• Monte Carlo advantage (\(G_t - V(s_t)\)): Unbiased (uses actual returns), but high variance — the sum of many random rewards fluctuates wildly between episodes.

GAE (Schulman et al. 2016) (Schulman et al., 2016) provides a smooth interpolation between these extremes via a single parameter \(\lambda \in [0, 1]\). It takes an exponentially-weighted average of \(n\)-step advantage estimates for all \(n\), giving a principled way to trade bias for variance.

Core idea: Compute the 1-step TD error \(\delta_t\) at each timestep, then blend them with exponentially decaying weights \((\gamma\lambda)^l\) — recent TD errors get full weight, distant ones are down-weighted:

\[ \boxed{\hat{A}_t^{\text{GAE}} = \sum_{l=0}^{T-t} (\gamma\lambda)^l \delta_{t+l}, \quad \delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)} \]

Tip

What $\lambda$ Controls --- Bias-Variance Tradeoff

• \(\lambda = 0\): \(\hat{A}_t = \delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)\). Trust value function completely. Low variance, but biased if \(V\) is inaccurate.

• \(\lambda = 1\): \(\hat{A}_t = \sum_l \gamma^l r_{t+l} - V(s_t)\). Full Monte Carlo return minus baseline. Unbiased but very high variance.

• \(\lambda = 0.95\) (standard): Sweet spot. Mostly trusts \(V\) but corrects with actual returns for distant effects. Works because value head becomes accurate after initial training.

For LLMs specifically: \(\gamma = 1.0\) (no time discounting — all tokens matter equally in single-turn), \(\lambda = 0.95\).

Intuitive Mapping of Bias and Variance in GAE

In supervised learning, bias and variance stem from structural model assumptions. In reinforcement learning via GAE, they stem from how much you trust a flawed model versus how much you trust a chaotic environment:

• Bias (Systemic Misalignment): Arises when the estimator relies on the structural assumptions and imperfect predictions of the value network \(V_\theta\). If \(\theta\) is under-trained or lacks capacity, the baseline guesses are systematically wrong.

• Variance (Sample Jitteriness): Arises when the estimator relies on long, unconstrained environmental trajectories. Stochastic transitions, random seeds, and policy execution noise accumulate over long horizons, causing empirical sample rewards to swing wildly between rollouts.

The Architectural Spectrum: Boundary Analysis

The hyperparameter \(\lambda\) serves as a slide-rule between two fundamental estimation paradigms.

Important

High Bias / Low Variance Limit ($\lambda = 0$)

\[ \hat{A}_t^{\text{GAE}(\gamma, 0)} = \delta_t^V = r_t + \gamma V_\theta(s_{t+1}) - V_\theta(s_t) \]

• Behavior: The advantage is heavily dictated by the current state of parameters \(\theta\).

• Intuition: Highly biased because the network is grading its own performance over a 1-step window; if \(V_\theta\) is inaccurate, the gradient step is corrupted. Low variance because it ignores future stochastic events beyond step \(t+1\), leading to smooth, stable parameter updates.

• Risk: Policy traps in sub-optimal local minima — never discovers complex delayed reward sequences.

Important

Low Bias / High Variance Limit ($\lambda = 1$)

When \(\lambda = 1\), intermediate value terms telescopically cancel, reducing GAE to Monte Carlo return minus baseline: \[ \hat{A}_t^{\text{GAE}(\gamma, 1)} = \sum_{l=0}^{\infty} \gamma^l r_{t+l} - V_\theta(s_t) \]

• Behavior: Discards bootstrap look-aheads and sums up the literal reality of the entire episode.

• Intuition: Completely unbiased with respect to true environment dynamics — measures actual rewards instead of neural approximations. However, exhibits extreme variance: minor perturbations early in an episode can result in completely divergent total returns, causing policy updates to become erratic.

• Risk: Destructive gradient updates; training explosions.

The Trade-off Matrix

By selecting \(\lambda \in [0.95, 0.99]\), GAE minimizes the total mean squared error of the advantage estimate:

Operational comparison of GAE parameter choices.

Diagnostics for Tuning \(\lambda\)

Monitoring training curves yields direct insight into whether bias or variance dominates:

1. High Variance Indicators: Policy entropy drops precipitously while explained variance of the value function becomes highly negative or erratic \(\rightarrow\) policy updates are noisy. Remedy: Lower \(\lambda\) to smooth target updates.

2. High Bias Indicators: Agent achieves early stable training but completely fails to discover complex delayed reward sequences \(\rightarrow\) under-estimating long-horizon dependencies due to bootstrapping. Remedy: Raise \(\lambda\) closer to \(1.0\) to expose policy to real downstream trajectory signals.

On-Policy vs Off-Policy — Detailed Comparison

Tip

On/Off-Policy for RLHF Methods

PPO/GRPO are on-policy: Generate responses with current policy, compute advantages, update, discard data, generate again. This is why generation is 60% of compute — you regenerate every step.

DPO is off-policy: Train on a fixed preference dataset. No generation during training. Much cheaper but suffers from distribution shift (data becomes stale as policy changes).

Online DPO is a hybrid: Generates fresh data (on-policy generation) but uses DPO’s supervised loss (off-policy-style optimization). Gets benefits of both.

PPO’s cleverness: Uses the clip ratio \(r = \pi_\text{new}/\pi_\text{old}\) to squeeze multiple gradient steps from one batch of on-policy data (4 epochs), making it “slightly off-policy” in a controlled way.

Model-Based vs Model-Free

Tip

Why LLM RL is Model-Free

Language generation dynamics are trivial (append token to sequence — deterministic transitions). The “model” of the environment is not the bottleneck. What’s hard is the reward — predicting what humans will prefer. This makes model-based methods unnecessary for LLM RL.

The reward model in RLHF could be seen as a “model” in some sense (it predicts human preference), but it’s used as a reward signal, not for planning/simulation. LLM RL is fundamentally model-free policy optimization.

Reward Shaping

Reward shaping (Ng et al. 1999) is a technique where a developer modifies or supplements the environment’s original reward function. Its primary objective is to transform a sparse reward scenario — where the agent receives feedback only upon final task completion — into a dense reward scenario with intermediate feedback signals to accelerate convergence.

The Mathematical Framework

Let the original reward at time step \(t\) be \(R_t(s, a, s')\). The reshaped reward adds an auxiliary shaping function \(F\):

\[ \boxed{R'_t(s, a, s') = R_t(s, a, s') + F(s, a, s')} \]

Warning

The Risk of Naive Reshaping: Reward Hacking

If \(F(s, a, s')\) is arbitrarily designed, the agent will find structural loopholes to maximize auxiliary signals while ignoring the global objective.

Example: A navigation agent rewarded for reaching intermediate landmarks might learn to loop indefinitely around a single checkpoint to accumulate infinite rewards — without ever reaching the destination.

In LLMs: a model rewarded for “sounding confident” might learn to always start with “Absolutely!” regardless of accuracy.

Potential-Based Reward Shaping (PBRS)

To mathematically guarantee that reshaping does not alter the optimal policy, use Potential-Based Reward Shaping. The shaping function \(F\) is constrained to the difference in a scalar potential function \(\Phi\) across states:

\[ \boxed{F(s, a, s') = \gamma, \Phi(s') - \Phi(s)} \]

where \(\Phi: \mathcal{S} \to \mathbb{R}\) is a real-valued potential function evaluating the desirable proximity of a state to the goal, and \(\gamma\) is the discount factor.

The complete PBRS reward:

\[ R'(s, a, s') = R(s, a, s') + \gamma, \Phi(s') - \Phi(s) \]

Theoretical Guarantees

Important

PBRS Policy Invariance Theorem

• Policy Invariance: The optimal policy \(\pi^*\) under the reshaped reward \(R'\) is identical to the optimal policy under the original reward \(R\). The shaping cannot introduce sub-optimal behaviors.

• Loop Immunity: Any cyclic trajectory starting and ending at the same state results in net potential change of exactly zero (\(\Phi(s) - \Phi(s) = 0\)). The agent cannot exploit loops to hack the reward.

• Convergence Acceleration: While the optimal policy is unchanged, the shaped reward provides denser gradient signals, enabling the agent to converge 5–50\(\times\) faster in sparse reward environments.

RL Foundations for Language Models

Supervised fine-tuning (SFT) teaches a model to imitate demonstrations, but imitation has a ceiling: the model can never exceed the quality of its training data. Reinforcement learning breaks this barrier. By generating novel text, receiving reward feedback, and updating toward higher-reward behaviours, an RL-trained model can discover strategies that no human demonstrator wrote—producing outputs that are more helpful, more accurate, and better aligned with human preferences (Ouyang et al. 2022).

This is the mechanism behind every frontier model: GPT-4 (OpenAI 2023), Claude, Llama-3 (Grattafiori et al. 2024), and DeepSeek-R1 (DeepSeek-AI et al. 2025) all apply RL after SFT as the critical step that transforms a capable but unsteered model into an aligned assistant.

Two Paradigms for RL in LLMs

RL methods for language models fall into two broad paradigms, each suited to different goals:

Paradigm 1: Alignment via Human Preferences (RLHF/DPO).

The original motivation for applying RL to LLMs was alignment—making models helpful, harmless, and honest. Reinforcement Learning from Human Feedback (RLHF) (Ouyang et al. 2022; Ziegler et al. 2019; Christiano et al. 2017) trains a reward model from pairwise human judgments (“which response is better?”) and then optimizes the policy to maximize that learned reward. DPO (Rafailov et al. 2023) simplifies this by eliminating the reward model entirely, converting preferences directly into a supervised loss. Both approaches produce aligned assistants that follow instructions and respect safety constraints.

Paradigm 2: Capability Enhancement via Verifiable Rewards (RLVR).

More recently, RL has been used not just for alignment but for teaching new capabilities—particularly reasoning, mathematics, and code generation. Here the reward comes not from human preferences but from verifiable outcomes: did the model produce the correct answer? Did the code pass all tests? DeepSeek-R1 (DeepSeek-AI et al. 2025) demonstrated that GRPO with rule-based rewards (format correctness + answer accuracy) can train models to develop sophisticated chain-of-thought reasoning without any human preference data. This paradigm—RL from Verifiable Rewards (RLVR)—is now the dominant approach for building reasoning models and agentic systems.

Important

The Shared Foundation

Despite their different goals, both paradigms share the same core machinery:

• A policy \(\pi_\theta\) (the LLM) that generates text autoregressively

• A reward signal \(r(x, y)\) (learned from preferences or computed from verification)

• A KL constraint against a reference policy to prevent degenerate solutions

• Policy gradient optimization (PPO or GRPO) to update the model toward higher reward

The chapters in this part develop each component in detail.

Text Generation as an MDP

The key insight that makes RL applicable to language models is recasting autoregressive generation as a Markov Decision Process:

Tip

The LLM-as-Agent Analogy

Think of the LLM as an agent writing a response one token at a time. At each step, it looks at everything written so far (the state), chooses the next word (the action), and the page grows by one token (the transition). When the response is complete, a judge scores it (the reward). The goal: learn a writing strategy (a policy) that consistently earns high scores.

Formally, the MDP for text generation is:

• State \(s_t = (x, y_1, \ldots, y_{t-1})\): the prompt concatenated with all tokens generated so far.

• Action \(a_t \in \{1, \ldots, \vert \mathcal{V}\vert \}\): choosing the next token from the vocabulary (32K–128K options).

• Transition \(P(s_{t+1}\vert s_t, a_t)\): deterministic—just append the chosen token. No environment stochasticity.

• Reward \(r\): typically given only at the end of generation (sparse). For RLHF: the reward model score. For RLVR: correctness of the final answer.

• Policy \(\pi_\theta(a_t\vert s_t)\): the LLM’s next-token probability distribution—exactly what the softmax output already computes.

• Discount \(\gamma = 1.0\): episodes are finite (one response), so no discounting needed.

This mapping is powerful because the LLM already is a policy—its softmax output defines \(\pi_\theta(a_t\vert s_t)\) for every state. We don’t need to build a separate policy network; we just need to adjust the weights \(\theta\) so the model assigns higher probability to token sequences that earn higher reward.

The RLHF Pipeline

The classic RLHF pipeline (Ouyang et al. 2022) consists of four stages:

1. Supervised Fine-Tuning (SFT): Train a base model on high-quality demonstrations to produce a policy \(\pi_{\text{SFT}}\) that can follow instructions.

2. Reward Model Training: Collect human preference comparisons (\(y_w \succ y_l\) for the same prompt) and train a reward model \(R_\phi(x, y)\) using the Bradley-Terry objective.

3. RL Optimization: Use the reward model as a signal to optimize the policy via PPO or GRPO, subject to a KL constraint against \(\pi_{\text{SFT}}\).

4. Evaluation and Iteration: Evaluate the aligned model, collect new failure cases, and iterate.

For RLVR (reasoning/agentic training), stages 1–2 are replaced: the SFT model is trained on reasoning traces, and the reward model is replaced by a verifier (e.g., checking mathematical correctness). Stage 3 remains the same—PPO or GRPO optimization against the reward signal.

Tip

How LLM RL Differs from Classical RL

The LLM setting differs from classical RL in important ways:

• Deterministic transitions: The “next state” is just the concatenation of previous tokens—no stochastic environment.

• Sparse reward: Feedback is typically given once at the end of generation (outcome reward) or at key steps (process reward).

• Massive action space: 32K–128K possible tokens at every step, but exploration is implicit via temperature sampling.

• KL anchor: LLM RL is constrained to stay close to the SFT policy, preventing reward hacking at the cost of reduced exploration.

• No value function needed: GRPO eliminates the critic network entirely, using group-relative normalization of rewards instead.

These differences explain why PPO and GRPO dominate over DQN-style approaches for LLMs.

Roadmap of This Part

The chapters ahead build the complete RL-for-LLMs toolkit:

1. PPO (Chapter 5) — The clipped surrogate objective, GAE for advantage estimation, the critic network, and the full RLHF training loop. The workhorse behind GPT-4 and Claude.

2. DPO (Chapter 6) — Bypassing RL entirely by converting preferences into a contrastive supervised loss. Simpler but less flexible than online RL.

3. GRPO (Chapter 7) — DeepSeek’s critic-free algorithm that uses group-level reward normalization. The method behind DeepSeek-R1 and the dominant choice for reasoning model training.

4. Preference optimization variants (Chapter 8) — Online DPO, KTO, Best-of-N, and guidance on method selection.

5. Reward modeling (Chapter 9) — Bradley-Terry models, process vs. outcome rewards, rule-based rewards for RLVR, and multi-objective combinations.

6. SFT best practices (Chapter 10) — Sequence packing, chat templates, data mixing, and how SFT quality determines the RL ceiling.

7. Systems engineering (Chapter 11) — Distributed training at scale: parallelism strategies, generation–training decoupling, and infrastructure for hundreds of GPUs.

PPO — Proximal Policy Optimization

Motivation and History

Problem: Vanilla policy gradient updates have no constraint on step size. A single unlucky batch can push the policy into a region where it generates garbage \(\rightarrow\) garbage gets low rewards \(\rightarrow\) next gradient makes things worse \(\rightarrow\) unrecoverable collapse.

Solution history:

1. TRPO (Schulman et al. 2015) (2015): Constrain KL divergence between old and new policy. Works perfectly but requires expensive second-order optimization (Fisher information matrix, conjugate gradients).

2. PPO (2017) (Schulman et al. 2017): Achieve similar stability with a simple first-order clipped objective. 10\(\times\) simpler to implement, works almost as well, scales to distributed training trivially.

The Clipped Objective

The core innovation of PPO is a clipped surrogate objective that prevents destructively large policy updates while remaining simple to implement.

\[ \boxed{L^{\text{CLIP}}(\theta) = \mathbb{E}_t\left[\min\left(r_t(\theta)\hat{A}_t,; \text{clip}(r_t(\theta), 1{-}\epsilon, 1{+}\epsilon)\hat{A}_t\right)\right]} \]

where \(r_t(\theta) = \frac{\pi_\theta(a_t\vert s_t)}{\pi_{\theta_\text{old}}(a_t\vert s_t)}\) is the probability ratio.

Tip

Clipping Intuition — The Key Insight

The \(\min\) operator creates a pessimistic bound:

• Good action (\(\hat{A} > 0\)): We want to increase its probability. The surrogate \(r\hat{A}\) grows as \(r\) increases. But clip caps benefit at \(r = 1 + \epsilon\). “Don’t get greedy on one good example.”

• Bad action (\(\hat{A} < 0\)): We want to decrease its probability. \(r\hat{A}\) improves as \(r\) decreases. But clip caps benefit at \(r = 1 - \epsilon\). “Don’t forget too aggressively based on one bad example.”

Net effect: policy changes by at most \(\pm\)20% per update step. Prevents both catastrophic collapse and overconfident specialization.

Full PPO Loss

\[ L = L^{\text{CLIP}} - c_1 \underbrace{(V_\theta(s_t) - V^{\text{target}}_t)^2}_{\text{value loss}} + c_2 \underbrace{H[\pi_\theta(\cdot|s_t)]}_{\text{entropy bonus}} \]

• Value loss (\(c_1 = 0.1\)): Trains the critic to predict returns. Also clipped for stability.

• Entropy bonus (\(c_2 = 0.01\)): Prevents premature convergence to deterministic policy. Critical for exploration.

Derivation of the PPO Gradient and Update Rule

This section traces the mathematical path from the RL objective to the PPO update rule, showing why the clipped surrogate works.

Step 1: The RL Objective

The goal is to maximize expected cumulative reward under the policy:

\[ J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}\left[\sum_{t=0}^T r_t\right] \]

Step 2: Policy Gradient Theorem

The gradient of \(J(\theta)\) with respect to policy parameters:

\[ \boxed{\nabla_\theta J(\theta) = \mathbb{E}_{\pi_\theta}\left[\sum_{t=0}^T \nabla_\theta \log \pi_\theta(a_t|s_t) \cdot \hat{A}_t\right]} \]

where \(\hat{A}_t\) is the advantage function (how much better action \(a_t\) was compared to the average action in state \(s_t\)). This replaces the full return with the advantage to reduce variance.

Step 3: Importance Sampling for Off-Policy Data

PPO collects data using \(\pi_{\theta_{\text{old}}}\) but updates \(\pi_\theta\). To correct for this distribution mismatch, apply importance sampling:

\[ \nabla_\theta J(\theta) = \mathbb{E}_{\pi_{\theta_{\text{old}}}}\left[\frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{\text{old}}}(a_t|s_t)} \nabla_\theta \log \pi_\theta(a_t|s_t) \cdot \hat{A}_t\right] \]

Define the probability ratio \(r_t(\theta) = \frac{\pi_\theta(a_t\vert s_t)}{\pi_{\theta_{\text{old}}}(a_t\vert s_t)}\). Using the identity \(\nabla_\theta \log f = \frac{\nabla_\theta f}{f}\), we get:

\[ \nabla_\theta J(\theta) = \mathbb{E}_{\pi_{\theta_{\text{old}}}}\left[\nabla_\theta, r_t(\theta) \cdot \hat{A}_t\right] \]

This means maximizing the surrogate objective:

\[ L^{\text{CPI}}(\theta) = \mathbb{E}_t\left[r_t(\theta) \cdot \hat{A}_t\right] \]

Step 4: The Problem with Unconstrained Surrogates

\(L^{\text{CPI}}\) is a valid objective, but without constraints, a single gradient step can push \(r_t(\theta)\) far from 1.0, causing:

• Importance weights become extreme \(\rightarrow\) high variance

• Policy enters untested regions \(\rightarrow\) reward model gives unreliable scores

• Catastrophic collapse: policy generates garbage, can’t recover

TRPO solution: Constrain \(D_{\text{KL}}(\pi_{\theta_{\text{old}}} \\vert \pi_\theta) \leq \delta\). Requires second-order methods (expensive).

Step 5: PPO’s Clipped Surrogate (First-Order Approximation)

PPO replaces the hard KL constraint with a clipped objective that achieves similar behavior using only first-order gradients:

\[ \boxed{L^{\text{CLIP}}(\theta) = \mathbb{E}_t\left[\min!\left(r_t(\theta)\hat{A}_t,;\text{clip}(r_t(\theta), 1{-}\epsilon, 1{+}\epsilon)\hat{A}_t\right)\right]} \]

Derivation of the gradient:

Let \(L_t = \min(r_t \hat{A}_t,; \bar{r}_t \hat{A}_t)\) where \(\bar{r}_t = \text{clip}(r_t, 1{-}\epsilon, 1{+}\epsilon)\).

\[ \nabla_\theta L_t = \begin{cases} \nabla_\theta r_t(\theta) \cdot \hat{A}_t & \text{if } r_t \hat{A}_t < \bar{r}_t \hat{A}_t \text{ (unclipped term is smaller)} \\ 0 & \text{if } r_t \hat{A}_t \geq \bar{r}_t \hat{A}_t \text{ (clipped term is smaller, gradient = 0)} \end{cases} \]

Expanding the conditions:

• When \(\hat{A}_t > 0\) and \(r_t < 1+\epsilon\): Gradient flows normally — policy is encouraged to increase \(\pi_\theta(a_t\vert s_t)\).

• When \(\hat{A}_t > 0\) and \(r_t \geq 1+\epsilon\): Gradient is zero — policy has already increased enough, stop pushing.

• When \(\hat{A}_t < 0\) and \(r_t > 1-\epsilon\): Gradient flows normally — policy is encouraged to decrease \(\pi_\theta(a_t\vert s_t)\).

• When \(\hat{A}_t < 0\) and \(r_t \leq 1-\epsilon\): Gradient is zero — policy has already decreased enough, stop pushing.

Step 6: The Complete PPO Update Rule

Combining the clipped policy loss, value loss, and entropy bonus:

\[ \boxed{\theta_{k+1} = \theta_k + \alpha \cdot \nabla_\theta \left[L^{\text{CLIP}}(\theta) - c_1 L^{\text{VF}}(\theta) + c_2 H[\pi_\theta]\right]} \]

where:

\[ \begin{aligned} L^{\text{VF}}(\theta) &= \left(V_\theta(s_t) - V_t^{\text{target}}\right)^2 & &\text{(value function regression loss)} \\ H[\pi_\theta] &= -\sum_a \pi_\theta(a|s_t)\log\pi_\theta(a|s_t) & &\text{(entropy of the policy)} \end{aligned} \]

Tip

Summary: Why This Works

1. Policy gradient theorem gives us the direction to improve the policy.

2. Importance sampling lets us reuse data from \(\pi_{\theta_{\text{old}}}\) across multiple epochs.

3. Clipping prevents the importance weights from becoming extreme, keeping updates safe.

4. The min operator ensures we always take the more conservative of (clipped, unclipped) — a pessimistic lower bound on improvement.

5. Result: Monotonic improvement with probability 1, using only first-order gradients. No Hessians, no conjugate gradients, no line searches.

Rollout Buffer and Rollouts

In PPO, data management relies on a specialized, short-term storage system known as a Rollout Buffer. Unlike off-policy algorithms (DQN) that store experiences indefinitely in a replay buffer, PPO requires an ephemeral structure to satisfy its on-policy mathematical constraints.

What is a Rollout?

A rollout (trajectory) is a sequence of interactions generated by the agent running its current policy in the environment:

• The process: The agent observes a state, selects an action, receives a reward, and moves to the next state. It repeats for a fixed number of steps or until the episode ends.

• In LLMs/RLHF: A rollout consists of taking a prompt from a dataset and letting the language model generate a complete sequence of tokens token-by-token until an end-of-text marker is hit. Each token is one “step.”

The Rollout Buffer

The rollout buffer temporarily stores all data collected during the rollout phase. For every generated token/step, it records:

\[ \boxed{\mathcal{B} = \left\{ \left(s_t,; a_t,; \log\pi_{\theta_{\text{old}}}(a_t|s_t),; r_t,; V(s_t)\right) \right\}_{t=1}^{T}} \]

• \(s_t, a_t, r_t\): State, action taken, and reward at step \(t\).

• \(\log\pi_{\theta_{\text{old}}}(a_t\vert s_t)\): Log-probability of taking that action under the exact policy that generated it (needed for ratio computation).

• \(V(s_t)\): Value function’s baseline prediction (needed for GAE advantage computation).

The Rollout Buffer Lifecycle

The buffer operates in a strict three-phase clockwork cycle:

1. Collect: The active policy interacts with the environment to fill the buffer with fresh trajectories (for a 70B model with batch=128, max_tokens=512: up to 65K token-level transitions per rollout).

2. Train: Compute GAE advantages across trajectories. Run \(K\) epochs (typically 3–10) of mini-batch gradient descent to update policy weights using the clipped objective.

3. Purge: The entire buffer is completely wiped clean. Because PPO is on-policy, data generated by the old policy cannot be safely reused for the next update cycle — the ratio \(r_t(\theta)\) would become stale and the clipping guarantees would break.

Warning

Rollout Buffer vs Replay Buffer

Replay Buffer (DQN, SAC): Off-policy. Stores millions of transitions indefinitely. Random sampling. Data reused across many updates.

Rollout Buffer (PPO, GRPO): On-policy. Stores one batch of trajectories. Used for a few epochs, then discarded entirely. Fresh data required every cycle.

This is why PPO requires continuous generation — the buffer is emptied after every update, demanding fresh rollouts. This makes the generation bottleneck (60–70% of wall-clock time) particularly painful.

Tip

vLLM in RLHF Context

In RLHF training, vLLM is used for the generation phase (60–70% of wall-clock time). The policy model generates rollouts that are then scored by the reward model. Key benefits:

• Batched generation: Generate 256+ responses in parallel across prompts.

• Memory efficiency: Fit more concurrent generations \(\rightarrow\) higher GPU utilization during the generation bottleneck.

• Prefix sharing: When generating \(N=8\) responses per prompt (GRPO), the prompt KV is computed once and shared across all 8 — no redundant prefill.

• Integration: Frameworks like OpenRLHF and TRL use vLLM as the generation backend, separating generation workers (vLLM) from training workers (DeepSpeed/FSDP).

PPO for RLHF: The Full Loop

Note

Concrete PPO Step for a 70B Chat Model

Setup: Batch of 128 prompts, Llama-3-70B policy, 512 max tokens.

Step 1 — Generate: Sample 128 responses (temperature=0.7, top-p=0.9). This takes 60% of time.

Step 2 — Score: Reward model scores each (prompt, response) pair. Range: 0.2–0.95.

Step 3 — KL: Compute per-token KL: \(\text{KL}_t = \log\pi_\theta(y_t\vert y_{<t}) - \log\pi_\text{ref}(y_t\vert y_{<t})\). Mean KL across tokens: typically 3–8.

Step 4 — Final reward: \(R = r_\text{RM} - 0.05 \times \text{mean_KL}\) (only at last token).

Step 5 — GAE: Compute \(\hat{A}_t\) for each token position using value head predictions. Whiten advantages (zero mean, unit variance).

Step 6 — Update: 4 epochs of SGD on mini-batches of 16. Clip ratio \(\epsilon = 0.2\). Gradient norm clipping at 1.0.

Result: Policy improves by \(\sim\)0.005 win-rate per step. After 10K steps: 5–10% absolute improvement over SFT.

Warning

Tokenization Pitfalls in RL for LLMs

When computing per-token KL penalties and advantages, remember that tokenization determines what a “step” is. A single conceptual action (e.g., outputting “2024”) might span 1–4 tokens depending on the tokenizer. This creates subtle issues:

• KL accounting: Per-token KL sums to different totals for the same semantic content tokenized differently (e.g., rare words split into more subwords get higher total KL penalty).

• Credit assignment: GAE assigns advantage per token position—but semantic “decisions” often span multiple tokens. The model only truly “decides” at the first token of a word; subsequent subword tokens are largely deterministic.

• Reward placement: Placing reward only at the final token means all preceding tokens must propagate credit backward through GAE—longer responses suffer from more diluted signal.

Mitigation: Some systems normalize KL by sequence length, use word-level reward shaping, or apply reward at semantic boundaries rather than the final token.

Detailed Mechanics: Logits and Policy Updates

PPO manages two distinct parameter states in memory, which share the same neural network topology but hold different weight values during optimization:

Important

Core Architecture: Two Networks

1. The Policy Network (\(\pi_\theta\)): The active, live network parameterized by weights \(\theta\). Continuously updated via backpropagation during optimization.

2. The Old Policy Network (\(\pi_{\theta_{\text{old}}}\)): A frozen snapshot parameterized by weights \(\theta_{\text{old}}\). Acts as a static anchor during a single optimization cycle to prevent the policy from shifting too drastically.

Phase 1: Rollout (Data Collection)

During data collection, the agent interacts with the environment for \(T\) steps. At each time-step \(t\):

1. The environment yields the current state/observation \(s_t\) (for LLMs: prompt + tokens generated so far).

2. State \(s_t\) is passed through the current network snapshot (\(\theta_{\text{old}}\)).

3. The network outputs raw unnormalized values — logits \(z_{\text{old}}\) — a vector of size \(\vert V\vert\) (vocabulary size 32K–128K).

4. Probabilities are computed via Softmax:

\[ \boxed{P(a \mid s_t) = \text{Softmax}(z_{\text{old}}) = \frac{\exp(z_{\text{old}, a})}{\sum_{j=1}^{|V|} \exp(z_{\text{old}, j})}} \]

5. An action \(a_t\) (next token) is sampled from \(P(a \mid s_t)\), and the transition tuple \(\langle s_t, a_t, r_t, s_{t+1} \rangle\) along with \(\log \pi_{\theta_{\text{old}}}(a_t \mid s_t)\) is stored in the rollout buffer.

Tip

Why Store Log-Probabilities?

Storing \(\log \pi_{\theta_{\text{old}}}(a_t \mid s_t)\) as a scalar during rollout avoids re-running the frozen network during optimization. This saves one full forward pass per mini-batch — significant for 70B models.

Phase 2: Optimization Loop (Mini-Batch Updates)

Once the rollout buffer is full, PPO runs \(K\) epochs (typically 3–10) over mini-batches. For every gradient step, logits are generated for both policies using the stored state \(s_t\):

Old Policy Evaluation (frozen):

\[ z_{\text{old}} = f(s_t; \theta_{\text{old}}) \quad \longrightarrow \quad \log \pi_{\theta_{\text{old}}}(a_t \mid s_t) = \text{LogSoftmax}(z_{\text{old}})[a_t] \]

Implementation shortcut: reuse the stored scalar from rollout instead of re-computing.

Live Policy Evaluation (updating):

\[ z_{\text{new}} = f(s_t; \theta) \quad \longrightarrow \quad \log \pi_\theta(a_t \mid s_t) = \text{LogSoftmax}(z_{\text{new}})[a_t] \]

Because \(\theta\) updates after every mini-batch gradient step, \(z_{\text{new}}\) changes continuously throughout the optimization loop, whereas \(z_{\text{old}}\) remains perfectly static.

From Logits to Probability Ratio

The core PPO ratio measures how much more or less likely an action is under the new policy vs the old:

\[ \boxed{r_t(\theta) = \frac{\pi_\theta(a_t \mid s_t)}{\pi_{\theta_{\text{old}}}(a_t \mid s_t)}} \]

To avoid catastrophic numerical underflow/overflow from dividing raw probabilities, the calculation is performed in log-space:

\[ \begin{aligned} \log \pi_\theta(a_t \mid s_t) &= \text{LogSoftmax}(z_{\text{new}})[a_t] \\ \log \pi_{\theta_{\text{old}}}(a_t \mid s_t) &= \text{LogSoftmax}(z_{\text{old}})[a_t] \end{aligned} \]

The ratio is recovered via exponentiation of the difference:

\[ \boxed{r_t(\theta) = \exp!\left(\log \pi_\theta(a_t \mid s_t) - \log \pi_{\theta_{\text{old}}}(a_t \mid s_t)\right)} \]

This ratio is injected into the PPO clipping objective:

\[ \boxed{\mathcal{L}^{\text{CLIP}}(\theta) = \hat{\mathbb{E}}_t \left[ \min!\left(r_t(\theta)\hat{A}_t, ;\text{clip}(r_t(\theta),, 1{-}\epsilon,, 1{+}\epsilon),\hat{A}_t\right) \right]} \]

Tip

How Clipping Works

• If \(\hat{A}_t > 0\) (good action): ratio is clipped at \(1+\epsilon\) — cannot over-exploit good actions.

• If \(\hat{A}_t < 0\) (bad action): ratio is clipped at \(1-\epsilon\) — cannot over-penalize bad actions.

• The \(\min(\cdot)\) ensures we always take the more conservative estimate.

Result: monotonic improvement within a trust region — no catastrophic collapses.

The PPO Weight Lifecycle

Evolution of \(\theta\) and \(\theta_{\text{old}}\) across PPO training phases.

Continuous Action Spaces Extension

For continuous action spaces (not typical for LLMs, but important for robotics RL), the network outputs distribution parameters instead of discrete logits:

• Predicted mean vector \(\mu\)

• Predicted standard deviation vector \(\sigma\)

Log-probabilities are computed via the Gaussian log-PDF:

\[ \boxed{\log \pi(a_t \mid s_t) = -\frac{1}{2}\left(\frac{a_t - \mu}{\sigma}\right)^{!2} - \log(\sigma) - \frac{1}{2}\log(2\pi)} \]

The ratio \(r_t(\theta) = \exp(\log \pi_\theta - \log \pi_{\theta_{\text{old}}})\) is then computed identically and fed into the same clipping objective.

TRL Implementation

The HuggingFace TRL library (Werra et al. 2022) provides production-ready implementations of all major RL methods for LLMs.

from trl import PPOConfig, PPOTrainer, AutoModelForCausalLMWithValueHead
from transformers import AutoTokenizer
from peft import LoraConfig

# Model setup
model = AutoModelForCausalLMWithValueHead.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct",
    torch_dtype=torch.bfloat16, device_map="auto",
    peft_config=LoraConfig(r=64, lora_alpha=16, target_modules=["q_proj","v_proj","k_proj","o_proj"])
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")

# PPO config with all critical hyperparameters
ppo_config = PPOConfig(
    learning_rate=1.5e-6,        # Low LR for stability
    batch_size=128,              # Prompts per step
    mini_batch_size=16,          # Gradient accumulation unit
    ppo_epochs=4,                # Epochs per batch (reuse data)
    gamma=1.0,                   # No discounting (single turn)
    lam=0.95,                    # GAE lambda
    cliprange=0.2,               # PPO epsilon
    cliprange_value=0.2,         # Value function clipping
    vf_coef=0.1,                 # Value loss coefficient
    init_kl_coef=0.05,           # Initial KL penalty
    target_kl=6.0,               # Adaptive KL target
    whiten_rewards=True,         # Normalize advantages
    gradient_accumulation_steps=4,
    max_grad_norm=1.0,
)

ppo_trainer = PPOTrainer(config=ppo_config, model=model, tokenizer=tokenizer,
    dataset=prompt_dataset, data_collator=collator)

# Training loop
for batch in ppo_trainer.dataloader:
    # 1. Generate responses
    query_tensors = batch["input_ids"]
    response_tensors = ppo_trainer.generate(
        query_tensors, max_new_tokens=512, temperature=0.7, top_p=0.9, do_sample=True
    )
    # 2. Score with reward model
    texts = [tokenizer.decode(r, skip_special_tokens=True) for r in response_tensors]
    rewards = [torch.tensor(reward_model.score(q, r)) for q, r in zip(batch["query"], texts)]
    # 3. PPO update (handles KL, GAE, clipping internally)
    stats = ppo_trainer.step(query_tensors, response_tensors, rewards)
    # Monitor: stats["ppo/mean_scores"], stats["ppo/policy/approx_kl"]

Critical Hyperparameters

DPO — Direct Preference Optimization

Motivation

PPO requires 4 models in memory (policy, reference, reward model, value head), complex RL infrastructure, and is notoriously unstable. DPO (Rafailov et al. 2023) asks: can we skip the RL and learn directly from preferences?

Key insight: The optimal policy under the RLHF objective (reward maximization + KL penalty) has a closed-form solution. We can derive a supervised loss that implicitly optimizes the same objective.

Mathematical Derivation

Step 1: RLHF objective: \(\max_\pi \mathbb{E}_{x,y\sim\pi}[r(x,y)] - \beta D_\text{KL}[\pi\\vert \pi_\text{ref}]\)

Step 2: The optimal solution is: \(\pi^*(y\vert x) = \frac{1}{Z(x)} \pi_\text{ref}(y\vert x) \exp\left(\frac{r(x,y)}{\beta}\right)\)

Step 3: Rearrange to express reward in terms of policy: \(r(x,y) = \beta \log \frac{\pi^*(y\vert x)}{\pi_\text{ref}(y\vert x)} + \beta \log Z(x)\)

Step 4: Substitute into Bradley-Terry preference model \(P(y_w \succ y_l) = \sigma(r(y_w) - r(y_l))\). The \(Z(x)\) cancels!

\[ \boxed{\mathcal{L}_\text{DPO}(\theta) = -\mathbb{E}_{(x, y_w, y_l)}\left[\log\sigma\left(\beta\log\frac{\pi_\theta(y_w|x)}{\pi_\text{ref}(y_w|x)} - \beta\log\frac{\pi_\theta(y_l|x)}{\pi_\text{ref}(y_l|x)}\right)\right]} \]

Tip

What DPO Actually Does

Define the implicit reward as \(\hat{r}(x,y) = \beta\log\frac{\pi_\theta(y\vert x)}{\pi_\text{ref}(y\vert x)}\).

DPO minimizes the cross-entropy loss where the “label” is: chosen should have higher implicit reward than rejected. The margin is controlled by \(\beta\):

• Large \(\beta\): need large margin \(\rightarrow\) policy moves aggressively \(\rightarrow\) risk forgetting

• Small \(\beta\): small margin suffices \(\rightarrow\) policy stays close to reference \(\rightarrow\) conservative

The reference model acts as a regularizer: the policy must “justify” any deviation from it by showing preference alignment.

Gradient Analysis

The DPO gradient decomposes as:

\[ \nabla_\theta \mathcal{L} = -\beta \cdot \underbrace{\sigma(-\hat{r}_w + \hat{r}_l)}_{\text{weight: higher when model is wrong}} \cdot \left[\nabla_\theta \log\pi_\theta(y_w|x) - \nabla_\theta \log\pi_\theta(y_l|x)\right] \]

Interpretation: The gradient increases probability of chosen and decreases rejected. The weight is largest when the model currently prefers the wrong answer — it focuses learning on “confusing” pairs.

Note

Concrete DPO Example

Prompt: “Explain quantum entanglement to a 10-year-old.”

Chosen (\(y_w\)): “Imagine you have two magic coins. When you flip one and it’s heads, the other one instantly becomes tails, no matter how far apart they are!”
\(\log\pi_\theta(y_w\vert x) = -15.3\), \(\log\pi_\text{ref}(y_w\vert x) = -16.1\)

Rejected (\(y_l\)): “Quantum entanglement is a phenomenon where two particles become correlated such that the quantum state of one particle cannot be described independently.”
\(\log\pi_\theta(y_l\vert x) = -12.8\), \(\log\pi_\text{ref}(y_l\vert x) = -12.5\)

Implicit rewards: \(\hat{r}_w = 0.1 \times ((-15.3) - (-16.1)) = 0.08\), \(\hat{r}_l = 0.1 \times ((-12.8) - (-12.5)) = -0.03\)

Loss input: \(\sigma(0.08 - (-0.03)) = \sigma(0.11) = 0.527\)

Loss: \(-\log(0.527) = 0.64\) — The model barely prefers the chosen. Gradient will push hard.

After training: chosen probability increases, rejected decreases, until margin stabilizes around \(1/(2\beta)\).

TRL Implementation

The following shows a minimal working example using HuggingFace TRL.

from trl import DPOConfig, DPOTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig
from datasets import load_dataset

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B-Instruct",
    torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")

# Dataset format: {"prompt": str, "chosen": str, "rejected": str}
dataset = load_dataset("argilla/ultrafeedback-binarized-preferences")

lora_config = LoraConfig(r=64, lora_alpha=16, lora_dropout=0.05,
    target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"])

dpo_config = DPOConfig(
    output_dir="./dpo_output",
    beta=0.1,                    # KL regularization strength
    learning_rate=5e-7,          # Very low LR for stability
    loss_type="sigmoid",         # Standard DPO loss
    max_length=2048,             # Max sequence length
    max_prompt_length=1024,      # Truncation for prompts
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,  # Effective batch = 16
    gradient_checkpointing=True,
    bf16=True,
    num_train_epochs=1,          # DPO overfits fast - 1 epoch!
    warmup_ratio=0.1,
    logging_steps=10,
    eval_strategy="steps",
    eval_steps=200,
    save_strategy="steps",
    save_steps=500,
)

trainer = DPOTrainer(
    model=model,
    ref_model=None,             # With LoRA, ref = base model (no copy needed!)
    args=dpo_config,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    tokenizer=tokenizer,
    peft_config=lora_config,
)
trainer.train()
# Key metrics to monitor: train/rewards/chosen, train/rewards/rejected, train/rewards/margins

How DPO Works: Full Mechanics

This section provides the complete computational details of DPO — what happens at the token level during training.

Sequence-Level Log-Probabilities

The key quantity in DPO is the log-probability of an entire sequence \(y = (y_1, y_2, \ldots, y_T)\) given prompt \(x\). This is computed as the sum of per-token log-probabilities:

\[ \boxed{\log \pi_\theta(y|x) = \sum_{t=1}^{T} \log \pi_\theta(y_t \mid x, y_{<t})} \]

Each term \(\log \pi_\theta(y_t \vert x, y_{<t})\) is the log-softmax output at position \(t\) for the actual token \(y_t\) in the sequence. This is identical to the cross-entropy loss used in standard language modeling — but here we sum rather than average.

Critical detail: The gradient flows through every token position in both \(y_w\) and \(y_l\). There is no masking of intermediate tokens — every token contributes to the sequence-level log-probability.

The DPO Loss Decomposed

Starting from the loss:

\[ \mathcal{L}_{\text{DPO}}(\theta) = -\mathbb{E}_{(x, y_w, y_l) \sim \mathcal{D}}!\left[\log \sigma!\left(\beta \cdot h_\theta(x, y_w, y_l)\right)\right] \]

where the “implicit reward margin” \(h_\theta\) is:

\[ h_\theta(x, y_w, y_l) = \underbrace{\log \frac{\pi_\theta(y_w|x)}{\pi_{\text{ref}}(y_w|x)}}_{\text{chosen reward proxy}} - \underbrace{\log \frac{\pi_\theta(y_l|x)}{\pi_{\text{ref}}(y_l|x)}}_{\text{rejected reward proxy}} \]

Expanding into token-level terms:

\[ \boxed{h_\theta = \sum_{t=1}^{|y_w|}!\left[\log\pi_\theta(y_w^t | x, y_w^{<t}) - \log\pi_{\text{ref}}(y_w^t | x, y_w^{<t})\right] - \sum_{t=1}^{|y_l|}!\left[\log\pi_\theta(y_l^t | x, y_l^{<t}) - \log\pi_{\text{ref}}(y_l^t | x, y_l^{<t})\right]} \]

Forward Pass: Step by Step

For one training example \((x, y_w, y_l)\):

1. Concatenate: Form two sequences: \([x; y_w]\) and \([x; y_l]\). Pad to equal length within the batch.

2. Forward pass (policy \(\pi_\theta\)): Run both sequences through the model. Collect logits at every response position.

3. Extract log-probs: At each position \(t\) in the response, take \(\log\text{softmax}(\text{logits}_t)[y_t]\) — the log-probability of the actual token.

4. Sum over tokens:

\[ \begin{aligned} \text{logp_chosen} &= \sum_{t \in \text{response positions}} \log\pi_\theta(y_w^t | x, y_w^{<t}) \\ \text{logp_rejected} &= \sum_{t \in \text{response positions}} \log\pi_\theta(y_l^t | x, y_l^{<t}) \end{aligned} \]

5. Subtract reference (pre-computed or from second forward pass):

\[ \begin{aligned} \text{ratio_w} &= \text{logp_chosen} - \text{ref_logp_chosen} \\ \text{ratio_l} &= \text{logp_rejected} - \text{ref_logp_rejected} \end{aligned} \]

6. Compute loss: \(\mathcal{L} = -\log\sigma(\beta \cdot (\text{ratio_w} - \text{ratio_l}))\)

7. Backward pass: Gradients flow back through steps 5 \(\rightarrow\) 4 \(\rightarrow\) 3 \(\rightarrow\) 2 to update \(\theta\).

Token-Level Gradient Analysis

Does every token get a gradient? Yes. The gradient with respect to the logits at position \(t\) in the chosen sequence is:

\[ \frac{\partial \mathcal{L}}{\partial \text{logits}_t^{(w)}} = -\underbrace{\sigma(-\beta \cdot h_\theta)}_{\text{scaling factor}} \cdot \beta \cdot \frac{\partial \log\pi_\theta(y_w^t | \cdot)}{\partial \text{logits}_t^{(w)}} \]

Key insight: The scaling factor \(\sigma(-\beta \cdot h_\theta)\) is shared across all tokens in both sequences. It acts as an adaptive learning rate:

• When \(h_\theta\) is small (model can’t distinguish chosen from rejected): scaling \(\approx 0.5\) — strong gradient, learn aggressively.

• When \(h_\theta\) is large (model already prefers chosen): scaling \(\approx 0\) — negligible gradient, don’t over-fit.

Effect on chosen tokens: Probability is increased (log-prob pushed up).
Effect on rejected tokens: Probability is decreased (log-prob pushed down).
Relative to reference: Only the difference from \(\pi_{\text{ref}}\) matters. If the model already assigns high probability to the chosen response (matching the reference), there’s little gradient.

Per-Token vs. Sequence-Level: Length Normalization

A subtle issue: longer sequences naturally have lower log-probabilities (more terms summed, each \(\leq 0\)). If \(\vert y_w\vert \gg \vert y_l\vert\), the loss can be biased toward preferring shorter responses.

Solutions:

• Length-normalized DPO: Replace \(\log\pi_\theta(y\vert x)\) with \(\frac{1}{\vert y\vert }\sum_t \log\pi_\theta(y_t\vert \cdot)\). Used in some implementations (SimPO adopts this).

• Standard DPO: Uses raw sum (no normalization). This implicitly penalizes verbosity — the model must assign high probability to every token in the chosen response.

• Practical impact: On benchmarks, length-normalized DPO reduces length gaming but can hurt instruction-following quality. Standard (unnormalized) is more common in production.

Label Masking: What Gets Gradients

Important

Which Tokens Receive Gradient in DPO

• Prompt tokens (\(x\)): NO gradient. The loss is computed only over response positions. Prompt tokens provide context but their logits don’t contribute to \(\log\pi(y\vert x)\).

• Chosen response tokens (\(y_w\)): ALL tokens get gradient. Each \(y_w^t\) contributes to the sum. Gradient pushes their probabilities up.

• Rejected response tokens (\(y_l\)): ALL tokens get gradient. Each \(y_l^t\) contributes to the sum. Gradient pushes their probabilities down.

• Padding tokens: NO gradient. Masked out with attention mask.

Pseudocode: DPO Training Step

Note

DPO Forward + Backward (PyTorch-style)

def dpo_loss(model, ref_model, batch, beta=0.1):
    """One DPO training step."""
    # batch contains: input_ids_chosen, input_ids_rejected,
    #                 labels_chosen, labels_rejected (prompt masked to -100)

    # 1. Forward pass: get per-token log-probs
    logps_chosen = get_sequence_logprob(model, batch["chosen"])
    logps_rejected = get_sequence_logprob(model, batch["rejected"])

    # 2. Reference log-probs (pre-computed or computed here)
    with torch.no_grad():
        ref_logps_chosen = get_sequence_logprob(ref_model, batch["chosen"])
        ref_logps_rejected = get_sequence_logprob(ref_model, batch["rejected"])

    # 3. Compute implicit reward margins
    chosen_rewards = beta * (logps_chosen - ref_logps_chosen)
    rejected_rewards = beta * (logps_rejected - ref_logps_rejected)

    # 4. DPO loss = -log(sigmoid(chosen_reward - rejected_reward))
    loss = -F.logsigmoid(chosen_rewards - rejected_rewards).mean()
    return loss

def get_sequence_logprob(model, sequences):
    """Sum of log-probs over response tokens only."""
    outputs = model(sequences["input_ids"], attention_mask=sequences["mask"])
    logits = outputs.logits[:, :-1, :]  # Shift for next-token prediction

    # Gather log-prob of actual tokens
    labels = sequences["labels"][:, 1:]  # Shifted labels
    log_probs = F.log_softmax(logits, dim=-1)
    token_logps = log_probs.gather(-1, labels.unsqueeze(-1)).squeeze(-1)

    # Mask: only sum over response tokens (labels != -100)
    mask = (labels != -100).float()
    return (token_logps * mask).sum(dim=-1)  # Shape: [batch_size]

Common Pitfalls

Warning

DPO Implementation Pitfalls

• Forgetting to mask the prompt: If prompt tokens are included in the log-prob sum, the model optimizes for prompt likelihood (useless) and the effective \(\beta\) is wrong.

• Using mean instead of sum: \(\frac{1}{T}\sum_t \log\pi\) vs. \(\sum_t \log\pi\) — these give different implicit length penalties. Must be consistent between \(\pi_\theta\) and \(\pi_{\text{ref}}\).

• Stale reference model: If \(\pi_{\text{ref}}\) is too far from \(\pi_\theta\) (e.g., base model vs. fine-tuned), the KL term dominates and gradients vanish. Solution: use the SFT checkpoint (not base) as reference.

• \(\beta\) too large: Magnifies log-prob differences \(\rightarrow\) sigmoid saturates \(\rightarrow\) zero gradients. Start with \(\beta = 0.1\), tune in \([0.05, 0.5]\).

• \(\beta\) too small: Theoretically allows more freedom from reference (weaker KL constraint), but the gradient \(\propto \beta \cdot \sigma(-\beta h)\) becomes vanishingly small \(\rightarrow\) loss landscape is flat \(\rightarrow\) extremely slow convergence. The model has “permission” to move far but receives almost no signal telling it where to move.

DPO Variants and When Each Fails

Warning

When DPO Fails

1. Distribution shift: Preference data from old model. Current policy generates different text \(\rightarrow\) loss is optimizing on irrelevant examples.

2. No exploration: Can’t discover behaviors not in dataset. Stuck in local optimum.

3. Reference collapse: If reference is too strong, policy can’t move. If too weak, no regularization.

4. Data quality: Noisy labels poison training. Unlike PPO which averages over many samples, DPO memorizes individual pairs.

5. Preference data diversity: Ensure chosen/rejected pairs cover the full spectrum of quality differences (not just good-vs-terrible). Pairs that differ in approach, not just quality, teach richer policy distinctions.

\(\beta\) Selection Guide

DPO Batch Size Configuration and Scaling

Unlike standard SFT which operates on single-sequence token predictions, DPO leverages a pairwise loss comparing a preferred sequence against a dispreferred sequence. This fundamentally alters memory utilization and optimization stability.

Global Batch Size Target

Empirical evidence across DPO implementations establishes an optimal global batch size range:

\[ \boxed{B_{\text{global}} \in [32, 128]} \]

• \(B_{\text{global}} < 32\): Severe gradient noise in implicit reward estimation \(\rightarrow\) policy oscillates destructively between alignment goals (helpfulness vs. safety).

• \(B_{\text{global}} > 128\): Diminishing returns on convergence velocity; high communication overhead across distributed compute.

Mathematical Decomposition

Because DPO loads two model copies simultaneously (active policy \(\pi_\theta\) + frozen reference \(\pi_{\text{ref}}\)), per-sequence memory is doubled. The global batch size is decomposed as:

\[ \boxed{B_{\text{global}} = B_{\text{micro}} \times N_{\text{GPUs}} \times K_{\text{accum}}} \]

• \(B_{\text{micro}}\): Per-device micro-batch size (preference pairs per forward pass).

• \(N_{\text{GPUs}}\): Number of parallel data-processing devices.

• \(K_{\text{accum}}\): Gradient accumulation steps before weight update.

The pairing multiplier: A single DPO data instance contains a prompt (\(x\)), chosen (\(y_w\)), and rejected (\(y_l\)). The actual tensor load per micro-batch:

\[ T_{\text{sequences}} = 2 \times B_{\text{micro}} \]

For models \(>\)7B parameters on 80GB GPUs with context lengths 4096–8192 tokens, the physical limit is rigidly constrained to \(B_{\text{micro}} \in [1, 2]\).

Distributed Scaling Configurations

Distributed scaling profiles for DPO training (\(B_{\text{global}} = 64\) target).

VRAM Optimization: Pre-computing Reference Log-Probabilities

The DPO loss:

\[ \mathcal{L}_{\text{DPO}}(\theta) = -\mathbb{E}_{(x, y_w, y_l)}!\left[\log \sigma!\left(\beta \log \frac{\pi_\theta(y_w|x)}{\pi_{\text{ref}}(y_w|x)} - \beta \log \frac{\pi_\theta(y_l|x)}{\pi_{\text{ref}}(y_l|x)}\right)\right] \]

Because \(\pi_{\text{ref}}\) is completely static throughout training, its outputs can be pre-computed:

Important

Reference Model Eviction Strategy

1. Execute a forward pass over dataset \(\mathcal{D}\) using only \(\pi_{\text{ref}}\) before training begins.

2. Cache the scalars \(\log \pi_{\text{ref}}(y_w\vert x)\) and \(\log \pi_{\text{ref}}(y_l\vert x)\) to disk.

3. Evict \(\pi_{\text{ref}}\) completely from GPU memory.

Result: Available GPU memory doubles \(\rightarrow\) can increase \(B_{\text{micro}}\) from 1–2 to 4–8, maximizing hardware utilization and training throughput.

Implementation: In TRL, set precompute_ref_log_probs=True in DPOConfig. For 70B models, this saves \(\sim\)140GB of GPU memory across the cluster.

DPO Extensions and Variants

Direct Preference Optimization (DPO) reformulates RLHF as a supervised learning problem by deriving a closed-form mapping between the reward function and the optimal policy. The standard DPO loss is:

\[ \mathcal{L}_{\text{DPO}}(\theta) = -\mathbb{E}_{(q,y_w,y_l)}!\left[ \log \sigma!\left( \beta \log \frac{\pi_\theta(y_w|q)}{\pi_{\text{ref}}(y_w|q)} - \beta \log \frac{\pi_\theta(y_l|q)}{\pi_{\text{ref}}(y_l|q)} \right) \right], \]

where \(y_w\) is the preferred (winning) response, \(y_l\) is the dispreferred (losing) response, and \(\beta\) controls the strength of the KL penalty. The following subsections cover the most important extensions and variants.

f-DPO – Generalised f-Divergence DPO

Tip

Beyond Reverse KL

Standard DPO uses reverse KL divergence as the regulariser between policy and reference. Reverse KL is mode-seeking: it prefers to concentrate probability mass on a few high-reward responses. Forward KL is mode-covering: it spreads probability mass to cover all plausible responses. f-DPO (J. Wang et al. 2024) generalises to any f-divergence, allowing practitioners to trade off these behaviours.

The f-DPO loss replaces the log-ratio with the derivative of the f-divergence generator:

\[ \mathcal{L}_{f\text{-DPO}} = -\mathbb{E}!\left[ f'!\left(\frac{\pi_\theta(y_w|q)}{\pi_{\text{ref}}(y_w|q)}\right) - f'!\left(\frac{\pi_\theta(y_l|q)}{\pi_{\text{ref}}(y_l|q)}\right) \right], \]

where \(f'\) is the derivative of the f-divergence generator function.

Important

f-Divergence Options in TRL

• reverse_kl: \(f'(u) = \log u\). Standard DPO. Mode-seeking.

• forward_kl: \(f'(u) = -1/u\). Mode-covering. Better diversity.

• js_divergence: \(f'(u) = \log(2u/(u+1))\). Balanced mode-seeking/covering.

• alpha_divergence: \(f'(u) = u^{\alpha-1}\). Interpolates between forward and reverse KL.

Note

f-DPO in TRL

from trl import DPOConfig, DPOTrainer

# Jensen-Shannon divergence (balanced)
config = DPOConfig(
    f_divergence_type="js_divergence",
    beta=0.1,
)

# Alpha divergence (alpha=0: forward KL, alpha=1: reverse KL)
config_alpha = DPOConfig(
    f_divergence_type="alpha_divergence",
    f_alpha_divergence_coef=0.5,   # alpha parameter
    beta=0.1,
)

trainer = DPOTrainer(
    model=model,
    ref_model=ref_model,
    args=config,
    train_dataset=dataset,
)

Important

When to Use f-DPO

• Use JS divergence when you want a balance between diversity and quality.

• Use forward KL for creative tasks where diversity is paramount.

• Use reverse KL (standard DPO) for tasks with a single correct answer.

• Use alpha divergence to continuously interpolate and tune the trade-off.

Robust DPO

Tip

Noisy Labels in Preference Data

Human preference annotations are noisy. Annotators disagree, make mistakes, and sometimes flip the preferred/dispreferred labels. Standard DPO treats all labels as ground truth, which can cause the model to overfit to noise. Robust DPO (Chowdhury et al. 2024) analytically debiases the loss under a known noise model.

Assume each label is flipped with probability \(\epsilon\) (the noise rate). The debiased loss is:

\[ \boxed{ \mathcal{L}_{\text{robust}} = \frac{(1-\epsilon),\mathcal{L}_{\text{DPO}}(y_w, y_l) - \epsilon,\mathcal{L}_{\text{DPO}}(y_l, y_w)}{1 - 2\epsilon}, } \]

where \(\mathcal{L}_{\text{DPO}}(y_w, y_l)\) is the standard DPO loss treating \(y_w\) as preferred, and \(\mathcal{L}_{\text{DPO}}(y_l, y_w)\) is the loss with labels flipped. This correction removes the bias introduced by label noise.

Important

Intuition for Robust DPO

The formula is a linear combination that “subtracts out” the contribution of flipped labels. When \(\epsilon=0\), it reduces to standard DPO. When \(\epsilon=0.5\), the denominator goes to zero – the labels are pure noise and no learning is possible. In practice, \(\epsilon \in [0.05, 0.2]\) covers most real annotation noise levels.

Note

Robust DPO in TRL

from trl import DPOConfig, DPOTrainer

config = DPOConfig(
    loss_type="robust",
    label_smoothing=0.1,   # corresponds to epsilon = 0.1
    beta=0.1,
)

trainer = DPOTrainer(
    model=model,
    ref_model=ref_model,
    args=config,
    train_dataset=dataset,
)

TR-DPO – Trust Region DPO

Tip

Stale Reference Model Problem

Standard DPO uses a fixed reference model \(\pi_{\text{ref}}\) throughout training. As the policy \(\pi_\theta\) improves, the KL penalty \(\beta \log(\pi_\theta/\pi_{\text{ref}})\) grows, eventually dominating the loss and preventing further improvement. TR-DPO (Gorbatenko 2024) periodically updates the reference model to track the current policy.

TR-DPO updates the reference model using an exponential moving average (EMA):

\[ \pi_{\text{ref}}^{(t+1)} \leftarrow \alpha \cdot \pi_\theta^{(t)} + (1-\alpha) \cdot \pi_{\text{ref}}^{(t)}, \]

where \(\alpha \in (0,1)\) is the mixup coefficient. This is applied every \(T_{\text{sync}}\) gradient steps.

Note

TR-DPO in TRL

from trl import DPOConfig, DPOTrainer

config = DPOConfig(
    loss_type="sigmoid",        # standard DPO loss
    sync_ref_model=True,        # enable TR-DPO
    ref_model_mixup_alpha=0.6,  # alpha: how much of current policy to mix in
    ref_model_sync_steps=512,   # T_sync: sync every 512 steps
    beta=0.1,
)

trainer = DPOTrainer(
    model=model,
    ref_model=ref_model,
    args=config,
    train_dataset=dataset,
)

Important

When to Use TR-DPO

• Long training runs where the policy drifts far from the initial reference.

• When you observe the DPO loss plateauing early due to KL penalty domination.

• Iterative DPO pipelines where new preference data is collected from the current policy.

• Set \(\alpha\) close to 1 for fast reference updates; close to 0 for slow updates.

EXO – Exact Optimisation

Tip

DPO’s KL Direction Problem

DPO is derived by solving for the optimal policy under a reverse KL constraint. However, the resulting loss actually optimises a forward KL objective in the reward space, which is the wrong direction. EXO (Ji et al. 2024) corrects this by using reverse KL probability matching, which is the theoretically correct objective for alignment.

EXO minimises the reverse KL between the model distribution and the target (reward-optimal) distribution:

\[ \mathcal{L}_{\text{EXO}} = \mathbb{E}_{y \sim \pi_\theta}!\left[ \log \frac{\pi_\theta(y|q)}{p^*(y|q)} \right], \]

where \(p^*(y\vert q) \propto \pi_{\text{ref}}(y\vert q) \exp(r(y,q)/\beta)\) is the optimal policy. In practice, EXO approximates this using the available preference pairs:

\[ \mathcal{L}_{\text{EXO}} \approx -\mathbb{E}!\left[ \log \sigma!\left( \beta \log \frac{\pi_{\text{ref}}(y_w|q)}{\pi_\theta(y_w|q)} - \beta \log \frac{\pi_{\text{ref}}(y_l|q)}{\pi_\theta(y_l|q)} \right) \right]. \]

Note the swapped roles of \(\pi_\theta\) and \(\pi_{\text{ref}}\) compared to DPO.

Note

EXO in TRL

from trl import DPOConfig, DPOTrainer

config = DPOConfig(
    loss_type="exo_pair",
    beta=0.1,
)

trainer = DPOTrainer(
    model=model,
    ref_model=ref_model,
    args=config,
    train_dataset=dataset,
)

NCA – Noise Contrastive Alignment

Tip

Likelihood Collapse in DPO

A known failure mode of DPO is likelihood collapse: the model learns to decrease the probability of the losing response but also decreases the probability of the winning response (since the loss only cares about the difference). NCA (H. Chen et al. 2024) adds an absolute likelihood term to prevent this.

NCA reframes alignment as noise-contrastive estimation. The loss has three terms:

\[ \boxed{ \mathcal{L}_{\text{NCA}} = -\log \sigma(r_w) - \tfrac{1}{2}\log \sigma(-r_w) - \tfrac{1}{2}\log \sigma(-r_l), } \]

where \(r_y = \beta \log(\pi_\theta(y\vert q)/\pi_{\text{ref}}(y\vert q))\) is the implicit reward. The first term encourages high reward for \(y_w\); the second and third terms penalise high reward for both \(y_w\) and \(y_l\) (preventing collapse).

Note

NCA in TRL

from trl import DPOConfig, DPOTrainer

config = DPOConfig(
    loss_type="nca_pair",
    beta=0.01,   # small beta: absolute likelihood term dominates
)

trainer = DPOTrainer(
    model=model,
    ref_model=ref_model,
    args=config,
    train_dataset=dataset,
)

Important

When to Use NCA

• When you observe the winning response probability decreasing during DPO training.

• For tasks where absolute response quality matters, not just relative ranking.

• Use small \(\beta\) (e.g., 0.01) to give the absolute likelihood term more weight.

SLiC-HF – Sequence Likelihood Calibration

Tip

Hinge Loss as a Simpler Alternative

The log-sigmoid loss in DPO is smooth but can be slow to converge when the margin is large. SLiC-HF (Yao Zhao et al. 2023) uses a hinge loss, which is zero when the margin exceeds a threshold and linear otherwise. This is simpler, faster, and surprisingly competitive.

The SLiC-HF loss is:

\[ \mathcal{L}_{\text{SLiC}} = \max!\left(0,; \delta - \beta\log\frac{\pi_\theta(y_w|q)}{\pi_{\text{ref}}(y_w|q)} + \beta\log\frac{\pi_\theta(y_l|q)}{\pi_{\text{ref}}(y_l|q)} \right), \]

where \(\delta\) is the margin threshold. When the model already assigns a margin of \(\delta\) between winning and losing responses, the loss is zero.

Note

SLiC-HF in TRL

from trl import DPOConfig, DPOTrainer

config = DPOConfig(
    loss_type="hinge",
    beta=0.1,
)

trainer = DPOTrainer(
    model=model,
    ref_model=ref_model,
    args=config,
    train_dataset=dataset,
)

Iterative RPO – Reasoning Preference Optimisation

Tip

DPO Forgets How to Generate

Standard DPO trains the model to discriminate between winning and losing responses. But for reasoning tasks, the model also needs to generate correct reasoning traces. A model that can discriminate but not generate is useless at inference time. RPO adds an NLL (negative log-likelihood) term on the winning response to ensure the model learns to generate it.

The RPO loss combines DPO and SFT:

\[ \mathcal{L}_{\text{RPO}} = \lambda_1 \mathcal{L}_{\text{DPO}}(y_w, y_l) + \lambda_2 \mathcal{L}_{\text{NLL}}(y_w), \]

where \(\mathcal{L}_{\text{NLL}}(y_w) = -\log \pi_\theta(y_w\vert q)\) is the standard language modelling loss on the winning response.

Note

Iterative RPO in TRL

from trl import DPOConfig, DPOTrainer

config = DPOConfig(
    loss_type="sigmoid",            # Standard DPO loss
    rpo_alpha=1.0,                  # NLL regularisation weight (RPO)
    beta=0.1,
)

trainer = DPOTrainer(
    model=model,
    ref_model=ref_model,
    args=config,
    train_dataset=dataset,
)

Important

When to Use RPO

• Reasoning tasks (math, code, logic) where the model must generate step-by-step solutions.

• When DPO training causes the model to lose fluency or generation quality.

• Iterative pipelines: generate rollouts, label them, train with RPO, repeat.

• The NLL term acts as a regulariser, preventing the policy from collapsing.

SimPO – Simple Preference Optimisation

Tip

Reference-Free Preference Learning

DPO requires a reference model to compute the implicit reward. This doubles memory usage and adds complexity. SimPO (Meng et al. 2024) eliminates the reference model by using the average log-probability of the response as an implicit reward, with a length normalisation term to prevent the model from preferring short responses.

SimPO defines the implicit reward as:

\[ r_{\text{SimPO}}(y|q) = \frac{\beta}{|y|} \log \pi_\theta(y|q), \]

and the loss as:

\[ \boxed{ \mathcal{L}_{\text{SimPO}} = -\mathbb{E}!\left[ \log \sigma!\left( \frac{\beta}{|y_w|}\log\pi_\theta(y_w|q) - \frac{\beta}{|y_l|}\log\pi_\theta(y_l|q) - \gamma \right) \right], } \]

where \(\gamma > 0\) is a target reward margin that ensures the winning response has strictly higher reward than the losing response by at least \(\gamma\).

Important

SimPO vs DPO vs ORPO

• DPO: uses reference model; ratio-based implicit reward.

• ORPO: reference-free; adds odds-ratio term to SFT loss.

• SimPO: reference-free; length-normalised log-prob reward + margin.

• SimPO is simpler than DPO (no reference model) and more principled than ORPO.

• The length normalisation in SimPO is critical: without it, the model prefers long responses.

Note

SimPO in TRL

from trl import DPOConfig, DPOTrainer

config = DPOConfig(
    loss_type="simpo",
    simpo_gamma=0.5,   # target reward margin gamma
    beta=2.5,          # length normalisation coefficient
    # No ref_model needed!
)

trainer = DPOTrainer(
    model=model,
    ref_model=None,    # SimPO is reference-free
    args=config,
    train_dataset=dataset,
)