Adaptive Critical Token-Aware Retrieval for Repository-Level Code Generation
2026-09-01 • Software Engineering
Software EngineeringArtificial IntelligenceComputation and Language
AI summaryⓘ
The authors look at how to improve code generation when working with very large code repositories, which are too big for current language models to handle all at once. They find that mistakes often happen at a few key spots in the generated code, which they call critical tokens. To fix this, the authors created a system called ACToR that focuses on these critical tokens by retrieving relevant repository info only when needed, making the code generation more accurate. Their tests show that this method works better than previous approaches on popular benchmarks. They also analyze how important these critical tokens are in causing code generation errors.
repository-level code generationlarge language models (LLMs)retrieval-augmented generation (RAG)critical tokensautoregressive generationdense retrieversposition-aware weightingRepoExec benchmarkCoderEval benchmark
Authors
Kefeng Duan, Dewu Zheng, Yanlin Wang, Terry Yue Zhuo, Mingwei Liu, Jianxing Yu, Jiachi Chen, Ensheng Shi, Xilin Liu, Yuchi Ma, Zibin Zheng
Abstract
The repository-level code generation task requires synthesizing code that satisfies task requirements while remaining consistent with the target repository context. Since real-world repositories often exceed the input length limits of LLMs, existing approaches commonly adopt retrieval-augmented generation (RAG) to provide repository-specific context. Despite improving repository-context retrieval, existing methods typically provide context as task-level support, without explicitly identifying the critical tokens that require fine-grained repository context during generation. During the autoregressive generation process of LLMs, errors often concentrate at a small number of decisive positions: once such tokens are generated incorrectly, subsequent code may follow an incorrect semantic path and eventually lead to functional failure. We refer to these positions as "critical tokens". In this paper, we propose ACToR, an adaptive critical token-aware retrieval framework for repository-level code generation. ACToR identifies critical tokens during generation and triggers targeted retrieval on demand to provide repository context at these decisive positions. In addition, we design a position-aware weighting method for dense retrievers to prioritize context that is more informative for generation. We evaluate ACToR on two representative repository-level benchmarks, RepoExec and CoderEval. Experimental results show that ACToR consistently outperforms state-of-the-art methods, achieving relative improvements of 8.4% on RepoExec and 15.4% on CoderEval. Beyond performance gains, we systematically quantify the impact of critical tokens, revealing their central role in major generation failures and highlighting the necessity of targeted retrieval strategies. We provide the code and data at https://github.com/DeepSoftwareAnalytics/ACToR.