🤖 machine learning

LLM Program Optimization via Retrieval Augmented Search

This paper introduces Retrieval Augmented Search (RAS) and AEGIS, novel blackbox adaptation methods that leverage LLM-generated natural language descriptions and atomic edits to significantly outperform state-of-the-art strategies in optimizing C++ and Python programs while enhancing interpretability and minimizing code changes.

Original authors: Sagnik Anupam, Alexander Shypula, Osbert Bastani

Published 2026-06-25

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Sagnik Anupam, Alexander Shypula, Osbert Bastani

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a very talented but inexperienced chef (the Large Language Model, or LLM) who is trying to cook a meal. The goal isn't just to make the food taste good; it's to make it cook faster without changing the flavor.

The paper presents two new ways to help this chef work smarter, using a "library of past recipes" (a dataset of slow vs. fast code) to guide them.

The Problem: The Chef is Stuck

Usually, if you ask this chef to speed up a recipe, they might guess randomly or try to rewrite the whole thing at once. They often fail because they haven't seen enough examples of how to make things faster. They need a coach to show them specific tricks.

Method 1: RAS (Retrieval Augmented Search)

The Analogy: The "Smart Librarian" vs. The "Keyword Search"

Imagine the chef needs to find a similar recipe in a massive library to get inspiration.

The Old Way (Code Retrieval): The chef looks at the ingredients list (the actual code) of the current dish and searches the library for recipes with similar ingredients. This is like searching for "flour and eggs" and getting back a cake recipe when you actually need a soup recipe that just happens to use flour. It's too literal.
The New Way (Contextual Retrieval): The chef asks the librarian to write a summary of the dish first (e.g., "This is a soup that simmers vegetables to extract flavor"). The librarian then searches the library for other dishes that do the same thing, regardless of the specific ingredients.
- Why it works: It finds recipes that solve the same problem, not just ones that look the same on paper.

The Search Process (Beam Search):
Instead of asking the chef to fix the recipe in one giant leap, RAS breaks it down.

The chef makes a small improvement.
The librarian finds a new, relevant example based on this new version.
The chef tries again, using that new example.
They repeat this loop, like climbing a mountain step-by-step, always picking the path that goes up the fastest.

The Result: This method made C++ programs run up to 2 times faster than previous best methods and improved Python program speed significantly.

Method 2: AEGIS (Atomic Edit Guided Search)

The Analogy: The "Lego Master" vs. The "Demolition Crew"

Even with the smart librarian, the chef might still make huge, confusing changes to the recipe (like swapping the entire cooking method). This is hard to understand and risky.

AEGIS changes the library itself. Instead of storing whole "Slow vs. Fast" recipe pairs, it breaks them down into Atomic Edits.

The Process: The system takes a slow recipe and a fast recipe and asks an expert to explain the difference in tiny, single steps.
- Step 1: "Switch from a wooden spoon to a metal one (faster heat transfer)."
- Step 2: "Chop the onions finer (cooks faster)."
- Step 3: "Use a lid to trap steam."
The Generalization: It then rewrites these steps into general rules (e.g., "Use metal utensils for high heat") so they can be applied to any recipe, not just that one specific dish.

How it helps: When the chef needs to optimize a new dish, the system doesn't say "Here is a whole new recipe." It says, "Here is a specific, tiny trick: 'Switch to a metal spoon.' Try that." The chef applies one tiny trick, checks the result, then tries the next.

The Result:

It makes the changes smaller and easier to understand (like swapping one Lego brick instead of rebuilding the castle).
It reduced the size of the changes by 17% to 30% compared to the first method.
It still made programs faster, though slightly less than the first method, but with much more control and clarity.

Summary of Results

RAS is like having a smart coach who guides you through a series of small, informed steps, finding the best path to a faster program. It beat the previous best method by 2x for C++ and improved Python speed by 10% on average.
AEGIS is like having a coach who breaks the problem down into tiny, manageable Lego blocks. It makes the changes 1.37x better in terms of efficiency and makes the edits much smaller and safer, ensuring the program doesn't break while trying to get faster.

In short, the paper teaches us that to make AI code faster, we shouldn't just ask it to "guess." We should give it a librarian who understands the story of the code (Contextual Retrieval) and a set of tiny, specific tools (Atomic Edits) to make small, safe improvements one by one.

Technical Summary: LLM Program Optimization via Retrieval Augmented Search

Problem Statement

Large Language Models (LLMs) have demonstrated significant potential in code generation but struggle with program optimization out-of-the-box. Traditional training datasets lack sufficient data on program performance, making direct adaptation difficult. While recent work (Shypula et al., 2024) introduced the "Performance Improving Edits" (PIE) benchmark and demonstrated that dynamic code retrieval (retrieving slow-fast program pairs based on code embedding similarity) improves performance, existing approaches remain "end-to-end." They ask an LLM to directly output an optimized version of a program based on a few examples. This contrasts with modern compilers, which systematically modify programs through a series of incremental passes. Furthermore, end-to-end approaches often produce large, difficult-to-interpret changes.

