PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Kefeng Duan, Dewu Zheng, Yanlin Wang, Terry Yue Zhuo, Mingwei Liu, Jianxing Yu, Jiachi Chen, Ensheng Shi, Xilin Liu, Yuchi Ma, Zibin Zheng

arXiv:2609.01601v1cs.SEcs.AIcs.CL

TL;DR

Repository-level code generation needs repository-specific context, but existing methods do not explicitly model which generation positions require different context. ACToR identifies critical tokens during generation, triggers targeted retrieval, and uses position-aware weighting; it reports relative gains of 8.4% on RepoExec and 15.4% on CoderEval.

  • Problem

    Existing repository-level retrieval methods provide task-level context without explicitly modeling critical token-level decisions during path-dependent generation.

  • Method

    ACToR identifies critical tokens during generation, triggers targeted on-demand retrieval, and uses position-aware weighting for dense retrievers.

  • Results

    8.4% on RepoExec and 15.4% on CoderEval are ACToR’s reported relative gains over strong repository-level baselines.

  • Takeaways & Limitations

    Critical tokens occupy about 5–11% of generated positions while concentrating much of the model’s error and uncertainty, supporting sparse generation-time retrieval updates.

  • Takeaways & Limitations

    Evaluation focuses on RepoExec and CoderEval, and critical-token labels rely on observable proxy signals because downstream impact is not directly observable at generation time.

Abstract

from arXiv · show

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.

I. INTRODUCTION

Repository-level code generation must satisfy task requirements while remaining consistent with project-specific repository context, but existing retrieval methods lack token-level adaptation. ACToR addresses this gap by identifying critical tokens during generation and retrieving targeted context on demand, improving benchmark performance.

  • Repository-level generation must preserve consistency with project-specific APIs, dependencies, data structures, coding conventions, and broader repository context.
  • Existing retrieval methods mainly provide holistic task-level context, although different autoregressive generation positions may require different contextual information.
  • Critical tokens are decisive generation positions where incorrect API names, variable references, indices, or control-flow tokens can redirect subsequent code and cause implementation failure.
  • ACToR performs dynamic, fine-grained retrieval during generation when critical tokens require additional context, alongside position-aware weighting for denser retrieval.
  • 8.4% on RepoExec and 15.4% on CoderEval are ACToR’s reported relative gains over strong repository-level baselines.
  • Critical tokens are syntactically diverse, spanning comments, API calls, variable references, and index expressions, so fixed syntactic rules or one pre-generation retrieval step are insufficient.

III. METHODOLOGY

ACToR’s methodology constructs training data and uses a two-phase pipeline to identify critical tokens and provide targeted retrieval during inference. Its pipeline combines offline discriminator training with position-aware retrieval and critical-token-guided inference.

  • ACToR divides its methodology into offline training and online inference phases.
  • The offline phase trains a lightweight discriminative ensemble model to identify critical tokens.
  • The online phase uses a Position-Aware Weighted Retriever and Critical Token-Guided Inference to enhance code generation.
  • Training-data construction comprises Repository Filtering, Function Sampling, and Prompt Construction.
  • Prompts concatenate retrieved contexts in score order with the initial code sequence, using separators and file-path annotations to preserve context boundaries.

2) Token Evaluation:

ACToR evaluates token criticality using prediction mismatch, uncertainty, and subsequent attention influence. Teacher forcing reduces contamination from the model’s own earlier errors when judging each position.

  • Token criteria: Critical tokens are positions where local prediction errors can substantially affect subsequent tokens and increase functional-failure risk.The classification target captures both misprediction tendency and downstream semantic impact.
  • Token criteria: Token mismatch marks positions where the model’s top-1 prediction differs from the ground-truth token.A mismatch indicates prediction difficulty and becomes more critical when accompanied by downstream impact.
  • Token criteria: Entropy quantifies uncertainty at each generation position, with larger entropy indicating greater uncertainty and higher critical-token likelihood.The entropy is computed from the model’s token-probability distribution.
  • Token criteria: Subsequent attention influence estimates how strongly later positions attend to the current token and therefore how errors may propagate.The method uses Top-5 subsequent attention with K = 5 rather than averaging all later positions.
  • Evaluation procedure: Teacher forcing replaces the model’s prediction at position i with the ground-truth token before evaluating later positions.This suppresses downstream error propagation and makes each token’s criticality more reflective of its individual value.

3) Label Balance:

