RL for LLM10 / 13

Module 10. Search, verification, and test-time compute

Read the explanations and code in the existing notebooks, or download the files. Code does not run on this site.

Download course files (.zip) 382.1 KB

The archive preserves the source folder structure, both languages, and supporting files.

Module 10. Test-Time Compute: Sampling, Verification, and Search

Main module notebook

Download .ipynb module-10-test-time-search.en.ipynb · 19.6 KB

Show notebook text

Module 10. Test-Time Compute: Sampling, Verification, and Search

This notebook connects the module’s mathematics to a reproducible CPU laboratory. Coverage, selection quality, verifier uncertainty, proxy pressure, and search efficiency are measured separately. The core route needs neither network access nor a GPU; the optional real-Qwen3 branch is disabled and does not contribute to the release results.

0. Reproducible setup

smoke executes the complete required route. full increases the number of independent runs. The code locates the unpacked course root automatically.

Python

from __future__ import annotations
from pathlib import Path
import sys
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from IPython.display import display

PROFILE = "smoke"  # "smoke" or "full"
RUN_REAL_MODEL = False
SEEDS = range(40 if PROFILE == "smoke" else 100)

ROOT = Path.cwd()
if not (ROOT / "modules").exists():
    if (ROOT.parent / "modules").exists():
        ROOT = ROOT.parent
    else:
        raise RuntimeError("Run the notebook from the course root or notebooks directory")
MODULE = ROOT / "modules" / "module-10"
sys.path.insert(0, str(MODULE))

from module_10_reference import (
    pass_at_n, pass_curve, unbiased_pass_at_k,
    plurality_accuracy_exact, plurality_accuracy_mc,
    binary_verifier_success, process_score, prefix_success_estimate,
    wilson_interval, proxy_overoptimization_curve,
    allocate_budget_greedy, mean_coverage, pareto_frontier,
    break_even_queries, cost_regularized_utility,
)
from module_10_search import (
    ReasoningMaze, repeated_sampling, best_of_n_proxy, beam_search,
    best_first_search, mcts_search, evaluate_search_methods, maze_to_array,
)

LABELS = {'coverage': 'probability of at least one success', 'budget': 'budget N', 'exact': 'exact', 'mc': 'Monte Carlo', 'vote': 'voting accuracy', 'proxy': 'proxy score', 'goal': 'independent objective', 'score': 'score', 'success': 'successful-run rate', 'calls': 'model calls', 'good': 'good scorer', 'biased': 'biased scorer', 'allocation': 'allocated samples', 'prompt': 'prompt', 'quality': 'quality', 'cost': 'cost', 'skipped': 'The real-model branch is disabled. Enable it only in a prepared environment with model access.'}
np.set_printoptions(precision=4, suppress=True)
print(f"profile={PROFILE}, independent seeds={len(list(SEEDS))}")

1. pass@N coverage and finite-sample pass@k

The formula $1-(1-p)^N$ describes the presence of at least one success under conditional independence and constant success probability $p$ . We also compute the unbiased finite-sample pass@k estimate for an already generated pool.

Python

budgets = np.array([1, 2, 4, 8, 16, 32, 64])
fig, ax = plt.subplots(figsize=(8.8, 5.0))
for p in (0.0, 0.05, 0.2, 0.5):
    ax.plot(budgets, pass_curve(p, budgets), marker="o", label=f"p={p:g}")
ax.set_xscale("log", base=2)
ax.set_xticks(budgets, labels=budgets)
ax.set_xlabel(LABELS["budget"])
ax.set_ylabel(LABELS["coverage"])
ax.set_ylim(-0.02, 1.02)
ax.grid(True, alpha=0.25)
ax.legend()
plt.show()

