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

Authors: Yifeng Ding, Hantian Ding, Shiqi Wang, Qing Sun, Varun Kumar, Zijian Wang
Primary area: foundation or frontier models, including LLMs
Keywords: Language Models, Fill-in-the-Middle, Planning, Training

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.

Decision: reject
Reviewer scores: 6, 3, 5

Reviews and discussion:
<thread>
creation date:
1730084704761

thread_type:
official_review

<thread_object>
creator: reviewer_ogj5

<content>
summary:
This paper discusses the problem of fill-in-the-middle capabilities of code generation. The next-token-prediction way tends to make the model fail in long-length code generation to combine with both the left and right code sides. The paper proposes the method Horizon-Length Prediction (HLP) to tackle this issue. During the training, the model is also required to predict the ratio of the left token length to be generated to the total token length to be generated. This method is supposed to enhance the ability of LLMs to take care of the code generation length.

The experiments on extensive settings prove the effectiveness. The authors also use probing method to show that HLP trained model tends to have better capability on code generation length prediction. I am not an expert in code LLMs. Hence, I am not sure whether the proposed question and method are novel. Based on the current experiments, I think the results are effective and the paper is clear. However, the paper lacks extensive analysis on the proposed method like why should set like M-T/M for length prediction not other settings. I am giving a borderline in this case and will surely look through other reviews and authors' comments for final evaluation.

strengths:
The issue and the method are clearly explained, the experiments are carried out extensively and achieved notably better performance. Meanwhile, the method will nearly not increase the burden of LLM training. The method is intuitively reasonable.

weaknesses:
1) The techniques used are relatively simple. Most contents in the paper are well-known. However, it makes sense.

2) The setting of target as M-t/M is not such solid. In this equation, the target y also depends on the total length M. Whether it is optimal setting should be discussed. For example, ablation studies on other settings.

3) The paper also lacks the thorough analysis like why enhancing the ability to predict code generation length is such effective, especially in fill-in-the-middle problem. Will it also work in uni-directional code generation?

typo: Line 132 are are -> are

questions:
As discussed in the above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730683080929

thread_type:
official_review

<thread_object>
creator: reviewer_s26y

<content>
summary:
The paper considers the task Fill-in-the-Middle (FIM) in code generation. Instead of completing the code, which is a common task in code generation, the authors consider the task of completing missing code given the left and right contexts.

They consider adding an auxiliary objective, predicting the number of missing lines, to the training of a language model. The model is trained to optimize next-token prediction as well as the number of missing lines. Empirically, this method improves performance on file-level and repository-level benchmarks.

strengths:
FIM is an important problem that is relatively underexplored in the literature. The method proposed in this paper is simple and straightforward. It shows significant improvements on diverse benchmarks.

The paper is well-written and easy to follow.

weaknesses:
**Contributions.**

Generalization: The proposed method appears to target at the FIM task. It limits its generalization to other code generation tasks.

Cost: The method requires fine-tuning a model specifically for FIM, which could be costly. Whenever a code model or a generalized model is released, it needs to be finetuned and maintained solely for this task.

**Evaluation metrics.** Table 4 uses Exact Match (EM) and Edit Similarity (ES) as evaluation metrics, which are not standard in code generation. This choice seems to be consistent with prior work. Is it possible to evaluate using pass@1 / pass@k? Or is it justifiable to measure EM and ES for codes?

questions:
My questions are in the weaknesses section above. The authors are welcome to correct any possible misunderstandings.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730683860290

thread_type:
official_review

<thread_object>
creator: reviewer_gjnh

<content>
summary:
The paper highlights the current limitations of Fill-in-Middle (FIM) evaluations of popular coding benchmarks due to the post-process that takes places, which often boost evaluation performance artificially. The main contribution of the paper is introducing a new auxiliary loss that predicts the "horizon" length (i.e. the portion of tokens left to predict) on top of the typical Prefix-Suffix-Middle next token prediction loss to perform FIM tasks. With this training objective, LLMs perform better at FIM tasks, notably, the repository-level cross-fill code tasks.