The critical-token dataset is highly imbalanced because non-critical tokens dominate. ACToR balances training data by retaining informative hard negatives and uses lightweight MLP classifiers over final-layer hidden states.

  • Label Balance: Non-critical tokens constitute the majority of the dataset, producing pronounced class imbalance.This imbalance motivates filtering low-value negative samples before classifier training.
  • Label Balance: The method ranks non-critical tokens by self-information and preferentially retains higher-information samples until positive and negative classes are balanced.These retained negatives are treated as hard negatives.
  • Classifier Training: Each critical-token judger is a 3-layer MLP that takes a base Code-LLM hidden state for a token and predicts its binary class probability.The classifier is trained using hidden states from the Code-LLM’s final decoder layer.
  • Classifier Training: The classifiers contain 4.01–10.01 million parameters, substantially fewer than the evaluated 1B and 13B backbone LLMs.This reports the classifier scale relative to the backbone models.

C. Online Inference

ACToR combines generation-relevant weighted pooling with online retrieval. The dense-retriever weighting emphasizes span endpoints while adding negligible per-query runtime cost.

  • Weighted Dense Retriever: The weighted pooling scheme uses a double-endpoint Gaussian shape because average pooling ignores relative token importance.It emphasizes positions near both the sequence head and tail.
  • Weighted Dense Retriever: The unnormalized weight for position i is formed from two Gaussian kernels centered at the span’s endpoints.For a span of length L, token indices run from 0 through L − 1.
  • Weighted Dense Retriever: Smaller σ concentrates weighting at the endpoints, whereas larger σ makes the profile smoother or nearly flat.σ controls the sharpness of the peaks and their overlap.
  • Weighted Dense Retriever: Softmax normalization converts the position weights into a proper pooling distribution, which is then used for the weighted sum of token embeddings.The pooled vector is intended to improve retrieval of generation-informative context.
  • Weighted Dense Retriever: The weights depend only on σ and span length, so they can be precomputed and cached with negligible online overhead.This is the method’s stated low-cost benefit.

2) Critical Token Guidance Inference:

Critical-token guidance monitors generation online and retrieves additional context when a token is judged critical. The updated context is retained for subsequent steps until another critical token triggers correction.

  • Critical Token Guidance Inference: ACToR performs real-time monitoring of generated tokens and applies targeted retrieval when it identifies a critical token.This is the core Critical Token-Guided Dynamic Correction strategy.
  • Critical Token Guidance Inference: Initial retrieval supplies context before token-by-token generation, while each hidden state is evaluated by a pretrained classifier ensemble.The ensemble determines whether the current token is normal or critical.
  • Critical Token Guidance Inference: For a unanimously classified critical token, ACToR queries with the prompt, generated code, and token, retrieves updated context, and redecodes the current step.The redecoding produces a corrected token under the more targeted context.
  • Critical Token Guidance Inference: After correction, the updated context remains active for subsequent generation steps until a new critical token appears.This preserves the retrieved guidance rather than discarding it immediately after one correction.

IV. EXPERIMENTAL SETUP

The evaluation uses RepoExec and CoderEval to test repository-level code generation, comparing ACTOR with non-retrieval, vanilla-RAG, and state-of-the-art repository-level baselines.

  • Benchmarks: RepoExec contains 355 Python tasks requiring cross-file dependency integration and functional correctness.Each task includes comprehensive tests constructed through unit-test generation and coverage enhancement.
  • Benchmarks: CoderEval evaluates code generation under varying contextual dependencies using standardized Docker-based execution.The Python version contains 230 tasks from real-world open-source projects.
  • Baselines: ACTOR is compared with RawPrompt, RawRAG, RepoCoder, and RLCoder.These baselines cover non-retrieval prompting, one-shot retrieval, iterative retrieval, and reinforcement-learning-based retrieval optimization.

C. Evaluation Metrics

The study evaluates correctness with Pass@k across multiple model families and scales, using unified retrieval, generation, and critical-token detection settings. ACTOR consistently improves performance across the evaluated benchmarks and models while limiting retrieval overhead through sparse correction.

  • Evaluation Metrics: Pass@k for k = 1, 3, 5 measures whether at least one generated sample passes all associated unit tests.A task is correct only when the generated code passes every unit test.
  • Evaluation Configuration: The experiments use UniXcoder as the dense retriever and DeepSeekCoder and CodeLlama models as generators.The evaluated generators range from 1.3B to 13B parameters.
  • Efficiency: ACTOR adds 3.3 ms to per-token latency but reduces end-to-end time to 2.66 s per sample versus RepoCoder’s 4.50 s.Correction is invoked only at sparsely identified critical-token positions.

B. RQ2: Ablation Study

