Training LLMs to Better Self-Debug and Explain Code

2024-05-28

Neural Information Processing Systems

AI Summary

Existing large language models (LLMs) exhibit weak self-debugging capabilities in complex code generation, especially among small open-source models. Method: We propose an explanation-guided multi-step self-debugging framework. It constructs “generate–explain–correct” execution-feedback trajectories; introduces the first explanation-driven joint supervised fine-tuning and reinforcement learning (RL) paradigm; designs a composite reward function integrating explanation accuracy and correction correctness; and employs execution-based filtering and teacher distillation to enhance data quality. Results: On four major benchmarks, the framework achieves up to a 15.92% absolute improvement in pass@1; RL further boosts performance by 3.54%. Human evaluation confirms that generated explanations are more accurate, interpretable, and actionable for debugging—demonstrating significant advances in both code correctness and explainable self-correction.

Technology Category

• Natural Language Processing: Code Generation / Program Synthesis from Natural Language
• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

In the domain of code generation, self-debugging is crucial. It allows LLMs to refine their generated code based on execution feedback. This is particularly important because generating correct solutions in one attempt proves challenging for complex tasks. Prior works on self-debugging mostly focus on prompting methods by providing LLMs with few-shot examples, which work poorly on small open-sourced LLMs. In this work, we propose LeDex, a training framework that significantly improves the self-debugging capability of LLMs. Intuitively, we observe that a chain of explanations on the wrong code followed by code refinement helps LLMs better analyze the wrong code and do refinement. We thus propose an automated pipeline to collect a high-quality dataset for code explanation and refinement by generating a number of explanations and refinement trajectories from the LLM itself or a larger teacher model and filtering via execution verification. We perform supervised fine-tuning (SFT) and further reinforcement learning (RL) on both success and failure trajectories with a novel reward design considering code explanation and refinement quality. SFT improves the pass@1 by up to 15.92% and pass@10 by 9.30% over four benchmarks. RL training brings additional up to 3.54% improvement on pass@1 and 2.55% improvement on pass@10. The trained LLMs show iterative refinement ability and can keep refining code continuously. Lastly, our human evaluation shows that the LLMs trained with our framework generate more useful code explanations and help developers better understand bugs in source code.

Problem

• Enhance LLMs' self-debugging and code explanation capabilities
• Improve code refinement through automated data collection and training
• Boost performance on benchmarks via novel reward design

Innovation

• Automated pipeline for dataset collection
• Supervised fine-tuning with execution verification
• Reinforcement learning with novel reward design