Syntax-Aware Retrieval Augmented Code Generation
Authors: Xiangyu Zhang
Yu Zhou
Guang Yang
Taolue Chen
Venue: EMNLP 2023 Findings (2023)
Keywords: Code Generation
Retrieval Augmented Generation
Neural-Symbolic
TL;DR: A novel retrieval augmented code generation approach.

Abstract: Neural code generation models are nowadays widely adopted to generate code from natural language descriptions automatically. Recently,  pre-trained neural models equipped with token-level retrieval capabilities have exhibited great potentials in neural machine translation. However, applying them directly to code generation experience challenges: the use of the retrieval-based mechanism inevitably introduces extraneous noise to the generation process, resulting in even syntactically incorrect code. Computationally, such models necessitate frequent searches of the cached datastore, which turns out to be time-consuming. To address these issues, we propose $k$NN-TRANX, a token-level retrieval augmented code generation method. $k$NN-TRANX allows for searches in smaller datastores tailored for the code generation task. It leverages syntax constraints for the retrieval of datastores, which reduces the impact of retrieve noise. We evaluate $k$NN-TRANX on two public datasets and the experimental results confirm the effectiveness of our approach.

Syntax-Aware Retrieval Augmented Code Generation
Xiangyu Zhang' Yu Zhou' Guang Yang' Taolue Chen 2+
Nanjing University of Aeronautics and Astronautics
2 Birkbeck, University of London
Sahangx1angyyu zhouyu, gunngguanggye Jenuaa. edu.
t. chen@bbk. uk
Abstract
NL
Check if object obj is a string.
ANN-MT MT
all (aisinstance isinstance str)
Neural code generation models are nowadays widely adopted to generate code from natural language descriptions automatically. Recently pre-trained neural models equipped with token- level retrieval capabilities have exhibited great potentials in neural machine translation. How- ever, applying them directly to code generation experience challenges: the use of the retrieval- based mechanism inevitably introduces extra- neous noise to the generation process, result- ing in even syntactically incorrect code. Com- putationally, such models necessitate frequent searches of the cached datastore, which turns out to be ttme-consuming To address these is- sues, we propose ANN-TRANX, a token-level retrieval augmented code generation method. ANN-TRANX allows for searches in smaller datastores tailored for the code generation task. It leverages syntax constraints for the retrieval of datastores, which reduces the impact of re- trieve noise. We evaluate ANN-TRANX on public datasets and the experimental re- sults confirm the effectiveness of our approach.
The correct one
Table I: An example code generated by ANN-MT.
the models with a token-level retriever. In par- ticular, this neurall-errieval--n-te--ooo approach facilitates the integration of external knowledge into the pre-trained model and provides a simple yet effective method to update the model by switch- ing the retrieval datastore, without fine-tuning the model parameters.
Can such neuraa-retrievva--n-tt------ approach benefit neural code generation? Our preliminary experiments reveal three main issues (cf. the exam- ple in Table 1) if it is adopted outright. Firstly, the model performance may be negatively affected by the noise in the retrieved knowledge. For example, "all" does not match the intention of the descrip- tion, but it is recognized as the target token by the retriever, resulting in the generation of incorrect code. Secondly, the code generated by kNN-MT cannot guarantee syntactic correctness, as demon- strated by the mismatching parentheses in the given example. Thirdly, the token-level retrieval method requires similarity search of the entire datastore at each time step of inference, which hinders the deployment of such approach.
1 Introduction
Neural code generation aims to map the input natu- ral language (NL) to code snippets using deep learn- ing. Due to its great potential to streamline soft- ware development, it has gamnered significant atten- tions from both natural language processing and software engineering communitics. Various meth- ods have been explored to facilitate code genera- tion (Yin and Neubig, 2018; Wang et al., 2021; 1; Guo et al., 2022). Recent progress in neural machine translation (NMT) shows that the non- non-parametric k- neirest- machine translation (NNN-MT) approach may significantly boost the performance of standard NMT models (Khandelwal et al., 2021) and other text generation models (Kassner and Schütze, 2020; Shuster et al., 2021) by equipping
In this paper, we propose a novel code gener- ation approach, i.e. kNN-TRANX, to overcome the limitations of the neuraa--erieeall-in--elloooo paradigm for code generation tasks. The basic idea is to integrate symbolic knowledge to ensure the quality of the generated code and expedite the re- trieval process. To achieve this, leverage the sequence-to-te (seq2tree) model to generate ab- stract syntax tree (AST), which is a hierarchical tree-like structure used to represent the code, rather than generate target code snippet directly. This en- ables us to use AST construction rules to guarantee
"Corresponding author.

---

the syntactic correctness of the generated code and filter out retrieval noise.
target token. Formally, given a bilingual sentence pair (a, y) from the training corpus (X,)), pre- trained NMT model SNMr(() generates the i-th context representation h; fNATT( (z, V<i), then the datastore D is constructed as follows.
We design kNN-TRANX as a two-step process (cf. Figure 2). In the first step, we construct two separated datastores, i.e., the syntactic datastore and the semantic datastore, based on the type of AST nodes. This allows to determine the type of the next node be predicted according to the grammar rules and query a specific datastore. In the second step, we utilize syntactic rules to filter out irrelevant knowledge and convert the similarity retrieval results of the current target token into a probability distribution, i.c., the ANN probability. This probability, together with probability from the neural network, yields the probability of the ac- tion be used for AST generation via learnable confidence parameter. It can help minimize re- trieval noise and dynamically exploit combinations the two probabilities, resulting in improved code generation performance.
D (K, V)
U {(hs, (i), Vy (x,))((X)))
Model Inferring. During inference, time step given the already generated token der and the contextual information hj, the ANN-MT model gen- erates yi by retrieving the datastore, which can be calculated as
d T
PANN x Ves) & Σ 1 11== exp
(hz:85)
(1)
where T is the temperature and dj indicates the l2 distance between query hę and the retrieved key hj.
To evaluate the effectiveness of kNN-TRANX, we perform experiments on two publicly avail- able code generation datasets (i.e., CoNaLa and Django). The experimental results show 27.6% improvement in the exact match metric on the CoNaLa dataset and a 4.2% improvement in the BLEU metric on the Django dataset, surpassing five state-of-the-art models under comparison. Ad- ditionally, we conduct an experiment on model canonical incremental adaptation, which updates ANN-TRANX by switching the datastore. The experimental results demonstrate that our model can achieve performance comparable fully fine- tuned models and reduce the trainable parameters by 7,000 times.
2.2 Seq2tree Model
The purpose of the seq2tree code generation mod- els is to generate ASTs instead of directly out- putting code snippets. Compared to the sequence- to-sequence (seq2seq) models, the seq2tree models ensure the syntactic correctness of the generated code. Among the seq2tree models, BertranX (Beau and Crabbé, 2022) was recently proposed and rep- resented the state-of-the-art architecture. BertranX employs BERT to process the input natural lan- guage and features a grammar-based decoder.
expr BinOp ( expr left, operator op, expr right)
Call (expr func, expr* args, keyword* keywords)
Constant ( coontant value)
2
Background
Name Identifler Id, expr_context ctx)
In this section, we provide an overview of the kNN- MT paradigm and the seq2tree model.
Figure I: Example of ASDL. for Python. ASDI defines a set of grammatical symbols, which are denoted in orange and distinguished by a unique constructor name highlighted in blue. Each rule assigns names to its fields or the symbols marked in black. The grammatical symbols be classified into two types: nonterminals (e.g., expr) and terminals or primitives (e.g., identifier), Some the grammatical symbols may have qualifiers (*) that allow for zero more iterations of the symbol.
2.1 ANN-MT
The &NN-MT paradigm (Khandelwal et al., 2021) is a translation mechanism that enhances the qual- ity of model generation by incorporating an ad- ditional translation retriever. This allows NMT models to benefit from the retrieved knowledge. The paradigm comprises two main parts, namely, datastore building and model inferring.
BertranX describes ASTs using sequences of actions based on ASDL (Wang et al., 1997) gram- mar, which gives concise notations for describing the abstract syntax of programming languages (cf. Figure as an example). With ASDL, BertranX
Datastore Building. The datastore consists of a set of kaluv pairs, where the key is the decoder hidden state and the value is the corresponding

---

