Repoformer: Selective Retrieval for Repository-Level Code Completion

• Repository-level completion requires cross-file knowledge such as user-defined APIs and inter-module dependencies, which current-file context alone may miss.RAG supplies relevant cross-file snippets or documentation to code LMs.Source: §1 · p. 1
• Always retrieving is often unnecessary: up to 80% of standard-RAG retrievals do not improve performance, and some introduce irrelevant information that degrades it.Always retrieving also adds latency, including for sparse retrieval on moderately sized repositories.Source: §1 · p. 1
• REPOFORMER estimates whether retrieval could improve its prediction, weighing both whether it knows the answer and whether the task depends on cross-file information.The framework is trained with self-supervision to assess retrieval need and complete code robustly with optional retrieved context.Source: §1 · p. 2
• The training approach simulates RAG on sampled repository blanks, pairing them with retrieved cross-file contexts to derive retrieval labels.This supports training the code LM to self-evaluate retrieval need and complete code with optional retrieval augmentation.Source: §1 · p. 2
• REPOFORMER-16B achieves state-of-the-art performance across tasks, while selective retrieval provides up to 70% inference speedup without harming accuracy.The same-sized REPOFORMER outperforms always-retrieving StarCoderBase by more than 3 absolute points for edit similarity across multiple tasks.Source: §1 · p. 2
• Across languages and retrievers, REPOFORMER makes precise retrieval-abstention decisions and remains robust to retrieved contexts.The framework can also serve as a plug-and-play selective retrieval policy for other generation models.Source: §1 · p. 2

• Earlier repository-level completion work trained n-gram LMs, RNNs, and Transformers or used retrieved code and documentation, including with black-box LMs [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13].The paper positions selective retrieval through self-assessment as orthogonal to those approaches.Source: §2 · p. 3
• Adaptive-RAG policies have adjusted retrieval weighting by LM performance or retrieval quality, used predictive uncertainty, or avoided retrieval for popular facts [1, 2, 3, 4].Source: §2 · p. 3
• Unlike SKR and Self-RAG, this approach fine-tunes a code LM for self-evaluation without extra modules, a knowledge store, or oracle-LM-generated labels [1, 2].The paper reports no performance harms, substantial speedup, and high decision accuracy.Source: §2 · p. 3

3.1 Background

• A repository-level completion instance comprises left and right same-file context, a target completion, and the repository’s other files.The retriever uses the in-file context to return cross-file code chunks, which a code LM may use with the same-file context to generate a completion.Source: §3.1 · p. 3

3.2 Self-selective RAG for Code Completion

• Trial-retrieval and trial-generation heuristics can preserve performance with a 50% retrieval budget on some RepoEval tasks, but generalize poorly across tasks and still incur retrieval-decision latency.Source: §3.2 · pp. 3–4
• Self-selective RAG has the LM trigger retrieval with <cc> or abstain with an empty token after observing the current-file context.The LM then completes code with or without retrieved contexts; the decision and generation run in a single left-to-right pass, with a single-forward-pass decision overhead.Source: §3.2 · p. 4

3.3 Self-supervised Multi-task Learning

• The training objective combines a retrieval-assessment loss with a generation loss for completion using either in-file context alone or retrieved context.Instances are verbalized differently according to their retrieval label, and the two losses are weighted together.Source: §3.3 · p. 4
• Training instances are built from sampled code chunks or function bodies, retrieved contexts, and labels based on whether retrieval improves completion quality by more than a threshold.The data use permissively licensed repositories; retrieval queries include the target for half the data to align with noniterative and iterative RAG use cases.Source: §3.3 · p. 4

4.1 REPOFORMER Implementation Details

• 18k filtered Python repositories supplied 240k chunk and 120k function instances, labeled with StarCoderBase-1B at T=0 and split into 500 validation repositories and the remainder for training.Repositories were required to have at least five Python files, three imports per file, and two local imports per file.Source: §4.1 · p. 5
• StarCoderBase 1B, 3B, 7B, and 16B variants were fine-tuned for two epochs with λ=1.0 and a maximum sequence length of 2048.Training used a learning rate of 2e-5, batch size 512, 50 warmup steps, and linear learning-rate decay.Source: §4.1 · p. 5
• Hyperparameters were selected by grid search on StarCoderBase-1B using validation code-completion performance.The search varied learning rate, λ, training epochs, and warmup steps.Source: §4.1 · p. 5

4.2 Evaluation Setup

• Evaluation covered RepoEval and CrossCodeEval, plus CrossCodeLongEval, a new chunk- and function-completion benchmark built from 1,500 raw Python repositories.RepoEval covers line, API, and function completion across 32 Python repositories; CrossCodeEval covers line completion in four languages.Source: §4.2 · p. 5
• Reference-based evaluation reports exact match and edit similarity, while execution-based evaluation reports unit-test pass rate.Edit similarity is normalized by the longer of the prediction and reference lengths; unit tests pass when replacing the reference with the prediction causes no test to fail.Source: §4.2 · p. 5
• The experiments use CodeGen-Mono and StarCoder/StarCoderBase, with StarCoder trained from StarCoderBase on an additional Python corpus.Source: §4.2 · p. 5

