# Adaptive Critical Token-Aware Retrieval for Repository-Level Code Generation

arXiv: https://arxiv.org/abs/2609.01601 · HTML: https://mlsanity.com/paper/2609.01601/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 54 of its 67 references that are in the map.
- First posted on arXiv 1 Sep 2026.

## Builds on (most cited first, 30 of 54)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Code Llama: Open Foundation Models for Code](https://mlsanity.com/paper/2308.12950/) (2023)
- [Competition-Level Code Generation with AlphaCode](https://mlsanity.com/paper/2203.07814/) (2022)
- [Is Your Code Generated by ChatGPT Really Correct? Rigorous Evaluation of Large Language Models for Code Generation](https://mlsanity.com/paper/2305.01210/) (2023)
- [DeepSeek-Coder: When the Large Language Model Meets Programming -- The Rise of Code Intelligence](https://mlsanity.com/paper/2401.14196/) (2024)
- [CodeGen: An Open Large Language Model for Code with Multi-Turn Program Synthesis](https://mlsanity.com/paper/2203.13474/) (2022)
- [StarCoder: may the source be with you!](https://mlsanity.com/paper/2305.06161/) (2023)
- [BigCodeBench: Benchmarking Code Generation with Diverse Function Calls and Complex Instructions](https://mlsanity.com/paper/2406.15877/) (2024)
- [StarCoder 2 and The Stack v2: The Next Generation](https://mlsanity.com/paper/2402.19173/) (2024)
- [InCoder: A Generative Model for Code Infilling and Synthesis](https://mlsanity.com/paper/2204.05999/) (2022)
- [Large Language Models for Software Engineering: A Systematic Literature Review](https://mlsanity.com/paper/2308.10620/) (2023)
- [CodeT: Code Generation with Generated Tests](https://mlsanity.com/paper/2207.10397/) (2022)
- [RepoCoder: Repository-Level Code Completion Through Iterative Retrieval and Generation](https://mlsanity.com/paper/2303.12570/) (2023)
- [UniXcoder: Unified Cross-Modal Pre-training for Code Representation](https://mlsanity.com/paper/2203.03850/) (2022)
- [IntelliCode Compose: Code Generation Using Transformer](https://mlsanity.com/paper/2005.08025/) (2020)
- [CacheBlend: Fast Large Language Model Serving for RAG with Cached Knowledge Fusion](https://mlsanity.com/paper/2405.16444/) (2024)
- [CacheGen: KV Cache Compression and Streaming for Fast Large Language Model Serving](https://mlsanity.com/paper/2310.07240/) (2023)
- [Compressing Context to Enhance Inference Efficiency of Large Language Models](https://mlsanity.com/paper/2310.06201/) (2023)
- [Large Language Models for Software Engineering: Survey and Open Problems](https://mlsanity.com/paper/2310.03533/) (2023)
- [MAGIS: LLM-Based Multi-Agent Framework for GitHub Issue Resolution](https://mlsanity.com/paper/2403.17927/) (2024)
- [Repository-Level Prompt Generation for Large Language Models of Code](https://mlsanity.com/paper/2206.12839/) (2022)
- [LLM Hallucinations in Practical Code Generation: Phenomena, Mechanism, and Mitigation](https://mlsanity.com/paper/2409.20550/) (2024)
- [Towards an Understanding of Large Language Models in Software Engineering Tasks](https://mlsanity.com/paper/2308.11396/) (2023)
- [A Survey of Large Language Models for Code: Evolution, Benchmarking, and Future Trends](https://mlsanity.com/paper/2311.10372/) (2023)
- [CoCoMIC: Code Completion By Jointly Modeling In-file and Cross-file Context](https://mlsanity.com/paper/2212.10007/) (2022)
- [Agents in Software Engineering: Survey, Landscape, and Vision](https://mlsanity.com/paper/2409.09030/) (2024)
- [Repoformer: Selective Retrieval for Repository-Level Code Completion](https://mlsanity.com/paper/2403.10059/) (2024)
- [No Need to Lift a Finger Anymore? Assessing the Quality of Code Generation by ChatGPT](https://mlsanity.com/paper/2308.04838/) (2023)
- [RLCoder: Reinforcement Learning for Repository-Level Code Completion](https://mlsanity.com/paper/2407.19487/) (2024)

## Abstract (arXiv)

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.