arman1o1 blog
AI/ML Reasoning Test-Time Compute PRM & MCTS PyTorch

Test-Time Compute & Reasoning Verification: PRMs, MCTS & Search Scaling Laws

• ☕ 14 min read

For the first decade of deep learning, scaling laws meant one thing: spend more FLOPs during pre-training ($\mathcal{C}_{\text{train}} \approx 6 N D$). Yet human cognition does not allocate equal time to every question – you answer "What is $2+2$?" in $200\text{ ms}$, while spending $20\text{ minutes}$ working through an Olympiad proof. Standard autoregressive decoding forces an LLM to commit to every token in a constant $\mathcal{O}(1)$ forward pass regardless of logical difficulty.

Test-Time Compute Scaling (Snell et al., 2024; Lightman et al., 2023) introduces a second orthogonal scaling axis: allocating adaptive inference compute $\mathcal{C}_{\text{test}}$ via Process Reward Models (PRMs), Best-of-$N$ weighted verification, and Monte Carlo Tree Search (MCTS). On verifiable reasoning benchmarks, a smaller model equipped with a PRM and tree search routinely outperforms a 14× larger model running greedy decoding at a fraction of total FLOPs.

0. The Two Axes of Scaling: Pre-Training vs. Test-Time Search

Suppose a base policy $\pi_\theta(y \mid x)$ has a single-sample pass rate $p = \mathbb{P}(\text{correct} \mid x)$ on a reasoning problem. If we draw $N$ independent trajectories $y^{(1)}, \dots, y^{(N)} \sim \pi_\theta(\cdot \mid x)$ and possess an ideal verifier, the probability of finding at least one valid proof is the Pass@$N$ scaling curve:

Pass@$N$ Coverage vs. Unbiased Empirical Estimator
$$\text{Pass@}N = 1 - (1 - p)^N, \qquad \widehat{\text{Pass@}k} = 1 - \frac{\binom{n - c}{k}}{\binom{n}{k}}$$

where $n \ge k$ total samples are drawn and $c$ are correct (Chen et al., 2021). Even if a single rollout only succeeds with $p = 0.05$, sampling $N = 64$ trajectories raises coverage to $1 - 0.95^{64} = \mathbf{96.3\%}$! The central engineering challenge of test-time compute is bridging the gap between Pass@$N$ (oracle selection) and Pass@$1$ (verifier selection) without ground-truth labels at inference time.

1. Outcome Reward Models (ORM) vs. Process Reward Models (PRM)

A reasoning trajectory $y = (s_1, s_2, \dots, s_K)$ consists of $K$ discrete logical steps delimited by newline or step tokens (e.g., `\n\n`). How we supervise the verifier determines whether search can prune bad branches early:

  • Outcome Reward Model (ORM): Predicts a single scalar probability $R_{\text{ORM}}(x, y) \in [0, 1]$ only at the final token of the complete trajectory $y$. If a 12-step mathematical proof makes a sign error at step 3 but arrives at a plausible-looking expression at step 12, the ORM suffers from coarse credit assignment and false-positive hallucinations.
  • Process Reward Model (PRM): Evaluates every prefix step $s_{1:k}$, predicting $P_\psi(z_k = 1 \mid x, s_{1:k}) \in [0, 1]$ at the end-of-step token of each step $k \in \{1, \dots, K\}$, where $z_k \in \{0, 1\}$ indicates whether step $s_k$ is both logically valid and remains on a solvable path (Lightman et al., 2023; Wang et al., 2024).
Step-Level PRM Objective & Trajectory Aggregation
$$\mathcal{L}_{\text{PRM}}(\psi) = -\frac{1}{K}\sum_{k=1}^{K} \Bigl[ z_k \log r_\psi(x, s_{1:k}) + (1 - z_k) \log\!\bigl(1 - r_\psi(x, s_{1:k})\bigr) \Bigr]$$ $$\mathcal{V}_{\text{prod}}(x, y) = \prod_{k=1}^{K} r_\psi(x, s_{1:k}), \qquad \mathcal{V}_{\min}(x, y) = \min_{1 \le k \le K} r_\psi(x, s_{1:k})$$

In practice, $\mathcal{V}_{\min}(x, y)$ ("chain is only as strong as its weakest step") or geometric mean $\bigl(\prod_{k=1}^K r_\psi(x, s_{1:k})\bigr)^{1/K}$ avoids penalizing longer multi-step proofs.

