Module 13: Reasoning & Test-Time Compute

Introduction

Every module so far spends a fixed amount of compute per answer: run the model once, decode a token at a time (m08), stop. The answer you get is the answer the model happened to produce on its single pass. But hard problems — multi-step arithmetic, logic, code — are exactly the ones where a single pass is most likely to slip.

Test-time compute is the frontier’s newest lever: spend more compute at inference to get a better answer, with the weights frozen. No retraining, no bigger model — just let the model think longer, or think several times and reconcile. This is the idea behind OpenAI’s o1 and DeepSeek-R1, and this module builds its most important, most learnable form from scratch:

  • Chain-of-thought (CoT): let the model write intermediate steps before the final answer, turning one hard leap into several easy ones.
  • Self-consistency: sample many independent chains, read off each final answer, and take the majority vote. One idea, a large, reliable gain.
  • Best-of-N: sample N candidates and let a verifier pick the best.

Why it matters for LLMs:

  • Self-consistency alone lifted GSM8K accuracy by +17.9 points over a single chain-of-thought pass (Wang et al., 2022) — no new parameters.
  • The reasoning models topping today’s benchmarks are, at their core, models that were taught (with RL) to use test-time compute well. Understanding the parallel, no-training version first makes the RL version legible.

What You’ll Learn

After this module, you can:

  • Explain the difference between train-time and test-time compute, and why spending more of the latter raises accuracy.
  • Build self-consistency from scratch: extract answers, tally a plurality vote, and aggregate many sampled chains.
  • Prove why voting works with Condorcet’s jury theorem, and see the exact condition it needs (the correct answer is the mode, which is weaker than “right more than half the time”).
  • Implement best-of-N with a verifier and see when it beats voting.
  • Read the compute–accuracy curve — the empirical signature of test-time scaling — and place chain-of-thought and RL-for-reasoning (o1 / R1) on it.
  • Build GRPO from scratch — the critic-free RL recipe behind DeepSeek-R1 — and train a toy policy to answer correctly from a verifiable reward alone, watching the group-relative advantage do the work of a value network.

Prerequisites

This module requires familiarity with:

  • Module 08: Generation — temperature sampling and the generate loop that produces the diverse chains we vote over.
  • Module 12: Alignment — the aligned model that actually produces answers to reason about.

Intuition: Think Longer, Answer Better

Ask a person a hard arithmetic question and make them answer instantly, and they’ll often be wrong. Give them scratch paper — chain-of-thought — and they do better. Let them solve it three different ways and go with the answer they reached most often — self-consistency — and they do better still. None of this changes who is solving the problem; it changes how much thinking they spend.

Language models are the same. A single greedy decode is the instant answer. Two cheap upgrades, neither touching the weights:

  1. Chain-of-thought — prompt the model to emit reasoning steps before the answer. Each token now conditions on written-out intermediate work instead of leaping straight to the result.
  2. Self-consistency — because sampling (m08) makes each run take a different reasoning path, run it many times and vote. A correct answer is a fixed target many paths land on; wrong answers scatter.

Step through it: one prompt, many sampled chains, one vote.

NoteKey Insight

Self-consistency needs the model to be unreliable but not adversarial: as long as the correct answer is the single most likely one per sample, more samples sharpen the vote toward it. The reasoning path is a means to an end — we marginalize it out and keep only where the chains agree.

The Math: Why Voting Helps (Condorcet)

Model each sampled chain as an independent voter that is correct with probability \(p\). If it is wrong, take the worst case first: every wrong vote lands on the same wrong answer. Then the majority is correct exactly when more than half of \(n\) voters are correct:

\[ P_\text{correct}(n) = \sum_{k > n/2} \binom{n}{k}\, p^{k} (1-p)^{n-k}. \]

This is Condorcet’s jury theorem. Its behavior is a sharp phase transition at \(p = \tfrac12\):

  • If \(p > \tfrac12\), \(P_\text{correct}(n) \to 1\) as \(n \to \infty\) — more votes, more accuracy.
  • If \(p < \tfrac12\), it goes to \(0\) — voting amplifies a bad model’s errors.
  • If \(p = \tfrac12\), it stays at \(\tfrac12\) forever.

Real self-consistency does better than this bound, because wrong answers do not collude — they scatter across many values. Then the correct answer only has to beat each wrong answer individually (be the mode), a much weaker condition than \(p > \tfrac12\). Drive \(p\) and watch the curve swing from “amplifies errors” to “converges to certain”:

Code: Self-Consistency from Scratch

