🤖 AI

CoBa: Cost-Effective Test-Time Scaling via Compute-Balanced Routing

This paper introduces CoBa, a compute-balanced routing policy that optimizes test-time scaling by dynamically allocating inference resources between generation and verification, achieving state-of-the-art accuracy on mathematical reasoning benchmarks while significantly reducing computational costs compared to traditional scaling methods.

Original authors: Yan Zhou, Yue Ouyang, Kaiyang Zheng, Suncheng Xiang

Published 2026-08-11

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Yan Zhou, Yue Ouyang, Kaiyang Zheng, Suncheng Xiang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a really tricky riddle. You have a limited amount of energy to spend, and you have three ways to use it: you can think of a new answer, you can double-check an answer you already have, or you can just say, "Okay, I'm done, here is my best guess." For a long time, computers trying to solve hard problems (like math puzzles) have been told to just "try harder" by doing one of these things over and over again. They might generate a hundred different answers and pick the most popular one, or they might spend a huge amount of energy checking every single answer with a super-smart, slow robot judge. The problem is, this is like using a sledgehammer to crack a nut. If the answer is obvious, checking it a hundred times is a waste. If the answer is super hard, just guessing randomly might never find it. The big question for scientists is: how do we spend our limited energy so we get the right answer without burning out the battery?

This is where a new idea called CoBa comes in. Think of CoBa as a smart manager for a team of workers. Instead of telling every worker to do the exact same thing, CoBa watches the situation and decides the very best next step. If the workers are already agreeing on an easy answer, CoBa says, "Great, stop! We're done." If they are confused, CoBa might say, "Let's ask two more people for ideas." But if they are stuck on a really hard problem, CoBa says, "Okay, let's send the best two ideas to the super-expert judge to check." The researchers tested this on thousands of math problems and found that by being a smart manager, CoBa could get the same high scores as the "brute force" methods (which try everything) but used 49.1% fewer of the expensive "energy units" (called parameter-weighted tokens). It's like getting the same delicious meal but using half the ingredients.

The Problem: The "More is Better" Trap

For a while, the best way to make AI smarter at solving problems was to just throw more computing power at it. If you wanted a better answer, you asked the AI to generate more solutions. If you wanted to be sure, you asked a second AI to check them all. This is like trying to find a needle in a haystack by bringing in a thousand people to look at the whole haystack, even if the needle is right on top. It works, but it's incredibly expensive and slow.

The paper argues that this "one-size-fits-all" approach is wasteful. Some problems are easy and need a quick glance. Others are hard and need deep thinking. But the old methods treated every problem the same. They would spend the same huge amount of energy on a simple math question as they would on a complex Olympiad-level puzzle. The authors of this paper wanted to fix this by turning test-time scaling (making the AI think harder during the test) into a resource allocation problem. They asked: "Given a fixed budget of energy, should we spend the next bit of energy on generating a new idea, checking an old one, or just stopping?"

The Solution: The Smart Manager (CoBa)

The authors introduced CoBa (Compute-Balanced test-time scaling), which acts like a traffic controller for the AI's brain. Here is how it works, step-by-step:

The Warm-up: First, CoBa asks the AI to generate a small, diverse set of answers (like asking two friends for their ideas). This is the "warm-up."
The Cheap Check: Next, it runs a quick, low-energy check on all the answers. This is like asking a friend, "Does this sound right?" It's fast and cheap.
The Decision: Based on what the cheap check found, CoBa makes a choice:
- Stop: If everyone agrees and the cheap check is confident, CoBa stops immediately. No need to waste energy.
- Generate More: If the answers are all over the place, CoBa asks for a few more ideas to get more options.
- Strong Verify: If there are a few promising answers that are still a bit uncertain, CoBa sends only those specific ones to a "Super Judge" (a much more powerful, expensive AI model) for a deep, thorough check.

This is the key difference: instead of checking everything with the Super Judge, CoBa only sends the most interesting candidates. It saves the expensive energy for the moments where it actually matters.

The Results: Smarter, Not Harder

