Оптимизация LLM-программ посредством поиска с дополнением извлечения

Аннотация

Недавние работы продемонстрировали потенциал больших языковых моделей (БЯМ) для оптимизации программ — ключевой задачи в области языков программирования. Мы предлагаем метод адаптации «черного ящика» под названием Retrieval Augmented Search (RAS), который выполняет лучевой поиск по кандидатам оптимизаций; на каждом шаге он извлекает контекстные примеры из заданного обучающего набора пар медленных и быстрых программ для направления БЯМ. Критически важным является то, что выполнение контекстного извлечения на основе сгенерированного БЯМ описания на естественном языке значительно превосходит извлечение на основе исходного кода. Мы также предлагаем AEGIS — метод повышения интерпретируемости путем разложения обучающих примеров на «атомарные правки», которые по своей природе являются значительно более инкрементальными. Мы показываем, что RAS работает до 2,06 раза лучше, чем предыдущие передовые стратегии адаптации «черного ящика» при оптимизации программ на C++, и что AEGIS работает до 1,37 раза лучше, внося при этом значительно меньшие правки. Мы также показываем, что использование RAS улучшает средний процентиль времени выполнения программ на Python на 10,27 по сравнению с базовыми методами.

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.