arXiv · 2501.18916
LLM Program Optimization via Retrieval Augmented Search
Abstract
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$\times$ better than prior state-of-the-art blackbox adaptation strategies on optimizing C++ programs, and that AEGIS performs up to 1.37$\times$ 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.
Explore related subjects
Keep this discovery
Explore connections, maps & timelines
Sagnik Anupam, Alexander Shypula, Osbert Bastani. 2025-01-31. LLM Program Optimization via Retrieval Augmented Search. https://arxiv.org/abs/2501.18916
Cite the original work for its findings. Save a collection to share your selection of sources.