Source-linked AI summary
Domain Adaptive Code Completion via Language Models and Decoupled Domain Databases
Ze Tang, Jidong Ge, Shangqing Liu, Tingwei Zhu, Tongtong Xu, Liguo Huang, Bin Luo
TL;DR
Domain-specific code knowledge can limit language-model completion, while continuously changing projects and inaccessible model weights complicate fine-tuning. kNM-LM retrieves from a decoupled database of model mistakes and combines it with the language model through Bayesian inference. It improves CodeGPT and UniXcoder on intra-project and intra-scenario completion, and does not require access to model weights.
Problem
General language models may lack domain-specific code knowledge, and frequent domain changes or inaccessible model weights can make fine-tuning infeasible.
Source: §I · p. 1Method
kNM-LM stores code tokens the language model mispredicts in a decoupled database and uses Bayesian inference to combine retrieved information with the model.
Source: §I · p. 2Results
kNM-LM improves both CodeGPT and UniXcoder on intra-project and intra-scenario code completion.
Source: §I · p. 2Takeaways & Limitations
Because it does not require access to language-model weights, kNM-LM is suitable for settings with black-box model access.
Source: §X · p. 11The evaluation used CodeGPT and UniXcoder rather than black-box systems such as Codex or ChatGPT, and did not verify ReACC and Hybrid on Java.
Source: §VIII · pp. 9–10
Abstract
from arXiv · showhide
Large Language Models (LLMs) have demonstrated remarkable performance in code completion. However, due to the lack of domain-specific knowledge, they may not be optimal in completing code that requires intensive domain knowledge for example completing the library names. Although there are several works that have confirmed the effectiveness of fine-tuning techniques to adapt language models for code completion in specific domains. They are limited by the need for constant fine-tuning of the model when the project is in constant iteration. To address this limitation, in this paper, we propose $k$NM-LM, a retrieval-augmented language model (R-LM), that integrates domain knowledge into language models without fine-tuning. Different from previous techniques, our approach is able to automatically adapt to different language models and domains. Specifically, it utilizes the in-domain code to build the retrieval-based database decoupled from LM, and then combines it with LM through Bayesian inference to complete the code. The extensive experiments on the completion of intra-project and intra-scenario have confirmed that $k$NM-LM brings about appreciable enhancements when compared to CodeGPT and UnixCoder. A deep analysis of our tool including the responding speed, storage usage, specific type code completion, and API invocation completion has confirmed that $k$NM-LM provides satisfactory performance, which renders it highly appropriate for domain adaptive code completion. Furthermore, our approach operates without the requirement for direct access to the language model's parameters. As a result, it can seamlessly integrate with black-box code completion models, making it easy to integrate our approach as a plugin to further enhance the performance of these models.
I. INTRODUCTION
Domain-specific code changes can make general language models and repeated fine-tuning a poor fit for adaptive completion. kNM-LM combines a decoupled database of model mistakes with a language model, and improves CodeGPT and UniXcoder on intra-project and intra-scenario completion.
I. INTRODUCTION
- General language models may miss domain-specific features, while black-box access and continuously changing code can make fine-tuning infeasible.Code in IDEs includes changing third-party libraries and intra-project references.Source: §I · p. 1
- Unlike kNN-LM's large datastore and manually selected interpolation weight, kNM-LM stores only tokens the language model mispredicts and combines the systems with Bayesian inference.The decoupled database reduces storage and search resources.Source: §I · pp. 1–2
- The method retrieves similar mistaken tokens, normalizes them into a distribution, and uses the model's domain error rate and preceding-token predictions to calculate a posterior for merging distributions.The posterior combines the datastore and language-model distributions.Source: §I · p. 2
- The authors propose kNM-LM as retrieval-augmented domain adaptation without fine-tuning and assess plugin-relevant speed, space, code-type, and API-call performance.They report satisfactory performance across the considered user-experience factors.Source: §I · p. 2
- kNM-LM improves both CodeGPT and UniXcoder on intra-project and intra-scenario code completion.The paper also reports satisfactory performance across plugin factors including completion speed, space use, specific code types, and third-party API calls.Source: §I · p. 2
A. k-Nearest Neighbors Language Model
The section describes kNN-LM retrieval and interpolation, then introduces Bayesian inference as a way to update beliefs using evidence.
A. k-Nearest Neighbors Language Model
- kNN-LM retrieves nearest-neighbor examples for the current context and interpolates their next-token distribution with the language model’s prediction.The interpolation coefficient λ ranges from 0 to 1 and requires careful selection.Source: A. k-Nearest Neighbors Language Model · p. 2
B. Bayesian Inference
- Bayesian inference updates beliefs about an event using observed evidence, with the posterior depending on the prior and likelihood.The section introduces a Beta prior and posterior expectation for observations in which the event occurs α times among N.Source: B. Bayesian Inference · pp. 2–3
- The primary workflow figure presents the proposed kNM-LM and marks components requiring access to the black-box language model.Source: Figure 1 · p. 3
III. APPROACH
The approach separates the retrieval module from the language model and uses Bayesian inference to combine them, grounding that design in distinct correct- and incorrect-prediction events.
III. APPROACH
- The task predicts the next token or line from known code context, using in-domain code to adapt general or black-box language models.The approach separately predicts next-token probabilities with the LM and retrieval module before combining them through Bayesian inference.Source: §III · p. 3
A. Analysis of R-LM from a Statistical Perspective
- The retrieval module should leave correct LM predictions unchanged and use similar past mistakes to correct erroneous predictions.The approach frames these as distinct events: E for a correct prediction and E′ for an incorrect one.Source: A. Analysis of R-LM from a Statistical… · p. 3
- Decoupling lets kNM-LM compute the combination coefficient p(E′|ct) mathematically, avoiding the manual empirical selection required by prior R-LM ensembles.The coefficient depends on both domain-specific context and the language model, supporting adaptation across models and domains.Source: A. Analysis of R-LM from a Statistical… · pp. 3–4
B. Step 1. Decoupled Database Construction
The method builds a decoupled datastore from in-domain code, obtains LM predictions and context embeddings, and retrieves similar stored contexts for inference.
B. Step 1. Decoupled Database Construction
- The datastore stores embedded preceding contexts as keys and tokens associated with LM-error event E′ as values; its smaller size than kNN-LM saves storage and search time.Its reduction ratio is inversely proportional to the LM error rate on the in-domain database.Source: B. Step 1. Decoupled Database… · p. 4
C. Step 2. Inference with LM
- The LM inference step obtains the next-token probability and embedded context through the language model’s APIs, without accessing its parameters.Source: C. Step 2. Inference with LM · p. 4
D. Step 3. Inference with DB
- The retrieval module uses the embedded context to find nearest datastore keys and form a distribution over retrieved tokens; known-context token checks provide likelihood evidence for updating beliefs.For decoder-only models such as GPT-3, checking predictions within the known context incurs no additional computational cost.Source: D. Step 3. Inference with DB · p. 4, Figure 2 · p. 4
E. Step 4. Combination of Inference Results
The method combines language-model and retrieval predictions by using Bayesian inference to estimate how much the model is likely to err in the current context. The resulting coefficient shifts reliance between the model and retrieval according to their error rates.
E. Step 4. Combination of Inference Results
- Bayes’ rule estimates p(E′|ct), the combination coefficient, from a prior and the likelihood of observing the context given E′.The direct probability is unavailable because it concerns the next token, which is not yet known.Source: E. Step 4. Combination of Inference… · p. 4
- The likelihood uses the latest N−1 observations: α are classified as E′ and N−α as E.For each context token, the language model prediction is compared with the observed token to classify the event.Source: E. Step 4. Combination of Inference… · p. 4
- The model treats p(E′) as a Beta-distributed variable with expectation err, then updates its distribution using the new context evidence.The posterior parameters add α to err·N and N−α to (1−err)·N.Source: E. Step 4. Combination of Inference… · pp. 4–5
- k sets the number of neighbors used for retrieval, while N sets the number of observations used for the likelihood.Source: E. Step 4. Combination of Inference… · p. 5
- When the language model’s error rate is high, λ grows and favors retrieval; when it is low, λ shrinks and favors the model.The coefficient depends on the code context, language model, and database.Source: E. Step 4. Combination of Inference… · p. 5
IV. EXPERIMENTAL SETUP
The experiments ask whether kNM-LM improves intra-project and intra-scenario completion, with additional investigation of code types and Android API-call lines. The evaluation compares two pretrained language models and several retrieval-augmented alternatives.
IV. EXPERIMENTAL SETUP
- The study evaluates kNM-LM on intra-project completion and intra-scenario completion, with further analyses of specific code types and Android API-call lines.Source: §IV · p. 5
A. Baselines
- CodeGPT and UniXcoder serve as base models for Java and Python; UniXcoder uses one cross-language model for both languages.The selected CodeGPT variants are Java- and Python-adapted models.Source: A. Baselines · p. 5
- Retrieval baselines include block methods BM25, ReACC, and Hybrid, plus token-level kNN-LM.ReACC and Hybrid are evaluated only for Python because the released ReACC retrieval model is Python-only.Source: A. Baselines · p. 5
B. Implementation Details
The evaluation uses historical project snapshots and later-created methods to test intra-project completion, alongside token-type analyses. Across projects, kNM-LM improves completion and has the largest gains among the compared baselines.
B. Implementation Details
- Implementation sets k to 8 for RQ1 and 1024 for RQ2, λ to 0.1, and window size N to 8; token accuracy and line EM and ES are the evaluation metrics.Source: B. Implementation Details · p. 5
B. Dataset
- The project dataset uses 10 large and 10 small Java-mid projects, with the snapshot at 80% of method creation forming the database and later methods forming the test set.Methods are ordered by Git commit time.Source: A. Study Design · pp. 5–6, B. Dataset · p. 6
C. Results
- 26.67% with CodeGPT and 35.99% with UniXcoder on hover; kNM-LM improves all projects and shows its largest gains on larger projects.The figure orders projects from large to small by size.Source: Figure 3 · p. 6, C. Results · pp. 6–7
- For punctuation, kNM-LM improves base-model accuracy by 17.66%–21.71%; identifier accuracy gains are 9.31% and 8.90% in large projects, and 6.72% and 4.14% in small projects.Identifier and literal completion are project-specific, while punctuation relates to programming-language syntax.Source: C. Results · pp. 6–7, Table II · p. 7
VI. RQ2: INTRA-SCENARIO CODE COMPLETION
The intra-scenario study builds a database from projects using the same third-party libraries, whose code contains scenario-specific API usage patterns. It evaluates token completion on API-Bench and separately tests completion of Android API-call lines.
A. Study Design
- Projects using the same third-party libraries supply the scenario database, capturing scenario-specific knowledge such as API usage patterns.Source: A. Study Design · p. 7
B. Dataset
- API-Bench covers four Java and Python scenarios; each scenario’s train set builds the datastore and its test set supports token-completion evaluation.Source: B. Dataset · p. 7
- The separate Android API-call test set contains 49,344 lines, completed using the code context before each call; it reuses the Android dataset’s database.The test lines include the API name and input parameters.Source: B. Dataset · p. 7
C. Results
Across intra-scenario tasks, kNM-LM improves token completion over base language models, and it performs strongly on lines containing Android APIs. The discussion also compares its resource use and hyper-parameter sensitivity with kNN-LM.
C. Results
- kNM-LM improves intra-scenario token completion by approximately 6–12% over the base LMs.Table IV covers token completion accuracy on the API-Bench intra-scenario dataset.Source: C. Results · p. 7, Table IV · p. 8
- Bayesian inference performs best in most ablation cases, while manually setting λ has the least favorable effect.The comparison includes kNM-LM without Bayesian inference and without new observations.Source: C. Results · p. 8, Table IV · p. 8
- On Android API lines, kNM-LM improves EM by 7.75% and 8.94%, with 4,245 more correctly predicted lines than all other models combined.The Venn diagram examines correctly completed Android-API lines.Source: C. Results · p. 8, Table V · p. 8, Figure 4 · p. 9
VII. DISCUSSION
- The discussion assesses kNM-LM through comparisons with kNN-LM and fine-tuned models, plus case studies.The case studies provide concrete examples of the approach.Source: §VII · p. 8
A. Compared with kNN-LM.
- kNM-LM saves 284.83MB and 389.36MB in large projects and 55.57MB and 67.61MB in small projects versus kNN-LM.It is less sensitive to hyper-parameter settings; even the worst-case completion of 937.06 words takes under one second.Source: Table VI · p. 8, A. Compared with kNN-LM · p. 9
B. Compared with Fine-tuning.
The paper compares kNM-LM with fine-tuning, presents qualitative Android API completion examples, and describes threats involving model selection and evaluation settings.
B. Compared with Fine-tuning.
- Fine-tuning is more effective than kNM-LM on intra-scenario completion, though applying kNM-LM after fine-tuning adds improvement.The comparison is presented in Fig. 6.Source: B. Compared with Fine-tuning · p. 9, Figure 6 · p. 10
C. Qualitative Analysis.
- Examples show kNM-LM completing complex Android API names and parameters, while other models make API or irrelevant-code errors.A further example shows kNM-LM ignoring irrelevant code when completing the intended code.Source: C. Qualitative Analysis · p. 9, Figure 7 · p. 10
VIII. THREATS TO VALIDITY
- The evaluation uses CodeGPT and UniXcoder rather than black-box systems for the fine-tuning comparison, though GPT-3 API use is described.The paper notes that its 80%-written commit point may not reflect practical scenarios and plans to study database-size effects.Source: §VIII · pp. 9–10
- The intra-project database and test set come from the same project, split by commit history; ReACC and Hybrid were not verified on Java.The authors use BM25 as a comparison for Java, and note that the 80%-written commit point may not reflect practice.Source: §VIII · p. 10
IX. RELATED WORK
The related-work discussion distinguishes code-completion categories and retrieval-augmented language models, then positions kNM-LM as a domain-adaptive approach. The conclusion describes its decoupled datastore and Bayesian combination strategy.
A. Code Completion
- Code completion includes API-name, variable-name, and arbitrary-token tasks; this approach completes arbitrary tokens in code sequences.The paper distinguishes these task types by the code being completed and the techniques used.Source: A. Code Completion · pp. 10–11
- Domain-adaptive completion switches databases across domains and does not require access to LM parameters, enabling use with black-box LLMs such as GPT-4.The approach is presented as generic across domains.Source: A. Code Completion · p. 11
B. Retrieval-augment Language Model
- Token R-LMs retrieve tokens during generation, but can incur high storage costs and require hyper-parameter selection.The paper uses the token R-LM architecture.Source: B. Retrieval-augment Language Model · p. 11
- Unlike other approaches described, kNM-LM adds no trained network module; datastore decoupling reduces storage cost and Bayesian inference selects hyper-parameters.The paper contrasts its design with trained weighting, retrieval gating, and contextual-graph approaches.Source: B. Retrieval-augment Language Model · p. 11
X. CONCLUSION
- The conclusion reports best results on intra-project and intra-scenario completion, with a decoupled database storing tokens the LM cannot predict correctly.The approach combines database and LM results using Bayesian inference and does not require LM weights or additional neural modules.Source: §X · p. 11