AB-MCTS: Wider or Deeper?

Chapter 1: The Key Insight

AB-MCTS's central idea is beautifully simple: let the search tree decide its own shape.

The mechanism is a special node called the GEN node. Every node in the tree has a GEN node as a virtual child. When the search selects a GEN node, a brand new LLM response is generated from that node's prompt — widening the tree. When the search selects an existing child instead, it descends deeper. The choice between GEN and existing children is made by Thompson sampling over Bayesian posterior distributions of expected scores.

This means the branching factor is unbounded. A node might get 1 child or 50 children, depending on whether new samples keep producing promising results. Traditional MCTS cannot handle this because UCB/UCT requires knowing the number of arms in advance. Thompson sampling naturally accommodates new arms appearing dynamically.

Chapter 2: Repeated Sampling as a Flat Tree

Before we build AB-MCTS, let's understand exactly why repeated sampling works and where it breaks.

You have a coding problem. You call the LLM 128 times with the same prompt at temperature 0.6. Each call returns a complete solution. You run the tests on each, pick the one that passes the most. This is best-of-n.

Why does this work? Because LLMs at non-zero temperature explore a vast output space. Even if the probability of any single sample being correct is only 2%, with 128 samples the probability that at least one is correct is 1 − 0.98¹²⁸ ≈ 92%. Large Language Monkeys formalized this: the number of correct samples grows roughly as a coverage function — more samples, more coverage of the solution space.

Where it fails

Consider a problem where the LLM's approach is fundamentally correct but it keeps making a subtle off-by-one error in the boundary condition. Sampling 128 times might produce 128 solutions that all fail the same edge case in slightly different ways. What you actually need is to take the best attempt, show the LLM the failing test output, and ask it to fix that specific line. Repeated sampling cannot do this — it has no feedback loop.

In tree terms, repeated sampling is a tree of depth 1 and width n. Every node is a direct child of the root. The tree is maximally wide and minimally deep.

Chapter 3: Standard MCTS Limitations

Standard Monte Carlo Tree Search, as used in LATS (Language Agent Tree Search), applies the classic MCTS algorithm to LLM-generated solutions. Here is how it works:

1. Selection: Start at the root. At each node, use UCT (Upper Confidence bound for Trees) to select a child, balancing exploitation (high-scoring children) and exploration (rarely-visited children). Descend until you reach a leaf.
2. Expansion: At the leaf, generate exactly k new children by calling the LLM k times. Each child gets the parent's solution plus feedback (e.g., test errors) as context.
3. Evaluation: Score each new child by running tests.
4. Backpropagation: Propagate the scores up the tree to update node statistics.

The critical limitation is the fixed branching factor k. In LATS, k = 5. Every node gets exactly 5 children, no more, no less. This creates two problems:

The fundamental issue is that the tree shape is a hyperparameter, chosen before seeing any data. But the optimal shape depends on the problem: easy problems want width, hard-but-refinable problems want depth, and most problems want a mix. Standard MCTS cannot adapt.

UCT and the fixed-arm assumption

UCT assumes a fixed number of arms (children). The formula is:

where Q̄ is the average reward and N is the visit count. This breaks when the number of children can grow unboundedly — the exploration bonus for a "new child that doesn't exist yet" is undefined. AB-MCTS replaces UCT with Thompson sampling, which handles dynamic arm creation naturally.

Chapter 4: The AB-MCTS Algorithm

Now we build the full algorithm. AB-MCTS runs for n_nodes iterations. Each iteration adds exactly one new node to the tree. The process has three steps:

Step 1: Selection

Start at the root. At each node, choose among its children plus a GEN node. The GEN node represents the option of generating a brand new child right here. Use Thompson sampling (not UCT) to choose. If you select an existing child, descend into it and repeat. If you select the GEN node, stop — this node will be expanded.

Step 2: Expansion

Generate a new LLM response. If the selected node is the root, the prompt is just the task description. If it is a non-root node, the prompt includes the node's existing answer plus external feedback (compiler errors, test results). The new response becomes a new child node.

Step 3: Score Backup

Run the scoring function on the new answer (e.g., execute the code, count passing tests). Propagate this score up the tree to update the Bayesian posteriors at each ancestor.

