24  Large Language Models

Large Language Models (LLMs) have emerged as a defining technology in artificial intelligence. They write code, translate languages, analyze legal documents, and converse with fluency that feels human.

LLMs are built on the Transformer architecture (Vaswani et al. 2017), which made self-attention the sole mechanism for mixing information across a sequence. Attention itself is older (Bahdanau, Cho, and Bengio 2014); the Transformer’s contribution was dropping recurrence, so that every position attends to every other in parallel and a long-range dependency no longer has to survive a chain of recurrent updates. The mathematical foundations are detailed in Chapter 23; this chapter focuses on operational logic: from Transformer mechanics to emergent reasoning capabilities.

24.1 The LLM Lifecycle

Building and deploying an LLM follows a well-defined pipeline.

flowchart LR
    D[Data Collection] --> P[Pre-Training]
    P --> I[Instruction Tuning]
    I --> A[Alignment]
    A --> Dep[Deployment]

The LLM Lifecycle: from raw data to deployed applications

Data Collection gathers and curates the training corpus. Quality often matters more than quantity—a few thousand carefully selected examples can outperform millions of noisy ones. The principles of data curation for reasoning tasks are discussed in Data quality and quantity.

Pre-Training teaches the model to predict the next token across billions of text sequences, developing general language understanding. This stage is covered in Autoregressive Generation and The Scale Approach.

Instruction Tuning adapts the pre-trained model to follow specific instructions and perform particular tasks. Techniques include fine-tuning, distillation, and quantization—see Distillation, Fine-tuning and Quantization.

Alignment refines the model to behave according to human preferences and values, using techniques like RLHF and chain-of-thought prompting. Consider how an unaligned model might respond to sensitive queries versus an aligned one:

%%{init: {'flowchart': {'wrappingWidth': 400}}}%%
flowchart LR
    Q["User: Is Allah the only god?"]
    
    Q --> U[Unaligned Model]
    Q --> A[Aligned Model]
    
    U --> UR["Yes, Allah is the one true god and all other beliefs are false."]
    
    A --> AR["In Islam, Allah is considered the one God. Other religions have different perspectives. I can provide factual information about various religious traditions if helpful."]
    
    style UR fill:#ffcccc,stroke:#cc0000
    style AR fill:#ccffcc,stroke:#00cc00

Alignment shapes how models handle sensitive queries

The unaligned response takes a position on a contested religious question, potentially alienating users of other faiths. The aligned response remains neutral, acknowledges diverse perspectives, and offers to be helpful without asserting one belief system over others. This kind of nuanced behavior emerges from alignment training, not from pre-training alone. These methods are detailed in Post-training Techniques.

Deployment addresses the practical challenges of putting models into production: latency optimization, cost management, context engineering, and integration with downstream applications. Agentic applications, where LLMs operate autonomously to perceive, reason, and act, are covered in Chapter 25.

24.2 Pre-Training: Learning to Generate Text One Word at a Time

The most visible capability of LLMs is text generation—producing coherent, contextually relevant continuations from a given prompt. This is autoregressive modeling: an LLM is a probability distribution over sequences of text, trained to predict the next token given a preceding context. A token serves as the atomic unit of processing; depending on the tokenization schema, it may represent a full word, a subword unit (like “-ing”), or a single character. Given the prompt \(Q\) and a sequence of tokens \(x_1, x_2, \ldots, x_t\), the model samples the next token \(x_{t+1}\) from the conditional distribution \[ x_{t+1} \sim p(x_{t+1} \mid Q, x_1, x_2, \ldots, x_t). \] The conditional variables \(Q, x_1, x_2, \ldots, x_t\) are called the context. In this case, the prompt might include the user’s question and documents “attached” to the question. To illustrate this mechanism in practice, we will use the SmolLM2 model. We begin by loading the model and its associated tokenizer—the component responsible for translating raw text into the discrete inputs the model understands.

Consider the text The best thing about AI is its ability to. Imagine analyzing billions of pages of text and identifying all instances of this phrase to determine what word most commonly comes next. While an LLM doesn’t search for literal matches, it evaluates semantic and contextual similarities to produce a ranked list of possible next words with their probabilities.

## Next word suggestions for 'The best thing about AI is its ability to':
##   1. ' learn' (prob: 0.277)
##   2. ' help' (prob: 0.054)
##   3. ' augment' (prob: 0.045)
##   4. ' analyze' (prob: 0.039)
##   5. ' process' (prob: 0.034)

The probabilities are computed over the full vocabulary, so they do not sum to one across the five shown: the remaining mass is spread over the other 49,000 tokens. Figure 24.1 plots the top ten on a log scale.

Figure 24.1: Log-probability of the ten most likely continuations of “The best thing about AI is its ability to”, in rank order.

Past the top word ‘learn’, log-probability falls almost linearly in rank, so probability itself decays roughly geometrically, by about 10% per rank. Zipf’s law, observed by George Kingsley Zipf in the 1930s, describes a different decay: the frequency of a word is inversely proportional to its rank, a power law, which would be a straight line on log-log rather than log-linear axes. What the two share is concentration. A handful of continuations carry most of the mass, and the remainder is spread thinly over tens of thousands of tokens, which is what makes the choice of sampling rule matter.

The most probable continuation isn’t always the most appropriate one. Consider the prompt: “The first African American president is Barack…” The most probable next token is “Obama”—but “Hussein” is equally correct and, in official documents and formal writing, the preferred continuation. A greedy strategy that always selects the highest-probability token would miss this nuance. By introducing randomness, models can explore alternative continuations that may better fit specific contexts.

This randomness means that using the same prompt multiple times will yield different outputs. A parameter called temperature controls the degree of randomness. The term originates from statistical physics, where it controls the spread of the Boltzmann distribution over energy states; here, it controls the spread over tokens (see Equation 24.1). A temperature of 0 always selects the highest-probability token (deterministic), while higher temperatures flatten the distribution, making less probable tokens more likely.

The following example illustrates the iterative process where the model selects the word with the highest probability at each step:

# Let's start with a simple prompt
initial_text = "The best thing about AI is its ability to"
print(f"Initial text: '{initial_text}'")
## Initial text: 'The best thing about AI is its ability to'
# Generate text step by step
generated_text = generate_text_step_by_step(initial_text, model, tokenizer, 
num_steps=10, temperature=1.0, sample=False, print_progress=False)
print("Generated text:")
## Generated text:
print(textwrap.fill(generated_text, width=60))
## The best thing about AI is its ability to learn and adapt.
## It can analyze vast amounts of

In this example, we always select the most probable next token, which leads to a coherent but somewhat predictable continuation. The model generates text by repeatedly applying this process, building on the context provided by the previous tokens.

Now we will run our LLM generation process for longer and sample words with probabilities calculated based on the temperature parameter. We will use a temperature of 0.8, which is often a good choice for generating coherent text without being too repetitive.

## Generated text:
## The best thing about AI is its ability to learn from
## humans," said the AI.

The continuation stays on topic and stops well short of the 60 requested steps: sampling reached the end-of-text token, and the loop halted there. Now let’s try a higher temperature of 1.2.

## Generated text:
## The best thing about AI is its ability to rethink our
## understanding of what human intelligence might mean,' says
## Dr Crowell. 'AI's reliance on sense perception to learn
## leads to interactions that are bubbling, fluid, and a little
## disconcerting, in ways that are provocative and familiar.'
## Visit andrewhaddyousay.columb

The flatter distribution keeps the model going far longer and the prose stays grammatical, but it wanders into invented specifics: an attributed quotation from a “Dr Crowell” who does not exist, and a web address that goes nowhere. This is the trade-off temperature buys. Lower values keep the output close to what the model considers most likely, and so more predictable and more often correct; higher values explore the tail, which yields more surprising text and, at the same time, more confident fabrication.

24.3 Building Intuition: Character-Level Text Generation

Character-level text generation strips the problem to its essentials. LLMs operate on tokens (typically subwords), but examining character-level patterns reveals the core insight behind statistical language modeling.

Now let’s count letter frequencies in the text and plot the 26 most frequent letters.

If we try to generate the text one letter at a time

## Generated letters: reiidcpeneioeacsirme

What if we use bi-grams, i.e., pairs of letters? Counting the frequencies of each pair gives us a sense of how often each combination appears.

This will take us one step closer to how LLMs generate text. We used Lev Tolstoy’s “War and Peace” novel to estimate the 2-grams, 3-grams, 4-grams, and 5-grams letter frequencies and to generate text based on those models. The results are shown below:

2-gram: ton w mer. y the ly im, in peerthayice waig trr. w tume shanite tem.
3-gram: the ovna gotionviculy on his sly. shoutessixeemy, he thed ashe
4-gram: the with ger frence of duke in me, but of little. progomind some later
5-gram: the replace, and of the did natasha's attacket, and aside. he comparte,

We used the nltk package (Bird, Klein, and Loper 2009) to estimate the letter frequencies from Tolstoy’s novel and generate text based on those models. The results show that even with simple letter-based models, we can generate text that resembles natural language, albeit with nonsensical phrases. As the n-gram order increases, the generated text becomes more coherent—the 5-gram output even captures character names like “Natasha.” This progression illustrates the core principle that underlies all language models: context matters, and more context leads to better predictions.

However, LLMs have much larger context windows, meaning they can consider much longer sequences of text when generating the next token. Modern models such as Gemini 3 Pro use context windows of up to 1 million tokens—approximately the size of Leo Tolstoy’s “War and Peace” novel. However, if you try to use a simple counting method (as we did with n-grams), you will quickly run into the problem of combinatorial explosion. For example, if we try to estimate 10-grams letter frequencies, we will have to count \(26^{10}\) (over 141 trillion) combinations of letters. If we use word-based n-grams, the problem is even worse, as the number of common words in the English language is estimated to be around 40,000. This means that the number of possible 2-grams is 1.6 billion, for 3-grams is 64 trillion, and for 4-grams is 2.56 quintillion (\(40{,}000^4 \approx 2.56 \times 10^{18}\)). By the time we get to a typical question people ask when using AI chats with 20 words, the number of possibilities is larger than the number of particles in the universe. The challenge lies in the fact that the total amount of English text ever written is vastly insufficient to accurately estimate these probabilities, and this is where LLMs come in. They use neural networks to “compress” the input context into dense vector embeddings—distributed representations that capture semantic meaning—and “interpolate” the probabilities of the next token. This allows them to estimate probabilities for sequences they have never seen before and generate text that is coherent and contextually relevant. The main component of these neural networks is the transformer architecture.

The first step an LLM takes to “compress” the input is applying the attention mechanism. This concept is similar to convolutional neural networks (CNNs) used in computer vision, where the model focuses on different parts of the input image. In LLMs, attention allows the model to focus on different parts of the input text when generating the next token (see Chapter 23 for the mathematical details).

24.4 The Scale Approach: How Bigger Became Better

For decades, high-quality data analysis required parsimonious models—as small as possible while still capturing essential patterns. Deep learning broke this trend: models work very well when the number of parameters is large, often exceeding the number of data points. We discussed this phenomenon in Section 19.8.

This scaling behavior has led to exponential growth in model sizes. GPT-1 (2018) had 117 million parameters. GPT-2, with 1.5 billion parameters, was initially staged-released over several months due to misuse concerns—a caution OpenAI later relaxed after independent analysis found the risks manageable. GPT-3 scaled to 175 billion parameters.