defines two distinct types of actions that generate ASTs, i.e., PREDRULE and GENERATE. The first type is used for initiating the generation a new node from its parent node, which we mark as syn- tactic nodes in this paper; the second type on the other hand, is used to produce terminal or primitive symbols that we mark semantic nodes.
As result, two separated symbolic datastores can be constructed based on the various types of target actions within the training set. Constructing datastores in this manner is more effective than storing both types of actions in a single datastore since it helps reduce noise during retrieval. More- over, the subsequent token type can be determined based on grammar rules, facilitating the retrieval of a specific datastore and accelerating the retrieval process.
3 ANN-TRANX
The workflow of kNN-TRANX is depicted in Fig- ure 2. It consists of two main components: datas- tore building and model inferring.
3.2 Model Inferring
The process of model inference can be divided into three main phases, as shown in Figure 2. First, the code fragment x is put into the trained model to generate the context representation h2 and com- pute the neural network distribution (PNN). Then, we query the datastore using this representation to obtain the k- neighbr distribution (pENN). Finally, we combine these two distributions to pre- dict the target token. In the subsequent sections, we will discuss three pivotal components of ÀNN- TRANX: syntax-constrainet token-level retrieval, meta-k network, and confidence network.
3.1 Datastore Building
Given pre-trained seq2tree model and the training dataset, we first parse code snippets to ASTs and generate all instances in the training corpus. This process allows us to capture and store the decoder representations along with their corresponding tar- get tokens as key-value pairs. The actions that con- stitute AST can be categorized into two groups: rules and primitives. These categories align with the actions of GENERATE and PREDRULE, re- spectively. As shown in Figure 3, the two types of nodes have significant differences in terms of type and quantity. Combining them into the same datastore could potentially reduce retrieval accu- racy. Therefore, we employ separated datastores for each node type. referred to as the syntactic and semantic datastores respectively. Nodes represent- ing the structural information (e.g., Expr and Call) are put into the syntactic datastore, while nodes representing the semantic information (c.g., text and split) the code are put into the semantic one.
token-level retrieval. Given the current context representation h4 generated by the model, first calculate the l2 distance dj 12(h,, h2) between the context representa- tion h; and each neighbor (h3, 2,) in the datas- tore to determine the k nearest neighbors. Previ- ous studies (Meng et al., 2022; Dai et al., 2023) have restricted the search space based on poten- tial input-output patterns to improve decoding ef- ficiency and reduce the impact of noise. However, the restricted search space may also exclude some valuable knowledge.
Given an NL-code pair (Ir, y) from the training corpus (X', "Y), we first transform the code snippets y into AST representations Z. Next, we calculate the i-th context representation h4 fe(æ, 2<3)), where fo(:) refers to the trained seq2tree model and Z. The datastore is constructed by taking ht's as keys and 2,'s 's as values. Namely,
To mitigate this problem, our approach features syntax-aware retrieval capability. In contrast conventional seq2seq models, our model aims generate ASTs that allow incorporate symbolic knowledge and determine the retrieved tokens by means of syntactic rules. During the retrieval pro- cess, we can determine the type of the subsequent token based on the tokens already produced. If the next token is expected to represent the syntactic information, we just retrieve the syntactic datastore to accelerate the retrieval process, and vice versa. Additionally, we can also use the ASDL rules exclude illegitimate tokens to reduce the amount of irrelevant information. For example, as seen in Figure 2, our model has already generated the
D(era)= K, v(srra)
U {1h,, 2:) I 2, E 2 & 21 E rules), (2,2)
and
D(pri)_ = (K, vlori)
U ((14) 2i E x, p
(2,))

---

Regrementation
as features and - felt,aa)
Sunductic
NL: Split String text by space.
Trained M odel
in rules
p(,),.P,)) ApPax AdP
Code: teat.spitt
Og. memotic dittattoee
Cill sibs
Abstract Symtax True
Transition Systemm PREDAULEEExpr
Call
test
ASOL Grammar
Confidencee Network
Apt-
GENERATEESphn
spit
ASDL Grommar Rules
Bx07
cko
prediction
Results
Results
Ret (Ber,
Representation A fe(t
10
Call
NL: Split String text by space.
Network
Infemiin
Figure 2: Workflow of kNN-TRANX
the range is S 10, 2 (K) K]). To evaluate the weight of each of the values, we use distance dj and the count of distinct values in top í neighbors obtain a normalized weight by Pg(k) = softmax "e([di,..,,dkee). di.................. ck])) where Sp(.) denotes the Meta-k Network. The prediction of KNN can be obtained by PANN (2/ &ci)
(a) Syntactic tokens
(b) Semantic tokens
Figure 3: I-SNE visualization of CoNaLa data features,
with dark colored dots indicating more frequently oc- curring vocabulary.
Σ Pa (ki,) PK.NN (=)a0, where Pk,NN K.ES
indicates the ha-ncarest-neighhbee prediction results calculated as Equation (1). In this way, the &NN model can expand the search space while reducing the impact of the retrieval noise.
node Expr in the previous time step. should be noticed that alias nodes according to the distances. However, the child nodes of Expr do not support alias in the ASDL grammar. In this way, we filter out these nodes from the search results to reduce noise and avoid excluding valuable information.
have retrieved Call and
Confidence network. In order to utilize the knowl- edge of the symbolic datastore while maintaining the generalization ability of the neural network, we combine two probability distributions by means of weighting. Previous studies have integrated the distributions of &NN and NN by using fixed or learnable parameter À to measure their respec- tive weights. Khandelwal et al. (2021) combine the probability distributions using fixed weights, but this approach fails to dynamically adjust the weights based on the distance retrieved. Zhu et al. (2023) adjust the weights based the retrieval distance, but they overlook the confidence of the neural network output. Khandelwal al. (2021) utilize the probability of the neural network and retrieval distance as features to dynamically select the value of À which consider the confidence
Meta-k network. We retrieve relevant pieces of knowledge from the datastore, and then map the distances between the query vector and cached representation as probabilities. Empirically, the number of retrievals, k. is crucial in our model because too few retrievals may result in valuable in- formation being ignored, while too many retrievals may introduce noise. To alleviate this problem, employ the meta-k network (Zheng al., 2021) to dynamically evaluate the weight of the retrieved knowledge. Meta-k network considers range values that are smaller than the upper bound K, instead of using a fixed value of k. Typically

---

