Adaptive Critical Token-Aware Retrieval for Repository-Level Code Generation
ACToR identifies critical tokens during code generation and triggers on-demand retrieval, boosting accuracy by 8.4% on RepoExec and 15.4% on CoderEval.
Key Findings
Methodology
ACToR introduces a dynamic framework that detects critical tokens during autoregressive code generation using a trained discriminator model based on hidden states. When a token is identified as critical—via mismatch, entropy, or attention influence—the system triggers targeted repository context retrieval. This process employs a position-aware weighted dense retriever utilizing a double-Gaussian pooling scheme to emphasize important sequence regions. The approach combines offline training of the discriminator with online real-time detection and retrieval, enabling precise context supplementation at key generation steps. Experimental results on RepoExec and CoderEval benchmarks demonstrate consistent performance gains, with improvements of 8.4% and 15.4%, respectively.
Key Results
- On RepoExec, ACToR outperforms previous state-of-the-art methods, achieving an 8.4% relative increase in BLEU scores, with notable improvements in functional correctness and dependency handling.
- On CoderEval, the method achieves a 15.4% relative boost, especially excelling in complex dependency scenarios and syntactic diversity, confirming the effectiveness of critical token targeting.
- Analysis shows critical tokens constitute only 5-11% of generated positions but account for a majority of errors and uncertainty, validating the targeted retrieval strategy.
Significance
This work addresses fundamental limitations in large language model-based code generation, particularly the inability to dynamically focus on crucial parts of large codebases. By identifying and retrieving context only at pivotal positions, ACToR significantly enhances code correctness, robustness, and project-specific adaptation. It bridges the gap between static context provision and the need for fine-grained, real-time context management, paving the way for scalable, industrial-grade automated software development. The approach also offers theoretical insights into error propagation and decision-making in autoregressive models, with broad implications for AI-assisted programming.
Technical Contribution
The paper introduces a novel critical token discriminator trained on labeled data, integrating entropy, mismatch, and attention metrics. It proposes a position-aware weighted pooling mechanism for dense retrievers, emphasizing sequence ends and signature regions. The end-to-end system dynamically detects critical tokens during generation, triggering targeted retrieval, and re-decoding tokens with enriched context. This architecture surpasses static retrieval methods both in performance and efficiency, providing a new paradigm for context-aware code synthesis.
Novelty
This is the first systematic quantification of critical tokens in repository-level code generation, linking token-level errors to functional failures. It innovates by combining real-time critical token detection with targeted retrieval, moving beyond prior static or task-level context support. The position-aware weighted pooling further distinguishes this work, enabling more relevant context embedding. These innovations collectively represent a significant leap in fine-grained, adaptive code generation techniques.
Limitations
- The approach may struggle with extremely long sequences or highly sparse dependencies, where critical tokens are difficult to detect accurately.
- Retrieval frequency and thresholds require careful tuning; excessive retrieval could impact efficiency, while too sparse may miss critical tokens.
- Generalization across diverse programming languages and project styles remains to be validated, especially in unseen or domain-specific contexts.
Future Work
Future directions include integrating multi-modal signals like code comments and execution traces to improve critical token detection, applying reinforcement learning to optimize retrieval policies, and extending the framework to cross-project and cross-language scenarios. Additionally, efforts will focus on reducing computational overhead and exploring unsupervised or semi-supervised training to enhance generalization.
AI Executive Summary
The rapid advancement of large language models (LLMs) has revolutionized code generation, yet challenges remain in handling large-scale repositories with complex dependencies. Existing methods often rely on static retrieval of repository context, which is insufficient for ensuring correctness across diverse code segments. Recognizing that errors tend to concentrate at a small number of pivotal positions—termed critical tokens—this research introduces ACToR (Adaptive Critical Token-aware Retrieval), a novel framework that dynamically identifies these tokens during generation and triggers targeted context retrieval.
ACToR's core innovation lies in training a discriminator model that assesses token importance based on hidden states, entropy, mismatch, and attention influence. During inference, the system continuously evaluates each generated token, and upon detecting a critical token, it performs a focused retrieval of relevant repository snippets. This targeted approach ensures that the model always has access to the most pertinent information at crucial decision points, significantly reducing the likelihood of cascading errors.
To enhance retrieval relevance, the authors propose a position-aware weighted dense retriever utilizing a double-Gaussian pooling scheme, emphasizing sequence ends and signature regions. Extensive experiments on RepoExec and CoderEval benchmarks demonstrate that ACToR outperforms existing state-of-the-art methods, achieving 8.4% and 15.4% relative improvements respectively. Analysis reveals that critical tokens, though sparse, are responsible for most errors, underscoring the importance of targeted retrieval.
This work offers a significant step forward in scalable, accurate repository-level code generation, addressing key limitations of static context strategies. It opens avenues for more robust, fine-grained AI-assisted programming, with potential impacts spanning industrial automation, software maintenance, and intelligent development tools. Future research will explore multi-modal signals, reinforcement learning, and cross-language generalization to further enhance the framework's effectiveness and efficiency.
Deep Analysis
Background
近年来，随着大规模预训练模型（如GPT-3、Codex）的出现，代码生成技术取得了巨大突破。代表性工作包括OpenAI的Codex、DeepMind的AlphaCode以及基于检索增强的CodeT5。这些模型在短文本和单一任务中表现优异，但在仓库级代码生成中仍面临巨大挑战。仓库规模庞大，依赖关系复杂，输入长度受限，导致模型难以获取完整上下文。检索增强技术（如RepoCoder、DRACO）被引入以缓解这一问题，但多为静态或任务级支持，未能动态识别关键位置的上下文需求。随着软件系统复杂度提升，如何在保证生成质量的同时，有效利用仓库中的丰富信息，成为研究重点。
Core Problem
仓库级代码生成的核心难题在于输入长度限制和不同位置对上下文的不同依赖。静态提供上下文容易忽略关键Token，导致错误集中在少数位置，影响整体功能正确性。错误一旦发生，可能引发级联失误，导致程序逻辑崩溃。现有方法缺乏实时识别关键Token的机制，难以在生成过程中动态补充最需要的上下文信息。这限制了模型在复杂场景中的表现，亟需一种能够在生成时刻识别并补充关键上下文的技术方案。
Innovation
本研究提出ACToR框架，创新点包括：1）定义并量化仓库级代码中的关键Token，揭示其对最终代码正确性的决定性作用；2）设计基于生成动态识别关键Token的目标导向检索机制，结合位置感知的高斯加权池化，提升检索相关性；3）训练轻量级判别模型，实时判断Token关键性，触发定向检索，增强上下文信息。该方案突破了以往静态、任务级支持的局限，实现了端到端的动态、粒度更细的上下文补充，显著改善了生成的准确性和鲁棒性。
Methodology
- �� 离线阶段：
- 构建训练数据：筛选高质量仓库，采样完整函数，构建检索提示。
- 训练判别模型：利用多层感知器（MLP）对每个Token的关键性进行分类，输入为模型隐藏状态。
- 标注关键Token：结合预测误差、不确定性（熵）和后续注意力指标，标记出对生成影响最大的Token。
- 采样平衡：利用信息熵筛选“难负样本”，确保正负样本平衡。
- �� 在线推理阶段：
- 目标检索：利用位置感知的高斯加权池化机制，对上下文进行加权池化，提升相关性。
- 关键Token识别：根据判别模型输出，判断当前Token是否关键。
- 定向检索：若关键，触发目标检索，补充上下文。
- 逐步生成：在每个Token生成时动态判断，确保关键位置获得充分信息。
Experiments
采用RepoExec和CoderEval两个仓库基准，比较ACToR与SOTA方法的性能。指标包括BLEU、准确率和功能完整性。超参数如检索阈值、加权参数通过交叉验证调优。设置对比模型包括静态检索和无检索基线。通过消融实验验证关键Token识别和位置感知机制的贡献。结果显示，ACToR在两个基准上均优于对比方法，性能提升分别为8.4%和15.4%。分析关键Token的分布和错误影响，验证了目标检索的有效性。
Results
实验显示，ACToR在仓库级代码生成中实现了显著性能提升，尤其在复杂依赖和多样语法结构中表现优越。关键Token占生成位置的5-11%，但其错误率和不确定性极高，验证了目标检索策略的重要性。消融分析表明，关键Token识别和位置感知机制共同贡献了性能提升，模型在超长序列和稀疏依赖场景中表现出更强鲁棒性。
Applications
该技术适用于工业界自动化代码补全、依赖管理和持续集成等场景。只需提供部分函数签名和描述，即可实现高质量代码生成，减少人工干预。未来结合持续学习和多模态信息，有望实现跨项目、跨语言的智能代码生成，推动软件开发自动化的广泛应用。
Limitations & Outlook
模型在极端复杂或超长依赖场景下识别效果有限，检索频次和参数调优影响效率。对不同项目风格和编码规范的适应性有限，泛化能力有待验证。此外，实时检索带来的计算成本和延迟也是实际部署的挑战。未来需优化模型结构和检索策略，提升效率和适应性。
Plain Language Accessible to non-experts
想象你在厨房做菜，菜谱很长，有很多步骤。你不可能每次都看完整个菜谱，只在关键步骤时翻看相关的提示，比如“加盐”或“炒至金黄”。如果你漏掉了这些关键步骤，菜可能会变得难吃或失败。这个研究就像厨师在做菜时，能自动识别出哪些步骤最重要，然后只在这些关键点查阅提示，确保每一步都正确。这样，即使菜谱很长，也能做出好菜。它让电脑在写代码时，知道哪些地方最容易出错，然后重点“查资料”，避免出错，做出更好的程序。
ELI14 Explained like you're 14
想象你在学校写一篇长作文，里面有很多段落。有时候，写错一个关键的词，比如“因为”写成“虽然”，就会改变整篇文章的意思。这个研究就像老师告诉你：在写作过程中，要特别注意那些关键的词，一旦写错了，就会导致整篇文章出错。老师还会在你写作时，随时提醒你注意这些重要的词，帮你检查和补充信息。这样，你就能写出一篇没有大错的好作文。类似地，这个方法让电脑在自动写代码时，能找到那些最重要的词或符号，确保代码的逻辑和功能都正确。它就像一个聪明的老师，随时提醒你最关键的地方，帮你写出完美的程序。
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.