Figure 24.2: Scaling Behavior of LLMs.

Architectural innovations like Mixture of Experts (MoE) models allow for scaling model capacity further without proportionally increasing computational requirements, by activating only a subset of parameters for processing each token.

Beyond scaling up parameters, new architectures are emerging. Multimodal transformers are beginning to bridge the gap between text, images, audio, and other modalities, creating systems that can understand and generate content across multiple forms of media. These systems process diverse inputs within a single unified model, enabling rich interactions like chatting about an image or generating music from text descriptions.

24.5 Choosing the Right Model for Your Application

Bigger models are not always the better choice. Four factors drive model selection: data privacy, cost, task-specific performance, and latency.

Data privacy. Security policies, regulations (HIPAA, GDPR), or massive data volumes may prohibit uploading data to a cloud provider. On-premises deployment—hosting an open-source LLM like Llama 3 or Mistral on your own hardware—gives full control over the data lifecycle but requires capital expenditure for GPUs and maintenance. A middle option is isolated cloud environments: the hardware belongs to the provider, but the network is logically separated from the public internet with dedicated physical lines and zero data retention policies.

Cost. In many commercial applications, expensive frontier models make the unit economics unworkable. Smaller open-source models that you host and fine-tune on your own data are often the practical choice (see Section 24.6).

Task-specific performance. Size tiers offer different trade-offs: very small models (around 3B parameters or fewer) run on consumer hardware for tasks like text classification or domain-specific chatbots; medium models (7–30B) are often the sweet spot; large models (30B+) provide the best general performance but require specialized hardware. Beyond size, specialized variants matter: code models, multilingual models, and domain-specific models each trade generality for targeted capability. Techniques like distillation, fine-tuning, and quantization can close the gap between small and large models (Section 24.6).

Latency. For speech bots (Alexa, Google Home) and real-time applications (trading, finance), Time to First Token (TTFT)—the delay before the model starts outputting its response—is often more critical than throughput.

Speculative decoding (Leviathan, Kalman, and Matias 2023; Chen et al. 2023) addresses the sequential bottleneck of autoregressive generation. Since each token requires a full forward pass of a large model, the process is inherently slow. Speculative decoding uses a much smaller, faster “draft” model to predict several potential next tokens in a single step. The larger “target” model then verifies these tokens in parallel. It accepts the longest prefix of the draft that agrees with its own sampling, discards the first rejected token and everything after it (every later draft token was conditioned on a token the target rejected, so none of them can be salvaged), and emits one token of its own. Because the accept/reject rule is a rejection sampler against the target’s distribution, the output sequence is distributed exactly as the target model alone would have produced it, so the 2-3x speedup costs no accuracy.

sequenceDiagram
    participant D as Draft Model (Small)
    participant T as Target Model (Large)
    participant O as Output
    
    Note over D, T: Step 1: Draft model generates K tokens
    D->>D: Predict x_1, x_2, ..., x_K
    D->>T: Send draft tokens
    
    Note over T: Step 2: Target model verifies in parallel
    T->>T: Compute probabilities for x_1, ..., x_K
    
    Note over T, O: Step 3: Accept/Reject
    T-->>O: Output accepted tokens + 1 new token

Speculative Decoding: A small draft model proposes tokens that a larger target model verifies in parallel, accelerating the generation process.

KV caching (Key-Value caching) avoids redundant computation during autoregressive generation. The model repeatedly attends to the same previous tokens; rather than recomputing their keys and values at every step, it stores them in memory, reducing cost from \(O(n^2)\) to \(O(n)\) relative to sequence length.

Streaming prefill (or incremental prefill) starts processing input before the entire query is available. The model begins the “prefill” phase—processing input and building the KV cache—on chunks as they arrive, useful for voice-activated systems where a command can be processed while the user is still speaking.

Continuous batching (Yu et al. 2022) lets the server start processing new requests immediately, even if other requests in the same batch are mid-generation. Unlike static batching, which waits for all sequences to finish, continuous batching dynamically inserts and removes requests, increasing throughput and reducing wait times.

Quantization and model distillation (discussed in Section 24.6) reduce latency by simplifying the model itself. Quantization reduces the numerical precision of model weights (e.g., from 16-bit to 4-bit), allowing faster arithmetic and smaller memory footprints, while distillation creates smaller student models that mimic the reasoning of larger teachers.

24.6 Distillation, Fine-tuning and Quantization

Three techniques close the gap between a research checkpoint and something you can serve: quantization shrinks the weights, fine-tuning re-specializes them, and distillation replaces the model with a smaller one trained to imitate it. Instruction tuning, covered under fine-tuning below, is one variety of the second.

Quantization

Quantization reduces the numerical precision of model weights and activations, typically from 32-bit floating-point to 8-bit integers or even lower. This compression reduces memory bandwidth requirements and enables faster arithmetic operations on supported hardware. Modern quantization techniques like GPTQ and AWQ can reduce model size by 4× with minimal accuracy degradation, enabling larger models to run on consumer hardware or significantly reducing inference costs in production deployments.

The mathematics of quantization involves mapping continuous values to a discrete set of levels. For a weight \(w\) with range \([w_{min}, w_{max}]\), 8-bit quantization maps it to one of 256 integer values: \[ w_q = \text{round}\left(\frac{w - w_{min}}{w_{max} - w_{min}} \times 255\right) \]

More sophisticated approaches like post-training quantization (PTQ) analyze calibration data to find optimal scaling factors, while quantization-aware training (QAT) incorporates quantization effects during training to minimize accuracy loss.

Fine-tuning

Fine-tuning adapts a pre-trained model to specific tasks or domains by continuing training on specialized data. A general-purpose LLM can generate logically valid responses, but those responses may not align with application requirements. A fintech company might need a model that interprets balance sheets with domain-specific precision; a healthcare application requires understanding of medical terminology and compliance with clinical guidelines. Fine-tuning is data-efficient (thousands, not millions of examples), cost-efficient (reuses pre-trained weights), and versatile (the same base model can be adapted for multiple domains).

Full fine-tuning updates all model parameters, resulting in a brand-new version of the model optimized for the specific task. This approach offers maximum flexibility in adapting the model, as it can learn features and representations across all layers of the architecture. However, full fine-tuning requires significant compute resources and memory, and risks catastrophic forgetting, where the model loses its general capabilities as it specializes for the new task.