the two distributions, but this approach neglects the correlation between the two distributions.
To expedite the implementation. leverage ANN- box (Zhu et al., 2023), an open-source toolkit for building ANN-MT, to implement kNN-TRANX. As explained in Section 3, ANN-TRANX creates two datastores. Due to the considerable difference in vocabulary sizes between the two datastores, we construct separate settings for the syntactic and semantic datastores. For the syntactic datastore, set the upper limit Kyule to 4 to account for its limited token variety. For the semantic datastore, we set Kers to 64. To train the meta-k network and confidence network, we employ the Adam optimizer with a learning rate 3e-4. To accelerate the datastore retrieval process, we incorporate the FAISS library (Johnson et al., 2019) for similarity retrieval. All experiments are performed on single NVIDIA 2080Ti.
To address this issue, we propose a confidence network that estimates the confidence of both probabilities based on ANN and NN distributions. In addition, we incorporate the weights of each k-value as features into the confidence network, ensuring that the model is aware of the number of tokens that require attention. As such our model can capture the relationship between the two distributions. In cases where the two distributions conflict, we assign a higher weight to the distribution with higher confidence. The confidence A is calculated from W as A S S(W[PANN (z: zz & PNN (zila, & (p)))) where S denotes the sigmoid activation function.
The final distribution at prediction 2; calcu- lated as a weighted sum of two distributions with X i.e., p (z 2es) = A:PANN + (1 >;) PNN-
Baselines. We compare kNN-TRANX against five state-of-the-art code generation models.
TRANN (Yin and Neubig, 2018) is a seq2tree model consisting of bidirectional LSTM en- coder for learning the semantic representation and a decoder for outputting a sequence of ac- tions for constructing the tree.
4 Experiments
In this section, we first introduce the datasets and evaluation metrics. Then, conduct a compre- hensive study and analysis on code generation and model canonical incremental adaptation.
Reranker (Yin and Neubig, 2019) reorders a set of N-best candidates to improve the quality of the generated results.
4.1 Datasets and evaluation metrics
Ext-codegen (Xu et al., 2020) incorporates API documentation as external knowledge into the model, thus enabling data augmentation.
We evaluate &NN-TRANX on two code genera- tion datasets, namely, CoNaLa dataset (Yin et al.,
2018) and Django dataset (Oda et al., 2015). The
TAE (Norouzi et al., 2021) uses BERT and a transformer decoder to auto-encoding monolin- gual data.
CoNaLa dataset comprises 600k NL-code pairs col- lected from StackOverflow, out of which 2,879 NL
were rewritten by developers. This dataset con-
BertranX (Beau and Crabbé, 2022) uses BERT encoder and serves as the base model for our &NN-TRANX.
tains questions that programmers encounter in their
real-world projects. On the other hand, the Django
dataset consists of 18,805 examples, where each example consists of one line of Python code accom- panied by corresponding comments. Compared to CoNaLa, approximately 70% of the examples in Django are simple tasks that include variable as- signment, method definition, and exception han- dling, easily inferred from the corresponding NL predictions. We employed BLEU (Papineni et al., 2002), CodeBLEU (Ren et al., 2020), and exact match (EM) metrics assess the performance of our experiments.
REDCODER (Parvez al., 2021) retrieves rele- vant code from retrieval database and provides them as a supplement to code generation models. CodeT5 (Wang et al., 2021) builds on the sim- ilar architecture of T5 (Raffel et al., 2020) but incorporates code-specific knowledge to endow the model with better code understanding.
Main results.
The experimental results are presented in Table 2. Our proposed performance over BertranX on the CoNaLa dataset by 3.11 BLEU (9.1%), 2.95 CodeBLEU (8.2%). and 1.6 EM (27.6%). On the Django dataset, we observed improvements 3.34 BLEU (4.2%), 2.81 Code (3.6%), and 2.39 EM (3.0%). These re-
exhibits superior
4.2 Code Generation
Implementation details. We use BertranX as the base seq2tree model for our experiments, which is trained on annotated data and 100k mined data.

---

CoNaLa CodeBLEU
Django CodeBLEU
Model
BLEU
EM
BLEU
EM
TRANX
28.11
29.01
2.5
62.31
64.35
73.70
Reranker
30.11
28.98
2.8
73.26
71.29
80.18
Ext-codegen
32.26
33.23
3.0
TAE
33.41
32.87
3.4
68.39
70.27
81.03
BertranX
34.18
36.09
5.8
79.86
78.85
79.77
REDCODER
35.12
36.82
6.4
75.36
73.01
78.82
CodeT5
36.28
35.34
6.8
76.50
71.92
80.93
ANN-TRANX
37.29
39.04
7.4
83.20
81.66
82.16
Table 2: Comparative results of models trained the CoNaLa and Django test datasets.
same datastore. The result shows that this approach significantly reduces the quality of the generated code by introducing a substantial amount of noise. Morcover, we analyze the effect of Krule and K Kpri on the experimental results, presented in Table 4. The results align with our conjecture that retrieving small number of the syntactic nearest neighbors and a relatively large number of semantic entries leads better code generation. Furthermore, Fig- ure shows that increasing the size of datastore can improve the quality of generated code. Ad- ditionally, Figure 5(a) and Figure 5(d) depict that even a small amount of syntactic knowledge stored can achieve a high quality of code generation. In contrast, the quality of generated code keeps im- proving as the size of semantic datastore increases. We believe this is because semantic knowledge is relatively scarce compared to syntactic knowledge, which can be demonstrated by Figure 3.
Syntax Match
0.23
0.26
TRANX
Reranker
Brh-codlegon
TAE
BertranX
REDCODER
CodeTS
Our
Figure 4: The syntax match score of the generated code the CoNaLa dataset.
sults indicate that token-level retrieval improve the performance of code generation models, and the improvement is more evident for challenging tasks. To demonstrate that our model can gener- ate code with more accurate syntax, we provide the syntax match score (Ren et al., 2020) in Fig- ure 4, which reflects the degree of syntax matching in code. The results indicate that our model out- performs the baseline model in terms of syntax accuracy.
4.3 Model Canonical Incremental Adaptation
Implementation details. Although the recent pro- posed large-scale models such Codex (Chen et al., 2021) and GPT-4 (OpenAl, 2023) have demonstrated its powerful code generation capabili- ties, one the challenges that the trained models are difficult to update. Models need to be continu- ously trained via incremental learning, which con- sumes huge computing resources. To make matters worse, incremental learning with new data can lead to catastrophic forgetting problems (Li and Hoiem, 2016). To address this, validate model using incremental learning by only updating the datas- tore without adapting the model parameters. We use BertranXt as our base model to simulate the
Analysis. We conduct an ablation study of kNN- TRANX as shown Table The experimental results demonstrate that retrieval filtering method can significantly enhance the performance of the code generation models. For the method that com- bines NN distribution and &NN distribution, we compared the method of measuring retrieval dis- tance proposed in adaptive ANN-box (Zhu al., 2023). Experimental results show that our ap- proach of considering both distributions compre- hensively achieves better results. also evaluate the effect of placing both types of action in the
BertranXt is trained on annotated data and Sk mined data on CoNaLa dataset.

---

CoNaLa CodeBLEU
Django CodeBLEU
Model
BLEU
EM
BLEU
EM
BertranX
34.18 18
36.09
5.8
79.86
78.85
79.77
ANN-TRANX
37.29
39.04
7.4
83.20
81.66
82.16
w/o noise filtering
36.51
38.23
7.0
81.32
79.92
81.54
w/o Confidence Network
36.89
37.82
5.8
81.23
79.15
81.98
w/o separate datastore
35.62
37.21
6.2
80.74
79.23
80.08
Table 3: Ablation study of different strategies and networks on CoNaLa dataset and Django dataset.
Kyude
K Kpri
BLEU
CodeBLEU
EM
periment. Figure 5 confirms that providing small amount of syntactic knowledge can improve the performance the model. Therefore, we consider two ways to update the datastores in our experi- ments, i.e., updating both datastores and updating the semantic datastore only.
36.05
37.74
6.4
1
16
36.33
37.46
7.0
1
64
36.51
37.69
7.4
1
128
36.63
37.84
7.2
4
35.76
37.59
6.4
Main results. We adopte the same evaluation met- rics code generation. As shown in Table 5, firstly, without using additional datastore, ANN- TRANXI can outperform BertranXt, As the knowl- edge the datastore is constantly updated, we can see that kNN-TRANX has improved on three evaluation criteria. terms of BLEU evaluation, ANN-TRANXi with 95k external data can achieve performance comparable to that of training-basee BertranX. Furthermore, we update only the seman- tic datastore, which can also be effective. We also provide examples demonstrate how k- nearest-neighbor retrieval assists in model decision- making in Appendix A. 1. should be noted that the CoNaLa dataset was obtained through data min- ing, and the majority of the NL-code pairs obtained through mining are irrelevant, which greatly adds noise to our retrieval. Therefore, we believe that ENN-TRANX perform better more reliable data sources through incremental learning.
4
16
37.09
38.72
7.0
4
64
37.29
39.04
7.4
4
128
37.06
38.80
7.2
16
36.02
37.81
6.4
16
16
37.01
38.70
7.2
16
64
36.89
38.76
7.4
16
128
37.34
38.99
7.2
Table 4: We studied the impact of different K values on the generation results on CoNaLa dataset. Krate {1, 4, 16] and Kpri (1, 64, 128).
scenario for incremental learning. Then, we update the datastore using (0k, 10k, 50k, 95k) mined data in CoNaLa dataset respectively. In our experiment, we update the datastore to perform efficient fine- tuming. Compared to the method of fine-tuning all parameters (which requires training 122,205k parameters), our method only needs to train 17k, greatly reducing the GPU memory consumption required for training, and achieving comparable results to fine-tuning all parameters.
Related Work
Code generation. Code generation aspires gen- erate target code through natural language to im- prove programmers' development efficiency. Yin and Neubig (2018) propose TRANX to generate ASTs instead of generating code snippets directly. Based on TRANX, Beau and Crabbé (2022) pro- pose BertranX relying BERT encoder and grammar-based decoder. Poesia al. (2022) pro- pose framework for substantially improving the reliability of pre-trained models for code genera- tion. Chen et al. (2022) propose CodeT to lever- age pre-trained language models to generate both
is worth noting that, compared to the previous KNN generative model, our model includes two datastores for syntactic and semantic information, respectively. There are 109 kinds of token in the syntactic datastore, and the number of correspond- ing datastore entries 1,603k. However, the types of token in the semantic datastore may be infinite. depending on the actual defined vocabulary, so the knowledge corresponding to each vocabulary is relatively scarce, with only 518k entries in this ex-

---

CoNaLa dataset
40
46
39
31
34
18
18
37
37
56
16
0L.0
mil 01EU
25
100%
075
200
(a) Ratio of Syntactic Datastore
(b) Ratio of Semantic Datastore
(c) Ratio of both Datastores
Django dataset
60
84
64
85
33
85
80
32
82
81
31
11
50
80
10
77
76
1U8U
79
75
76
2006
4.0%
600
1.09
09
2006
4000
100%
2006
4.0%
(d) Ratio of Syntactic Datastore
(e) Ratio of Semantic Datastore
(1) Ratio of both Datastores
Figure 5: The impact of datastore size on the BLEU and CodeBLEU scores of CoNaLa a dataset and Django dataset. We used three strategies to reduce the datastore size, which are reducing the storage of syntactic knowledge, semantic knowledge, and both types of knowledge.
the code snippets along with test cases, and select the best code based on the number of test cases passed. Besides, researchers used pre-training methods to incorporate more external knowledge into the model, effectively improving its perfor- mance downstream tasks (Feng et al., 2020; Wang et al., 2021; Ahmad al., 2021). Recently, large-scale pre-training models have demonstrated remarkable capabilities in code generation (Li et al., 2023; Wang et al., 2023; Touvron al., 2023).
ing a datastore using a decoder representation as a key and using the corresponding target charac- ter as a value. The generation done by retriev- ing the top h neighbors as the result. Based on this paradigm, number of optimization methods have also been proposed. Zheng et al. (2021) meta-k network to dynamically adjust the retrieval weights. Jiang al. (2022) improve the robustness of i-N--TT in terms of both model structure and training methods. Meng et al. (2022) propose Fast kNN-MT, which improves retrieval efficiency by constructing multiple smaller datastores.
Retrieval-augmented models. Retrieving and in- legrating auxiliary sentences has shown effective- ness in enhancing the generation of class models. Farajian et al. (2017) propose retrieving similar sentences from the training set to adapt to differ- ent inputs. The work of Hayati et al. (2018) is built on the neural network model of AST driver and generates code by searching action subtrees. Parvez al. (2021) propose using retrieved natu- ral language code to improve the performance of code generation and code translation models using REDCODER.
Conclusion
In this paper, we propose ENN-TRANX. By provid- ing syntactic and semantic datastores for seq2tree model, we are able to outperform the baselines. In addition, we provide more knowledge for the model by switch the datastores without fine-tuning the neural network. Experimental results show that ÀNN-TRANX exhibits competitive performance against learning- methods through incremen- tal learning. the future, we plan to construct a smaller and more finennraiee syntactic datas-
Recently, Khandelwal et al. (2021) propose ANN-MT, non- parametric paradigm for construct-

---

Only Semanties
External Data
Training Cost
Extra
Model
BLEU
CodeBLEU
EM
(FLOPs)
Storage
BertranXi
29.22
29.89
4.0
1.09 10
ANN-TRANXI
0
32.25
32.08
4.2
1.11
10
0.2G
10k
32.95
32.87
4.4
1.14
10
0.5G
ÀNN-TRANXI
False
50k
33.71
33.58
4.4
1.23
10
1.3G
95k
34.04
34.69
4.8
1047 1.35
2.2G
10k
32.28
32.14
4.2
1.14 1
0.3G
ANN-TRANXI
True
50k
33.25
33.87
4.2
1.23
10
0.5G
95k
33.89
34.43
4.6
1.35
10
0.8G
BertranX
34.18
36.09
5.8
1.55
Table The results of model canonical incremental adaptation. BertranX/ is trained on cleaned 2k and mined data. KNN-TRANXÍ is built the same size data datastore top of BertranXt, When Only Semantics is set to False, both datastores are updated simultaneously. while True means only semantic datastore is updated. External data refers the quantity of training data updating the datastore.
tore to reduce the search space of the model and accelcrate model inference.
References
Wasi Ahmad, Saikat Chakraborty, Baishakhi Ray, and Kai-Wei Chang. 2021. Unified pre-training for pro- gram understanding and generation. In Proceedings of the 2021 Conference of the North American Chap- ter of the Association for Computational Linguistics:
Limitations
Although our proposed approach can enhance the generalizability of the seq2tree model and enable rapid updates by switching datastores, incorporat- ing extra datastores necessitates a similarity search at each time step of inference. Even though only one of the two datastores needs to be retrieved, the inference time for the model may still increase considerably (ef. Appendix A.2). Furthermore, the incorporation of additional datastores will con- sume storage space and may escalate with growing training data. Although the previous work has pro- posed approuches to reduce the content of datastore through pruning. this approach also leads to model performance degradation.
Human Language Technologies, pages 2655-2668, Online. Association for Computational Linguistics.
Nathanaël Beau and Benoit Crabbé. 2022. The impact of lexical and grammatical processing on generating code from natural language. In Findings of the As- sociation for Computational Linguistics: ACL 2022, pages 2204-2214, Dublin, Ireland. Association for Computational Linguistics.
Bei Chen, Fengji Zhang, Anh Nguyen, Daoguang Zan, Zeqi Lin, Jian-Guang Lou, and Weizhu Chen. 2022. Codet: Code generation with generated tests. CoRR. abs/2207.10397.
Mark Chen, Jenry Tworek, Heewoo Jun, Qiming Yuan, Henrique Pondé de Oliveira Pinto, Jared Kaplan, Harrison Edwards. Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sas- try, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cum- mings, Matthias Plappert, Fotios Chantzis, Eliza- beth Barnes, Ariel Herbert- Voss, William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor Babuschkin. Suchir Balaji. Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leike, Joshua Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder. Bob McGrew, Dario Amodei, Sam McCandlish, Ilya
Acknowledgements
This work supported by the National Natural Science Foundation of China (No. 61972197, No. 62372232), the Natural Science Foundation of Jiangsu Province (No. BK20201292), and the Col- laborative Innovation Center of Novel Software Technology and Industrialization. T. Chen is par- tially supported by overseas grant from the State Key Laboratory of Novel Software Technology, Nanjing University (KFKT2022A03) Birkbeck BEI School Project (EFFECT) and National Natu- ral Science Foundation of China (No. 62272397).

---

Sutskever, and Wojeiech Zaremba. 2021. Evaluat- ing large language models trained on code. CoRR, abs/2 107.03343.
Raymond Li, Loubna Ben Allal, Yangtian Zi, Niklas Muennighoff, Denis Kocetkov, Chenghao Mou, Marc Marone, Christopher Akiki, Jia Li, Jenny Chim, et al. 2023. Starcoder: may the source be with you! arXiv preprint arXiv:2305.0616..
Yu-Hsiu Dai, Zhirui Zhang. Qiuzhi Liu, Qu Cui, Wei- Hong Li, Yichao Du, and Tong Xu. 2023. Simple and scalable nearest neighbor machine translation. ArXiv, abs/2302. 12188.
Zhizhong Li and Derek Hoiem. 2016. Learning without forgetting. In Computer Vision ECCV 2016 14th European Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part IV, volume 9908 of Lecture Notes in Computer Science, pages 614-629. Springer.
M. Amin Farajian, Marco Turchi, Matteo Negri, and Marcello Federico. 2017. Multi-domain neural ma- chine translation through unsupervised adaptation. In Proceedings of the Second Conference on Machine Translation, pages 127-137, Copenhagen, Denmark. Association for Computational Linguistics.
Yuxian Meng, Xiaoya Li, Xiayu Zheng, Fei Wu, Xi- aofei Sun, Tianwei Zhang, and Jiwei Li. 2022. Fast nearest neighbor machine translation. In Findings of the Association for Computational Linguistics: ACI. 2022, pages 555 255-555, Dublin, Ireland. Association for Computational Linguistics.
Zhangyin Feng. Daya Guo, Duyu Tang. Nan Duan, Xi- aocheng Feng, Ming Gong, Linjun Shou, Bing Qin, Ting Liu, Daxin Jiang, and Ming Zhou. 2020. Code- BERT: A pre-trained model for programming and natural languages. In Findings of the Association for Computational Linguistics: EMNLP 2020, pages 1536-1547, Online. Association for Computational Linguistics.
Sajad Norouzi, Keyi Tang, and Yanshuai Cao. 2021. Code generation from natural language with less prior knowledge and more monolingual data. In Proceed- ings of the 59th Annual Meeting of the Association for Computational Linguistics and the 1 11th International Joint Conference on Natural Language Processing ( (Volume 2: Short Pupers), pages 776-785, Online. Association for Computational Linguisties.
Daya Guo, Shuai Lu, Nan Duan, Yanlin Wang, Ming Zhou, and Jian Yin. 2022. UniXcoder: Unified cross- modal pre- pre-training for code representation. In Pro- ceedings of the 60th Annual Meeting of the Associa- tion for Computational Linguistics (Volume I: Long Papers), pages 7212-7225. Dublin, Ireland. Associa- tion for Computational Linguistics.
Yusuke Oda Hiroyuki Fudaba, Graham Neubig, Hideaki Hata, Sakriani Sakti, Tomoki Toda, and Satoshi Nakamura. 2015. Learning to generate pseudo-code from source code using statistical ma- chine translation. In 2015 30th IEEF/ACM Interna- tional Conference on Automated Software Engineer- ing (ASE), pages 574-584.
Shirley Anugrah Hayati, Raphael Olivier, Pravalika Av- varu, Pengcheng Yin, Anthony Tomasic, and Graham neural code gener-
Neubig. 2018 Retrieval-hasee ation. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pages 925-930, Brussels, Belgium. Association for Computational Linguistics.
OpenAL. 2023. Gpt-4 technical report. ArXiv, abs/2303.08774.
Kishore Papineni, Salim Roukos, Todd Ward, and Wei- Jing Zhu. 2002. Bleu: a method for automatic evalu- ation of machine translation. In Proceedings of the 40th Annual Meeting of the Association for Compu- tational Linguistics, pages 311-318, Philadelphia. Pennsylvania, USA. Association for Computational Linguisties.
Hui Jiang, Ziyao Lu, Fandong Meng, Chulun Zhou, Jie Zhou, Degen Huang, and Jinsong Su. 2022. To- wards robust k-nearest-neighhor machine translation. In Proceedings of the 2022 Conference on Empiri- cal Methods in Natural Language Processing, pages Abu Dhabi, United Arab Emirates. As- sociation for Computational Linguistics.
Md Rizwan Parvez, Wasi Ahmad, Saikat Chakraborty, Baishakhi Ray, and Kai-Wei Chang. 2021. Retrieval augmented code generation and summarization. In Findings of the A Association for Computational Lin- guistics: EMNLP 2021, pages 2719-2734, Punta Cana, Dominican Republic. Association for Compu- tational Linguistics.
Jeff Johnson, Matthijs Douze, and Hervé Jégou. 2019. Billion-scale similarity search with GPUs. IEFF Transactions on Big Data, 73):535-547.
Nora Kassner and Hinrich Schütze. 2020. BERT-KNN:
Adding a KNN search component to pretrained lan- guage models for better QA. In Findings of the Asso- ciation for Computational Linguistics: EMNLP 2020, pages 3424-342233 3430, Online. Association for Computa- tional Linguistics.
Gabriel Poesia, Oleksandr Polozov, Vu Le, Ashish Ti- wari, Gustavo Soares, Christopher Meek, and Sumit Gulwani. 2022. Synchromesh: Reliable code gen- eration from pre-trained language models. arXiv preprint 22X1::2201.11222.
Urvashi Khandelwal, Angela Fan. Dan Jurafsky, Luke Zettlemayer, and Mike Lewis. 2021. Nearest neigh- bor machine translation. In 9th International Confer- ence on Learning Representations, ICLR 2021, Vir- tual Event, Austria, May 3-7, 2021. OpenReview.net.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. 2020. Exploring the limits

---

of transfer learning with a unified text-to-text trans- former. The Journal of Machine Learning Research, (1))):5485-551
Pengcheng Yin and Graham Neubig. 2018. TRANX: A transition- based neural abstract syntax parser for semantic parsing and code generation. In Proceed- ings of the 2018 Conference on Empirical Methods in Natural Language Processing: System Demonstra- tions, pages 7-12, Brussels, Belgium. Association for Computational Linguisties.
Shuo Ren, Daya Guo, Shuai Lu, Long Zhou, Shujie Liu, Duyu Tang. Neel Sundaresan, Ming Zhou, Ambrosio Blanco, and Shuai Ma. 2020. Codebleu: a method for automatic evaluation of code synthesis. CoRR, abs/2009.10297
Pengcheng Yin and Graham Neubig. 2019. Reranking for neural semantic parsing. In Proceedings of the 57th Annual Meeting of the Association for Computa- tional Linguistics, pages 4553 4559, Florence, Italy. Association for Computational Linguisties.
Kurt Shuster, Spencer Poff, Moya Chen, Douwe Kiela, and Jason Weston. 2021. Retrieval augmentation reduces hallucination in conversation, In Findings of the Association for Computational Linguistics: EMNLP 2021, pages 3784-3803, Punta Cana, Do- minican Republic. Association for Computational Linguistics.
Xin Zheng, Zhirui Zhang. Junliang Guo, Shujian Huang, Boxing Chen, Weihua Luo, and Jiajun Chen. 2021. Adaptive nearest neighbor machine translation. In Proceedings of the 59th Annual Meeting of the Asso- ciation for Computational Linguistics and the 1Ith International Joint Conference on Natural Language Processing ( (Volume 2: Short Papers), pages 368-374, Online. Association for Computational Linguistics.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Al- bert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. 2023. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2307.09288
Wenhao Zhu, Qianfeng Zhao, Yunzhe Lv, Shu- jian Huang, Siheng Zhao, Sizhe Liu, and Jia- jun Chen. 2023. knn-box: unified framework for nearest neighbor generation. arXiv preprint 13355
Daniel C. Wang. Andrew W. Appel, Jeffrey I.. Korn, and Christopher S. Serra. 1997. The zephyr ab- stract syntax description language. In Proceedings of the Conference
Domain- Domain-Specifie Languages, DSL '97, Santa Barbara, California, USA, October 15-17, 1997, pages 213-228. USENIX.
A
Appendix
Yue Wang, Hung Le, Akhilesh Deepak Gotmare, Nghi DQ Bui, Junnan Li, and Steven CH Hoi. 2023. Codet5+: Open code large language models for code understanding and generation. arXiv preprint arXiv:2305.07922.
A.1 Case Study on Model Canonical Incremental Adaptation
We showcase selected instances from the CoNaLa dataset in Table 6. For each instance, we present the output and probability distribution of the model when generating incorrect behavior. We also demonstrate how our method utilizes the kNN algorithm to enhance the model's decision-making process, thereby proving the effectiveness of our method. As shown in the first example, BertranXt demonstrates the ability construct code with ap- propriate syntax, but it inaccurately generates the primitive legend. However, with the assistance of KNN-TRANXI, the final decision of the model changes through retrieval, resulting in the desired code. In the second example, BertranXt makes an error when generating the second token. ANN- TR TRAA, utilizing an updated syntactic datastore, produces code that is semantically close to the tar- get code. In contrast, KNN-TRANXI, which did not update the syntactic datastore, has drawbacks in the generation of correct code due to insuffi- cient syntactic node information, despite having improved retrieval efficiency and datastore size.
Yue Wang, Weishi Wang. Shafiq Joty, and Steven C.H. Hoi. 2021. CodeT5: Identifier-awar unified pre- trained encoder-decoder models for code understand- ing and generation. In Proceedings of the 2021 Conference on Empirical Methods Natural Lan- guage Processing, pages 8696-8708, Online and Punta Cana, Dominican Republic. Association for Computational Linguistics.
Frank F. Xu, Zhengbao Jiang. Pengcheng Yin. Bogdan Vasilescu, and Graham Ncubig. 2020. Incorporating external knowledge through pre-training for natural language to code generation. Proceedings of the 58th Annual Meeting of the Association for Compu- tational L Linguistics, pages 6045-6052, Online. Asso- ciation for Computational Linguistics.
Pengcheng Yin, Bowen Deng, Edgar Chen, Bogdan Vasilescu, and Graham Neubig. 2018. Learning to mine parallel natural language/source code сorporа from stack overflow. In Proceedings of the 40th International Conference on Software Engineering: Companion Proceeedings, ICSE 2018, Gothenburg, Sweden, May 27- June 03, 2018, pages 388-389. ACM.

---

NL, Plot datatrame wilhou legend.
Target of plot (lagedd-Falae)
Target action sequeaee: [Expr, call, Attributo, hame, plot, Load. Reduce, keyword, legend, NareConstant, Model
Praditsed Cude
Prefficed Action Seqpence
Pone
p(alr)
(oli@ legendrö)
p(kind)
p(kind)
C. keyword, kind
...]
p(legend) 0.21
p(legend) 0:21
p(kind) 0.63
p(kind)
p(kind) 0.44
af (f.plee((lognneeFass)))
D-.. keyword, legend ..-]
0.38
p(tegend) p(tegend)
p(legeud)
p(kind) 0.53
p(kind) = 0.12
p(kend) 0.42
of (e.pplll((egnnnFallll))
keyaord, legond ..-
p(legend) pilegena))
pilegend)
NL: Get rid none values in dictionary kwargs,
Target code: res [k: For in List(kkargs, (3) If is not None)
Targel sequenee: Drssign, Name, res, Store, Reduce,, Dietc, Nore,
v. Lead, comprehension.
Model
Predicted Code
Predicied Action Sequence
Phw
A
P(a(=)
P(Cult) 0.43
P(Call) 0.48
kwarge, vallus())
[Fapr
p(DietC) 0.38
p(DiaC)) 0.38
p(Bubs) 0.14
((Subs) 0.14
P(Call) P(D(etC) 0.96
p(Call)
P(Col)
(ist(kkargs (tomss) 1f is not None)
(Expr, Dictt ...
p(DiatC) 0.38
0.47
p(DietC) 0.05
p(Suba) 0.14
y(Subx)
p(Suba) 0.07
P(Cul)) M(Call)
P(Cal) p(Dicc) 0.14
(tors
['spr,
Subss
...
((Did))) 0.38
p(Diet)
0.36
p(8mbx) 0.14
(uuhha) 0.79
p(Sube) 0.37
Table The case study of model canonical incremental adaptation. represents correct target code. the generated result of BertranXt, & and are the generated results of updating with 95k external data by ANN-TRANXI. The actions enclosed in the red boxes indicate the generated errors. We present the probability distribution of model decisions during the generation of these erroneous tokens.
updating the syntactic datastore
Model
Decoding Time (ms/code)
BertranX
357.5(x 1.00)
ANN-MT
(255.4(x3.52) 722.2(x2.02)
&NN-TRANX
Table 7: Decoding time of different models.
A.2 Decoding Time
We compared the decoding time of BertranX, A-N-- MT, and KNN-TRANX using the CoNaLa test set. During decoding, the beam size was set to 15. The results are presented in Table 7,


Reviews:
[{"review_id": "LAHxeR4U1d", "reviewer": "Reviewer_dTew", "summary": "", "questions": "- A. Why was BertanX chosen as the base model for this work on not a more recent NMT system like T5, BART.\n- B. Why are BLEU and CodeBLEU missing in Table 2 for Reranker?\n- C. Do the authors know how the performance might vary based on the underlying NMT model? Are there specific domains or tasks where this can be more impactful? maybe low-resource languages?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the comments.  \n\n## Questions for the authors\n\n* **Q1. Why was BertanX chosen as the base model for this work on not a more recent NMT system like T5, BART.** \n\n\t**A1:** Our approach aims to address the token-level retrieval challenge by leveraging the AST feature of programming languages. To achieve this, we adopt the BertranX seq2tree model as the basis. However, it should be emphasized that our approach can also be applied to seq2seq models. One viable option is to enhance the vocabulary table of conventional seq2seq models (such as T5 and BART) by incorporating the actions required for constructing ASTs, which would equip these models with comparable capabilities.\n\n* **Q2. Why are BLEU and CodeBLEU missing in Table 2 for Reranker?** \n\n    **A2:** To ensure a comprehensive comparison using more evaluation metrics, we tried to replicate the results by utilizing the checkpoints provided by the authors. However, the checkpoint for the Reranker on the Django dataset was unavailable, so we had to rely on the results from the original paper for this particular component. Nevertheless, we used their method to reproduce the results on the Django dataset, which are shown below.\n\n\t|  |  | Django | | \n\t| ------ | ------ | ------ | ------ |\n\t|  | BLEU | CodeBLEU | EM |\n\t| Reranker | 73.26 | 71.29 | 80.18 |\n\t| Ours | **83.20** | **81.66** | **82.16** |\n\n* **Q3. Do the authors know how the performance might vary based on the underlying NMT model? Are there specific domains or tasks where this can be more impactful? maybe low-resource languages?** \n\n\t**A3:** Although the underlying model can indeed impact performance, prior research indicates that powerful models can still reap benefits from token-level retrieval. Furthermore, by creating multiple domain-specific datastores using the target programming language, we can achieve domain adaptation by seamlessly switching between these datastores. For low-resource languages, they are often poorly optimized in the model, so we believe that introducing external knowledge will be helpful for the model to generate these languages.\n\n## Further clarification \n\n* **Novelty:** We agree that token-level retrieval methods have shown excellent performance in the field of NLP, especially in the NLG domain. However, our method aims to generate code written in programming languages, which is different from natural language. Programming languages need to meet strict syntactic constraints (otherwise the code would not compile), which traditional text generation models struggle to ensure (as shown in Table 1 of the paper). In addition, retrieval-based generation methods require higher inference latency and additional storage, which hinders the development of this approach. To address these issues, we leverage domain knowledge and program analysis techniques to ensure the syntactic correctness of the generated code while improving inference efficiency and reducing storage consumption.\n\n* **Ablation study:** While the ablation experiments conducted so far may not fully offer conclusive evidence about the significance of the components we propose, they indeed show a substantial impact on evaluation metrics. It is intriguing to consider testing this concept in a low-resource programming language, as it would provide valuable insights. Notably, as illustrated in Section 4.3, we simulated an environment with limited resources and the experimental results convincingly demonstrate that the model still benefits from retrieval.\n\n* **Evaluation:**\n\n   * **Add baselines**: To further demonstrate the effectiveness of our approach, it would better compare our model with the SoTA code pre-trained models and LLM. Based on the suggestions of the reviewers, we added four baselines for comparison. The experimental results are presented below. \n\n    |  |  | CoNaLa ||| Django || \n    | ------ | ------ | ------ | ------ | ------ | ------ | ------ |\n    |  | BLEU | CodeBLEU | EM | BLEU | CodeBLEU | EM |\n    | REDCODER | 35.12 | 36.82 | 6.4 | 75.36 | 73.01 | 78.82 |\n    | CodeT5 base | 36.28 | 35.34 | 6.8 | 76.50 | 71.92 | 80.93 |\n    | CodeT5+ base | 35.84 | 34.76 | 6.4 | 77.32 | 72.28 | 80.98 |\n    | gpt-3.5 3-shots | 36.25 | 38.03 | **7.4** | 68.97 | 72.58 | 58.85 |\n    | Ours | **37.29** | **39.04** | **7.4** | **83.20** | **81.66** | **82.16** |\n\n\tOne can see that our model performs better on three evaluation metrics. Furthermore, it is important to note that our method can be integrated with existing models.\n\n    * **Training data leakage**: Our base model utilizes BERT as the encoder and a decoder that is randomly initialized. Compared to pre-trained encoder-decoder and decoder-only architectures, we believe that our model has a lower risk of information leakages.\n\n\t* **Evaluation metrics:** To demonstrate that our model can generate code with more accurate syntax, we provide the syntax match score, which reflects the degree of syntax matching in code. \n\n\t| Model | TRANX | Reranker | Ext-codegen | TAE | BertranX | REDCODER | CodeT5 base | CodeT5+ base | gpt-3.5 3-shots| Ours|\n    | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ |\n\t|Syntax Match| .287 | .284 | .322| .318 | .347 | .344 | .339 | .341 | .357 | **.368** |\n\n\tCodeBERTScore is also a widely adopted automated code evaluation metric, and we provide its evaluation below.\n\n    |  | CoNaLa |||| Django ||||\n    | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ |     \n\t|  | $P$ | $R$ | $F_{1}$ | $F_{3}$ | $P$ | $R$ | $F_{1}$ | $F_{3}$ |\n    | TRANX | .811 | .796 | .802 | .797 | .915 | .902 | .905 | .908 |\n    | Reranker | .823 | .798 | .809 | .800 | .932 | .925 | .932 | .930 |\n    | Ext-codegen | .809 | .805 | .806 | .805 | .935 | .931 | .933 | .932 |\n    | TAE | .812 | .808 | .804 | .803 | .941 | .934 | .936 | .935 |\n    | BertranX | .821 | .812 | .815 | .812 | .987 | .982 | .982 | .981 |\n    | REDCODER | .832 | .825 | .822 | .818 | .986 | .983 | .985 | .981 |\n    | CodeT5 base | .831 | .823 | .826 | .824 | .985 | .982 | .984 | .983 |\n    | CodeT5+ base | .831 | .815 | .822 | .818 | .986 | .981 | .986 | .982 |\n    | gpt-3.5 3-shots | .838 | **.835** | .826 | .827 | .974 | .972 | .974 | .969 |\n    | Ours | **.842** | .827 | **.833** | **.828** | **.989** | **.985** | **.987** | **.985** |", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "The strengths of this paper are:\n1. Adapting token level retrieval for generation tasks to the domain of code generation and discussing the insights and challenges.\n2. Well written and easy to follow. It is a simple system that works well for the domain.\n3. Improving over existing retrieval augmented generation systems with meta-K networks and confidence networks.", "reasons_to_reject": "The paper proposes a very nice system for improving code generation performance using token level retrieval but it has some weaknesses that need to be addressed.\n1. The idea around token level retrieval for improved augmentation in itself is not novel and has been extensively used in NLP.\n2. The novel components in BertranX seem to be the separate data store, constrained generation and the confidence network for combining model and retrieval predictions. It is not clear how significant these are as the ablations are not very indicative. Perhaps it would be more convincing if the authors test this on a low resource programming language which might benefit from retrieval.\n3. The evaluation is not very thorough. Some suggestions\n    - There needs to be some comparison with SOTA code generation models like T5, Llama2, GPT-3.5\n    - Perhaps, Conala and Django are not the best languages to test this on as a lot of models have seen these in their pretrasining.\n    - The metrics reported are standard for code generation but specifically for this work, it would be interesting to see the overall syntactic generation fraction (since the authors claim that they generate more syntactic code) and CodeBERTScore since Conala is known to be a noisy dataset with no normalization for variable names and other constants.", "questions_to_authors": "- A. Why was BertanX chosen as the base model for this work on not a more recent NMT system like T5, BART.\n- B. Why are BLEU and CodeBLEU missing in Table 2 for Reranker?\n- C. Do the authors know how the performance might vary based on the underlying NMT model? Are there specific domains or tasks where this can be more impactful? maybe low-resource languages?", "paper_topic_and_main_contributions": "The paper explores the problem of improving the quality of language guided code generation using token level retrieval from an offline datastore that is separately constructed at train time. The paper improves upon this idea of token level retrieval by making the following contributions:\n(1) Instead of directly generating code, the paper generates AST which allows the system to constrain generation and improve quality in terms of generating more syntactic code.\n(2) The paper extends the meta-K networks for choosing the optimal retrieval cache size.\n(3) Instead of using a fixed confidence threshold (as done in previous works), the paper employs a confidence network for combining the predictions from the neural and retrieval models.\n\nThe paper improves traditional code generation via token level retrieval which is particularly useful in low-resource domains where training data is scarce.", "excitement": 4, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Missing_References": "https://arxiv.org/abs/2201.11227 -- introduces constraint guided generation"}, "extra_text": {"Typos_Grammar_Style_and_Presentation_Improvements": "1. The datastore construction equations in section 3.1 need to be in their own environment\n2. Incremental learning capability of kNN-TRANX are very abruptly introduced in section 4.3. Perhaps, a discussion early on might be useful.\n3. It would be nice if the example in Figure-2 would be used as a running example through the paper especially in section 3."}}, {"review_id": "FbIuOqt6R7", "reviewer": "Reviewer_VgbA", "summary": "", "questions": "How is the proposed approach compared with large language model (LLM) based code generation methods?", "limitations": "", "rating": 3, "confidence": 4, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments. \n\n## Questions for the authors\n\n* **Q. How is the proposed approach compared with large language model (LLM) based code generation methods?** \n\n    **A:**  To demonstrate the effectiveness of our work, we agree that it would better compare our model with the SoTA code pre-trained models and LLM. We present experimental results that encompass retrieval-based REDCODER, as well as code pre-training models CodeT5 and CodeT5+, alongside the large language model gpt-3.5.\n\n    |  |  | CoNaLa ||| Django || \n    | ------ | ------ | ------ | ------ | ------ | ------ | ------ |\n    |  | BLEU | CodeBLEU | EM | BLEU | CodeBLEU | EM |\n    | REDCODER | 35.12 | 36.82 | 6.4 | 75.36 | 73.01 | 78.82 |\n    | CodeT5 base | 36.28 | 35.34 | 6.8 | 76.50 | 71.92 | 80.93 |\n    | CodeT5+ base | 35.84 | 34.76 | 6.4 | 77.32 | 72.28 | 80.98 |\n    | gpt-3.5 3-shots | 36.25 | 38.03 | **7.4** | 68.97 | 72.58 | 58.85 |\n    | Ours | **37.29** | **39.04** | **7.4** | **83.20** | **81.66** | **82.16** |\n\n\tOne can see that our model performs better on three evaluation metrics. Furthermore, it is important to note that our method can be integrated with existing language models.\n\n## Further clarification \n\n* **Related works to be added:** In the era of large language models, we fully understand the necessity of relevant discussions. We will include additional research on large language models in the revised version of the paper.\n\n* **Metrics**: We acknowledge the effectiveness of using pass@k as an indicator for assessing code correctness. However, these two datasets do not provide standardized variable names and other constants, so the pass@k metric cannot be effectively applied. As for the datasets that can be evaluated using the pass@k metric, due to limited computational resources, small models struggle to generate high-quality method-level code. Hence, for our experiments on these two code datasets, we performed evaluations at the line level. In addition to BLEU and EM, which have been used in previous studies, we have introduced the CodeBLEU metric in our paper and presented CodeBERTScore as an automatic evaluation metric in the following table, showcasing the efficacy of our approach.\n\n    |  | CoNaLa |||| Django ||||\n    | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ |     \n\t|  | $P$ | $R$ | $F_{1}$ | $F_{3}$ | $P$ | $R$ | $F_{1}$ | $F_{3}$ |\n    | TRANX | .811 | .796 | .802 | .797 | .915 | .902 | .905 | .908 |\n    | Reranker | .823 | .798 | .809 | .800 | .932 | .925 | .932 | .930 |\n    | Ext-codegen | .809 | .805 | .806 | .805 | .935 | .931 | .933 | .932 |\n    | TAE | .812 | .808 | .804 | .803 | .941 | .934 | .936 | .935 |\n    | BertranX | .821 | .812 | .815 | .812 | .987 | .982 | .982 | .981 |\n    | REDCODER | .832 | .825 | .822 | .818 | .986 | .983 | .985 | .981 |\n    | CodeT5 base | .831 | .823 | .826 | .824 | .985 | .982 | .984 | .983 |\n    | CodeT5+ base | .831 | .815 | .822 | .818 | .986 | .981 | .986 | .982 |\n    | gpt-3.5 3-shots | .838 | **.835** | .826 | .827 | .974 | .972 | .974 | .969 |\n    | Ours | **.842** | .827 | **.833** | **.828** | **.989** | **.985** | **.987** | **.985** |\n\n* **Novelty:** The success of $k$NN-MT in enhancing text generation models and its ease of model updating surprised us. However, despite its potential benefits, this token-level retrieval-enhanced model is not widely adopted in practice. The primary concern is the increasing inference latency and additional storage resources. Furthermore, retrieval-based approaches introduce noise that can negatively impact generation, a critical issue in code generation tasks where a single inappropriate token may lead to the failure of an entire code block. In contrast to traditional text generation tasks, our work relies on domain knowledge and program analysis techniques. By leveraging these techniques, we construct multiple datastores that encompass both syntax and semantics information. This approach improves retrieval efficiency and reduces noise through restrictive retrieval. Additionally, our model introduces a self-adaptive combination method that significantly enhances the effectiveness of model generation. We have also investigated the impact of domain knowledge on model generation and devised a convenient approach for model updating, preserving competitive performance while minimizing additional storage consumption.\n\n* **Reproducibility concern:** In the event that the paper is accepted, we are committed to providing source code, training parameters, and stored datastores.\n\n* **Missing References:**  We will include more relevant papers in the revised version.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- This paper is generally well written with a good structure. It is easy to follow.\n\n- The idea of combining kNN-MT and Seq2Tree for code generation is intriguing, and the authors proposed three mechanisms to enhance performance: Syntax-constrained token-level retrieval, Meta-k network, and Confidence network.\n\n", "reasons_to_reject": "- Even though the paper has proposed several novel techniques to improve code generation performance, my major concern lies in the lack of comparison with recent large language models (LLMs) such as ChatGPT, StarCoder, and CodeT5+. Existing LLMs have demonstrated significant performance in code generation, and it is suggested that the authors compare their model with these LLMs. At least a discussion on this comparison is required.\n\n- In the Related Work Section, many works on code generation have not been included, especially those LLM-based models.\n\n- In many recent code generation papers, the Pass@K metric is consistently adopted for evaluation. The authors are strongly encouraged to take this metric into consideration during evaluation.\n\n- The core idea of this work is primarily built upon the existing kNN-MT and Seq2Tree models. The authors are encouraged to emphasize the main innovative points and distinct contributions in comparison to these models.", "questions_to_authors": "How is the proposed approach compared with large language model (LLM) based code generation methods?", "paper_topic_and_main_contributions": "To reduce extraneous noise in the generation process of pretrained neutral machine translation models, the authors propose kNN-TRANX, a token-level retrieval augmented code generation method. kNN-TRANX enables searches in smaller datastores tailored for the code generation task, leveraging syntax constraints for the retrieval of datastores, thereby mitigating the impact of retrieval noise.", "excitement": 3, "reproducibility": 2, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {"Missing_References": "As mentioned, there are many recent works on LLM-based code generation models. Also, there are other retrieval-augmented code intelligence works such as the following:\nJian Zhang et al., 2020. Retrieval-based neural source code summarization. In Proceedings of the ACM/IEEE 42nd International Conference on Software Engineering (ICSE '20),1385–1397. "}}, {"review_id": "W2nJuj3pyF", "reviewer": "Reviewer_Av9a", "summary": "", "questions": "* Why do you need to use seq2tree models to extract AST instead of using a static parser?\n* Can you provide comparison with SoTA retrieval-augmented methods such as REDCODER and the recent CodeT5+ models?", "limitations": "", "rating": 2, "confidence": 4, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the comments. \n\n## Questions for the authors\n\n* **Q1. Why do you need to use seq2tree models to extract AST instead of using a static parser?** \n\n    **A1:** A static parser can easily obtain AST information. However, in our work, the seq2tree model is actually used as a generator rather than a parser. In fact, the static parser is also part of the seq2tree architecture, used to parse the target code into AST. The main functionality of the seq2tree model, as described in section 2.2 of this paper, is to provide a standardized method for generating ASTs using ASDL (To make it easier to understand, we provided a toy example in Figure 2 of the paper to illustrate the generation process of $k$NN-TRANX.). In this way, we can combine program analysis techniques during the generation process to accelerate inference and filter noise.\n\n* **Q2. Can you provide comparison with SoTA retrieval-augmented methods such as REDCODER and the recent CodeT5+ models?**\n\n    **A2:**  We have conducted a detailed evaluation of the performance of four SoTA baselines in the following table. \n\n    |  |  | CoNaLa ||| Django || \n    | ------ | ------ | ------ | ------ | ------ | ------ | ------ |\n    |  | BLEU | CodeBLEU | EM | BLEU | CodeBLEU | EM |\n    | REDCODER | 35.12 | 36.82 | 6.4 | 75.36 | 73.01 | 78.82 |\n    | CodeT5 base | 36.28 | 35.34 | 6.8 | 76.50 | 71.92 | 80.93 |\n    | CodeT5+ base | 35.84 | 34.76 | 6.4 | 77.32 | 72.28 | 80.98 |\n    | gpt-3.5 3-shots | 36.25 | 38.03 | **7.4** | 68.97 | 72.58 | 58.85 |\n    | Ours | **37.29** | **39.04** | **7.4** | **83.20** | **81.66** | **82.16** |\n\n    As for the APPS and HumanEval datasets, due to limited computational resources, small-scale models struggle to perform well on these two datasets. In future work, we hope to experiment with large language models on these two datasets.\n\n## Further clarification \n\n* **Inference efficiency:** While we acknowledge the importance of inferring efficiency in deployment, it is true that introducing a retrieval module will unavoidably lead to an increase in inference latency as we discussed in the Limitation part of the paper. However, in our research, we still strive to solve this issue, which can be divided into two parts. Firstly, we have constructed fine-grained datastores based on the kind of AST nodes to narrow down the retrieval space, thereby minimizing the potential impact on latency. This allows for more precise retrieval of information during the code generation process. Secondly, we have ensured that the retrieved target tokens comply with the construction rules of the AST. By doing so, we strike a balance between efficient retrieval and maintaining the structural integrity of the code being generated. While it is true that there may still be some increase in inference latency (shown in Table 7 of the paper), our approach ultimately yields better results in terms of code generation. Additionally, our proposed method offers a more efficient incremental learning capability. This enables the model to have online learning ability without additional training requirements. We believe that these advantages are also significant in real-world applications where efficiency and adaptability are crucial.\n\n* **Missing References:**  We will include more relevant papers in the revised version.\n\n* **Reproducibility concern:** In the event that the paper is accepted, we are committed to providing source code, training parameters, and stored datastores.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "* The paper is well written and easy to read. \n* The proposed method achieves better results than baselines on two benchmarks. Additional experimental analysis is well conducted to better understand the models.", "reasons_to_reject": "* One major limitation of this paper is that this work does not compare with SoTA retrieval-augmented methods such as REDCODER and the recent CodeT5+ models or more widely adopted code generation benchmarks such as APPS and HumanEval. These SoTA retrieval-augmented methods often employ an emsembling approach to use a sparse retriever (e.g., BM25) for syntatic matching and a deep retriever (e.g., CodeBERT) for semantic matching. Without comparison to these baselines, it is difficult to determine the technical significance of the proposed model.\n* The use of seq2tree models is a bit strange as the AST information can be easily obtained via off-the-shelf parsers.\n* The retrieval component of kNN-TRANX is a bit too complicated which might sacrifice its inference efficiency. In Table 5, the authors show some statistics of the training cost, but for code generation systems, the inference efficiency in deployment is often more crucial.", "questions_to_authors": "* Why do you need to use seq2tree models to extract AST instead of using a static parser?\n* Can you provide comparison with SoTA retrieval-augmented methods such as REDCODER and the recent CodeT5+ models?", "paper_topic_and_main_contributions": "This paper proposes kNN-TRANX, a token-level retrieval augmented code generation method. kNN-TRANX allows for searches in smaller datastores tailored for the code generation task. It utilizes seq2tree model to get AST information and leverages such syntax constraints for the retrieval of datastores, which reduces the impact of retrieve noise. Results on two public datasets verify the effectiveness of the proposed approach. However, this work does not compare with SoTA retrieval-augmented methods such as REDCODER and the recent CodeT5+ models or more widely adopted code generation benchmarks such as APPS and HumanEval. Besides, the use of seq2tree models is a bit strange as the AST information can be easily obtained via off-the-shelf parsers.", "excitement": 2, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Missing_References": "A recent work of CodeT5+ also includes evaluation on retrieval-augmented code generation tasks and should be cited and compared."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper investigates methods for improving the performance of token-level retrieval augmented models for code generation.

The proposed framework seems to work well and achieves good results, compared to the chosen baselines.
All reviewers comment that the paper is very well written and easy to read.

Reviewers highlighted concerns about lack of comparison to state-of-the-art models and LLMs. Authors have provided some additional results in the comments which seem to demonstrate the proposed method outperforms several LLMs.
Reviewers also voiced concerns about the novelty and motivation of the method.
The evaluation could be improved, with better ablation tests to clearly show whether the proposed components are responsible for performance improvements.
Decision: Accept-Findings