Прогнозирование горизонта предсказания: развитие возможностей заполнения пропусков для генерации кода с планированием предпросмотра.

Horizon-Length Prediction: Advancing Fill-in-the-Middle Capabilities for Code Generation with Lookahead Planning

October 4, 2024

Авторы: Yifeng Ding, Hantian Ding, Shiqi Wang, Qing Sun, Varun Kumar, Zijian Wang

cs.AI

Аннотация

Заполнение-середина (Fill-in-the-Middle, FIM) стало неотъемлемой частью языковых моделей кода, позволяя генерировать отсутствующий код на основе контекста слева и справа. Однако текущая парадигма обучения FIM, которая переупорядочивает исходные обучающие последовательности и затем выполняет обычное предсказание следующего токена (Next-Token Prediction, NTP), часто приводит к тому, что модели испытывают затруднения при генерации контента, который гармонично сочетается с окружающим контекстом. Критически важно отметить, что существующие работы полагаются на правила основанного на правилах постобработки для обхода этого недостатка, однако такие методы не могут быть практически использованы в задачах открытой области завершения кода, поскольку они зависят от ограничительных, специфичных для набора данных предположений (например, генерация того же количества строк, что и в истине). Более того, производительность модели на задачах FIM значительно ухудшается без этих нереалистичных предположений. Мы предполагаем, что NTP в одиночку недостаточен для того, чтобы модели могли эффективно учиться планированию, зависящему от дальнего правого контекста, что является критическим фактором для успешного заполнения кода. Для преодоления этого мы предлагаем Предсказание Длины Горизонта (Horizon-Length Prediction, HLP), новую цель обучения, которая учит модели предсказывать количество оставшихся средних токенов (т.е. длину горизонта) на каждом шаге. HLP продвигает FIM с планированием просмотра вперед, позволяя моделям внутренне учиться границам заполнения для произвольных левых и правых контекстов без использования специфичной для набора данных постобработки. Наша оценка на различных моделях и размерах показывает, что HLP значительно улучшает производительность FIM на до 24% относительно на различных бенчмарках, на уровне файлов и репозиториев, и без прибегания к нереалистичным методам постобработки. Более того, улучшенная способность планирования, полученная благодаря HLP, повышает производительность модели в рассуждениях о коде. Важно отметить, что HLP обладает незначительными накладными расходами на обучение и не приводит к дополнительным затратам на вывод, обеспечивая его практичность для реальных сценариев.

English

Fill-in-the-Middle (FIM) has become integral to code language models, enabling generation of missing code given both left and right contexts. However, the current FIM training paradigm, which reorders original training sequences and then performs regular next-token prediction (NTP), often leads to models struggling to generate content that aligns smoothly with the surrounding context. Crucially, while existing works rely on rule-based post-processing to circumvent this weakness, such methods are not practically usable in open-domain code completion tasks as they depend on restrictive, dataset-specific assumptions (e.g., generating the same number of lines as in the ground truth). Moreover, model performance on FIM tasks deteriorates significantly without these unrealistic assumptions. We hypothesize that NTP alone is insufficient for models to learn effective planning conditioned on the distant right context, a critical factor for successful code infilling. To overcome this, we propose Horizon-Length Prediction (HLP), a novel training objective that teaches models to predict the number of remaining middle tokens (i.e., horizon length) at each step. HLP advances FIM with lookahead planning, enabling models to inherently learn infilling boundaries for arbitrary left and right contexts without relying on dataset-specific post-processing. Our evaluation across different models and sizes shows that HLP significantly improves FIM performance by up to 24% relatively on diverse benchmarks, across file-level and repository-level, and without resorting to unrealistic post-processing methods. Furthermore, the enhanced planning capability gained through HLP boosts model performance on code reasoning. Importantly, HLP only incurs negligible training overhead and no additional inference cost, ensuring its practicality for real-world scenarios.