PaperlayerView exact revision on arXiv

Source-linked AI summary

Wider or Deeper? Scaling LLM Inference-Time Compute with Adaptive Branching Tree Search

Yuichi Inoue, Kou Misaki, Yuki Imajuku, So Kuroki, Taishi Nakamura, Takuya Akiba

arXiv:2503.04412v5cs.AI

TL;DR

Repeated sampling benefits from candidate diversity but lacks refinement through external feedback, despite such feedback being available in tasks like coding. AB-MCTS adaptively balances generating new candidates with refining existing ones through multi-turn tree search. Across complex coding and engineering evaluations, it outperformed repeated sampling and standard MCTS, while relying on a reliable score evaluator.

  • Problem

    Repeated sampling explores multiple candidates but lacks an explicit mechanism to exploit external feedback for refining promising solutions.

  • Method

    AB-MCTS uses Bayesian decision-making at each tree node to choose between generating new candidate responses and refining existing ones with external feedback.

  • Results

    AB-MCTS outperformed repeated sampling and standard MCTS across evaluations on complex coding, engineering, and other challenging tasks using frontier models.

  • Takeaways & Limitations

    Adaptive handling of unbounded branching combines LLM response diversity with multi-turn solution refinement for inference-time scaling.

  • Takeaways & Limitations

    AB-MCTS assumes a reliable score evaluator, which can be difficult to develop depending on the task.

Abstract

from arXiv · show

Recent advances demonstrate that increasing inference-time computation can significantly boost the reasoning capabilities of large language models (LLMs). Although repeated sampling (i.e., generating multiple candidate outputs) is a highly effective strategy, it does not leverage external feedback signals for refinement, which are often available in tasks like coding. In this work, we propose Adaptive Branching Monte Carlo Tree Search (AB-MCTS), a novel inference-time framework that generalizes repeated sampling with principled multi-turn exploration and exploitation. At each node in the search tree, AB-MCTS dynamically decides whether to "go wider" by expanding new candidate responses or "go deeper" by revisiting existing ones based on external feedback signals. We evaluate our method on complex coding and engineering tasks using frontier models. Empirical results show that AB-MCTS consistently outperforms both repeated sampling and standard MCTS, underscoring the importance of combining the response diversity of LLMs with multi-turn solution refinement for effective inference-time scaling. Code is available at https://github.com/SakanaAI/treequest .

1 Introduction

Inference-time scaling benefits from generating diverse candidates, but repeated sampling does not refine them with external feedback. AB-MCTS addresses this gap by adaptively choosing wider exploration or deeper refinement and outperforms repeated sampling and standard MCTS under matched computation.

  • Inference-time scaling allocates additional computation during inference to improve LLM performance on complex tasks.
  • Repeated sampling generates multiple independent candidates and selects a promising answer, exploiting LLM output diversity but focusing exclusively on exploration.
  • External feedback, such as code-test results, enables promising solutions to be selected and refined across multiple turns.
  • Fixed-width methods use a predetermined number of generated answers or child nodes, which hinders scaling to diverse and vast LLM output spaces.
  • AB-MCTS adaptively chooses to generate new responses or refine existing ones at each tree node using external feedback and Bayesian posterior updates.
  • Under the same computational budget, AB-MCTS achieved better results than repeated sampling and standard MCTS on complex coding, engineering, and ARC-AGI evaluations.

2 Related Work

Related inference-time scaling methods include fine-tuning, reward-guided sequential reasoning, repeated sampling, task-specific search, and MCTS-based approaches. AB-MCTS builds on this landscape by combining broad candidate generation with feedback-based refinement while addressing fixed-width limitations.

  • Post-training fine-tuning deepens reasoning through reinforcement learning or supervised chain-of-thought training, whereas AB-MCTS generates and refines multiple candidates with external feedback.
  • Reward-guided chain-of-thought methods search one step at a time, primarily targeting improved single-answer quality in mathematical tasks.
  • Repeated sampling is widely used because it generates many candidates and selects the best one, but it does not refine candidates using external feedback.
  • Progressive widening gradually increases actions per node using visit-count heuristics, while AB-MCTS exploits same-LLM generation homogeneity to choose statistically between widening and deepening.

3 Method

