通过预测剩余代码长度,让模型更准地补全代码。
Planning-Aware Code Infilling via Horizon-Length Prediction
- 让模型预测当前距离右端还有多少代码要生成,实现前瞻规划。
- 在多个基准上提升代码补全效果最高达24%,且不增加推理成本。
- 适合需要精准补全、对上下文理解要求高的编程场景。
代码补全中的填空(FIM)已成为代码语言模型的核心能力,可在左右上下文给定情况下生成缺失代码。然而,现有基于重排序序列的下一个词预测(NTP)训练方式,使模型难以充分考虑远距离右端上下文,导致生成内容与上下文不一致。本文提出一种新的训练目标——视野长度预测(HLP),让模型在每一步预测剩余中间 token 的数量,从而实现前瞻规划。该方法使模型能自适应地学习任意左右上下文下的补全边界,无需依赖数据集特定后处理。实验表明,不同模型家族和规模下,HLP在文件级和仓库级多个基准上相对性能提升最高达24%。此外,其增强的规划能力也提升了代码推理表现。关键优势在于训练开销极低,推理无额外负担,适用于实际应用。
原文摘要 · Abstract (English)
Fill-in-the-Middle (FIM), or infilling, 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 performs next-token prediction (NTP) over reordered sequence often leads to models struggling to generate content that aligns well with the surrounding context. 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 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 model families and sizes shows that HLP significantly improves FIM performance by up to 24% relatively on diverse benchmarks, across file-level and repository-level. Furthermore, the enhanced planning capability gained through HLP boosts model performance on code reasoning. Importantly, HLP incurs negligible training overhead and no additional inference cost, ensuring its practicality for real-world scenarios.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。