Methodology

The paper proposes two primary methods to address these limitations: Retrieval Augmented Search (RAS) and Atomic Edit GuIded Search (AEGIS).

1. Retrieval Augmented Search (RAS)

RAS is a blackbox adaptation method that performs a beam search over candidate optimizations, integrating two key insights to improve dynamic retrieval:

Contextual Retrieval: Instead of retrieving examples based on the raw source code, RAS uses an LLM ( $F_{context}$ ) to generate a natural language description of the program's core algorithms and data structures. This description is then embedded ( $\psi$ ) to retrieve the most relevant slow-fast training pairs. The authors hypothesize that this abstraction allows the system to retrieve examples based on intent and logic rather than superficial syntactic similarities.
Iterative Beam Search: Rather than retrieving a fixed set of examples once, RAS performs a retrieve-optimize-evaluate loop. At each step $i$ , given the current program $p_{i-1}$ , the system retrieves the top- $k$ training pairs most similar to $p_{i-1}$ (based on the contextual embedding). The LLM ( $F_{opt}$ ) then generates $k$ candidate optimizations. The system selects the fastest correct program as $p_i$ and repeats the process for $m$ steps.

2. Atomic Edit GuIded Search (AEGIS)

To address the interpretability issue of large, monolithic changes, AEGIS introduces a preprocessing step to decompose training examples into "atomic edits."

Decomposition: An LLM ( $F_{decomp}$ ) analyzes a slow-fast program pair $(p, p')$ and decomposes the transformation into a sequence of natural language descriptions $[s_1, ..., s_r]$ , where each $s_i$ explains a specific incremental change and its performance rationale.
Generation and Generalization: The system uses an LLM ( $F_{edit}$ ) to apply each description $s_i$ sequentially to the slow program to generate intermediate programs. Then, a generalization LLM ( $F_{gen}$ ) rewrites each specific edit description into a generalizable "atomic edit" $e_i$ that applies to a wider variety of programs.
Search with Atomic Edits: The resulting dataset $\Pi_{atomic}$ consists of tuples $(e_i, (p_{i-1}, p_i))$ . RAS is then run on this dataset, where the retrieval step finds the most relevant atomic edit, and the LLM is prompted to apply that specific incremental edit to the current program. This forces the optimization process to be more incremental and interpretable.

Key Contributions

Contextual Retrieval Superiority: The paper demonstrates that retrieving examples based on LLM-generated natural language descriptions significantly outperforms retrieval based on raw source code embeddings. This suggests that abstracting the algorithmic intent is crucial for effective few-shot learning in optimization.
RAS Framework: The authors propose RAS, which combines contextual retrieval with iterative beam search. This approach treats optimization as a sequential search problem rather than a single-step generation task.
AEGIS for Interpretability: The paper introduces AEGIS, a method to distill training data into atomic, generalizable edits. This allows the system to perform optimizations in smaller, more interpretable steps while maintaining performance gains.
Comprehensive Evaluation: The methods are evaluated on the PIE benchmark (C++) and the Mercury benchmark (Python), comparing against state-of-the-art baselines including dynamic retrieval, instruction-only prompting, and human performance.

Results

C++ Optimization (PIE Benchmark)

RAS Performance: Using Qwen3-Coder, RAS achieved an average speedup of 8.70×, compared to 4.23× for the previous state-of-the-art dynamic retrieval baseline. This represents a 2.06× improvement over the baseline.
AEGIS Performance: Using GPT-4o, AEGIS achieved a 6.08× average speedup. While slightly lower than full RAS (which achieved higher speedups with specific models), AEGIS produced significantly smaller edits.
Edit Size: AEGIS reduced the average edit size (measured by string edit distance) by 17% compared to RAS (with GPT-4o) and by 30% when considering only the first edit in the search process.
Accuracy: In GPT-4o experiments, RAS maintained a 99.5% accuracy rate (programs passing all test cases), while AEGIS achieved 100%. Formal verification using CBMC on a subset of DeepSeek 3.2 outputs confirmed semantic equivalence in 8/10 cases, with the remaining 2 differing only in minor formatting (newlines).

Python Optimization (Mercury Benchmark)

Using RAS with Qwen2.5-7B-Instruct, the mean runtime percentile metric (Beyond@1) improved by 10.27 compared to baselines, significantly narrowing the performance gap between this model and larger reasoning models.

Significance and Claims

The paper claims that breaking down program optimization into incremental steps via retrieval-augmented search significantly improves the ability of LLMs to optimize code. The authors emphasize that contextual retrieval (using natural language descriptions) is a critical factor in outperforming code-based retrieval. Furthermore, they argue that AEGIS offers a viable path to making LLM-driven optimizations more interpretable and safer by enforcing incremental changes, without sacrificing substantial performance gains. The work positions these blackbox adaptation strategies as a practical alternative to fine-tuning, capable of leveraging existing training datasets of slow-fast pairs to achieve state-of-the-art results in program optimization.