Under the Hood
Book ↗
Projects/Reasoning, Tools, and Retrieval
Project 25

Test-Time Reasoning (CoT, Self-Consistency, Best-of-N).

The Concept

The model has the capability. It does not have the deliberation. A student who knows how to add can still mess up a long arithmetic problem if they try to do every step in their head and only write down the final number. The same student, given a piece of scratch paper and told "show your work," makes fewer mistakes. The information needed to be correct was always in their head. The scratch paper changed what they could keep track of.

A language model has the same problem. When it generates a single answer token in response to a hard problem, it is doing every intermediate calculation in one forward pass, with no way to write anything down. Chain-of-thought, often shortened to CoT, is the scratch paper. You prompt the model with something like "Let's think step by step," and the model produces a sequence of intermediate reasoning tokens before giving its final answer. The final answer is now generated with all those intermediate tokens already in its context, so it has a much easier job.

That is one move. There are three more, and each one buys accuracy by spending more inference compute.

Self-consistency is the second move. Generate not one chain of reasoning, but eight. Sample them at a temperature high enough that the chains genuinely differ — temperature 0.7 is the usual setting. Read off the final answer from each chain. Then take a majority vote: whichever final answer appears most often across the eight chains is the answer you keep. The intuition is that there are many wrong ways to reason about a hard problem, but most of the right ways tend to converge on the same number. If five out of eight chains say "47" and the other three say three different wrong things, "47" is your best guess.

Best-of-N, written BoN, is the third move. Sample N candidate answers — full chains plus final answers — and pick the best one using a verifier. The verifier is a separate model trained to score answers, often called a reward model or outcome reward model (ORM) when it scores final outcomes. An ORM takes a question and a candidate solution and produces a number: higher is better. You take the candidate with the highest ORM score and return that one. Majority vote uses agreement to pick. BoN uses an external judge.

Process reward models (PRMs) are the fourth move. An ORM scores the final answer. A PRM scores each individual step in the reasoning chain: was step one correct, was step two correct, and so on. With a PRM you can do something more surgical than picking among finished candidates. At each reasoning step, you sample K continuations, score each one with the PRM, keep the best step, and continue. The search is now step-level rather than candidate-level. It is more compute, but the verifier can reject bad reasoning before it derails the whole chain.

PRMs are also where my opinion on this whole space gets sharpest. I think PRM training is one of the most fragile pipelines in the test-time playbook, and the literature understates how easy it is to ship a PRM that grades the wrong feature. The BREAK IT later in the chapter exists because I have seen this failure mode in the wild, on fusion candidates where a verifier was rewarding chain length and we mistook it for rewarding correctness for almost a full day.

A useful analogy to hold all four together: a student doing a math problem. The student who shows their work catches their own arithmetic mistakes; that is CoT. Three students working in parallel and majority-voting on the final answer do better than any one of them; that is self-consistency. A teacher who reads each student's final answer and picks the best one does better still; that is BoN with an ORM. A teacher who reads each step as the student writes it and tells them when they have gone wrong does the best of all; that is a PRM doing step-level search.

There is a search-tree version of the last move that is worth naming separately. MCTS stands for Monte Carlo Tree Search, and it generalizes step-level BoN into a real tree search. From the current partial reasoning, expand several candidate next-steps, score them with the PRM, recurse into the most promising one, occasionally back up and try a different branch when scores stop improving. The math is more involved, but the structure is the same: PRM-guided search through the space of possible reasoning chains. A rollout in this context is one full path from the current node down to a final answer. MCTS does many rollouts and uses their outcomes to decide which branches to develop further.

All five methods (direct answering, CoT, self-consistency, BoN, step-level BoN with a PRM, and MCTS) trade inference compute for accuracy. Direct answering is the cheapest and least accurate. MCTS is the most expensive and, when the PRM is good, the most accurate. There is no free improvement here. You are paying for accuracy in GPU-seconds at inference time, the same way you paid for it in GPU-hours at training time.

I think of this as the budget version of adaptive complexity routing. Multi-model routing work I have done across single, two, and three-model ensembles was solving the same problem at a higher level: when does it pay to spend more compute on a question? Test-time reasoning is the same question one rung down, inside a single model. The answer is also the same: sometimes a lot, sometimes nothing, and you cannot tell ahead of time without measuring.

Why It Matters

Without test-time reasoning, your model is leaving capability on the table. The training run already happened. The weights already encode whatever they encode. You can either ask the model to one-shot every hard problem and live with the failure rate, or you can spend a little more compute at inference and recover accuracy you already paid for during training.

