PaperlayerView exact revision on arXiv

Source-linked AI summary

LLM Program Optimization via Retrieval Augmented Search

Sagnik Anupam, Alexander Shypula, Osbert Bastani

arXiv:2501.18916v2cs.LG

TL;DR

LLMs struggle with program optimization, motivating methods that can adapt them using performance-oriented examples and more incremental edits. RAS uses retrieval-guided beam search, while AEGIS decomposes examples into atomic edits; RAS performs up to 2.06× better than prior blackbox strategies.

  • Problem

    LLMs have difficulty optimizing programs out of the box, and existing end-to-end approaches motivate testing more incremental optimization steps.

  • Method

    RAS retrieves examples using LLM-generated program descriptions and searches candidate optimizations with beam search, while AEGIS decomposes examples into incremental atomic edits.

  • Results

    RAS performs up to 2.06× better than dynamic retrieval, while AEGIS performs 1.37× better and produces smaller edits.

  • Takeaways & Limitations

    The methods provide a blackbox strategy for adapting LLMs to code optimization without fine-tuning, with AEGIS supporting more incremental edits.

  • Takeaways & Limitations

    Both methods are more computationally expensive because they use beam search, while AEGIS also requires additional training-time computation.

Abstract

from arXiv · show

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.

1 Introduction

LLMs show promise for program optimization but struggle out of the box, motivating retrieval-based adaptation that better reflects compiler-style optimization. The paper proposes RAS for contextual retrieve-optimize-evaluate search and AEGIS for smaller, more interpretable atomic edits.

  • Motivation: LLMs have difficulty optimizing programs out of the box because performance data is not widely available in traditional training datasets.This makes adaptation necessary for program optimization.
  • Motivation: Existing dynamic retrieval prompts an LLM with relevant slow-fast program pairs, but its end-to-end edits differ from systematic compiler passes.The retrieved examples are selected at test time using embedding similarity.
  • RAS: RAS combines contextual retrieval based on an LLM-generated natural-language program description with a retrieve-optimize-evaluate search loop.The description abstracts the program’s core algorithms and data structures from their implementation.
  • AEGIS: AEGIS decomposes slow-fast training pairs into sequences of incremental atomic edits associated with natural-language descriptions and explanations.It is designed to address the interpretability problems caused by large changes.
  • Evaluation: 8.70× average speedup compared to 4.23× for dynamic retrieval using Qwen3-Coder; AEGIS achieves a 6.08× average speedup using GPT-4o.These results are reported on the PIE benchmark for C++ program optimization.

2 Retrieval Augmented Search

RAS is a blackbox program-optimization framework that retrieves similar slow-fast training pairs during beam search and uses an LLM to generate optimized candidates. Its contextual retrieval embeds LLM-generated natural-language program descriptions rather than source code.

  • General framework: RAS assumes slow-fast training pairs, a new program, a retrieval distance, and blackbox LLM access for generating optimized program samples.The LLM receives an in-context slow-fast pair and a new program, then samples optimized versions without changing model weights.
  • General framework: At each beam-search step, RAS retrieves the top-k training pairs most similar to the current program and supplies them as in-context examples.Similarity is computed from distances between the current slow program and training programs, with FAISS used for vector search.
  • General framework: RAS selects the fastest candidate that passes all test cases, retains the previous program when candidates fail or time out, and returns the final program after m iterations.If the final program is invalid or times out, the procedure returns the source program p0.
  • Contextual retrieval: Contextual retrieval generates a natural-language description with Fcontext, embeds that description with ψ, and uses the resulting vector as ϕ(p).Training-example embeddings can be precomputed, and descriptions are prompted to capture algorithms and data structures.

3 Atomic Edit Guided Search

AEGIS preprocesses slow-fast training pairs into atomic edits to improve RAS interpretability. It uses incrementally applied, generalized edits as retrieval context so the optimizer can make smaller changes.

  • Method overview: AEGIS is a dataset preprocessing step designed to improve RAS interpretability.It is described with an overview in Figure 2 and pseudocode in Algorithm 2.
  • Atomic edits: Atomic edits are semantically equivalent program pairs that roughly differ by a single code optimization.The design is inspired by compiler passes that incrementally transform programs, making changes easier to understand.
  • Atomic-edit dataset: AEGIS replaces Πtrain with Πatomic, an aggregated dataset of edits decomposed from each slow-fast pair by Fdecomp.Retrieving atomic edits guides Fopt toward incremental optimizations rather than large changes.
  • Atomic-edit construction: Fedit applies each natural-language edit sequentially to p, producing programs p0 = p through pr that should resemble the original optimized program p′.Fdecomp describes pairwise differences as a list [s1, . . . , sr], and Fgen generalizes each edit into an atomic edit ei.
  • Integration with RAS: RAS uses each atomic edit (e, π) as additional Fopt input, where e gives optimization instructions and π provides an example application.The modified sampler produces an optimized version p′ from atomic edit (e, π) and program p.

