수평 길이 예측: 전방 계획을 활용한 코드 생성을 위한 가운데 채우기 능력 발전

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

October 4, 2024

저자: Yifeng Ding, Hantian Ding, Shiqi Wang, Qing Sun, Varun Kumar, Zijian Wang

cs.AI

초록

Fill-in-the-Middle (FIM)은 코드 언어 모델에 필수적이 되어, 왼쪽과 오른쪽 컨텍스트를 모두 고려하여 누락된 코드를 생성할 수 있게 합니다. 그러나 현재의 FIM 훈련 패러다임은 원본 훈련 시퀀스를 재배열한 후 정규 다음 토큰 예측(NTP)을 수행하는데, 종종 모델이 주변 컨텍스트와 원활하게 일치하는 콘텐츠를 생성하는 데 어려움을 겪게 합니다. 기존 연구는 이러한 약점을 우회하기 위해 규칙 기반의 후처리에 의존하고 있지만, 이러한 방법은 제한적이고 데이터셋 특정 가정에 의존하기 때문에 오픈 도메인 코드 완성 작업에서 실용적으로 사용하기 어렵습니다 (예: 그라운드 트루스와 동일한 라인 수를 생성). 게다가, 이러한 비현실적인 가정 없이 FIM 작업에서 모델 성능이 상당히 저하됩니다. 우리는 NTP만으로는 모델이 먼 오른쪽 컨텍스트에 의존한 효과적인 계획 학습을 할 수 없다고 가설을 세웁니다. 이것은 성공적인 코드 채워넣기에 중요한 요소입니다. 이를 극복하기 위해 우리는 Horizon-Length Prediction (HLP)을 제안합니다. 이것은 모델이 각 단계에서 남은 중간 토큰의 수(즉, 수평 길이)를 예측하도록 가르치는 새로운 훈련 목표입니다. HLP는 미래를 내다보는 계획을 통해 FIM을 발전시켜, 데이터셋 특정 후처리에 의존하지 않고 임의의 왼쪽과 오른쪽 컨텍스트에 대한 채워넣기 경계를 내재적으로 학습할 수 있게 합니다. 다양한 벤치마크에서 다양한 모델과 크기에 걸쳐 우리의 평가 결과는 HLP가 FIM 성능을 상당히 향상시키며, 파일 수준 및 저장소 수준에서 24%까지 상대적으로 향상시킵니다. 또한 비현실적인 후처리 방법에 의존하지 않고 모델 성능을 향상시키는 HLP를 통해 향상된 계획 능력은 코드 추론에 모델 성능을 향상시킵니다. 중요한 점은, HLP는 실제 시나리오에서 실용적이며 추가적인 훈련 부담이 거의 없으며 추가적인 추론 비용도 없어 실용성을 보장합니다.

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.