PaperlayerView exact revision on arXiv

Source-linked AI summary

NExT: Teaching Large Language Models to Reason about Code Execution

Ansong Ni, Miltiadis Allamanis, Arman Cohan, Yinlin Deng, Kensen Shi, Charles Sutton, Pengcheng Yin

arXiv:2404.14662v1cs.LGcs.CLcs.PLcs.SE

TL;DR

LLMs often lack runtime-semantic reasoning despite needing it for complex coding tasks and useful debugging explanations. NExT finetunes them with execution-aware chain-of-thought rationales bootstrapped through self-training, improving repair rates and rationale quality while generalizing to trace-free test-time settings.

  • Problem

    Code LLMs struggle with runtime execution reasoning, although program-repair models benefit from understanding variable states and explaining bugs and fixes.

    Source: §1 · p. 1

  • Method

    NExT repeatedly samples execution-aware rationales and fixes, keeps fixes verified by unit tests, and finetunes the model on these synthetic examples.

    Source: §1 · pp. 1–2, §4 · p. 5

  • Results

    NExT improves PaLM 2-L fix rates by 26.1% on Mbpp-R and 14.3% on HumanEvalFix-Plus, with significantly better rationale quality and generalization when traces are absent at test time.

    Source: §1 · p. 3, §5.1 · p. 8

  • Takeaways & Limitations

    Execution-aware self-training can improve both program repair success and the usefulness of natural-language explanations within the evaluated tasks.

    Source: §1 · p. 3, §7 · p. 12

    Rationale and fix quality may plateau at different training iterations, and the evaluation is conducted on two Python program-repair tasks.

    Source: §1 · p. 3, §4 · p. 5

Abstract

from arXiv · show

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.

1. Introduction

NExT targets LLMs’ difficulty reasoning about program runtime behavior by combining execution traces with natural-language rationales and self-training. It improves repair performance and rationale quality while also generalizing beyond trace-available testing.

Motivation

  • LLMs trained on code still struggle with runtime-dependent software tasks, while execution-aware explanations can help developers understand and assess proposed fixes.Human debugging uses variable states and mental execution simulation, motivating analogous reasoning in repair models.Source: §1 · p. 1

NExT

  • NExT trains LLMs to inspect execution traces and generate natural-language chain-of-thought rationales grounded in program behavior.Traces provide line-by-line variable states or exceptions, represented as inline comments without disrupting program structure.Source: §1 · pp. 1–2

Self-training

  • NExT bootstraps execution-aware rationales by repeatedly sampling solutions, retaining those that pass unit tests, and finetuning without manual rationale annotation.The self-training loop is intended to improve both task success and rationale quality across iterations.Source: §1 · p. 2

Results

  • 26.1% and 14.3% absolute gains in PaLM 2-L fix rate were reported on Mbpp-R and HumanEvalFix-Plus, respectively, alongside higher-quality rationales and trace-free test-time generalization.Proxy evaluation and human annotation supported the rationale-quality improvement.Source: §1 · p. 3

2. Task: Program Repair with Traces

The task is to repair buggy programs using execution traces that expose intermediate variable states and test outcomes. NExT presents these traces compactly so models can reason about bugs and fixes while preserving program structure.

Program Repair with Execution Traces

  • Program repair generates a fixed program that passes all test cases while the buggy input fails, using traces collected during execution against tests.A trace is a sequence of intermediate variable states after each executed statement and can include debugging information such as exceptions.Source: §2 · p. 3

An LLM-friendly Trace Representation

  • The prompt concatenates the task instruction, buggy code, tests, and execution traces, with traces formatted as compact inline code comments.This representation lets the model use its learned code structure while incorporating execution effects.Source: §2 · pp. 3–4

CoT Reasoning with Execution

  • Models generate both a fixed program and a rationale that identifies bugs, proposes repairs, and can contrast expected with actual variable values.These execution-grounded explanations are intended to help developers understand both the defect and the proposed solution.Source: §2 · pp. 3–4

An LLM-friendly Trace Representation

  • Inline traces record changed variables, order states by execution, and omit intermediate loop states to reduce token use while retaining the original code layout.The format is designed to make repeated and non-obvious execution order intelligible without fully unrolling the program.Source: §2 · p. 4

3. Preliminary Study: Can LLMs reason with program traces in natural language?

The preliminary study tests whether prompting alone enables LLMs to reason naturally about execution traces. Results show that traces help, but vanilla models often hallucinate execution explanations and struggle with trace-based chain-of-thought reasoning.

Study Setup

  • The study few-shot prompts general-purpose and code-specific LLMs to solve program repair tasks using the inline trace representation.Models include PaLM 2, GPT, Mixtral, StarCoder, and DeepSeek Coder, evaluated on two Python repair datasets.Source: §3 · p. 4

Results

  • A 7.8% average drop on HumanEvalFix-Plus occurred when comparing chain-of-thought prompting with traces against vanilla prompting, including large declines for code-specific models.DeepSeek Coder fell from 57.3 to 30.5, while StarCoder fell from 29.3 to 16.5.Source: §3 · p. 5
  • Removing traces caused average performance drops of 4.8% on Mbpp-R and 4.3% on HumanEvalFix-Plus, indicating that traces were useful even though models struggled to reason with them.Sampled PaLM 2-L rationales also showed hallucinations, and only about 30% correctly explained input-program errors.Source: §3 · p. 5

4. NExT: Naturalized Execution Tuning

NExT finetunes LLMs through iterative self-training on execution-aware rationales and fixes sampled for difficult repair problems. It filters candidates with unit tests, retrains from the base checkpoint, and extends rationale bootstrapping to execution-centered programming tasks.

NExT: Naturalized Execution Tuning

  • NExT repeatedly samples rationales and code fixes, filters candidates by unit-test execution, and finetunes the model on solutions that pass all tests.This sample-filter-train loop is repeated across iterations to improve repair success and rationale quality.Source: §4 · p. 5

Sampling rationales and code solutions

  • Sampling focuses on problems that the base model finds difficult, reducing sampling cost while directing learning toward harder repairs.A task-level metric and threshold identify the hard subset before iterative training begins.Source: §4 · pp. 5–6

Filtering candidate solutions

  • Candidate rationales and fixes are sampled jointly from the current model, then retained as training data when their fixes pass the available tests.The algorithm represents each training item with its instruction, buggy program, tests, and execution trace.Source: §4 · p. 6

Model training

  • NExT finetunes from the initial checkpoint at each iteration to reduce overfitting to potentially lower-quality samples from early iterations.The filtering reward directly tests fix correctness rather than rationale quality, although rationale quality improves empirically during training.Source: §4 · pp. 5–6

Discussion

  • NExT treats rationales and code fixes as latent variables supervised by unit-test outcomes, and its framework can extend beyond repair to other execution-reasoning programming tasks.The authors identify broader programming-task coverage and richer multilingual trace representations as future directions.Source: §4 · p. 6

5. Experiments

Experiments evaluate NExT on Python program repair using Mbpp-R and HeFix+, measuring fix rates, rationale quality, ablations, generalization, and human judgments. NExT improves repair performance and rationale quality, while traces and natural-language reasoning support generalization and sample diversity.

5. Experiments

  • The evaluation uses 10,047 Mbpp-R training tasks, 1,468 development tasks, and 164 HumanEvalFix-derived HeFix+ programs, with pass@k measuring tested functional correctness.HeFix+ adds the more rigorous EvalPlus test suites and measures generalization without further finetuning.Source: §5 · pp. 6–7
  • Rationale quality is evaluated both by smaller PaLM 2 variants solving tasks conditioned on rationales and by human ratings of bug explanations and fix suggestions.The proxy metric and human judgments are reported as consistent.Source: §5 · p. 7, §5.2 · p. 10

5.1. Main Results

  • 26.1% absolute pass@1 improvement on Mbpp-R shows NExT substantially raises PaLM 2-L’s end-to-end program repair fix rate.PaLM 2-L+NExT also outperforms GPT-3.5 across all pass@k metrics.Source: §5.1 · pp. 8, 10

5.2. Human Evaluation of Rationale Quality

  • NExT consistently improves proxy-based rationale quality, with gains also supported by human ratings showing stronger bug explanations than base PaLM 2-L.Annotators preferred NExT rationales overall, while GPT-3.5 produced more partially correct fix suggestions.Source: §5.1 · p. 8, §5.2 · pp. 10–11

5.1. Main Results

  • Natural-language reasoning improves sample diversity: NExT matches the no-rationale ablation at pass@1 but increasingly leads at pass@10 as training continues.The no-rationale variant converges faster on training data, whereas NExT generalizes better on unseen Mbpp-R problems.Source: §5.1 · pp. 8–9
  • Execution traces are critical: removing them makes PaLM 2-L consistently worse on development-set fix rates and rationale-quality metrics across training iterations.Similar training-set scores suggest traces primarily improve generalization to unseen tasks.Source: Figure 4 · p. 8, §5.1 · p. 9
  • 40.8% end-to-end fix rate without traces at test time remains 21.8% above the 3-shot PaLM 2-L baseline.This result is only 3.3% below NExT trained without traces, and proxy-based rationale evaluation follows the same pattern.Source: §5.1 · p. 9, Table 3 · p. 9
  • 14.3% end-to-end fix-rate improvement on HeFix+ demonstrates generalization from Mbpp-R training to an out-of-distribution repair benchmark.The model is trained on Mbpp-R without additional dataset-specific finetuning.Source: §5 · p. 6, §5.1 · p. 9

6. Related Work

NExT differs from prior execution-reasoning and program-repair methods by taking runtime traces as input and learning flexible natural-language rationales for downstream repair.

Reasoning about Program Execution

  • NExT takes existing execution traces as input rather than predicting trace-containing reasoning chains, producing compact natural-language rationales tailored to downstream tasks.This contrasts with Scratchpad and Self-Debugging, which predict execution information or its summaries as reasoning chains.Source: §6 · p. 11

Program Repair

  • NExT inlines indexed variable states as code comments, preserving code structure while improving token efficiency compared with directly pairing unrolled lines and states.Source: §6 · p. 11
  • NExT generates both natural-language rationales and fixed programs through self-training, unlike related approaches that use test results without producing interpretable rationales.Source: §6 · p. 11

Supervised CoT Reasoning

  • NExT uses self-training to bootstrap execution-aware rationales without manually labeled or distilled chain-of-thought data.Correct program fixes filter the synthetic rationale data, while execution information is taken as given rather than predicted.Source: §6 · p. 12

7. Conclusion

The paper presents NExT as a self-training method for helping LLMs reason about program execution from traces, and reports stronger rationale quality and repair success on two tasks.

7. Conclusion

  • NExT finetunes LLMs to reason with program execution given traces and improves PaLM 2-L’s success rates on two program repair tasks.Source: §7 · p. 12
  • The authors identify broader program-understanding tasks and support for more programming languages as future directions.Source: §7 · p. 12

A.1. Details for Inline Trace Representation

The inline trace representation compresses execution states into code comments while retaining execution order, but its tracing scope and object representations remain limited.

Definitions

  • Execution traces are sequences of program states ordered by executed statements, with each state mapping identifiers to values and repeated statements allowed through loops or recursion.Source: §A.1 · p. 18

Program state representation

  • Variable states record only changed values, use Python initialization syntax, mark no-effect statements with “NO_CHANGE,” and prepend execution-order numbers.Source: §A.1 · p. 18

Inline trace representation

  • Inline traces group states by statement, retain the first two and last state when more than three occur, and append semicolon-delimited states as hash-marked comments.Source: §A.1 · p. 18

Limitations

  • The tracing framework currently supports native Python only and does not capture explicit condition outcomes for if statements, although execution order can reveal which branch was taken.Source: §A.1 · p. 18
  • Collection traces currently show all elements, while complex objects and longer programs may require heuristics or selective inspection to improve token efficiency.Source: §A.1 · pp. 18–19

A.2. Details for Iterative Self-Training

NExT bootstraps rationale-and-fix training data through temperature sampling, test-based filtering, and per-problem subsampling.

Bootstrapping rationales and fixes via temperature sampling

  • The first self-training iteration uses three few-shot exemplars and 96 samples to avoid cold start, while later iterations use zero-shot prompting with 32 samples at temperature T=0.8.Source: §A.2 · p. 19

Filtering rationales and fixes

  • Sampled rationales and fixes are separated with Python regular expressions, filtered by test cases, and capped at three correct fixes per Mbpp-R problem.The cap balances examples across problems and limits overrepresentation of easier cases.Source: §A.2 · p. 19

A.3. Discussion with Previous Work

NExT differs from prior execution-reasoning approaches by using existing execution traces as inputs while training natural-language rationales, rather than generating traces or simulating them step by step. It also improves repair performance without gold traces at test time, suggesting broader self-debugging potential.

Scratchpad and NExT

  • NExT feeds existing execution traces to the model and uses natural language as the primary reasoning format, making its representation more compact than Scratchpad.Scratchpad generates traces as intermediate reasoning steps, whereas NExT grounds rationales in execution states supplied as input.Source: §A.3 · pp. 19–20

Self-Debugging and NExT

  • NExT uses concrete step-wise execution traces, unlike Self-Debugging’s optional error messages and model-generated natural-language trace-like rationales.Because execution information is already captured, NExT rationales can be more succinct without reconstructing execution step by step.Source: §A.3 · p. 20
  • NExT greatly improves program repair even without gold execution traces at test time, indicating potential to strengthen self-debugging through iterative training.The paper identifies this setting as similar in spirit to Self-Debugging and leaves iterative self-debugging as future work.Source: §A.3 · p. 20

B. Experiment Setup Details