The numbers in the literature are large enough to take seriously. On GSM8K, a grade-school math benchmark, naive prompting of a mid-sized model gets you maybe 30 percent accuracy. CoT prompting on the same model jumps to 50 percent or higher with no weight changes. Self-consistency on top of CoT, with eight samples, pushes that another five to ten points. BoN with a trained verifier can push it again. The o1 family of models from OpenAI, and the public discussion around them, leans heavily on this trick. The headline "reasoning" capability is in large part a story about spending more compute per query at inference time, using methods that are direct descendants of the ones in this chapter.

There is a second reason this chapter sits where it does. The previous chapters in this book taught you to spend training compute carefully. Project 23: Reward Models and RLHF taught you to train a reward model on preference data. Project 24: DPO and Preference Optimization taught you a cheaper alternative that bakes preferences into the policy directly. Test-time reasoning re-opens the question those chapters seemed to close. Maybe you do not need to bake every behavior into the weights. Maybe some behaviors are cheaper to get by spending compute at inference time, with a smaller policy and a smaller reward model. That tradeoff is one of the live research questions of the period this book is being written in.

A third reason, less direct but worth saying. Test-time methods make capability cheaper to study. You do not need to retrain anything to see whether the model "knows" how to do a task. Crank up the inference compute, watch what happens to accuracy, and you have a much clearer answer than any single-shot eval. Project 22: Evaluation Methodology taught you to be careful with benchmarks; test-time reasoning is one of the reasons that care matters.

The Build

build.py

$python projects/25_test-time-reasoning-cot-self-consistency-best-of-n/build.py --tiny

"""
Project 25: Test-Time Reasoning (CoT, Self-Consistency, Best-of-N)

Canonical runnable build for this project. Generated by tools/extract_code.py
from the book's chapter markdown — re-run that tool to regenerate.
"""

# === Step 1: Direct prompting and CoT prompting on GSM8K ===

def run_one(prompt, max_new_tokens=256, temperature=0.0):
    out = model.generate(prompt, max_new_tokens=max_new_tokens,
                         temperature=temperature, do_sample=temperature > 0)
    return out

def extract_answer(text):
    nums = re.findall(r"-?\d+\.?\d*", text.replace(",", ""))
    return float(nums[-1]) if nums else None

# === Step 2: Self-consistency ===

def self_consistency(question, n=8, temperature=0.7):
    answers = []
    for _ in range(n):
        text = run_one(cot_prompt(question), temperature=temperature)
        ans = extract_answer(text)
        if ans is not None:
            answers.append(ans)
    if not answers:
        return None
    counts = Counter(answers)
    return counts.most_common(1)[0][0]

# === Step 3: Train a tiny outcome reward model ===

class ORM(nn.Module):
    def __init__(self, base_model):
        super().__init__()
        self.base = base_model
        self.head = nn.Linear(base_model.hidden_size, 1)

    def forward(self, input_ids, attention_mask):
        h = self.base(input_ids, attention_mask=attention_mask).last_hidden_state
        last = h[torch.arange(h.size(0)), attention_mask.sum(dim=1) - 1]
        return self.head(last).squeeze(-1)

# === Step 4: Best-of-N with the ORM ===

def bon(question, n=8, temperature=0.7):
    candidates = []
    for _ in range(n):
        text = run_one(cot_prompt(question), temperature=temperature)
        score = orm.score(question, text)
        candidates.append((score, text))
    candidates.sort(key=lambda x: -x[0])
    best = candidates[0][1]
    return extract_answer(best)

# === Step 5: Train a tiny PRM ===

def rollout_label(question, prefix, m=8):
    correct = 0
    for _ in range(m):
        continuation = run_one(prefix, temperature=0.7)
        ans = extract_answer(prefix + continuation)
        if ans == ground_truth[question]:
            correct += 1
    return correct / m

# === Step 6: Step-level Best-of-N with the PRM ===

def step_level_bon(question, k=4, max_steps=12):
    chain = ""
    for _ in range(max_steps):
        candidates = []
        for _ in range(k):
            step = sample_step(question, chain, temperature=0.7)
            score = prm.score(question, chain + step)
            candidates.append((score, step))
        candidates.sort(key=lambda x: -x[0])
        best_step = candidates[0][1]
        chain += best_step
        if has_final_answer(chain):
            break
    return extract_answer(chain)

# === Step 7: A tiny MCTS over reasoning steps ===

def mcts(question, n_iter=32, k_expand=4):
    root = Node(prompt=cot_prompt(question))
    for _ in range(n_iter):
        node = select(root)
        children = expand(node, k=k_expand)
        for child in children:
            rollout_value = rollout_and_grade(child)
            backup(child, rollout_value)
    return best_leaf_answer(root)

