SlowGuess/NAACL_2022_OCR huggingface dataset:
Models	De-En	En-Vi
BLEU	upd/s	token/s	BLEU	upd/s	token/s
Transformer	34.11	2.02(1.00×)	3148.10(1.00×)	30.76	2.55(1.00×)	2870.07(1.00×)
Word-KD	34.26	1.77(0.88×)	3291.28(1.06×)	30.98	2.14(0.84×)	2782.53(0.97×)
Seq-KD	34.60	2.14(1.06×)	3409.86(1.08×)	31.20	2.80(1.10×)	2855.77(1.00×)
BERT-KD	35.63	1.70(0.84×)	3275.43(1.04×)	31.51	2.14(0.84×)	2785.69(0.97×)
Selective-KD	34.38	1.72(0.85×)	3365.68(1.07×)	31.48	2.09(0.82×)	3044.68(1.06×)
kNN-MT	36.17	-	920.72(0.29×)	32.08	-	617.88(0.22×)
kNN-KD	36.30	2.14(1.06×)	3321.24(1.05×)	32.27	2.60(1.02×)	2879.68(1.00×)
Table 2: Experimental results on IWSLT'14 De-En and IWSLT'15 En-Vi translation tasks. “-” means “not applicable” since vanilla kNN-MT can only be adopted in the decoding phase. “upd/s” and “token/s” are abbreviated notations for “training updates per second” and “generated tokens per second”.
Models	Koran	Medical	Law
BLEU	token/s	BLEU	token/s	BLEU	token/s
Pre-trained Model	16.26	1038.97(1.00×)	39.91	1765.56(1.00×)	45.71	2404.31(1.00×)
kNN-MT	19.45	246.17(0.24×)	54.35	701.29(0.40×)	61.78	853.66(0.36×)
Transformer	13.84	1297.45(1.25×)	27.51	1073.53(0.61×)	60.77	1689.89(0.70×)
kNN-KD	24.86	1236.23(1.19×)	56.50	1853.58(1.05×)	61.89	2456.62(1.02×)
Table 3: Experimental results on multi-domains translation datasets. We leave out the metric for training efficiency (upd/s) since it is only applicable for Transformer and kNN-KD, and the training efficiency of these two models are basically the same.
IWSLT datasets, the base NMT model is used to obtain the context vectors, while for multi-domains translation datasets, we follow Khandelwal et al. (2021) to build datastores by the pre-trained model<sup>5</sup>. According to the model configuration, the keys are 512-dimensional and 1024-dimensional for IWSLT datasets and multi-domains translation datasets, respectively. We use FAISS (Johnson et al., 2017) for the nearest neighbor search. And we conduct grid searches over  $k \in \{4,8,16,32,64,128,256,512,1024\}$  and  $\tau \in \{1,10,50,100,200,500,1000\}$ , and choose the final settings according to the best BLEU score on the validation set. The final hyper-parameter settings are shown in Table 1.
Evaluation For all the datasets, we use the beam search with beam size 5. We evaluate the translation in terms of quality and efficiency.
- Quality. For IWSLT'14 De-En and IWSLT'15 En-Vi, following the common practice, we measure case sensitive BLEU by multi-bleu.perl $^{6}$ . For multi-domains trans
lation datasets, we closely follow Khandelwal et al. (2021) to evaluate the results by Sacre-BLEU (Post, 2018) for a fair comparison.
- Efficiency. We evaluate the efficiency of training and inference by the training updates per second (upd/s) and the generated tokens per second (token/s), respectively.
Results of IWSLT Datasets We first compare kNN-KD with vanilla kNN-MT and other KD methods on the two IWSLT translation tasks. Note that there are several hyper-parameters in vanilla kNN-MT: tunable weight  $(\lambda)$ , number of neighbors per query  $(k)$ , and temperature  $(\tau)$ . These hyper-parameters have great effects on the translation results. We also conduct grid searches over these hyper-parameters, and find the best settings according to BLEU score on the validation set.
As shown in Table 2,  $k\mathrm{NN}$ -KD outperforms all the other strong baselines on both IWSLT datasets, e.g., an improvement of +2.14 and +1.51 BLEU score over Transformer. Moreover, we observe that our proposed  $k\mathrm{NN}$ -KD can even perform better than vanilla  $k\mathrm{NN}$ -MT, while gaining a significant
https://github.com/pytorch/fairseq/tree/main/examples/wmt19
$^{6}$ https://github.com/moses-smt/mosesdecoder/blob/master/scripts/
generic/multi-bleu.perl
Models	Law→Medical	Medical→Law
Transformer	18.73	2.07
kNN-KD	22.31	14.82
Table 4: Generalizability Evaluation. "Law  $\rightarrow$  Medical" means that we train the model on the Law domain and directly apply it to Medical domain, and vice versa. The results are BLEU scores.
speedup. On the one hand,  $k$ NN-KD, like other KD methods, maintains the standard Seq2Seq manner at inference time, thus keeping the same decoding speed as Transformer. On the other hand,  $k$ NN-KD also keeps the same training speed as Transformer, and it is more efficient than Word-KD, BERT-KD and Selective-KD. This is because the calculation of the teacher model distribution  $p_{\mathrm{kNN}}^{\mathrm{T}}(y_i \mid \mathbf{x}, \mathbf{y}_{<i}^{\star})$  only needs to be performed on a relatively small  $k$ NN retrieved set  $\mathcal{K}^i$ , while word-level KD have to compute the teacher distribution over the whole vocabulary  $\mathcal{V}$ .
Results of Multi-domains Datasets Apart from IWSLT datasets, we further compare our  $k$ NN-KD with  $k$ NN-MT on multi-domains translation datasets. First, we follow Khandelwal et al. (2021) to conduct inference with the pre-trained model and vanilla  $k$ NN-MT. Then, we train the base NMT model using standard CE and  $k$ NN-KD on each domain's training data, and report the results in Table 3 as a comparison. In all domains,  $k$ NN-KD again outperforms all the baselines. Most importantly, our proposed  $k$ NN-KD can achieve a consistent improvement over vanilla  $k$ NN-MT (+2.56 BLEU score on average) with a significant speedup. This further confirms the effectiveness and efficiency of our method.
Generalizability To verify the generalizability of our method, we further conduct experiments on the scenario that we train a NMT model on a specific domain and evaluate it on the out-of-domain test set. As shown in Table 4, our kNN-KD performs significantly better than Transformer trained by standard CE. It proves the statement in Section 1 that compared with standard CE, kNN-KD can improve the generalizability of NMT models.
There are two key hyper-parameters in our kNN-KD: number of neighbors per query  $(k)$ , and temperature  $(\tau)$ . In this section, we investigate the effects of these two hyper-parameters on the vali
Figure 2: BLEU scores with different  $k$  and fixed  $\tau$ $(\tau = 100)$  on the validation set of IWSLT'14 De-En dataset.
Figure 3: BLEU scores with different  $\tau$  and fixed  $k$ $(k = 64)$  on the validation set of IWSLT'14 De-En dataset.
dition set of IWSLT'14 De-En.
Effect of  $k$  We fix the temperature  $\tau$  to 100, and train the model using kNN-KD with different  $k$ . As shown in Figure 2, the BLEU score first rises with the increase of  $k$ , and reaches the best performance peak when  $k = 64$ . And then, performance deteriorates with a larger  $k$ . This suggests that, the retrieved results of kNN search can substantially improve training when  $k$  is relatively small, but it will also introduce more noise when  $k$  gets larger.
Effect of  $\tau$  We train the model using  $k\mathrm{NN}$ -KD with different  $\tau$  and fixed  $k$  ( $k = 64$ ). As shown in Figure 3, a temperature of 1 results in a significantly lower BLEU score than those greater than 1. This is because a large temperature value can flatten the  $k\mathrm{NN}$  teacher distribution in Equation 10 to prevent assigning most of the probability mass to a single neighbor. The results show that for  $k = 64$ , the optimal temperature is 100.
Figure 4: Predicted probabilities output from the base NMT model, kNN-MT and our kNN-KD, given the translation context "es gibt eine mange geschickt davon, warum wir dies getan haben. || there are"
In this section, we show how our proposed method works by presenting a real case. There exists an example in the test set of IWSLT'14 De-En that the source sentence is "es gibt eineMMC geschickt darüber, warum wir dies getan haben." and the corresponding target sentence is "there are a lot of stories about why we did this". Given the source sentence and target subsequence "there are" as the translation context, "many...", "lots of...", and "a lot of..." are all correct translations. We input this translation context to the base NMT model, kNN-MT, and our model, and observe the predicted probabilities over the vocabulary. As shown in Figure 4, all the models predict "a" with the maximal probability that matches the ground-truth. However, since the base model is trained by CE loss using one-hot vector as supervision, it suffers from a serious overcorrection problem that the model assigns an overconfident probability to the token "a" and almost none to other reasonable target tokens such as "lots" and "many". On the contrary, both kNN-MT and our kNN-KD increase the probabilities of the reasonable target tokens, and these two models have similar predicted probabilities. Note that kNN-MT obtains this probability distribution by interpolating the base NMT probability with a kNN search probability at decoding time, while our kNN-KD directly outputs this distribution without any additional operations. This further confirms that kNN-KD can train the model to learn the knowledge of kNN that prevents the overconfidence of the model on the one-hot label, thus
leading to the better generalizability for inference.
5.1 Neural Machine Translation
Machine translation has developed rapidly in recent years. The early models were mainly based on statistical machine learning (Brown et al., 1990; Och, 2003; Koehn et al., 2007). Then, with the development of deep learning technology, many models used RNN(Sutskever et al., 2014; Bahdanau et al., 2015), CNN(Gehring et al., 2017), or Transformer(Vaswani et al., 2017) as their backbones.
Recently, a few studies have combined  $k$  nearest neighbors algorithm closely with NMT models to improve performance. Khandelwal et al. (2021) used a nearest neighbor classifier to predict tokens on a large datastore of cached examples and proposed kNN-MT. However, Meng et al. (2021) pointed out that kNN-MT is two-order slower than vanilla MT models, which limits the deployment for real-world applications. They proposed Fast kNN-MT to solve this problem. Wang et al. (2021b) also noticed the low-efficiency problem of kNN-MT. Thus, they used a hierarchical clustering strategy and proposed Faster kNN-MT. Although the above studies have made feasible fixes, kNN search is still required in the decoding phase, which dramatically increases the difficulty of practical applications compared to standard MT models.
5.2 Knowledge Distillation
Knowledge distillation (KD) introduces teacher network and student network to help knowledge transfer and it was widely used in NMT (Hinton et al., 2015a). Kim and Rush (2016) introduced two sequence-level KD methods to improve the performance of NMT. Barone et al. (2017) used KD to address the problem of catastrophic forgetting in the fine-tuning stage. Tan et al. (2019) used KD to enable the multilingual model to fit the training data and to match the outputs of the teacher models. Clark et al. (2019) distilled single-task models into one multi-task model. Chen et al. (2020) used BERT as the teacher model after fine-tuning on the target generation tasks to improve the conventional Seq2Seq models. Wang et al. (2021a) proposed batch-level and global-level selection strategies to choose appropriate samples for knowledge distillation. We focus on using KD to leverage the knowledge retrieved by kNN search to enhance a base NMT model.
In this paper, we introduce  $k$  NN-KD that distills the knowledge retrieved by  $k$  NN search to prevent the base NMT model from overcorrection. Experiments show that  $k$  NN-KD can improve over vanilla  $k$  NN-MT and other baselines without any additional cost for training and decoding. In the future, we will apply  $k$  NN-KD to many other tasks. We will also explore the effect of  $k$  NN-KD on improving the diversity of text generation.
Roee Aharoni and Yoav Goldberg. 2020. Unsupervised domain clusters in pretrained language models. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, ACL 2020, Online, July 5-10, 2020, pages 7747-7763. Association for Computational Linguistics.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. 2015. Neural machine translation by jointly learning to align and translate. In 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings.
Antonio Valerio Miceli Barone, Barry Haddow, Ulrich Germann, and Rico Sennrich. 2017. Regularization techniques for fine-tuning in neural machine translation. In Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing, EMNLP 2017, Copenhagen, Denmark, September 9-11, 2017, pages 1489-1494. Association for Computational Linguistics.
Peter F. Brown, John Cocke, Stephen A. Della Pietra, Vincent J. Della Pietra, Fredrick Jelinek, John D. Lafferty, Robert L. Mercer, and Paul S. Roossin. 1990. A statistical approach to machine translation. Computational Linguistics, 16(2):79-85.
Yen-Chun Chen, Zhe Gan, Yu Cheng, Jingzhou Liu, and Jingjing Liu. 2020. Distilling knowledge learned in BERT for text generation. In Proceedings of the
58th Annual Meeting of the Association for Computational Linguistics, ACL 2020, Online, July 5-10, 2020, pages 7893-7905. Association for Computational Linguistics.
Kevin Clark, Minh-Thang Luong, Urvashi Khandelwal, Christopher D Manning, and Quoc Le. 2019. Bam! born-again multi-task networks for natural language understanding. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 5931-5937.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers), pages 4171-4186.
Jonas Gehring, Michael Auli, David Grangier, Denis Yarats, and Yann N Dauphin. 2017. Convolutional sequence to sequence learning. In Proceedings of the 34th International Conference on Machine Learning-Volume 70, pages 1243-1252.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. 2015a. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531.
Geoffrey E. Hinton, Oriol Vinyals, and Jeffrey Dean. 2015b. Distilling the knowledge in a neural network. CoRR, abs/1503.02531.
Jeff Johnson, Matthijs Douze, and Hervé Jégou. 2017. Billion-scale similarity search with gpus. CoRR, abs/1702.08734.
Urvashi Khandelwal, Angela Fan, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. 2021. Nearest neighbor machine translation. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net.
Yoon Kim and Alexander M. Rush. 2016. Sequence-level knowledge distillation. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, EMNLP 2016, Austin, Texas, USA, November 1-4, 2016, pages 1317-1327. The Association for Computational Linguistics.
Philipp Koehn, Hieu Hoang, Alexandra Birch, Chris Callison-Burch, Marcello Federico, Nicola Bertoldi, Brooke Cowan, Wade Shen, Christine Moran, Richard Zens, Chris Dyer, Ondrej Bojar, Alexandra Constantin, and Evan Herbst. 2007. Moses: Open source toolkit for statistical machine translation. In ACL 2007, Proceedings of the 45th Annual Meeting of the Association for Computational Linguistics, June 23-30, 2007, Prague, Czech Republic. The Association for Computational Linguistics.
Xiang Lin, Simeng Han, and Shafiq R. Joty. 2021. Straight to the gradient: Learning to use novel tokens for neural text generation. In Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 6642-6653. PMLR.
Minh-Thang Luong and Christopher D. Manning. 2015. Stanford neural machine translation systems for spoken language domains. In Proceedings of the 12th International Workshop on Spoken Language Translation: Evaluation Campaign@IWSLT 2015, Da Nang, Vietnam, December 3-4, 2015.
Yuxian Meng, Xiaoya Li, Xiayu Zheng, Fei Wu, Xiaofei Sun, Tianwei Zhang, and Jiwei Li. 2021. Fast nearest neighbor machine translation. abs/2105.14528.
Franz Josef Och. 2003. Minimum error rate training in statistical machine translation. In Proceedings of the 41st annual meeting of the Association for Computational Linguistics, pages 160-167.
Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. 2019. *fairseq: A fast, extensible toolkit for sequence modeling*. In *Proceedings of NAACL-HLT* 2019: Demonstrations.
Matt Post. 2018. A call for clarity in reporting BLEU scores. In Proceedings of the Third Conference on Machine Translation: Research Papers, pages 186-191, Belgium, Brussels. Association for Computational Linguistics.
Rico Sennrich, Barry Haddow, and Alexandra Birch. 2016. Neural machine translation of rare words with subword units. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics, ACL 2016, August 7-12, 2016, Berlin, Germany, Volume 1: Long Papers. The Association for Computer Linguistics.
Ilya Sutskever, Oriol Vinyals, and Quoc V Le. 2014. Sequence to sequence learning with neural networks. In Advances in neural information processing systems, pages 3104-3112.
Xu Tan, Yi Ren, Di He, Tao Qin, Zhou Zhao, and TieYan Liu. 2019. Multilingual neural machine translation with knowledge distillation. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In *In Advances in Neural Information Processing Systems*, pages 5998-6008.
Fusheng Wang, Jianhao Yan, Fandong Meng, and Jie Zhou. 2021a. Selective knowledge distillation for neural machine translation. In Proceedings of the
59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing, ACL/IJCNLP 2021, (Volume 1: Long Papers), Virtual Event, August 1-6, 2021, pages 6456-6466. Association for Computational Linguistics.
Shuhe Wang, Jiwei Li, Yuxian Meng, Rongbin Ouyang, Guoyin Wang, Xiaoya Li, Tianwei Zhang, and Shi Zong. 2021b. Faster nearest neighbor machine translation. CoRR, abs/2112.08152.
Wen Zhang, Yang Feng, Fandong Meng, Di You, and Qun Liu. 2019. Bridging the gap between training and inference for neural machine translation. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics. Association for Computational Linguistics.
Xin Zheng, Zhirui Zhang, Junliang Guo, Shujian Huang, Boxing Chen, Weihua Luo, and Jiajun Chen. 2021. Adaptive nearest neighbor machine translation. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing, ACL/IJCNLP 2021, (Volume 2: Short Papers), Virtual Event, August 1-6, 2021, pages 368-374. Association for Computational Linguistics.
The dataset statistics for all the datasets are reported in Table 5. It is worth to mention that IWSLT datasets are under the Creative Commons BY-NC-ND license, and the multi-domains translation datasets are under the BSD license.
Train	Valid	Test
IWLST‘14 De-En	160,239	7,283	6,750
IWLST‘15 En-Vi	133,166	1,553	1,268
Koran	17,982	2,000	2,000
Medical	248,099	2,000	2,000
Law	467,309	2,000	2,000
Table 5: The number of examples for different datasets.
A.2 Hyper-parameters Setting
All the algorithms are implemented in Pytorch with fairseq toolkit (Ott et al., 2019), and all the experiments are conducted on a machine with 8 NVIDIA GTX 1080Ti GPUs with the hyperparameters reported in Table 6.
Note that during training, we are using the dynamic batching provided by fairseq, and choose the max tokens according to the GPU memory constraint. We train the model for 200 epochs on
Hyperparameters	IWSLT	Multi-domains
Max tokens	8192	1280
Learning rate	5e-4	5e-4
LR scheduler	Inverse sqrt	Inverse sqrt
Minimal LR	1e-9	1e-9
Warm-up LR	1e-7	1e-7
Warm-up steps	4000	4000
Gradient clipping	0.0	0.0
Weight decay	0.0	0.0001
Droupout	0.3	0.2
Attention dropout	0.0	0.1
Activation dropout	0.0	0.1
α in Equation 12	0.5	0.5
Optimizer	Adam	Adam
-β1	0.9	0.9
-β2	0.98	0.98
-ε	1e-8	1e-8
Table 6: Hyperparameter settings for different datasets.
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.
B Limitation and Potential Risks
Although  $k$  NN-KD is efficient in both training and inference, it will take a relatively long time for preprocessing to build the datastore and conduct  $k$  NN searches, and it also requires large disk space to store all these results. However, since the preprocessing can be done offline, it does not limit the deployment of  $k$  NN-KD in real-world applications.
Our model is trained on open source datasets, and thus if there exists toxic text in the training data, our model may have the risk of producing toxic content.