The researchers tested this system on over 3,000 math problems, ranging from standard school math to very hard competition-level puzzles (like AIME and AMC). They compared CoBa against the old "brute force" methods.

The Big Win: The best version of CoBa (called CoBa-Routed-Strong) achieved an accuracy of 85.13%. This is almost exactly the same as the most expensive methods, which scored 85.20% (using a self-evaluation weighted voting system) and 85.12% (using a "best-of-16" majority vote).
The Savings: Here is the magic part. To get that same high score, the expensive methods used a massive amount of computing power. CoBa-Routed-Strong used 49.1% fewer parameter-weighted tokens than the self-evaluation method and 58.9% fewer than the best-of-16 method.
The Trade-off: The paper notes that the "best-of-16" method (which just generates 16 answers and picks the most common one) was still slightly more accurate by a tiny fraction (0.01 points), but it cost 2.43 times more to run. CoBa suggests that for most real-world uses, the tiny accuracy gain isn't worth the huge cost.

What This Means for the Future

The paper suggests that the future of AI reasoning isn't just about making bigger, stronger models. It's about being smarter with the resources we have. The authors found that the biggest bottleneck isn't always the checking; sometimes, it's that the AI just didn't come up with the right answer in the first place. On the hardest problems (like AIME 2025), even the smartest manager couldn't find the right answer if the initial "pool" of ideas didn't contain it.

However, for the vast majority of problems, CoBa showed that we can stop wasting energy. By using a "triage" system—cheap checks for everyone, and expensive checks only for the tricky ones—we can get the best results without breaking the bank. The paper concludes that this approach turns the "oracle gap" (the difference between what the AI could do if it knew the answer and what it actually does) into a useful signal. If the AI fails, we now know exactly where to look: did we stop too early? Did we check the wrong candidates? Or did we just need to generate better ideas in the first place?

In short, CoBa teaches us that in the world of AI, timing and targeting are just as important as raw power. It's not about how hard you think; it's about knowing when to think hard and when to stop.

Technical Summary: CoBa – Cost-Effective Test-Time Scaling via Compute-Balanced Routing

1. Problem Formulation

The paper addresses the challenge of test-time scaling in Large Language Models (LLMs). Current strategies typically scale computation along a single axis: generating more solutions (sampling), extending reasoning chains, or applying stronger evaluators. Under a fixed inference budget, these strategies compete; spending compute on one action precludes it from being used elsewhere.

The authors formalize test-time reasoning as a compute-allocation problem. At any decision state $t$ , a system must determine the optimal next action from a space including:

Sampling: Generating a new candidate solution.
Verification: Applying lightweight ( $V_1$ ) or strong ( $V_L$ ) verifiers to existing candidates.
Stopping: Finalizing the answer.

The objective is to maximize task utility (accuracy) while minimizing cost (measured in tokens, model calls, and parameter-weighted tokens) within a fixed budget $B$ . The core hypothesis is that accuracy gains and cost savings are best achieved not by uniformly applying expensive actions, but by dynamically routing compute based on the "value" of the next unit of computation for a specific example.

2. Methodology: CoBa Routing Policy

The proposed solution, CoBa (Compute-Balanced test-time scaling), implements a staged routing policy (Algorithm 1) that treats inference as a compute allocation process. It operates without ground-truth labels during the routing decision.

Key Components:

State Representation: The system state $s_t$ includes the current candidate set $C_t$ , observed verifier scores $S_t$ , remaining budget $B_t$ , and history features (e.g., answer diversity, token caps).
Warm-up Phase: The system begins with a fixed, small number of generations ( $k=2$ ) to establish initial candidate diversity.
Triage Layer (Cheap Verification): All candidates are scored by a lightweight verifier (e.g., a smaller LLM or rule-based frequency check). This provides "cheap evidence" regarding answer agreement and confidence.
Routing Logic:
- Stop: If the top answer is stable (high agreement) and the lightweight score is high, the system stops.
- Sample: If the system is uncertain, it generates an additional candidate.
- Route to Strong Verification: If specific candidates are uncertain or high-value, they are routed to a stronger, more expensive verifier ( $V_L$ ).
Final Ranking: The final answer is selected using a fused ranking score $R(c)$ that combines answer frequency, lightweight scores, and strong verifier scores (Eq. 6).
Variants: The paper evaluates three configurations:
- CoBa-Routed-Light: Warm-up + cheap verification only.
- CoBa-Routed: Warm-up + up to 4 candidates + routing 2 to strong verification.
- CoBa-Routed-Strong: Warm-up + up to 8 candidates + routing 4 to strong verification.

