PaperlayerView exact revision on arXiv

Source-linked AI summary

AlignCoder: Aligning Retrieval with Target Intent for Repository-Level Code Completion

Tianyue Jiang, Yanli Wang, Yanlin Wang, Daya Guo, Ensheng Shi, Yuchi Ma, Jiachi Chen, Zibin Zheng

arXiv:2601.19697v2cs.SEcs.AI

TL;DR

Repository-level completion is challenging because code LLMs have limited repository-specific context, and existing retrieval methods face query-target misalignment and do not effectively use inference information. AlignCoder enhances retrieval queries with sampled candidate completions and trains a retriever with reinforcement learning. It outperforms baselines in the reported evaluations, with an 18.1% EM-score improvement on CrossCodeEval Python.

Abstract

from arXiv · show

Repository-level code completion remains a challenging task for existing code large language models (code LLMs) due to their limited understanding of repository-specific context and domain knowledge. While retrieval-augmented generation (RAG) approaches have shown promise by retrieving relevant code snippets as cross-file context, they suffer from two fundamental problems: misalignment between the query and the target code in the retrieval process, and the inability of existing retrieval methods to effectively utilize the inference information. To address these challenges, we propose AlignCoder, a repository-level code completion framework that introduces a query enhancement mechanism and a reinforcement learning based retriever training method. Our approach generates multiple candidate completions to construct an enhanced query that bridges the semantic gap between the initial query and the target code. Additionally, we employ reinforcement learning to train an AlignRetriever that learns to leverage inference information in the enhanced query for more accurate retrieval. We evaluate AlignCoder on two widely-used benchmarks (CrossCodeEval and RepoEval) across five backbone code LLMs, demonstrating an 18.1% improvement in EM score compared to baselines on the CrossCodeEval benchmark. The results show that our framework achieves superior performance and exhibits high generalizability across various code LLMs and programming languages.

I. INTRODUCTION

Repository-level completion is difficult when models lack project-specific context and retrieval queries do not align with target code. AlignCoder enhances queries with candidate completions and trains a retriever to use inference information; it outperforms baselines in the reported evaluation.

I. INTRODUCTION

  • Repository-specific context is difficult for models to acquire from newly created, proprietary, or work-in-progress projects, and concatenating all files introduces irrelevant information.RAG instead retrieves code snippets as cross-file context for completion.Source: §I · p. 1

P1 The misalignment between query and target code

  • Existing RAG methods use unfinished code as the query, which can omit target-related tokens, while dense retrievers may not leverage code LLM inference capabilities.RepoCoder's iterative approach can propagate incorrect completions and reduce efficiency through repeated retrieval.Source: §I · p. 1, P1 The misalignment between query and… · p. 1
  • AlignCoder appends sampled candidate completions to unfinished code to form an enhanced query, then trains AlignRetriever with reinforcement learning to use that inference information for retrieval.The reward model evaluates target-code generation perplexity using retrieved snippets to update the retriever.Source: P1 The misalignment between query and… · p. 2
  • 18.1% improvement in EM score compared with baselines on CrossCodeEval Python, with effectiveness across code LLMs and programming languages.The evaluation uses five backbone LLMs and two benchmarks.Source: P1 The misalignment between query and… · p. 2

1) Retrieve-then-Generate:

  • Retrieve-then-generate methods query the repository with unfinished code, then provide the top-k snippets as context for generating the target completion.ReACC and RLCoder follow this paradigm, which may miss code relevant to the target because the query lacks key tokens.Source: 1) Retrieve-then-Generate · p. 2

2) Iterative Generate-and-Retrieve:

Iterative retrieval can propagate errors from a single intermediate completion, while unfinished code and target code can occupy different semantic spaces. The section motivates multiple sampling as a way to increase the chance of a correct completion and reports rising EM@k with k on both benchmarks.

2) Iterative Generate-and-Retrieve:

  • Iterative generate-and-retrieve methods reuse each candidate completion as the next query, so an incorrect intermediate completion can propagate errors through later retrieval rounds.The approach alternates generation and retrieval across multiple iterations.Source: 2) Iterative Generate-and-Retrieve · p. 2

