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

Nan Jiang, Xiaopeng Li, Shiqi Wang, Qiang Zhou, Soneya Binta Hossain, Baishakhi Ray, Varun Kumar, Xiaofei Ma, Anoop Deoras
AWS AI Lab
NeurIPS 2024
MY ALT TEXT

Pipeline of letting LLM generate code and self-debug

MY ALT TEXT

Overview of LeDex

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.

Poster

BibTeX

@inproceedings{jiang2024selfdebug,
    title={Training LLMs to Better Self-Debug and Explain Code}, 
    author={Nan Jiang and Xiaopeng Li and Shiqi Wang and Qiang Zhou and Soneya Binta Hossain and 
            Baishakhi Ray and Varun Kumar and Xiaofei Ma and Anoop Deoras},
    year={2024},
    booktitle={Advances in Neural Information Processing Systems}
}