reasoning.py builds the whole pipeline as small, model-agnostic functions — they wrap any generator, so we can test them without a trained model. First, pull the answer out of a chain. Following GSM8K, the model marks its final answer with ####:

from reasoning import extract_answer

print(extract_answer("6 eggs/day * 7 days = 42 eggs.\n#### 42"))
print(extract_answer("...so there are 1,024 bytes. The answer is 1,024."))
print(extract_answer("I'm not sure."))   # nothing to extract
42
1024
None

Then tally a plurality vote — the most common answer wins, Nones ignored:

from reasoning import majority_vote, vote_distribution

answers = ["42", "36", "42", "42", "17", "42"]
print("distribution:", vote_distribution(answers))
print("vote winner: ", majority_vote(answers))
distribution: {'42': 4, '36': 1, '17': 1}
vote winner:  42

self_consistency puts them together: sample n chains, extract each answer, return the vote and the tally. Here we drive it with NoisyReasoner, a controllable stand-in whose per-sample accuracy is exactly p — no model needed to see the effect (swap in a temperature-sampled generate call for real use):

from reasoning import self_consistency, NoisyReasoner

gen = NoisyReasoner(correct="42", p=0.55, num_distractors=4, seed=0)

for n in (1, 5, 41):
    answer, dist = self_consistency(gen.sample, n)
    correct = "✓" if answer == "42" else "✗"
    print(f"n={n:>2}: voted {answer!r} {correct}   tally={dist}")
n= 1: voted '3' ✗   tally={'3': 1}
n= 5: voted '42' ✓   tally={'42': 2, '3': 1, '2': 2}
n=41: voted '42' ✓   tally={'1': 6, '42': 22, '2': 5, '3': 4, '0': 4}

A single sample is a coin-flip near p; by 41 samples the correct answer is the clear plurality even though the generator is right only 55% of the time. That is the phase transition, made concrete.

Watch Accuracy Climb with Compute

The headline of test-time compute: accuracy is a rising function of how many samples you draw. demonstrate_self_consistency runs thousands of problems at each sample count and reports the fraction solved — the empirical compute–accuracy curve — beside the Condorcet bound.

from reasoning import demonstrate_self_consistency

accuracy = demonstrate_self_consistency(p=0.55, num_distractors=4, trials=3000, seed=0)
============================================================
SELF-CONSISTENCY: accuracy vs. samples (test-time compute)
============================================================
  per-sample accuracy p = 0.55, distractors = 4, trials = 3000

     n   empirical   Condorcet
     1       0.560       0.550
     3       0.655       0.575
     5       0.779       0.593
    11       0.931       0.633
    21       0.992       0.679
    41       1.000       0.741

  More samples -> higher accuracy (since p > 1/2); the empirical
  curve beats the binary Condorcet bound because wrong answers scatter.
from reasoning import condorcet_majority_prob

# Bridge the empirical curve (and the theoretical bound) to the plot below.
sc_points = [
    {"n": n, "empirical": acc, "condorcet": condorcet_majority_prob(0.55, n)}
    for n, acc in accuracy.items()
]
ojs_define(sc_points = sc_points)
TipTry This

The measured curve (orange) sits above the Condorcet bound (blue). That gap is the scattering effect: with four distractors, the wrong 45% splits four ways (~11% each), so the correct 55% is the runaway plurality. Real reasoning benchmarks have many possible wrong answers, so self-consistency is even more forgiving than the binary theorem predicts.

Best-of-N and Verifiers

Voting assumes the right answer is common. But sometimes the model finds the right answer rarely — it just needs help recognizing it. Best-of-N samples N candidates and keeps the one a verifier (a learned scorer / reward model, m12) rates highest, so a single good sample can win even if it is outvoted:

from reasoning import best_of_n

# Toy verifier: prefers solutions that "show their work" (longer, with an '=').
def verifier(chain: str) -> float:
    return len(chain) + (5.0 if "=" in chain else 0.0)

candidates = iter([
    "#### 42",
    "6 * 7 = 42, so the total is 42.\n#### 42",
    "idk maybe 40\n#### 40",
])
best = best_of_n(lambda: next(candidates), score_fn=verifier, n=3)
print("verifier picked:\n", best)
verifier picked:
 6 * 7 = 42, so the total is 42.
#### 42

Voting and best-of-N are the two parallel ways to spend test-time compute: voting is a verifier-free majority; best-of-N trades the vote for a scorer that can spot a rare gem. Production systems blend them — e.g. weight each vote by its verifier score.