The algorithm in pseudocode

function AB-MCTS(n_nodes):
    T = initialize_tree()
    for n = 1 to n_nodes:
        N = select_expansion_target(T)   # Step 1
        N_new = expand(N, T)              # Step 2
        score_backup(N_new, T)             # Step 3
    return select_best(T)

function select_expansion_target(T):
    N = root
    while not is_leaf(N):
        N_next = thompson_select_child(N)  # includes GEN
        if N_next is GEN_NODE:
            break                          # expand here
        N = N_next
    return N

Chapter 5: The Branching Decision

The heart of AB-MCTS is deciding: wider or deeper? This decision is made at every node during selection. The authors propose two approaches, both using Thompson sampling over Bayesian posteriors.

AB-MCTS-M: The Mixed Model

At node N with children N₁, ..., N_k, model the score of a future expansion as:

Each existing child j has a "group-level intercept" α_j capturing how good that branch is. The GEN node (action a₀) gets no direct observations, but its α₀ is drawn from the same shared distribution — so the model infers what a new branch might score based on what existing branches scored.

If existing branches score well with low variance, the GEN node's expected score is moderate with high variance — Thompson sampling will sometimes still pick it (exploration). If existing branches score poorly, the GEN node looks relatively more promising — the tree widens.

This uses MCMC (Markov Chain Monte Carlo) to sample from the posterior. Computationally heavier, but statistically principled.

AB-MCTS-A: Node Aggregation

A lighter-weight alternative. Instead of a hierarchical mixed model, AB-MCTS-A introduces a CONT node (continue) alongside the GEN node. At each node N:

• GEN node: represents "generate a brand new child." Its posterior is updated with scores of nodes created via GEN at this level.
• CONT node: represents "continue into an existing child." Its posterior is updated with scores of all descendant refinements.

Thompson sampling draws from both posteriors. If the GEN draw is higher, widen. If the CONT draw is higher, descend into the best existing child and repeat.

AB-MCTS-A uses conjugate priors (Beta or Gaussian), so posterior updates are analytical — no MCMC needed. Much faster per iteration.

Thompson Sampling (intuition)

Thompson sampling works by maintaining a probability distribution over each arm's expected reward. To choose an arm, you draw one random sample from each arm's distribution and pick the arm with the highest draw. Arms with high expected reward get picked often (exploitation). Arms with high uncertainty occasionally produce very high draws (exploration). This naturally balances both — and it handles the GEN node elegantly because a new arm naturally has high uncertainty.

Chapter 6: External Feedback Integration

AB-MCTS's power comes from combining diverse LLM generation with external feedback signals — information the LLM cannot generate on its own but that tells it exactly what went wrong.

Feedback in coding tasks

When an LLM generates a Python solution for a competitive programming problem, the system can:

1. Compile the code — catching syntax errors and import failures immediately.
2. Run public test cases — revealing which inputs produce wrong outputs, runtime errors, or timeouts.
3. Report the score — fraction of tests passed, giving a continuous signal in [0, 1].

This feedback serves two roles in AB-MCTS:

Feedback for other domains

The framework generalizes beyond coding:

• MLE-Bench (Kaggle competitions): The score is the competition metric (accuracy, F1, RMSE). Feedback includes validation set performance and error logs from training.
• ARC-AGI (pattern puzzles): The score is whether the predicted output grid matches the expected one. Feedback is binary (correct/incorrect per example).
• Any task with verifiable outputs: Mathematical proofs (checked by a proof assistant), SQL queries (checked against expected results), API interactions (checked against test harnesses).

The key requirement is a surrogate evaluator R available at search time. This need not be the final hidden-test evaluator — even a partial signal (public tests only) is enough to guide the search. The paper acknowledges that developing such evaluators is itself a challenge for some domains.

Chapter 7: Results

AB-MCTS was evaluated on four benchmarks using GPT-4o and DeepSeek-V3, with a generation budget of 128 LLM calls (2⁷).

Main results (Table 1 from the paper)

Values: GPT-4o / DeepSeek-V3. Higher is better. Budget = 128 calls.

Benchmark-specific insights

