Prédiction de la Longueur de l'Horizon : Amélioration des Capacités de Complétion Automatique pour la Génération de Code avec Planification Anticipée

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

Résumé

Le Fill-in-the-Middle (FIM) est devenu essentiel pour les modèles de langage de code, permettant la génération de code manquant en tenant compte à la fois des contextes gauche et droit. Cependant, le paradigme actuel d'entraînement FIM, qui réordonne les séquences d'entraînement originales puis effectue une prédiction régulière du jeton suivant (NTP), conduit souvent à des modèles ayant du mal à générer du contenu qui s'aligne harmonieusement avec le contexte environnant. De manière cruciale, alors que les travaux existants reposent sur un post-traitement basé sur des règles pour contourner cette faiblesse, de telles méthodes ne sont pas praticables dans les tâches d'achèvement de code en domaine ouvert car elles dépendent d'hypothèses restrictives et spécifiques au jeu de données (par exemple, générer le même nombre de lignes que dans la vérité terrain). De plus, les performances du modèle sur les tâches FIM se détériorent significativement sans ces hypothèses irréalistes. Nous émettons l'hypothèse que NTP seul est insuffisant pour que les modèles apprennent une planification efficace conditionnée par le contexte droit éloigné, un facteur critique pour un remplissage de code réussi. Pour surmonter cela, nous proposons la Prédiction de Longueur d'Horizon (HLP), un nouvel objectif d'entraînement qui enseigne aux modèles à prédire le nombre de jetons de milieu restants (c'est-à-dire la longueur de l'horizon) à chaque étape. HLP fait progresser FIM avec une planification à long terme, permettant aux modèles d'apprendre de manière inhérente les limites de remplissage pour des contextes gauche et droit arbitraires sans dépendre d'un post-traitement spécifique au jeu de données. Notre évaluation sur différents modèles et tailles montre que HLP améliore significativement les performances de FIM jusqu'à 24% relativement sur divers benchmarks, au niveau des fichiers et des dépôts, et sans recourir à des méthodes de post-traitement irréalistes. De plus, la capacité de planification améliorée acquise grâce à HLP renforce les performances du modèle en matière de raisonnement sur le code. Importamment, HLP n'entraîne qu'une surcharge d'entraînement négligeable et aucun coût d'inférence supplémentaire, garantissant sa praticité pour les scénarios réels.

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.