strengths:
- The paper soundly points out a glaring problem with current evaluation of FIM tasks that convincingly concludes that the post-processing only artificially boosts model scores while not providing additional insights into performance towards practical settings
- The paper provides a simple but novel idea that shows model improvements in the more rigorous evaluation (without any post-processing) of FIM tasks that would be of interest for the research community to apply to other domains
- The paper explores a breadth of different evaluation tasks for assessing their methods, including direct FIM tasks, repository-level cross-file FIM tasks, code fixing tasks, and reasoning tasks, demonstrating the effectiveness of their method
- That paper is well written and easy to understand/follow

weaknesses:
- Lack of rigorous confidence interval analysis - all the of experimental results lack statistical significance numbers, making it hard to judge if the performance improvements are due to noise or if they are statistically significant.
- Lack of theoretical/empirical evidence for why HLP works - it is not clear to me why this method works (assuming the experimental results are statistically significant). I believe the authors should add a section explaining (at least intuitively) why this method should work.
- Lack of additional baselines + ablations - adding some other strong baseline results would further validate this method. For instance, authors mention multi-token prediction in their related works. This methods performance should be reported as a strong baseline. For ablations, one idea could be explore the affect of increasing the complexity of the $hlp$_$head$ (e.g. using a MLP of increasing layers).

questions:
- Could you report the statistical significance for each of the evaluations?
- I found Section 5 particularly interesting and believe a natural question to ask is if "Horizon Awareness" under NTP increases with model parameters?
- Why is $HLP_{L2R} + HLP_{FIM}$ not presented as the standard approach?
- Is there a reason that during training the HLP objective is applied for each token instead of just the first token? My concern is that remaining tokens after the first token do not really provide any "additional" signal for the horizon length.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732171089077

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Responses

comment:
We thank all the reviewers for their insightful comments and suggestions for improving the paper! We have conducted additional ablations and analyses following reviewers’ suggestions, and we address the main weaknesses and questions in the response to individual reviewers below. We will add these ablations, analysis, and discussions to a later revision of the paper.

Please let us know if you have additional questions or need more clarification for our responses. We are happy to discuss further with reviewers during the reviewer-author discussion period.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732411162872

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Submission Revision

comment:
We thank the reviewers for the inspiring comments and we have updated our submission by making the following notable changes (highlighted in blue):

- **(Experiment) Adding statistical significance for each of the evaluations (Section 4 / Table 3,4,5)**: our analysis reveals strong statistical significance for the majority of our evaluations. (gJnh)

- **(Experiment) Adding multi-token prediction as a baseline (Section 5.1 / Table 6)**: multi-token prediction performs much worse than HLP due to limited planning horizon. (gJnh)

- **(Experiment) Conducting additional probing experiments on all the models (Section 5.2 / Table 7 / Figure 4)**: while “Horizon Awareness” under NTP increases with more model parameters, they are still very weak, showing that training with NTP only cannot bring “Horizon Awareness” to language models effectively, even with a larger model. (gJnh)

- **(Experiment) Studying the effect of the complexity of `hlp_head` (Appendix A.2 / Table 9)**: we find that the complexity of `hlp_head` does not have a major impact on performance. (gJnh)

- **(Experiment) Comparing against applying HLP to the first token only (Appendix A.2 / Table 10)**: we find that applying HLP loss for each token can achieve better performance than applying it to just the first token, as having HLP loss on every token provides denser and consistent supervision signals which makes learning easier. (gJnh)

- **(Experiment) Comparing against using `M-t` as the target in HLP (Appendix A.2 / Table 11)**: we find that setting the target as `M-t` fails to improve FIM performance of baselines, likely due to the large-scale HLP loss after using `M-t` as the target interferes with NTP pre-training. (ogj5)

- **(Discussion) Providing a preliminary theoretical explanation for HLP (Section 3).** (gJnh)

- **(Writing) Fixing the typo (Section 2).** (ogj5)

We are looking forward to discussing further with reviewers!
</content>
</thread_object>

</thread>
<thread>
creation date:
1733294619884

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Author-Reviewer Discussion

comment:
We thank all the reviewers for their thoughtful feedback and constructive suggestions. Below, we summarize the main discussion points and how they were addressed through the rebuttal process:

