Nearest Neighbor Machine Translation is Meta-Optimizer on Output Projection Layer (2023)

Abstract: Nearest Neighbor Machine Translation ( k NN-MT) has achieved great success in domain adaptation tasks by integrating pre-trained Neural Machine Translation (NMT) models with domain-specific token-level retrieval. However, the reasons underlying its success have not been thoroughly investigated. In this paper, we comprehensively analyze k NN-MT through theoretical and empirical studies. Initially, we provide new insights into the working mechanism of k NN-MT as an efficient technique to implicitly execute gradient descent on the output projection layer of NMT, indicating that it is a specific case of model fine-tuning. Subsequently, we conduct multi-domain experiments and word-level analysis to examine the differences in performance between k NN-MT and entire-model fine-tuning. Our findings suggest that: ( i ) Incorporating k NN-MT with adapters yields comparable translation performance to fine-tuning on in-domain test sets, while achieving better performance on out-of-domain test sets; ( ii ) Fine-tuning significantly outperforms k NN-MT on the recall of in-domain low-frequency words, but this gap could be bridged by optimizing the context representations with additional adapter layers.

## introduction in recent years, nearest neighbor machine translation (knn-mt) and its variants @xcite @xcite have provided a new paradigm and achieved strong performance for fast domain adaptation through retrieval pipelines. unlike model fine-tuning, which requires additional parameter updates or introduces external adapter layers, knn-mt combines traditional neural machine translation (nmt) models @xcite @xcite with a token-level k-nearest-neighbour retrieval mechanism. this allows for direct access to domain-specific datastores, improving translation accuracy without the need for supervised finetuning. although knn-mt has achieved great success in domain adaptation tasks, its working mechanism is still an open problem that has not been thoroughly investigated. in this paper, we propose a novel perspective to understand knn-mt by describing it as a special case of fine-tuning, specifically a process of meta-optimization on the output projection layer (opl) of nmt, and establish connections between knn-mt and model fine-tuning (section 3). our novel perspective on knn-mt posits that (i) the working mechanism of knn-mt is to implicitly execute gradient descent on opl, producing metagradients via forward computation based on knearest-neighbors, and (ii) explicit fine-tuning on opl shares a similar gradient format with the metagradients obtained by knn-mt, according to the derivation of back-propagation. as illustrated in figure 1 , knn-mt and explicit opl fine-tuning share a dual view of gradient descent-based optimization. the key difference between them lies in the method for computing gradients: knn-mt produces meta-gradients through forward computation and interpolation, while fine-tuning method computes gradients of opl via back-propagation. hence, it is reasonable to understand knn-mt as an implicit form of model fine-tuning. to provide empirical evidence for our understanding, we carry out experiments based on multidomain datasets (section 4.1). specifically, we compare the model predictions of knn-mt and explicit opl fine-tuning on five domain adaptation tasks. as expected, the predictions of knn-mt is highly similar to that of explicit opl fine-tuning. these findings support our understanding that knn-mt performs implicit opl fine-tuning. next, we conduct comprehensive multi-domain experiments and word-level analysis to examine the differences in translation performance between knn-mt and other popular fine-tuning methods, such as entire-model fine-tuning and adapter-based fine-tuning (section 4.2 and 4.3). our empirical results suggest that: (i) introducing knn-mt on top of adapter-based fine-tuning obtains comparable translation performance to entire-model finetuning on in-domain test sets, while achieving better performance on out-of-domain test sets. (ii) the entire-model fine-tuning significantly outperforms knn-mt in terms of the recall of in-domain low-frequency words, but this difference can be mitigated by optimizing the context representations with lightweight adapter layers. ## knn-mt performs implicit gradient descent on output projection layer in this section, we first demonstrate that probability distribution in knn-mt, including p knn and p nmt , is equivalent to transformer attention. on top of that, we argue that knn-mt implicitly performs gradient descent on opl, producing meta-gradients via forward computation and interpolation based on k-nearest-neighbors. next, we draw comparisons between knn-mt and explicit opl fine-tuning, establishing connections between these two forms. ## experiments in this section, we begin by comparing the model predictions of knn-mt and explicit opl finetuning (opl-ft) using multi-domain datasets to verify our earlier analysis. then we carry out comprehensive multi-domain experiments and wordlevel analysis to gain a better understanding of the translation performance differences between knn-mt and current popular fine-tuning methods. ## related work retrieval-augmented methods have attracted much attention from the community and achieved remarkable performance on various tasks, including language modeling @xcite @xcite @xcite , machine translation @xcite @xcite @xcite , question answering @xcite @xcite , and dialogue generation @xcite . for the nmt system, @xcite propose knn-mt that utilizes a knn classifier over a large datastore with traditional nmt models @xcite @xcite to achieve significant improvements. recently, several attempts have been made by most researchers to improve the robustness and scalability. @xcite and @xcite propose fast versions of knn-mt. @xcite develop adaptive knn-mt by dynamically determining the number of retrieved tokens k and interpolation λ at each step, while @xcite attempt to retrieve chunks of tokens from the datastore instead of a single token. @xcite adopt a lightweight neural network and the cluster-based pruning method to reduce retrieval redundancy. @xcite improve both decoding speed and storage overhead by dynamically constructing an extremely small datastore and introducing a distance-aware adapter for inference, and further observe the similar behaviours between knn-based methods and translation memory approaches @xcite @xcite . despite the great success of the knn-mt family, the working mechanism of these methods remains an open question. @xcite analyze the relationship between the datastore and nmt model to better understand the behaviour of knn-mt. to the best of our knowledge, we are the first to provide a meta-optimization perspective for knn-mt, i.e., knn-mt performs implicit gradient descent on the output projection layer. ## conclusion in this paper, we present a new meta-optimization perspective to understand knn-mt and establish connections between knn-mt and model finetuning. our results on multi-domain datasets provide strong evidence for the reasonability of this perspective. additional experiments indicate that (i) incorporating knn-mt with adapter-based finetuning achieves comparable translation quality to entire-model fine-tuning, with better performance on out-of-domain test sets; (ii) knn-based models suffer from the low recall of in-domain lowfrequency words, which could be mitigated by optimizing the representation vectors with lightweight adapter layers. we hope our understanding would have more potential to enlighten knn-based applications and model design in the future.