Adaptive Nearest Neighbor Machine Translation (2021)

Abstract: kNN-MT, recently proposed by Khandelwal et al. (2020a), successfully combines pre-trained neural machine translation (NMT) model with token-level k-nearest-neighbor (kNN) retrieval to improve the translation accuracy. However, the traditional kNN algorithm used in kNN-MT simply retrieves a same number of nearest neighbors for each target token, which may cause prediction errors when the retrieved neighbors include noises. In this paper, we propose Adaptive kNN-MT to dynamically determine the number of k for each target token. We achieve this by introducing a light-weight Meta-k Network, which can be efficiently trained with only a few training samples. On four benchmark machine translation datasets, we demonstrate that the proposed method is able to effectively filter out the noises in retrieval results and significantly outperforms the vanilla kNN-MT model. Even more noteworthy is that the Meta-k Network learned on one domain could be directly applied to other domains and obtain consistent improvements, illustrating the generality of our method. Our implementation is open-sourced at https://github.com/zhengxxn/adaptive-knn-mt .

## introduction retrieval-based methods @xcite @xcite are increasingly receiving attentions from the machine translation (mt) community recently. these approaches complement advanced neural machine translation (nmt) models @xcite @xcite to alleviate the performance degradation when translating out-of-domain sentences @xcite , rare words @xcite , etc. the ability of accessing any provided datastore during translation makes them scalable, adaptable and interpretable. to tackle this problem, we propose adaptive knn-mt that determines the choice of k regarding each target token adaptively. specifically, instead of utilizing a fixed k, we consider a set of possible k that are smaller than an upper bound k. then, given the retrieval results of the current target token, we propose a light-weight meta-k network to estimate the importance of all possible k-nearest neighbor results, based on which they are aggregated to obtain the final decision of the model. in this way, our method dynamically evaluate and utilize the neighbor information conditioned on different target tokens, therefore improve the translation performance of the model. we conduct experiments on multi-domain machine translation datasets. across four domains, our approach can achieve 1.44∼2.97 bleu score improvements over the vanilla knn-mt on average when @xmath0. the introduced light-weight meta-k network only requires thousands of parameters and can be easily trained with a few training samples. in addition, we find that the meta-k net- in this section, we will briefly introduce the background of knn-mt, which includes two steps: creating a datastore and making predictions depends on it. datastore creation. the datastore consists of a set of key-value pairs. formally, given a bilingual sentence pair in the training set (x, y) ∈ (x , y), a pre-trained autoregressive nmt decoder translates the t-th target token y t based on the translation context (x, @xmath1. denote the hidden representations of translation contexts as f (x, @xmath2, then the datastore is constructed by taking f (x, @xmath3, therefore, the datastore can be created through a single forward pass over the training set (x , y). prediction. while inference, at each decoding step t, the knn-mt model aims to predict ŷt given the already generated tokens ŷ<t as well as the context representation f (x, ŷ<t ), which is utilized to query the datastore for k nearest neighbors w.r.t the l 2 distance. denote the retrieved neighbors as n @xmath4, v i ), @xmath5, 2, ..., k}}, their distribu-tion over the vocabulary is computed as: ## adaptive knn-mt the vanilla knn-mt method utilizes a fixed number of translation contexts for every target token, which fails to exclude noises contained in retrieved neighbors when there are not enough relevant items in the datastore. we show an example with @xmath6. the correct prediction spreadsheet has been retrieved as top candidates. however, the model will finally predict table instead because it appears more frequently in the datastore than the correct prediction. a naive way to filter the noises is to use a small k, but this will also cause over-fitting problems for other cases. in fact, the optimal choice of k varies when utilizing different datastores in vanilla knn-mt, leading to poor robustness and generalizability of the method, which is empirically discussed in section 4.2. to tackle this problem, we propose a dynamic method that allows each untranslated token to utilize different numbers of neighbors. specifically, we consider a set of possible ks that are smaller than an upper bound k, and introduce a lightweight meta-k network to estimate the importance of utilizing different ks. practically, we consider the powers of 2 as the choices of k for simplicity, as well as @xmath7, i.e., @xmath8, k @xmath9. then the meta-k network evaluates the probability of different knn results by taking retrieved neighbors as inputs. concretely, at the t-th decoding step, we first retrieve k neighbors n t from the datastore, and for each neighbor (h i , v i ), we calculate its distance from the current context representation as distances and @xmath10, ..., c k ) as counts of values for all retrieved neighbors, we then concatenate them as the input features to the meta-k network. the reasons of doing so are two-fold. intuitively, the distance of each neighbor is the most direct evidence when evaluating their importance. in addition, the value distribution of retrieved results is also crucial for making the decision, i.e., if the values of each retrieved results are distinct, then the knn predictions are less credible and we should depend more on nmt predictions. we construct the meta-k network f meta (•) as two feed-forward networks with non-linearity between them. given [d; c] as input, the probability of applying each knn results is computed as: prediction. instead of introducing the hyperparameter λ as equation ( 2 ), we aggregate the nmt model and different knn predictions with the output of the meta-k network to obtain the final prediction: training. we fix the pre-trained nmt model and only optimize the meta-k network by minimizing the cross entropy loss following equation (4), which could be very efficient by only utilizing hundreds of training samples. ## conclusion and future works in this paper, we propose adaptive knn-mt model to dynamically determine the utilization of retrieved neighbors for each target token, by introducing a light-weight meta-k network. in the experiments, on the domain adaptation machine trans-lation tasks, we demonstrate that our model is able to effectively filter the noises in retrieved neighbors and significantly outperform the vanilla knn-mt baseline. in addition, the superiority of our method on generality and robustness is also verified. in the future, we plan to extend our method to other tasks like language modeling, question answering, etc, which can also benefit from utilizing knn searching @xcite .