5.10 Reasoning, Chain-of-Thought & Test-Time Compute¶

For the first few years of the transformer era, the recipe for better performance was simple: train a larger model on more data. The Scaling Laws: Kaplan, Chinchilla & Beyond chapter covers that axis in detail. But around 2022–2024, a second axis emerged: test-time compute — spending more computation at inference time to improve answer quality, independent of the model’s parameter count. Chain-of-thought prompting was the first crack in the wall. The o1/R1 family of models turned it into a new scaling regime. This chapter explains the full arc: from a three-word prompt trick to Monte Carlo Tree Search over reasoning trajectories, process reward models, and budget-forcing strategies.

Why Reasoning Needs More Than a Single Forward Pass¶

A 7B-parameter model and a 70B-parameter model both produce their answers in a single autoregressive forward pass per token. For factual recall or simple summarization that is often enough. But consider a multi-step math problem or a competitive programming task: the answer depends on a chain of interdependent deductions, and a single pass cannot “look back” and correct early errors. The architecture (see The Attention Mechanism From Scratch) does not have a recurrent error-correction mechanism; each token is generated by attending to the preceding context, not by iteratively refining a hypothesis.

The core insight behind test-time scaling is: reasoning is a search problem. Given a question, the solution is not at the top of a probability distribution — it is at the end of a path through a reasoning tree. More compute at test time lets us search that tree more thoroughly.

The rest of this chapter follows the historical progression: prompting tricks → voting → tree search → learned process rewards → trained reasoning models → scaling laws.

Chain-of-Thought Prompting¶

The Basic Technique¶

Wei et al. (Chain-of-Thought Prompting Elicits Reasoning in Large Language Models, 2022) showed that prepending a few worked-out examples — where each example shows the reasoning steps, not just the answer — dramatically improved accuracy on arithmetic and commonsense benchmarks. The prompt change is minimal:

# Standard prompt (answer only)
Q: Roger has 5 tennis balls. He buys 2 more cans of tennis balls.
   Each can has 3 balls. How many tennis balls does he have now?
A: 11

# Chain-of-thought prompt (steps shown)
Q: Roger has 5 tennis balls. He buys 2 more cans of tennis balls.
   Each can has 3 balls. How many tennis balls does he have now?
A: Roger starts with 5 balls. 2 cans × 3 balls/can = 6 balls.
   5 + 6 = 11. The answer is 11.

When the model sees several such examples in the few-shot context, it learns to “show its work.” Zero-shot CoT (Kojima et al., 2022) found that simply appending the phrase “Let’s think step by step.” to the question substantially improves accuracy even without worked examples.

Why It Works¶

The mechanistic explanation: CoT externalises intermediate state into the context window. A transformer cannot do multi-hop reasoning in depth in a single forward pass, but it can condition on previously generated text. By emitting intermediate calculations into the token stream, the model has those values available as attention targets for subsequent tokens. Each reasoning step is a shallow computation; chaining many shallow steps achieves deep computation.

This is directly analogous to the difference between a bounded-depth circuit (a single-pass model) and an unbounded sequential program. CoT effectively converts the fixed-depth network into a variable-depth one, governed by the length of the reasoning trace rather than the number of transformer layers.

import openai  # or any compatible client

def chain_of_thought_query(client, model: str, question: str) -> str:
    """
    Zero-shot chain-of-thought using the 'Let's think step by step' trick.
    The two-stage approach first elicits reasoning, then extracts the answer.
    """
    # Stage 1: elicit reasoning
    reasoning_prompt = f"{question}\n\nLet's think step by step."
    stage1 = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": reasoning_prompt}],
        temperature=0.0,
        max_tokens=512,
    )
    reasoning_trace = stage1.choices[0].message.content

    # Stage 2: extract final answer from the trace
    extraction_prompt = (
        f"{question}\n\n{reasoning_trace}\n\n"
        "Therefore, the final answer (just the number/word) is:"
    )
    stage2 = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": extraction_prompt}],
        temperature=0.0,
        max_tokens=32,
    )
    return stage2.choices[0].message.content.strip()

Self-Consistency: Marginalising Over Reasoning Paths¶

A single CoT trace can still hallucinate. Wang et al. (Self-Consistency Improves Chain of Thought Reasoning in Language Models, 2022) proposed self-consistency: sample \(N\) independent reasoning traces at temperature \(T > 0\), then take a majority vote over the final answers. Because different traces reach the correct answer via different paths, incorrect paths tend to disagree, while the correct answer clusters.

Formally, let \(\mathcal{R} = \{r_1, r_2, \ldots, r_N\}\) be \(N\) sampled reasoning chains and \(a_i = f(r_i)\) the extracted answer from chain \(i\). The self-consistent answer is:

This is equivalent to best-of-N or majority voting at the output level.

from collections import Counter
from typing import List