5.1 Is retrieval always helpful?

• Retrieved contexts improve performance on 20% or fewer instances, leave it unchanged on more than 60%, and harm it on about 20% across tasks and model sizes.The pattern holds for API and function completion and for small and moderate-to-large code LMs.Source: §5.1 · p. 5

5.2 REPOFORMER achieves strong code completion performance via selective RAG

• Threshold selection outperforms greedy selection in code completion accuracy on every task.The two policies offer different balances between accuracy and inference budget.Source: §5.2 · p. 6
• REPOFORMER-16B outperforms the strongest StarCoder baseline by 3%, averaged across all tasks, setting a new state of the art for repository-level code completion.Table 2 also reports that REPOFORMER outperforms same-sized StarCoderBase while using a smaller retrieval budget.Source: §5.2 · p. 6, Table 2 · p. 7

5.3 REPOFORMER improves inference efficiency

• Threshold selection improves both accuracy and latency over always retrieving, while greedy selection yields larger latency gains with around 1.0 ES performance degradation.The latency evaluation covers line and API completion.Source: §5.3 · p. 7, Table 3 · p. 7
• With dense retrieval, threshold selection brings more than 70% speedup because retrieval accounts for a larger share of total latency.The paper notes that dense retrieval on large repositories can consume more than 80% of the RAG pipeline’s latency.Source: §5.3 · p. 7
• REPOFORMER-1B’s selective decisions reduce latency by approximately 25% with larger generation models while improving their accuracy.The paper reports this result for Llama, CodeGen25, and ChatGPT.Source: Table 4 · p. 8, §5.3 · p. 8

Threshold sensitivity

• At threshold 0.4, REPOFORMER maintains performance similar to always retrieving while reducing latency by 50%.Increasing the threshold initially improves completion performance by avoiding potentially harmful retrievals; results are relatively insensitive to threshold settings.Source: §6 · p. 8, Figure 4 · p. 8

Selective decision quality

• Abstention precision exceeds 0.8 on all tasks except RepoEval Function, where it is 0.78.Abstentions generally occur when the prediction is already correct or retrieval cannot improve it; calibration is suboptimal for function completion with UT.Source: Figure 5 · p. 8, §6 · p. 8

Robustness to retrieval

• On instances where it requests retrieval, REPOFORMER shows more and greater performance gains than StarCoderBase, with fewer performance decreases.It also retains quality and latency improvements with dense retrieval despite a different context distribution.Source: §6 · pp. 8–9, Figure 6 · p. 9

Ablation Study

• Removing the self-evaluation loss slightly improves RAG performance, while removing training contexts worsens it; combining losses can make retrieval selection nearly constant.The combined-loss variant can assign <cc> probability near 1, and placing retrieved context inside the infilling region performs much worse on function completion than on chunk completion.Source: §6 · p. 9

7 Conclusion

• REPOFORMER identifies whether cross-file context is necessary and self-triggers retrieval, with evaluations reporting improved accuracy and reduced latency.The conclusion describes the framework as a potential approach for practical coding environments.Source: §7 · p. 9
• Future work could explore using REPOFORMER as a speculative-decoding draft model and developing more accurate, scalable labels for function completion.The current function-completion labels use lexical similarity as a heuristic.Source: §7 · p. 9
• The selective policy is uniform across repositories; adapting it to repository-specific RAG-friendliness is another proposed direction.The discussion notes that repositories may differ in their level of duplication.Source: §7 · p. 9

Impact Statement

• The impact statement calls for ongoing evaluation of code automation and mechanisms to prevent RAG systems from revealing sensitive data.It also flags possible implications for programming education and inadvertent bias.Source: Impact Statement · p. 10
• Experiments mainly use open-source, permissively licensed repositories and models; RepoEval includes some non-permissively licensed repositories, whose dataset and code are not redistributed or repurposed.The work relies on the original RepoEval authors’ distributed dataset and code.Source: Impact Statement · p. 10

A Detailed RAG Execution Setup

• Indexing uses overlapping fixed-size code chunks, while each query uses a fixed number of lines immediately before the completion blank.Chunk sizes are 20 lines for line, API, and chunk completion and 50 for function completion; stride is half the chunk size.Source: A Detailed RAG Execution Setup · p. 14
• Retrieval selects the top 10 chunks by Jaccard similarity and replaces each selected chunk with the immediately following chunk for the cross-file context.The paper reports that the explored cosine-similarity and CodeBLEU alternatives do not outperform Jaccard.Source: A Detailed RAG Execution Setup · p. 14
• Retrieved cross-file context is concatenated with in-file context; CodeGen and StarCoder use different prompt orderings, with StarCoder prompted in fill-in-the-middle format.The prompt formats place left, right, and cross-file contexts in different orders.Source: A Detailed RAG Execution Setup · p. 14
• Experiments use greedy decoding, truncate left, right, and cross-file contexts to specified token limits, and apply task-specific post-processing to predictions.Maximum generation length is 50 tokens for line, API, and chunk completion and 256 tokens for function completion.Source: A Detailed RAG Execution Setup · p. 14