finite = pd.DataFrame({
    "n": [100, 100, 100, 20],
    "c": [1, 5, 20, 3],
    "k": [1, 10, 10, 5],
})
finite["pass@k"] = [unbiased_pass_at_k(int(r.n), int(r.c), int(r.k)) for r in finite.itertuples()]
display(finite.round(6))
assert np.isclose(pass_at_n(0.2, 8), 1 - 0.8**8)

2. Voting and ties

Voting is useful only when the correct answer is the largest mode. Ties are split uniformly here, avoiding the hidden positional bias of a plain argmax.

Python

cases = {
    "scattered_errors": np.array([0.40, 0.15, 0.15, 0.15, 0.15]),
    "systematic_error": np.array([0.30, 0.45, 0.25]),
}
rows = []
for name, probs in cases.items():
    for k in ((1, 4, 8, 16) if len(probs) == 5 else (1, 8, 32, 64)):
        exact = plurality_accuracy_exact(probs, k)
        mc = plurality_accuracy_mc(probs, k, trials=20_000 if PROFILE == "smoke" else 80_000, seed=10+k)
        rows.append({"case": name, "k": k, "exact": exact, "Monte Carlo": mc})
votes = pd.DataFrame(rows)
display(votes.round(6))

fig, ax = plt.subplots(figsize=(8.6, 4.8))
for name, group in votes.groupby("case"):
    ax.plot(group["k"], group["exact"], marker="o", label=name)
ax.set_xscale("log", base=2)
ax.set_xlabel("k")
ax.set_ylabel(LABELS["vote"])
ax.set_ylim(0, 1.02)
ax.grid(True, alpha=0.25)
ax.legend()
plt.show()
assert abs(votes["exact"].iloc[0] - 0.4) < 1e-12

3. Binary verification and the fallback protocol

When no candidate is accepted, finite-budget accuracy depends on whether the system selects from the already rejected pool or draws a fresh unverified candidate. The two protocols require different formulas.

Python

p, alpha, beta = 0.2, 0.9, 0.1
ns = np.arange(1, 65)
rejected = np.array([
    binary_verifier_success(int(n), p, alpha, beta, fallback_mode="rejected_pool").success_probability
    for n in ns
])
fresh = np.array([
    binary_verifier_success(int(n), p, alpha, beta, fallback_mode="fresh_unverified").success_probability
    for n in ns
])
example_rej = binary_verifier_success(8, p, alpha, beta, fallback_mode="rejected_pool")
example_fresh = binary_verifier_success(8, p, alpha, beta, fallback_mode="fresh_unverified")
display(pd.DataFrame([example_rej.__dict__, example_fresh.__dict__]).round(6))

fig, ax = plt.subplots(figsize=(8.7, 4.8))
ax.plot(ns, rejected, label="fallback: rejected pool")
ax.plot(ns, fresh, label="fallback: fresh draw")
ax.axhline(example_rej.accepted_precision, linestyle="--", label="accepted-pool precision")
ax.set_xscale("log", base=2)
ax.set_xlabel(LABELS["budget"])
ax.set_ylabel(LABELS["success"])
ax.grid(True, alpha=0.25)
ax.legend()
plt.show()
assert np.isclose(example_rej.success_probability, 0.632486004809728)
assert np.isclose(example_fresh.success_probability, 0.6480396435591987)

4. Proxy overoptimization under Best-of-N

A larger pool improves the selected proxy score, while an independent objective may peak and then decline. This is a course-owned synthetic mechanism, not an external benchmark.

Python

budgets_proxy = np.array([1, 2, 4, 8, 16, 32, 64, 128, 256, 512])
proxy, independent = proxy_overoptimization_curve(
    budgets_proxy, trials=4_000 if PROFILE == "smoke" else 20_000, seed=202_610
)
proxy_table = pd.DataFrame({"N": budgets_proxy, "proxy": proxy, "independent_objective": independent})
display(proxy_table.round(4))