def self_consistent_answer(
    client,
    model: str,
    question: str,
    n_samples: int = 16,
    temperature: float = 0.8,
) -> str:
    """
    Sample N reasoning chains; return the most common final answer.
    For math: parse the last number. For MCQ: parse the letter.
    """
    answers: List[str] = []
    for _ in range(n_samples):
        response = client.chat.completions.create(
            model=model,
            messages=[
                {"role": "user", "content": question + "\nThink step by step."}
            ],
            temperature=temperature,
            max_tokens=512,
        )
        trace = response.choices[0].message.content
        # Heuristic: last boxed number or last line
        answer = extract_final_answer(trace)
        answers.append(answer)

    # Majority vote
    vote_counts = Counter(answers)
    winner, count = vote_counts.most_common(1)[0]
    print(f"Votes: {dict(vote_counts)}, winner={winner} ({count}/{n_samples})")
    return winner


def extract_final_answer(trace: str) -> str:
    """
    Attempt to find '\\boxed{...}' (LaTeX), then 'The answer is X',
    then fall back to the last line.
    """
    import re
    boxed = re.search(r"\\boxed\{([^}]+)\}", trace)
    if boxed:
        return boxed.group(1).strip()
    answer_line = re.search(r"[Tt]he answer is[:\s]+(\S+)", trace)
    if answer_line:
        return answer_line.group(1).strip(" .,")
    return trace.strip().splitlines()[-1]

Sampling N Traces in Practice: vLLM and math-verify¶

The loop above issues \(N\) sequential requests, which is the wrong shape for real work. Every serving engine exposes an \(n\)-samples-per-prompt parameter that prefills the prompt once and then decodes \(N\) continuations that share that KV cache, so best-of-64 costs far less than 64 independent requests (see Prefix Caching & KV-Cache Reuse):

from vllm import LLM, SamplingParams

llm = LLM(model="Qwen/Qwen2.5-Math-7B-Instruct", max_model_len=4096)
params = SamplingParams(n=64, temperature=0.8, top_p=0.95, max_tokens=1024)

outputs = llm.generate([f"{question}\nThink step by step."], params)
traces = [completion.text for completion in outputs[0].outputs]  # 64 traces

SGLang exposes the same knob (n= in its sampling parameters), and the OpenAI-compatible servers both engines ship accept n on /v1/completions.

The second practical piece is grading. The regex in extract_final_answer is a teaching device: it will call 1/2 and 0.5 different answers, and choke on \frac{1}{2}. Use a symbolic checker instead — Hugging Face’s math-verify (extracted from the open-r1 project) parses LaTeX/expressions and compares them for mathematical equivalence:

from math_verify import parse, verify

gold = parse("\\boxed{\\frac{1}{2}}")
pred = parse(trace)              # extracts the boxed/final expression from the trace
is_correct = verify(gold, pred)  # True for 0.5, 1/2, \frac{1}{2}, \dfrac12, ...

The same function is what you use as the reward in RLVR training (RL with Verifiable Rewards (RLVR) & The Reasoning Recipe) — verification at test time and verification in the training loop are literally the same code path.

Coverage vs. Selection: pass@k and maj@k¶

Repeated sampling produces two different quantities, and conflating them is the most common analysis error in test-time-compute work:

• Coverage — pass@k: does at least one of the \(k\) samples contain the correct answer? This is the ceiling any selection rule could ever reach.
• Selection — maj@k (self-consistency), best-of-N with a reward model: can we actually pick the right sample without an oracle?

pass@k must be estimated without bias. Drawing \(n \ge k\) samples and observing \(c\) correct ones, the unbiased estimator (Chen et al., Evaluating Large Language Models Trained on Code — the Codex paper, 2021) is

which is far lower-variance than literally sampling \(k\) traces once and checking:

import math

def pass_at_k(n: int, c: int, k: int) -> float:
    """Unbiased pass@k from n samples of which c are correct (Chen et al., 2021).

    1 - C(n-c, k) / C(n, k), written as a product to avoid huge binomials.
    """
    if n - c < k:           # fewer than k wrong samples => every k-subset has a correct one
        return 1.0
    return 1.0 - math.prod((n - c - i) / (n - i) for i in range(k))

# 64 samples of which 5 are correct: pass@1 is 5/64 = 7.8%, but pass@16 is ~78%.
assert abs(pass_at_k(64, 5, 1) - 5 / 64) < 1e-9
assert round(pass_at_k(64, 5, 16), 3) == 0.775

