Predizione della Lunghezza dell'Orizzonte: Avanzamento delle Capacità di Riempimento Centrale per la Generazione di Codice con Pianificazione Anticipata

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

October 4, 2024

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

cs.AI

Abstract

Il Fill-in-the-Middle (FIM) è diventato essenziale per i modelli di linguaggio di codice, consentendo la generazione di codice mancante dati contesti sia a sinistra che a destra. Tuttavia, il paradigma attuale di addestramento FIM, che riordina le sequenze di addestramento originali e poi esegue regolarmente la previsione del prossimo token (NTP), porta spesso i modelli a faticare nella generazione di contenuti che si allineano in modo fluido con il contesto circostante. In modo cruciale, mentre i lavori esistenti si basano su un post-processing basato su regole per aggirare questa debolezza, tali metodi non sono praticamente utilizzabili in compiti di completamento del codice in un dominio aperto poiché dipendono da assunzioni restrittive e specifiche del dataset (ad esempio, generando lo stesso numero di righe del ground truth). Inoltre, le prestazioni del modello nei compiti FIM peggiorano significativamente senza queste assunzioni irrealistiche. Ipotizziamo che la sola NTP non sia sufficiente affinché i modelli imparino una pianificazione efficace condizionata al contesto a destra distante, un fattore critico per un riempimento di codice di successo. Per superare ciò, proponiamo la Predizione della Lunghezza dell'Orizzonte (HLP), un nuovo obiettivo di addestramento che insegna ai modelli a prevedere il numero di token intermedi rimanenti (cioè, la lunghezza dell'orizzonte) ad ogni passo. HLP fa progredire il FIM con una pianificazione di previsione, consentendo ai modelli di imparare in modo innato i confini di riempimento per contesti arbitrari a sinistra e a destra senza dipendere da un post-processing specifico del dataset. La nostra valutazione su diversi modelli e dimensioni mostra che HLP migliora significativamente le prestazioni del FIM fino al 24% relativamente su diversi benchmark, a livello di file e di repository, e senza ricorrere a metodi di post-processing irrealistici. Inoltre, la capacità di pianificazione potenziata acquisita attraverso HLP migliora le prestazioni del modello nel ragionamento del codice. È importante sottolineare che HLP comporta solo un sovraccarico di addestramento trascurabile e nessun costo aggiuntivo di inferenza, garantendo la sua praticità per scenari reali.

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.