Automated PRM Labeling via Monte Carlo Rollouts (Math-Shepherd)

Human step-level annotation (PRM800K) is prohibitively expensive. Wang et al. (2024) showed that we can label any intermediate step $s_k$ automatically by running $M$ fast rollout completions $y_{k+1:K}^{(m)} \sim \pi_\theta(\cdot \mid x, s_{1:k})$ from step $s_k$ and checking the final answer against a deterministic verifier $V_{\text{ans}}$:

Monte Carlo Hard & Soft Step Targets
$$z_k^{\text{hard}} = \mathbb{I}\!\left(\exists m \in \{1,\dots,M\} : V_{\text{ans}}\!\left(y_{k+1:K}^{(m)}\right) = 1\right), \qquad z_k^{\text{soft}} = \frac{1}{M}\sum_{m=1}^{M} V_{\text{ans}}\!\left(y_{k+1:K}^{(m)}\right)$$
Outcome Reward Model (ORM): Sparse Terminal Supervision Only Step 1: Factor x²−4 No step score Step 2: (x−2)(x−2) Silent sign bug! Step 3: Solve x = 2 Wasted FLOPs ORM Score: r = 0.12 Where did it fail? Unknown. Process Reward Model (PRM): Dense Step-by-Step Verification & Early Pruning Step 1: Factor x²−4 ✓ r₁ = 0.98 (Keep) Step 2: (x−2)(x−2) ✗ r₂ = 0.04 (PRUNE!) Alt Step 2: (x−2)(x+2) ✓ r₂' = 0.97 (Expand) Step 3: x ∈ {−2, 2} ✓ r₃ = 0.99 (Solved)
Figure 1: Outcome Reward Model (ORM) vs. Process Reward Model (PRM). PRM pinpoints the exact reasoning step where an algebraic error occurs ($r_2 = 0.04$), allowing tree/beam search to prune the dead branch immediately and reallocate compute.
💻 PyTorch Implementation: Token-Masked Process Reward Model (PRM) & AggregatorPure PyTorch 
import torch
import torch.nn as nn
import torch.nn.functional as F