Brown et al. (Large Language Monkeys: Scaling Inference Compute with Repeated Sampling, 2024) show that coverage keeps climbing smoothly — roughly log-linearly in \(k\) — across several orders of magnitude of \(k\), while maj@k and reward-model selection plateau far earlier. The gap between those two curves is exactly what a verifier buys you. Where a sound verifier exists (unit tests for code, a proof checker, an exact-answer grader for competition math), test-time compute converts nearly all coverage into accuracy; where it does not, selection quality — not the model’s ability to ever find the answer — is the binding constraint. Report both numbers in any test-time-compute experiment: pass@k tells you whether to invest in a better verifier, maj@k tells you what you can ship today.

Tree-of-Thought and Graph-of-Thought¶

Self-consistency searches the answer space at the leaf level. Tree of Thoughts (Yao et al., 2023) searches the reasoning process itself. The key idea: a solution can branch at any intermediate step, and we can evaluate each partial reasoning path before committing to it.

Tree of Thoughts (ToT)¶

ToT treats reasoning as a tree where:

• Nodes are partial solution states (think: “after 2 reasoning steps”).
• Edges are model-generated next-step candidates.
• Value function \(V(s)\) scores whether a partial state looks promising.
• Search algorithm is BFS, DFS, or beam search over this tree.

import heapq
from dataclasses import dataclass, field
from typing import Optional

@dataclass(order=True)
class ThoughtNode:
    """A node in the Tree-of-Thought search. Lower neg_value = higher priority."""
    neg_value: float            # for min-heap: negate the value
    depth: int = field(compare=False)
    thought: str = field(compare=False)
    parent: Optional["ThoughtNode"] = field(compare=False, default=None)

    def trace(self) -> str:
        """Reconstruct the full reasoning chain from root to this node."""
        steps = []
        node = self
        while node is not None:
            steps.append(node.thought)
            node = node.parent
        return "\n".join(reversed(steps))


def tot_beam_search(
    client,
    model: str,
    problem: str,
    beam_width: int = 4,
    max_depth: int = 5,
    n_candidates_per_node: int = 3,
) -> str:
    """
    BFS-style Tree-of-Thought beam search.

    At each depth level we:
      1. Generate `n_candidates_per_node` next thoughts for each beam node.
      2. Score every candidate with the value model.
      3. Keep the top `beam_width` nodes to form the next beam.
    Returns the highest-scoring terminal node's full trace.
    """
    # Initialize beam with the problem as context
    root = ThoughtNode(neg_value=0.0, depth=0, thought=f"Problem: {problem}")
    beam = [root]

    for depth in range(1, max_depth + 1):
        candidates = []
        for parent_node in beam:
            # Generate candidate next thoughts
            for _ in range(n_candidates_per_node):
                thought = generate_next_thought(
                    client, model, parent_node.trace(), depth
                )
                value = score_thought(
                    client, model, parent_node.trace(), thought, problem
                )
                candidates.append(
                    ThoughtNode(
                        neg_value=-value,
                        depth=depth,
                        thought=thought,
                        parent=parent_node,
                    )
                )
        # Keep top beam_width by value
        candidates.sort()
        beam = candidates[:beam_width]

        # Early termination if top node has a final answer
        if is_terminal(beam[0].thought):
            break

    return beam[0].trace()


def generate_next_thought(client, model, context, depth):
    """Ask the model for the next reasoning step given current context."""
    prompt = (
        f"{context}\n\n"
        f"Step {depth}: What is the next logical reasoning step? "
        "Be concise and specific."
    )
    r = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}],
        temperature=0.8, max_tokens=200,
    )
    return r.choices[0].message.content.strip()


def score_thought(client, model, context, thought, problem):
    """
    Value model: ask the LM to rate the current partial solution on 0-10.
    In production, this would be a fine-tuned process reward model.
    """
    prompt = (
        f"Problem: {problem}\n\n"
        f"Partial solution so far:\n{context}\n\n"
        f"Proposed next step: {thought}\n\n"
        "On a scale of 0 to 10, how promising is this partial solution? "
        "Answer with just a number."
    )
    r = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}],
        temperature=0.0, max_tokens=5,
    )
    try:
        return float(r.choices[0].message.content.strip())
    except ValueError:
        return 5.0  # default if parse fails


def is_terminal(thought: str) -> bool:
    """Check if a thought contains a final answer marker."""
    markers = ["final answer", "therefore", "the answer is", "\\boxed"]
    return any(m in thought.lower() for m in markers)

Graph of Thought¶

Graph of Thoughts (Besta et al., 2023) generalises ToT by allowing reasoning states to merge: two branches that arrive at a compatible sub-conclusion can be combined into a single node, avoiding re-exploration. This is especially useful for problems with independent sub-problems (e.g., multi-step proofs where different lemmas can be proved in parallel).

The search graph is a DAG: nodes are reasoning states, directed edges are operations (Generate, Score, Aggregate), and the planner selects which operations to execute. In practice, for most LLM tasks ToT suffices; GoT matters when sub-problems are truly decomposable.

Process Reward Models vs. Outcome Reward Models¶