fig, ax = plt.subplots(figsize=(8.8, 4.9))
ax.plot(budgets_proxy, proxy, marker="o", label=LABELS["proxy"])
ax.plot(budgets_proxy, independent, marker="o", label=LABELS["goal"])
ax.set_xscale("log", base=2)
ax.set_xlabel(LABELS["budget"])
ax.set_ylabel(LABELS["score"])
ax.grid(True, alpha=0.25)
ax.legend()
plt.show()
peak_n = int(budgets_proxy[np.argmax(independent)])
print("independent-objective peak N =", peak_n)

5. Process scores and prefix uncertainty

A prefix label estimated by Monte Carlo continuations has finite-sample uncertainty. We compare several step-score aggregators and Wilson intervals.

Python

successes = np.arange(9)
rows = []
for s in successes:
    lo, hi = wilson_interval(int(s), 8)
    rows.append({"successes": s, "estimate": prefix_success_estimate(int(s), 8), "Wilson_low": lo, "Wilson_high": hi})
prefix_table = pd.DataFrame(rows)
display(prefix_table.round(4))

fig, ax = plt.subplots(figsize=(8.6, 4.8))
mid = prefix_table["estimate"].to_numpy()
lo = prefix_table["Wilson_low"].to_numpy()
hi = prefix_table["Wilson_high"].to_numpy()
ax.errorbar(successes, mid, yerr=np.vstack([mid-lo, hi-mid]), fmt="o", capsize=4)
ax.set_xlabel("successes out of 8")
ax.set_ylabel("prefix-success estimate and 95% Wilson interval")
ax.set_ylim(-0.03, 1.03)
ax.grid(True, alpha=0.25)
plt.show()

step_scores = [0.92, 0.78, 0.61, 0.85]
aggregation = {m: process_score(step_scores, m) for m in ("min", "mean", "product", "last", "log_product")}
display(pd.Series(aggregation, name="aggregate").to_frame())

6. Reasoning maze: environment and individual trajectories

The maze is a controlled analogue of a reasoning tree. The proposal model is frozen while the answer algorithm changes: repeated sampling, Best-of-N, beam, best-first, and MCTS spend the same kind of budget differently.

Python

env = ReasoningMaze.course_maze()
methods = {
    "repeated": repeated_sampling(env, 48, seed=7),
    "best_of_n_proxy": best_of_n_proxy(env, 48, seed=7),
    "beam": beam_search(env, 48, beam_width=3, score_mode="good", seed=7),
    "best_first": best_first_search(env, 48, score_mode="good", seed=7),
    "mcts": mcts_search(env, 48, score_mode="good", seed=7),
}
summary = pd.DataFrame([
    {"method": name, "success": r.success, "model_calls": r.model_calls,
     "generated_trajectories": r.generated_trajectories, "path_length": r.selected_state.depth}
    for name, r in methods.items()
])
display(summary)

chosen = methods["best_first"]
grid = maze_to_array(env, chosen.selected_state)
fig, ax = plt.subplots(figsize=(6.3, 5.5))
ax.imshow(grid)
for i in range(env.height):
    for j in range(env.width):
        if (i, j) not in env.walls:
            ax.text(j, i, "·", ha="center", va="center", fontsize=10)
ax.set_title(f"best-first: success={chosen.success}, calls={chosen.model_calls}")
ax.set_xticks(range(env.width)); ax.set_yticks(range(env.height))
plt.show()

7. Search curves with a good and a biased scorer

The same search procedure may improve rapidly under an informative process score and fail under a systematically attractive false branch. Search quality therefore cannot be separated from scorer quality.

Python

search_budgets = np.array([12, 16, 24, 32, 48])
good = evaluate_search_methods(search_budgets, seeds=SEEDS, score_mode="good")
biased = evaluate_search_methods(search_budgets, seeds=SEEDS, score_mode="biased")