B. Semantic Gap in Repository Code Retrieval

  • Unfinished code and target code occupy different semantic spaces, making the unfinished code alone a suboptimal retrieval query.Prior approaches either require aligning these spaces or risk cascading errors from a single sample in iterative retrieval.Source: B. Semantic Gap in Repository Code… · p. 2

1) Single Sample Unreliability:

2) Multiple Sampling Benefits:

  • Multiple samples are not truly independent, but stochastic sampling supports treating them as quasi-independent for analysis; correlation reduces the effective number of independent samples by a factor of 1 − ρ.The text identifies temperature-based, top-k, or nucleus sampling as sources of stochasticity and describes ρ as the inter-sample correlation coefficient.Source: 2) Multiple Sampling Benefits · p. 3
  • The analysis states that success probability increases substantially with sampling attempts and that multiple sampling consistently yields higher accuracy than single sampling.The stated benefit is qualified by a diminishing-returns threshold beyond which extra samples may add marginal utility or degrade overall performance through answer-selection complexity and computational overhead.Source: 2) Multiple Sampling Benefits · p. 3
  • For DeepSeekCoder-1B with the aforementioned prompt and 10 sampling attempts, EM@k (k=1–10) increases on CrossCodeEval and RepoEval, showing multiple sampling gives higher correct-completion probability than single sampling.Source: 2) Multiple Sampling Benefits · p. 4, Figure 2 · p. 4

3) Diminishing Returns and Optimal Sampling Threshold:

The section models the trade-off in repeated completion sampling: more samples can help find a correct completion, but accumulated errors and costs create an optimal threshold beyond which utility may fall.

3) Diminishing Returns and Optimal Sampling Threshold:

  • As sampling grows, incorrect responses can accumulate and potentially overwhelm correct completions during aggregation or selection.Source: 3) Diminishing Returns and Optimal… · p. 4
  • For random selection, the chance of choosing correctly depends on how many of n samples are correct, while more sophisticated selection mechanisms make the relationship more complex.Source: 3) Diminishing Returns and Optimal… · p. 4
  • The threshold balances a higher probability of at least one correct sample against accumulated errors and computational cost, weighted by α, β, and γ.The weights represent correctness benefit, error penalty, and computational cost, respectively.Source: 3) Diminishing Returns and Optimal… · p. 4
  • Beyond an optimal sampling threshold, additional samples may offer only marginal utility or degrade performance through answer-selection complexity and computational overhead.The section concludes that multiple sampling is beneficial, but not without limit.Source: 3) Diminishing Returns and Optimal… · p. 4

D. Perplexity in LLMs for Code Assessment

Perplexity measures how well a model predicts token sequences and is used here as an intrinsic confidence signal for code completion candidates. It can help rank candidates, but it does not directly measure correctness.

D. Perplexity in LLMs for Code Assessment

  • Perplexity measures how well a probability model predicts token sequences and serves as an intrinsic measure of LLM confidence.Lower perplexity indicates higher model confidence.Source: D. Perplexity in LLMs for Code… · p. 4
  • Perplexity can relate to code quality through statistical regularity, syntactic consistency, and semantic coherence, but it does not directly measure correctness.Source: D. Perplexity in LLMs for Code… · p. 4
  • Uncommon but correct styles, domain-specific patterns, and model-understanding gaps can cause perplexity to diverge from correctness; it is presented as a first-order ranking heuristic when combined with other validation methods.The cited studies report a moderate positive correlation between low perplexity and code correctness [27], [28].Source: D. Perplexity in LLMs for Code… · p. 4

III. METHODOLOGY

AlignCoder builds a retrieval codebase, uses retrieved context and sampled completions to create an enhanced query, then retrieves references for final completion; its codebase combines base and dependency snippets.

