PaperlayerView exact revision on arXiv

Source-linked AI summary

CodeTree: Agent-guided Tree Search for Code Generation with Large Language Models

Jierui Li, Hung Le, Yingbo Zhou, Caiming Xiong, Silvio Savarese, Doyen Sahoo

arXiv:2411.04329v2cs.CL

TL;DR

Code-generation agents still face challenges in planning, generating, and debugging across large search spaces. CodeTree uses a unified tree and agent-guided feedback to explore strategies and refine solutions. It reports consistent gains over strong baselines across benchmarks, including a 43.0% pass@1 result on CodeContests.

  • Problem

    Large code-generation search spaces challenge agentic approaches that must plan, generate, and debug across multiple stages.

    Source: §1 · p. 1

  • Method

    CodeTree uses a unified tree for strategy exploration and solution refinement, with Critic-guided decisions informed by execution and agent feedback.

    Source: §1 · pp. 1–2

  • Results

    CodeTree with Critic Agent achieved pass@1=43.0% on CodeContests, a 22.4% performance gain over the Resampling baseline, and outperformed all other baselines on 4 of 5 benchmarks for GPT-4o-mini and GPT-4o.

    Source: §4.2 · p. 6

  • Takeaways & Limitations

    Across diverse code-generation benchmarks, CodeTree exhibits consistent and strong performance against many baselines.

    Source: §5 · p. 8

    CodeTree requires LLMs with strong reasoning and instruction-following capabilities; smaller models such as 8B models struggled to follow complex agent-role instructions.

    Source: §6 · p. 9

Abstract

from arXiv · show

Pre-trained on massive amounts of code and text data, large language models (LLMs) have demonstrated remarkable achievements in performing code generation tasks. With additional execution-based feedback, these models can act as agents with capabilities to self-refine and improve generated code autonomously. However, on challenging coding tasks with extremely large search space, current agentic approaches still struggle with multi-stage planning, generating, and debugging. To address this problem, we propose CodeTree, a framework for LLM agents to efficiently explore the search space in different stages of the code generation process. Specifically, we adopted a unified tree structure to explicitly explore different coding strategies, generate corresponding coding solutions, and subsequently refine the solutions. In each stage, critical decision-making (ranking, termination, expanding) of the exploration process is guided by both the environmental execution-based feedback and LLM-agent-generated feedback. We comprehensively evaluated CodeTree on 7 code generation benchmarks and demonstrated the significant performance gains of CodeTree against strong baselines. Using GPT-4o as the base model, we consistently achieved top results of 95.1 on HumanEval, 98.7 on MBPP, and 43.0 on CodeContests. On the challenging SWEBench benchmark, our approach led to significant performance gains.

1 Introduction

Code generation requires executable, functionally correct code, but the large search space challenges both brute-force sampling and refinement from a small set of initial candidates. CodeTree uses a unified tree for strategy exploration and solution refinement, with agents and a Critic guiding search decisions; the paper reports consistent gains over strong baselines.

1 Introduction

  • Early approaches sampled up to 1 million code candidates per problem, while vertical refinement was bounded by the local optima of its chosen search space.Coding tasks require generated code to be executable and functionally correct.Source: §1 · p. 1
  • Thinker, Solver, and Debugger handle strategy planning, solution implementation, and solution improvement, respectively.The agents interact during tree expansion under Critic guidance.Source: §1 · pp. 1–2
  • CodeTree uses a tree to explore alternative strategies as sibling nodes and refine solutions through child nodes.The input problem is the root, and subsequent nodes represent code solutions.Source: §1 · pp. 1–2
  • At each expansion, the Critic evaluates test outputs and strategy fidelity, then decides whether to refine, abort, or accept the solution.For solutions that pass visible tests, it checks whether further improvement is warranted; for failing solutions, it judges whether debugging and refinement are promising.Source: §1 · p. 2
  • CodeTree showed significant, consistent performance gains against strong baselines across diverse code-generation benchmarks.The introduction reports top results on HumanEval+, MBPP+, and CodeContests, and significant gains on SWEBench.Source: §1 · p. 2
  • The authors also conducted ablation and qualitative analyses to identify best practices and limitations of the method.Source: §1 · p. 2

2 Related Work