class ProcessRewardHead(nn.Module):
    """
    Projects hidden states [B, S, d_model] into scalar step-validity logits at
    designated step-boundary tokens (where step_labels != -100).
    """
    def __init__(self, d_model: int, dropout: float = 0.1):
        super().__init__()
        self.head = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(d_model, d_model // 2),
            nn.GELU(),
            nn.Linear(d_model // 2, 1)
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.head(hidden_states).squeeze(-1)  # [B, S] raw step logits

def prm_step_loss(
    step_logits: torch.Tensor,  # [B, S]
    step_labels: torch.Tensor   # [B, S] with {0.0, 1.0} at step delimiter tokens, -100 elsewhere
) -> torch.Tensor:
    valid_mask = step_labels != -100
    active_logits = step_logits[valid_mask]
    active_targets = step_labels[valid_mask].float()
    return F.binary_cross_entropy_with_logits(active_logits, active_targets, reduction="mean")

def aggregate_prm_trajectory_score(
    step_logits: torch.Tensor,
    step_mask: torch.Tensor,
    reduction: str = "min"
) -> torch.Tensor:
    """
    Aggregates per-step PRM probabilities across a trajectory using 'min', 'prod', or 'geomean'.
    step_logits: [N_trajectories, Max_Steps], step_mask: [N_trajectories, Max_Steps] boolean.
    """
    probs = torch.sigmoid(step_logits)
    if reduction == "min":
        masked_probs = probs.masked_fill(~step_mask, 1.0)
        return masked_probs.min(dim=-1).values
    elif reduction == "prod":
        log_probs = F.logsigmoid(step_logits).masked_fill(~step_mask, 0.0)
        return torch.exp(log_probs.sum(dim=-1))
    else:  # geomean
        log_probs = F.logsigmoid(step_logits).masked_fill(~step_mask, 0.0)
        counts = step_mask.sum(dim=-1).clamp(min=1).float()
        return torch.exp(log_probs.sum(dim=-1) / counts)

2. Best-of-$N$ Verification vs. Weighted Majority Voting

Given $N$ sampled trajectories $\{y^{(i)}\}_{i=1}^N$ with extracted final answers $a^{(i)} = \text{ans}(y^{(i)})$ and verifier scores $v_i = \mathcal{V}(x, y^{(i)}) \in (0, 1)$, there are three distinct ways to select the final prediction:

Unweighted Majority, Best-of-$N$, and Verifier-Weighted Majority
$$\hat{a}_{\text{maj}} = \arg\max_{a} \sum_{i=1}^{N} \mathbb{I}(a^{(i)} = a), \qquad \hat{a}_{\text{BoN}} = a^{(i^*)}\;\text{where } i^* = \arg\max_i v_i$$ $$\hat{a}_{\text{weighted}} = \arg\max_{a} \sum_{i : a^{(i)} = a} \log \frac{v_i}{1 - v_i} \quad (\text{Log-Odds Naive Bayes Optimal Aggregator})$$

Why log-odds $\log \frac{v_i}{1 - v_i}$? Under conditional independence of verifier noise across traces sharing answer $a$, the posterior log-likelihood ratio updates additively by the log-odds of each verifier score, preventing a single overconfident outlier trace ($v_i = 0.999$, often a reward-hacked trajectory) from overriding 15 consistent proofs with $v_j = 0.92$.

3. Monte Carlo Tree Search (MCTS) over Reasoning Steps

Best-of-$N$ still generates all $N$ trajectories from start to finish independently. When reasoning chains span $K = 20$ steps, Step-Level MCTS treats each reasoning step $s_k$ as an action in a search tree, balancing exploitation of high-value prefixes $Q(s, a)$ against exploration of under-visited branches via the PUCT (Predictor + Upper Confidence Bound applied to Trees) formula from AlphaZero:

PUCT Action Selection & Backpropagation Update
$$a^* = \arg\max_{a \in \mathcal{A}(s)} \left[ Q(s, a) + c_{\text{puct}} \, \pi_\theta(a \mid s) \, \frac{\sqrt{\sum_{b} N(s, b)}}{1 + N(s, a)} \right]$$ $$N(s, a) \leftarrow N(s, a) + 1, \qquad Q(s, a) \leftarrow Q(s, a) + \frac{v_{\text{leaf}} - Q(s, a)}{N(s, a)}$$

where $\pi_\theta(a \mid s)$ is the length-normalized policy prior of generating candidate step $a$ from state $s$, $N(s, a)$ is the visit count, and $v_{\text{leaf}} = r_\psi(x, s_{1:k})$ is the PRM evaluation of the newly expanded node.

1. Selection (PUCT) s₀ Q=.4 Q=.8 Follows argmax(Q + U) to promising leaf node 2. Expansion (π_θ) s₁ Sample B_branch candidate next steps s₂^(1..B) ~ π_θ 3. PRM Verification s₂* v = r_ψ(x, s₁:₂) = 0.94 Score step prefix directly without full rollout to end! 4. Backprop Update N+1 N+1 v=.94 Q(s,a) ← Q + (v−Q)/N Updates all ancestors
Figure 2: Step-level Monte Carlo Tree Search (MCTS) guided by a Process Reward Model (PRM). Instead of running full Monte Carlo rollouts to terminal tokens at every node, the PRM directly provides value estimates $v = r_\psi(x, s_{1:k})$ at intermediate steps.
💻 PyTorch Implementation: PUCT Reasoning Tree Search & Log-Odds Verifier VotingPure PyTorch 
import math
from dataclasses import dataclass, field
import torch

@dataclass
class ReasoningNode:
    state_text: str
    parent: "ReasoningNode | None" = None
    prior: float = 1.0
    visit_count: int = 0
    value_sum: float = 0.0
    prm_score: float = 0.0
    children: list["ReasoningNode"] = field(default_factory=list)
    is_terminal: bool = False
    final_answer: str | None = None

    @property
    def q_value(self) -> float:
        return self.value_sum / self.visit_count if self.visit_count > 0 else 0.0

    def select_child_puct(self, c_puct: float = 1.25) -> "ReasoningNode":
        total_sqrt = math.sqrt(max(1, sum(child.visit_count for child in self.children)))
        best_score = -1e9
        best_child = self.children[0]
        for child in self.children:
            u_score = c_puct * child.prior * (total_sqrt / (1.0 + child.visit_count))
            puct = child.q_value + u_score
            if puct > best_score:
                best_score = puct
                best_child = child
        return best_child

def mcts_reasoning_search(
    root: ReasoningNode,
    expand_fn,   # Callable[[ReasoningNode], list[ReasoningNode]] returning children with prior & prm_score
    num_simulations: int = 32,
    c_puct: float = 1.25
) -> ReasoningNode:
    """Runs PUCT-guided MCTS over reasoning steps and returns the highest-value terminal or leaf node."""
    for _ in range(num_simulations):
        node = root
        # 1. Selection
        while node.children and not node.is_terminal:
            node = node.select_child_puct(c_puct=c_puct)

        # 2. Expansion & 3. PRM Evaluation
        if not node.is_terminal:
            new_children = expand_fn(node)
            node.children.extend(new_children)
            if new_children:
                node = max(new_children, key=lambda c: c.prm_score)

        leaf_value = node.prm_score

        # 4. Backpropagation
        cur = node
        while cur is not None:
            cur.visit_count += 1
            cur.value_sum += leaf_value
            cur = cur.parent

    # Select best visited path
    best = root
    while best.children:
        best = max(best.children, key=lambda c: (c.visit_count, c.q_value))
    return best

def verifier_weighted_majority_vote(answers: list[str], prm_scores: torch.Tensor) -> str:
    """Aggregates candidate answers using log-odds weighting log(v / (1 - v))."""
    probs = prm_scores.clamp(min=1e-4, max=1.0 - 1e-4)
    log_odds = torch.log(probs / (1.0 - probs))
    score_table: dict[str, float] = {}
    for ans, weight in zip(answers, log_odds.tolist()):
        score_table[ans] = score_table.get(ans, 0.0) + weight
    return max(score_table.items(), key=lambda kv: kv[1])[0]

4. Failure Modes & Compute Allocation Diagnostics

1. Verifier Overoptimisation (Goodhart's Law at $N \gt 256$): As Best-of-$N$ sample count $N$ scales past $256$, the maximum verifier score $\max_i r_\psi(x, y^{(i)})$ continues to rise monotonically, but actual ground-truth accuracy peaks and then degrades! Why? Because the tail of $N$ stochastic samples discovers adversarial trajectories that exploit calibration blind spots in $r_\psi$.
Fix: Switch from pure `argmax` Best-of-$N$ to Verifier-Weighted Majority Voting (which averages out adversarial singletons across equivalence classes of final answers) or ensemble 2 orthogonal verifiers.
2. Step-Splitting Manipulation in Product PRMs: If trajectory value is computed as $\prod_{k=1}^K r_k$, the search policy learns to merge all reasoning into $K=1$ giant step so it only pays a single multiplication penalty. Conversely, if value is $\sum_{k=1}^K r_k$, the policy emits dozens of trivial tautologies ("Note that $1=1$...") to farm positive step rewards.
Fix: Always use $\min_k r_k$ or last-step cumulative PRM score $r_K$, which are invariant to redundant step splitting.
Easy Prompts ($p \gt 0.5$)
Parallel **Self-Consistency / Majority Voting** ($N = 4\text{–}8$) is compute-optimal; tree search adds unnecessary overhead.
Medium Prompts ($p \in [0.05, 0.5]$)
**PRM-Guided Beam / MCTS Search** achieves the steepest accuracy gain per FLOP, matching a $10\times$ larger pre-trained model.
Hard Prompts ($p \lt 0.01$)
Parallel sampling fails when the base model lacks the primitives; requires **Sequential Revision / RL-trained System-2 CoT** (DeepSeek-R1 style) to alter the proposal distribution itself.

5. Runnable End-to-End Verification: PRM + MCTS vs. Greedy & Majority Vote

The runnable script below (executed alongside the implementations defined in Sections 1–3) tests `ProcessRewardHead`, `prm_step_loss`, `aggregate_prm_trajectory_score`, `verifier_weighted_majority_vote`, and `mcts_reasoning_search` on a multi-step symbolic derivation tree where a high-prior distractor branch ("trap step") fools greedy decoding, demonstrating how PRM-guided MCTS recovers the true derivation.

🧪 Runnable Verification Script: PRM Training Step, Weighted Voting & MCTS RecoveryVerified PyTorch 
# Requires ProcessRewardHead, prm_step_loss, ReasoningNode,
# mcts_reasoning_search, and verifier_weighted_majority_vote from above.
import torch

def verify_test_time_compute():
    torch.manual_seed(101)

    # 1. Verify PRM Head & Masked Step Binary Cross-Entropy
    B, S, d_model = 2, 12, 64
    hidden = torch.randn(B, S, d_model, requires_grad=True)
    step_labels = torch.full((B, S), -100.0)
    step_labels[0, [3, 7, 11]] = torch.tensor([1.0, 1.0, 0.0])  # Step 3 failed
    step_labels[1, [4, 9, 11]] = torch.tensor([1.0, 1.0, 1.0])  # All steps valid

    prm = ProcessRewardHead(d_model=d_model)
    logits = prm(hidden)
    loss = prm_step_loss(logits, step_labels)
    loss.backward()
    assert hidden.grad is not None and loss.item() > 0.0

    # 2. Verify Verifier-Weighted Majority Vote vs Raw Majority Vote
    # Suppose 3 noisy traces say "x=4" with low PRM scores (0.15), and 2 solid traces say "x=2" with high PRM (0.95)
    candidate_answers = ["x=4", "x=4", "x=4", "x=2", "x=2"]
    prm_scores = torch.tensor([0.15, 0.12, 0.18, 0.95, 0.96])
    weighted_winner = verifier_weighted_majority_vote(candidate_answers, prm_scores)
    assert weighted_winner == "x=2", f"Expected x=2, got {weighted_winner}"

    # 3. Verify MCTS recovers correct multi-step derivation when Greedy prior picks a trap
    root = ReasoningNode(state_text="Solve x^2 - 5x + 6 = 0", prm_score=1.0)

    def toy_expand(node: ReasoningNode) -> list[ReasoningNode]:
        if node.state_text == "Solve x^2 - 5x + 6 = 0":
            # Greedy prior prefers trap branch (0.65), but PRM catches sign error (0.12 vs 0.94)
            trap = ReasoningNode("Step 1A: (x - 6)(x + 1) = 0", parent=node, prior=0.65, prm_score=0.12)
            valid = ReasoningNode("Step 1B: (x - 2)(x - 3) = 0", parent=node, prior=0.35, prm_score=0.94)
            return [trap, valid]
        elif "Step 1A" in node.state_text:
            return [ReasoningNode("Ans: x in {6, -1}", parent=node, prior=1.0, prm_score=0.05, is_terminal=True, final_answer="{6,-1}")]
        elif "Step 1B" in node.state_text:
            return [ReasoningNode("Ans: x in {2, 3}", parent=node, prior=1.0, prm_score=0.98, is_terminal=True, final_answer="{2,3}")]
        return []

    best_leaf = mcts_reasoning_search(root, toy_expand, num_simulations=16, c_puct=1.25)
    assert best_leaf.final_answer == "{2,3}", f"MCTS failed: {best_leaf.final_answer}"

    print(f"[PASS] PRM Masked Step Loss            : {loss.item():.4f} (Grad norm: {hidden.grad.norm().item():.4f})")
    print(f"[PASS] Verifier-Weighted Majority Vote : Selected '{weighted_winner}' over 3 low-PRM distractors")
    print(f"[PASS] MCTS + PRM Search Winner        : '{best_leaf.state_text}' (Q={best_leaf.q_value:.3f}, visits={best_leaf.visit_count})")

if __name__ == "__main__":
    verify_test_time_compute()

6. Test-Time Compute Strategies: Head-to-Head

Strategy Verifier Needed? Early Pruning? Latency Profile Best Regime
Self-Consistency (Maj@$N$) No (uses answer frequency) No (full rollouts) Low (100% parallel batch) Closed-form math/QA with $p \gt 0.25$
ORM Best-of-$N$ Outcome Verifier $R_{\text{ORM}}$ No (scores at EOS only) Low (parallel rollout + 1 ORM pass) Code generation & open-ended tasks
PRM Step Beam Search Process Verifier $r_\psi(x, s_{1:k})$ Yes (prunes every step $k$) Moderate (synchronized step barriers) Multi-step math & logical deductions
PRM-Guided MCTS (PUCT) Process Verifier + Policy Prior Yes (adaptive lookahead) Higher (sequential tree updates) Hard Olympiad proofs & synthetic data generation