Planning-Aware Code Infilling via Horizon-Length Prediction

  • 2025-07-16 15:44:04
  • Yifeng Ding, Hantian Ding, Shiqi Wang, Qing Sun, Varun Kumar, Zijian Wang
  • 0

Abstract

Fill-in-the-Middle (FIM), or infilling, has become integral to code languagemodels, enabling generation of missing code given both left and right contexts.However, the current FIM training paradigm which performs next-token prediction(NTP) over reordered sequence often leads to models struggling to generatecontent that aligns well with the surrounding context. We hypothesize that NTPalone is insufficient for models to learn effective planning conditioned on thedistant right context, a critical factor for successful code infilling. Toovercome this, we propose Horizon-Length Prediction (HLP), a novel trainingobjective that teaches models to predict the number of remaining middle tokensat each step. HLP advances FIM with lookahead planning, enabling models toinherently learn infilling boundaries for arbitrary left and right contextswithout relying on dataset-specific post-processing. Our evaluation acrossdifferent model families and sizes shows that HLP significantly improves FIMperformance by up to 24% relatively on diverse benchmarks, across file-leveland repository-level. Furthermore, the enhanced planning capability gainedthrough HLP boosts model performance on code reasoning. Importantly, HLP incursnegligible training overhead and no additional inference cost, ensuring itspracticality for real-world scenarios.

 

Quick Read (beta)

loading the full paper ...