Prior code-generation work explored large sets of candidates, filtered them using execution outcomes, or refined generated code with feedback. Related work also used tree search and multi-agent approaches, while CodeTree presents a unified tree-based search space spanning code-generation stages.

2 Related Work

3 Method

CodeTree represents code generation as a tree of code solutions, using specialized agents to generate strategies, implement solutions, and refine code. A Critic agent guides tree expansion and search decisions using execution-based and LLM-generated feedback.

3.1 Coding Task-Specific Agents

  • Thinker generates a problem-dependent set of high-level strategies, Solver conditions code candidates on each strategy, and Debugger revises code using Thinker reflections.The number of strategies can vary with the problem, and reflections inform code modifications.Source: §3.1 · pp. 3–4

3 Method

  • The tree starts at the problem specification, with subsequent nodes representing code solutions; each node carries feedback, strategy, and reflection attributes.Adding a node involves generating code from a strategy and evaluating the resulting solution.Source: §3.2 · p. 4

3.2 Tree Expanding with Critic Agent

  • The Critic dynamically expands the tree, scoring execution outputs and strategy adherence, checking whether visible-test-passing solutions appear robust, and choosing to refine, abort, or accept.Accepting terminates the search; aborting prunes a node and returns exploration to its siblings.Source: §3.2 · pp. 4–5

3.3 Multi-agent Collaboration

  • Task-specific agents collaborate with the Critic during search, and inference limits exploration steps to control computation overhead.When the Critic accepts a solution or the step limit is reached, CodeTree selects a candidate by its evaluation score.Source: §3.3 · p. 5

4 Experiments

Across the evaluated benchmarks, CodeTree with a Critic Agent outperformed strong baselines in most settings, while results varied with search strategy and problem difficulty. Ablations and budget-limited comparisons highlight the roles of Critic tasks and continued tree exploration; CodeTree also gained on SWEBench.

4.1 Experimental Setup

  • Experiments used pass@1 and a 20-sample generation budget across function- and program-implementation benchmarks, comparing against Direct, CoT, Reflexion, MapCoder, and Resample.Benchmarks included HumanEval, MBPP, their EvalPlus variants, CodeContests, and APPS.Source: §4.1 · p. 5

4.2 Main Results

  • CodeTree with Critic Agent outperformed all baselines on 4 of 5 benchmarks for GPT-4o-mini and GPT-4o; on CodeContests it scored pass@1=43.0%, a 22.4% gain over Resampling.The reported CodeContests result is for competition-level coding tasks.Source: Table 2 · p. 6, §4.2 · p. 6
  • CodeTree-BFS almost always outperformed CodeTree-DFS, while Resampling achieved the best results on 4 benchmarks with Llama-3.1-8B.The authors suggest small models may be less suited to multi-agent frameworks requiring distinct task roles and instructions.Source: §4.2 · p. 6

4.4 Analysis by Problem Difficulty

  • On APPS, CodeTree performed better on simpler difficulty levels but still struggled with Competition-level problems.The authors hypothesize that the 20-sample generation budget is limited for highly complex problems.Source: Figure 3 · p. 7, §4.4 · p. 7

4.3 Analysis of Search Strategies

  • Compared with BFS/DFS settings of d=3 and w=3, wider BFS search and d=2 improved pass@1; deeper DFS refinement did not boost performance significantly.w=4 worked better for HumanEval and w=5 for CodeContests, which the authors connect to more complex problems benefiting from more coding strategies.Source: §4.3 · pp. 6–7, Table 3 · p. 7

4.5 Ablation Study

  • All Critic Agent tasks were important; node abort and solution verification had the largest impacts on final performance.The ablation study reported pass@1 on HumanEval and HumanEval+ using GPT-4o-mini.Source: Table 4 · p. 7, §4.5 · p. 7

4.6 Search Efficiency

  • Within a 20-solution budget, CodeTree outperformed other methods and maintained relatively high pass@1 with fewer than 9 samples.On CodeContests, it began lower below 5 samples but improved during later tree exploration.Source: Figure 4 · p. 8, §4.6 · p. 8

4.7 Results on SWEBench

  • On SWEBench, CodeTree led to significant performance gains over related approaches including CoT and Reflexion.The authors present this as evidence of versatility for repository-level code generation.Source: Table 5 · p. 8, §4.7 · p. 8

5 Conclusion