Chain-of-Thought & the Frontier

Two loose ends connect this from-scratch core to the models making headlines.

Chain-of-thought is the enabler. Self-consistency only helps if the sampled chains are diverse yet mostly-correct — which is exactly what CoT prompting produces. Asking for steps (“Let’s think step by step”) both raises per-sample accuracy \(p\) and creates the path diversity voting feeds on.

RL for reasoning is the sequential counterpart. Everything above is parallel test-time compute: draw independent samples, aggregate. The other axis is sequential — train the model (with RL against verifiable rewards) to produce one long, self-correcting chain that thinks longer on harder problems. That is the o1 / DeepSeek-R1 recipe; its GRPO objective is the natural from-scratch sequel to the DPO you built in m12 — and you build it next, below.

NoteKey Insight

There are two knobs for test-time compute: parallel (sample many, aggregate — this module) and sequential (one longer, self-correcting chain — RL-trained). Both trade inference FLOPs for accuracy, and both leave the pre-trained weights of the base model exactly where scaling laws (m07) left them.

The Sequential Branch: Training to Reason with GRPO

Self-consistency spends compute at inference on a frozen model. The reasoning models that top today’s benchmarks do something more: they are trained to spend that compute well — to emit one long, self-correcting chain and only then commit to an answer. The recipe that made this reproducible and open is GRPO (Group Relative Policy Optimization), introduced in DeepSeekMath (2024) and used to train DeepSeek-R1 (2025).

GRPO is reinforcement learning, and it is the direct sequel to the DPO you built in m12. Both push a policy toward better outputs using only a preference or reward signal, no labeled target text. The twist that makes GRPO special is how it estimates whether an answer was good.

The problem RL has to solve. To improve, the policy needs to know whether a sampled answer was better than expected. Classic PPO trains a second network — a critic / value model \(V(s)\), as big as the policy — just to predict that expected reward, so the “advantage” is \(r - V(s)\). That doubles the memory.

GRPO’s one idea: let the group be the baseline. For a prompt \(q\), sample a whole group of \(G\) answers \(o_1, \dots, o_G\) from the current policy, score each with a reward \(r_i\), and use the group’s own mean reward as the expectation. An answer that beat its group’s average gets pushed up; one below average gets pushed down. No critic, no value network — just the samples you already drew.

NoteKey Insight

The reward can be verifiable: for math or code you don’t need a learned reward model at all — you can check whether the final answer is correct. DeepSeek-R1 trained on exactly this (a correctness reward plus a format bonus), and R1-Zero reached frontier reasoning with no supervised reasoning traces at all — pure RL from a rule that grades the answer.

The Math: Group-Relative Advantage

Score the group, then turn raw rewards into advantages by standardizing within the group:

\[ A_i = \frac{r_i - \operatorname{mean}(r_1, \dots, r_G)}{\operatorname{std}(r_1, \dots, r_G) + \varepsilon}. \]

Subtracting the mean centers the signal (this is the critic-free baseline); dividing by the std puts every prompt on the same scale. A group that is all correct or all wrong has zero std → zero advantage — correctly, since there is nothing relative to learn from it.

Every token in output \(o_i\) shares that one scalar \(A_i\). The policy is then nudged with the same clipped surrogate PPO uses, with a KL leash to a frozen reference policy \(\pi_\text{ref}\):

\[ \mathcal{J}_\text{GRPO} = \frac{1}{G}\sum_{i=1}^{G} \min\!\Big(\rho_i A_i,\ \operatorname{clip}(\rho_i,\, 1-\varepsilon,\, 1+\varepsilon)\,A_i\Big) \;-\; \beta\, \mathbb{D}_\text{KL}\!\left(\pi_\theta \,\|\, \pi_\text{ref}\right), \qquad \rho_i = \frac{\pi_\theta(o_i \mid q)}{\pi_{\theta_\text{old}}(o_i \mid q)}. \]

Two guards keep the update honest. The clip is a trust region: once a sample’s probability has moved more than \(\varepsilon\) in the helpful direction, the clipped branch wins the \(\min\) and its gradient goes to zero, so one lucky group can’t yank the policy off a cliff. The KL penalty keeps \(\pi_\theta\) near the reference so it doesn’t forget how to write — or learn to reward-hack. GRPO uses Schulman’s unbiased, always-non-negative k3 estimator, \(\mathbb{D}_\text{KL} = e^{d} - d - 1\) with \(d = \log\pi_\text{ref} - \log\pi_\theta\).

Step by Step