fig, axes = plt.subplots(1, 2, figsize=(13.8, 4.9), sharey=True)
for ax, curves, title in zip(axes, (good, biased), (LABELS["good"], LABELS["biased"])):
    for method, values in curves.items():
        ax.plot(search_budgets, values, marker="o", label=method)
    ax.set_title(title)
    ax.set_xlabel(LABELS["calls"])
    ax.grid(True, alpha=0.25)
axes[0].set_ylabel(LABELS["success"])
axes[0].legend(fontsize=8)
plt.show()

comparison = pd.DataFrame({
    ("good", method): values for method, values in good.items()
} | {
    ("biased", method): values for method, values in biased.items()
}, index=search_budgets)
comparison.index.name = "model_calls"
display(comparison.round(3))

8. Adaptive allocation and the Pareto frontier

A uniform quota is simple but need not maximize mean coverage. We compare it with greedy marginal-gain allocation, then identify non-dominated quality–cost points.

Python

ps = np.array([0.02, 0.08, 0.20, 0.50])
total = 20
uniform = np.full(len(ps), total // len(ps), dtype=int)
greedy = allocate_budget_greedy(ps, total, minimum_per_prompt=1)
allocation = pd.DataFrame({"p": ps, "uniform": uniform, "greedy": greedy})
display(allocation)
print("uniform mean coverage:", round(mean_coverage(ps, uniform), 6))
print("greedy mean coverage:", round(mean_coverage(ps, greedy), 6))

costs = np.array([8, 12, 16, 24, 32, 48, 64], dtype=float)
quality = np.array([0.42, 0.55, 0.69, 0.76, 0.79, 0.80, 0.795])
front = pareto_frontier(quality, costs)
pareto_table = pd.DataFrame({"cost": costs, "quality": quality, "Pareto": front})
pareto_table["utility(lambda=0.004)"] = [cost_regularized_utility(q, c, 0.004) for q, c in zip(quality, costs)]
display(pareto_table)
print("break-even queries:", break_even_queries(12_000, 0.12, 0.07))

fig, ax = plt.subplots(figsize=(7.6, 4.8))
ax.scatter(costs[~front], quality[~front], label="dominated")
ax.scatter(costs[front], quality[front], label="Pareto")
ax.plot(costs[front], quality[front], linestyle="--")
ax.set_xlabel(LABELS["cost"]); ax.set_ylabel(LABELS["quality"])
ax.grid(True, alpha=0.25); ax.legend(); plt.show()

9. Optional real-Qwen3 branch

The template below performs real multi-candidate generation and extracts normalized transition log-probabilities. It was not executed in the release gate and supports no claims about quality, latency, or memory.

Python

if not RUN_REAL_MODEL:
    print('The real-model branch is disabled. Enable it only in a prepared environment with model access.')
else:
    # Install Transformers in a dedicated real-model environment and record
    # the actually resolved version before a reproducible run.
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed

    model_id = "Qwen/Qwen3-0.6B"
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype="auto",
        device_map="auto",
    )
    prompt = "Solve carefully and end with a boxed integer: 17 * 23"
    messages = [{"role": "user", "content": prompt}]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer([text], return_tensors="pt").to(model.device)
    set_seed(10)
    generated = model.generate(
        **inputs,
        do_sample=True,
        temperature=0.8,
        top_p=0.95,
        max_new_tokens=192,
        num_return_sequences=8,
        return_dict_in_generate=True,
        output_scores=True,
    )
    transition_scores = model.compute_transition_scores(
        generated.sequences,
        generated.scores,
        normalize_logits=True,
    )
    prompt_len = inputs.input_ids.shape[1]
    responses = tokenizer.batch_decode(generated.sequences[:, prompt_len:], skip_special_tokens=True)
    rows = []
    for i, response in enumerate(responses):
        rows.append({
            "candidate": i,
            "response": response,
            "sum_logprob": float(transition_scores[i].sum().cpu()),
            "mean_logprob": float(transition_scores[i].mean().cpu()),
        })
    display(pd.DataFrame(rows))