TLDR: ToolTrain is a new two-stage training framework that significantly improves how Large Language Models (LLMs) find and pinpoint bugs in code repositories. By combining initial supervised fine-tuning with advanced reinforcement learning, ToolTrain teaches LLMs to strategically use code retrieval tools. This approach has enabled models to achieve state-of-the-art performance in issue localization, even outperforming larger commercial AI models on specific tasks, and leads to better overall automated issue resolution.
Software development often involves the challenging task of issue localization, which means identifying the exact parts of code that need modification to fix a problem. This process is complex because natural language descriptions of issues need to be translated into specific code locations, often requiring multi-step reasoning through various code dependencies.
Large Language Models (LLMs) have shown great promise in software engineering, including code generation and test synthesis. Recently, LLM-based agents have emerged as a natural progression, allowing these models to use tools for exploring code repositories. However, this introduces a new challenge: what researchers call “Repo Deep Search.” This task demands that LLMs not only reason through multiple steps but also effectively utilize various repository retrieval tools without aimlessly wandering through the code.
To address this, researchers have introduced ToolTrain, a novel two-stage training framework designed to significantly enhance LLMs’ ability to use retrieval tools for issue localization. ToolTrain combines rejection-sampled supervised fine-tuning with tool-integrated reinforcement learning.
Introducing RepoSearcher and its Tools
At the core of ToolTrain is RepoSearcher, a lightweight and user-friendly issue localization agent. RepoSearcher is equipped with a set of simple yet efficient tools that allow LLMs to retrieve specific content from a repository. These tools include functions to get the repository’s file structure, imports of a file, and search for class, function, or class method content by name. Once the LLM has gathered enough information, it uses an ‘Exit’ tool to conclude the search and provide the localization results.
The Two-Stage Training Approach
ToolTrain’s effectiveness comes from its unique two-stage training process:
1. Rejection-Sampled Supervised Fine-Tuning (SFT): In this initial stage, the LLM generates multiple tool-use trajectories (sequences of tool calls). Only those trajectories that lead to correct answers are selected and used for fine-tuning the model. This stage acts as a warm-up, helping the LLM understand the task format, how to call the retrieval tools, and basic reasoning strategies for issue localization.
2. Tool-Integrated Reinforcement Learning (RL): After the SFT warm-up, the LLM enters the reinforcement learning stage. Here, it learns through trial and error by exploring different tool-use strategies. The model’s sampled trajectories are scored based on how well their predicted answers align with the correct solution. These scores act as a reward signal, guiding the model to generate more accurate trajectories and discouraging incorrect or illogical tool calls. This iterative process teaches the LLM to make more strategic tool calls, avoiding redundant explorations and focusing on promising code paths.
Building High-Quality Training Data
A crucial aspect of ToolTrain is its high-quality training data. The researchers constructed this data from over 28,000 pairs of resolved issues and their corresponding pull requests from open-source GitHub projects. This ensures that the training examples are authentic and relevant to real-world software development scenarios.
Also Read:
- Enhancing Code Generation with Explicit Data Dependencies from UML Diagrams
- The Hidden Flaw: How Large Language Models Handle Bad Code Instructions
Impressive Results and Impact
Experimental results demonstrate that models trained with ToolTrain achieve state-of-the-art performance in issue localization. For instance, a 32B model trained with ToolTrain even surpassed leading proprietary models like Claude-3.7 on certain function-level localization tasks. Notably, a smaller 7B model trained with ToolTrain outperformed other frameworks using larger 32B models, highlighting the efficiency of this training approach.
The research also validated that improved issue localization performance directly translates to better end-to-end issue resolution. This suggests that focusing on training for issue localization is a viable and effective strategy for advancing automated software development. The code for ToolTrain is available for further exploration. You can find more details in the research paper.
