ESMake Every Expansion Count

Make Every Expansion Count:Learning to Search Efficiently
in LLM Code Repair

Elad Sharony1,2, Ido Greenberg2, Jialin Song2, Shie Mannor1,2, Gal Chechik2, Eli Meirom2

1Technion2NVIDIA Research

arXiv Code coming soon Dataset coming soon

Abstract

LLM code-repair systems fix faulty programs by repeatedly generating revisions and running unit tests. Each revision is costly, making the choice of which repair to pursue important for search efficiency. The adaptive policies we compare against make this choice using test pass rates and search statistics, without reading the detailed test output or proposed edit descriptions. We introduce CodeTriage, a learned search policy whose small 0.6B-parameter model reads test output, failure diagnoses, and proposed edits to predict how many repair steps each edit could require to reach a fix. These predictions prioritize all pending edits across programs before generating and testing the revised code. On held-out repair tasks, CodeTriage finds fixes with 21% fewer program revisions and 12% fewer output tokens than the strongest adaptive baseline on each metric. The same model transfers without retraining to eleven unseen repair LLMs, achieving the highest success rate averaged across budgets for each one. Ablations show the value of interpreting proposed edits together with test output and diagnoses, and of learning to compare edits across different programs. Additionally, we release a corpus of 12.2 million program versions generated by four LLMs, organized into branching repair sequences with proposed edits and test results.

(a) Exhaustive-tree evaluation

Qwen3-Coder-30B trees from unseen tasks

(b) Cross-proposer deployment

Live search on eleven unseen proposers, averaged across proposers

Figure 1. More fixes with less search. Across budgets, CodeTriage solves more problems with fewer expansions than any baseline, including adaptive baselines. (a) Exhaustive-tree evaluation on Qwen3-Coder-30B trees from unseen tasks. (b) Live search on eleven unseen proposers, averaged across proposers. CodeTriage is shown in green; the shaded region is its gain over the strongest adaptive baseline at equal budgets, and B labels mark budgets along its curve. Live search is capped at 36 expansions. Curve coordinates reconstructed from the paper figures are approximate (≈); hover a marker for its values.

1 Problem setting

A repair LLM, the proposer, reads the test output (execution feedback), explains the failure, and describes candidate repairs as proposed edits; each specifies what to change and why, before the revised program is generated. The frontier contains every proposed edit that has yet to be tried, including alternatives for programs explored earlier. A search policy chooses one edge from the frontier. Expanding it generates the revised program, runs the full test suite, and adds the evaluated result as a child node, whose own k proposed edits join the frontier. This operation is one expansion, and the objective is to minimize the number of expansions required to find the first solution.

The fraction of passing tests does not reveal why a program fails or how difficult it will be to repair, and the best edit for one program may be less promising than the second-best edit for another. A search policy should therefore compare proposed edits across programs before paying to generate and test the revised code. CodeTriage changes the search policy while leaving the proposer unchanged. Below it is replayed on an exhaustive tree, in which every edge through depth D has already been expanded, so its decisions can be compared with every baseline on the same candidates and outcomes.

Recorded repair search
max_tasks_completed(n, T, task_durations, m)Scheduling tasks among adventurers within a time budget.

Given n task durations, a total time T, and m adventurers who each take at most one task, return the maximum number of tasks that can be completed. The seven supplied tests match KodCode’s reference solution: sort durations ascending, take tasks while the running total stays within T, and stop at m.

Prompt as seen by the repair LLM
def max_tasks_completed(n: int, T: int, task_durations: List[int], m: int) -> int:
    """
    Determines the maximum number of tasks that can be completed within a given total time
    by scheduling the tasks optimally among the available adventurers.

    Parameters:
    n (int): Number of tasks
    T (int): Total available time
    task_durations (list of int): List of durations for each task
    m (int): Number of available adventurers

    Returns:
    int: Maximum number of tasks that can be completed
    """

from solution import max_tasks_completed

def test_example_1():
    assert max_tasks_completed(4, 10, [2, 3, 5, 7], 3) == 3

def test_example_2():
    assert max_tasks_completed(5, 8, [1, 2, 3, 4, 5], 2) == 2

def test_single_task():
    assert max_tasks_completed(1, 5, [4], 1) == 1

