NExT: Teaching Large Language Models to Reason about Code Execution

TL;DR

NExT enhances large language models' ability to understand and reason about program execution by training them on execution traces, significantly improving their performance on program repair tasks.

Contribution

The paper introduces NExT, a self-training method that teaches LLMs to reason about code execution using synthetic, execution-aware rationales, improving their debugging capabilities.

Findings

NExT improves fix rates by 26.1% on MBPP and 14.3% on HumanEval.

NExT enhances rationale quality as per automated metrics and human raters.

The model generalizes to scenarios without program traces at test-time.

Abstract

A fundamental skill among human developers is the ability to understand and reason about program execution. As an example, a programmer can mentally simulate code execution in natural language to debug and repair code (aka. rubber duck debugging). However, large language models (LLMs) of code are typically trained on the surface textual form of programs, thus may lack a semantic understanding of how programs execute at run-time. To address this issue, we propose NExT, a method to teach LLMs to inspect the execution traces of programs (variable states of executed lines) and reason about their run-time behavior through chain-of-thought (CoT) rationales. Specifically, NExT uses self-training to bootstrap a synthetic training set of execution-aware rationales that lead to correct task solutions (e.g., fixed programs) without laborious manual annotation. Experiments on program repair tasks based on MBPP and HumanEval demonstrate that NExT improves the fix rate of a PaLM 2 model, by 26.1% and 14.3% absolute, respectively, with significantly improved rationale quality as verified by automated metrics and human raters. Our model can also generalize to scenarios where program traces are absent at test-time.