AB-MCTS adapts Monte Carlo tree search to unbounded LLM branching by choosing between generating new answers and refining existing ones. It uses external feedback, Bayesian posterior predictives, and Thompson sampling to balance exploration and exploitation.

  • Each iteration selects an expansion target, generates one child answer, backs up its score, and ultimately selects the best node.For non-root nodes, expansion generates a refinement; no separate rollout is needed because scores are directly evaluable.
  • Unlike standard MCTS, AB-MCTS permits repeated expansion of non-leaf nodes and represents new-child generation with a GEN node, enabling theoretically unbounded branching.A selected GEN node branches off from its parent, while existing child nodes support deeper refinement.
  • Thompson sampling chooses between the GEN action and existing child nodes by sampling posterior predictive scores and selecting the action with the highest sampled score.At each node, the action set contains one GEN action and actions for already-existing children.
  • AB-MCTS-M: AB-MCTS-M models each subtree as a group in a node-specific mixed Bayesian model, using observed scores to predict expected scores for new nodes and existing subtrees.MCMC estimates shared parameters; the GEN action receives an inferred group-level intercept informed by observed groups, and score backups update the posterior indirectly through shared parameters.

4 Experiments

AB-MCTS is evaluated across coding, abstract reasoning, and machine-learning engineering benchmarks using frontier language models and multiple inference-time generation budgets. It adaptively balances exploration and exploitation, generally outperforming baselines, scaling better at large budgets, and reaching strong performance efficiently.

  • Experimental Setup: AB-MCTS is evaluated on LiveCodeBench, CodeContest, ARC-AGI, and MLE-Bench using GPT-4o and DeepSeek-V3.Each model generates a complete solution per API call, with the generation budget defined as the maximum number of API calls.
  • Benchmark Results: AB-MCTS generally outperforms baselines on LiveCodeBench and CodeContest, pulling ahead at budgets of 23 and 25 or more, respectively.The comparison reports Pass@1 against generation budget for GPT-4o.
  • Benchmark Results: On ARC-AGI, AB-MCTS matches repeated sampling at the standard budget while standard MCTS yields only marginal improvements with larger budgets.ARC-AGI results use Pass@2 under the official evaluation protocol.
  • Benchmark Results: AB-MCTS-M consistently performs strongly across MLE-Bench competitions despite variation in which baseline performs best.The results highlight different task-specific exploration–exploitation trade-offs.
  • Search Behavior: AB-MCTS generates wider and more flexible trees than standard MCTS by choosing to explore from any existing node.Figure 5 analyzes average depth and average width at each depth to characterize this behavior.
  • Scaling with Increased Budget: With ARC-AGI budgets increasing from 200 to 500, AB-MCTS improves substantially while repeated sampling begins to plateau and standard MCTS remains less successful.The extended experiment uses DeepSeek-V3 and a generation budget up to 512.
  • Search Behavior: Search-tree visualizations show AB-MCTS-M dynamically allocating budget between diverse new candidates and refinement of promising candidates.The example tree colors nodes by evaluation score and marks failed-to-execute candidates in grey.
  • Efficiency and Performance: AB-MCTS can reach repeated sampling’s peak performance earlier on LiveCodeBench and CodeContest and eventually exceed repeated sampling on ARC-AGI.This comparison accounts for repeated sampling’s potential parallel-sampling efficiency and absence of feedback-computation costs.

5 Conclusions

AB-MCTS adaptively integrates multi-turn exploration and exploitation for inference-time scaling, outperforming repeated sampling and standard MCTS. Its main scope limitation is reliance on a reliable score evaluator and simplified cost accounting.

  • AB-MCTS dynamically chooses whether to “go wider” or “go deeper” using external feedback and Bayesian decision-making.
  • AB-MCTS outperforms repeated sampling and standard MCTS, demonstrating the value of adaptive handling of unbounded branching.
  • The approach assumes a reliable score evaluator, which can be difficult to develop for some tasks.
  • The paper identifies finer-grained real-world costs beyond API call counts as a direction for improving practical utility.

A Method Details

The method frames answer search as stochastic LLM generation followed by quantitative evaluation, then combines repeated sampling and sequential refinement through adaptive tree search. AB-MCTS uses reward-aware branching rather than a fixed branching factor, with Bayesian models supporting node selection.

  • Problem setup: An LLM maps a natural-language prompt to an output, and a score evaluator maps that output to a normalized reward where higher values are better.
  • Existing methods: Repeated sampling generates multiple candidates independently and selects one using criteria such as best-of-n, majority voting, or self-consistency.
  • Existing methods: Sequential refinement repeatedly constructs new inputs from prior input-output pairs and feedback, then generates refined answers.
  • Existing methods: Pure exploration and pure exploitation form limiting tree-search cases, while standard MCTS combines them with a fixed branching factor.
  • Adaptive branching: Progressive widening predetermines branching from visit counts and hyperparameters without using observed child rewards in the branching decision.
  • Adaptive branching: AB-MCTS adapts branching from observed rewards, deciding whether to generate new candidates or refine existing ones instead of fixing tree width.
  • AB-MCTS-M: AB-MCTS-M fits node-level mixed models, using group-level intercepts for base-solution quality and per-instance noise for generated rewards.
  • Hyperparameter sensitivity: A sensitivity study found stable Pass@1 performance across tested prior settings, indicating low sensitivity to initial hyperparameters.