def test_insufficient_time():
    assert max_tasks_completed(3, 5, [2, 4, 6], 3) == 1

def test_exact_fit_time():
    assert max_tasks_completed(3, 6, [2, 2, 2], 3) == 3

def test_more_adventurers_than_tasks():
    assert max_tasks_completed(3, 10, [1, 2, 3], 5) == 3

def test_large_input():
    assert max_tasks_completed(10, 20, [2, 2, 2, 2, 2, 2, 2, 2, 2, 2], 4) == 4
KodCode-V1 · Code_Contests_39361_C · held-out test split · Qwen3-Coder-30B repair tree · k = 3 · depth 4 · 121 programs · 2 fixes (identical code), both at depth 4 · CodeTriage checkpoint anchor-06b-s2

(a) Repair tree

Program → expanded edit (with its priority) → generated child, tested on the task’s supplied tests

Scroll horizontally to inspect all branches ↔

Ring: fraction of tests passed3Expanded edit, numbered by step, with its priority+2Proposals never expanded (listed in the frontier)Inspected branchFix: all tests pass
Expansion

(b) Frontier

Ranked by s(e)

Every unexpanded proposal across all evaluated programs, ordered by the recorded priority s(e). Select one to inspect it; its parent program highlights in the tree.

(d) Footprint on the full tree

Programs one policy generated before its first fix, on the same 121-node tree · select a policy

(c) Program and tests

Figure 2. A recorded CodeTriage search on one held-out repair tree. The task is KodCode-V1 Code_Contests_39361_C; the tree was generated by Qwen3-Coder-30B with k = 3 edits per program and depth 4, giving 121 programs of which 2 (identical code) are fixes, both at depth 4. (a) Programs the policy expanded, p₀ … p₁₄ in expansion order, with rings showing their test pass rate; numbered boxes are the expanded edits ej (j is the child’s node id in the tree) with their recorded priority, and a +n badge counts a program’s proposals that were never expanded. The two depth-1 edits that raise the pass rate to 5/7 score highest but both drop the adventurer limit and lead nowhere; the 3/7 edit that keeps it is the live one. Ten of fourteen expansions are dead ends, yet every baseline needs at least 35 on this tree except one AB-MCTS seed that needs 13. (b) The unexpanded proposals, ranked by s(e). (c) Program diff, recorded test outcomes, pytest feedback, and the generator’s diagnosis for the selected program or proposal; for proposals, what the exhaustive tree knows about the branch is shown separately. (d) One policy’s footprint at a time on the same tree: CodeTriage, inference seed 0 of each stochastic baseline, and the deterministic BFS, DFS, and pass-rate. Recorded costs are AB-MCTS 13–65, REx 42–58, BG-MCTS 52, DFS and pass-rate 84, BFS 94, and random 31–92 across seeds. Shading encodes expansion order from light to dark. DFS and pass-rate orders were not exported and are replayed on the exhaustive tree with deterministic tie-breaking, reproducing their recorded cost of 84; on this tree they expand the same 84 programs in a different order. The random replay differs from the recorded runs, whose random stream depended on batch composition. This root sits in the favourable tail; the test-split means are 11.54 expansions for CodeTriage and 14.59 for the best baseline. All 121 programs terminate and were re-executed against the task’s tests; check.py repeats this for the expanded ones.

2 Method

Exploring every proposed edit is costly because each expansion requires generating a revised program and running its unit tests. Under a limited budget, the policy must prioritize edits across programs and repair depths before their outcomes are known. CodeTriage makes these comparisons by predicting a distribution over each edit’s solve distance: the fewest successive expansions needed to reach a fix through that edit. The model reads the proposed edit together with the parent state, which contains the execution feedback, the failure diagnosis, and four numerical features: pass rate, normalized depth, failed-test fraction, and normalized program length. The edge’s remaining depth enters the priority rule (Eq. 1) but is not an input to the model.

Priority-model architecture

Parent state and proposed edits

Priority model

Frontier priorities

Ranked frontier

Each frontier edge is serialized with its parent’s execution feedback and diagnosis; the distributions and priorities shown are the recorded ones.