4 Experiments

Experiments evaluate RAS and AEGIS on C++ PIE and Python Mercury program-optimization benchmarks against retrieval, search, and instruction-only baselines. Results show strong performance gains, while additional analyses examine correctness and AEGIS’s incremental edits.

  • Benchmarks: The evaluation uses PIE, a slow-fast C++ program-pair benchmark from CodeNet, and Mercury, a LeetCode dataset with Python solutions.C++ execution time is measured with the gem5 simulator.
  • Baselines: Baselines include dynamic retrieval, a no-contextual code-retrieval search ablation, an instruction-only approach, and human speedup references.Dynamic retrieval returns the fastest correct program among its sampled choices, whereas the no-contextual ablation performs multiple iterations without contextual retrieval.
  • Experimental setup: RAS and AEGIS use k = 8 retrievals, m = 4 beam-search steps, and h = 1 sample per generated prompt in PIE experiments.For Mercury, RAS uses m = 2 iterations with k = 8 and h = 1, while no contextual uses k = 4 and h = 8.
  • Metrics and analysis: Evaluation reports mean speedup and % Optimized for PIE, Pass@1 and Beyond@1 for Mercury, and analyzes correctness and edit-distance-based interpretability.The interpretability analysis treats lower character-level edit distance as more incremental change.
  • Results: RAS significantly improves performance over all baselines on C++, including the original PIE and atomic-edit training sets, and doubles dynamic retrieval’s mean best speedup.AEGIS outperforms the evaluated ablations except full RAS, with the gap attributed to decomposed pairs’ speedup and retrieval-diversity limitations.

5 Conclusion

The paper proposes RAS and AEGIS for LLM-guided program optimization, combining iterative beam search with retrieval to achieve blackbox speedups over existing LLM-based techniques. It also identifies computational cost and scalability to complex object-oriented codebases as key limitations.

  • Contributions: RAS and AEGIS are proposed as methods for LLM-guided program optimization.Both methods target optimization without fine-tuning.
  • Method: Beam search and retrieval iteratively optimize a given program in the proposed methods.The methods incorporate both components to guide successive optimization steps.
  • Results: The methods achieve significant blackbox speedups and outperform existing LLM-based program optimization techniques.The blackbox setting does not use fine-tuning.
  • Limitations: Beam search makes both approaches more computationally expensive to execute.AEGIS additionally requires training-time compute to construct its atomic dataset from LLM-generated code.
  • Limitations: Scaling the results to complex object-oriented codebases with several components may be challenging.Such scaling would likely require intermediate steps to identify specific components.

A Dataset Construction

The study constructs high-quality training and held-out evaluation sets for PIE and Mercury from benchmark program pairs and reported solution runtimes. PIE uses a problem-based train-test split, while Mercury evaluates on held-out problems using the slowest provided reference solutions.

  • PIE: PIE uses 4,080 training pairs and 973 held-out test pairs, selecting up to four highest-speedup pairs per competitive programming problem.The train-test split is based on the competitive programming problem being solved.
  • Mercury: Mercury uses 6,372 training pairs from 1,633 Leetcode training problems and evaluates 256 held-out problems using their slowest-provided reference solutions.Training pairs are constructed using Leetcode’s reported runtimes for the solutions.

B Compute · C Comparing Instruction Prompting and Expert Programmer System Roles