**Statistical Significance**
- **Reviewer’s Concern**: Reviewer `gJnh` questioned whether the performance improvements of HLP are statistically significant.
- **Our Response**: We performed additional statistical significance analysis and revealed strong significance (p<0.05 through paired t-test) across most evaluations. 
    
**Design Choices & Ablations**
   - **Reviewer’s Concern**: Reviewer `gJnh` questioned (1) whether the complexity of `hlp_head` affects the effectiveness of HLP, (2) whether applying HLP to just the first token can achieve similar performance as applying HLP for each token, and (3)  whether HLP can outperform multi-token prediction. Additionally, Reviewer `ogj5` questioned the necessity of using `M-t/M` rather than `M-t` as the training target.
   - **Our Response**: We performed ablation studies for all these asks. We show that:
        - the complexity of `hlp_head` has a minimum impact on HLP’s performance;
        - applying HLP for each token can achieve better performance than applying it to just the first token, as having HLP loss on every token provides denser and more consistent supervision signals which makes learning easier;
        - multi-token prediction performs much worse than HLP. The reason behind this is that multi-token prediction only enhances models’ planning capability for a short and limited horizon, which is usually insufficient for FIM as the connection from middle to suffix happens over a long horizon;
        - setting the target as `M-t` fails to improve the FIM performance of baselines, likely due to the large-scale HLP loss after using `M-t` as the target interferes with NTP pre-training. 

        These results successfully justify the original design choices of HLP.

**Theoretical Understanding/Why HLP works**
  - **Reviewer’s Concern**: Reviewers `gJnh`, `s26Y`, and `ogj5` asked about the theoretical/intuitive explanation for why HLP works. Reviewer `gJnh` further questioned the connection between HLP and “planning”.
  - **Our Response**: We provided a preliminary theoretical explanation showing how HLP addresses myopic generation in standard next-token prediction. We demonstrated the concrete connection to planning through a novel decoding algorithm that enables verbosity control using HLP predictions (thank you Reviewer `gJnh` for the interesting idea!).

**Scaling of Horizon Awareness**
  - **Reviewer’s Concern**: Reviewer `gJnh` was curious about whether "Horizon Awareness" under NTP increases with model parameters.
  - **Our Response**: We performed additional experiments to reveal that, while "Horizon Awareness" increases with model size under NTP, it remains consistently weak without HLP, demonstrating the fundamental limitation of standard training.

**Generalizability**
   - **Reviewer’s Concern**: Reviewers `s26Y` and `ogj5` expressed concerns about limited generalization beyond FIM tasks and integration with uni-directional generation in code LLMs.
   - **Our Response**: We highlighted that:
        - FIM plays a fundamental role in modern code development and is a standard pre-training task in all code LLMs.
        - HLP is designed to be broadly applicable for pre-training general-purpose code LLMs that can do both left-to-right generation and fill-in-the-middle and demonstrates the monotonic improvement on FIM without regression on left-to-right capability, as supported by empirical evidence in Appendix A.1.

**Evaluation Metric**
- **Reviewer’s Concern**: Reviewer `s26Y` questioned whether Exact Match (EM) and Edit Similarity (ES) should be used as evaluation metrics for repository-level FIM tasks.
- **Our Response**: We note that it is the standard practice in the literature to use EM/ES for repository-level FIM tasks from multiple previous works in top conferences like NeurIPS/ICLR. We further highlighted that our evaluation uses a complementary set of metrics appropriate for different scenarios following past literature and standard practices, e.g., execution-based metrics for function-level tasks, and test suite pass rates for code fixing, which well captures the performance of the resulting models.

Finally, we sincerely thank all reviewers again for their dedication and thoughtful evaluations, and we are also deeply appreciative of the Area Chair’s guidance and support.

Thanks,

Authors of HLP
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces horizon-length prediction, a novel training objective designed to improve fill-in-the-middle tasks by enabling models to predict the number of remaining middle tokens, thus facilitating effective planning based on distant right context. Despite of the contributions, the paper still needs to be improved on its lack of statistical significance analysis, insufficient justification for the method's effectiveness, limited generalization, high fine-tuning costs, reliance on simple techniques, questionable target setting, and inadequate analysis of code generation length prediction.