Figure 3. CodeTriage priority model on the recorded frontier. One candidate per remaining depth among the 29 frontier edges before expansion 14 in Figure 2: the highest-priority edge with b = 3 (e₂₈), with b = 2 (e₁₁₃), and with b = 1 (e₈₅, the fix the search reaches next), each with its recorded five-bin solve-distance distribution and priority. Each frontier edge is serialized as parent execution feedback (at most 375 tokens), parent diagnosis (470), and edit text (175), separated by section markers, and passed through Qwen3-0.6B-Base[1], a decoder-only language model whose base weights are frozen and whose q/k/v/o projections carry rank-8 LoRA adapters[2]. The last-token hidden state is concatenated (⊕) with four standardized parent features, and an MLP head predicts a (D+1)-bin distribution over the solve distance t. The Priority step shows how the distribution becomes a score: bins within the remaining depth b accumulate into F(m) = P(t ≤ m), bins beyond b are discarded, and s(e) is the mean of those F(m) (Eq. 1); best-first search expands the argmax. Select a candidate on either side to trace it. Feature values are shown before standardization.

From solve distance to search priority

For an edge e = (v, u), the solve distance t(e) counts the expansion of e plus the shortest path from its child u to an observed solution. If u is at depth d in a tree of maximum depth D, then b = D − d + 1 is the largest solve distance observable for e. The model outputs probabilities (p1, …, pD, p>D), and F(m) = Σj ≤ m pj is the predicted probability of finding a fix within m expansions. The priority should favor edits that are likely to lead to a fix within the remaining depth and require fewer repair steps, so it averages these probabilities over each available number of steps, the normalized area under F up to b:

s(e) =1bb∑m = 1F(m).(1)

Averaging gives more weight to fixes predicted sooner and no weight to distances beyond the remaining depth, and dividing by b keeps scores on the same [0, 1] scale across depths. Search is deterministic: at each step, it expands the frontier edge with the largest score. Training uses a discrete-survival likelihood[3]. When a solution is observed within the remaining depth, the exact target t ≤ b contributes −log pt; otherwise the tree reveals only that t > b, a censored observation, and the loss sums the probability mass beyond b, −log(1 − F(b)), rather than treating the branch as permanently unsolvable.

3 Experiments

Exhaustive-tree evaluation. Policies are evaluated on 5,001 complete Qwen3-Coder-30B repair trees from 982 held-out tasks; every policy encounters the same candidates and outcomes, isolating the quality of its search decisions. Every tree contains a solution, so the primary cost is the mean number of expansions to a solution, E[T], reported alongside output tokens through the first fix. CodeTriage is trained with multiple seeds and a single checkpoint is selected on validation performance; stochastic baselines are averaged over inference seeds; ablations and diagnostics use the 4,996-tree validation split. Baselines are random, BFS, DFS, and pass-rate, and the adaptive REx[4], AB-MCTS[5], and BG-MCTS[6], each with its published or validation-tuned settings. Test uncertainties are bootstrap standard errors clustered by the 982 source tasks.

Cross-proposer deployment. The same policy is deployed without retraining or proposer-specific tuning in live search with eleven unseen proposers. Live roots are not screened for solvability, so results report success S(B) and realized cost C(B) under a budget B, where an unsuccessful run consumes the full budget, and AUSC(B), the normalized area under the success curve. Every deployment root is retained and each proposer is weighted equally.

Reported values from the manuscript tables. Pareto and calibration curves are reconstructed from the paper figures; their coordinate readouts are approximate (≈).

On the exhaustive-tree test set, CodeTriage uses 21% fewer expansions than AB-MCTS, the strongest adaptive baseline (11.54 compared with 14.59), and 30% fewer than random search (16.54). It uses 12% fewer output tokens than REx, the most token-efficient adaptive baseline (14.96k compared with 17.03k), and 21% fewer than random search. Averaged across the eleven unseen proposers, it matches or exceeds every adaptive baseline’s success rate at every reported budget at an equal or lower mean expansion cost; at a budget of 16 expansions it solves 84.8% of problems compared with 83.8% for REx, using 4.60 expansions on average compared with 4.79. Its AUSC(36) is the highest for all eleven proposers, with gains of up to 2.13 percentage points that are statistically significant on nine, and the mean across proposers improves from 81.60 to 82.50.

