Compute-Optimal Inference and Evaluation

The question is never simply "is the model smart enough?" It is always "is the model smart enough given the time and money I am willing to spend?"

Chinchilla, Cost Conscious AI Agent

1. The Compute-Optimal Inference Frontier

The compute-optimal inference problem, formalized by Snell et al. (2024), asks: given a total inference compute budget of $C$ FLOPs for a single query, what combination of model size and test-time strategies maximizes expected accuracy?

1.1 Formalizing the Trade-off

Let $C$ be the total inference compute budget. A model with $N$ parameters requires approximately $2N$ FLOPs per token generated (for dense transformers; MoE models use fewer active parameters). If the model generates $T$ total tokens (including thinking tokens), the compute cost is roughly $2NT$. For best-of-N with $K$ samples, the cost is $2NTK$.

The optimization problem is:

where $N$ is model size, $T$ is tokens per sample, and $K$ is number of samples.

A concrete comparison illustrates the trade-off. Suppose you have a budget of $C = 2.8 \times 10^{13}$ FLOPs. You can spend it two ways:

# Numeric example: comparing two strategies under the same FLOP budget
N_large = 70e9 # 70B parameter model
T_large = 200 # 200 tokens output
K_large = 1 # single pass
cost_large = 2 * N_large * T_large * K_large
print(f"Strategy A: 70B x 200 tokens x 1 sample = {cost_large:.1e} FLOPs")

N_small = 7e9 # 7B parameter model
T_small = 200 # 200 tokens per sample
K_small = 100 # best-of-100 with verifier
cost_small = 2 * N_small * T_small * K_small
print(f"Strategy B: 7B x 200 tokens x 100 samples = {cost_small:.1e} FLOPs")
print(f"Same budget: {cost_large == cost_small}")
# Both cost 2.8e13 FLOPs, but Strategy B wins on hard problems

1.2 Key Findings

Snell et al. established several important results through experiments on mathematical reasoning benchmarks:

1. Difficulty-dependent optimality: The optimal allocation shifts dramatically with problem difficulty. For easy problems (base accuracy >90%), a single pass through the largest affordable model is optimal. For hard problems (base accuracy 10% to 50%), using a smaller model with many samples and a reward model beats the larger model at the same total FLOPs.
2. The 14x equivalence: On hard problems, a model with test-time compute scaling can match a model 14x larger (by parameter count) at the same total compute budget. This means that a 7B model with best-of-N (N=128) and a PRM verifier can match a ~100B model on specific hard reasoning tasks.
3. Reward model quality matters: The benefit of test-time compute depends critically on having a good verifier. Without a reliable reward model, additional samples do not help because you cannot distinguish correct from incorrect solutions. The quality of the reward model sets the ceiling on test-time scaling benefits.
4. Diminishing returns: Like train-time scaling, test-time scaling exhibits diminishing returns. Doubling inference compute typically yields a 5% to 15% accuracy improvement, with the gain decreasing as compute increases.

Why this matters for production systems. The compute-optimal frontier turns model selection from a single choice into a routing problem. Instead of picking one model for all queries, you classify each query by difficulty and route easy queries to a large model with a single pass (cheapest per correct answer) and hard queries to a smaller model with many samples plus a verifier (best accuracy per FLOP). This is the same tiered routing pattern discussed in Section 10.3, but applied at the inference-compute level rather than the model-selection level. In practice, the difficulty classifier itself can be surprisingly simple: a lightweight model that predicts whether the base model's first attempt will be correct.

Figure 8.5.1: Compute-optimal inference frontier. On easy tasks (green), larger models are optimal. On hard tasks (red), smaller models with test-time compute scaling cross over and outperform larger models at matched total FLOPs. The crossover point depends on task difficulty and reward model quality.

2. Monte Carlo Tree Search for Language

Monte Carlo Tree Search (MCTS), the algorithm behind AlphaGo's superhuman Go play, can be adapted for structured reasoning over language. Instead of searching a tree of board positions, MCTS for language searches a tree of partial reasoning chains, where each node represents the reasoning so far and edges represent possible next reasoning steps.

2.1 Adapting MCTS for Reasoning

The four phases of MCTS (selection, expansion, simulation, backpropagation) map to language reasoning as follows:

1. Selection: Starting from the root (the problem statement), traverse the tree using the UCB formula to balance exploration and exploitation. Visit the child node with the highest UCB score: $UCB(node) = V(node)/N(node) + c * sqrt(ln(N(parent))/N(node))$, where V is cumulative value, N is visit count, and c is an exploration constant.
2. Expansion: At a leaf node, generate one or more candidate next reasoning steps using the language model. Each new step becomes a child node in the tree.
3. Simulation (rollout): From the new node, complete the reasoning chain to produce a final answer. This can use the same model with lower quality settings (higher temperature, fewer tokens) for speed, or a smaller model.
4. Backpropagation: Evaluate the completed solution (correct/incorrect), then propagate the result back up the tree. Increment visit counts and update value estimates for all nodes along the path from root to the expanded leaf.

After a fixed number of iterations (or a compute budget), the algorithm returns the solution found by following the most-visited child at each level from the root.

2.2 Key Design Choices for Language MCTS

• Step granularity: What constitutes a "step"? Options include sentence-level, paragraph-level, or reasoning-block-level (delimited by newlines or special tokens). Finer granularity allows more precise pruning but increases the tree's branching factor.
• Value function: MCTS requires a value function to evaluate partial solutions. Options include: a trained PRM (best quality), the LLM itself prompted as a self-evaluator (convenient but less reliable), or the rollout success rate (accurate but slow).
• Branching factor: How many candidate steps to generate at each expansion. Higher branching explores more options but dilutes compute. Typical values are 3 to 8 candidates per expansion.
• Stopping criterion: When to stop expanding. Options include a fixed compute budget (iterations or tokens), convergence detection (the best solution has not changed for N iterations), or confidence threshold (the highest-value leaf exceeds a score threshold).

2.3 AlphaProof: MCTS for Formal Mathematics

Google DeepMind's AlphaProof (2024) is the most impressive application of tree search to language reasoning. The system combines:

• A language model that proposes proof steps in the Lean 4 formal verification language
• MCTS to search over the proof tree, exploring different proof strategies
• Lean 4's type checker as a perfect verifier (a proof either compiles or it does not)
• Reinforcement learning to improve the language model's proof-step proposals over time

AlphaProof solved 4 of 6 problems from the 2024 International Mathematical Olympiad, earning a score equivalent to a silver medal. This is the strongest result to date in automated mathematical reasoning. The system's success demonstrates the power of combining language models with formal verification and structured search.

Why formal verification changes the game. The critical ingredient in AlphaProof is not the language model or the tree search; it is the Lean 4 type checker acting as a perfect verifier. In most reasoning tasks, verifying correctness is nearly as hard as generating the solution. But formal proof systems flip this asymmetry: checking a proof is trivial (the compiler does it), while finding one is extremely hard. This is why MCTS works so well here: each candidate proof step can be checked instantly, enabling the tree search to prune incorrect branches with zero false positives. For practical applications outside formal math, the absence of a perfect verifier is the main barrier to adopting tree search, which is why the reward model quality findings from Section 8.3 are so important.

3. Reasoning Benchmarks

Evaluating reasoning models requires benchmarks that specifically test multi-step logical reasoning, mathematical problem-solving, and complex analysis. Standard LLM benchmarks (MMLU, HellaSwag) are not sufficiently challenging to differentiate reasoning models. The following benchmarks have become the standard evaluation suite for reasoning capabilities.

3.1 Mathematics

3.2 Coding

3.3 Science and General Reasoning

Benchmark scores tell us what reasoning models can do, but they do not tell us what they cost. Before committing to these models in production, we need to quantify the additional expense and decide when the performance gains justify the price.

4. The Reasoning Tax: Cost Analysis

The "reasoning tax" is the additional cost of using reasoning models compared to standard models for the same task. Understanding this tax is essential for determining whether reasoning models are economically viable for a given application.

4.1 Cost Breakdown Example

Consider a task where you need to process 10,000 complex analysis queries per day. We compare three approaches:

The "routed (optimal)" row uses a difficulty classifier to send 70% of queries to GPT-4o and 30% to o4-mini. It achieves 90% accuracy (close to the reasoning model's 88% overall, higher on the hard subset) at a fraction of the cost of sending everything to a reasoning model.

4.2 When the Reasoning Tax Pays Off

The reasoning tax is justified when the value of the accuracy improvement exceeds the additional cost. Concretely:

For example, if each correct analysis saves $50 in manual review and accuracy improves from 72% to 90% on 10,000 daily queries, the additional 1,800 correct analyses per day save $90,000. The additional daily cost of $120 (routed approach) is trivially justified. But if each correct analysis saves only $0.01, the $120 additional cost far exceeds the $18 in value from improved accuracy.

5. Research Frontiers

Exercises