The experiments define the models, embedding system, simulator, compilation settings, and runtime environment used for optimization evaluation. The paper compares instruction prompting with an expert-programmer system role, using explicitly specified prompts and program placeholders.

  • B Compute: Qwen-3-Coder (480B parameters) and DeepSeek V3.2 (685B parameters) are the models used in the experiments.OpenAI’s text-embedding-3-large is used as the embedding model ψ.
  • B Compute: The gem5 simulator runs on a server with 2× Intel(R) Xeon(R) Gold 6342 CPUs and 96 cores total.C++ programs are compiled with g++ using the -O3 flag, while Mercury runtime percentiles are measured on an AWS t2.2xlarge instance.
  • C Comparing Instruction Prompting and Expert Programmer System Roles: The “Instruct Only” baseline tests two prompts: instruction prompting from the original PIE benchmark and an “expert programmer” system role.The paper provides the exact prompts for both approaches.
  • C Comparing Instruction Prompting and Expert Programmer System Roles: When referring to programs or retrieved natural language optimizations, the prompts enclose them in braces.This notation is used in the prompt specifications.
  • C.1 Instruction Prompting (IP): The instruction-prompting prompt asks the model to improve the performance of the given program.It introduces the input with “Given the program below, improve its performance:” and supplies the program to be optimized.
  • C.2 Expert Programmer System Role (EPSR): The expert-programmer system role requests faster rewritten source code in JSON, with the rewritten code stored under “optimized_code.”The prompt identifies the source code as the user input and instructs the model not to output anything other than C++ code.

C.3 Prompt Result Comparison · D Prompts for Experimental Results

The paper compares prompt variants on 973 programs and adopts an expert-programmer system role after observing a slight Mean Best Speedup increase. It also specifies prompts for description generation, contextual retrieval, dynamic retrieval, and AEGIS’s natural-language and atomic-edit workflows.

  • C.3 Prompt Result Comparison: The prompt comparison evaluates two prompts on 973 programs using k = 32 samples for m = 1 iteration of search, with results reported in Table 4.Table 4 compares metric differences between prompts in the Instruct Only setting.
  • C.3 Prompt Result Comparison: A slight increase in Mean Best Speedup with an expert-level system role leads the authors to use that role in subsequent prompts and reported Instruct Only results.The expert-programmer role is used by F′′ in the reported Tables 1 and 2 settings.
  • D Prompts for Experimental Results: The experimental prompts are presented for PIE, with “C++” replaced by “Python” for Mercury experiments.This establishes the language-specific adaptation between the two experiment suites.
  • D.1.1 Program Description Generation: Fcontext prompts the model to identify the source program’s algorithm and return a one-sentence description in a JSON object under the key “algorithm”.The prompt uses an expert-programmer system role and supplies the program to be optimized as the source program.
  • D.1.2 Generating Programs With Contextual Retrieval: Fopt provides a slow-fast program pair as context and asks the model to rewrite the source program into an optimized version.The prompt labels the retrieved faster program, the slower source program, and the requested optimized version.
  • D.1.3 Generating Programs With Dynamic Code Retrieval: RAS’s No Contextual and Dynamic Retrieval settings, and AEGIS’s No Contextual setting, use several slow-fast examples and request JSON output containing “optimized_code”.The final slower-version program is the source program, and the prompt asks the model to incorporate the optimizations from the examples.
  • D.2.1 Generating Natural Language Edits; D.2.2 Generating Program Sequence from Natural Language Edits; D.2.3 Generating Atomic Edits from Natural Language Edits: AEGIS decomposes training-pair changes into a JSON sequence of edits, explains why each edit may improve performance, and applies edits to generate program sequences.The workflow includes prompts for describing edits, applying them to source code, and generalizing them into atomic edits.
  • D.2.4 Generating Programs With Contextual Retrieval: AEGIS’s contextual-retrieval prompt supplies an atomic edit plus an example source-target pair and asks the model to incorporate the optimization into the given source code.The retrieved fields are labeled Optimization, Example Source, and Example Target.

E Additional Experimental Results · E.1 Metrics Across Beam Search Iterations · E.2 Comparison Between RAS and AEGIS

Additional experiments show that beam-search iterations continue improving optimization, particularly with contextual retrieval, while RAS and AEGIS differ in how edits are distributed across steps. Results across models and metrics indicate rapid early gains, with AEGIS continuing to improve despite more evenly distributed edits.

  • E.1 Metrics Across Beam Search Iterations: Figure 3 reports Mean Best Speedup, %Optimized, and Mean Edit Distance across beam-search steps for GPT-4o, Qwen-3-Coder, and Deepseek 3.2.The comparison focuses on the proposed approach versus the “No Contextual” ablation.
  • E.1 Metrics Across Beam Search Iterations: The first beam-search step provides the greatest benefit, but continued search improves all approaches, especially with contextual retrieval.The results are described as converging substantially more quickly when the first iteration already yields strong gains.
  • E.1 Metrics Across Beam Search Iterations: 1.1× speedup is reached by most programs after the first iteration, explaining why the results converge substantially more quickly.This condition is reported for the experiments discussed in Figure 3.
  • E.1 Metrics Across Beam Search Iterations: AEGIS continues gaining from additional search iterations, suggesting that continued search may close its performance gap.The passage presents this as evidence from the beam-search results rather than a quantified gap.
  • E.2 Comparison Between RAS and AEGIS: In the GPT-4o PIE experiment, RAS concentrates a large number of edits in the first optimization step.The optimization trajectory is illustrated in Figures 4 and 5.
  • E.2 Comparison Between RAS and AEGIS: AEGIS spreads its edits more evenly across different optimization steps than RAS.This contrast is reported for the GPT-4o experiment on the PIE dataset.