B.1 Tasks and Datasets

The evaluation spans coding, abstract reasoning, and machine-learning benchmarks, using task-specific scores and shared, minimally informed Bayesian priors. Figures compare performance and search-tree shape across budgets, algorithms, and tasks.

  • Tasks and datasets: The evaluation covers LiveCodeBench, CodeContest, ARC-AGI, and MLE-Bench, with repeated runs used to account for stochastic generation.
  • Tasks and datasets: CodeContest and LiveCodeBench use public tests for node scores and hidden tests for final evaluation, counting success only when all hidden cases pass.
  • Tasks and datasets: LiveCodeBench uses problems released between August and November 2024 to prevent data contamination.
  • Tasks and datasets: ARC-AGI scores nodes by the fraction of examples whose transformations are correct, using generated Python code to express inferred rules.
  • Tasks and datasets: MLE-Bench evaluation uses three low-complexity competitions, with validation data used to score nodes and select solutions.
  • Priors: The shared priors place most or all probability mass within [0, 1], use mean score 0.5, and avoid excessive concentration.
  • Figures: Figure 8 plots success rate against generation budget for AB-MCTS methods and baselines on LiveCodeBench, CodeContest, and ARC-AGI with DeepSeek-V3.
  • Figures: Figure 9 relates performance to log-ratio tree depth over width, where larger values indicate deeper searches and smaller values indicate wider searches.

C.1 Results with DeepSeek-V3 on Competitive Programming and ARC-AGI

With DeepSeek-V3, AB-MCTS remains strong across competitive programming and ARC-AGI despite changes in baseline strengths. Its trees are generally wider than standard MCTS, while still adapting depth to promising branches and task characteristics.

  • DeepSeek-V3 changes the relative strengths of baseline methods, including standard MCTS achieving the highest LiveCodeBench success rate.
  • AB-MCTS retains strong performance when the underlying model changes, although CodeContest differences from top baselines are less pronounced than with GPT-4o.
  • Repeated sampling produces wide trees, sequential refinement produces deep trees, and AB-MCTS consistently produces wider trees than standard MCTS.
  • On LiveCodeBench, AB-MCTS remains strong even though sequential refinement outperforms repeated sampling, indicating adaptive selection across competing search patterns.
  • Across three MLE-Bench competitions, the strongest baseline varies by task, while AB-MCTS-M consistently delivers strong performance.

C.3 Example search trees generated by each methods on MLE-Bench

AB-MCTS-M produces search trees that combine broader exploration with focused exploitation more flexibly than standard MCTS. Progressive widening can match AB-MCTS when tuned adequately, but its performance is sensitive to hyperparameters and can become unstable.

  • AB-MCTS-M search trees visually combine broader exploration with focused exploitation across all three MLE-Bench competitions.
  • Adequate progressive-widening parameters can produce performance comparable to AB-MCTS on LiveCodeBench.
  • Progressive widening is highly sensitive to hyperparameters, producing its worst result at (k, α) = (1, 0.45).
  • The (k, α) = (10, 0.55) setting shows search instability and the highest variance, whereas AB-MCTS is robust without such tuning.

C.5 AB-MCTS-M vs. AB-MCTS-A: Analysis and Selection

AB-MCTS-M and AB-MCTS-A offer different trade-offs in adaptive search. AB-MCTS-M is generally favored for outcome quality, while AB-MCTS-A is lighter and tends to support wider exploration.

  • AB-MCTS-M is often preferred when outcome quality is the main priority because of its consistently strong performance.Its node selection uses MCMC, which improves effectiveness but adds computational cost per selection.
  • AB-MCTS-A offers a lighter alternative for applications with strict time constraints through analytically tractable Gaussian and Beta posterior updates.
  • AB-MCTS-A tends to construct wider search trees because reaching depth d requires d consecutive CONT choices, making deeper paths geometrically less likely.This design favors wider expansions over deep paths.
  • AB-MCTS-A may be advantageous for tasks such as ARC-AGI where broader exploration is considered particularly beneficial.
  • The choice depends on the application: AB-MCTS-M targets outcome quality, whereas AB-MCTS-A emphasizes computational efficiency and broader exploration.
  • The ranking of methods remains unchanged between Pass@1 and Pass@2, indicating that the main conclusions are robust to the evaluation metric.