III. METHODOLOGY

  • AlignCoder first constructs the retrieval codebase, then uses coarse retrieval and sampled completions to form an enhanced query for fine retrieval; retrieved references support the generator's final completion.During training, an evaluator-based reward updates the retriever parameters.Source: §III · p. 5

A. Codebase Construction

  • The codebase contains base snippets representing initial cross-file code segments and dependency snippets representing class hierarchies and API interactions.Source: A. Codebase Construction · p. 5

1) Base Code Snippets Construction:

  • Base snippets use Split-Aggregate: files are split at blank lines into contiguous nonempty mini-blocks, then aggregated into snippets within a predefined maximum line count.The approach is motivated as preserving code continuity compared with fixed windows.Source: 1) Base Code Snippets Construction · pp. 5–6

2) Dependency Code Snippets Construction:

AlignCoder constructs dependency snippets from intra-repository imports and uses sampled completions to enhance the query before fine-grained retrieval.

2) Dependency Code Snippets Construction:

  • Dependency construction parses imports, filters out standard-library and third-party references, and extracts matching cross-file code bodies for imported entities.Class imports yield class and method signatures; imported functions or methods are represented by their code bodies.Source: 2) Dependency Code Snippets Construction · p. 6
  • Imported classes become individual dependency snippets, while imported methods or functions have their signatures aggregated into a single snippet; these join the base snippets in the retrieval codebase.Source: 2) Dependency Code Snippets Construction · p. 6

B. Query Enhancemant Mechanism

  • Coarse BM25 retrieval supplies context for sampling candidate completions, which are appended to the unfinished code as an enhanced query for fine-grained retrieval.In the example, sampled candidates 2 and 3 contain key tokens related to the target code.Source: B. Query Enhancemant Mechanism · p. 6

1) Training data construction:

AlignCoder’s training data is built from cross-file repository clusters, and its retriever objective rewards retrieval aligned with the enhanced query and target code. The evaluation setup uses two repository-level benchmarks and five code LLM backbones.

1) Training data construction:

  • Training data comes from 10,000 pre-March-2023 Python and Java repositories excluded from the evaluation benchmarks, with targets sampled from multi-file dependency clusters.Clusters are topologically sorted, and target segments are sampled from non-first files at lengths of 16–96 tokens.Source: 1) Training data construction · p. 6

2) The training objective function:

  • AlignRetriever training maximizes a reward based on retrieved-snippet correctness and similarity to the enhanced query.The objective identifies a candidate with the lowest target-code perplexity among retrieved snippets; target-code length is measured in tokens.Source: 2) The training objective function · p. 7

A. Benchmarks and Backbone Models

  • Evaluation uses CrossCodeEval’s Python and Java sets and RepoEval’s line-level and API-invocation tasks.The benchmarks assess repository-level completion, including cross-file context on CrossCodeEval.Source: A. Benchmarks and Backbone Models · p. 7

IV. EXPERIMENTAL SETUP

  • The experimental setup evaluates five code LLM generators: CodeLlama-7B, StartCoder-7B, StarCoder2-7B, DeepSeekCoder-1B, and DeepSeekCoder-7B.The experiments are framed around repository-level completion benchmarks and backbone models.Source: A. Benchmarks and Backbone Models · p. 7, §IV · p. 7

B. Baseline Methods

The experiments compare AlignCoder with three RAG-based baselines and evaluate completion accuracy using EM and ES. The section also specifies implementation settings and poses questions about effectiveness, sampling, components, and parameter robustness.

B. Baseline Methods

  • AlignCoder is compared with ReACC, RepoCoder, and RLCoder, which use hybrid retrieval, iterative retrieval and generation, and reinforcement-learning retrieval, respectively.RLCoder is described as learning to retrieve useful snippets without labeled data.Source: B. Baseline Methods · p. 7

C. Evaluation Metrics

  • Completion accuracy is evaluated with Exact Match (EM) and Edit Similarity (ES).The metrics follow established practice cited by the paper.Source: C. Evaluation Metrics · p. 7