CodeTree uses a unified tree-based search space and a Critic Agent to guide node scoring, expansion, and termination. Its collaboration among Thinker, Solver, and Debugger agents achieved consistent, strong performance against baselines across diverse code-generation benchmarks.

5 Conclusion

  • CodeTree combines a unified tree search and Critic Agent decisions with collaboration among Thinker, Solver, and Debugger agents.The Critic guides scoring, expansion, and termination to find solutions within a limited generation budget.Source: §5 · p. 8
  • Across diverse code-generation benchmarks, CodeTree showed consistent and strong performance against many baselines.Source: §5 · p. 8

6 Limitations

CodeTree requires models with strong reasoning and instruction-following capabilities, and smaller models can provide noisy agent feedback. The approach also incurs additional LLM costs and focuses on functional correctness rather than other code qualities.

6 Limitations

  • Smaller models such as those with 8B parameters may struggle with complex agent instructions, producing noisy Critic feedback for scoring or self-reflection.The paper also identifies a broader requirement for strong reasoning and instruction-following capabilities.Source: §6 · p. 9
  • Agent-generated guidance adds cost because LLMs must process long code contexts and generate more output tokens.Source: §6 · p. 9
  • CodeTree targets functional correctness and does not address qualities such as readability or efficiency of generated code.The paper proposes incorporating more holistic feedback and diverse exploration to improve these qualities.Source: §6 · p. 9

7 Ethical Considerations

The paper identifies information leakage and system-security risks from executing the pipeline. It describes risks from sending execution errors to commercial LLMs and running generated programs locally, and says a virtual execution environment was planned for the codebase.

7 Ethical Considerations

  • Sending execution-error details to commercial LLMs might expose private information, including execution paths, usernames, and package versions.Source: §7 · p. 9
  • Running AI-generated programs locally might put a user’s system and data at risk.Source: §7 · p. 9
  • The authors said they would add a virtual environment for code execution by the time they published the codebase.Source: §7 · p. 9

A.1 Qualitative Results

A HumanEval-36 example shows the Critic rejecting a solution that passed visible tests, then the Debugger using the Critic’s feedback to produce a solution that passed hidden tests. The example illustrates strategy selection, verification, and refinement in sequence.

A.1 Qualitative Results

  • W1 passed visible but failed hidden tests, while Critic-guided refinement produced W2, which passed hidden tests.The authors say this example demonstrates the Critic’s correct judgment of a false-positive solution and guidance toward a correct one.Source: §A.1 · p. 12
  • The Thinker generated five strategies, and the Solver implemented the first as W1 for execution on visible tests.Source: §A.1 · p. 12
  • After W1 passed visible tests, the Critic recommended considering cases such as zero and negative integers; hidden-test evaluation confirmed W1 was incorrect.Source: §A.1 · p. 12
  • The Debugger used the Critic’s judgment and reason as environmental feedback, alongside the problem and W1, to produce W2.Source: §A.1 · p. 12

A.2 Other Prompts

This appendix presents prompts for CodeTree’s agents and a HumanEval-36 example. The prompts specify strategy generation, code implementation, solution verification and scoring, and reflection; the example shows the Critic’s feedback informing a Debugger revision.

A.2 Other Prompts

  • Figure 8 presents prompts for Critic test-case and solution scoring and Thinker reflection, which were not used in the HumanEval-36 example.Source: §A.2 · p. 12, Figure 8 · p. 16
  • The Thinker prompt requests multiple distinct natural-language strategies, while its reflection prompt asks for self-contained alternatives based on code and test feedback.Source: §A.2 · pp. 13, 16
  • The Solver prompt requests a complete Python implementation only, while the Critic prompts score test outputs and solutions or verify whether a visible-test-passing solution needs modification.Source: §A.2 · pp. 13, 14, 16
  • The accompanying figures show strategy generation, visible-test execution, Critic rejection, and Debugger revision, including a correction that passes hidden tests.Source: Figure 5 · p. 13, Figure 6 · p. 14, Figure 7 · p. 15
  • In the example, the Critic flags missed divisibility and input edge cases, and the Debugger prompt uses the implementation, tests, execution feedback, and a hint to request a revised implementation.Source: §A.2 · pp. 14–15, Figure 6 · p. 14
Loading 2411.04329v2…