Self-consistency and ToT still need a quality signal to guide search. Two families of reward models have emerged:

Outcome Reward Models (ORM)¶

An ORM takes the full solution trace and scores it based on the final answer alone:

Training is simple: label full solutions as correct/incorrect (binary) and train a scalar-output classifier. ORMs are easier to collect data for — you just need a ground-truth answer checker. However, they provide no intermediate signal; they cannot distinguish a flawed-reasoning-but-lucky answer from a sound proof.

Process Reward Models (PRM)¶

A PRM assigns a scalar reward to each reasoning step:

Lightman et al. (Let’s Verify Step by Step, 2023) trained a step-level verifier on human-annotated mathematics solutions and showed it substantially outperforms ORM for guiding best-of-N search. The intuition: a bad reasoning step early can doom a solution even if the final answer is guessed correctly, and PRM catches this.

import torch
import torch.nn as nn
from transformers import AutoModel, AutoTokenizer

class ProcessRewardModel(nn.Module):
    """
    Lightweight PRM: finetune a language model backbone to predict
    per-step correctness. At inference, we score each step independently.
    """
    def __init__(self, backbone_name: str = "meta-llama/Llama-3-8B"):
        super().__init__()
        self.backbone = AutoModel.from_pretrained(backbone_name)
        hidden_dim = self.backbone.config.hidden_size
        # Scalar head: predicts P(step is correct | problem, steps so far)
        self.head = nn.Linear(hidden_dim, 1)

    def forward(
        self,
        input_ids: torch.Tensor,       # [B, T]
        attention_mask: torch.Tensor,  # [B, T]
        step_boundary_positions: torch.Tensor,  # [B] — index of last token in current step
    ) -> torch.Tensor:  # [B] logits
        outputs = self.backbone(
            input_ids=input_ids,
            attention_mask=attention_mask,
        )
        hidden = outputs.last_hidden_state  # [B, T, H]
        # Extract hidden state at the step boundary token
        step_hidden = hidden[
            torch.arange(hidden.size(0)), step_boundary_positions
        ]  # [B, H]
        return self.head(step_hidden).squeeze(-1)  # [B]

    @torch.no_grad()
    def score_steps(
        self,
        tokenizer: AutoTokenizer,
        problem: str,
        steps: list[str],
        device: str = "cuda",
    ) -> list[float]:
        """
        Score each prefix (problem + steps[:k]) and return per-step probabilities.
        """
        scores = []
        context = f"Problem: {problem}\n"
        for step in steps:
            context += f"Step: {step}\n"
            enc = tokenizer(context, return_tensors="pt").to(device)
            boundary = enc["input_ids"].shape[1] - 1  # last token position
            logit = self.forward(
                enc["input_ids"],
                enc["attention_mask"],
                torch.tensor([boundary], device=device),
            )
            prob = torch.sigmoid(logit).item()
            scores.append(prob)
        return scores

Where Step Labels Come From¶

The module above is useless without training data, and “who labels every step?” is the black box that stops most people from building a PRM. There are two answers.

Human annotation. OpenAI released PRM800K, roughly 800K step-level labels (positive / neutral / negative) over solutions to MATH problems, collected for Let’s Verify Step by Step. It is open and is still the standard supervised starting point — but it exists only for one domain, and you cannot afford to recreate it.

Automatic annotation by Monte-Carlo rollout. Math-Shepherd (Wang et al., 2024) removes the human entirely: the “correctness” of a prefix is defined as the empirical probability that continuing from that prefix reaches the gold answer. Roll out \(M\) completions from each step boundary, grade them with your verifier, and use the success fraction as the label.

def mc_step_labels(rollout_fn, is_correct, problem, steps, gold, n_rollouts=8):
    """Math-Shepherd-style automatic PRM labels.

    rollout_fn(prefix, n) -> list[str]: n independent continuations (batch this on vLLM).
    Returns one label per step: the fraction of completions from that prefix
    that reach the gold answer. Train the PRM head with BCE against these.
    """
    labels, prefix = [], f"Problem: {problem}\n"
    for step in steps:
        prefix += f"Step: {step}\n"
        completions = rollout_fn(prefix, n_rollouts)
        n_ok = sum(is_correct(c, gold) for c in completions)
        labels.append(n_ok / n_rollouts)   # soft label; hard variant = float(n_ok > 0)
    return labels

Two consequences worth internalising. First, the cost: labelling one solution costs n_rollouts × n_steps generations, so a modest 20K-solution PRM dataset is millions of completions — expensive, but embarrassingly parallel and therefore a pure throughput problem for a vLLM fleet. Second, the semantics: an MC label is a value function estimate under the current policy, not ground truth about logical validity. A step that is logically wrong but that the policy usually recovers from gets a high label. This is why MC-labelled PRMs drift as the policy improves and generally need re-labelling after major policy updates.

