🧑‍💻 Test Time Scaling for Code Generation

Раз уж пошел в стартап про кодогенерацию, то и статьи нужно читать и разбирать релевантные.

Сегодня очень прикладная работа про то, как быстро (точнее, просто за много сожжённых токенов у провайдеров) можно бустануть качество кодогенерации.

S∗: Test Time Scaling for Code Generation.

По факту адаптируют более общую идею из всем известной работы s1: Simple test-time scaling к домену кодогенерации.

Вся суть укладывается в одну картинку к посту. Дальше — детали:

1️⃣ Хотим генерировать код к некоторой задаче. Пусть есть доступный набор тестов (public tests), на которых можно тестировать программу. Есть предобученная модель, которая в целом может сгенерировать правильное решение, но всё-таки слишком стохастична, чтобы делать это с первой попытки. Хочется понять, как масштабировать доступный compute на инференсе, чтобы повысить качество генерируемых программ.

2️⃣ Начинаем с Parallel Branching — будем генерировать N программ одновременно. Очень понятный рабочий подход, более известный как Best-of-N.

3️⃣А теперь давайте добавим внутрь каждой ветки итеративный self-debugging. То есть модель сгенерировала код. Запускаем тесты. Какие-то из них падают. Используем это как фидбэк и просим модель поправить себя. Снова запускаем тесты — и так далее, в цикле M раз (пока есть бюджет). Это пример Sequential Scaling-а.

✅Шаги 2 и 3 вместе образуют Stage 1: Generation. Такая схема — это пример Hybrid Scaling-а, так как мы одновременно масштабируемся в обоих направлениях: в ширину (создаём много новых сэмплов) и в глубину (итеративно улучшаем существующие).

4️⃣После того как генерация отработала, у нас есть N версий кода, которые как-то работают на публичных тестах. Возможно, некоторые из них проходят все тесты. Появляется необходимость определить, какую версию кода выбрать в качестве финального ответа. Поэтому авторы вводят вторую стадию — Stage 2: Selection. Предлагается алгоритм:
a) кластеризуем все доступные сэмплы,
b) запускаем цикл по всем парам кластеров,
c) сэмплируем по одному решению из каждого кластера,
d) генерируем синтетические тесты для решений из двух кластеров,
e) запускаем эти тесты и добавляем +1 тому кластеру, который показывает лучший результат,
f) в качестве ответа выбираем кластер, набравший больше голосов, и из него сэмплируем финальное решение.

⚡️Дальше практические интересности:

✅Метод показывает существенно лучшие результаты, чем бейзлайны — обычный self-debugging и majority voting. Ablation studies подтверждают, что каждая фаза даёт вклад. Stage 2 в том виде, как у авторов, не всегда нужен.

✅Подход обобщается и отлично бустит все модели — как публичные, так и закрытые, включая reasoning-модели.

✅Температура сэмплирования играет значимую роль. Высокие значения приводят к деградации. Авторы выбрали 0.7 как оптимальное значение.

Итого: получаем очень простой фреймворк, который можно легко адаптировать под свою задачу и модифицировать в зависимости от бюджета и доступных инструментов для self-debugging-а или voting-а при выборе лучшего сэмпла. Опробовал сам — результаты хорошо совпадают. Много зон, где можно дальше улучшать.

#статья

@max_dot_sh