Otimização de Programas LLM via Busca Aumentada por Recuperação

Resumo

Trabalhos recentes demonstraram o potencial dos modelos de linguagem de grande porte (LLMs) para otimização de programas, um desafio fundamental em linguagens de programação. Propomos um método de adaptação de caixa-preta chamado Retrieval Augmented Search (RAS), que realiza busca em feixe sobre otimizações candidatas; a cada etapa, ele recupera exemplos no contexto de um conjunto de dados de treinamento de pares de programas lentos-rápidos para orientar o LLM. Criticamente, descobrimos que a recuperação contextual baseada em uma descrição em linguagem natural gerada pelo LLM supera significativamente a recuperação baseada no código-fonte. Também propomos AEGIS, um método para melhorar a interpretabilidade ao decompor exemplos de treinamento em "edições atômicas" que são significativamente mais incrementais por natureza. Mostramos que o RAS apresenta desempenho até 2,06 vezes superior ao das estratégias de adaptação de caixa-preta do estado da arte anterior na otimização de programas em C++, e que o AEGIS apresenta desempenho até 1,37 vezes superior, realizando edições significativamente menores. Também mostramos que o uso do RAS melhora o percentil do tempo de execução médio de programas Python em 10,27 em comparação com as linhas de 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.