Code: GRPO from Scratch

The whole algorithm is three small functions plus a loop. It lives in grpo.py; here we build the pieces inline. Start with the advantage — the one line that replaces PPO’s critic:

import torch
from grpo import group_relative_advantages, kl_divergence

# A group of 4 answers to one prompt; only the first was correct (verifiable reward).
rewards = torch.tensor([1.0, 0.0, 0.0, 0.0])
adv = group_relative_advantages(rewards)
print("rewards   :", rewards.tolist())
print("advantages:", [round(a, 3) for a in adv.tolist()])
print("mean ~ 0  :", round(float(adv.mean()), 4), " std ~ 1:", round(float(adv.std(unbiased=False)), 4))

# All-correct (or all-wrong) group -> no relative signal -> zero advantage.
print("all-correct advantages:", group_relative_advantages(torch.ones(4)).tolist())
rewards   : [1.0, 0.0, 0.0, 0.0]
advantages: [1.732, -0.577, -0.577, -0.577]
mean ~ 0  : 0.0  std ~ 1: 1.0
all-correct advantages: [0.0, 0.0, 0.0, 0.0]

The correct answer gets a positive advantage, the wrong ones negative, and an all-equal group yields zeros. Now the KL leash — the k3 estimator, non-negative by construction:

same = kl_divergence(torch.tensor(-1.0), torch.tensor(-1.0))   # policy == reference
moved = kl_divergence(torch.tensor(-2.0), torch.tensor(-1.0))  # policy drifted away
print(f"KL when unchanged: {float(same):.3f}")
print(f"KL after drifting: {float(moved):.3f}  (always >= 0)")
KL when unchanged: 0.000
KL after drifting: 0.718  (always >= 0)

Put them together with the clipped surrogate and you have the GRPO loss — exactly the objective loss.backward() differentiates. To keep it hand-followable we make each “answer” a single token: a categorical policy over K candidate answers, where an output’s log-prob is one entry of log_softmax. (For a real chain you sum the per-token log-probs of the generated sequence; the loss above is unchanged.)

from grpo import ToyReasoningPolicy, grpo_step

# A uniform policy over 6 possible answers; answer 3 is correct.
policy = ToyReasoningPolicy(num_answers=6, seed=0)
print(f"before: p(correct) = {float(policy.probabilities()[3]):.3f}  (chance = {1/6:.3f})")

for _ in range(30):
    grpo_step(policy, correct=3, group_size=32, lr=0.5)

print(f"after : p(correct) = {float(policy.probabilities()[3]):.3f}  "
      "(learned from a verifiable reward — no critic, no labels)")
before: p(correct) = 0.167  (chance = 0.167)
after : p(correct) = 0.993  (learned from a verifiable reward — no critic, no labels)

Each grpo_step samples a group from the current policy, grades it, standardizes the rewards into advantages, and takes a few clipped-surrogate gradient steps toward the above-average answers — the full loop in grpo.py.

Watch a Policy Learn to Reason

demonstrate_grpo runs that loop from a uniform start and records the curve. The policy climbs from chance (\(1/K\)) toward near-certainty on the correct answer, driven only by a reward that says right or wrong — a from-scratch miniature of the R1 recipe.

from grpo import demonstrate_grpo

grpo_hist = demonstrate_grpo(num_answers=6, correct=3, steps=40,
                             group_size=32, lr=0.5, seed=0, verbose=False)

# Bridge the learning curve to the interactive plot below.
ojs_define(grpo_curve = grpo_hist)
TipTry This
  1. Kill the reward signal. In the demonstrate_grpo call, set correct to an answer, then imagine every reward were identical — the advantages would all be zero and the curve would stay flat. This is why a group that is all right or all wrong teaches nothing.
  2. Shrink the group. Re-run with group_size=4 vs group_size=64. Smaller groups give noisier advantage estimates (the mean baseline is shakier), so the climb is bumpier — the same bias/variance trade a critic would smooth.
  3. Turn up the leash. Raise beta toward 1.0. The KL penalty fights the reward, slowing how fast the policy departs its uniform reference — a stronger leash trades learning speed for staying close to the starting model.

Common Pitfalls

