Chapter 4: Improving reasoning with inference-time scaling

Introduction

The chapter introduction establishes that reasoning accuracy can improve during text generation without retraining or modifying the model.

• Inference-time reasoning improvement: Explain that reasoning performance and answer accuracy can improve through methods used while the model generates text, without retraining or modifying the model itself.
  Why it matters: This distinction explains why the chapter can improve results using an existing model.

Key terms: inference-time scaling, inference

4.1 Introduction to inference-time scaling

The section defines inference-time scaling and introduces three ways to spend more computation after training to improve reasoning responses.

• Inference-time scaling and tradeoff: Explain how inference-time scaling differs from increasing training compute: a trained model does more work while generating each answer, without changing its weights. Describe the resulting tradeoff between response quality and inference cost.
  Why it matters: This distinction explains how a model can produce better answers without being retrained.
• Three inference-time method patterns: Describe the three methods introduced here: prompting a longer chain-of-thought response, sampling multiple responses and selecting the most frequent one, and iteratively reviewing and improving an answer. Distinguish generating longer responses from generating multiple responses.
  Why it matters: These patterns show the different ways additional answer-generation work can be used.

Key terms: training compute, inference-time scaling, test-time scaling, inference compute, chain-of-thought response, parallel sampling, self-consistency, iterative self-refinement

4.2 Loading a pretrained model

The section sets up a pretrained model and a flexible generation wrapper for comparing inference-time strategies, then shows that the baseline answer to a MATH-500 prompt is wrong.

• Interchangeable generation wrapper: Explain how generate_text_stream_concat_flex differs from the earlier wrapper and why that change helps compare generation strategies.
  Why it matters: Keeping the surrounding pipeline fixed makes differences between generation strategies easier to interpret.
• Baseline math response: Describe what happened when the pretrained model was tried on the MATH-500 prompt.
  Why it matters: The incorrect baseline answer gives a concrete result against which later approaches can be considered.
• CPU reproducibility recommendation: Explain why the section recommends using the CPU for a first run.
  Why it matters: A first run on the CPU makes it easier to compare your output with the book's examples.

Key terms: generate_text_stream_concat_flex, generate_text_stream_concat, generate_text_basic_stream_cache, MATH-500, pretrained base model, cpu, GPU

4.3 Generating better responses with chain-of-thought prompting

The section shows how a step-by-step prompt can improve a base model’s response, why it may help, and when its costs or limitations matter.

• Chain-of-thought prompt instruction: Describe how chain-of-thought prompting changes the input prompt to ask a model to explain its reasoning before giving a final answer, using an instruction such as “Explain step by step.”
  Why it matters: A small change to the prompt can change the kind of response a base model produces.
• Reasons for improved accuracy: Explain the two reasons given for why writing intermediate steps can improve an answer: opportunities for self-correction and alignment with detailed solutions seen in training.
  Why it matters: These reasons explain why asking for steps may help on multistep tasks without changing the model’s knowledge.
• Reasoning token cost: Explain why a longer chain-of-thought response increases inference time and cost, even when it produces a correct answer.
  Why it matters: The possible accuracy gain must be weighed against the work required to generate extra tokens.
• Limits of step-by-step prompting: Distinguish cases where chain-of-thought prompting may help a base model from cases where it may hurt or be unnecessary.
  Why it matters: A step-by-step prompt should not be treated as an automatic accuracy improvement for every problem or model.

Key terms: chain-of-thought prompting, chain of thought, reasoning chain, intermediate steps, training examples, existing knowledge, inference-time-scaling tradeoff, tokens, forward pass, latency

4.4 Controlling output diversity with temperature scaling

This section contrasts sequential chain-of-thought prompting with independently generated answers and explains why self-consistency needs diverse responses to the same prompt.

• Sequential and parallel generation: Distinguish chain-of-thought prompting, which extends the sequence of next-token predictions, from generating independent answers that can be sampled in parallel when resources permit.
  Why it matters: The distinction explains why generating several answers need not mean a longer wait for the user.
• Need for diverse candidate answers: Explain why the text generation function must produce different answers to the same prompt for self-consistency sampling, and identify temperature scaling and top-p filtering as the controls introduced for that purpose.
  Why it matters: Sampling diverse responses supplies the candidate answers that self-consistency requires.

Key terms: chain-of-thought prompting, sequential technique, parallel sampling, self-consistency sampling, text generation function, temperature scaling, top-p filtering, output diversity

4.4.1 Understanding the process of selecting the next token

This section traces how a model selects the next token from vocabulary-wide scores and illustrates those scores with an example.

• Next-token selection process: Trace how input text becomes token IDs, how the model's logits for the final input position are selected, and how the vocabulary entry with the largest score becomes the next token.
  Why it matters: This is the selection process that the following temperature-scaling discussion builds on.
• Vocabulary-wide logit scores: Describe what the final-position logit scores represent and how the plotted slice illustrates the highest-scoring token in the example.
  Why it matters: Seeing the scores as vocabulary-wide values makes clear what the model compares when it chooses a token.

Key terms: token IDs, logits, [:, -1], torch.argmax, vocabulary size, next_token_logits

4.4.2 Rescaling token scores (logits) via a temperature parameter

This subsection explains how dividing next-token logits by a positive temperature changes the differences between scores and, in turn, the likelihood of selecting different tokens.

• Temperature rescaling of logits: Describe how a temperature parameter rescales next-token logits before they are converted to probabilities: divide each logit by the temperature, leave scores unchanged at 1.0, and reject temperatures of zero or less.
  Why it matters: This is the operation the text generation function will use to control next-token selection.
• Score gaps and token choice: Explain why temperatures below 1.0 make score differences more pronounced, while temperatures above 1.0 shrink them and make lower-ranked tokens more likely to be sampled.
  Why it matters: Comparing score gaps explains why temperature can make next-token selection more confident or more diverse.

Key terms: temperature parameter, logits, temperature scaling, sharper, flatter, probability gap, lower-ranked tokens

4.4.3 Sampling the next token from a probability distribution

The section shows how softmax turns temperature-scaled logits into token probabilities and how sampling from those probabilities produces varied next-token choices.

• Softmax Token Probabilities: Explain how softmax converts rescaled logits into probabilities that can be interpreted and sampled.
  Why it matters: Sampling needs probabilities rather than raw token scores.
• Weighted Next Token Sampling: Distinguish multinomial sampling from greedy decoding and describe how repeated draws illustrate token probabilities.
  Why it matters: This distinction explains how a model can generate different completions from the same probability distribution.
• Temperature And Sampling Diversity: Explain how the section’s high- and low-temperature examples affect sampled tokens, and why sampling diversity can be useful.
  Why it matters: Temperature helps balance useful variety against the risk of unlikely, incoherent next tokens.

Key terms: softmax, logits, probability distribution, torch.softmax, multinomial sampling, greedy decoding, torch.multinomial, count_samples, temperature, sampling diversity

4.4.4 Adding temperature scaling to the text generation function

The section adds temperature scaling and sampling to a cached text-generation function and explains how to choose a temperature for different goals.

• Temperature sampling integration: Describe how the text-generation function adds temperature rescaling and sampling inside its token-generation loop, while a flexible wrapper lets the decoding strategy change without changing outer prompt handling.
  Why it matters: It shows where the new decoding behavior fits without requiring changes to prompt handling.
• Temperature choice tradeoffs: Explain how temperature settings trade off output diversity against reliability, including what zero, small nonzero, and much higher temperatures do.
  Why it matters: Choosing a temperature depends on whether the task benefits more from exploration or from a reliable likely answer.

Key terms: generate_text_temp_stream_cache, temperature rescaling, sampling, token-generation loop, generate_text_stream_concat_flex, temperature, greedy decoding, diversity, reliability

4.5 Balancing diversity and coherence with top-p sampling

Introduces top-p sampling as a way to limit unlikely tokens while retaining sampling diversity, and outlines the filtering steps.

• Top-p sampling motivation: Explain why top-p sampling is added to temperature-based sampling: greater diversity can also allow tokens unrelated to the user query, so top-p filters out very low-confidence tokens.
  Why it matters: It explains the problem the filter is meant to address.
• Top-p filtering sequence: Describe the high-level top-p filtering sequence and explain why the remaining probabilities must be renormalized before sampling.
  Why it matters: The outline shows how filtering can still leave a valid distribution for sampling.

Key terms: top-p filter, nucleus sampling, temperature scaling, token probabilities, cumulative sum, top-p cutoff, renormalizing

4.5.1 Selecting a subset of top-p tokens

This section shows how top-p filtering selects a probability-based subset of next tokens, then prepares that subset for sampling.

• Sorted cumulative token probabilities: Describe how sorting next-token probabilities from highest to lowest and taking their cumulative sum prepares them for top-p filtering.
  Why it matters: The cumulative totals make it possible to find a cutoff based on probability rather than a fixed number of tokens.
• Top-p threshold cutoff variants: Distinguish the simple cutoff that retains only tokens whose cumulative probability is at most p from the more common variant that also retains the token that crosses p.
  Why it matters: Whether the threshold-crossing token is included determines which tokens remain available to sample.
• Filtered sampling distribution: Explain what happens to token probabilities after the top-p cutoff and why filtering is useful before sampling.
  Why it matters: Filtering changes which tokens can be sampled while leaving a normalized distribution over those that remain.

Key terms: token probabilities, descending order, cumulative sum, torch.cumsum, top-p, top_p, keep_mask, cumulative probability mass, threshold, top-p filtering

4.5.2 Adding a top-p filter to the text generation function

This section adds top-p filtering to temperature-scaled text generation, shows how it changes sampled outputs, and contrasts it with top-k filtering.

• Top-p filter operation: Trace how top_p_filter selects and prepares token probabilities for sampling, including why the token that crosses the threshold remains eligible.
  Why it matters: This explains which tokens can still be sampled after filtering.
• Top-p placement and effects: Describe where top-p filtering enters temperature-scaled text generation and what the examples show about its effects and limits.
  Why it matters: Filtering changes the pool of possible next tokens, not whether the resulting response is correct.
• Top-k versus top-p selection: Distinguish the candidate sets retained by top-k and top-p filtering.
  Why it matters: The distinction clarifies whether the candidate pool has a fixed size or adapts to the probability distribution.

Key terms: top_p_filter, prefix cumulative mass, top_p threshold, renormalizes, temperature scaling, top-p filter, generate_text_top_p_stream_cache, top-k filtering, top-p, cumulative mass

4.6 Improving response accuracy with self-consistency

This section implements self-consistency voting and examines how sampling choices affect its accuracy, compute cost, and applicability.

• Self-consistency voting pipeline: Trace how self-consistency generates multiple responses with temperature and top-p sampling, extracts each final boxed answer, and selects the most frequent answer. Explain why this is inference-time scaling rather than a model update.
  Why it matters: Voting can turn several imperfect responses into a more accurate final answer without retraining the model.
• Sampling diversity balance: Explain how to adjust temperature and top-p for self-consistency by inspecting full responses: encourage more diversity when they are nearly identical, but reduce randomness when they become nonsensical.
  Why it matters: Sampling settings determine whether voting has useful alternatives or merely collects poor responses.
• Accuracy and compute tradeoff: Describe what the MATH-500 comparisons show about self-consistency, chain-of-thought prompting, and sampling alone. Explain why more samples can improve accuracy yet increase running time, with little additional gain in the reported 10-versus-3-sample comparison.
  Why it matters: The results help you decide whether an accuracy gain is worth the additional inference cost.
• Voting answer constraints: Describe the implementation’s behavior when final answers tie and explain why voting is harder for problems without extractable short final answers.
  Why it matters: Knowing what can be counted—and what happens in a tie—clarifies where this implementation can be used.

Key terms: self-consistency sampling, temperature scaling, top-p filtering, final boxed answer, majority voting, inference-time scaling, temperature, top-p, long answers, diversity

4.7 Summary

The chapter summary connects prompting and sampling techniques to the central inference-time scaling tradeoff: better answers can require more compute without retraining.