If you do not want to train one, several PRMs ship as open weights — for example peiyi9979/math-shepherd-mistral-7b-prm (the original MC-labelled model), Qwen’s Qwen2.5-Math-PRM-7B, and Skywork’s Skywork-o1-Open-PRM series. All are plain AutoModelForCausalLM/AutoModel checkpoints with a scoring head or a designated step-separator token whose logit you read; check each model card for the exact step delimiter, because feeding the wrong separator silently produces meaningless scores.

PRM-Guided Best-of-N¶

Given a PRM, best-of-N search becomes:

1. Sample \(N\) full solution traces.
2. Score each trace: \(s_i = \min_{t} r_\text{PRM}(x, y_{1:t})\) (the minimum step score — the weakest link).
3. Return the trace with the highest \(s_i\).

Using minimum is conservative; alternatives include product of step scores, or the score of the final step only.

MCTS for Language Model Reasoning¶

Monte Carlo Tree Search (MCTS) provides a principled way to allocate search budget. It is the algorithm behind AlphaGo/AlphaZero, and it adapts naturally to reasoning trees.

The Four Phases¶

The selection uses the UCB1 formula adapted for trees (PUCT, as in AlphaZero):

where \(Q(s,a)\) is the mean value of subtree rooted at child \(a\), \(P(a|s)\) is the LM’s prior probability for that action, \(N(s)\) is the visit count of the parent, and \(c_\text{puct}\) is an exploration constant (typically 1–5).

import math
from dataclasses import dataclass, field
from typing import Optional

@dataclass
class MCTSNode:
    state: str            # reasoning context so far
    parent: Optional["MCTSNode"] = None
    children: list["MCTSNode"] = field(default_factory=list)
    N: int = 0            # visit count
    Q: float = 0.0        # mean value
    P: float = 1.0        # prior from LM (log-prob of this branch)

    def ucb_score(self, c_puct: float = 2.0) -> float:
        if self.N == 0:
            return float("inf")
        parent_N = self.parent.N if self.parent else self.N
        return self.Q + c_puct * self.P * math.sqrt(parent_N) / (1 + self.N)

    def is_leaf(self) -> bool:
        return len(self.children) == 0


def mcts_search(
    lm_policy,       # callable(state) -> list[(next_thought, log_prob)]
    value_fn,        # callable(state) -> float  (PRM or rollout)
    root_state: str,
    n_iterations: int = 50,
    expansion_width: int = 3,
    c_puct: float = 2.0,
) -> MCTSNode:
    """
    Minimal MCTS for LLM reasoning. Returns the root node; caller can
    extract the best path by following max-Q children.
    """
    root = MCTSNode(state=root_state)

    for _ in range(n_iterations):
        # --- SELECT ---
        node = root
        path = [node]
        while not node.is_leaf():
            node = max(node.children, key=lambda c: c.ucb_score(c_puct))
            path.append(node)

        # --- EXPAND ---
        if node.N > 0 or node is root:  # expand visited nodes or root (identity, not ==)
            candidates = lm_policy(node.state)  # [(thought, log_prob), ...]
            for thought, log_prob in candidates[:expansion_width]:
                child = MCTSNode(
                    state=node.state + "\n" + thought,
                    parent=node,
                    P=math.exp(log_prob),
                )
                node.children.append(child)
            if node.children:
                # Descend into first unexplored child
                node = node.children[0]
                path.append(node)

        # --- SIMULATE (evaluate) ---
        value = value_fn(node.state)

        # --- BACKPROP ---
        for n in reversed(path):
            n.N += 1
            n.Q += (value - n.Q) / n.N  # running mean update

    return root

MCTS is notably more sample-efficient than naive best-of-N: it focuses computation on promising subtrees rather than wasting samples on dead ends. The tradeoff is latency — MCTS is sequential whereas best-of-N is trivially parallelisable.

Test-Time Scaling Laws¶

The observation that “more test-time compute improves accuracy” raises the natural question: at what rate? Snell et al. (Scaling LLM Test-Time Compute Optimally, 2024) and related work characterise this empirically.

The key findings:

1. Power-law scaling. For best-of-N, accuracy improves roughly as \(1 - \epsilon \cdot N^{-\alpha}\), where \(\alpha\) depends on problem difficulty and the model’s base capability. Typical values of \(\alpha\) are on the order of 0.1–0.3.

2. Compute-optimal frontier. For a fixed inference FLOPs budget, there is an optimal allocation between model size and number of samples. Smaller models with more samples can outperform larger models with fewer samples on many reasoning tasks.

3. PRM unlocks steeper scaling. Best-of-N with an ORM saturates quickly (the verifier cannot discriminate among many similar wrong answers). PRM-guided search exhibits steeper and longer scaling before saturation.