E.3 Failure Category Analysis of RAS and AEGIS · E.4 Mean Edit Distance · E.5 Impact of Code Embedding Models

The paper analyzes where AEGIS and RAS fail, reports mean edit distances, and tests whether specialized code embeddings affect contextual retrieval performance. These analyses identify failure categories and compare retrieval configurations using concrete test-set statistics.

  • E.3 Failure Category Analysis of RAS and AEGIS: 9.35% of the entire test set belongs to AEGIS’s unoptimized program set.These programs are test-set cases whose final best speedup after AEGIS is less than 1.1×.
  • E.3 Failure Category Analysis of RAS and AEGIS: Analyzing LLM-generated program descriptions reveals failure-associated algorithms and data structures that can guide targeted retrieval-data augmentation.The analysis compares term frequencies in each method’s unoptimized set with frequencies in the overall test set.
  • E.3 Failure Category Analysis of RAS and AEGIS: AEGIS struggles primarily with programs involving dynamic programming and binary search over sorted lists or trees.Dynamic-programming descriptions comprise 51.59% of the complete test set, while binary-search descriptions comprise 4.32 of programs.
  • E.3 Failure Category Analysis of RAS and AEGIS: 3.08% of the entire test set belongs to RAS’s unoptimized program set.For RAS, dynamic-programming problems comprise 16.67% and binary-search problems 6.67% of its unoptimized set.
  • E.3 Failure Category Analysis of RAS and AEGIS: 20% of RAS’s unoptimized program set involves descriptions mentioning Kruskal’s algorithm.Such problems constitute 1.03% of the total test set.
  • E.4 Mean Edit Distance: Mean edit distances are reported in Table 5.The supplied passage does not include the table’s numerical values.
  • E.5 Impact of Code Embedding Models: The code-embedding experiment replaces text-embedding-3-large with Codestral-Embed-2505 to test whether specialized embeddings reduce the contextual-versus-dynamic retrieval gap.It evaluates one iteration of RAS’s “No Contextual” ablation using DeepSeek 3.2, the best-performing model on PIE.
  • E.5 Impact of Code Embedding Models: 8.03 mean best speedup was achieved by RAS in the first iteration of the code-embedding experiment.The passage also reports 96.30% of the test set, but the supplied text truncates the corresponding outcome.

E.6 Mercury Results for Larger Models · F LLM Usage

The paper reports larger-model Mercury optimization results in the Instruct Only setting and illustrates RAS and AEGIS optimization trajectories from GPT-4o experiments. It also notes that LLM-generated experiment code was reviewed by the authors.

  • E.6 Mercury Results for Larger Models: Larger models were evaluated for optimizing programs in the Mercury dataset under the Instruct Only setting.The results are presented in Table 6.
  • F LLM Usage: LLMs generated code used to execute some experiments, and the authors reviewed that code.
  • F LLM Usage: Table 5 compares mean edit distances over steps between AEGIS and RAS on the PIE Benchmark.The supplied passage identifies the comparison but does not provide the table’s numerical values.
  • F LLM Usage: 10.06× is the final RAS speedup versus 9.58× for AEGIS in a randomly selected GPT-4o optimization trajectory.The human speedup for this example is 1.8×.
  • F LLM Usage: The RAS and AEGIS trajectories in one GPT-4o example implement similar optimizations and achieve similar speedups.The figure highlights changes between successive optimization steps.
  • F LLM Usage: 7.34× is the final RAS speedup versus 2.35× for AEGIS in a randomly selected GPT-4o trajectory where RAS significantly outperforms AEGIS.The human speedup for this example is 1.37×.
  • E.6 Mercury Results for Larger Models: Table 6 reports RAS experiments on Mercury using larger models.
Loading 2501.18916v2…