Predicción de Longitud de Horizonte: Avanzando las Capacidades de Rellenado en el Medio para la Generación de Código con Planificación de Anticipación

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

Resumen

Fill-in-the-Middle (FIM) se ha vuelto fundamental para los modelos de lenguaje de código, permitiendo la generación de código faltante dado tanto el contexto izquierdo como el derecho. Sin embargo, el paradigma actual de entrenamiento de FIM, que reordena las secuencias de entrenamiento originales y luego realiza una predicción regular del siguiente token (NTP), a menudo conduce a que los modelos tengan dificultades para generar contenido que se alinee suavemente con el contexto circundante. Es crucial destacar que, si bien los trabajos existentes dependen de un postprocesamiento basado en reglas para evitar esta debilidad, dichos métodos no son prácticamente utilizables en tareas de completado de código de dominio abierto, ya que dependen de suposiciones restrictivas y específicas del conjunto de datos (por ejemplo, generar el mismo número de líneas que en la verdad fundamental). Además, el rendimiento del modelo en tareas de FIM se deteriora significativamente sin estas suposiciones poco realistas. Hacemos la hipótesis de que NTP por sí solo es insuficiente para que los modelos aprendan una planificación efectiva condicionada al contexto derecho distante, un factor crítico para el éxito del relleno de código. Para superar esto, proponemos Predicción de Longitud de Horizonte (HLP), un objetivo de entrenamiento novedoso que enseña a los modelos a predecir el número de tokens medios restantes (es decir, longitud de horizonte) en cada paso. HLP avanza FIM con una planificación de anticipación, lo que permite a los modelos aprender inherentemente los límites de relleno para contextos izquierdos y derechos arbitrarios sin depender de un postprocesamiento específico del conjunto de datos. Nuestra evaluación en diferentes modelos y tamaños muestra que HLP mejora significativamente el rendimiento de FIM hasta un 24% en diversos puntos de referencia, tanto a nivel de archivo como de repositorio, y sin recurrir a métodos poco realistas de postprocesamiento. Además, la capacidad de planificación mejorada obtenida a través de HLP impulsa el rendimiento del modelo en el razonamiento de código. Es importante destacar que HLP solo conlleva un costo de entrenamiento insignificante y ningún costo adicional de inferencia, asegurando su practicidad para escenarios del 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.