4. Difficulty modulates the return. Easy problems saturate quickly; hard problems continue to benefit from more compute. This motivates adaptive test-time compute: route easy queries to cheap paths, hard ones to expensive tree search.

5. Two axes: parallel vs. sequential. You can spend a token budget in parallel (N independent samples, selected by a verifier) or sequentially (one trace that revises itself — the model reads its own draft and corrects it). Snell et al. find the better choice depends on difficulty: easy problems, where the model’s first guess is nearly right, favour sequential revision; hard problems, where the model needs a different idea rather than a fix, favour parallel sampling. Sequential scaling is what long-thinking models internalise; parallel scaling is what best-of-N does externally. They compose — sample N long-thinking traces and vote.

Long-Thinking Models: o1, DeepSeek-R1, and the Reasoning Era¶

OpenAI’s o1 (September 2024) was the first publicly released model designed around extended internal reasoning — a “thinking” trace visible to the model but (in the initial release) hidden from users. DeepSeek-R1 (January 2025) reproduced and open-sourced the training recipe. Within a year the approach became the industry default: OpenAI’s o3 (April 2025) and unified GPT-5 (August 2025), Anthropic’s Claude, and Google’s Gemini all ship extended reasoning as a built-in mode rather than a separate product. The key innovations:

Training Reasoning Models¶

Rather than prompting an existing model to produce CoT at inference time, reasoning models are trained to generate long, high-quality internal reasoning traces, guided by RL. The pipeline typically combines:

1. SFT warm-up: fine-tune on human-written or distilled CoT traces to establish the format.
2. RL with verifiable rewards (RLVR): the policy is reinforced when the final answer (extracted from the reasoning trace) is verifiable-correct. See RL with Verifiable Rewards (RLVR) & The Reasoning Recipe for the training details.
3. PRM integration: some systems include a learned process reward to shape the quality of intermediate steps, not just final correctness.

The RL algorithm itself can be PPO (see Policy Gradients & PPO for Language Models), GRPO (see GRPO, RLOO & Critic-Free RL), or simpler outcome-only methods.

What Emerges: “Aha” Moments¶

DeepSeek-R1 reports that with sufficient RL training, models develop self-reflection behaviors: recognising mid-trace that a path is wrong, backtracking, and retrying — without being explicitly supervised to do so. The model also learns to allocate more tokens to harder sub-problems. These behaviors are not hand-coded; they emerge from the gradient signal.

# Illustration: a reasoning trace from an R1-style model
# (simplified; real traces are much longer and more varied)

EXAMPLE_REASONING_TRACE = """
<think>
Let me work through this step by step.

The problem asks for the number of ways to arrange 5 distinct books
on a shelf such that two specific books (A and B) are never adjacent.

First, let me count total arrangements: 5! = 120.

Now, let me count arrangements where A and B ARE adjacent.
Treat A+B as a single "super-book": 4! arrangements × 2 (AB or BA) = 48.

Wait, let me double-check: treating [AB] as one unit, we have 4 units total.
4! = 24 arrangements, and we can have AB or BA so ×2 = 48. Yes, that's right.

Therefore, arrangements where A and B are NOT adjacent = 120 - 48 = 72.
</think>

The answer is **72**.
"""

Budget Forcing¶

Budget forcing was introduced by Muennighoff et al. (s1: Simple Test-Time Scaling, 2025) and is startlingly crude for how well it works. It is decoding-time control over the length of the thinking block, in two directions:

• To cut thinking short, force-emit the end-of-thinking delimiter (e.g. </think>) plus a lead-in like Final Answer: once the trace reaches \(B\) tokens. The model then has to commit with whatever it has.
• To extend thinking, do the opposite: suppress the end-of-thinking delimiter when the model tries to emit it and append the string Wait (or Hmm, But) to the trace. The model reliably picks the thread back up, often catching its own error — this is the whole mechanism, no training required.

Because the “Wait” move can be repeated, thinking length becomes an actual dial rather than something you merely ask the model to respect, and accuracy on hard math climbs as you turn it up — though only so far: past some point the extended trace starts looping and the curve flattens, so budget forcing extrapolates capability, it does not create it. The s1 recipe pairs this dial with SFT on just ~1K carefully selected reasoning traces (s1K), which is the cheapest credible route to a long-thinking model of your own.

Two implementation notes. Genuine budget forcing lives in the decoding loop (a logit processor that bans the stop token, plus injected text), so it needs a local engine — in vLLM you implement it with a logits_processor or by generating in segments with SamplingParams(stop=...) and re-issuing the continuation. The prompt-level variant below is the portable approximation that works through any hosted API, and modern APIs also expose an explicit thinking-token budget parameter (Anthropic’s budget_tokens, Gemini’s thinking budget, OpenAI’s reasoning.effort levels) — prefer those when available:

