Elamur Research⌕
☰
home / papers / 2024 / NExT: Teaching Large Language Models to Reason about Code Execution

NExT: Teaching Large Language Models to Reason about Code Execution

Paper · 2024 · Ansong Ni, Miltiadis Allamanis, Arman Cohan et al.

Summary

Code LLMs learn mostly from the static text of programs, so they often fail to anticipate what code actually does at run time. This weakens them on tasks like debugging. NExT (Naturalized Execution Tuning) teaches a model to read execution traces and explain program behaviour in natural language before proposing a fix. Traces record the variable values after each executed line, and the authors encode them compactly as numbered inline comments attached to the source code. Off-the-shelf models gain little from traces through few-shot prompting alone. NExT therefore finetunes the model through iterative self-training: it samples chain-of-thought rationales with candidate fixes, keeps only those whose fixed program passes the unit tests, and retrains on them. This removes the need for hand-written or teacher-distilled rationales. Applied to PaLM 2-L on Python repair benchmarks, the method substantially raises fix rates. Automatic and human evaluations show that the generated explanations also improve. This points to execution-grounded reasoning as a route to more accurate and more interpretable coding assistants.

Key points

  • Base model: PaLM 2-L. Absolute fix-rate gains are 26.1% on Mbpp-R and 14.3% on HumanEval-Fix-Plus.
  • Few-shot prompting baselines span PaLM 2-L, GPT-3.5, GPT-4, Mixtral 8x7B, DeepSeek Coder 33B and StarCoder 15.5B. Simply adding traces or CoT to the prompt often fails to help, and sometimes lowers accuracy.
  • Accuracy roughly matches a strong self-training repair baseline without rationales (Ye et al., 2022), while producing noticeably more diverse samples.
  • The model still works when traces are unavailable at test time, which is an out-of-distribution setting.
  • Rationale quality is measured with a new proxy method: smaller LLMs attempt the task while following NExT's explanations. Human raters confirm that the rationales are better than those from base PaLM 2-L.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓University of Illinois Urbana-ChampaignUrbana · United States228 works · 74,591 citations🎓Yale UniversityNew Haven · United States112 works · 14,009 citations

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksInformation Systems · 1,888 worksSoftware · 150 worksComputer Science Applications · 184 works

🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 4,145 works

Cite

Cite as Ansong Ni, Miltiadis Allamanis, Arman Cohan, Yinlin Deng, Kensen Shi, Charles Sutton and Pengcheng Yin. 2024. NExT: Teaching Large Language Models to Reason about Code Execution. arXiv:2404.14662.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2404.14662 · arXiv 2404.14662

@misc{ni2024next,
  title = {{NExT: Teaching Large Language Models to Reason about Code Execution}},
  author = {Ansong Ni, Miltiadis Allamanis, Arman Cohan, Yinlin Deng, Kensen Shi, Charles Sutton and Pengcheng Yin},
  year = {2024},
  doi = {10.48550/arxiv.2404.14662},
  eprint = {2404.14662},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2404.14662}
}

References

75 references, 28 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. R. Anil, A. M. Dai, O. Firat, M. Johnson, D. Lepikhin, A. Passos, S. Shakeri, E. Taropa, P. Bailey, Z. Chen, et al. PaLM 2 technical report. arXiv preprint arXiv:2305.10403, 2023.
  2. T. W. Anthony, Z. Tian, and D. Barber. Thinking fast and slow with deep learning and tree search. In Neural Information Processing Systems (NeurIPS), 2017.
  3. J. Austin, A. Odena, M. Nye, M. Bosma, H. Michalewski, D. Dohan, E. Jiang, C. Cai, M. Terry, Q. Le, et al. Program synthesis with large language models. arXiv preprint arXiv:2108.07732, 2021.
  4. D. Bieber, C. Sutton, H. Larochelle, and D. Tarlow. Learning to execute programs with instruction pointer attention graph neural networks. In Advances in Neural Information Processing Systems (NeurIPS), Oct. 2020. find
  5. D. Bieber, R. Goel, D. Zheng, H. Larochelle, and D. Tarlow. Static prediction of runtime errors by learning to execute programs with external resource descriptions. ArXiv, abs/2203.03771, 2022. — [arXiv:2203.03771](https://arxiv.org/abs/2203.03771) find
  6. M. Bosnjak, T. Rocktäschel, J. Naradowsky, and S. Riedel. Programming with a differentiable Forth interpreter. ArXiv, abs/1605.06640, 2016.
  7. I. Bouzenia, Y. Ding, K. Pei, B. Ray, and M. Pradel. TraceFixer: Execution trace-driven program repair. ArXiv, abs/2304.12743, 2023. — [arXiv:2304.12743](https://arxiv.org/abs/2304.12743) find
  8. M. Chen, J. Tworek, H. Jun, Q. Yuan, H. P. d. O. Pinto, J. Kaplan, H. Edwards, Y. Burda, N. Joseph, G. Brockman, et al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374, 2021a.
  9. X. Chen, D. X. Song, and Y. Tian. Latent execution for neural program synthesis beyond domain-specific languages. ArXiv, abs/2107.00101, 2021b. find
  10. X. Chen, M. Lin, N. Schärli, and D. Zhou. Teaching large language models to self-debug. arXiv preprint arXiv:2304.05128, 2023. find
  11. Z. Chen, S. Kommrusch, M. Tufano, L.-N. Pouchet, D. Poshyvanyk, and M. Martin. SequenceR: Sequence-to-sequence learning for end-to-end program repair. IEEE Transactions on Software Engineering, 47:1943–1959, 2018. find
  12. H. W. Chung, L. Hou, S. Longpre, B. Zoph, Y. Tay, W. Fedus, E. Li, X. Wang, M. Dehghani, S. Brahma, A. Webson, S. S. Gu, Z. Dai, M. Suzgun, X. Chen, A. Chowdhery, D. Valter, S. Narang, G. Mishra, A. W. Yu, V. Zhao, Y. Huang, A. M. Dai, H. Yu, S. Petrov, E. H. hsin Chi, J. Dean, J. Devlin, A. Roberts… find
  13. J. Cito, I. Dillig, V. Murali, and S. Chandra. Counterfactual explanations for models of code. International Conference on Software Engineering (ICSE), 2022. find
  14. P. Dasigi, M. Gardner, S. Murty, L. Zettlemoyer, and E. H. Hovy. Iterative search for weakly supervised semantic parsing. In North American Chapter of the Association for Computational Linguistics (NAACL), 2019. find
  15. S. Fakhoury, S. Chakraborty, M. Musuvathi, and S. K. Lahiri. Towards generating functionally correct code edits from natural language issue descriptions. ArXiv, abs/2304.03816, 2023. — [arXiv:2304.03816](https://arxiv.org/abs/2304.03816) find
  16. Z. Fan, X. Gao, M. Mirchev, A. Roychoudhury, and S. H. Tan. Automated repair of programs from large language models. 2023 IEEE/ACM 45th International Conference on Software Engineering (ICSE), pages 1469–1481, 2022. find
  17. Y. Fu, H.-C. Peng, L. Ou, A. Sabharwal, and T. Khot. Specializing smaller language models towards multi-step reasoning. In International Conference on Machine Learning (ICML), 2023.
  18. A. L. Gaunt, M. Brockschmidt, N. Kushman, and D. Tarlow. Differentiable programs with neural libraries. In International Conference on Machine Learning (ICML), 2016. find
  19. A. Graves, G. Wayne, and I. Danihelka. Neural turing machines. ArXiv, abs/1410.5401, 2014. find
  20. S. Gunasekar, Y. Zhang, J. Aneja, C. C. T. Mendes, A. D. Giorno, S. Gopi, M. Javaheripi, P. C. Kauffmann, G. de Rosa, O. Saarikivi, A. Salim, S. Shah, H. S. Behl, X. Wang, S. Bubeck, R. Eldan, A. T. Kalai, Y. T. Lee, and Y.-F. Li. Textbooks are all you need. ArXiv, abs/2306.11644, 2023. find
  21. D. Guo, Q. Zhu, D. Yang, Z. Xie, K. Dong, W. Zhang, G. Chen, X. Bi, Y. Wu, Y. K. Li, F. Luo, Y. Xiong, and W. Liang. DeepSeek-Coder: When the large language model meets programming – the rise of code intelligence, 2024. find
  22. K. Gupta, P. E. Christensen, X. Chen, and D. X. Song. Synthesize, execute and debug: Learning to repair for neural program synthesis. ArXiv, abs/2007.08095, 2020. find
  23. A. Heinonen, B. Lehtelä, A. Hellas, and F. Fagerholm. Synthesizing research on programmers’ mental models of programs, tasks and concepts - a systematic literature review. Inf. Softw. Technol., 164: 107300, 2022. find
  24. M. D. Hoffman, D. Phan, D. Dohan, S. Douglas, T. A. Le, A. T. Parisi, P. Sountsov, C. Sutton, S. Vikram, and R. A. Saurous. Training Chain-of-Thought via Latent-Variable inference. In Conference on Neural Information Processing Systems (NeurIPS), 2023. find
  25. X. Hu, G. Li, X. Xia, D. Lo, and Z. Jin. Deep code comment generation. 2018 IEEE/ACM 26th International Conference on Program Comprehension (ICPC), pages 200–20010, 2018. find
  26. A. Hunt and D. Thomas. The pragmatic programmer: From journeyman to master. 1999. find
  27. A. Q. Jiang, A. Sablayrolles, A. Roux, A. Mensch, B. Savary, C. Bamford, D. S. Chaplot, D. d. l. Casas, E. B. Hanna, F. Bressand, et al. Mixtral of experts. arXiv preprint arXiv:2401.04088, 2024. find
  28. N. Jiang, K. Liu, T. Lutellier, and L. Tan. Impact of code language models on automated program repair. 2023 IEEE/ACM 45th International Conference on Software Engineering (ICSE), pages 1430–1442, 2023. find
  29. S. Kang, B. Chen, S. Yoo, and J.-G. Lou. Explainable automated debugging via large language model-driven scientific debugging. ArXiv, abs/2304.02195, 2023. — [arXiv:2304.02195](https://arxiv.org/abs/2304.02195) find
  30. S. Kulal, P. Pasupat, K. Chandra, M. Lee, O. Padon, A. Aiken, and P. S. Liang. SPoC: Search-based pseudocode to code. Advances in Neural Information Processing Systems (NeurIPS), 32, 2019.
  31. C. Le Goues, M. Pradel, and A. Roychoudhury. Automated program repair. Communications of the ACM, 62(12):56–65, 2019. find
  32. R. Li, L. B. Allal, Y. Zi, N. Muennighoff, D. Kocetkov, C. Mou, M. Marone, C. Akiki, J. Li, J. Chim, et al. StarCoder: may the source be with you! arXiv preprint arXiv:2305.06161, 2023. find
  33. Y. Li, S. Wang, and T. N. Nguyen. DLFix: Context-based code transformation learning for automated program repair. 2020 IEEE/ACM 42nd International Conference on Software Engineering (ICSE), pages 602–614, 2020. find
  34. Z. Li, S. Lu, D. Guo, N. Duan, S. Jannu, G. Jenks, D. Majumder, J. Green, A. Svyatkovskiy, S. Fu, and N. Sundaresan. Automating code review activities by large-scale pre-training. Proceedings of the 30th ACM Joint European Software Engineering Conference and Symposium on the Foundations of Software… find
  35. C. Liang, M. Norouzi, J. Berant, Q. V. Le, and N. Lao. Memory augmented policy optimization for program synthesis and semantic parsing. Advances in Neural Information Processing Systems (NeurIPS), 31, 2018. find
  36. H. Lightman, V. Kosaraju, Y. Burda, H. Edwards, B. Baker, T. Lee, J. Leike, J. Schulman, I. Sutskever, and K. Cobbe. Let’s verify step by step. arXiv preprint arXiv:2305.20050, 2023.
  37. J. Liu, C. S. Xia, Y. Wang, and L. Zhang. Is your code generated by chatgpt really correct? rigorous evaluation of large language models for code generation. arXiv preprint arXiv:2305.01210, 2023.
  38. S. Longpre, L. Hou, T. Vu, A. Webson, H. W. Chung, Y. Tay, D. Zhou, Q. V. Le, B. Zoph, J. Wei, and A. Roberts. The flan collection: Designing data and methods for effective instruction tuning. In International Conference on Machine Learning (ICML), 2023.
  39. W. Ma, S. Liu, W. Wang, Q. Hu, Y. Liu, C. Zhang, L. Nie, and Y. Liu. ChatGPT: Understanding code syntax and semantics. 2023. find
  40. A. Madaan, N. Tandon, P. Gupta, S. Hallinan, L. Gao, S. Wiegreffe, U. Alon, N. Dziri, S. Prabhumoye, Y. Yang, S. Welleck, B. P. Majumder, S. Gupta, A. Yazdanbakhsh, and P. Clark. Self-refine: Iterative refinement with self-feedback. ArXiv, abs/2303.17651, 2023. find
  41. A. Mitra, L. D. Corro, S. Mahajan, A. Codas, C. Simoes, S. Agrawal, X. Chen, A. Razdaibiedina, E. Jones, K. Aggarwal, H. Palangi, G. Zheng, C. Rosset, H. Khanpour, and A. Awadallah. Orca 2: Teaching small language models how to reason. ArXiv, abs/2311.11045, 2023.
  42. N. Muennighoff, Q. Liu, A. Zebaze, Q. Zheng, B. Hui, T. Y. Zhuo, S. Singh, X. Tang, L. Von Werra, and S. Longpre. OctoPack: Instruction tuning code large language models. arXiv preprint arXiv:2308.07124, 2023. find
  43. S. Mukherjee, A. Mitra, G. Jawahar, S. Agarwal, H. Palangi, and A. H. Awadallah. Orca: Progressive learning from complex explanation traces of GPT-4. ArXiv, abs/2306.02707, 2023.
  44. A. Ni, P. Yin, and G. Neubig. Merging weak and active supervision for semantic parsing. In Proceedings of the AAAI Conference on Artificial Intelligence (AAAI), volume 34, pages 8536–8543, 2020. find
  45. A. Ni, J. P. Inala, C. Wang, A. Polozov, C. Meek, D. Radev, and J. Gao. Learning math reasoning from self-sampled correct and partially-correct solutions. In The Eleventh International Conference on Learning Representations (ICLR), 2022.
  46. A. Ni, S. Iyer, D. Radev, V. Stoyanov, W.-t. Yih, S. Wang, and X. V. Lin. LEVER: Learning to verify language-to-code generation with execution. In International Conference on Machine Learning (ICML), pages 26106–26128. PMLR, 2023.
  47. M. Nye, A. J. Andreassen, G. Gur-Ari, H. Michalewski, J. Austin, D. Bieber, D. Dohan, A. Lewkowycz, M. Bosma, D. Luan, et al. Show your work: Scratchpads for intermediate computation with language models. arXiv preprint arXiv:2112.00114, 2021.
  48. T. X. Olausson, J. P. Inala, C. Wang, J. Gao, and A. Solar-Lezama. Demystifying GPT self-repair for code generation. arXiv preprint arXiv:2306.09896, 2023.
  49. OpenAI. GPT-4 technical report, 2023.
  50. R. Paul, M. M. Hossain, M. L. Siddiq, M. Hasan, A. Iqbal, and J. C. S. Santos. Enhancing automated program repair through fine-tuning and prompt engineering. 2023. find
  51. K. Pei, D. Bieber, K. Shi, C. Sutton, and P. Yin. Can large language models reason about program invariants? In International Conference on Machine Learning (ICML), 2023. find
  52. A. Prasad, S. Saha, X. Zhou, and M. Bansal. Receval: Evaluating reasoning chains via correctness and informativeness. arXiv preprint arXiv:2304.10703, 2023.
  53. N. F. Rajani, B. McCann, C. Xiong, and R. Socher. Explain yourself! leveraging language models for commonsense reasoning. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics (ACL), pages 4932–4942, Florence, Italy, July 2019. Association for Computational Lingu…
  54. S. I. Ross, F. Martinez, S. Houde, M. J. Muller, and J. D. Weisz. The programmer’s assistant: Conversational interaction with a large language model for software development. Proceedings of the 28th International Conference on Intelligent User Interfaces, 2023. find
  55. B. Roziere, J. Gehring, F. Gloeckle, S. Sootla, I. Gat, X. E. Tan, Y. Adi, J. Liu, T. Remez, J. Rapin, et al. Code Llama: Open foundation models for code. arXiv preprint arXiv:2308.12950, 2023.
  56. K. Shi, H. Dai, K. Ellis, and C. Sutton. CrossBeam: Learning to search in bottom-up program synthesis. In International Conference on Learning Representations (ICLR), 2022. find
  57. K. Shi, J. Hong, Y. Deng, P. Yin, M. Zaheer, and C. Sutton. ExeDec: Execution decomposition for compositional generalization in neural program synthesis. In International Conference on Learning Representations (ICLR), 2024. find
  58. E. C. Shin, I. Polosukhin, and D. Song. Improving neural program synthesis with inferred execution traces. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems (NeurIPS), volume 31. Curran Associates, Inc., 2… find
  59. V. Shwartz, P. West, R. Le Bras, C. Bhagavatula, and Y. Choi. Unsupervised commonsense question answering with self-talk. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 4615–4629, Online, Nov. 2020. Association for Computational Linguistics.… find
  60. B. Siegmund, M. Perscheid, M. Taeumel, and R. Hirschfeld. Studying the advancement in debugging practice of professional software developers. In 2014 IEEE International Symposium on Software Reliability Engineering Workshops, pages 269–274, 2014. doi: 10.1109/ISSREW.2014.36. — [doi:10.1109/issrew.2014.36](https://doi.org/10.1109/issrew.2014.36) find
  61. D. Sobania, M. Briesch, C. Hanna, and J. Petke. An analysis of the automatic bug fixing performance of ChatGPT. 2023 IEEE/ACM International Workshop on Automated Program Repair (APR), pages 23–30, 2023. find
  62. H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023.
  63. J. Uesato, N. Kushman, R. Kumar, F. Song, N. Siegel, L. Wang, A. Creswell, G. Irving, and I. Higgins. Solving math word problems with process- and outcome-based feedback. ArXiv, abs/2211.14275, 2022.
  64. C. Wang, K. Tatwawadi, M. Brockschmidt, P.-S. Huang, Y. Mao, O. Polozov, and R. Singh. Robust text-to-SQL generation with execution-guided decoding. arXiv: Computation and Language, 2018. find
  65. P. Wang, L. Li, Z. Shao, R. X. Xu, D. Dai, Y. Li, D. Chen, Y. Wu, and Z. Sui. Math-shepherd: Verify and reinforce llms step-by-step without human annotations. ArXiv, abs/2312.08935, 2024.
  66. J. Wei, X. Wang, D. Schuurmans, M. Bosma, E. Chi, Q. Le, and D. Zhou. Chain of thought prompting elicits reasoning in large language models. "arXiv preprint arXiv:2201.11903", Jan. 2022a.
  67. J. Wei, X. Wang, D. Schuurmans, M. Bosma, F. Xia, E. Chi, Q. V. Le, D. Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems (NeurIPS), 35:24824–24837, 2022b.
  68. C. Xia and L. Zhang. Less training, more repairing please: revisiting automated program repair via zero-shot learning. Proceedings of the 30th ACM Joint European Software Engineering Conference and Symposium on the Foundations of Software Engineering, 2022. find
  69. C. Xia and L. Zhang. Keep the conversation going: Fixing 162 out of 337 bugs for $0.42 each using ChatGPT. ArXiv, abs/2304.00385, 2023. — [arXiv:2304.00385](https://arxiv.org/abs/2304.00385) find
  70. C. Xia, Y. Wei, and L. Zhang. Automated program repair in the era of large pre-trained language models. 2023 IEEE/ACM 45th International Conference on Software Engineering (ICSE), pages 1482–1494, 2023. find
  71. S. Yao, J. Zhao, D. Yu, N. Du, I. Shafran, K. Narasimhan, and Y. Cao. ReAct: Synergizing reasoning and acting in language models. ArXiv, abs/2210.03629, 2022.
  72. H. Ye, M. Martinez, X. Luo, T. Zhang, and M. Martin. SelfAPR: Self-supervised program repair with test execution diagnostics. Proceedings of the 37th IEEE/ACM International Conference on Automated Software Engineering (ASE), 2022. find
  73. P. Yin, W.-D. Li, K. Xiao, A. Rao, Y. Wen, K. Shi, J. Howland, P. Bailey, M. Catasta, H. Michalewski, O. Polozov, and C. Sutton. Natural language to code generation in interactive data science notebooks. In A. Rogers, J. Boyd-Graber, and N. Okazaki, editors, Proceedings of the 61st Annual Meeting of…
  74. W. Zaremba and I. Sutskever. Learning to execute. ArXiv, abs/1410.4615, 2014. find
  75. E. Zelikman, Y. Wu, J. Mu, and N. Goodman. STaR: Bootstrapping reasoning with reasoning. Advances in Neural Information Processing Systems (NeurIPS), 35:15476–15488, 2022.

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/ni-2024-next-teaching-large-language-models/