Horizontlängenvorhersage: Fortschrittliche Fill-in-the-Middle-Fähigkeiten für die Codegenerierung mit Vorausplanung

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

October 4, 2024

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

cs.AI

Zusammenfassung

Das Ausfüllen der Lücke (Fill-in-the-Middle, FIM) ist integraler Bestandteil von Code-Sprachmodellen, da es die Generierung fehlenden Codes ermöglicht, basierend auf sowohl dem linken als auch dem rechten Kontext. Allerdings führt das aktuelle FIM-Trainingsparadigma, das die ursprünglichen Trainingssequenzen neu anordnet und dann reguläre Vorhersagen des nächsten Tokens (Next-Token Prediction, NTP) durchführt, oft dazu, dass Modelle Schwierigkeiten haben, Inhalte zu generieren, die nahtlos zum umgebenden Kontext passen. Entscheidend ist, dass bestehende Arbeiten auf regelbasierte Nachbearbeitung angewiesen sind, um diese Schwäche zu umgehen. Solche Methoden sind jedoch in offenen Code-Vervollständigungsaufgaben nicht praktikabel, da sie auf restriktiven, datensatzspezifischen Annahmen beruhen (z. B. Generierung derselben Anzahl von Zeilen wie im Ground Truth). Darüber hinaus verschlechtert sich die Leistung der Modelle bei FIM-Aufgaben signifikant ohne diese unrealistischen Annahmen. Wir vermuten, dass NTP allein nicht ausreicht, damit Modelle effektive Planung basierend auf dem entfernten rechten Kontext erlernen, was ein entscheidender Faktor für erfolgreiches Code-Ausfüllen ist. Um dies zu überwinden, schlagen wir Horizon-Length Prediction (HLP) vor, ein neuartiges Trainingsziel, das den Modellen beibringt, die Anzahl der verbleibenden mittleren Tokens (d. h. die Horizontlänge) in jedem Schritt vorherzusagen. HLP verbessert FIM durch vorausschauende Planung und ermöglicht es den Modellen, Füllgrenzen für beliebige linke und rechte Kontexte inhärent zu erlernen, ohne auf datensatzspezifische Nachbearbeitung angewiesen zu sein. Unsere Evaluation über verschiedene Modelle und Größen hinweg zeigt, dass HLP die Leistung von FIM signifikant um bis zu 24 % relativ auf verschiedenen Benchmarks verbessert, sowohl auf Dateiebene als auch auf Repository-Ebene, und ohne auf unrealistische Nachbearbeitungsmethoden zurückzugreifen. Darüber hinaus verbessert die erweiterte Planungsfähigkeit, die durch HLP erlangt wird, die Leistung der Modelle bei der Code-Argumentation. Wichtig ist, dass HLP nur einen vernachlässigbaren Schulungsaufwand und keine zusätzlichen Inferenzkosten verursacht, was seine Praktikabilität für Szenarien in der realen Welt sicherstellt.

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.