Optimización de Programas LLM mediante Búsqueda Aumentada por Recuperación

Resumen

Trabajos recientes han demostrado el potencial de los grandes modelos de lenguaje (LLM) para la optimización de programas, un desafío clave en lenguajes de programación. Proponemos un método de adaptación de caja negra denominado Búsqueda Aumentada por Recuperación (RAS, por sus siglas en inglés) que realiza una búsqueda en haz sobre optimizaciones candidatas; en cada paso, recupera ejemplos en contexto de un conjunto de datos de entrenamiento de pares de programas lentos-rápidos para guiar al LLM. De manera crucial, descubrimos que realizar una recuperación contextual basada en una descripción en lenguaje natural generada por el LLM supera significativamente a la recuperación basada en el código fuente. También proponemos AEGIS, un método para mejorar la interpretabilidad mediante la descomposición de ejemplos de entrenamiento en "ediciones atómicas" que son sustancialmente más incrementales en su naturaleza. Demostramos que RAS se desempeña hasta 2,06 veces mejor que las estrategias previas de adaptación de caja negra de última generación en la optimización de programas en C++, y que AEGIS logra hasta 1,37 veces mejor rendimiento al realizar ediciones significativamente más pequeñas. Asimismo, mostramos que el uso de RAS mejora el percentil de tiempo de ejecución medio de programas en Python en 10,27 en comparación con las líneas base.

English

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.06times better than prior state-of-the-art blackbox adaptation strategies on optimizing C++ programs, and that AEGIS performs up to 1.37times 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.