B Why infilling?

• Directly including the right context preserves information about future execution and, for fill-in-the-middle-trained models such as StarCoder, is preferred to splitting it for retrieval.The comparison evaluates current-file and retrieval-augmented completion on RepoEval.Source: B Why infilling? · p. 14
• Right context improves RepoEval performance with or without cross-file context, with gains consistent for API and function completion and larger for StarCoder than CodeGen.The paper consequently focuses on infilling with the StarCoder family.Source: B Why infilling? · p. 15, Figure 7 · p. 15

C.1 Trial Retrieval

• Trial retrieval compares a query with the top-ranked retrieved chunk and uses a threshold to decide whether to include the top 10 chunks as cross-file context.The section tests Jaccard, weighted n-gram, and UniXcoder similarity scores.Source: §C.1 · p. 15
• For most retrievers and tasks, at most 60% retrieval budget reaches full-retrieval performance, but this strategy does not reduce latency because it still retrieves.The method also does not use whether the language model itself would benefit from the retrieved context.Source: §C.1 · p. 15

C.2 Trial Generation

• In most cases, 40% retrieval can be saved without sacrificing code-completion performance using retriever similarity scores to select instances for retrieval.The budget is the percentage of instances on which retrieval is performed.Source: Figure 8 · p. 16
• For line completion, token uncertainty outperforms entropy; for API completion, entropy is slightly better, while uncertainty metrics fail for long function generation.The plotted scores are measured across the entire test dataset at each retrieval budget.Source: Figure 9 · p. 16
• Uncertainty-based selection becomes less effective as output length increases from line to API to function completion, and it does not outperform trial retrieval.The comparison evaluates entropy and token uncertainty across retrieval budgets.Source: §C.2 · p. 16, Figure 9 · p. 16

D Data Creation for REPOFORMER Training and CrossCodeLongEval

• Chunk-completion examples sample sub-chunks from clustered 10-line repository chunks, retrieve top-three cross-file contexts, and label retrieval using the ES score difference against a threshold.The query uses surrounding in-file context, optionally including the target; retrieved contexts are obtained using Jaccard similarity.Source: D Data Creation for REPOFORMER… · p. 17
• Function-completion examples sample bodies from clustered functions of 3–30 lines, retrieve top-three cross-file contexts, and use the thresholded ES score difference as the retrieval label.The query combines in-file context around the function body, optionally including that body.Source: D Data Creation for REPOFORMER… · p. 17
• 30.18% of chunk instances and 35.16% of function instances receive positive retrieval labels in the training data.The analyzed sets contain 240k chunk-completion and 120k function-completion instances; Figure 10 plots performance changes after retrieved cross-file context is added.Source: Figure 10 · p. 18, D Data Creation for REPOFORMER… · p. 18
• CrossCodeLongEval addresses RepoEval’s limited repository coverage using newly sampled blanks from collected repositories rather than CrossCodeEval’s filtered instances.The cited setup excludes CrossCodeEval instances that StarCoderBase-1B can answer without retrieved context.Source: D Data Creation for REPOFORMER… · p. 18

E.1 Calibration of REPOFORMER’s Selective Retrieval Prediction

• REPOFORMER-1B is generally near-calibrated for line and API completion under ES, but its function-completion predictions are not calibrated under UT.The authors suggest ES as a possible reason and call for improved function-completion labels and multiple self-assessments for long generations.Source: Figure 11 · p. 19, §E.1 · p. 19

E.2 CrossCodeEval and Multilingual REPOFORMER

• REPOFORMER trained on Python data transfers across multiple languages, while REPOFORMER-M outperforms same-sized StarCoderBase and is slightly ahead of REPOFORMER at 1B and 7B.At 3B, REPOFORMER-M and REPOFORMER give similar performance.Source: §E.2 · p. 19, Table 7 · p. 20

E.3 REPOFORMER’s Robustness to the Retriever Choice

• More than 70% speedup is reported with threshold selection using dense retrieval, alongside accuracy and latency gains for selective retrieval.Compared with always retrieving, SelectiveT gains in both accuracy and latency, while SelectiveG yields larger latency gains with a small performance degradation.Source: Table 8 · p. 20, §E.3 · p. 20

E.4 Full Latency-Accuracy Visualizations

• Across REPOFORMER-1B, REPOFORMER-3B, StarCoderBase-7B, and StarCoder, selective RAG improves accuracy and latency, with clearer gains for line and API completion.Figures 12–15 show these latency-accuracy trade-offs; for function completion, retrieval overhead is generally low on RepoEval’s small repositories.Source: §E.4 · p. 20, Figure 12 · p. 21, Figure 13 · p. 21