Towards Robust k-Nearest-Neighbor Machine Translation (2022)

Abstract: k-Nearest-Neighbor Machine Translation (kNN-MT) becomes an important research direction of NMT in recent years. Its main idea is to retrieve useful key-value pairs from an additional datastore to modify translations without updating the NMT model. However, the underlying retrieved noisy pairs will dramatically deteriorate the model performance. In this paper, we conduct a preliminary study and find that this problem results from not fully exploiting the prediction of the NMT model. To alleviate the impact of noise, we propose a confidence-enhanced kNN-MT model with robust training. Concretely, we introduce the NMT confidence to refine the modeling of two important components of kNN-MT: kNN distribution and the interpolation weight. Meanwhile we inject two types of perturbations into the retrieved pairs for robust training. Experimental results on four benchmark datasets demonstrate that our model not only achieves significant improvements over current kNN-MT models, but also exhibits better robustness. Our code is available at https://github.com/DeepLearnXMU/Robust-knn-mt.

## introduction as a commonly-used paradigm of retrieval-based neural machine translation (nmt), k-nearest-neighbor machine translation (knn-mt) has proven to be effective in many studies @xcite @xcite @xcite , and thus attracted much attention in the community of machine translation. the core of knn-mt is to use an auxiliary datastore containing cached decoder representations and corresponding target tokens. this datastore can flexibly guide the nmt model figure 1 : an example of knn-mt model using dynamically estimation of λ @xcite . to make better predictions, especially for domain adaptation. compared with other retrieval-based paradigm @xcite , knn-mt has two advantages: 1) it is more scalable because we can directly improve the nmt model by just manipulating the datastore. 2) it is more interpretable due to its observable retrieved pairs. generally, knn-mt mainly involves two stages: datastore establishment and candidate retrieval. during the first stage, a pre-trained nmt model is used to construct a datastore containing key-value pairs, where the key is the decoder representation and the value is the corresponding target token. at the second stage, given the current decoder representation as a query at each time step, the k nearest key-value pairs are retrieved from the datastore. then, according to the query-key distances, the retrieved values are converted into a translation probability distribution over candidate target tokens, denoted as knn distribution. finally, the predicted distribution of the nmt model is interpolated by knn distribution with a hyper-parameter λ. along this line, many efforts have been made to improve knn-mt @xcite . particularly, as shown in figure 1 , adaptive knn-mt @xcite uses the query-key distances and retrieved pairs to dynamically estimate λ, exhibiting better performance than most knn-mt models. however, we find that existing knn-mt models often suffer from a serious drawback: the model performance will dramatically deteriorate due to the underlying noise in retrieved pairs. for example, in figure 1 , the retrieved results may contain unrelated tokens such as "no", which leads to a harmful knn distribution. besides, for estimating λ, previous studies @xcite only consider the retrieved pairs while ignoring the nmt distribution. back to figure 1 , compared with the knn distribution, the nmt model gives a much higher probability to the ground-truth token "been". although the knn distribution is insufficiently accurate, it is still assigned with a greater weight than the nmt distribution. obviously, this is inconsistent with our intuition that when the nmt model has high confidence in its prediction, it needs less help from others, and thus the knn distribution should have a lower weight. moreover, we find that during training, a non-negligible proportion of the retrieved pairs from the datastore do not contain groundtruth tokens. this can cause insufficient training of knn modules. to sum up, conventional knn-mt models are vulnerable to noise in datastore, for which we further conduct a preliminary study to validate the above issues. therefore, dealing with the noise for the knn-mt model remains to be a significant task. in this paper, we explore a robust knn-mt model. in terms of model architecture, we explore how to more accurately estimate the knn distribution and better combine it with the nmt distribution. concretely, unlike previous studies @xcite that only use retrieved pairs to dynamically estimate λ, we additionally use the confidence of nmt prediction to calibrate the calculation of λ where confidence is the predicted probability on each retrieved token. meanwhile, we improve the knn distribution by integrating the confidence to reduce the effect of noise. besides, we propose to boost the robustness of our model by randomly adding perturbations to retrieved key representations and augmenting retrieved pairs with pseudo ground-truth tokens. by these means, our proposed approach can enhance the knn-mt model to better cope with the noise in retrieved pairs, thus improving its robustness. to investigate the effectiveness and generality of our model, we conduct experiments on several commonly-used benchmarks. experimental results show that our model significantly outperforms the adaptive knn-mt, which is the state-of-the-art knn-mt model, across most domains. moreover, our model exhibits better performance than adaptive knn-mt on pruned datastores. ## related work retrieval-based approaches leveraging auxiliary sentences have shown effectiveness in improving nmt models. usually, they first retrieve relevant sentences from translation memory and then exploit them to boost nmt models during making a translation. for example, @xcite maintains a continuous cache storing attention vectors as keys and decoder representations as values. the retrieved values are then used to update the decoder representations. @xcite preform n-gram retrieval to identify similar source n-grams from the translation memory, where the corresponding target words are then encoded to update decoder representations. @xcite pack the retrieved target sentences into a compact graph which is then incorporated into decoder representations. @xcite propose several transformer-based encoding methods to vectorize retrieved target sentences. @xcite propose a cross-lingual memory retriever to leverage target-side monolingual translation memory, showing effectiveness in low-resource and domain adaption scenarios. finally, it should be noted that there have been many nlp studies @xcite @xcite on exploring robustness of nlp models. in comparison with the above-mentioned studies, our work is the first to improve the robustness of knn-mt approaches. ## preliminary study to investigate the impact of noise on knn-mt, we conduct a preliminary experiment in this section. we use the nmt confidence to represent the predicted probability on the target token from the nmt model. by this way, we remove the pairs of datastore within different confidence intervals to investigate the impact of noisy datastore in different degrees of nmt confidence. specifically, during the datastore establishment, besides the key-value pairs, we additionally record the nmt confidence of each target token and use it to rank these pairs. then we split the datastore into multiple partitions, we alternatively remove each partition of the datastore and observe the performance of vanilla knn-mt @xcite and adaptive knn-mt @xcite on the it training dataset @xcite . to ensure the persuasiveness of our experiments, we directly use the setting of adaptive knn-mt. in this way, we remove the key-value pairs within a specific interval of ranking to see the model performance. furthermore, since adaptive knn-mt ignores nmt distribution, it may give a large weight to the incorrect knn distribution. to verify this, we collect the λ generated by adaptive knn-mt with respect to different nmt confidences in figure 2 . looking at the orange curve, we find that when the adaptive knn-mt fails to retrieve the groundtruth token, it gives a similar weight λ regardless of the nmt confidence. besides, the blue curve shows the situation when the ground-truth token is successfully retrieved. we can see that adaptive knn-mt gives a relatively small λ=0.52 even when the nmt model fails to predict the groundtruth (see [0-0.2) interval in figure 2 ). intuitively, the performance of the model would be further improved if it can generate a larger λ when the model has unconfident nmt distribution and high-quality knn distribution. the above experimental results indicate that the knn-mt models are sensitive to the quality of the datastore, which limits their applicability to a noisy datastore. therefore, it is of great significance to explore robust knn-mt. ## experiments to investigate the effectiveness and robustness of our model, we carry out experiments on several commonly-used datasets. ## conclusion in this paper, via preliminary study, we first point out that existing knn-mt models are very susceptible to the quality of retrieved pairs. then, we explore robust knn-mt, which improves knn-mt models in the aspects of model architecture and training. concretely, we incorporate the confidence of nmt prediction into modeling knn distribution and dynamic estimation of λ. besides, during the model training, we inject two types of perturbations into the retrieved pairs, which can effectively enhance the generalization of the model. extensive results and in-depth analysis strongly demonstrate the effectiveness of our model. to further verify the generality of our model, we will extend our model to other conditional text generation tasks, such as speech translation. besides, we will try to combine knn-mt with topic information, which has been successfully applied in previous studies @xcite @xcite , to constraint retrieval in the future.