Cost Metrics:
The paper distinguishes between total token cost and parameter-weighted token cost ( $C_{ptok}$ ), which accounts for the computational expense of larger models (e.g., a 14B model is weighted higher than an 8B model).

3. Experimental Setup

Datasets: Evaluations span 3,129 example-generator pairs across MATH-500, AIME 2024/2025, AMC 2023, and a procedural symbolic reasoning subset (Reasoning Gym).
Models: Generators include Qwen3-14B, Phi-4-reasoning, and Qwen3-8B. Qwen3-8B serves as the lightweight judge; Qwen3-14B serves as the strong verifier.
Protocol: An offline replay protocol is used. All methods operate on a shared pool of $N=16$ pre-generated candidates per problem. This ensures fair comparison by isolating the allocation strategy from the generation capability.
Baselines: Comparisons include greedy decoding, fixed-length reasoning, Best-of-N majority voting, self-consistency, and various adaptive allocation proxies (e.g., self-evaluation weighted voting).

4. Key Results

Accuracy vs. Cost: CoBa-Routed-Strong achieves 85.13% macro accuracy. This is statistically indistinguishable from the Self-Evaluation Weighted-Voting proxy (85.20%) but uses 49.1% fewer parameter-weighted tokens.
Comparison to Best-of-N: CoBa-Routed-Strong matches Best-of-16 majority voting (85.12% vs. 85.13%) while using 58.9% fewer parameter-weighted tokens.
Significance: Paired bootstrap tests confirm CoBa-Routed-Strong significantly outperforms greedy decoding (+3.74 points) and is statistically comparable to high-cost baselines.
Dataset Specifics:
- On AIME 2024, routing significantly improves accuracy (from 65.6% greedy to 82.2%).
- On AIME 2025, a large gap remains between CoBa and the "Pool Oracle" (the theoretical maximum if a correct answer exists in the pool). This suggests that for the hardest problems, the bottleneck shifts from allocation to candidate generation (i.e., the correct answer is not in the generated pool).
Action Distribution: Analysis shows CoBa adapts its behavior: easy examples stop early after cheap verification, while hard examples trigger additional sampling and strong verification. This contrasts with uniform baselines that apply the same compute budget to all examples.

5. Contributions and Significance

The paper claims three primary contributions:

Formalization: It reframes test-time reasoning as a unified compute-allocation problem over generation, verification, and stopping actions.
Methodology: It introduces a reproducible, local routing policy (CoBa) that utilizes a two-tier verification system (lightweight and strong) to make adaptive allocation decisions.
Empirical Evidence: Through controlled replay experiments, it demonstrates that adaptive routing can reach the accuracy of much more expensive sampling and evaluator-scaling baselines with significantly lower cost.

Significance and Claims:

Pareto Improvement: The paper argues that CoBa achieves a Pareto improvement over common moderate-cost baselines, offering high accuracy with substantially reduced compute.
Actionable Failure Modes: The "oracle gap" analysis provides a diagnostic tool. If the gap is due to allocation errors (correct answer in pool but missed), better routing is needed. If the gap is due to generation errors (correct answer not in pool), better candidate generation is required.
Operational Scaling: The authors suggest that future test-time systems should treat generation and verification as a shared budget. The most effective scaling strategy is not "more compute everywhere," but "compute where it is most valuable."
Reporting Norms: The paper advocates for reporting the "action mix" (how many samples vs. verifiers) alongside accuracy and cost to make allocation mechanisms falsifiable and transparent.

The paper concludes that while CoBa effectively optimizes the use of available compute, the remaining performance ceiling on the hardest datasets indicates that future gains will require a combination of sharper uncertainty estimation, richer process signals, and improved candidate generation.