# LLM Program Optimization via Retrieval Augmented Search

arXiv: https://arxiv.org/abs/2501.18916 · HTML: https://mlsanity.com/paper/2501.18916/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Placed in RAG for code, one of 15 directions in Retrieval-augmented generation.
- Builds on 15 of its 23 references that are in the map.
- First posted on arXiv 31 Jan 2025.

## Placed in

- Retrieval-augmented generation / [RAG for code](https://mlsanity.com/rag/rag-for-code/)

## Builds on (most cited first, 15 of 15)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Competition-Level Code Generation with AlphaCode](https://mlsanity.com/paper/2203.07814/) (2022)
- [CodeNet: A Large-Scale AI for Code Dataset for Learning a Diversity of Coding Tasks](https://mlsanity.com/paper/2105.12655/) (2021)
- [MapCoder: Multi-Agent Code Generation for Competitive Problem Solving](https://mlsanity.com/paper/2405.11403/) (2024)
- [Reframing Instructional Prompts to GPTk's Language](https://mlsanity.com/paper/2109.07830/) (2021)
- [Synchromesh: Reliable code generation from pre-trained language models](https://mlsanity.com/paper/2201.11227/) (2022)
- [Learning Performance-Improving Code Edits](https://mlsanity.com/paper/2302.07867/) (2023)
- [CodeRAG-Bench: Can Retrieval Augment Code Generation?](https://mlsanity.com/paper/2406.14497/) (2024)
- [How Efficient is LLM-Generated Code? A Rigorous & High-Standard Benchmark](https://mlsanity.com/paper/2406.06647/) (2024)
- [Search-Based LLMs for Code Optimization](https://mlsanity.com/paper/2408.12159/) (2024)
- [CoaCor: Code Annotation for Code Retrieval with Reinforcement Learning](https://mlsanity.com/paper/1904.00720/) (2019)
- [Finding Missed Code Size Optimizations in Compilers using LLMs](https://mlsanity.com/paper/2501.00655/) (2024)

## Built on by (most cited first, 4 of 4)

- [Toward Agentic AI: Generative Information Retrieval Inspired Intelligent Communications and Networking](https://mlsanity.com/paper/2502.16866/) (2025)
- [MaxCode: A Max-Reward Reinforcement Learning Framework for Automated Code Optimization](https://mlsanity.com/paper/2601.05475/) (2026)
- [Tape: A Cellular Automata Benchmark for Evaluating Rule-Shift Generalization in Reinforcement Learning](https://mlsanity.com/paper/2601.04695/) (2026)
- [RASP-Tuner: Retrieval-Augmented Soft Prompts for Context-Aware Black-Box Optimization in Non-Stationary Environments](https://mlsanity.com/paper/2604.18026/) (2026)

## Abstract (arXiv)

Recent work has demonstrated the potential of large language models (LLMs) for program optimization, a key challenge in programming languages. We propose a blackbox adaptation method called Retrieval Augmented Search (RAS) that performs beam search over candidate optimizations; at each step, it retrieves in-context examples from a given training dataset of slow-fast program pairs to guide the LLM. Critically, we find that performing contextual retrieval based on an LLM-generated natural language description significantly outperforms retrieval based on the source code. We also propose AEGIS, a method for improving interpretability by decomposing training examples into ''atomic edits'' that are significantly more incremental in nature. We show that RAS performs up to 2.06× better than prior state-of-the-art blackbox adaptation strategies on optimizing C++ programs, and that AEGIS performs up to 1.37× better while making significantly smaller edits. We also show that using RAS improves the mean runtime percentile of Python programs by 10.27 compared to baselines.