Matched ablations on the validation split (Table 1) locate the gains. Removing parent-state text raises cost by 16% (11.80 to 13.64 expansions), removing the proposed-edit text by 22% (14.38), and removing both by 45% (17.14). A ranker trained only on sibling pairs needs 13% more expansions than the distributional model (13.32 compared with 11.79), and extending the same loss across the mixed-parent, mixed-depth frontier recovers most of that gap. Holding the trained model fixed, scoring by the one-step probability p1 alone raises cost by 15% and ranking separately within each depth by 18%. Full baseline configurations, uncertainty estimates, and priority-model overhead are reported in the manuscript and appendix.

4 Repair-tree corpus

The corpus consists of exhaustive repair trees that record the proposed edits, the programs they produce, and the test results for those programs. For each proposed edit with a successful repair path, the tree gives its solve distance; branches cut off at the depth limit provide lower bounds. The tasks are approximately 8,700 Python programming problems from the hard subset of KodCode-V1[7], spanning Codeforces[8], TACO[9], Code Contests[10], and APPS[11]. Qwen3-Coder-30B[1] generates diagnoses and up to k = 3 proposed edits per node and up to five initial solutions per problem, with maximum depth D = 4; the resulting corpus contains 45,424 trees. Matched corpora from Devstral-Small-2[12], Qwen3-Coder-Next[13], and Nemotron-Cascade2[14] expand the same initial programs under the same task-level splits, isolating the effect of proposer-generated repair trajectories. Together, the four corpora contain 12.2 million executed nodes.

Table 2. Corpus statistics before research-specific filtering.
Repair proposerExecuted nodesMean nodes / root
Qwen3-Coder-30B3,967,51687.3
Devstral-Small-23,634,26280.0
Qwen3-Coder-Next2,329,48151.3
Nemotron-Cascade22,256,18549.7
Total12,187,444—

Data are split by programming problem: all trees derived from a problem belong to the same train, validation, or test split. A node solves the task when its pass rate is 1.0. Trees with at least one reachable fix are retained intact, so branches that do not reach a solution within their remaining depth are preserved as censored observations rather than labeled unsolvable.

The corpus and its dataset card will be released together with the code.

Citation

@article{sharony2026codetriage,
  title   = {Make Every Expansion Count: Learning to Search Efficiently in LLM Code Repair},
  author  = {Sharony, Elad and Greenberg, Ido and Song, Jialin and Mannor, Shie and Chechik, Gal and Meirom, Eli},
  journal = {arXiv preprint},
  year    = {2026}
}

References

  1. Yang, A., et al. (2025). Qwen3 technical report. arXiv:2505.09388.
  2. Hu, E. J., et al. (2021). LoRA: Low-rank adaptation of large language models. arXiv:2106.09685.
  3. Tutz, G., & Schmid, M. (2016). Modeling discrete time-to-event data. Springer.
  4. Tang, H., Hu, K., Zhou, J., Zhong, S., Zheng, W.-L., Si, X., & Ellis, K. (2024). Code repair with LLMs gives an exploration–exploitation tradeoff. NeurIPS.
  5. Inoue, Y., Misaki, K., Imajuku, Y., Kuroki, S., Nakamura, T., & Akiba, T. (2026). Wider or deeper? Scaling LLM inference-time compute with adaptive branching tree search. NeurIPS.
  6. Miyamoto, S., Oba, D., & Okazaki, N. (2026). Aligning tree-search policies with fixed token budgets in test-time scaling of LLMs. arXiv:2602.09574.
  7. Xu, Z., Liu, Y., Yin, Y., Zhou, M., & Poovendran, R. (2025). KodCode: A diverse, challenging, and verifiable synthetic dataset for coding. Findings of ACL 2025.
  8. Codeforces. Codeforces. codeforces.com.
  9. Li, R., et al. (2023). TACO: Topics in algorithmic code generation dataset. arXiv:2312.14852.
  10. Li, Y., et al. (2022). Competition-level code generation with AlphaCode. Science.
  11. Hendrycks, D., et al. (2021). Measuring coding challenge competence with APPS. arXiv:2105.09938.
  12. Mistral AI (2025). Devstral Small 2. huggingface.co.
  13. Cao, R., et al. (2026). Qwen3-Coder-Next technical report. arXiv:2603.00729.
  14. Yang, Z., et al. (2026). Nemotron-Cascade 2: Post-training LLMs with cascade RL and multi-domain on-policy distillation. arXiv:2603.19220.