Don't Just Prompt an LLM to Optimize Code.
Retrieve+Search First.
We found that combining two strategies supercharges LLM-generated code optimization: retrieval and iterative search. While using Qwen3-Coder on the PIE C++ program optimization benchmark [1]: basic prompting gets 1.73x mean best speedup, code-based retrieval gets 4.23x, and four rounds of contextual retrieval-augmented search (RAS) gets 8.70x.
In this post, we explain why it works and how to outperform naive prompting.
Retrieval-Augmented Search (RAS)
Evaluations on the PIE benchmark have shown that using examples of slow programs and their faster equivalents in-context helps LLMs optimize programs better. However, a naive retrieval strategy simply retrieves the slow programs whose embedding is closest to the test set program. This gets us to 4.23x.
This strategy often misses relevant optimizations, because programs implementing similar algorithms or solving related problems can look completely different. Instead, we first prompt an LLM to describe what each program does (contextual retrieval). We then retrieve programs with similar descriptions rather than similar code. This helps us find examples that share algorithms and data structures rather than superficial similarities such as variable names.
Instead of just prompting once, we perform multiple iterations of search, prompting the LLM to repeatedly optimize the best-performing program it has generated. At each iteration:
Retrieve relevant examples.
Generate candidate programs by prompting LLM with in-context examples.
Evaluate and keep the best one.
Repeat.
This gets us to 8.70x within 4 iterations.
Distributing Edits
Here is a RAS-optimized program with a GPT-4o LLM backbone, sped up by 10.06x by iteration 4. Orange = changes made from the previous step, red = lines removed or modified in the next step.
RAS makes many changes in steps 1 and 2 and smaller modifications in later iterations, which makes the initial steps’ output hard to review for human programmers. In step 1, it implements two optimizations:
1) Faster method of reading input.
2) Skip updating arrays to avoid adding and subtracting the same numbers.
It makes sense to differentiate these optimizations and review them separately. We now ask a simple question: can we prompt the LLM to implement them in separate iterations?
Atomic-Edit Guided Search (AEGIS)
Core insight: if we retrieve examples containing only incremental edits, the LLM will likely also make an incremental edit.
Modification: Prompt an LLM to decompose program pairs in the retrieval dataset into a sequence of intermediate programs such that each program only implements one atomic edit. Then run RAS on the retrieval dataset with intermediate programs.
We call this method AEGIS.
When AEGIS is run with GPT-4o, it achieves a 9.58x speedup for the program we looked at earlier after four iterations, which is slower than RAS - there is an interpretability-performance tradeoff. But now, the two optimizations are spread out: one occurs in Step 1 and one in Step 4.
Practical Takeaways:
Retrieval of relevant examples + Search before optimization outperforms naive prompting.
If readability matters, encourage smaller edits rather than large refactors by decomposing retrieval dataset.
For more details, check out our code and our ACL 2026 Findings Paper.
If you found this helpful, please cite:
@article{anupam2026llm,
author = {Anupam, Sagnik and Shypula, Alexander and Bastani, Osbert},
title = {LLM Program Optimization via Retrieval Augmented Search},
journal = {Findings of the Association for Computational Linguistics: ACL 2026},
year = {2026},
url = {https://arxiv.org/abs/2501.18916},
}References:
Alexander Shypula, Aman Madaan, Yimeng Zeng, Uri Alon, Jacob R. Gardner, Yiming Yang, Milad Hashemi, Graham Neubig, Parthasarathy Ranganathan, Osbert Bastani, and Amir Yazdanbakhsh. 2024. Learning performance-improving code edits. In The Twelfth International Conference on Learning Representations.