Domain-Aware k -Nearest-Neighbor Knowledge Distillation for Machine Translation (2024)

Abstract: k NN-MT has utilized neighborhood knowledge for auxiliary decoding, significantly improving translation performance. Subsequently, k NN-KD transitions the use of neighborhood knowledge from the decoding phase to the training phase, to address the temporal and spatial inefficiencies inherent in k NN-MT. However, k NN-KD transfers all the k NN knowledge arbitrarily, which has the potential to restrict the learning of student models. In this paper, we propose a novel domain-aware k NN-KD method, which filters out domain-relevant neighborhood knowledge for learning in the distillation process. Notably, this entire process exclusively utilizes the neighborhood knowledge of the original model, eliminating the need for establishing any additional datastores. Experiments on four domain translation tasks demonstrate that our method achieves state-of-the-art performance, realizing an average gain of 1.55 COMET and 1.42 BLEU scores, by further enhancing the translation of rare words. Source code can be accessed at https://github.com/wangzx1219/Dk-KD.

## introduction the field of neural machine translation (nmt, @xcite has witnessed significant advancements, resulting in noteworthy improvements in various translation tasks. among these innovations, the introduction of knn-mt @xcite stands out as a pioneering approach. this method leverages neighborhood knowledge for assisted decoding, enhancing the translation capabilities of the model. in knn-mt, the translation process benefits from an expanded contextual understanding, allowing for more accurate and contextually relevant translations. this integration effectively bridges the * equal contribution † corresponding author gap between raw translation output and nuanced, context-aware language understanding. although the enhancement of model translation capability by knn-mt is significant, its retrieval cost in the decoding process and consumption of storage space are non-negligible, which limits the application of knn-mt in practical scenarios. the existing methods try to transfer knowledge from the knn datastore into new models. knn-kd @xcite employs knowledge distillation (kd, @xcite for training the network from scratch, which, in some cases, leads to suboptimal and unstable outcomes. ink @xcite propose to migrate knowledge into the adapter module @xcite , which is more efficient due to its smaller parameter size and generally yields better results. however, existing methods indiscriminately transfer knowledge from the datastore, a practice we believe may have its limitations. building on the pros and cons of the previous methods, this paper introduces domain-aware knn-kd (dk-kd), aimed at extracting the most valuable knowledge from the knn datastore within the domain. the approach begins by training a more refined teacher model, specifically optimized for storing knn knowledge, which facilitates the learning process of the adapter. subsequently, knn knowledge is selectively distilled from the teacher model into the adapter, ensuring targeted and efficient knowledge transfer. specifically, we first construct a datastore using the representations from the original nmt model. subsequently, we utilize domain-relevant knowledge to aid in training the teacher model. finally, we leverage both the teacher and domain-relevant knowledge to assist in training the final model. domain-aware knowledge selection is employed to filter domain-relevant knowledge from the retrieved representations. experiments on four domain translation tasks indicate that dk-kd outperforms other advanced models that utilize knn knowledge for kd. our primary contributions are as follows: building knn datastore the datastore is a crucial aspect of the knn-mt system, where it stores the knowledge of a pre-trained nmt model in the form of key-value pairs. in this context, the key represents the output representation at each time step, and the value corresponds to the target token, which is the accurate translation. specifically, given a set of training data (x , y), we process each sentence pair (x, y) to construct the datastore (k, v). this process can be understood as structuring the knowledge learned by the model in a way that facilitates efficient retrieval and utilization later on. the construction for the datastore is: model inference during the inference process, at each decoding step i, knn-mt transforms the current translation context into a representation, f (x, @xmath0. this representation, f , is then used to query the k nearest neighbors n i @xmath1, v j ) | @xmath2, 2, . . . , k}} from the datastore by comparing the l 2 distance, effectively identifying the k most similar historical contexts to the current one. the distribution for knn is: apply knn to kd let p knn represent the knn distribution retrieved from the datastore. to train an nmt model from scratch using kd, the loss during distillation is: distant teacher ce loss + kl loss relevant teacher (fine-tuned model) relevant teacher retrieve kl loss ce loss distilled student + relevant teacher fine-tuning domain-aware adapter distillation medical it koran law input save build datastore domain-aware knn datastore medical it koran law domain-aware knowledge selection update update distant teacher (vanilla mt model) retrieve medical it koran law adapter figure 1: an overview of our proposed dk-kd method. ## conclusion and future work in this paper, we introduce a method named domain-aware knn-kd, which extracts domainrelevant knowledge from a domain-aware datastore constructed by a domain-agnostic model. in dk-kd, we design a two-stage knowledge distillation process that initially trains a domain-relevant teacher with the aid of domain-relevant knowledge, and then utilizes the teacher model to distill domain-relevant knowledge into the adapter layers, achieving improvements in translation performance. notably, dk-kd significantly enhances the source representation and the translation accuracy of domain-specific low-frequency words. future work includes: 1) attempting to shift the filtering of domain-focused knowledge from the training phase to the datastore construction phase, aiming to reduce the storage footprint of the datastore and alleviate retrieval times; 2) identifying more effective and efficient methods for filtering domain-relevant knowledge to enhance the efficacy of knn knowledge distillation.