Nearest Neighbor Knowledge Distillation for Neural Machine Translation (2022)

Abstract: k-nearest-neighbor machine translation ( k NN-MT), proposed by Khandelwal et al. (2021), has achieved many state-of-the-art results in machine translation tasks. Although effective, k NN-MT requires conducting k NN searches through the large datastore for each decoding step during inference, prohibitively increasing the decoding cost and thus leading to the difficulty for the deployment in real-world applications. In this paper, we propose to move the time-consuming k NN search forward to the preprocessing phase, and then introduce k Nearest Neighbor Knowledge Distillation ( k NN-KD) that trains the base NMT model to directly learn the knowledge of k NN. Distilling knowledge retrieved by k NN can encourage the NMT model to take more reasonable target tokens into consideration, thus addressing the overcorrection problem. Extensive experimental results show that, the proposed method achieves consistent improvement over the state-of-the-art baselines including k NN-MT, while maintaining the same training and decoding speed as the standard NMT model.

## introduction neural machine translation (nmt) has shown impressive progress with the prevalence of deep neural networks @xcite @xcite . recently, @xcite have proposed k-nearest-neighbor machine translation (knn-mt) that first stores context representations and target tokens into a large datastore, and then retrieves k possible target tokens by conducting nearest search from the datastore to help with the final next-token decision. the results show that knn-mt can significantly improve the performance over the base nmt model. despite the outstanding performance, knn-mt will drastically increase the testing runtime since each decoding step needs to conduct knn searches @xcite . how to speed up the decoding of knn-mt without degrading performance still remains an open question. several recent works @xcite introduce some elaborate strategies to compress the datastore in which knn searches are conducted, thus improving decoding efficiency to some extent. however, we argue that, where there is a time-consuming knn search in the decoding phase, there is the prohibitive decoding cost, which makes it hard to be deployed on real-world applications. in order to address the aforementioned issue more thoroughly, it is necessary to figure out why knn-mt performs so well. the standard nmt models are typically trained with cross-entropy (ce) loss with teacher forcing technique, which requires a strict word-by-word matching between the model prediction and the ground-truth. in natural language, a sentence usually has more than one expression. however, even when the model predicts a word that is reasonable but deviates from the ground-truth, the ce loss will treat it as an error and punish the model. this phenomenon is called overcorrection @xcite , which often seriously harms the generalizability of nmt models. we conclude that knn-mt can alleviate the problem of overcorrection by retrieving more reasonable target tokens in the decoding phase. one natural question can be raised: can we train the model to directly learn the knowledge of knn in the training phase, thus maintaining the standard decoding process without any additional decoding cost? to answer this question, we propose k nearest neighbor knowledge distillation (knn-kd) to distill the knowledge of the non-parametric model, i.e., knn, into the base nmt model in the training phase. in detail, we first construct the datastore and then conduct knn searches immediately. these two steps can be done offline in the preprocessing phase. during training, a teacher distribution p t knn can be easily computed using the pre-stored results of knn searches to train the nmt model to directly learn the knowledge of knn. at inference time, knn searches are not required, so the decoding speed is as fast as the base nmt model. therefore, knn-kd can achieve two desirable goals simultaneously: preventing overcorrection (effectiveness) and reducing decoding cost (efficiency). we conduct experiments on two widely acknowledged nmt benchmarks: iwslt'14 german-english and iwslt'15 english-vietnamese. experimental results show that our knn-kd maintains the same training and decoding speed as the standard nmt model, while it outperforms vanilla knn-mt and all the other kd methods, and yields an improvement of +2.14 and +1.51 bleu points over the transformer baseline. we further verify that knn-kd can be adapted to diverse domains by performing experiments on multi-domains translation datasets @xcite ) and achieving 2.56 bleu improvement over vanilla knn-mt on average. in summary, the contributions of our work are as follows: ## methodology the core idea of our work is to enhance the nmt model with a nearest neighbor retrieval mechanism in a training manner, and thus quantitatively evaluated, the model can perform as well or better than vanilla knn-nmt without any additional decoding cost. in section 3.1, we first introduce k nearest neighbor knowledge distillation (knn-kd) to distill the knowledge of knn into a base nmt model. and then, we provide the theoretical analysis in section 3.2 to support that our method can help to address the overcorrection issue. ## conclusion in this paper, we introduce knn-kd that distills the knowledge retrieved by knn search to prevent the base nmt model from overcorrection. experiments show that knn-kd can improve over vanilla knn-mt and other baselines without any additional cost for training and decoding. in the future, we will apply knn-kd to many other tasks. we will also explore the effect of knn-kd on improving the diversity of text generation. iwslt datasets, 250 epochs on koran domain, 100 epochs on medical domain, 120 epochs on law domain, while the early-stop mechanism is also adopted with patience set to 20.