LiveCodeBench & CodeContest (competitive programming): AB-MCTS pulls ahead at budgets as small as 8 LLM calls. On CodeContest, AB-MCTS-A (Beta) with DeepSeek-V3 achieves 44.8 — the single highest score in the table. These problems benefit from a mix of diverse approaches (width) and debugging specific solutions (depth).

ARC-AGI (pattern puzzles): Repeated sampling is strongest here, reflecting ARC's nature — the LLM either "gets" the pattern or doesn't, and feedback (correct/incorrect) is too coarse for targeted refinement. Still, AB-MCTS is competitive and never catastrophically worse.

MLE-Bench (Kaggle): AB-MCTS-M achieves the best average rank (1.3) across three ML competitions. These tasks especially benefit from iterative refinement — adjusting hyperparameters, fixing data processing bugs — making depth particularly valuable.

Chapter 8: Analysis

The paper provides a fascinating analysis of tree shapes that emerge from AB-MCTS compared to baselines.

Width vs. depth in practice

The authors measure each algorithm's tree shape using the log-ratio of mean depth to mean width. This single number captures the tree's character: negative values mean wider (more exploration), positive values mean deeper (more refinement).

When does depth help most?

Depth is most valuable when:

• The LLM can generate a near-correct solution but struggles with specific details (off-by-one errors, edge cases, import issues).
• External feedback is rich and specific — compiler errors with line numbers, failing test cases with expected vs. actual output. This gives the LLM actionable information for refinement.
• The solution space is structured such that small edits can fix a mostly-correct solution (typical of coding tasks).

When does width help most?

Width is most valuable when:

• The problem requires a fundamentally different approach that incremental edits cannot reach. Sampling many diverse starting points is the only way to discover it.
• External feedback is coarse or binary (correct/incorrect, like ARC-AGI). With no detailed error signal, refinement is essentially blind rewriting.
• The per-sample success rate is low but nonzero — the coverage argument from Large Language Monkeys applies.

Scaling to large budgets

On ARC-AGI with budgets up to 512, AB-MCTS continues improving substantially between 200–500 samples, while repeated sampling begins to plateau. This is remarkable: even on a task that favors width, AB-MCTS's adaptive depth provides gains at scale that pure width cannot.

Qualitative tree analysis

The paper visualizes a real search tree from MLE-Bench (Random Acts of Pizza). The tree shows clearly adaptive branching: the root has many children (wide exploration), but certain promising branches develop significant depth (focused refinement). Failed branches (grey nodes with no score) are abandoned quickly. This is exactly the behavior the algorithm was designed to produce.

Chapter 9: Connections

AB-MCTS sits at the intersection of several important research threads in inference-time compute scaling. Let's map where each connection leads.

LATS (Language Agent Tree Search)

Zhou et al. (2024) applied standard MCTS to LLM-generated solutions, using UCT for selection and a fixed branching factor. AB-MCTS directly extends LATS by replacing UCT with Thompson sampling and introducing the GEN node for adaptive branching. LATS corresponds to the "Standard MCTS" baseline in the experiments.

Large Language Monkeys

Brown et al. (2024) formalized the power of repeated sampling: if you sample enough times, even hard problems get solved. Their "coverage" analysis shows that success probability grows as a function of the number of samples. AB-MCTS generalizes this by treating repeated sampling as a special case (flat tree) and showing that adaptive depth can improve upon pure width.

Scaling Test-Time Compute

Snell et al. (2025) analyzed how to optimally allocate inference-time compute between different strategies (revision, verification, search). AB-MCTS provides a concrete algorithm for this allocation: the Bayesian branching decision automatically distributes budget between exploration (width) and exploitation (depth) based on observed scores.

Self-Refine and Reflexion

Madaan et al. (Self-Refine, 2024) and Shinn et al. (Reflexion, 2024) pioneered iterative self-refinement with feedback. These correspond to the "Sequential Refinement" baseline — depth-only search. AB-MCTS incorporates their insight (feedback is valuable) while adding width (diversity is also valuable).

Progressive Widening

Coulom (2007) and Chaslot et al. (2008) proposed increasing the number of actions per node over time in MCTS. AB-MCTS differs because all new branches come from the same LLM (homogeneous generation), enabling principled statistical comparison between "sample new" and "refine existing." Progressive widening was designed for games with unique action semantics.

Cheat Sheet