Nearest Neighbor Knowledge Distillation for Neural Machine Translation

Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 5546 - 5556 July 10-15, 2022 ©2022 Association for Computational Linguistics Nearest Neighbor Knowledge Distillation for Neural Machine Translation Zhixian Yang and Renliang Sun and Xiaojun Wan Wangxuan Institute of Computer Technology, Peking University Center for Data Science, Peking University The MOE Key Laboratory of Computational Linguistics, Peking University yangzhixian@stu.pku.edu.cn sunrenliangpku@gmail.com wanxiaojun@pku.edu.cn Abstract k-nearest-neighbor machine translation (kNN- MT), proposed by Khandelwal et al. (2021), has achieved many state-of-the-art results in machine translation tasks. Although effective, kNN-MT requires conducting kNN searches through the large datastore for each decoding step during inference, prohibitively increasing the decoding cost and thus leading to the dif- ficulty for the deployment in real-world appli- cations. In this paper, we propose to move the time-consuming kNN search forward to the preprocessing phase, and then introduce k Nearest Neighbor Knowledge Distillation (kNN-KD) that trains the base NMT model to directly learn the knowledge of kNN. Distilling knowledge retrieved by kNN can encourage the NMT model to take more reasonable tar- get tokens into consideration, thus addressing the overcorrection problem. Extensive experi- mental results show that, the proposed method achieves consistent improvement over the state- of-the-art baselines including kNN-MT, while maintaining the same training and decoding speed as the standard NMT model.1 1 Introduction Neural machine translation (NMT) has shown im- pressive progress with the prevalence of deep neu- ral networks (Vaswani et al., 2017; Zhang et al., 2019; Chen et al., 2020). Recently, Khandelwal et al. (2021) have proposed k-nearest-neighbor ma- chine translation (kNN-MT) that first stores context representations and target tokens into a large data- store, 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 1Our code is available at https://github.com/ FadedCosine/kNN-KD each decoding step needs to conduct kNN searches (Meng et al., 2021). How to speed up the decoding of kNN-MT without degrading performance still remains an open question. Several recent works (Meng et al., 2021; Wang et al., 2021b) introduce some elaborate strategies to compress the datastore in which kNN searches are conducted, thus improv- ing decoding efficiency to some extent. However, we argue that, where there is a time-consuming kNN search in the decoding phase, there is the pro- hibitive 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 nat- ural language, a sentence usually has more than one expression. However, even when the model predicts a word that is reasonable but deviates …

Content truncated to 500 words