Source-linked AI summary
Thinking Before Running! Efficient Code Generation with Thorough Exploration and Optimal Refinement
Xiaoqing Zhang, Yuhan Liu, Flood Sung, Xiuying Chen, Shuo Shang, Rui Yan
TL;DR
Code-generation refinement and test-time computation can incur substantial token and computational overhead. ThinkCoder combines exploration with optimal refinement, and ReST trains LLMs on successful exploration trajectories. The authors report superior performance over current composite models across LLM sizes while optimizing computational efficiency.
Problem
Code-generation refinement can incur substantial token overhead and trial-and-error costs, while multi-agent methods such as MapCoder consume considerable computational resources.
Source: §1 · pp. 1–2Method
ThinkCoder combines LLM-based exploration with a non-LLM CodeVerifier for refinement, while ReST uses successful exploration trajectories to train LLMs and reduce test-time computation overhead.
Source: §1 · p. 2Results
Takeaways & Limitations
Fine-tuning on successful ThinkCoder trajectories is reported to enhance code-generation performance while reducing refinement iterations.
Source: §7 · p. 8The testing pool lacks effective noise-cleaning strategies, and its test diversity relies heavily on LLM generation capabilities.
Source: Limitations · p. 9
Abstract
from arXiv · showhide
Code generation is crucial in software engineering for automating the coding process efficiently. While test-time computation methods show promise, they suffer from high latency due to multiple computation rounds. To overcome this, we introduce \textbf{ThinkCoder}, a framework that combines thorough exploration with optimal refinement. The exploration phase diversifies the solution space by searching for potential solutions, followed by a refinement phase that enhances precision. This approach allows us to select the best solution through careful consideration before taking action, avoiding excessive trial and error. To further minimize test-time computation overhead, we introduce preference-driven optimization with Reinforced Self-Training (ReST), which uses exploration trajectories from ThinkCoder to guide LLM's evolution. This approach enhances LLM's exploration efficiency via preference learning, cutting costs while maintaining accuracy. ThinkCoder boosts the performance with a single LLM, excelling on benchmarks like HumanEval and MBPP. Compared to SOTA models, it improves Pass@1 by 3.0\% over MapCoder with just 6.4\% of the computation cost. Against AgentCoder, ThinkCoder achieves a 0.5\% higher Pass@1 after 2 rounds, outperforming AgentCoder's 5 rounds. Additionally, ReST with success trajectories enhances efficiency, allowing models like LLaMA2-7B to achieve competitive results using only 20\% of the computational resources. These results highlight the framework's effectiveness and scalability.
1 Introduction
Code-generation refinement can incur substantial token overhead and trial-and-error costs, while multi-agent approaches such as MapCoder consume considerable computational resources. ThinkCoder combines exploration with refinement, and ReST uses exploration trajectories to reduce test-time computation costs.
1 Introduction
- Existing code-generation refinement can incur substantial token overhead and trial-and-error costs, and multi-agent methods such as MapCoder use considerable computational resources.Source: §1 · pp. 1–2
- The framework combines thorough exploration and optimal refinement, with the CodeVerifier independently guiding optimizations using an evolving pool of tests.The Exploration Agent generates diverse code and test cases; the verifier uses the pool to guide subsequent optimizations.Source: §1 · p. 2
- ThinkCoder integrates an LLM-based exploration agent with a non-LLM CodeVerifier and a self-evolving testing pool to speed code verification and improvement.The framework is reported to achieve state-of-the-art code-generation performance while optimizing computational efficiency.Source: §1 · p. 2
- ReST uses successful exploration trajectories for preference-driven LLM training to reduce computational overhead during testing.The authors report notable performance and efficiency improvements for models such as LLaMA2-7B.Source: §1 · p. 2
- ThinkCoder demonstrates state-of-the-art code-generation performance across multiple benchmark datasets and LLM scales while optimizing computational efficiency.Source: §1 · p. 2
2 Related Work
Related work covers test-time computation and verification, multi-agent code generation, and instruction tuning. ThinkCoder is presented against approaches that can incur high compute costs or lack comprehensive fine-tuning guidance.
2 Related Work
- Test-time methods use self-play and additional verifiers to improve performance despite computational overhead; task-based testing is presented as a simpler verifier aligned with task requirements.Source: §2 · p. 2
- Multi-agent systems divide code generation across specialized agents, while MapCoder's overhead and reliance on sample I/O can limit test coverage and robustness.ThinkCoder is contrasted as using efficient self-verification without excessive computational costs; its process proceeds from exploration to recursive refinement.Source: §2 · p. 2, Figure 1 · p. 3
- Instruction-tuned code models draw on real-world or synthetic data, but existing methods often lack comprehensive guidance on data collection, volume management, and fine-tuning strategies.Source: §2 · p. 2
- The paper presents a detailed Llama2-7B fine-tuning strategy using instruction-tuning data collected via CodeQwen as a transparent, reproducible methodology.Source: §2 · pp. 2–3
3 ThinkCoder
ThinkCoder generates candidate solutions and tests, then uses a CodeVerifier and a testing pool to guide iterative refinement. Exploration expands the candidate space, while testing, feedback, and verification shape subsequent improvements.
3.1 Overall Architecture
- The architecture generates candidate solutions and tests, selects the code with the highest testing-pool pass rate, and updates instructions with failed-test feedback; refinement stops when the pass rate exceeds θ.Source: §3.1 · p. 3
3.2 Thorough Exploration
- Thorough exploration generates k answers simultaneously at increased temperature, with initial iterations producing diverse code and tests and later iterations targeting specific errors.The first iteration aims to include regular, boundary, and performance tests.Source: §3.2 · p. 3
3.3 Optimal Refinement
- Optimal refinement selects a best solution through self-verification, updates instructions from the problem, current best code, feedback, and tests, and retains solutions with improved pass rates.Offline success trajectories are verified with ground-truth tests, while LLM-generated tests support reflection on errors.Source: §3.3 · pp. 3–4, Figure 2 · p. 4
3.4 Exploration Agent& CodeVerifier
- The Exploration Agent generates diverse code and tests, while the CodeVerifier executes snippets and tests locally and uses multiprocessing for parallel validation.Source: §3.4 · p. 4
3.5 Testing Pool
- The testing pool contains regular, boundary, and performance tests to assess complex cases; AST-based deduplication removes redundant tests.Source: §3.5 · p. 4
4 Preference-driven Optimization
ThinkCoder collects successful exploration and reflection trajectories, then uses them to fine-tune models for more efficient code generation. Its data pipeline samples diverse candidates, verifies them, and tests failed candidates again with verifier feedback.
4 Preference-driven Optimization
- ThinkCoder collects improving code-and-test trajectories to fine-tune the model and refine exploration direction.The paper describes collecting improvements across iterations to help the LLM generate better code and tests and reduce exploration time.Source: §4 · p. 4
4.1 Data Collection
- Temperature-based generation samples diverse candidates, retaining verifier-passing solutions in DM0.The verifier returns pass rate and feedback; candidates with pass_rate == 1 enter DM0.Source: §4.1 · pp. 4–5
- Reflection-based generation adds verifier feedback to failed samples and retains corrected outputs that pass in RM0.The reflection trajectory uses the feedback-augmented input and verifier results to select successful corrections.Source: §4.1 · p. 5
4.2 Model Training and Inference
- ReST trains on successful exploration and reflection trajectories; inference tests both direct generation and ThinkCoder refinement.The framework configuration uses k explorations and n refinements with varying temperatures.Source: §4.2 · p. 5
5 Experiments
The experiments evaluate ThinkCoder on standard and extended code-generation benchmarks, compare it with existing methods, and specify separate configurations for base-model and SOTA comparisons.
5.1 Dataset
- Evaluation covers MBPP, HumanEval, extended variants, LiveCodeBench, and CodeContests, with extended sets containing substantially more test cases.MBPP-ET and HumanEval-ET have 80 and 35 times more test cases than their original datasets, respectively.Source: §5.1 · pp. 5–6
5.2 Baselines
- ThinkCoder is compared with MapCoder, AgentCoder, LDB, and LPW, using several instruction-tuned general and code-specific base models.The evaluation includes GPT-4-Turbo and GPT-4o as well as LLaMA2-7B-Chat, Kimi2, and CodeQwen1.5-7B-Chat.Source: §5.2 · p. 6
5.3 Implementation details
- Base-model experiments use k=5, t=0.5, n=5, and θ=0.8; SOTA comparisons use k=20, t=1.0, n=2, and θ=1.0.The exploration agent generates m=3 test cases per code in both settings.Source: §5.3 · p. 6
6 Evaluation
ThinkCoder improves code-generation results, with larger gains on standard test sets than extended ones, and combines exploration with refinement. ReST trajectories improve LLama2-7B, while generated test-pool accuracy rises alongside solution performance.
6.1 Performance on Code Generation
- 4%-70%: ThinkCoder improves MBPP and HumanEval performance after five refinement iterations, while gains on extended variants are 1%-35%.With k=20, ThinkCoder surpasses AgentCoder, which requires five refinement rounds, in two iterations; performance saturates in complex scenarios.Source: §6.1 · p. 7
6.2 Ablation Study of Self-exploration and Self-refinement
- Self-exploration yields greater performance gains than self-refinement, but refinement improves efficiency when paired with exploration.An excessive exploration budget such as k=20 is inefficient, while a moderate increase can accelerate refinement.Source: §6.2 · p. 7
6.3 Performance with ReST Training
- ReST fine-tuning LLama2-7B on successful ThinkCoder trajectories improves over SFT and the untuned model; at t=0.5, one refinement round is optimal.The authors report that ReST reduces computational overhead.Source: §6.3 · pp. 7–8
6.4 Evolution of Generated Code and Tests
- Code and test-pool accuracy improve simultaneously across refinement rounds on MBPP with CodeQwen1.5-7B-Chat.The ground-truth test set evaluates solutions, while ground-truth code evaluates the testing pools.Source: Figure 4 · p. 7, §6.4 · p. 8
6.5 Hyperparameters Study
- Exploration gains are more significant than refinement gains, and Pass@1 rises with more refinement iterations; Pass@k eventually stabilizes as k increases.Optimal temperature differs by model: it rises for LLama2-7B-Chat, CodeQwen1.5-7B-Chat, and Kimi, but lower temperatures suit GPT-4-Turbo and GPT-4o.Source: §6.5 · p. 8
6.6 Cost Analysis
- ThinkCoder uses one LLM-based agent and a non-LLM CodeVerifier, with lower time overhead and more efficient token use than MapCoder at similar Pass@1.MapCoder uses four LLM-based agents executed serially; token usage is used to measure cost.Source: §6.6 · p. 8
7 Conclusion
ThinkCoder combines LLM exploration with a non-LLM CodeVerifier, then uses successful exploration trajectories to fine-tune the LLM. Across varying model sizes, it outperforms current composite models while reducing token use and refinement iterations.
7 Conclusion
- ThinkCoder combines LLM exploration with a non-LLM CodeVerifier to improve code-generation efficiency.The conclusion says this design greatly reduces token usage.Source: §7 · p. 8
- Successful exploration trajectories are captured to fine-tune the LLM, reducing refinement iterations for optimal solutions.Source: §7 · p. 8
- Evaluations across LLMs of varying sizes show superior performance over current composite models.Source: §7 · p. 8
Limitations
The testing pool lacks effective noise-cleaning strategies, and its diversity depends heavily on LLM-generated tests. ReST fine-tuning converges rapidly, while gains in exploring the overall solution space remain limited.
Limitations
- The testing-pool approach selectively expands tests but lacks an effective strategy for cleaning noise.Source: Limitations · p. 9
- Testing-pool diversity relies heavily on the LLM’s test-generation capabilities.Source: Limitations · p. 9
- ReST validation shows rapid convergence, with fine-tuning teaching the LLM to find good answers but only limited improvement in overall solution-space exploration.Source: Limitations · p. 9
- The authors identify on-policy training that combines data collection with fine-tuning as future work for increasing answer diversity and online exploration.Source: Limitations · p. 9
A.1 Algorithm of ThinkCoder
Algorithm 1 presents ThinkCoder’s pseudocode.
A.1 Algorithm of ThinkCoder
- Algorithm 1 shows the pseudo-code of ThinkCoder.Source: §A.1 · p. 11
A.2 Exploration Budget Control
The appendix describes budget control, test selection, agent and verifier inputs, and examples of exploration, refinement, and testing-pool contents. The examples show refinement improving a solution’s testing-pool success rate from 0.217 to 0.356.
A.2 Exploration Budget Control
- With θ = 0.8, tasks requiring correction per iteration fall by nearly half while Pass@1 progressively improves; higher thresholds reduce runtime but risk leaving incorrectly validated tasks uncorrected.Source: §A.2 · p. 11, Figure 5 · p. 11
A.3 Exploration for Tests
- Exploration runs for k iterations, deduplicates generated tests by AST, and adds the tests to the pool when the iteration improves the global best pass rate; refinement stops when that rate exceeds θ.Source: §A.3 · p. 11
- Increasing the number of test cases m improves Pass@1 but lengthens execution time; the experiments choose m to balance time efficiency and performance.The observation uses 60 MBPP tasks and the Kimi base model.Source: §A.3 · p. 11, Table 8 · p. 11
A.4 Prompts for Exploration Agent
- The exploration agent receives the problem, best solution, ground-truth and generated tests, and test-failure feedback.Source: §A.4 · p. 11
A.5 Inputs for CodeVerifier
- The executable Python environment receives the code and all test cases in the testing pool.Source: §A.5 · p. 11
A.6 Example of Thorough Exploration
- The exploration agent explores multiple solutions.Source: §A.6 · p. 12
A.7 Example of Optimal Refinement
- 0.217 to 0.356: a refinement iteration uses the prior best solution, failed tests, and unmet feedback, then revalidates against the testing pool.The updated solution is selected as best after its success rate improves.Source: §A.7 · p. 12
A.8 Example of Testing Pool
- The testing pool contains test cases generated through multiple ThinkCoder explorations.Source: §A.8 · p. 12