Thinking Before Running! Efficient Code Generation with Thorough Exploration and Optimal Refinement

2024-12-30

arXiv.org

AI Summary

To address high test-time computational overhead and low trial-and-error efficiency—leading to inference latency in large language model (LLM)-based code generation—this paper proposes ThinkCoder, a two-stage “explore-refine” paradigm. The first stage performs deep forward reasoning via diverse solution-space exploration and optimal-solution selection. The second stage introduces preference-driven Reinforced Self-Training (ReST), leveraging successful execution trajectories to optimize the generation policy *prior* to inference. ThinkCoder enables efficient single-LLM deployment and substantially reduces test-time computation. Experiments show that on HumanEval and MBPP, it achieves a 3.0% Pass@1 improvement over MapCoder while consuming only 6.4% of its compute cost; it outperforms AgentCoder’s 5-turn results within just 2 generations; and LLaMA2-7B reaches the performance ceiling of comparable methods using only 20% of their computational budget.

Technology Category

• Natural Language Processing: Code Generation / Program Synthesis from Natural Language
• Machine Learning: Large Multimodal Models (LMMs)
• Planning, Routing, and Scheduling: Planning with Language Models

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

Abstract

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 extbf{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.

Problem

• Reducing high latency in test-time computation for code generation
• Enhancing precision in code generation through exploration and refinement
• Minimizing computational costs while maintaining accuracy in LLMs

Innovation

• ThinkCoder combines exploration and refinement phases
• Preference-driven optimization with Reinforced Self-Training
• Efficient LLM evolution via exploration trajectories