def forced_budget_inference(
    client,
    model: str,
    question: str,
    budget_tokens: int = 512,
) -> str:
    """
    Budget-forced inference: append a budget instruction to the system prompt.
    The model will attempt to reason within the token limit.

    For thinking models that expose a `thinking_tokens` parameter,
    use that directly. This shows the prompting-based variant.
    """
    system = (
        "You are a careful reasoning assistant. "
        f"You have a thinking budget of approximately {budget_tokens} tokens. "
        "Think concisely, prioritise the most important reasoning steps, "
        "and then give the final answer."
    )
    response = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": question},
        ],
        max_tokens=budget_tokens + 128,  # +128 for the answer after thinking
        temperature=0.0,
    )
    return response.choices[0].message.content


def adaptive_budget_routing(
    client,
    model_fast: str,
    model_slow: str,
    question: str,
    difficulty_threshold: float = 0.7,
) -> tuple[str, str]:
    """
    Route to fast (small budget) or slow (large budget) based on estimated
    problem difficulty. Returns (answer, model_used).

    A real system would use a small classifier to estimate difficulty.
    Here we use a heuristic: question length + keyword presence.
    """
    hard_keywords = ["prove", "derive", "optimal", "complexity", "algorithm"]
    difficulty_score = min(
        1.0,
        len(question) / 500 + 0.15 * sum(k in question.lower() for k in hard_keywords)
    )
    if difficulty_score < difficulty_threshold:
        answer = forced_budget_inference(client, model_fast, question, budget_tokens=256)
        return answer, f"{model_fast} (fast path)"
    else:
        answer = forced_budget_inference(client, model_slow, question, budget_tokens=4096)
        return answer, f"{model_slow} (slow path)"

The Scaling Picture¶

Test-time compute as a scaling axis has now been confirmed empirically: holding model parameters fixed, reasoning accuracy can be improved by increasing the number of tokens spent thinking. The shape of this scaling curve differs from the training-time curve (loss ∝ \(C^{-0.07}\) in Chinchilla). Test-time scaling is typically steeper initially (large gains from 1→16 samples) but saturates more quickly at the extreme (diminishing returns after thousands of tokens).

Critically, training-time and test-time scaling compose: a model trained with more compute also benefits more from test-time search, so the frontier model uses both axes.

Putting It All Together: A Taxonomy of Test-Time Strategies¶

Gains are highly task-dependent and model-dependent; treat them as order-of-magnitude intuitions, not precise figures.

Infrastructure and Production Considerations¶

Long reasoning traces change the inference serving problem. A trace that is 4,000 tokens long generates a KV cache 20× larger than a 200-token response. See PagedAttention & KV-Cache Memory Management for how vLLM handles this, and Inference Economics: Latency, Throughput & Cost for cost modeling.

Key practical points:

• Parallelise best-of-N across model replicas, not within a single model. Each sample is fully independent.
• Disaggregate thinking from answering: in prefill-decode disaggregated systems (see Disaggregated Prefill/Decode & Chunked Prefill), the thinking phase is a long decode run. Route thinking requests to throughput-optimised nodes.
• Caching reasoning prefixes: if many queries share a problem preamble, prefix KV caching (see Prefix Caching & KV-Cache Reuse) provides significant savings.
• Budget forcing at the API level: expose thinking_tokens as a first-class parameter so product teams can tune the cost-accuracy tradeoff per use case.

Key Takeaways¶

Further Reading¶

• Wei et al., Chain-of-Thought Prompting Elicits Reasoning in Large Language Models, NeurIPS 2022.
• Wang et al., Self-Consistency Improves Chain of Thought Reasoning in Language Models, ICLR 2023.
• Kojima et al., Large Language Models are Zero-Shot Reasoners, NeurIPS 2022.
• Yao et al., Tree of Thoughts: Deliberate Problem Solving with Large Language Models, NeurIPS 2023.
• Besta et al., Graph of Thoughts: Solving Elaborate Problems with Large Language Models, AAAI 2024.
• Lightman et al., Let’s Verify Step by Step, ICLR 2024.
• Snell et al., Scaling LLM Test-Time Compute Optimally, arXiv 2024.
• Chen et al., Evaluating Large Language Models Trained on Code (the Codex paper), arXiv 2021 — origin of the unbiased pass@k estimator.
• Brown et al., Large Language Monkeys: Scaling Inference Compute with Repeated Sampling, arXiv 2024.
• Wang et al., Math-Shepherd: Verify and Reinforce LLMs Step-by-step without Human Annotations, ACL 2024.
• Muennighoff et al., s1: Simple Test-Time Scaling, arXiv 2025.
• DeepSeek-AI, DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning, arXiv 2025.
• Schulman et al., Proximal Policy Optimization Algorithms, arXiv 2017 (the RL backbone for many reasoning training pipelines).
• Silver et al., Mastering the Game of Go without Human Knowledge (AlphaGo Zero), Nature 2017 (MCTS/PUCT foundations).