D. Experimental Details

  • Retriever training uses UniXcoder initialization and DeepSeekCoder-1B as evaluator for 20 epochs, with 3,000 samples per epoch and a 5e-5 learning rate.At inference, vLLM accelerates the approach’s multiple sampling.Source: D. Experimental Details · p. 7

V. EXPERIMENTAL RESULTS

  • The experiments ask whether AlignCoder is effective, how sampling strategies and counts perform, how components contribute, and whether parameter settings are robust.These are posed as four research questions; the supplied passages do not report their answers.Source: §V · p. 7

A. RQ1: Overall Performance of AlignCoder

AlignCoder outperforms the baselines across the reported benchmarks and backbone models, with its largest gains on CrossCodeEval. Multiple sampling generally helps, but excessive sampling can reduce performance; ablations also show a larger impact from removing retriever training than from removing other components.

A. RQ1: Overall Performance of AlignCoder

  • AlignCoder outperforms previous methods across all evaluation settings; on CrossCodeEval Python, its EM gains over RLCoder range from 12.7% to 18.1%.On CrossCodeEval Java, EM gains over RLCoder range from 7.2% to 14.0% across backbone models.Source: A. RQ1: Overall Performance of… · pp. 7–8, Table I · p. 8

B. RQ2: Mutiple Sampling Strategy

  • With four samples, AlignCoder achieves its best average EM and ES performance; multiple sampling generally beats single sampling, while performance declines on some datasets above four.At five samples, CrossCodeEval Java and both RepoEval tasks fall below single-sampling results.Source: B. RQ2: Mutiple Sampling Strategy · p. 8, Table II · p. 9

C. RQ3: Ablation Studies

  • Removing dependency context or query enhancement generally lowers EM and ES across benchmarks, while removing reinforcement learning causes more pronounced performance degradation.The study concludes that context and query enhancement alone are insufficient without retriever training to use information in the enhanced query.Source: Table III · p. 9, EM ES EM ES · p. 9

D. RQ4: Sampling Parameter Stability

AlignCoder's performance is generally stable across the tested temperature and top-p settings. Most configurations change EM and ES by less than 1%, although some show larger differences.

D. RQ4: Sampling Parameter Stability

  • Most tested temperature and top-p configurations change EM and ES by less than 1% across both benchmarks.Five configurations show differences; temperature 0.8 with top-p 0.7 decreases CrossCodeEval Python EM by 2.4%.Source: D. RQ4: Sampling Parameter Stability · pp. 9–10

E. Case Study

AlignCoder combines inference-informed retrieval with retriever training; a case study contrasts its dependency-focused retrieval with surface-similar snippets retrieved by RLCoder and RepoCoder. The conclusion reports an 18.1% EM improvement on the CrossCodeEval Python subset and generalizability across models and programming languages.

E. Case Study

  • AlignCoder retrieved a Validator dependency snippet that supported the correct completion, while RLCoder and RepoCoder retrieved surface-similar base code and produced an incorrect result.The case study concerns completing the update function of the BuildConfiguration class.Source: E. Case Study · p. 10

A. Code Completion

  • Code completion has progressed from rule-based and statistical methods to neural and LLM systems, with RAG approaches retrieving code snippets, documentation, or examples.The cited related work includes Red-Coder, DocPrompting, and AceCoder.Source: A. Code Completion · p. 10

B. Repository-Level Code Completion

  • Prior repository-level approaches include dependency learning, static analysis, code graphs, iterative retrieval, and tool invocation; AlignCoder combines inference-derived reference snippets with reinforcement-learning retriever training.The conclusion describes query enhancement with multiple candidate completions and an AlignRetriever trained to use inference information.Source: B. Repository-Level Code Completion · pp. 10–11, §VII · p. 11

VII. CONCLUSION

  • 18.1% improvement in EM score on the Python subset of CrossCodeEval, with strong generalizability across code LLMs and programming languages.The conclusion reports experiments on CrossCodeEval and RepoEval across five backbone models.Source: §VII · p. 11
Loading 2601.19697v2…