LeDex: Training LLMs to Better Self-Debug and Explain Code

TL;DR

LeDex is a training framework that enhances large language models' ability to self-debug and explain code through supervised fine-tuning and reinforcement learning, leading to more accurate and useful code generation and debugging.

Contribution

This work introduces LeDex, a novel training pipeline that improves LLM self-debugging by leveraging explanation and refinement trajectories generated and filtered automatically.

Findings

Supervised fine-tuning improves pass@1 by up to 15.92%

Reinforcement learning further boosts pass@1 by up to 3.54%

Trained models demonstrate iterative code refinement and generate more helpful explanations.

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.