Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

LLM Program Optimization via Retrieval Augmented Search

Sagnik Anupam, Alexander Shypula, Osbert Bastani

TL;DR

The paper tackles the problem of improving program performance using LLMs in settings lacking abundant performance data. It introduces Retrieval Augmented Search (RAS), which uses contextual NL descriptions and beam search to iteratively retrieve and apply related slow-fast program pairs, achieving substantial speedups on the PIE benchmark. To address interpretability, it also proposes Atomic Edit Guided Search (AEGIS), which decomposes optimizations into incremental atomic edits and trains the system to apply them sequentially. The combination of contextual retrieval, iterative search, and incremental edits yields significant performance gains over prior blackbox methods, with practical implications for safer and more transparent LLM-assisted code optimization, albeit with higher computational cost.

Abstract

With the advent of large language models (LLMs), there has been a great deal of interest in applying them to solve difficult programming tasks. Recent work has demonstrated their potential at program optimization, a key challenge in programming languages research. 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. In addition, we propose a method called AEGIS for improving interpretability by decomposing training examples into "atomic edits" that are significantly more incremental in nature. We show that RAS performs 1.8$\times$ better than prior state-of-the-art blackbox adaptation strategies, and that AEGIS performs 1.37$\times$ better while performing significantly smaller edits.

LLM Program Optimization via Retrieval Augmented Search

TL;DR

The paper tackles the problem of improving program performance using LLMs in settings lacking abundant performance data. It introduces Retrieval Augmented Search (RAS), which uses contextual NL descriptions and beam search to iteratively retrieve and apply related slow-fast program pairs, achieving substantial speedups on the PIE benchmark. To address interpretability, it also proposes Atomic Edit Guided Search (AEGIS), which decomposes optimizations into incremental atomic edits and trains the system to apply them sequentially. The combination of contextual retrieval, iterative search, and incremental edits yields significant performance gains over prior blackbox methods, with practical implications for safer and more transparent LLM-assisted code optimization, albeit with higher computational cost.

Abstract

With the advent of large language models (LLMs), there has been a great deal of interest in applying them to solve difficult programming tasks. Recent work has demonstrated their potential at program optimization, a key challenge in programming languages research. 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. In addition, we propose a method called AEGIS for improving interpretability by decomposing training examples into "atomic edits" that are significantly more incremental in nature. We show that RAS performs 1.8 better than prior state-of-the-art blackbox adaptation strategies, and that AEGIS performs 1.37 better while performing significantly smaller edits.

Paper Structure

This paper contains 26 sections, 9 equations, 4 figures, 4 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Retrieval Augmented Search
  4. Problem Formulation
  5. General Framework
  6. Contextual Retrieval
  7. Atomic Edit Guided Search
  8. Experiments
  9. Experimental Setup
  10. Results
  11. Conclusion
  12. Comparing Instruction Prompting and Expert Programmer System Roles
  13. Instruction Prompting (IP)
  14. Expert Programmer System Role (EPSR)
  15. Prompt Result Comparison
  16. ...and 11 more sections

Figures (4)

  • Figure 1: Overview of RAS
  • Figure 2: Overview of Aegis
  • Figure 3: Results across beam search steps.
  • Figure 4: We show a randomly selected example optimization trajectory where we demonstrate the improvements made at each step of RAS vs. Aegis. We have highlighted lines that have changed from the previous step in orange, while lines that change in the next step have been highlighted in red.