When implementing test-time compute, watch out for:

  1. Voting with greedy decoding. Self-consistency needs stochastic sampling (temperature > 0, top-p/top-k from m08). Greedy makes every chain identical, so the “vote” is n copies of one answer — no gain.
  2. Voting when p < ½ and the answer space is small. With a single dominant wrong answer, more samples make you more confidently wrong. Voting helps only when the correct answer is the mode.
  3. Comparing raw strings. "42", "42.", and "1,024" vs "1024" must be normalized before tallying, or correct chains split their own vote. That is why extract_answer strips commas and pulls a clean number.
  4. Counting compute as free. N-sample self-consistency costs ~N× the FLOPs of one pass. The compute–accuracy curve has diminishing returns; past some N the next point of accuracy is not worth the tokens.
  5. Trusting a weak verifier in best-of-N. Best-of-N is only as good as its scorer — an exploitable verifier makes the model reward-hack (m12), picking samples that fool the scorer rather than solve the problem.
  6. A degenerate GRPO group teaches nothing. If every sampled answer earns the same reward (all correct, or all wrong), the group’s std is zero and every advantage is zero — no gradient. Reward design and a policy that isn’t already saturated matter: you need a spread of outcomes for the group baseline to have signal.
  7. Dropping GRPO’s KL leash. Without the \(\beta\,\mathbb{D}_\text{KL}\) term the policy is free to chase the reward straight into reward-hacking or degenerate text. The clipped surrogate bounds one step; the KL keeps the model near a sane reference across many steps.

Exercises

Exercise 1: The break-even sample count

from reasoning import condorcet_majority_prob

# For p = 0.55, how many samples n does the Condorcet bound need to first exceed
# 0.90? (Loop odd n and check.) Then confirm real scattered-distractor accuracy
# reaches it sooner.

# Your implementation here:

Exercise 2: Self-consistency vs. best-of-N

from reasoning import self_consistency, best_of_n, NoisyReasoner

# Build a generator that is correct only 30% of the time but whose correct
# chains are always the longest. Show best_of_n(score_fn=len) beats
# self_consistency at the same N. When does voting win instead?

# Your implementation here:

Exercise 3: Weighted voting

from collections import Counter

# Combine voting and verification: instead of one vote per chain, weight each
# chain's vote by a verifier score, then take the argmax. Implement
# weighted_vote(answers, scores) and compare to plain majority_vote.

# Your implementation here:

Exercise 4: GRPO with a shaped reward

from grpo import ToyReasoningPolicy, group_relative_advantages, grpo_step

# The verifiable_reward above is 0/1. Give partial credit instead: reward the
# correct answer 1.0, a "close" answer (index correct-1 or correct+1) 0.5, else 0.
# Write shaped_reward(answers, correct, num_answers) -> tensor, then run GRPO with
# it (adapt grpo_step, or roll your own loop). Does the policy still converge on
# the exactly-correct answer, or does partial credit slow it down?

# Your implementation here:

Summary

Key takeaways:

  1. Test-time compute is a third axis. Beyond bigger models and more data (m07 scaling laws), you can spend more compute at inference to raise accuracy — with the weights frozen.
  2. Self-consistency is sample-and-vote. Draw many chain-of-thought samples, extract each answer, take the plurality. It lifted GSM8K by +17.9 points with no new parameters.
  3. Condorcet explains it. If each sample is correct with probability \(p > \tfrac12\), the majority of \(n\) is correct with probability \(\to 1\); below \(\tfrac12\) voting backfires. The real requirement is weaker — the correct answer just has to be the mode, since wrong answers scatter.
  4. Accuracy rises with samples. The compute–accuracy curve climbs (with diminishing returns), and empirically beats the binary Condorcet bound.
  5. Best-of-N uses a verifier instead of a vote, so it can surface a rare correct answer — as good as its scorer, and vulnerable to reward hacking.
  6. Parallel vs. sequential. Sampling-and-aggregating is the parallel branch; RL-trained long self-correcting chains (o1 / R1) are the sequential branch. Both are test-time compute.
  7. GRPO trains reasoning without a critic. Sample a group of answers, use their mean reward as the baseline (\(A_i = (r_i - \text{mean})/\text{std}\)), and nudge the policy with a clipped surrogate plus a KL leash. With a verifiable reward (is the answer correct?) it needs no reward model and no labeled traces — the DeepSeek-R1 recipe, built from scratch in grpo.py.

What’s Next

You now have both halves of the reasoning story: the parallel branch (sample-and-aggregate) and the sequential branch (GRPO — training a model to reason from verifiable rewards, no critic). From here the book turns to serving these compute-hungry models efficiently — Module 14: Quantization shrinks the weights, and speculative decoding speeds the decode — and to Module 15: Retrieval-Augmented Generation and agents that put reasoning to work on live knowledge.

Going Deeper

Core Papers:

Practical Resources: