Learning Performance-Improving Code Edits

By Aman Madaan, Alexander Shypula, Uri Alon, Milad Hashemi, Parthasarathy Ranganathan, Yiming Yang, Graham Neubig, Amir Yazdanbakhsh. 2023

155 citations recorded by the source corpus.

The waning of Moore's Law has shifted the focus of the tech industry towards alternative methods for continued performance gains. While optimizing compilers are a standard tool to help increase program efficiency, programmers continue to shoulder much responsibility in crafting and refactoring code with better performance characteristics. In this paper, we investigate the ability of large language models (LLMs) to suggest functionally correct, performance improving code edits. We hypothesize that language models can suggest such edits in ways that would be impractical for static analysis alone. We investigate these questions by curating a large-scale dataset of Performance-Improving Edits, PIE. PIE contains trajectories of programs, where a programmer begins with an initial, slower version and iteratively makes changes to improve the program's performance. We use PIE to evaluate and improve the capacity of large language models. Specifically, use examples from PIE to fine-tune multiple variants of CODEGEN, a billion-scale Transformer-decoder model. Additionally, we use examples from PIE to prompt OpenAI's CODEX using a few-shot prompting. By leveraging PIE, we find that both CODEX and CODEGEN can generate performance-improving edits, with speedups of more than 2.5x for over 25% of the programs, for C++ and Python, even after the C++ programs were compiled using the O3 optimization level. Crucially, we show that PIE allows CODEGEN, an open-sourced and 10x smaller model than CODEX, to match the performance of CODEX on this challenging task. Overall, this work opens new doors for creating systems and methods that can help programmers write efficient code.

Subfields: Computer Networks and Communications, Hardware and Architecture, Information Systems, Natural Language Processing

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Computer Networks and Communications

1. Yiming Yang
2. Alhussein Fawzi
3. Bernardino Romera-Paredes
4. Xinyun Chen
5. Shin Yoo
6. Pushmeet Kohli
7. Mark Harman
8. Niket Tandon

In Hardware and Architecture

1. Xinyun Chen
2. Shin Yoo
3. Pushmeet Kohli
4. Mark Harman
5. Sean Welleck
6. Alhussein Fawzi
7. Nouha Dziri
8. Uri Alon

In Information Systems

1. Shin Yoo
2. Mark Harman
3. Denny Zhou
4. Xinyun Chen
5. Weizhu Chen
6. Yiming Yang
7. Aman Madaan
8. Shubho Sengupta

In Natural Language Processing

1. Jianfeng Gao
2. Weizhu Chen
3. Denny Zhou
4. Mohit Bansal
5. Nan Duan
6. Sarah Wiegreffe
7. Luyu Gao
8. Xinyun Chen

Citations in the corpus

• Self-Refine: Iterative Refinement with Self-Feedback (2023, methodology): “We evaluate SELF-REFINE on 7 diverse tasks: Dialogue Response Generation (Appendix M; Mehri and Eskenazi, 2020), Code Optimization (Appendix N; Madaan et al., 2023), Code Readability Improvement (Appendix L; Puri et al.”
• ReEvo: Large Language Models as Hyper-Heuristics with Reflective Evolution (2024, unknown): “…code generation capabilities of LLMs are utilized for code debugging (Chen et al., 2023c; Liventsev et al., 2023), enhancing code performance (Madaan et al., 2023a), solving algorithmic competition challenges (Li et al., 2022; Shinn et al., 2023), robotics (Liang et al., 2023; Wang et al.,…”
• Mathematical discoveries from program search with large language models (2023, methodology): “…(LLMs) have recently seen notable improvements in their coding capabilities 4–8 , with applications including debugging 9,10 , solving code competitions 11,12 and improving code performance 13 , synthesizing ‘solve’ programs for open problems requires finding new ideas that are verifiably correct.”
• CRITIC: Large Language Models Can Self-Correct with Tool-Interactive Critiquing (2023, background): “[17] Aman Madaan, Alexander Shypula, Uri Alon, Milad Hashemi, Parthasarathy Ranganathan, Yiming Yang, Graham Neubig, and Amir Yazdanbakhsh.”
• WizardCoder: Empowering Code Large Language Models with Evol-Instruct (2023, unknown): “Introduction of a heuristic emphasizing time and space complexity leverages insights from previous studies Madaan et al. (2023), providing a valuable avenue for improving task complexity.”
• NoFunEval: Funny How Code LMs Falter on Requirements Beyond Functional Correctness (2024, methodology): “(Madaan et al., 2023), which comprises 1K pairs of slow-code and fast-code for problems in the CodeNet challenge (Puri et al., 2021).”
• Teaching Large Language Models to Self-Debug (2023, methodology): “For code generation, a number of works have trained models to perform code optimization [34], interact with users for multi-turn code generation [62, 58, 38], and fix generated programs based on human feedback [25, 16, 6].”
• Large Language Models for Software Engineering: Survey and Open Problems (2023, methodology): “In the work by Madaan et al. [169], the authors use CODEGEN and CodeX to suggest functionally correct, Performance-Improving Edits (PIEs), improving execution time of Python and C++ (already pre-optimised with the maximally optimising compiler option -O3 ).”