Horizon-Length Prediction: Advancing Fill-in-the-Middle Capabilities for
  Code Generation with Lookahead Planning
Authors: name: Yifeng Ding; name: Hantian Ding; name: Shiqi Wang; name: Qing Sun; name: Varun Kumar; name: Zijian Wang
Published: 2024-10-04 02:53:52

Abstract: 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.