Previsão de Comprimento de Horizonte: Avançando as Capacidades de Preenchimento no Meio para Geração de Código com Planejamento de Antecipação

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

October 4, 2024

Autores: Yifeng Ding, Hantian Ding, Shiqi Wang, Qing Sun, Varun Kumar, Zijian Wang

cs.AI

Resumo

O Fill-in-the-Middle (FIM) tornou-se fundamental para os modelos de linguagem de código, permitindo a geração de código ausente considerando os contextos à esquerda e à direita. No entanto, o paradigma atual de treinamento do FIM, que reordena sequências de treinamento originais e depois realiza previsão regular do próximo token (NTP), frequentemente resulta em modelos com dificuldades para gerar conteúdo que se alinhe suavemente com o contexto circundante. De forma crucial, enquanto trabalhos existentes dependem de pós-processamento baseado em regras para contornar essa fraqueza, tais métodos não são praticamente utilizáveis em tarefas de conclusão de código em domínio aberto, pois dependem de suposições restritivas e específicas do conjunto de dados (por exemplo, gerar o mesmo número de linhas que na verdade real). Além disso, o desempenho do modelo em tarefas de FIM deteriora significativamente sem essas suposições irreais. Hipothetizamos que NTP sozinho é insuficiente para os modelos aprenderem um planejamento eficaz condicionado ao contexto distante à direita, um fator crítico para o preenchimento bem-sucedido de código. Para superar isso, propomos a Predição do Comprimento do Horizonte (HLP), um objetivo de treinamento inovador que ensina os modelos a prever o número de tokens intermediários restantes (ou seja, comprimento do horizonte) em cada etapa. HLP avança o FIM com planejamento de antecipação, permitindo que os modelos aprendam de forma inerente os limites de preenchimento para contextos arbitrários à esquerda e à direita sem depender de pós-processamento específico do conjunto de dados. Nossa avaliação em diferentes modelos e tamanhos mostra que o HLP melhora significativamente o desempenho do FIM em até 24% relativamente em diversos benchmarks, em nível de arquivo e de repositório, e sem recorrer a métodos de pós-processamento irreais. Além disso, a capacidade aprimorada de planejamento obtida por meio do HLP impulsiona o desempenho do modelo no raciocínio de código. Importante ressaltar que o HLP incorre apenas em um custo de treinamento negligenciável e nenhum custo adicional de inferência, garantindo sua praticidade para cenários do mundo real.

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.