The ablation studies show that ACTOR’s token-labeling signals, weighted retrieval, and dynamic inference each contribute to performance. Dynamic inference has the clearest quality impact, while weighted retrieval provides complementary gains.

  • Token-Level Labeling: Removing any single token-labeling criterion lowers Pass@k on both RepoExec and CoderEval.Mismatch, uncertainty, and subsequent-attention signals provide non-overlapping supervision, with relative importance varying by benchmark.
  • Runtime Components: Disabling dynamic critical-token inference produces the clearest quality loss, especially on CoderEval.The strongest relative drops are roughly in the mid-teens across Pass@k.
  • Runtime Components: Both dynamic inference and weighted retrieval are required because neither runtime stage fully substitutes for the other.The ablation therefore rejects retrieval-only and inference-only versions of ACTOR.
  • Runtime Components: Replacing position-aware weighted retrieval with mean-pooled embeddings causes smaller but consistent regressions.Weighted retrieval improves the context supplied before and between dynamic corrections.

C. RQ3: Sensitivity Analysis

ACTOR remains stable across tested threshold configurations, while critical tokens occupy only a small fraction of generated positions. Their composition and prevalence also vary across model families and parameter scales.

  • Hyperparameter Sensitivity: ACTOR is insensitive to the specific tested uncertainty and subsequent-attention thresholds.The authors adopt the selected configuration as the standard setting for subsequent experiments.
  • Hyperparameter Sensitivity: Across 25 configurations, Pass@1, Pass@3, and Pass@5 variances remain approximately 2.83, 3.42, and 2.65, respectively.The corresponding standard deviations are approximately 1.68, 1.85, and 1.63.
  • Critical Token Composition: Critical tokens constitute 5.08% to 10.62% of positions across models, with mismatch and high-uncertainty tokens consistently overlapping.This distribution indicates that most generated positions are non-critical.
  • Model-Scale Patterns: Increasing DSCoder scale from 1B to 7B is associated with lower mismatch rates, from 3.29% to 2.43%, and uncertainty rates, from 5.26% to 2.56%.CodeLlama shows broadly stable rates from 7B to 13B.

2) Critical Token Syntactic Analysis:

Critical tokens are syntactically diverse but concentrate in particular token types and error-related conditions. Their analysis supports targeted retrieval strategies, illustrated by ACTOR’s dynamic recovery of repository-specific information.

  • Critical Token Syntactic Analysis: Critical-token composition is diverse but non-random: mismatch and uncertainty concentrate in keywords, whereas attention emphasizes structural operators and delimiters.Figures 7 and 8 provide absolute and relative syntax distributions after filtering dedent and indent types.
  • Critical Token Syntactic Analysis: Identifiers and keywords predominate among critical tokens in error-related conditions such as “Mismatch Only” and “Mismatch + Uncertainty”.The “Attention Only” condition has a more balanced syntax composition.
  • Critical Token Syntactic Analysis: Keywords are over-represented by 7.64x and 6.05x the baseline in “Mismatch + Uncertainty” and “Mismatch Only”, respectively.This makes keywords a frequent source of combined error and uncertainty despite identifiers being more numerous.
  • Critical Token Syntactic Analysis: Operators are over-represented by 1.86x the baseline in “Attention Only”, indicating disproportionate model focus on structural elements.The summary analysis also characterizes attention as centering on structural operators and delimiters.
  • Case Study: In a handler-registration decorator case, ACTOR dynamically retrieves the repository-defined attribute self.registry, unlike methods producing self.listeners, self.hooks, or redundant logic.The example shows how critical-token identification enables concise code aligned with the ground truth.

VII. THREATS TO VALIDITY

Repository-level generation must preserve project-specific semantics despite context limits, motivating retrieval methods and token-specific retrieval for decisive generation points. The paper reports broad benchmark gains but identifies proxy-label, computational, model-selection, and benchmark-scope constraints.

  • Internal Threats: Critical-token labels approximate downstream impact using mismatch, uncertainty, and attention proxies, which may miss some functionally decisive tokens.The authors mitigate this threat by combining signals, using teacher forcing, and validating contributions through ablations.
  • External Threats: Iterative context updates and KV-cache recalculation may challenge deployment in large-scale, real-time systems, while evaluation is limited to RepoExec and CoderEval.The authors report controllable time costs and plan broader benchmarks and additional KV-cache optimization.
  • Background and Motivation: Repository-level code generation requires implementations consistent with repository contents, environments, requirements, APIs, dependencies, and conventions.Repository context is essential, but entire repositories often exceed LLM input limits.
  • Related Work: Existing retrieval methods provide contextual support but generally treat context as separate from token generation, overlooking that different tokens may require different contexts.The paper addresses this limitation with retrieval triggered by critical tokens.
  • Error Analysis: The paper analyzes token-level precursors to generation failures through mismatch, uncertainty, and subsequent attention influence, linking them to syntactic categories.This complements prior work focused on errors visible in final outputs.
  • Contributions: ACTOR introduces critical-token-triggered, on-demand retrieval and a weighted mechanism, achieving state-of-the-art repository-level code-generation performance.The framework is presented as dynamic and fine-grained rather than static.
Loading 2609.01601v1…