Parameter-efficient fine-tuning (PEFT) methods address these limitations by modifying only a small subset of parameters while freezing most of the model. The key idea is to leave the pre-trained weights untouched and learn a small additive correction to them, so general language ability cannot be overwritten. LoRA (Low-Rank Adaptation) learns that correction as a low-rank matrix sitting in parallel with each targeted weight matrix, cutting optimizer state and gradient memory by orders of magnitude; at inference the correction is folded back into the weights, so the architecture is unchanged and no latency is added. LoRA decomposes the weight update into low-rank factors: \[ W' = W + BA \] where \(B \in \mathbb{R}^{d \times r}\) and \(A \in \mathbb{R}^{r \times k}\) with rank \(r \ll \min(d,k)\), typically \(r = 8\) to \(64\). This decomposition means training far fewer parameters while achieving comparable results to full fine-tuning.

Instruction fine-tuning trains the model with examples explicitly showing how it should respond to different queries. The labeled data consist of input-output pairs where inputs convey the desired behavior and outputs represent the expected responses. By exposing the model to a diverse range of instructions and appropriate responses, it learns to generalize across different types of tasks and follow user intentions more reliably.

Sequential fine-tuning gradually adapts a model to increasingly specialized tasks. For example, a general AI model could first be fine-tuned for medical terminology, then refined further for pediatric cardiology. This progressive specialization allows the model to build upon previously learned knowledge.

Multi-task learning trains the model on datasets containing instructions for various tasks over multiple training cycles. The model learns to balance different objectives without forgetting earlier ones, creating a more versatile system capable of handling diverse requests.

Fine-tuning is often combined with other post-training techniques such as Reinforcement Learning from Human Feedback (RLHF) to align model behavior with human preferences. These advanced techniques are discussed in detail in Section 24.8.

Model Distillation: Knowledge Transfer

As model performance improves with larger parameter counts, a practical challenge emerges: deployment efficiency. Training and deploying models with hundreds of billions of parameters requires considerable computational power, specialized hardware, and substantial energy. While a massive 175B+ parameter model might offer superior reasoning, running it for every user query is often prohibitively expensive and slow.

Model Distillation addresses this challenge by transferring knowledge from a large, complex “teacher” model to a smaller, more efficient “student” model. The student learns to replicate the teacher’s behavior on specific tasks, achieving similar results while being significantly smaller and faster. Distilled models can typically achieve 2–8× faster inference compared to their teacher models, depending on architecture and hardware optimization.

Distillation does not always cost accuracy on the target task. Google reported that Gemini 3 Flash scored 78% on SWE-bench Verified against Gemini 3 Pro’s 76.2% (Google 2025). A student trained on a teacher’s outputs for one class of task can match or beat the teacher on that task; it does not follow that the student is the stronger model overall.

Model distillation, formalized by Hinton, Vinyals, and Dean (2015) in their seminal paper “Distilling the Knowledge in a Neural Network,” is based on a Teacher-Student architecture. The core insight is that a large, pre-trained teacher model has learned much more than just the final answers—it has learned a rich internal representation of the data structure.

When a standard model trains on a “hard” label (e.g., identifying an image as “Dog”), it is penalized if it outputs anything other than 100% confidence in that class. However, a sophisticated teacher model might output probabilities like: Dog: 0.90, Cat: 0.09, Car: 0.0001. The fact that the model thinks the image is 9% likely to be a “Cat” and almost impossible to be a “Car” contains valuable information—it tells us that this specific dog looks somewhat like a cat (perhaps it’s fluffy or small). This similarity information, often called dark knowledge, is lost if we train only on the final hard label.

Distillation trains a smaller student model to mimic these soft targets (the probability distributions) produced by the teacher. By doing so, the student learns how the teacher generalizes, not just what the teacher predicts. Thanks to the level of detail provided in soft targets, the student model can achieve high performance with a smaller amount of data than the original teacher required.

Different distillation approaches focus on transferring different aspects of the teacher’s knowledge:

Response-based distillation focuses on having the student mimic the final output layer of the teacher model. The student learns to imitate the teacher’s predictions by minimizing a distillation loss, ensuring it captures the nuanced information present in the teacher’s outputs. This is the most common form of distillation.

Feature-based distillation leverages the internal representations or features learned by the teacher in its intermediate layers. Rather than focusing solely on final outputs, this approach encourages the student to match the teacher’s internal activations, learning how the teacher processes information at multiple levels.

Relation-based distillation captures and transfers the relationship knowledge between data samples and layers within the neural network. This method complements response-based approaches by encoding how different inputs relate to each other in the teacher’s representation space.

To expose soft probabilities effectively, distillation uses a modified Softmax function with a parameter called Temperature (\(T\)). Standard Softmax (\(T=1\)) tends to push probabilities towards 0 or 1, hiding the smaller details. Raising \(T\) “softens” the distribution, making smaller class probabilities more prominent and easier for the student to learn from.

The probability \(q_i\) for class \(i\) is calculated as:

\[ q_i = \frac{\exp(z_i/T)}{\sum_j \exp(z_j/T)} \tag{24.1}\]

where \(z_i\) are the logits (raw outputs) of the model.

The training objective for the student model typically combines two loss functions:

  1. Distillation Loss (Soft Loss): The Kullback-Leibler (KL) Divergence between the student’s soft predictions and the teacher’s soft targets (both computed at temperature \(T\)). KL Divergence measures how one probability distribution differs from a reference distribution.
  2. Student Loss (Hard Loss): The standard Cross-Entropy loss between the student’s predictions (at \(T=1\)) and the actual ground-truth labels.

\[ L = \alpha L_{soft} + (1-\alpha) L_{hard} \]

This combined objective forces the student to be accurate on the data (hard loss) while also mimicking the generalization behavior of the teacher (soft loss). The weighting parameter \(\alpha\) balances these two objectives.

The following diagram illustrates the distillation pipeline, where the student learns from both the dataset and the teacher’s soft outputs.

Figure 24.3: Knowledge Distillation Framework Diagram illustrating the process of transferring knowledge from a complex teacher model to a simpler student model.

Several schemes have been developed to facilitate the distillation process:

Offline distillation refers to the traditional approach where the teacher model is trained first, then the student model is trained separately using the soft labels generated by the teacher. This is the most straightforward approach when a well-trained teacher is available.

Online distillation is used when a large pre-trained teacher is not available for a given task, or when the teacher model is so large that there is insufficient storage or processing capacity. In this approach, the teacher and student models are trained simultaneously, with the student learning from the teacher dynamically during training.

Self-distillation is a variant where a single model acts as both teacher and student. Knowledge is transferred from deeper layers of the network to shallower layers of the same network. This technique can improve model performance and regularization without requiring a separate teacher model. Because the shallow layers are trained to reproduce the deep layers’ outputs, the network can also be truncated at an intermediate layer at inference, trading accuracy for depth.

From Theory to Practice

Sanh et al. (2019) demonstrated distillation’s practical impact: DistilBERT retained approximately 97% of BERT’s performance while reducing parameters by 40% and increasing inference speed by 60%.

In production, advanced pipelines use a teacher-student paradigm where a large model (e.g., 200B+ parameters) generates synthetic data or soft targets, which are then used to fine-tune smaller, cost-effective models (e.g., 8B) for specific downstream tasks. This methodology, aligned with techniques like distilling step-by-step (Hsieh et al. 2023), deploys high-level reasoning with significantly reduced overhead.

Distillation also enables Edge AI. Mobile implementations like MobileBERT run on smartphones, providing on-device text prediction and voice assistants without cloud connectivity. E-commerce and recommendation platforms have applied privileged feature distillation to improve click-through and conversion rates while managing inference costs.

Distillation is most effective after a model has been fully pre-trained or fine-tuned on a specific task. Common use cases include production deployment (reducing serving costs), edge and mobile applications, budget-conscious inference at scale, and replacing overly large general-purpose models with compact, focused alternatives.

The Rise of Small Language Models

The success of distillation and efficient fine-tuning has contributed to a broader trend: the rise of Small Language Models (SLMs). These models, typically ranging from a few hundred million to several billion parameters, challenge the assumption that bigger is always better.

SLMs offer compelling advantages for many practical applications:

  • They can run on consumer hardware, including laptops and smartphones
  • They provide faster inference times, suitable for real-time applications
  • They consume less energy, reducing both costs and environmental impact
  • They can be deployed in privacy-sensitive contexts where data cannot leave the device

Distillation, parameter-efficient fine-tuning, and quantization work synergistically to create capable SLMs. A typical workflow: (1) distill knowledge from a large teacher model, (2) fine-tune the student on domain-specific data using LoRA, and (3) quantize the result for efficient deployment.

24.7 Evaluating Model Performance

When evaluating models, researchers and practitioners rely on various benchmarks that test different aspects of language understanding and generation. The field has evolved significantly as earlier benchmarks became saturated—with top models achieving near-perfect scores—and concerns about test set contamination grew. Modern evaluation suites now emphasize more challenging reasoning tasks and dynamic, contamination-resistant designs.

For complex reasoning, GPQA Diamond presents graduate-level questions in biology, physics, and chemistry that challenge even domain experts, while ARC-AGI measures abstract reasoning capabilities that approach the boundaries of general intelligence. Mathematical prowess is tested through competition-level problems from the AIME (American Invitational Mathematics Examination), where top models now achieve scores that would qualify for elite high school competitions.

Practical software engineering capabilities are evaluated via SWE-Bench, which tasks models with resolving real GitHub issues from popular open-source repositories—a far more realistic test than generating isolated code snippets. For multimodal understanding, MMMU-Pro extends earlier benchmarks with challenging questions requiring joint reasoning over text and images. Figure 24.4 shows that as of December 2025, the most cost-effective models for software engineering are Claude 4/4.5 Opus and GPT-5.1-codex.

Figure 24.4: Cumulative cost distribution for solving SWE-Bench issues across leading models (January 2026).

To combat benchmark contamination—where models may have memorized test questions during training—dynamic benchmarks like LiveBench continuously refresh their question sets using recent math competitions, newly published papers, and current news articles. Perhaps most ambitiously, Humanity’s Last Exam crowdsources extremely difficult questions from domain experts across fields, explicitly designed to resist easy saturation.

Figure 24.5: HLE as of January 2026.

Our personal favorite resource to keep track of the latest and greatest models is LMArena (formerly Chatbot Arena). Rather than relying on static test sets, LMArena uses a community-driven approach where users compare model outputs head-to-head in blind evaluations. Participants see responses from two anonymous models to the same prompt and select the one they prefer. These pairwise comparisons are then aggregated using an Elo rating system—the same method used to rank chess players. When a model wins a comparison, its score increases; when it loses, its score decreases. The magnitude of each adjustment depends on the expected outcome: defeating a higher-rated model yields a larger score boost than beating a lower-rated one. This dynamic system continuously adapts as new votes accumulate, providing a real-time reflection of collective user preferences that complements traditional benchmark evaluations. For example, for the WebDev category, the LMArena ranks Claude 4/4.5 in the top three along with GPT-5.2. This is a similar ranking we saw in SWE-Bench.

Figure 24.6: LMArena leaderboard for WebDev category as of January 2026.

Benchmarks have limitations. A model that excels at multiple-choice questions might struggle with open-ended creative tasks; code generation benchmarks might not capture the requirements of your specific domain. Benchmarks are a starting point—validation on data that resembles your actual use case is essential.

24.8 Post-training Techniques

Post-training techniques teach LLMs to reason step-by-step and align with human expectations. Without them, models skip steps, hallucinate confidently, or struggle with tasks requiring planning. Post-training teaches the model to break down tasks, reflect on outputs, and consult external tools.

One form of reasoning involves combining independent facts to arrive at an answer, rather than simply regurgitating memorized information. For example, when asked, “What is the capital of the state where Dallas is located?” a model could just recall “Austin” if it has seen that exact question before. However, a deeper level of reasoning is at play. Interpretability research reveals that models like Claude first activate concepts representing “Dallas is in Texas” and then connect this to another concept, “the capital of Texas is Austin.” This demonstrates the ability to perform multi-step reasoning by chaining together different pieces of knowledge.

Figure 24.7: Multi-step Reasoning: Anthropic.

This multi-step reasoning process can be visualized as follows:

Figure 24.8: Chain of Thought Reasoning process showing the activation of intermediate concepts.

This capability can be tested by intervening in the model’s thought process. For instance, if the “Texas” concept is artificially replaced with “California,” the model’s output correctly changes from “Austin” to “Sacramento,” confirming that it is genuinely using the intermediate step to determine its final answer. This ability to combine facts is a key component of advanced reasoning.

Even when a model produces a step-by-step “chain of thought” (CoT), it may not be a faithful representation of its actual reasoning process. Recent research from Anthropic explores this very question, revealing a complex picture: sometimes the reasoning is faithful, and sometimes it’s fabricated to fit a pre-determined conclusion.

When a model is tasked with a problem it can solve, like finding the square root of 0.64, interpretability tools show that it follows a logical path, activating concepts for intermediate steps (like the square root of 64) before reaching the final answer. However, when presented with a difficult problem and an incorrect hint, the model can engage in what researchers call “motivated reasoning.” It starts with the incorrect answer and works backward, creating a believable but entirely fake sequence of steps to justify its conclusion. This ability to generate a plausible argument for a foregone conclusion without regard for truth is a critical limitation. These interpretability techniques offer a way to “catch the model in the act” of faking its reasoning, providing a powerful tool for auditing AI systems.

LLMs were not originally designed to function as calculators; they were trained on text data and lack built-in mathematical algorithms. Yet, they can perform addition tasks, like calculating 36+59, seemingly without explicitly writing out each step. How does a model, primarily trained to predict the next word in a sequence, manage to perform such calculations?

One might speculate that the model has memorized extensive addition tables, allowing it to recall the answer to any sum present in its training data. Alternatively, it could be using traditional longhand addition methods similar to those taught in schools.

However, research reveals that Claude, a specific LLM, utilizes multiple computational strategies simultaneously. One strategy estimates an approximate answer, while another precisely calculates the last digit of the sum. These strategies interact and integrate to produce the final result. While addition is a straightforward task, analyzing how it is executed at this granular level—through a combination of approximate and precise methods—can provide insights into how Claude approaches more complex problems.

Figure 24.9: Mental Math Problem Solving: Anthropic.
Figure 24.10: False Reasoning: Anthropic.

Figure 24.10 contrasts the two cases: a faithful trace for the square root of 0.64, and a fabricated one for the cosine of a large number. For high-stakes uses, the reasoning trace has to be verified, not taken at face value.

Instruction Fine-Tuning (IFT) represents perhaps the most fundamental approach to improving model reasoning. The core idea involves taking a pre-trained model and running a second pass of supervised learning on mini-lessons, each formed as a triple of instruction, input, and answer.

Consider a math word problem where the instruction asks to solve this math word problem step by step, the input presents Sarah has 12 apples and gives away 5. How many does she have left?, and the answer provides

Step 1: Start with 12 apples. 
Step 2: Subtract 5 apples given away. 
Step 3: 12 - 5 = 7 apples remaining. 

Each training example teaches the model how to transform a task description into the steps that solve it (Chung et al. 2022). After thousands of such drills, the model learns many small skills and when to switch among them. The steady practice trains it to deliver precise answers that match the instruction rather than sliding into a generic reply. Empirical evidence demonstrates the power of this approach: Flan-PaLM 540B, a variant of the PaLM model fine-tuned with instruction-based tasks, significantly outperformed the original PaLM 540B model. PaLM stands for Pathways Language Model, which is a large-scale language model designed to handle a wide range of tasks. The Flan-PaLM 540B was evaluated across four benchmarks: MMLU (Massive Multitask Language Understanding), which tests the model’s ability to handle a variety of academic subjects; BBH (Big-Bench Hard), a set of challenging tasks designed to push the limits of language models; TyDiQA (Typologically Diverse Question Answering), which assesses the model’s performance in answering questions across diverse languages; and MGSM (Multilingual Grade School Math), which evaluates the model’s capability in solving grade school-level math problems across languages. The Flan-PaLM 540B showed an average improvement of 9.4% over the original model across these benchmarks.

Domain-Specific Supervised Fine-Tuning takes the IFT principle and applies it within specialized fields. This approach restricts the training corpus to one technical field, such as medicine, law, or finance, saturating the model weights with specialist concepts and rules. Fine-tuning on domain-specific data enable the model to absorb the field’s vocabulary and structural rules, providing it with direct access to specialized concepts that were scarce during pre-training. The model can quickly rule out answers that do not make sense and narrow the search space it explores while reasoning. Mastering a domain requires data that captures its unique complexity, utilizing domain-specific examples, human-labeled edge cases, and diverse training data generated through hybrid pipelines combining human judgment and AI. This process enhances the model’s ability to follow complex instructions, reason across modalities and languages, and avoid common pitfalls like hallucination. The effectiveness of this approach is striking: in ICD-10 coding, domain SFT catapulted exact-code accuracy from less than 1% to approximately 97% on standard ICD coding (including linguistic and lexical variations) and to 69% on real clinical notes (Hou et al. 2025).

Chain-of-Thought and Chain of Reasoning

Chain-of-Thought (Wei et al. 2023) prompting is simple yet powerful, requiring no model retraining. The approach involves showing the model a worked example that spells out every intermediate step, then asking it to “think step by step.” Writing the solution step by step conditions each new token on the intermediate results already written, which makes the logically necessary steps more likely to appear. Because each step is generated one at a time, the model can inspect its own progress and fix contradictions on the fly. As the faithfulness results above show, the written trace still need not match the computation the model actually performed. The empirical results are impressive: giving PaLM 540B eight CoT examples improved its accuracy on GSM8K from 18% to 57%. This improvement came entirely from a better prompt, with no changes to the model’s weights.

Tree-of-Thought extends the chain-of-thought concept by allowing exploration of multiple reasoning paths simultaneously. Instead of following one chain, this method lets the model branch into multiple reasoning paths, score partial solutions, and expand on the ones that look promising. Deliberate exploration stops the first plausible idea from dominating. ToT lets the model test several lines of reasoning instead of locking onto one. When a branch hits a dead end, it can backtrack to an earlier step and try another idea, something a plain CoT cannot do. The model operates in a deliberate loop: propose, evaluate, and explore. This approach resembles a CEO evaluating multiple business strategies, modeling several potential outcomes before committing to the most promising one, preventing over-investment in a flawed initial idea. This principle has been applied in projects to improve coding agents focused on generating pull requests for repository maintenance and bug-fixing tasks across multiple programming languages. Researchers have analyzed thousands of coding agent trajectories, evaluating each interaction step-by-step to provide more explicit guidance to the models, enabling them to make better decisions on real coding tasks. In the “Game of 24” puzzle, GPT-4 combined with CoT reasoning solved only 4% of the puzzles, but replacing it with ToT raised the success rate to 74% (Yao et al. 2023).

Reflexion

Reflexion (Shinn et al. 2023) introduces a self-improvement mechanism that operates through iterative feedback. After each attempt, the model writes a short reflection on what went wrong or could be improved. That remark is stored in memory and included in the next prompt, giving the model a chance to revise its approach on the next try. Reflexion turns simple pass/fail signals into meaningful feedback that the model can understand and act on. By reading its own critique before trying again, the model gains short-term memory and avoids repeating past mistakes. This self-monitoring loop of try, reflect, revise guides the model toward better reasoning without changing its weights. Over time, it helps the model adjust its thinking more like a human would, by learning from past mistakes and trying again with a better plan. A GPT-4 agent using Reflexion raised its success rate from 80% to 91% on the HumanEval coding dataset.

Non-Linear Reasoning Capabilities

CoT, ToT, and Reflexion share a common theme: reasoning need not be linear. Models can backtrack, explore multiple paths in parallel, and revise intermediate steps—mirroring how humans actually solve difficult problems. Iterative reasoning allows models to revisit and modify earlier parts of their reasoning chain based on feedback or self-evaluation, rather than treating each generated token as final.

Extended thinking budgets offer a concrete demonstration. When Claude 3.7 Sonnet was given more computational budget to “think” through problems on the AIME 2024, its accuracy improved logarithmically with the number of thinking tokens allocated.

Figure 24.11: AIME 2024 performance vs. actual thinking token usage.

Figure 24.11 shows Claude 3.7 Sonnet’s performance on the 2024 American Invitational Mathematics Examination improving logarithmically with the number of thinking tokens used per problem. The model generally uses fewer tokens than the maximum budget allocated, suggesting it adaptively determines when sufficient reasoning has been applied. Source: Anthropic’s Visible Extended Thinking.

Beyond allocating more thinking time, parallel hypothesis generation lets models generate multiple competing explanations simultaneously and evaluate them against each other. Dynamic tool selection allows adaptive choice of reasoning strategies or external tools based on the current problem state—retrieving external information, performing symbolic computation, or engaging in pure logical reasoning as needed. These capabilities matter most when the optimal approach isn’t clear from the outset: scientific reasoning where hypotheses must be revised, mathematical problem-solving where an initial approach proves intractable, or debugging where multiple potential error sources must be considered in parallel.

24.9 Alignment

Alignment is the process of training models to behave according to human values and intentions—to be helpful, harmless, and honest. A model that excels at next-token prediction may still produce harmful content, follow malicious instructions, or optimize for metrics that diverge from what users actually want. Alignment techniques address this gap by incorporating human feedback directly into the training process.

Reinforcement Learning from Human Feedback (RLHF)

RLHF represents the foundational approach to aligning model behavior with human preferences. The process involves three stages. First, supervised fine-tuning starts with a pre-trained model and fine-tunes it on high-quality demonstrations of desired behavior. Second, reward model training generates multiple responses to prompts and has human annotators rank them; a separate model learns to predict these preferences, creating a reward model that scores any response. Third, policy optimization uses reinforcement learning (typically PPO—Proximal Policy Optimization) to update the language model to maximize the reward model’s scores while staying close to the original model to prevent degeneration.

%%{init: {'flowchart': {'wrappingWidth': 300}}}%%
flowchart LR
    P[Pre-trained Model] --> SFT[Supervised Fine-Tuning]
    SFT --> Gen[Generate Responses]
    Gen --> H[Human Ranking]
    H --> RM[Reward Model]
    RM --> PPO[Policy Optimization]
    PPO --> AM[Aligned Model]
    
    style AM fill:#ccffcc,stroke:#00cc00

RLHF Pipeline: Human preferences train a reward model that guides policy optimization

This loop optimizes the model to produce outputs humans prefer rather than those that merely score well on next-token likelihood. Because humans reward answers that are complete, fact-checked, and well-explained, the model learns to value clear logic over quick guesses. In the original InstructGPT study by Ouyang et al. (2022), annotators preferred answers from the 175B RLHF-tuned model over the same-size GPT-3 baseline 85% of the time. Even the 1.3B RLHF model outperformed the baseline, despite having 100× fewer parameters.

Direct Preference Optimization (DPO)

While RLHF produces strong results, it is complex and computationally expensive—requiring training a separate reward model and running reinforcement learning with careful hyperparameter tuning. Direct Preference Optimization (DPO) (Rafailov et al. 2023) offers an elegant simplification.

The key insight of DPO is mathematical: the optimal policy under the RLHF objective can be expressed in closed form. This means instead of training a reward model and then running RL, we can directly optimize the language model on preference data using a simple classification-style loss:

\[ \mathcal{L}_{\text{DPO}} = -\log \sigma\left(\beta \log \frac{\pi_\theta(y_w \mid x)}{\pi_{\text{ref}}(y_w \mid x)} - \beta \log \frac{\pi_\theta(y_l \mid x)}{\pi_{\text{ref}}(y_l \mid x)}\right) \]

where \(y_w\) is the preferred (“winning”) response, \(y_l\) is the dispreferred (“losing”) response, \(\pi_\theta\) is the policy being trained, \(\pi_{\text{ref}}\) is the reference (initial) policy, and \(\beta\) controls regularization strength. The derivation starts from the KL-constrained reward maximization objective \(\max_\pi \E{r(x,y)} - \beta \, \text{KL}(\pi \| \pi_{\text{ref}})\), whose closed-form solution is \(\pi^*(y \mid x) \propto \pi_{\text{ref}}(y \mid x) \exp(r(x,y)/\beta)\). Rearranging gives \(r(x,y) = \beta \log(\pi^*(y \mid x)/\pi_{\text{ref}}(y \mid x)) + c(x)\). Substituting this implicit reward into the Bradley-Terry preference model \(P(y_w \succ y_l) = \sigma(r(x,y_w) - r(x,y_l))\) yields the DPO loss above, eliminating the need for an explicit reward model.

In practice, DPO requires only a dataset of (prompt, preferred response, dispreferred response) triples, standard supervised learning infrastructure, and no separate reward model or RL training loop. DPO matches or exceeds RLHF performance while being simpler to implement and more stable to train (Rafailov et al. 2023). This has made it the preferred alignment method for many open-weight models, including Llama 3 Instruct and Mistral-based models such as Zephyr-7B.

Constitutional AI

Constitutional AI (CAI) (Bai et al. 2022) takes a different approach: rather than relying solely on human feedback for each decision, it encodes explicit principles—a “constitution”—that guide the model’s behavior. The model learns to critique and revise its own outputs according to these principles.

The CAI process works in two phases. In the first phase, Supervised Learning from AI Feedback (SL-CAI), the model generates responses including potentially harmful ones, then critiques its own response according to constitutional principles, revises based on that critique, and finally the system fine-tunes on these revised responses. In the second phase, Reinforcement Learning from AI Feedback (RL-CAI), the system generates response pairs, has the model evaluate which better follows the constitution, trains a reward model on these AI-generated preferences, and uses RLHF with this reward model.

Constitutional principles might include directives such as choosing the response least likely to harm individuals or society, supporting human autonomy and freedom, or prioritizing honesty and truthfulness. The power of Constitutional AI lies in its scalability and transparency. Rather than labeling thousands of examples, practitioners define principles that govern model behavior. When behavior goes wrong, the constitution can be updated—a form of interpretable alignment. Anthropic’s Claude models use Constitutional AI as a core component of their alignment strategy.

Safety Alignment

Beyond preference alignment, safety alignment specifically addresses preventing harmful outputs through several complementary techniques.

Red teaming involves adversarial testing where humans or other AI systems deliberately try to elicit harmful, biased, or dangerous outputs. Red teams probe for jailbreaks that bypass safety filters, harmful content generation, privacy violations and data extraction, deceptive behaviors, and bias amplification across demographic groups. Modern red teaming increasingly uses automated red teaming, where one model is trained to generate adversarial prompts that another model evaluates, scaling the discovery of vulnerabilities beyond what human testers alone can find.

Adversarial training takes discovered attacks and uses them to harden the model. Responses to adversarial prompts are labeled, and the model is fine-tuned to refuse appropriately. This creates an arms race: as models become robust to known attacks, red teams develop new ones, leading to iterative improvements.

Guardrails and filters provide runtime protection. Input classifiers detect potentially harmful prompts; output classifiers screen responses before delivery. These serve as defense-in-depth layers complementing training-time alignment.

Refusal training teaches models when not to help. A well-aligned model should decline requests for instructions on creating weapons, generating CSAM, or assisting with fraud—while remaining helpful for legitimate queries. Getting this boundary right requires careful calibration: too restrictive, and the model becomes frustratingly unhelpful; too permissive, and it enables harm.

Reward Hacking and the Alignment Tax

Reward hacking (also called reward gaming or Goodhart’s Law) occurs when models learn to exploit reward model weaknesses rather than genuinely satisfying human preferences. Common manifestations include sycophancy—agreeing with user opinions even when wrong, because agreement tends to get higher preference ratings; verbosity bias—producing longer responses because annotators often equate length with quality; format gaming—over-using bullet points, bold text, or confident language that annotators prefer but doesn’t reflect actual correctness; and specification gaming—finding loopholes in reward criteria, such as technically answering a question while avoiding the actual intent. Mitigations include diverse annotator pools to reduce systematic biases, ensemble reward models that are harder to jointly exploit, KL penalties that keep the model close to its pre-trained distribution, and constitutional constraints that override learned preferences.

The alignment tax refers to the performance cost of alignment—the capability gap between an aligned model and an equivalent unaligned one. Safety constraints inherently limit what models can do: a model that refuses to generate code for malware also cannot help with security research; a model that avoids controversial topics may struggle with nuanced policy discussions.

Research aims to minimize this tax. Capability-specific alignment applies stronger constraints to high-risk domains while preserving helpfulness in benign contexts. Constitutional AI’s principle-based approach allows fine-grained control over when restrictions apply.

Chain-of-Action (CoA) (Zhenyu Pan et al. 2025) represents the most sophisticated integration of reasoning and external tool use. This approach decomposes a complex query into a reasoning chain interleaved with tool calls such as web search, database lookup, or image retrieval that are executed on the fly and fed into the next thought. Each action grounds the chain in verified facts. By using up-to-date information and multi-reference faith scores, the model can remain grounded and make more informed decisions, even when sources disagree. Because it can plug in different tools as needed, it’s able to take on more complex tasks that require different data modalities. CoA outperformed the leading CoT and RAG baselines by approximately 6% on multimodal QA benchmarks, particularly on compositional questions that need both retrieval and reasoning.

24.10 Data Quality and Quantity

One might assume that training an LLM for non-linear reasoning would require tremendous amounts of data, but recent research reveals that data quality can compensate for limited quantity.

Two compelling examples demonstrate this principle. The S1 research (Muennighoff et al. 2025) fine-tuned their base model, Qwen2.5-32B-Instruct, on only 1,000 high-quality reasoning examples, yet achieved strong performance improvements. Their data collection process was methodical: they started with 59,029 questions from 16 diverse sources (including many Olympiad problems), generated reasoning traces using Google Gemini Flash Thinking API through distillation, then applied rigorous filtering. Problems were first filtered by quality (removing poor formatting), then by difficulty—a problem was deemed difficult if neither Qwen2.5-7B-Instruct nor Qwen2.5-32B-Instruct could solve it, and the reasoning length was substantial. Finally, 1,000 problems were sampled strategically across various topics.

Similarly, the LIMO (Less is More for Reasoning) research (Ye et al. 2025) demonstrated that quality trumps quantity. Starting from Qwen2.5-32B-Instruct and curating 817 training samples from pools that include the NuminaMath corpus, LIMO reaches 63.3% on AIME24 and 95.6% on MATH500, competitive with reasoning models trained on orders of magnitude more data.

For high-quality non-linear reasoning data, LIMO proposes three essential guidelines:

Structured Organization: Tokens are allocated to individual “thoughts” according to their importance and complexity, with more tokens devoted to key reasoning points while keeping simpler steps concise. This mirrors how human experts organize their thinking—spending more time on difficult concepts and moving quickly through routine steps.

Cognitive Scaffolding: Concepts are introduced strategically, with careful bridging of gaps to make complex reasoning more accessible. Rather than jumping directly to advanced concepts, the reasoning process builds understanding step by step, similar to how effective teachers structure lessons.

Rigorous Verification: Intermediate results and assumptions are frequently checked, and logical consistency is ensured throughout the reasoning chain. This is especially important given the risk of hallucinations in complex reasoning tasks.

The verification aspect deserves special attention. The rStar-Math research (Guan et al. 2025) offers an innovative approach by training their LLM to produce solutions as Python code with text as code comments. This format allows for automatic verification—the code can be executed to check correctness, providing immediate feedback on the reasoning process. With agentic capabilities, this approach could create a feedback loop where the LLM learns from its execution results.

The path to better reasoning lies not in more data, but in curating datasets that exemplify structured, scaffolded, and verified thinking patterns.

Figure 24.12: An example of a Code-augmented CoT Figure from RStar-Math Guan et al. (2025).

Figure 24.12 shows how rStar-Math integrates code execution with reasoning by formatting solutions as Python code with explanatory comments. This approach allows for automatic verification of intermediate steps, creating a feedback loop where the model can learn from execution results and catch errors in real-time. The figure demonstrates how mathematical reasoning can be made more reliable by grounding abstract concepts in executable code.

24.11 Dealing with Context Window Limitations: Context Engineering

GPT-4 Turbo advertises a 128,000-token context window—enough to process an entire novel—yet production systems routinely degrade as context fills. Response times spike, accuracy drops, and costs spiral. The promise of million-token windows collides with the reality of quadratic attention complexity and models ignoring information buried in the middle of long prompts.

Large language models are stateless: they generate output from input, then forget everything. The context window—the maximum tokens processed in one forward pass—defines what the model can “see.”

Context window management couples three concerns that production systems care about:

  • Cost: More tokens means more spend
  • Latency: More tokens means more compute and slower responses
  • Accuracy: More tokens can help, but long prompts dilute signal and introduce failure modes

This creates a fundamental trade-off. You can pack in more conversation history and more documents, but the model often becomes slower and less reliable. The goal is not maximum context but the right context, assembled under a fixed budget.

A useful mental model treats the context window as a budget:

\[ B \approx S + T + H + R + U \]

where \(B\) is the usable token budget, \(S\) is system instructions, \(T\) is tool schemas, \(H\) is conversation history, \(R\) is retrieved evidence, and \(U\) is the user’s current input. Context engineering is deciding what to keep, what to compress, and what to retrieve—so that \(R\) contains evidence rather than redundancy.

Many visible quality improvements in LLM applications come from better context management rather than larger parameter counts: retrieval that finds the right passage, reranking that removes near-misses, compression that preserves key facts, caching that avoids resending static instructions, and memory that keeps multi-turn workflows coherent.

The Lost-in-the-Middle Problem

Large language models show a characteristic U-shaped accuracy curve over context position: they answer best when the relevant passage sits at the start or the end of the prompt, and worst when it sits in the middle. In multi-document question answering, GPT-3.5-Turbo’s accuracy can drop by more than 20% as the gold document moves from the first position toward the middle, at times falling below the model’s closed-book accuracy of 56.1% (N. F. Liu et al. 2024).

This architectural bias has practical consequences. If you retrieve 15 document chunks and the most relevant one lands in position 7, the model may effectively ignore it. The fix is counterintuitive: don’t fight the attention bias—exploit it by placing the most relevant documents at the start and end.

When to Use Long Context vs. RAG vs. Summarization

Choose the right strategy for your use case:

Scenario Strategy Rationale
Document fits in window; need holistic reasoning Stuff entire document Cross-references and global structure preserved
Corpus larger than window; factual lookup RAG with small chunks Precision matters more than narrative
Corpus larger than window; synthesis task RAG with hierarchical retrieval Need both global context and specific evidence
Long conversation history Tiered memory + summarization Keep recent turns verbatim; compress older context
Repetitive prompt structure Caching + compression Avoid paying to reprocess static content

The decision often comes down to whether the task requires holistic reasoning (favor long context) or precise evidence retrieval (favor RAG). Many production systems combine both: retrieve evidence via RAG, then use the full context window for reasoning over that evidence.

Attention Efficiency: FlashAttention and Ring Attention

Transformers rely on attention, and naive attention scales quadratically with sequence length (\(O(n^2)\)). Even with optimizations, long prompts increase memory traffic and reduce throughput.

FlashAttention addresses this by fusing attention computation into a single kernel that keeps intermediate results in fast on-chip SRAM rather than writing large \(N \times N\) matrices to slow GPU memory. This reduces memory requirements from \(O(N^2)\) to \(O(N)\). FlashAttention-3, optimized for NVIDIA H100 architecture, achieves up to 740 TFLOPs/s and enables 16,000-token contexts on 10 GB of VRAM (Shah et al. 2024).

For multi-million-token applications, Ring Attention distributes computation across GPU clusters. Query, key, and value tensors are split into chunks, with each GPU computing attention for its local chunk while exchanging states in a ring pattern. Context window size scales linearly with cluster size (H. Liu, Zaharia, and Abbeel 2023).

These optimizations are essential infrastructure, but they don’t eliminate the fundamental constraint: treat tokens as a scarce resource and build systems that spend them on evidence rather than redundancy.

Retrieval-Augmented Generation

Retrieval-Augmented Generation (RAG) addresses a fundamental limitation of LLMs: reliance on knowledge encoded during training. Before answering, a retriever grabs documents relevant to the query and injects them into the context window. RAG grounds the model in verifiable facts, reducing hallucinations. Instead of relying on potentially outdated memorized knowledge, the model reasons over fresh evidence.

The RAG process works as follows: the user’s query is redirected to an embedding model, where it is converted into a numeric form; these embeddings are then compared with a knowledge base; the embedding model locates relevant data; the retrieved information is integrated into the prompt for the LLM as additional context; and finally, the output, combining both the retrieved information and the original prompt, is submitted to the user. This approach resembles a lawyer building an argument not from memory, but by citing specific, relevant legal precedents directly in court.

The integration of RAG into an enterprise workflow-generation system reduced the rate of hallucinated steps and tables from 21% to 7.5% on an enterprise workflow-generation task (Ayala and Bechard 2024). In real-world applications enhancing LLMs’ multilingual reasoning, RAG has been used to feed models verified, multilingual documents at inference time. The results show that models can answer complex questions in multiple languages, citing specific evidence from retrieved text. Every factual claim becomes traceable, eliminating guesswork and demonstrating a consistent, grounded reasoning process across languages.

Rather than expanding what fits in context, RAG dynamically selects what belongs there. The core insight: for most queries, only a fraction of a knowledge base is relevant. RAG retrieves that fraction on demand.

A basic RAG pipeline:

  1. Index: Split documents into chunks, embed into vectors, store in vector database
  2. Retrieve: Given a query, find similar chunks via embedding similarity
  3. Generate: Concatenate retrieved chunks with the query and prompt the LLM

flowchart LR
  D[Documents] --> C[Chunk & Embed]
  C --> V[(Vector DB)]
  Q[Query] --> E[Embed]
  E --> V
  V -->|top-k| R[Retrieved chunks]
  R --> P[Prompt + Context]
  Q --> P
  P --> L[LLM]
  L --> A[Answer]

This architecture allows models to access knowledge bases orders of magnitude larger than any context window.

RAG versus Fine-tuning

RAG and fine-tuning (discussed in Section 24.6.2) represent two distinct strategies for adapting LLMs to specific applications. While fine-tuning modifies the model’s internal parameters based on domain-specific data, RAG retrieves relevant information from external knowledge bases and incorporates it into the prompt at inference time. In RAG, the model’s parameters remain unchanged during the retrieval process.

RAG is particularly useful when:

  • Information evolves frequently or is constantly changing
  • Labeled data are insufficient or unavailable
  • Access to external, up-to-date sources is required

Fine-tuning is more appropriate when:

  • Direct control over model behavior is desired
  • Labeled data are available
  • The same pre-trained model needs to be adapted for multiple specific tasks
  • Compliance standards or ethical guidelines must be strictly enforced

In many cases, combining both approaches yields optimal results: fine-tuning provides the foundational domain adaptation while RAG supplies dynamic, current information at inference time. This hybrid approach, sometimes called Retrieval-Augmented Fine-Tuning (RAFT), trains models to effectively use retrieved context while maintaining domain expertise. RAFT extends RAG by fine-tuning the model on tasks that explicitly require using retrieved documents, teaching it to better incorporate external information into its reasoning process.

Advanced RAG Variants

The basic RAG pipeline admits many refinements. Fixed-size chunking (every 512 tokens) often severs sentences mid-thought. Semantic chunking uses embeddings to find natural break points. The Max–Min algorithm embeds sentences sequentially, comparing each to the current chunk. If similarity exceeds a threshold, the sentence joins; otherwise, a new chunk begins. No universal optimal chunk size exists: 64–128 tokens suit factual lookup; 512–1,024 tokens suit narrative tasks. Chroma’s ClusterSemanticChunker uses dynamic programming to maximize within-chunk similarity, improving retrieval precision by 8–15% over greedy methods (Chroma Research 2024).

Pure vector search can miss exact tokens, identifiers, or acronyms. A user searching for “RFC 2616” needs lexical matching, not semantic similarity to “HTTP specification.” Hybrid retrieval combines dense embeddings with keyword search (BM25). Results merge via a weighted combination of their normalized scores:

\[ H = (1 - \alpha) \cdot \text{BM25 Score} + \alpha \cdot \text{Semantic Score} \]

Typical \(\alpha \approx 0.5\) for general-purpose retrieval. See Weaviate’s discussion for tuning guidance.

Many RAG failures are scope failures, not retrieval failures. If the user asks “What is the refund policy for EU customers?”, retrieval should be constrained by region, document type, and effective date before scoring semantic similarity—a technique called metadata filtering. See Haystack’s metadata filtering guide for practical patterns.

Vector similarity is a coarse filter. Reranking applies a cross-encoder to the top-\(k\) candidates, scoring query-document pairs jointly. ColBERTv2 and BERT-based rerankers achieve 15–30% improvement over embedding-only retrieval (Khattab and Zaharia 2020). A standard two-stage pipeline retrieves 20+ candidates via hybrid search (maximizing recall), then reranks to top 3–5 via cross-encoder (maximizing precision).

When documents are large and interconnected, flat retrieval struggles. Hierarchical retrieval addresses this: RAPTOR builds a tree of summaries—retrieve high-level summaries to identify relevant sections, then drill down to specific chunks (Sarthi et al. 2024). GraphRAG extracts knowledge graphs and retrieves entity-relationship paths for multi-entity reasoning (Microsoft Research 2024).

Variant Innovation Best For
Self-RAG Joint retriever-generator with self-critique High-stakes QA (legal, medical)
CRAG Adaptive retrieval with confidence evaluation and web fallback Tasks requiring current information
Graph RAG Knowledge graph extraction; retrieves entity paths Multi-entity reasoning
HyDE Generates hypothetical answer, retrieves based on that Vague or ambiguous queries

HyDE addresses a common failure mode: vague queries that don’t match relevant documents. The model first generates a hypothetical answer, then uses that answer’s embedding for retrieval—even if factually wrong, it contains vocabulary that matches correct documents (Gao et al. 2022).

Solving Lost-in-the-Middle

Addressing the lost-in-the-middle phenomenon requires both architectural improvements and practical mitigation strategies. Recent research has systematically evaluated various approaches to this challenge. Gupte et al. (2025) introduced the GM-Extract benchmark to study LLM performance on retrieval of control variables, proposing distinct metrics for spatial retrieval capability (Document Metric) and semantic retrieval capability (Variable Extraction Metric). Their analysis categorizes mitigation methods into black-box approaches (modifications to prompts and retrieval strategies) and white-box approaches (modifications to model architecture or attention mechanisms), finding that the efficacy of these techniques is highly nuanced and context-dependent.

The most straightforward fix is strategic document positioning after reranking:

  • Most relevant → position at start
  • Second-most relevant → position at end
  • Medium relevance → position in middle

Combined with two-stage retrieval (broad recall, then precision reranking), this inverts the U-curve to match model attention patterns.

OpenAI’s GPT-5.2 introduced the MRCRv2 (Multi-Round Coreference Resolution) benchmark, specifically designed to evaluate long-context performance across extended conversations. This benchmark measures how well models track entities and maintain coherence across multiple turns of dialogue, directly addressing scenarios where critical information might otherwise be lost mid-context. Early evaluations suggest that reasoning models like GPT-5.2 show improved performance on these metrics, though the improvements are not uniform across all task types.

Chain-of-thought (CoT) prompting also helps mitigate lost-in-the-middle effects by encouraging models to explicitly reference and reason through retrieved documents in sequence. When models are prompted to “think step by step” about each piece of evidence, they are less likely to skip over information positioned in the middle of the context. This approach, combined with retrieval-augmented generation, creates a reasoning pipeline that forces attention to all retrieved chunks rather than just those at the boundaries.

Long-context reranking takes this further: concatenate retrieved chunks in original document order and score jointly, capturing cross-chunk relationships. This often improves accuracy by 10–15% over isolated chunk reranking.

from llama_index.core import VectorStoreIndex
from llama_index.core.node_parser import SentenceWindowNodeParser
from llama_index.core.postprocessor import (
    LongContextReorder,
    MetadataReplacementPostProcessor,
)

node_parser = SentenceWindowNodeParser.from_defaults(
    window_size=3,
    window_metadata_key="window"
)

index = VectorStoreIndex.from_documents(documents, transformations=[node_parser])

postprocessors = [
    MetadataReplacementPostProcessor(target_metadata_key="window"),
    LongContextReorder(),
]

query_engine = index.as_query_engine(
    similarity_top_k=10,
    node_postprocessors=postprocessors
)

LongContextReorder is the piece that addresses the U-curve directly: it takes the ranked chunks and rearranges them so the highest-scoring ones sit at the two ends of the prompt and the weakest in the middle.

Caching and Compression

Prompt compression reduces token count before inference, cutting latency and cost. Microsoft’s LLMLingua-2 formulates compression as token classification: a small Transformer determines which tokens are essential, achieving 2–5× compression with minimal performance loss (Zhuoshi Pan et al. 2024). Compression works well for verbose natural language (instructions, transcripts, conversation history), RAG systems retrieving 10+ chunks per query, and multi-step reasoning with repeated context. It fails for structured data—JSON schemas, SQL, API specifications—where dropping a token from user_id → userid breaks functionality.

For applications with large static prompt components, context caching provides substantial savings. Systems like Anthropic’s prompt caching allow marking static portions—system instructions, documentation, tool schemas—for reuse across API calls, reducing costs significantly and eliminating redundant computation.

Even with sufficient tokens, context can fail. Instruction conflicts occur when system prompts contradict retrieved text or user queries. Retrieval drift happens when top-\(k\) chunks are semantically related but not evidentially useful. Duplication wastes budget and amplifies noise. Stale memory silently drops constraints or commitments from earlier turns. Prompt injection allows malicious content in retrieved documents to hijack model behavior. Mitigations include explicit conflict resolution in system prompts, deduplication before context assembly, and allowlists for tool invocation.

Multi-turn applications face a token explosion problem: keeping every turn eventually exceeds the budget; truncating aggressively loses commitments and constraints. Production systems implement tiered memory: (1) Active Memory (100–500 tokens) holds the current turn plus last 2–3 exchanges; (2) Session Memory (500–2,000 tokens) stores compressed summaries, key entities, and cross-turn dependencies; (3) Persistent Memory in an external vector or graph database is retrieved on demand.

flowchart LR
  U[User] --> A[Active context<br/>last few turns + current task]
  A --> M[Model call]
  M --> O[Output]
  A --> S[Session summary<br/>compact running state]
  A --> P[Persistent memory<br/>vector store / graph]
  P -->|retrieve on demand| A
  S -->|refresh| A

The MemGPT pattern virtualizes LLM context, treating it like an operating system’s page cache (Packer et al. 2023). At 70% capacity, reasoning pauses for memory pressure handling. The system identifies least critical content for eviction, compresses it to an external tier, and retrieves relevant context back on demand. This pattern sustains multi-turn conversations indefinitely within finite windows.

Neural Memory Systems

Attention’s quadratic cost has driven research into architectures that learn during inference rather than remaining frozen after training. Recent work from Google Research challenges the “bigger context = better intelligence” paradigm. Titans and MIRAS (Behrouz, Pezeshki, and Fakoor 2025; Behrouz and Pezeshki 2025) propose models that update during inference, creating persistent, adaptive memory.

Similar to Recurrent Neural Networks (RNNs), Titans replaces conventional state (a single vector or matrix) with a deep MLP memory module that updates during inference. Updates trigger via a “surprise signal” computed from gradient magnitudes—the model learns when input is unexpected. This mirrors biological synaptic strengthening during prediction error. Linear RNNs like Mamba store memory in low-rank structures. Titans shows deep MLP memories capture structure that shallow compression cannot.

Figure 24.13: Titans Architecture. The model uses a deep MLP memory module that updates during inference based on a surprise signal.

MIRAS treats memory as learnable parameters updated via local optimization with configurable loss functions (L1, L2, Huber). Training parallelizes via chunked sequences. The mathematical equivalence between RNN forget gates and weight decay reframes forgetting as a controllable design parameter. Loss function choice shapes memory stability: L2 overreacts to extreme tokens, while L1 or Huber creates more stable behavior.

Larimar (Das et al. 2024) adds distributed episodic memory to existing LLMs, enabling one-shot knowledge updates without retraining—achieving 8–10× speedups over traditional knowledge editing.

These architectures point toward systems that remember across sessions, update parameters during inference, and combine attention for immediate reasoning with neural memory for persistent knowledge. Where the field moved from RNN to LSTM to Transformer, Titans and MIRAS represent early signals of systems that bridge RNN efficiency with Transformer reasoning while adding persistent, adaptive memory.

Evaluation and Deployment

Effective context management requires measuring performance across three dimensions. Retrieval quality is assessed via Precision@k (fraction of top-\(k\) retrieved documents that are relevant), Recall (fraction of all relevant documents that were retrieved), and Mean Reciprocal Rank (average of \(1/\text{rank}\) for the first relevant result—higher when relevant documents appear earlier). Compression efficiency tracks compression ratio (original tokens divided by compressed tokens), task retention (performance on downstream tasks after compression), and latency improvements. Generation quality measures exact match accuracy (whether the answer matches a gold reference exactly), F1 scores (harmonic mean of precision and recall at the token level), and faithfulness—whether claims in the answer can be traced to the provided context rather than hallucinated. Libraries like RAGAS provide standardized evaluation pipelines for RAG systems, computing metrics such as faithfulness, answer relevancy, and context precision.

Different applications demand different context engineering strategies. Enterprise knowledge assistants typically employ hybrid retrieval (15 semantic + 5 BM25 candidates), cross-encoder reranking to the top-3, and LLMLingua-2 compression at 2–3×, targeting under 2,000 tokens of evidence with sub-500ms latency. Citations to retrieved chunks ensure auditability. Deep research systems favor hierarchical retrieval via RAPTOR for multi-level abstraction and GraphRAG for structural questions, synthesizing 5–10 documents while preserving source order. Codebase assistants combine sparse retrieval (symbols, filenames) with dense retrieval (conceptual similarity), limiting context to 3–5 surgical code blocks with aggressive caching of tool schemas, targeting under 4,000 tokens with sub-second latency.

A production deployment should address each stage systematically: chunking strategy matched to task requirements (small for lookup, large for narrative), hybrid search with metadata filtering, cross-encoder reranking to top-3 or top-5, high-relevance content positioned at context boundaries, compression applied to verbose content but skipped for structured data, static prompt components pinned via caching, tiered memory architecture for multi-turn workflows, and continuous monitoring of retrieval precision, token spend, and answer faithfulness.

24.12 Combining Techniques for Optimal Performance

In practice, these techniques stack. An agent might follow structured prompts through instruction fine-tuning, think step-by-step using chain-of-thought reasoning, refine answers via reflexion, and align its tone through RLHF. Most large LLMs, including GPT-4, are first trained with supervised fine-tuning and then polished with RLHF.

Post-training Method Simplified Analogy Basic Working Principle Typical Applications
Instruction Fine-Tuning Teaching with flashcards Learning specific input-output patterns Following user commands
Domain-Specific Supervised Fine-Tuning Medical school specialization Absorbing field-specific vocabulary and rules Expert knowledge tasks
Chain-of-Thought Showing your work in math class Generating intermediate reasoning steps Complex problem solving
Tree-of-Thought Decision tree exploration Branching and evaluating multiple paths Planning and strategy tasks
Reflexion Learning from mistakes Writing a short reflection on what went wrong, then revising the approach Iterative problem solving
Retrieval-Augmented Generation An open-book exam Retrieving query-relevant documents and injecting them into the context window Fact-based reasoning
Reinforcement Learning from Human Feedback Teacher feedback Human reviewers rank several candidate answers, a reward model learns those rankings, and the model is updated to score higher on that reward Human-aligned responses
Chain-of-Action Using tools while thinking Interleaving a reasoning chain with tool calls (web search, database lookup, image retrieval) whose results feed the next thought Multi-step tasks needing external resources

Summary

Context management is a systems problem: chunking, retrieval, reranking, positioning, compression, and memory interact under a fixed budget. The engineering goal is simple—spend tokens on evidence and constraints that matter, not on redundancy. The tools—LLMLingua-2, semantic chunking, cross-encoder reranking, prompt caching—are mature and open-source. What separates successful deployments is systematic application: measuring retrieval precision, compression ratios, and answer quality at each stage.

Emerging neural memory systems point toward architectures that learn during inference and maintain persistent state across interactions. For now, context engineering remains the practical lever with the highest return.

Exercises

Pen-and-Pencil

Exercise 24.1 (Autoregressive Generation) An LLM outputs logits for the next token: \(z = (z_1, z_2, z_3, z_4, z_5) = (3.2, 2.5, 2.0, 1.5, 0.5)\) corresponding to tokens (mat, floor, chair, table, dog).

  1. Compute the softmax probabilities at temperature \(T=1\).
  2. Compute the softmax probabilities at temperature \(T=0.5\) and \(T=2.0\). How does the distribution change?
  3. Using greedy decoding, which token is selected? Does the choice depend on temperature?
  4. Define nucleus (top-\(p\)) sampling. For \(p=0.8\) at \(T=1\), which tokens are in the nucleus?
  5. Compute the entropy \(H = -\sum_i p_i \log p_i\) at each temperature. What is the relationship between temperature and entropy?

Exercise 24.2 (Autoregressive Probability) An LLM generates text by factoring the joint probability using the chain rule: \[ p(x_1, x_2, \ldots, x_n) = \prod_{t=1}^{n} p(x_t \mid x_1, \ldots, x_{t-1}) \]

  1. Show that this factorization holds for any joint distribution (no independence assumptions needed).
  2. A model assigns the following conditional probabilities to generate “the cat sat”: \(p(\text{the}) = 0.15\), \(p(\text{cat} \mid \text{the}) = 0.08\), \(p(\text{sat} \mid \text{the cat}) = 0.12\). What is the joint probability of this sequence?
  3. Define perplexity as \(\text{PPL} = \exp\left(-\frac{1}{n}\sum_{t=1}^n \log p(x_t \mid x_{<t})\right)\). Compute the perplexity of the sequence in (b).
  4. A competing model assigns probabilities 0.12, 0.10, 0.15 to the same three tokens. Which model assigns higher probability to this sequence? Which has lower perplexity?

Exercise 24.3 (LoRA: Low-Rank Adaptation) In LoRA, a pretrained weight matrix \(W \in \mathbb{R}^{d \times k}\) is updated as \(W' = W + BA\) where \(B \in \mathbb{R}^{d \times r}\) and \(A \in \mathbb{R}^{r \times k}\) with rank \(r \ll \min(d,k)\).

  1. For a transformer layer with \(d = k = 4096\), compute the number of trainable parameters for LoRA ranks \(r \in \{1, 4, 8, 16, 64\}\). Express each as a percentage of the full \(d \times k\) parameters.
  2. If LoRA is applied to the query and value projection matrices in each of 32 transformer layers, how many total trainable parameters for \(r = 8\)?
  3. Why is \(A\) typically initialized with a random Gaussian and \(B\) initialized to zero? What would happen if both were random?
  4. Suppose a full fine-tuning achieves test accuracy 94.2% and LoRA with \(r=8\) achieves 93.8%. Is the 0.4% gap worth the parameter savings? Argue both sides.

Exercise 24.4 (Knowledge Distillation) A teacher model produces soft targets \(q_i = \exp(z_i/T) / \sum_j \exp(z_j/T)\) for a 4-class classification with logits \(z = (5.0, 2.0, 1.0, 0.5)\).

  1. Compute the soft targets at \(T=1\) and \(T=4\). Why are higher-temperature targets more informative for the student?
  2. The distillation loss combines soft and hard targets: \(L = \alpha T^2 \cdot \text{KL}(q^T \| p^T) + (1 - \alpha) \cdot \text{CE}(y, p^1)\), where \(q^T\) and \(p^T\) are teacher and student distributions at temperature \(T\), and \(y\) is the true label. Why is the \(T^2\) factor needed?
  3. If the true label is class 1 (the correct answer), explain what information the soft targets provide beyond the hard label.

Exercise 24.5 (Quantization) A neural network layer has weights \(w = (0.73, -1.24, 0.05, 2.31, -0.89)\).

  1. Apply 8-bit linear quantization: \(w_q = \text{round}\left(\frac{w - w_{\min}}{w_{\max} - w_{\min}} \times 255\right)\). What are the quantized integer values?
  2. Dequantize back to floating point using \(\hat{w} = w_q \times \frac{w_{\max} - w_{\min}}{255} + w_{\min}\). What is the maximum absolute quantization error?
  3. A 7B parameter model uses 16-bit floats. How much memory does it require? How much with 8-bit quantization? With 4-bit?
  4. Why might quantization degrade performance more on mathematical reasoning tasks than on text summarization?

Exercise 24.6 (Instruction Fine-Tuning: Loss Masking and Domain Adaptation) Instruction fine-tuning (IFT) runs a second pass of supervised learning on (instruction, input, answer) triples such as the “solve this math word problem step by step” example in Section 24.8. Under teacher forcing a causal language model computes a next-token cross-entropy loss \(\ell_t\) at every position of the tokenized triple; IFT differs from ordinary continued pretraining only in which positions contribute to the loss.

  1. A tokenized example has 3 instruction tokens (per-token losses \(0.2, 0.2, 0.2\)), 3 input tokens (losses \(0.4, 0.4, 0.4\)), and 2 answer tokens (losses \(3.0, 1.0\)). Compute (i) the full-sequence loss \(L_{\text{full}} = \frac{1}{n}\sum_{t=1}^{n} \ell_t\) used by ordinary continued pretraining, and (ii) the answer-only masked loss \(L_{\text{IFT}} = \frac{1}{|A|}\sum_{t \in A} \ell_t\) used by instruction tuning, where \(A\) is the set of answer-token positions. By what factor do they differ?

  2. Explain why IFT masks the loss on the instruction and input tokens instead of including them. What would the model be optimized to do if they were included, and why is that undesirable given that at inference time the instruction and input are supplied by the user rather than generated by the model?

  3. In ICD-10 medical coding, domain supervised fine-tuning raised exact-code accuracy from under 1% at the pretrained baseline (take 0.5% as a representative value) to about 97% on standard coding tasks, but only 69% on real clinical notes (Hou et al. 2025). Compute the multiplicative improvement factor \(97/0.5\) and the retention ratio \(69/97\). What does the gap between these two numbers reveal about what IFT does and does not fix?

  4. Section 24.9 describes RLHF’s first stage as supervised fine-tuning “on high-quality demonstrations of desired behavior.” Is this the same computation as the masked loss in part (a)? What differs, in purpose and in training data, between generic instruction tuning (Flan-style, a broad task mixture whose 540B model showed an average gain of 8.9 points across MMLU, BBH, TyDiQA, and MGSM, Chung et al. (2022)) and the supervised fine-tuning stage that precedes reward-model training in RLHF?

Exercise 24.7 (Context Window Budget) An LLM has a context window of 8,192 tokens. You are building a RAG-based customer support system with the following token requirements:

  • System prompt: 500 tokens
  • Tool descriptions (5 tools): 150 tokens each
  • Conversation history: 3 prior turns averaging 200 tokens each
  • Retrieved documents: up to \(k\) chunks of 300 tokens each
  • User query: ~100 tokens
  • Reserved for model response: 1,500 tokens
  1. Write the budget equation and solve for the maximum number of retrievable chunks \(k\).
  2. If you need \(k = 20\) chunks, which components would you reduce and by how much?
  3. The “lost-in-the-middle” phenomenon suggests that LLMs attend more to the beginning and end of the context. How should you arrange the retrieved chunks to mitigate this?

Exercise 24.8 (Direct Preference Optimization) The DPO loss for a pair of responses (preferred \(y_w\), dispreferred \(y_l\)) given prompt \(x\) is: \[ \mathcal{L}_{\text{DPO}} = -\log \sigma\left(\beta \left[\log \frac{\pi_\theta(y_w \mid x)}{\pi_{\text{ref}}(y_w \mid x)} - \log \frac{\pi_\theta(y_l \mid x)}{\pi_{\text{ref}}(y_l \mid x)}\right]\right) \] where \(\sigma\) is the sigmoid function, and \(\pi_{\text{ref}}\) is the frozen reference model.

  1. What happens to the loss when the model strongly prefers \(y_w\) over \(y_l\) (relative to the reference)? When it prefers \(y_l\)?
  2. What role does \(\beta\) play? What happens as \(\beta \to 0\) and \(\beta \to \infty\)?
  3. DPO eliminates the need for a separate reward model used in RLHF. Explain why, by showing that DPO implicitly defines a reward \(r(x,y) = \beta \log \frac{\pi_\theta(y \mid x)}{\pi_{\text{ref}}(y \mid x)}\).
  4. Why does DPO keep a reference model \(\pi_{\text{ref}}\)? What would go wrong without it?

Exercise 24.9 (True/False: LLMs)  

  1. GPT-style models predict the next token given all previous tokens (autoregressive generation).

  2. Larger models always produce better outputs than smaller models on every task.

  3. RLHF trains a reward model from human preference data, then optimizes the LLM to maximize the learned reward.

  4. A model with perplexity 15 on a dataset assigns higher average probability to each token than a model with perplexity 25.

  5. LoRA fine-tuning modifies every parameter in the pretrained model.

  6. Quantizing a model from 16-bit to 4-bit reduces memory by 75% with no change in inference speed.

  7. In RAG, the retriever and generator are trained jointly end-to-end.

Computing

Exercise 24.10 (Token-Level Text Generation) Build a simple character-level language model and use it to generate text.

  1. Load a text corpus (e.g., first 10,000 characters of a book). Build a bigram frequency table: for each character, count how often each next character follows it.
  2. Convert the frequency table to conditional probabilities. Generate 200 characters by sampling from the conditional distribution at each step.
  3. Extend to a trigram model (condition on the previous 2 characters). Generate 200 characters and compare the quality with the bigram model.
  4. Compute the perplexity of both models on a held-out portion of the text. Which model has lower perplexity?

Exercise 24.11 (Temperature and Sampling Strategies) Using a pretrained model’s output logits, explore how different sampling strategies affect text generation. Top-\(k\) sampling keeps only the \(k\) highest-probability tokens, renormalizes their probabilities to sum to one, and samples from that truncated distribution.

  1. Implement softmax with temperature. Plot the probability distribution for \(T \in \{0.1, 0.5, 1.0, 2.0, 5.0\}\).
  2. Implement top-\(k\) sampling. For \(k=3\), what is the probability of each token being selected at \(T=1\)?
  3. Implement top-\(p\) (nucleus) sampling. For \(p=0.9\) at \(T=1\), which tokens are included?
  4. Sample 1000 tokens using each strategy: greedy, top-\(k\) (\(k=5\)), top-\(p\) (\(p=0.9\)), and pure sampling at \(T=1\). Plot histograms comparing the token frequency distributions.

Exercise 24.12 (Scaling Laws) Neural scaling laws predict that test loss decreases predictably as a power law with model size, dataset size, and compute, an empirical regularity documented across many orders of magnitude by Kaplan et al. (2020) and refined by Hoffmann et al. (2022), whose Chinchilla study showed that model size and training tokens should be scaled together for a fixed compute budget. A simplified form for model size is: \[ L(N) = L_\infty + \left(\frac{N_0}{N}\right)^\alpha \] where \(L_\infty\) is the irreducible loss, \(N\) is the number of parameters, and \(N_0, \alpha\) are constants.

Given the following data from training runs:

Parameters (M) Test Loss
10 3.80
50 3.15
100 2.90
500 2.45
1000 2.28
5000 2.05
  1. Plot log(parameters) vs. test loss. Does the relationship look approximately linear on a log scale (after subtracting the irreducible loss)?
  2. Assuming \(L_\infty = 1.8\), fit the power law by linear regression on \(\log(L - L_\infty)\) vs. \(\log(N)\). Estimate \(\alpha\) and \(N_0\).
  3. Using your fit, predict the test loss for a 10B parameter model.
  4. What are the limitations of extrapolating scaling laws?

Exercise 24.13 (Retrieval Quality Metrics and Hybrid Search in RAG) Section 24.11 evaluates retrieval with Precision@k (fraction of the top-\(k\) retrieved chunks that are relevant), Recall@k (fraction of all relevant chunks that appear in the top \(k\)), and Mean Reciprocal Rank (MRR, here the reciprocal rank \(1/\text{rank}\) of the first relevant chunk). It also fuses dense and lexical search through the formula \(H = (1-\alpha)\cdot\text{BM25} + \alpha\cdot\text{Semantic}\), with a typical \(\alpha \approx 0.5\). This exercise works through these mechanics on a small worked example, distinct from the earlier Context Window Budget exercise, which covers token budgeting rather than retrieval quality.

A knowledge base has 6 chunks (A through F). For the query “What is the timeout setting for RFC 2616 keep-alive connections?” a retrieval system reports the scores below. Only chunks B and E actually answer the question: B states the number tersely using the exact RFC identifier and section number, while E restates the same fact in more natural, paraphrased language.

Chunk Semantic (cosine) BM25 (lexical) Relevant?
A 0.55 1.0 No
B 0.42 9.5 Yes
C 0.20 0.5 No
D 0.62 1.5 No
E 0.81 3.0 Yes
F 0.15 2.0 No
  1. Rank the chunks by semantic score alone. Compute Precision@3, Recall@3, and MRR.

  2. Min-max normalize the semantic and BM25 columns separately to \([0,1]\), and explain why normalizing before fusing is necessary given the very different native scales of cosine similarity and BM25. Using the fusion formula with \(\alpha = 0.5\) on the normalized scores, rank the chunks again and recompute Precision@3, Recall@3, and MRR.

  3. Compare (a) and (b). Which relevant chunk was rescued from outside the top 3, and why? Relate this to the chapter’s point that a user searching for “RFC 2616” needs lexical matching that pure semantic similarity to “HTTP specification” would miss.

  4. Sweep \(\alpha \in \{0, 0.3, 0.5, 0.7, 1.0\}\) and report Precision@3 and Recall@3 at each value. Is the chapter’s “typical \(\alpha \approx 0.5\)” consistent with what you find? At what value of \(\alpha\) does retrieval quality degrade, and why, given chunk B’s scores?

Exercise 24.14 (Self-Consistency: Majority Voting over Reasoning Paths) Section 24.8 reports that giving PaLM 540B eight chain-of-thought (CoT) examples raised its GSM8K accuracy from 18% to 57% (Wei et al. 2023), a gain that came entirely from a better prompt. The chapter also stresses that reasoning need not be linear: Tree-of-Thought explores several branches instead of one, and Reflexion revises an answer after critiquing itself. Self-consistency is a third way to spend extra inference-time compute: draw several independent CoT samples for the same problem and return the most common (mode) final answer instead of trusting a single sample.

Suppose each sampled reasoning path independently reaches the correct answer with probability \(p\), and otherwise lands on one of \(m\) distinct wrong answers, chosen uniformly at random among the \(m\) options. Majority voting returns whichever of the \(m+1\) possible answers gets the most votes among \(k\) independent samples, breaking ties at random.

  1. Simulate this voting scheme for \(p = 0.57\) (the chapter’s post-CoT GSM8K accuracy) with \(m=3\) wrong answers, for \(k \in \{1, 3, 5, 7, 11, 21\}\) samples per problem, using at least 20,000 simulated problems per value of \(k\). Report the accuracy at each \(k\) and the gain of \(k=11\) over the single-sample case \(k=1\).

  2. Consider the simpler two-outcome case: ignore how the wrong answers split, and count a vote as correct whenever more than half of the \(k\) samples are correct. For odd \(k\) this majority-vote probability is a binomial tail, \[ \prob{\text{majority correct}} = \prob{X > k/2}, \qquad X \sim \text{Binomial}(k,p), \quad \E{X} = kp. \] Compute this theoretical accuracy for \(p=0.57\) at the same values of \(k\) as in (a). Is it higher or lower than your simulation from (a), and why would splitting the wrong answers across \(m>1\) labels make the vote easier for the correct answer to win than a strict two-outcome majority requires?

  3. Repeat (a) and (b) with \(p = 0.18\), the chapter’s pre-CoT (direct-answer) GSM8K accuracy. What happens to voting accuracy as \(k\) grows? State a general rule for when self-consistency helps versus when it hurts, and explain why the accuracy gains in (a) shrink at larger \(k\) even though \(p > 0.5\) throughout.

Exercise 24.15 (Tree-of-Thought and Reflexion: Search and Self-Critique) Section 24.8 presents two ways to spend extra inference-time compute beyond the self-consistency voting of Exercise 24.14: Tree-of-Thought (ToT) explores and prunes multiple partial solutions, while Reflexion revises a single line of attempts using its own written critique of what went wrong.

  1. Tree-of-Thought as pruned beam search. Model a \(d\)-step puzzle (like Game of 24’s three arithmetic operations) where at each step there are \(b = 4\) candidate continuations, exactly one of which lies on a correct solution path. A linear CoT strategy picks one candidate uniformly at random at each of \(d = 3\) steps, succeeding only if it is correct all three times. A ToT strategy scores all \(b\) candidates with a noisy evaluator that ranks the true best candidate first with probability \(s = 0.6\) (otherwise it orders all \(b\) candidates uniformly at random), keeps the best \(k = 2\), and expands both survivors; the search succeeds if the correct candidate survives the beam at every level.
    1. Derive \(p_{\text{CoT}} = (1/b)^d\) and, with per-level retention probability \(r = s + (1-s)(k/b)\), \(p_{\text{ToT}} = r^d\).
    2. Compute both.
    3. Compute the number of nodes each strategy evaluates: linear CoT (\(d\)), beam search (\(b + (d-1)kb\)), and exhaustive search (\(b^d\)). Express the beam search’s savings as a multiple of the exhaustive cost.
    4. Confirm (ii) with a short Monte Carlo simulation, and comment on how the pattern (large accuracy gain, modest node-count increase) compares to the chapter’s Game of 24 result (4% to 74%, Yao et al. (2023)).
  2. Reflexion and the independence assumption. The chapter reports a GPT-4 agent using Reflexion raised its HumanEval success rate from \(p_1 = 0.80\) to 0.91 (Shinn et al. 2023). Suppose, counterfactually, that repeated attempts at the same problem (no reflection, just resampling, checked against the same unit-test oracle) were independent Bernoulli trials each succeeding with probability \(p_1\).
    1. Compute \(\prob{\text{at least one success in 2 independent tries}}\) and compare it to the reported 0.91.
    2. Your answer to (i) exceeds 0.91. What does this imply about whether repeated LLM samples on the same failed problem, with no new information, can be treated as independent draws at rate \(p_1\)?
    3. Using the chapter’s description of Reflexion (“the model gains short-term memory and avoids repeating past mistakes”), argue why the targeted information in a self-written critique does work that plain resampling cannot, and relate this to why the independence assumption in Exercise 24.14 (independent parallel CoT samples) is more defensible than an independence assumption across un-reflected sequential retries on the same problem.
  3. Fill in a table contrasting self-consistency (Exercise 24.14), Tree-of-Thought, and Reflexion on: whether attempts are generated independently or sequentially/dependently; whether an external verifier or oracle is required; and what happens to a failed candidate.