All articles Reasoning & Research

Test-Time Compute & Inference Scaling for Reasoning Models

A deep technical analysis of the shift from pre-training compute scaling to inference-time scaling: Monte Carlo Tree Search (MCTS), Process Reward Models (PRMs), and self-correction loops.

For years, frontier AI progress was driven almost exclusively by pre-training compute scaling laws (more parameters, more GPUs, more tokens). But as high-quality web text data reaches saturation, the frontier of AI research has shifted to Test-Time Compute Scaling — trading latency and inference FLOPs for higher accuracy on complex math, coding, and reasoning tasks. Paradigm shifts demonstrated by models like OpenAI o1 and DeepSeek R1 prove that allocating extra compute budget during generation enables smaller models to match or exceed frontier pre-trained giants.

Pre-training scales knowledge; test-time compute scales thinking depth.

The shift from pre-training to test-time scaling laws

Traditional autoregressive generation samples tokens sequentially using greedy decoding or top-p sampling. If the model makes a logical mistake at step 3 of a 50-step mathematical proof, all subsequent tokens compound that initial hallucination.

Test-time compute scaling alters this dynamic by allowing the system to explore multiple candidate reasoning trajectories, critique intermediate steps, backtrack from dead ends, and verify conclusions before returning the final response to the user.

Core mechanisms of test-time compute

1. Parallel Sampling (Majority Voting / Best-of-N)

Generates N independent responses in parallel, scores each response using an evaluation model or verifier, and selects the highest-scoring candidate (or returns the majority consensus).

2. Sequential Refinement & Chain-of-Thought (CoT)

Encourages the model to output explicit internal reasoning tokens (e.g. <think> ... </think>) to break down multi-step problems, reflect on potential errors, and refine draft answers before producing output.

3. Search Tree Exploration (MCTS / Beam Search)

Frames text generation as a search tree over intermediate thoughts. A Process Reward Model evaluates every candidate step, allowing search algorithms to prune unpromising branches early.

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

The reliability of test-time search depends heavily on the reward signal used to guide exploration:

  • Outcome Reward Models (ORM): Evaluate only the final answer (e.g., +1 if the final code passes unit tests, 0 if it fails). ORMs provide sparse feedback and fail to pinpoint which step caused the failure.
  • Process Reward Models (PRM): Evaluate every step in the step-by-step reasoning process (e.g., step 1: 0.98, step 2: 0.95, step 3: 0.12 [flagged error]). PRMs enable precise step-level credit assignment and early tree pruning.
                     [Root Prompt]
                          |
             +------------+------------+
             |                         |
       [Step 1A (0.95)]          [Step 1B (0.30)]
             |                         x (Pruned)
     +-------+-------+
     |               |
[Step 2A (0.92)] [Step 2B (0.88)]
     |               |
  [Final A]       [Final B]

Python implementation of MCTS with Process Reward Verification

import math
import random
from typing import List, Optional

class ReasoningNode:
    def __init__(self, thought_step: str, parent: Optional['ReasoningNode'] = None):
        self.thought_step = thought_step
        self.parent = parent
        self.children: List['ReasoningNode'] = []
        self.visits = 0
        self.value = 0.0
        self.prm_score = 0.0 # Process Reward Model score

    def uct_score(self, c_param: float = 1.414) -> float:
        if self.visits == 0:
            return float('inf')
        return (self.value / self.visits) + c_param * math.sqrt(math.log(self.parent.visits) / self.visits)

class MCTSSearch:
    def __init__(self, root_prompt: str, prm_evaluator, llm_generator):
        self.root = ReasoningNode(thought_step=root_prompt)
        self.prm = prm_evaluator
        self.llm = llm_generator

    def select(self, node: ReasoningNode) -> ReasoningNode:
        while node.children:
            node = max(node.children, key=lambda n: n.uct_score())
        return node

    def expand(self, node: ReasoningNode, num_candidates: int = 3):
        candidates = self.llm.generate_candidate_steps(node.thought_step, n=num_candidates)
        for cand in candidates:
            child = ReasoningNode(thought_step=cand, parent=node)
            child.prm_score = self.prm.score_step(node.thought_step, cand)
            node.children.append(child)

    def backpropagate(self, node: ReasoningNode, score: float):
        curr = node
        while curr is not None:
            curr.visits += 1
            curr.value += score
            curr = curr.parent

    def run_search(self, iterations: int = 20) -> ReasoningNode:
        for _ in range(iterations):
            leaf = self.select(self.root)
            if leaf.visits > 0:
                self.expand(leaf)
                if leaf.children:
                    leaf = leaf.children[0]
            
            # Simulate and score rollout
            reward = leaf.prm_score
            self.backpropagate(leaf, reward)

        # Select best trajectory path
        best_child = max(self.root.children, key=lambda n: n.visits)
        return best_child

The future of inference scaling

  1. Inference compute is the new scaling frontier: Models can be made smaller during pre-training if given compute budget to think during inference.
  2. Process Reward Models are key: Step-level verification enables search algorithms (MCTS / Beam Search) to prune hallucinated branches early.
  3. Self-Correction requires explicit reasoning tokens: Training models on synthetic chain-of-thought data with reflection tags (<think> ... </think>) dramatically boosts zero-shot problem solving.
← Back to all articles