C.7 Analysis of Adaptive Search Behavior via Node Degree Distribution

The paper examines adaptive search through node-degree behavior, branching-factor sensitivity, and extensions to difficulty-aware and multi-LLM search. These analyses frame adaptive branching as a way to adjust exploration across tasks, models, and solution stages.

  • Node-degree behavior: AB-MCTS-M favors depth-focused search, with around 90% of non-leaf nodes having degree 1–3, while still broadening adaptively when required.Its long-tail degree distribution reaches up to 40.
  • Branching-factor sensitivity: Standard MCTS performs best at fixed branching factor w = 5 among tested values, while performance degrades at both smaller and larger widths.This highlights the baseline’s sensitivity to its branching-factor hyperparameter.
  • Adaptive branching: Unlike standard MCTS, whose branching factor must be predefined, AB-MCTS dynamically adjusts effective width across task types and LLMs.The paper reports robust AB-MCTS performance across varied task types, while other methods excel selectively.
  • Future directions: A proposed future direction is to estimate problem difficulty from collected rewards and switch between deep AB-MCTS-M and wide AB-MCTS-A strategies.The proposal identifies low rewards as a signal of difficult problems.
  • Multi-LLM search: Multi-LLM AB-MCTS extends search by selecting among multiple answer generators during node expansion.The selected generator is chosen using the current answer tree and can be adapted to the search stage.
  • Evaluation setting: The multi-LLM extension is evaluated on ARC-AGI-2, where Python-code solutions are refined using rewards based on correctly solved demonstrations.The public evaluation set contains 120 problems, with a generation budget of 250 per problem.
  • Evaluation setting: The ARC-AGI-2 experiment focuses on Pass@k search success rather than the official Pass@2 contest criterion because Pass@2 requires an additional candidate-selection mechanism.

D.2.2 Results

On ARC-AGI-2, AB-MCTS improves over repeated sampling, and combining Gemini-2.5-Pro with DeepSeek-R1-0528 improves performance further. The multi-LLM method finds correct solutions for more than 30% of the problems.

  • 23% Pass@k success was achieved by repeated sampling with o4-mini, compared with 27.5% for single-model AB-MCTS using o4-mini.The advantage of AB-MCTS over repeated sampling becomes more evident as the generation budget increases.
  • Figure 13 compares Pass@k coverage across generation budgets for repeated sampling, AB-MCTS, and Multi-LLM AB-MCTS.
  • Figure 14 sorts ARC-AGI-2 trials by maximum demonstration-case reward and marks trials where a correct final solution was found.
  • Over 30% of ARC-AGI-2 problems were solved by Multi-LLM AB-MCTS combining Gemini-2.5-Pro and DeepSeek-R1-0528.
  • Integrating DeepSeek-R1-0528 increased the number of solved problems despite its lower individual performance.

D.2.3 Analysis

Multi-LLM AB-MCTS adaptively allocates heterogeneous models across problems and can enable collaborative solutions when no single model succeeds. However, preliminary Pass@2 results reveal a substantial gap from the much larger Pass@250 rate.

  • Adaptive model allocation: Model usage varied by problem difficulty: proficient models tended to handle trials that quickly achieved high demonstration-case rewards, while harder trials used models more variably.The distribution of LLM usage differed across problems according to how readily trials obtained high rewards.
  • Collaborative solving: Some problems unsolvable by any single LLM were solved through collaboration among multiple models.An incorrect o4-mini solution served as a useful hint for DeepSeek-R1-0528 and Gemini-2.5-Pro in producing the correct solution.
  • Collaborative solving: Multi-LLM AB-MCTS facilitated flexible collaboration among heterogeneous frontier LLMs rather than merely matching each problem to its best individual model.The reported collaboration involved multiple models contributing to a correct solution.
  • Challenges and future work: 19.2% Pass@2 remained more than 10 percentage points below the 30% Pass@250 rate under the preliminary final-answer selection rule.The authors identify more sophisticated reward modeling or LLM-based judging as possible directions for closing this gap.
Loading 2503.04412v5…