Voorspelling van de Horizontale Lengte: Het verbeteren van de invulmogelijkheden in het midden voor Codegeneratie met plannen op basis van vooruitkijken

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

October 4, 2024

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

cs.AI

Samenvatting

Het invullen van de lege plekken (FIM) is essentieel geworden voor code-taalmodellen, waardoor het genereren van ontbrekende code mogelijk is gegeven zowel de linker- als rechtercontext. Echter, het huidige FIM-trainingsparadigma, dat de oorspronkelijke trainingssequenties herordent en vervolgens reguliere voorspelling van het volgende token (NTP) uitvoert, leidt vaak tot modellen die moeite hebben om inhoud te genereren die soepel aansluit bij de omringende context. Belangrijk is dat, terwijl bestaande werken vertrouwen op op regels gebaseerde post-processing om deze zwakte te omzeilen, dergelijke methoden niet praktisch bruikbaar zijn in open-domein code-completietaken omdat ze afhankelijk zijn van beperkende, dataset-specifieke aannames (bijv. het genereren van hetzelfde aantal regels als in de grondwaarheid). Bovendien verslechtert de modelprestatie op FIM-taken aanzienlijk zonder deze onrealistische aannames. We veronderstellen dat NTP alleen onvoldoende is voor modellen om effectieve planning te leren die is geconditioneerd op de verre rechtercontext, een cruciale factor voor succesvolle code-infilling. Om dit te overwinnen, stellen we Horizon-Lengte Voorspelling (HLP) voor, een nieuw trainingsdoel dat modellen leert om het aantal resterende middelste tokens (d.w.z. horizonlengte) bij elke stap te voorspellen. HLP bevordert FIM met vooruitkijkende planning, waardoor modellen inherent infillingsgrenzen kunnen leren voor willekeurige linker- en rechtercontexten zonder afhankelijk te zijn van dataset-specifieke post-processing. Onze evaluatie over verschillende modellen en groottes toont aan dat HLP de FIM-prestaties aanzienlijk verbetert met maximaal 24% relatief op diverse benchmarks, op zowel bestandsniveau als repositoryniveau, en zonder toevlucht te nemen tot onrealistische post-processing-methoden. Bovendien verbetert de verbeterde planningscapaciteit die wordt verkregen door HLP de modelprestatie op code-redenering. Belangrijk is dat HLP slechts verwaarloosbare trainingskosten met zich meebrengt en geen extra inferentiekosten, waardoor het praktisch toepasbaar is voor real-world scenario's.

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.