The Steps

7 reference files

Each step is the code for one piece, in isolation — read them as the book introduces them. They assemble into build.py above.

Step 01Direct Prompting And Cot Prompting On Gsm8kstep_01_direct-prompting-and-cot-prompting-on-gsm8k.py
"""
Project 25: Step 1 — Direct prompting and CoT prompting on GSM8K

Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""

def run_one(prompt, max_new_tokens=256, temperature=0.0):
    out = model.generate(prompt, max_new_tokens=max_new_tokens,
                         temperature=temperature, do_sample=temperature > 0)
    return out

def extract_answer(text):
    nums = re.findall(r"-?\d+\.?\d*", text.replace(",", ""))
    return float(nums[-1]) if nums else None
Step 02Self Consistencystep_02_self-consistency.py
"""
Project 25: Step 2 — Self-consistency

Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""

def self_consistency(question, n=8, temperature=0.7):
    answers = []
    for _ in range(n):
        text = run_one(cot_prompt(question), temperature=temperature)
        ans = extract_answer(text)
        if ans is not None:
            answers.append(ans)
    if not answers:
        return None
    counts = Counter(answers)
    return counts.most_common(1)[0][0]
Step 03Train A Tiny Outcome Reward Modelstep_03_train-a-tiny-outcome-reward-model.py
"""
Project 25: Step 3 — Train a tiny outcome reward model

Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""

class ORM(nn.Module):
    def __init__(self, base_model):
        super().__init__()
        self.base = base_model
        self.head = nn.Linear(base_model.hidden_size, 1)

    def forward(self, input_ids, attention_mask):
        h = self.base(input_ids, attention_mask=attention_mask).last_hidden_state
        last = h[torch.arange(h.size(0)), attention_mask.sum(dim=1) - 1]
        return self.head(last).squeeze(-1)
Step 04Best Of N With The Ormstep_04_best-of-n-with-the-orm.py
"""
Project 25: Step 4 — Best-of-N with the ORM

Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""

def bon(question, n=8, temperature=0.7):
    candidates = []
    for _ in range(n):
        text = run_one(cot_prompt(question), temperature=temperature)
        score = orm.score(question, text)
        candidates.append((score, text))
    candidates.sort(key=lambda x: -x[0])
    best = candidates[0][1]
    return extract_answer(best)
Step 05Train A Tiny Prmstep_05_train-a-tiny-prm.py
"""
Project 25: Step 5 — Train a tiny PRM

Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""

def rollout_label(question, prefix, m=8):
    correct = 0
    for _ in range(m):
        continuation = run_one(prefix, temperature=0.7)
        ans = extract_answer(prefix + continuation)
        if ans == ground_truth[question]:
            correct += 1
    return correct / m
Step 06Step Level Best Of N With The Prmstep_06_step-level-best-of-n-with-the-prm.py
"""
Project 25: Step 6 — Step-level Best-of-N with the PRM

Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""

def step_level_bon(question, k=4, max_steps=12):
    chain = ""
    for _ in range(max_steps):
        candidates = []
        for _ in range(k):
            step = sample_step(question, chain, temperature=0.7)
            score = prm.score(question, chain + step)
            candidates.append((score, step))
        candidates.sort(key=lambda x: -x[0])
        best_step = candidates[0][1]
        chain += best_step
        if has_final_answer(chain):
            break
    return extract_answer(chain)
Step 07A Tiny Mcts Over Reasoning Stepsstep_07_a-tiny-mcts-over-reasoning-steps.py
"""
Project 25: Step 7 — A tiny MCTS over reasoning steps

Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""

def mcts(question, n_iter=32, k_expand=4):
    root = Node(prompt=cot_prompt(question))
    for _ in range(n_iter):
        node = select(root)
        children = expand(node, k=k_expand)
        for child in children:
            rollout_value = rollout_and_grade(child)
            backup(child, rollout_value)
    return best_leaf_answer(root)

Break It

break_it.py

$python projects/25_test-time-reasoning-cot-self-consistency-best-of-n/break_it.py --tiny

"""
Project 25: BREAK IT experiment.

Deliberately sabotages one mechanism from build.py to show what happens
when it's removed. Compare outputs to the working version.
"""

def truncate_chain(chain):
    steps = chain.split("\n")
    keep = max(1, len(steps) // 2)
    return "\n".join(steps[:keep])

broken_data = [(q, truncate_chain(c), rollout_label(q, truncate_chain(c)))
               for q, c, _ in original_data]

Read it in the book

This page ships the code. The full derivation — why the chain rule is mechanical not magical, what the computational graph buys you, the long-form reasoning — lives in Chapter 25 of Under the Hood.

Get the book →