The experiments construct Mbpp-R from failed solutions, distinguish trace-generation tests from evaluation tests, and assess rationale quality through blinded side-by-side human annotation. The setup also documents dataset difficulty and possible test-case overfitting.

B.1. Creating Mbpp-R

  • Mbpp-R contains 10,047 training and 1,468 development repair tasks created from filtered failed solutions, with at most 20 buggy solutions retained per problem.The re-split Mbpp data uses 624/90/250 train/dev/test problems before collecting failed solutions for the training and development sets.Source: §B.1 · p. 20

B.2. Use of test cases.

  • Mbpp-R is biased toward harder problems because easier problems may lack 20 failed solutions among 100 samples, making repair more challenging than original Mbpp generation.This bias is identified as a possible reason for the higher difficulty of Mbpp-R.Source: §B.2 · p. 20
  • Traces use open tests, while a repair is correct only if it passes open and hidden tests; Mbpp instead uses all tests for both tracing and evaluation.The authors note that Mbpp’s lack of an open/hidden distinction can produce fixes overfit to the tests, including cases found during human annotation.Source: §B.2 · pp. 20–21

B.3. Details of Human Annotation of Rationale Quality

  • Human raters compare anonymized, randomized rationales from PaLM 2-L+NExT, PaLM 2-L, and GPT-3.5 on 104 sampled Mbpp-R tasks.Each task is rated for bug-explanation helpfulness and fix-suggestion helpfulness using three response levels.Source: §B.3 · p. 22
  • For fix-suggestion quality, a rationale can count as correct when its bug explanation makes the repair obvious even without an explicit fix.This annotation rule treats detailed explanations as sufficient when they enable the developer to infer the correction.Source: §B.3 · p. 22

C. Additional Experiment Results

Additional experiments show that rationale training improves high-sample repair performance, hard-only training improves efficiency and generalization, and proxy-based rationale evaluation is robust across proxy models.

Additional Experiment Results

  • 7.6% on Mbpp-R and 6.8% on HeFix+ are the pass@25 improvements from learning chain-of-thought rationales.The advantage over fix-only fine-tuning is larger at pass@25 than at pass@10.Source: C. Additional Experiment Results · p. 22
  • Training only on hard examples accelerates sampling and generally improves fix rates and rationale quality compared with training on both easy and hard examples.The w/o hard-only baseline performs worse in other evaluations, especially as k increases, although its end-to-end pass@1 is slightly better.Source: C. Additional Experiment Results · p. 22
  • Proxy-based rationale-quality results are largely consistent across PaLM 2-S and PaLM 2-S*, supporting the robustness of the proxy evaluation method.The stronger proxy model yields better proxy-based fix rates, while relative rationale-quality comparisons remain consistent.Source: C. Additional Experiment Results · p. 23

D. Case Study

The case study examines how execution traces, variable states, exceptions, and preserved values support diagnosing faulty programs and generating targeted repairs. It also illustrates how multi-step reasoning and trace reliance can lead to logical errors, incorrect fixes, or missed bugs.

D. Case Study

  • PaLM 2-L+NExT leverages exceptions, variable states, NO_CHANGE annotations, and multi-step trace computations to explain bugs and produce fixes.The examples include identifying an IndexError from a trace, comparing variable states with expected values, explaining preserved variables, and tracking computations across several execution steps.Source: D. Case Study · pp. 23, 25, 26
  • As reasoning chains become longer, the model is more likely to introduce minor logical errors while tracing multi-step execution.This limitation appears in the examples that track steps 2–4 and 2–9 to explain incorrect sums and missing list elements.Source: D. Case Study · pp. 23–24
  • The model can produce a valid-looking natural-language rationale alongside an incorrect code fix.Example 4 identifies the missing insertion after the final tuple element, yet the reported rationale-and-solution case is marked as having an incorrect fix.Source: D. Case Study · pp. 24, 28
  • When programs contain multiple bugs, the model may focus on an obvious traced failure while overlooking nuanced or hidden bugs not represented in the trace.It can notice an IndexError while missing incorrect outputs from zero-initialized freq_arr, or detect a TypeError while missing an off-by-one end index.Source: D. Case Study · pp. 24, 30, 31

E. Full Prompt

  • The trace-based analysis identifies that the non-prime checker returns None for n=2 because its loop range is empty, while n=35 proceeds into the divisibility test.It proposes reimplementing the logic with an is_prime() helper and returning not is_prime(x).Source: E. Full Prompt · p. 34
  • All three heap-selection assertions fail immediately because heapq.nlargest() is called without importing heapq, producing NameError: name 'heapq' is not defined.The suggested repair is to add the missing import before using heapq.nlargest().Source: E. Full Prompt · pp. 34–35
Loading 2404.14662v1…