Exercises¶

1. (Conceptual) The chapter argues that chain-of-thought “converts the fixed-depth network into a variable-depth one.” A plain transformer has a fixed number of layers \(L\), yet CoT lets it solve problems that seem to require more than \(L\) sequential steps of computation. Explain the mechanism that makes this possible. In your answer, address: (a) why a single forward pass is limited to bounded-depth computation, and (b) what resource CoT uses to store and reuse intermediate results across steps.

2. (Quantitative) A model answers a particular problem correctly with probability \(p\) on a single sampled trace, and errors are independent across samples. You apply self-consistency (majority vote) with \(N = 3\) samples; the vote is “correct” when at least 2 of the 3 traces are correct.

(a) For \(p = 0.70\), compute \(P(\text{majority correct})\) by hand and state the change relative to a single sample.

(b) Repeat for \(p = 0.40\). What does the sign of the change illustrate about a claim made in the chapter?

3. (Conceptual) The chapter’s “Common pitfall” warns that best-of-N with an outcome reward model (ORM) degrades as \(N\) grows large (e.g., \(N = 256\)), and recommends switching to a process reward model (PRM) once \(N \gtrsim 32\). (a) Explain why increasing \(N\) makes an ORM worse rather than monotonically better. (b) Explain concretely why a PRM is more robust to this failure, referencing how PRM scores are aggregated in the chapter’s PRM-guided best-of-N.

4. (Quantitative) Consider one MCTS selection step using the chapter’s PUCT rule, \(\text{score}(s,a) = Q(s,a) + c_\text{puct}\, P(a\mid s)\, \dfrac{\sqrt{N(s)}}{1 + N(s,a)}\), with \(c_\text{puct} = 2\) and parent visit count \(N(s) = 25\). The parent has two children:

• Child A (well-explored): \(Q = 0.80\), \(P = 0.30\), \(N(s,a) = 16\).
• Child B (barely explored): \(Q = 0.40\), \(P = 0.50\), \(N(s,a) = 1\).

Compute both PUCT scores and state which child is selected. Interpret the result in terms of the exploration/exploitation tradeoff.

5. (Quantitative) The chapter models best-of-N accuracy with the power law \(\text{acc}(N) = 1 - \epsilon\, N^{-\alpha}\). Suppose you measure \(\text{acc}(1) = 0.50\) and \(\text{acc}(16) = 0.75\).

(a) Solve for \(\epsilon\) and \(\alpha\).

(b) Predict \(\text{acc}(64)\).

© The chapter says typical \(\alpha\) is “on the order of 0.1–0.3.” Is your fitted value in range, and what would a larger \(\alpha\) mean for the payoff of extra samples?

6. (Implementation) The chapter’s PRM-guided best-of-N uses the minimum step score (the “weakest link”), while the boxed equation instead uses the product of step scores. (a) Implement a single function prm_best_of_n(traces, step_scores, aggregation) that selects the best trace index under aggregation in {"min", "product", "last"}, matching the chapter’s definitions. (b) Implement weighted_self_consistent_answer(answers, weights): a weighted majority vote that reduces to ordinary self-consistency when all weights are equal, so PRM/ORM confidence can be folded into voting. Keep the style consistent with the chapter’s code.

from typing import List

def prm_best_of_n(
    traces: List[List[str]],         # each trace = list of step strings
    step_scores: List[List[float]],  # per-step PRM probs, aligned with traces
    aggregation: str = "min",
) -> int:
    """
    Return the index of the best trace under the chosen PRM aggregation:
      - 'min'     : weakest-link score  min_t r(x, y_{1:t})
      - 'product' : product of step scores  prod_t r(x, y_{1:t})
      - 'last'    : score of the final step only
    """
    def aggregate(scores: List[float]) -> float:
        if not scores:
            return 0.0                      # empty trace: worst possible
        if aggregation == "min":
            return min(scores)
        if aggregation == "product":
            p = 1.0
            for s in scores:
                p *= s
            return p
        if aggregation == "last":
            return scores[-1]
        raise ValueError(f"unknown aggregation: {aggregation}")

    trace_values = [aggregate(s) for s in step_scores]
    return max(range(len(traces)), key=lambda i: trace_values[i])

from typing import List

def weighted_self_consistent_answer(
    answers: List[str],
    weights: List[float],
) -> str:
    """
    Majority vote weighted by per-trace quality (e.g. ORM/PRM score).
    Reduces to plain self-consistency when all weights are equal.
    """
    tally: dict[str, float] = {}
    for a, w in zip(answers, weights):
        tally[a] = tally.get(a, 0.0) + w
    # arg-max over accumulated weight; ties broken by first-seen order
    return max(tally, key=tally.get)