Venue: NAACL 2022

# Nearest Neighbor Knowledge Distillation for Neural Machine Translation

 Zhixian Yang and Renliang Sun & 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 (\(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.1

Footnote 1: Our code is available at [https://github.com/FadedCosine/kNN-KD](https://github.com/FadedCosine/kNN-KD)

+
Footnote †: _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

## 1 Introduction

Neural machine translation (NMT) has shown impressive progress with the prevalence of deep neural networks Vaswani et al. (2017); Zhang et al. (2019); Chen et al. (2020). Recently, Khandelwal et al. (2021) have proposed \(k\)-nearest-neighbor machine translation (\(k\)NN-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 \(k\)NN-MT can significantly improve the performance over the base NMT model.

Despite the outstanding performance, \(k\)NN-MT will drastically increase the testing runtime since each decoding step needs to conduct \(k\)NN searches Meng et al. (2021). How to speed up the decoding of \(k\)NN-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 \(k\)NN searches are conducted, thus improving decoding efficiency to some extent. However, we argue that, where there is a time-consuming \(k\)NN 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 \(k\)NN-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_Zhang et al. (2019), which often seriously harms the generalizability of NMT models. We conclude that \(k\)NN-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 \(k\)NN 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 (\(k\)NN-KD) to distill the knowledge of the non-parametric model, i.e., \(k\)NN, into the base NMT model in the training phase. In detail, we first construct the datastore and then conduct \(k\)NN searches immediately. These two steps can be done offline in the preprocessing phase. During training, a teacher distribution

can be easily computed using the pre-stored results of \(k\)NN searches to train the NMT model to directly learn the knowledge of \(k\)NN. At inference time, \(k\)NN searches are not required, so the decoding speed is as fast as the base NMT model. Therefore, \(k\)NN-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 \(k\)NN-KD maintains the same training and decoding speed as the standard NMT model, while it outperforms vanilla \(k\)NN-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 \(k\)NN-KD can be adapted to diverse domains by performing experiments on multi-domains translation datasets [3] and achieving \(2.56\) BLEU improvement over vanilla \(k\)NN-MT on average.

In summary, the contributions of our work are as follows:

* We propose \(k\)NN-KD that considers the \(k\)NN distribution as a teacher to guide the training of the base NMT model (Section 3.1).
* We prove that our proposed \(k\)NN-KD can help to address the overcorrection issue with theoretical analysis (Section 3.2).
* Quantitative and qualitative results on different translation tasks validate the effectiveness and efficiency of our method (Section 4).

## 2 Background

### Neural Machine Translation

The goal of the standard NMT model is to learn the conditional probability \(p_{\mathrm{MT}}\left(\mathbf{y}~{}|~{}\mathbf{x}\right)\) for translating a sentence \(\mathbf{x}=\left\{x_{1},\cdots,x_{m}\right\}\) in source language to a sentence \(\mathbf{y}=\left\{y_{1},\cdots,y_{n}\right\}\) in target language. Translation is usually performed in a autoregressive manner, and its probability can be factored as \(p_{\mathrm{MT}}\left(\mathbf{y}~{}|~{}\mathbf{x}\right)=\Pi^{n}_{i=1P}\left( y_{i}~{}|~{}\mathbf{x},\mathbf{y}_{<i}\right)\). When predicting \(i\)-th token in the target sentence given \(\left(\mathbf{x},\mathbf{y}_{<i}\right)\) as the _translation context_, the NMT model encodes the translation context into a hidden state \(h_{i-1}\), and outputs a probability distribution over vocabulary \(\mathcal{V}\) as follows:

\[p_{\mathrm{MT}}\left(y_{i}~{}|~{}\mathbf{x},\mathbf{y}_{<i}\right)=\frac{ \exp(\mathbf{o}_{y_{i}}^{\top}\mathbf{h}_{i-1})}{\sum_{w\in\mathcal{V}}\exp( \mathbf{o}_{w}^{\top}\mathbf{h}_{i-1})}, \tag{1}\]

where \(\mathbf{o}_{y}\) is the output embedding for \(w\in\mathcal{V}\).

We denote the ground-truth target sentence as \(\mathbf{y}^{\star}=\left\{y_{1}^{\star},\cdots,y_{n}^{\star}\right\}\), and for each \(y_{i}^{\star}\) in the training set, the CE loss is usually used for optimizing NMT models:

\[\mathcal{L}_{\mathrm{CE}}=-\sum_{y_{i}\in\mathcal{V}}\mathbbm{1}_{y_{i}=y_{i} ^{\star}}\log p_{\mathrm{MT}}\left(y_{i}~{}|~{}\mathbf{x},\mathbf{y}_{<i}^{ \star}\right), \tag{2}\]

where \(\mathbbm{1}\) is the indicator function, and the ground-truth target sequence \(\mathbf{y}_{<i}^{\star}\) is used in the conditions of \(p_{\mathrm{MT}}\) due to the teacher forcing technique.

### Nearest Neighbor Machine Translation

\(k\)NN-MT applies the nearest neighbor retrieval mechanism to the decoding phase of a NMT model, which allows the model direct access to a large-scale datastore for better inference. Specifically, \(k\)NN-MT includes two following steps:

**Datastore Building** Given a bilingual sentence pair in the training set \(\left(\mathbf{x},\mathbf{y}^{\star}\right)\in\left(\mathcal{X},\mathcal{Y}^{ \star}\right)\), \(k\)NN-MT first constructs a datastore \(\mathcal{D}\) as follows:

\[\small\begin{split}\left(\mathcal{K},\mathcal{V}\right)& =\bigcup_{\left(\mathbf{x},\mathbf{y}^{\star}\right)\in\left( \mathcal{X},\mathcal{Y}^{\star}\right)}\left\{\left(f\left(\mathbf{x}, \mathbf{y}_{<i}^{\star}\right),y_{i}^{\star}\right),\forall y_{i}^{\star} \in\mathbf{y}^{\star}\right\},\\ \end{split} \tag{3}\]

where the keys are the mapping representations of all the translation contexts in the training set using the projection \(f(\cdot)\), and the values are corresponding ground-truth tokens.

**Decoding** During inference, \(k\)NN-MT aims to interpolate the base NMT model's probability in Equation 1 with a \(k\)NN model. At each decoding step \(i\), \(k\)NN-MT maps the current translation context to a representation \(f\left(\mathbf{x},\mathbf{y}_{<i}\right)\), which is used to query the datastore for \(k\) nearest neighbors according to the \(l_{2}\) distances. Denote the retrieved neighbors as \(\mathcal{N}^{\tilde{\imath}}=\left\{\left(\mathbf{k}_{j},v_{j}\right),j\in \left\{1,2,\ldots,k\right\}\right\}\), and then a \(k\)NN distribution over vocabulary \(\mathcal{V}\) can be computed as:

\[\small\begin{split}& p_{\mathrm{kNN}}(y_{i}~{}|~{}\mathbf{x}, \mathbf{y}_{<i})\propto\\ &\sum_{\left(\mathbf{k}_{j},v_{j}\right)\in\mathcal{N}^{\tilde{ \imath}}}\mathbbm{1}_{y_{i}=v_{j}}\exp\left(\frac{-d\left(\mathbf{k}_{j},f \left(\mathbf{x},\mathbf{y}_{<i}\right)\right)}{\tau}\right),\end{split} \tag{4}\]

where \(\tau\) is the temperature, and \(d(\cdot,\cdot)\) is the \(l_{2}\) distance function. The final probability for the next token in \(k\)NN-MT is the interpolation 

\(p_{\mathrm{MT}}\left(y_{i}~{}|~{}\mathbf{x},\mathbf{y}_{<i}\right)\) and \(p_{\mathrm{kNN}}\left(y_{i}~{}|~{}\mathbf{x},\mathbf{y}_{<i}\right)\) with a tunable weight \(\lambda\):

\[\begin{split} p\left(y_{i}~{}|~{}\mathbf{x},\mathbf{y}_{<i}\right)& =(1-\lambda)p_{\mathrm{MT}}\left(y_{i}~{}|~{}\mathbf{x},\mathbf{ y}_{<i}\right)\\ &+\lambda p_{\mathrm{kNN}}\left(y_{i}~{}|~{}\mathbf{x},\mathbf{y} _{<i}\right).\end{split} \tag{5}\]

Note that each decoding step of each beam requires a \(k\)NN search over the whole datastore \(\mathcal{D}\), whose time complexity is \(\mathcal{O}(|\mathcal{D}|Bn)\) where \(B\) is the beam size, and \(n\) is the target length. The prohibitive decoding cost makes it hard for \(k\)NN-MT to be deployed on real-world applications.

### Knowledge Distillation

Knowledge Distillation (KD Hinton et al. (2015b) refers to the transfer of knowledge from one neural network T (called "teacher model") to another network S (called "student model").

For convenience, we introduce the details of KD from the perspective of machine translation. Let \(\mathbf{z}\in\mathcal{R}^{|\mathcal{V}|}\) denote the logits over \(\mathcal{V}\). Student model S outputs the probability:

\[p^{\text{S}}\left(y_{i}~{}|~{}\mathbf{x},\mathbf{y}_{<i}\right)=\frac{\exp \left(\mathbf{z}_{y_{i}}\right)}{\sum_{w\in\mathcal{V}}\exp\left(\mathbf{z}_{ w}\right)}, \tag{6}\]

where \(\mathbf{z}_{w}\) is the logit for token \(w\). Correspondingly, teacher model T also predicts the probability in the same way, and a temperature factor \(\tau\) can be introduced to soften the teacher's outputs as:

\[p^{\text{T}}\left(y_{i}~{}|~{}\mathbf{x},\mathbf{y}_{<i}\right)=\frac{\exp \left(\mathbf{z}_{y_{i}}/\tau\right)}{\sum_{w\in\mathcal{V}}\exp\left(\mathbf{ z}_{w}/\tau\right)}. \tag{7}\]

When \(\tau\rightarrow\infty\), \(p^{\text{T}}\) degenerates into a uniform distribution, and when \(\tau\to 0\), \(p^{\text{T}}\) becomes an one-hot vector. Specifically, KD defines the objective as:

\[\begin{split}\mathcal{L}_{\mathrm{KD}}=&-\sum_{y_{ i}\in\mathcal{V}}p^{\text{T}}\left(y_{i}~{}|~{}\mathbf{x},\mathbf{y}_{<i}^{ \star}\right)\\ &\times\log p^{\text{S}}\left(y_{i}~{}|~{}\mathbf{x},\mathbf{y}_ {<i}^{\star}\right).\end{split} \tag{8}\]

When we apply KD to improve the performance of machine translation, student model S is usually the NMT model that will be used for testing. And then, the overall training procedure is to minimize the summation of Equation 2 and Equation 8:

\[\mathcal{L}=(1-\alpha)\mathcal{L}_{\mathrm{CE}}+\alpha\mathcal{L}_{\mathrm{KD}}, \tag{9}\]

where \(\alpha\) is a weight to balance two losses.

## 3 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 \(k\)NN-NMT without any additional decoding cost. In Section 3.1, we first introduce \(k\) Nearest Neighbor Knowledge Distillation (\(k\)NN-KD) to distill the knowledge of \(k\)NN 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.

### Nearest Neighbor Knowledge Distillation

When we apply vanilla \(k\)NN-MT for testing using beam search with \(B\), the time complexity of it is \(\mathcal{O}(|\mathcal{D}|Bn)\). Compared with the standard beam search whose time complexity is \(\mathcal{O}(|\mathcal{V}|Bn)\), the decoding speed of vanilla \(k\)NN-MT is prohibitively slow. This is mainly because vanilla \(k\)NN-MT has to conduct a \(k\)NN search over an extremely large datastore \(\mathcal{D}\) for each decoding step of each beam. We propose to move this time-consuming search process forward to the preprocessing phase which can be done offline before training and inference. As shown in Figure 1, our proposed \(k\)NN-KD can be described as follows:

Figure 1: Illustration of \(k\)NN-KD. In the preprocessing phase, we finish the datastore building in Step \(1\), and conduct \(k\)NN search in advance in Step \(2\). These two steps can be done offline before training and inference. During training, we compute the \(k\)NN distribution as a teacher to train the base NMT model in Step \(3\). During inference, the model performs Step \(4\) to decode text in the standard Seq2Seq manner, which is omitted in this figure.

 



**Step 1: Datastore Building** We build the datastore \({\cal D}\) in the same way as vanilla \(k\)NN-MT Khandelwal et al. (2021) which has been described in Section 2.2, so we omit it here.

**Step 2: \(k\)NN Search in Advance** For all the translation contexts \(\left(\mathbf{x},\mathbf{y}_{<i}^{\star}\right)\) in the training set, we conduct a \(k\)NN search using \(f\left(\mathbf{x},\mathbf{y}_{<i}^{\star}\right)\) as a query to search through the datastore \({\cal D}\) built in Step 1, and then we obtain the retrieved results \(\mathcal{N}^{i}=\left\{\left(\mathbf{k}_{j},v_{j}\right),j\in\left\{1,2, \ldots,k\right\}\right\}\). Note that we are performing \(k\)NN search for **training set** translation contexts on the datastore built with the **training set**, which is equivalent to extending the training data by adding \(k\) reasonable target tokens for every translation context. Formally, by conducting \(k\)NN search in advance, we extend the target sentence in the training set from \(\mathbf{y}^{\star}=\left\{y_{1}^{\star},\cdots,y_{n}^{\star}\right\}\) to \(\mathbf{y}^{\star}=\left\{\left(y_{1}^{\star},\mathcal{K}^{1}\right),\cdots, \left(y_{n}^{\star},\mathcal{K}^{n}\right)\right\}\), where \(\mathcal{K}^{i}=\left\{\left(d\left(\mathbf{k}_{j},f\left(\mathbf{x},\mathbf{ y}_{<i}^{\star}\right)\right),v_{j}\right),j\in\left\{1,2,\ldots,k\right\}\right\}\).

**Step 3: \(k\)NN as a Teacher** In the training phase, a \(k\)NN distribution can be formulated as:

\[p_{\mathrm{kNN}}^{\mathrm{T}}\left(y_{i}~{}|~{}\mathbf{x}, \mathbf{y}_{<i}^{\star}\right)\propto\\ \sum_{(d_{j},v_{j})\in\mathcal{K}^{i}}\mathbbm{1}_{y_{i}=v_{j}} \exp\left(\frac{-d_{j}}{\tau}\right), \tag{10}\]

We then use \(p_{\mathrm{kNN}}^{\mathrm{T}}\) as a teacher to train the base NMT model, and the knowledge distillation objective in Equation 8 can be rewritten as:

\[\begin{split}\mathcal{L}_{\mathrm{kNN-KD}}=&-\sum_{y _{i}\in\mathcal{V}}p_{\mathrm{kNN}}^{\mathrm{T}}\left(y_{i}~{}|~{}\mathbf{x}, \mathbf{y}_{<i}^{\star}\right)\\ &\times\log p_{\mathrm{MT}}\left(y_{i}~{}|~{}\mathbf{x},\mathbf{ y}_{<i}^{\star}\right).\end{split} \tag{11}\]

And the final training objective in Equation 9 can be rewritten as:

\[\mathcal{L}=(1-\alpha)\mathcal{L}_{\mathrm{CE}}+\alpha\mathcal{L}_{\mathrm{ kNN-KD}}, \tag{12}\]

where \(\mathcal{L}_{\mathrm{CE}}\) can be calculated as Equation 2.

**Step 4: Decoding** During inference, our model remains in the standard Seq2Seq manner Vaswani et al. (2017), so we omit it here.

### Theoretical Analysis

In this section, we show that our proposed \(k\)NN-KD can help address the overcorrection issue from the perspective of gradients. The gradient of the final objective in Equation 12 with respect to the logit \(\mathbf{z}_{y_{i}},g_{i}\in\mathcal{V}\) is:

\[\begin{split}\frac{\partial\mathcal{L}}{\partial\mathbf{z}_{y_{ i}}}&=(1-\alpha)\left(p(y_{i})-\mathbbm{1}_{y_{i}=y_{i}^{\star}} \right)+\alpha\left(p(y_{i})-p^{\mathrm{T}}(y_{i})\right)\\ =&\begin{cases}p(y_{i})-\alpha p^{\mathrm{T}}(y_{i}) ,&\text{if $y_{i}\neq y_{i}^{\star}$ and $y_{i}\in\mathcal{ K}^{i}$}\\ p(y_{i})-\left(1-\alpha+\alpha p^{\mathrm{T}}(y_{i})\right),&\text{if $y_{i} \neq y_{i}^{\star}$ and $y_{i}\notin\mathcal{ K}^{i}$}\\ 0,&\text{otherwise}\end{cases}\end{split} \tag{13}\]

where we abbreviate \(p_{\mathrm{MT}}\left(y_{i}~{}|~{}\mathbf{x},\mathbf{y}_{<i}^{\star}\right)\) to \(p(y_{i})\) and \(p_{\mathrm{kNN}}^{\mathrm{T}}\left(y_{i}~{}|~{}\mathbf{x},\mathbf{y}_{<i}^{ \star}\right)\) to \(p^{\mathrm{T}}(y_{i})\).

For every gradient update in the training phase, the model is trained to decrease the gradient norm to \(0\) to reach a local minimum Lin et al. (2021). Therefore, for the tokens that are reasonable but not ground-truth (i.e., \(y_{i}\neq y_{i}^{\star}\) and \(y_{i}\in\mathcal{K}^{i}\)), the model has to learn to increase the probability \(p(y_{i})\) by the degree of \(\alpha p^{\mathrm{T}}(y_{i})\) so that the gradient norm \(|p(y_{i})-\alpha p^{\mathrm{T}}(y_{i})|\) can reach \(0\). For the other nonground-truth token (i.e., \(y_{i}\neq y_{i}^{\star}\) and \(y_{i}\notin\mathcal{K}^{i}\)), \(p^{\mathrm{T}}(y_{i})\) is equal to \(0\) since \(y_{i}\) is not included in the retrieved results of \(k\)NN search, and the model will learn to assign much lower probability \(p(y_{i})\) to reduce \(|p(y_{i})|\). Besides, since we build the datastore and conduct \(k\)NN search on the same training set data, for any translation context, its nearest neighbor over the datastore must be itself, which means if \(y_{i}=y_{i}^{\star}\), then \(y_{i}\in\mathcal{K}^{i}\). Then, for the ground-truth token (i.e., \(y_{i}=y_{i}^{\star}\)), the model is trained to increase the probability \(p(y_{i})\) by the degree of \(\big{(}1-\alpha+\alpha p^{\mathrm{T}}(y_{i})\big{)}\). Note that, the gradient norm of the standard CE loss is \(|p(y_{i})-1|\) for \(y_{i}=y_{i}^{\star}\), and thus that standard CE increases the probability \(p(y_{i})\) by the degree of \(1\). This demonstrates that our \(k\)NN-KD still makes the model learn to predict the ground-truth but with a relatively lower strength than the standard CE.

Taking the case in Figure 1 as an example, given the translation context _"Vielen Dark fair Ihren hilfreichen Vorschlag || Thanks for your"_, its ground-truth target token is _"useful"_, while _"helpful"_ is also reasonable for this translation. Assuming that we have conducted the \(k\)NN search with \(k=3\) in advance as shown in Figure 1, and set \(\tau\) to \(1\), we can then compute the \(k\)NN teacher distribution as:

\[p^{T}(y_{4})=\begin{cases}0.378,&\text{if $y_{4}$ is ``useful''}\\ 0.622,&\text{if $y_{4}$ is ``helpful''}\\ 0,&\text{otherwise}\end{cases} \tag{14}\]

According to Equation 13, the gradient norms are \(|p(\textit{``helpful''})-0.622\alpha|\) for _"helpful"_, and \(|p(\textit{``useful''})-(1-0.622\alpha)|\) for _"useful"_. Therefore, our \(k\)NN-KD can train the model to learn

from the \(k\)NN model to increase the probability of "_helpful"_ that is reasonable but not ground-truth, thus addressing the overcorrection issue.

## 4 Experiments

### Datasets

We conduct experiments on IWSLT'14 German-English (De-En, \(160k\) training samples), IWSLT'15 English-Vietnamese (En-Vi, \(113k\) training samples), and multi-domains translation datasets [2]De-En, \(733k\) training samples). For IWSLT'14 De-En, we follow the preprocessing steps provided by fairseq[9]Ott et al. (2019) to split the data, and process the text into bytepair encoding (BPE) [8]. For IWSLT'15 En-Vi, we use the pre-processed dataset3 provided by Luong and Manning (2015). We use test2012 as the validation set and tst2013 as the test set, which contains \(1,553\) and \(1,268\) sentences respectively. For multi-domains translation datasets, we use the pre-processed dataset4 provided by Zheng et al. (2021), and consider domains including _Koran_, _Medical_, and _Law_ in our experiments.

Footnote 2: [https://github.com/pytorch/fairseq/blob/main/examples/translation/prepare-iwslt14.shttps:nlp.stanford.edu/projects/nmt/](https://github.com/pytorch/fairseq/blob/main/examples/translation/prepare-iwslt14.shttps:nlp.stanford.edu/projects/nmt/)

Footnote 3: [https://github.com/zhengxxn/adaptive-knn-mt](https://github.com/zhengxxn/adaptive-knn-mt)

### Competitive Models

The proposed \(k\)NN-KD is an architecture-free method that can be applied to arbitrary Seq2Seq models, which is orthogonality to previous works that design delicate structures to improve performance. Therefore, we mainly compare \(k\)NN-KD with vanilla \(k\)NN-MT and some typical KD methods:

* **Word-KD**[10]. As described in Section 2.3, Word-KD is the standard KD that distills knowledge equally for each word.
* **Seq-KD**[11]. In this method, teacher model T first generates an extra dataset by running beam search and taking the highest-scoring sequence. Then student model S is trained on this teacher-generated data, and the training objective can be formulated as: \[{\cal L}_{\rm Seq-KD}=-\sum_{i=1}^{n}\sum_{y_{i}\in{\bf V}}\mathbbm{1}_{ y_{i}=\hat{y_{i}}}\] (15) \[~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{ }~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{ }~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{ }~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{ }~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} ~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} ~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} ~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} ~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} ~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} ~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} ~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} ~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} ~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} ~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} ~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} ~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} ~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} ~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} ~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} ~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} ~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} ~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} }~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} {}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} {}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} {}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} {}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} {}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} {}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} {}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} {}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} {}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} }~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{} }~{ 

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 model5. 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.

Footnote 5: [https://github.com/pytorch/fairseq/tree/main/examples/wmt19](https://github.com/pytorch/fairseq/tree/main/examples/wmt19)

**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.perl6_. For multi-domains translation datasets, we closely follow Khandelwal et al. (2021) to evaluate the results by SacreBLEU Post (2018) for a fair comparison. Footnote 6: [https://github.com/moses-smt/mosesdecoder/blob/master/scripts/generic/multi-bleu.perl](https://github.com/moses-smt/mosesdecoder/blob/master/scripts/generic/multi-bleu.perl)
* **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.

### Main Results

Results of IWSLT DatasetsWe first compare \(k\)NN-KD with vanilla \(k\)NN-MT and other KD methods on the two IWSLT translation tasks. Note that there are several hyper-parameters in vanilla \(k\)NN-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\)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\)NN-KD can even perform better than vanilla \(k\)NN-MT, while gaining a significant

\begin{table}
\begin{tabular}{l|c c c|c c c} \hline \hline Models & \multicolumn{3}{c|}{De-En} & \multicolumn{3}{c}{En-Vi} \\  & BLEU & upd/s & token/s & BLEU & upd/s & token/s \\ \hline Transformer & 34.11 & 2.02(1.00\(\times\)) & 3148.10(1.00\(\times\)) & 30.76 & 2.55(1.00\(\times\)) & 2870.07(1.00\(\times\)) \\ Word-KD & 34.26 & 1.77(0.88\(\times\)) & 3291.28(1.06\(\times\)) & 30.98 & 2.14(0.84\(\times\)) & 2782.53(0.97\(\times\)) \\ Seq-KD & 34.60 & 2.14(1.06\(\times\)) & 3409.86(1.08\(\times\)) & 31.20 & 2.80(1.10\(\times\)) & 2855.77(1.00\(\times\)) \\ BERT-KD & 35.63 & 1.70(0.84\(\times\)) & 3275.43(1.04\(\times\)) & 31.51 & 2.14(0.84\(\times\)) & 2785.69(0.97\(\times\)) \\ Selective-KD & 34.38 & 1.72(0.85\(\times\)) & 3365.68(1.07\(\times\)) & 31.48 & 2.09(0.82\(\times\)) & 3044.68(1.06\(\times\)) \\ \(k\)NN-MT & 36.17 & - & 920.72(0.29\(\times\)) & 32.08 & - & 617.88(0.22\(\times\)) \\ \hline \(k\)NN-KD & **36.30** & 2.14(1.06\(\times\)) & 3321.24(1.05\(\times\)) & **32.27** & 2.60(1.02\(\times\)) & 2879.68(1.00\(\times\)) \\ \hline \hline \end{tabular}
\end{table}
Table 2: Experimental results on IWSLT’14 De-En and IWSLT’15 En-Vi translation tasks. “-” means ”not applicable” since vanilla \(k\)NN-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”.

\begin{table}
\begin{tabular}{l|c c|c c|c c} \hline \hline \multirow{2}{*}{Models} & \multicolumn{2}{c|}{Koran} & \multicolumn{2}{c|}{Medical} & \multicolumn{2}{c}{Law} \\  & BLEU & token/s & BLEU & token/s & BLEU & token/s \\ \hline Pre-trained Model & 16.26 & 1038.97(1.00\(\times\)) & 39.91 & 1765.56(1.00\(\times\)) & 45.71 & 2404.31(1.00\(\times\)) \\ \(k\)NN-MT & 19.45 & 246.17(0.24\(\times\)) & 54.35 & 701.29(0.40\(\times\)) & 61.78 & 853.66(0.36\(\times\)) \\ Transformer & 13.84 & 1297.45(1.25\(\times\)) & 27.51 & 1073.53(0.61\(\times\)) & 60.77 & 1689.89(0.70\(\times\)) \\ \hline \(k\)NN-KD & **24.86** & 1236.23(1.19\(\times\)) & **56.50** & 1853.58(1.05\(\times\)) & **61.89** & 2456.62(1.02\(\times\)) \\ \hline \hline \end{tabular}
\end{table}
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 \(k\)NN-KD, and the training efficiency of these two models are basically the same.

 

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_{\text{kNN}}^{\mathrm{T}}\left(y_{i}~{}|~{}\mathbf{x},\mathbf{y}_{<i}^{\star}\right)\) 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 \(k\)NN-KD performs significantly better than Transformer trained by standard CE. It proves the statement in Section 1 that compared with standard CE, \(k\)NN-KD can improve the generalizability of NMT models.

### Analysis

There are two key hyper-parameters in our \(k\)NN-KD: number of neighbors per query (\(k\)), and temperature (\(\tau\)). In this section, we investigate the effects of these two hyper-parameters on the validation set of IWSLT'14 De-En.

**Effect of \(k\)** We fix the temperature \(\tau\) to \(100\), and train the model using \(k\)NN-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 \(k\)NN 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\)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\)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 3: BLEU scores with different \(\tau\) and fixed \(k\) (\(k=64\)) on the validation set of IWSLT’14 De-En dataset.

\begin{table}
\begin{tabular}{l|c c} \hline \hline Models & Law→Medical & Medical→Law \\ \hline Transformer & 18.73 & 2.07 \\ \(k\)NN-KD & 22.31 & 14.82 \\ \hline \hline \end{tabular}
\end{table}
Table 4: Generalizability Evaluation. “Law→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.

 



### Case Study

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 eine menge geschichten daruber , warm 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 off.."_, and _"a lot of..."_ are all correct translations. We input this translation context to the base NMT model, \(k\)NN-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 \(k\)NN-MT and our \(k\)NN-KD increase the probabilities of the reasonable target tokens, and these two models have similar predicted probabilities. Note that \(k\)NN-MT obtains this probability distribution by interpolating the base NMT probability with a \(k\)NN search probability at decoding time, while our \(k\)NN-KD directly outputs this distribution without any additional operations. This further confirms that \(k\)NN-KD can train the model to learn the knowledge of \(k\)NN that prevents the overconfidence of the model on the one-hot label, thus leading to the better generalizability for inference.

## 5 Related Works

### 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 RNNSutskever et al. (2014); Bahdanau et al. (2015), CNNGehring et al. (2017), or TransformerVaswani 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 datstore of cached examples and proposed \(k\)NN-MT. However, Meng et al. (2021) pointed out that \(k\)NN-MT is two-order slower than vanilla MT models, which limits the deployment for real-world applications. They proposed Fast \(k\)NN-MT to solve this problem. Wang et al. (2021b) also noticed the low-efficiency problem of \(k\)NN-MT. Thus, they used a hierarchical clustering strategy and proposed Faster \(k\)NN-MT. Although the above studies have made feasible fixes, \(k\)NN search is still required in the decoding phase, which dramatically increases the difficulty of practical applications compared to standard MT models.

### 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. (2015). 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 \(k\)NN search to enhance a base NMT model.

Figure 4: Predicted probabilities output from the base NMT model, \(k\)NN-MT and our \(k\)NN-KD, given the translation context _“es gibt eine menge geschichten derüber , warm wir dies getan haben ._|| there are”_



## 6 Conclusion

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.

## Acknowledgements

This work was supported by National Science Foundation of China (No. 62161160339), National Key R&D Program of China (No.2018YFB1005100), State Key Laboratory of Media Convergence Production Technology and Systems and Key Laboratory of Science, Techonology and Standard in Press Industry (Key Laboratory of Intelligent Press Media Technology). We appreciate the anonymous reviewers for their helpful comments. Xiaojun Wan is the corresponding auhor.

## References

* Aharoni and Goldberg (2020) 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, pp. 7747-7763. External Links: Link Cited by: SS1.
* D. Bahdanau, K. Cho, and Y. 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, Cited by: SS1.
* A. V. Barone, B. Haddow, U. Germann, and R. 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, pp. 1489-1494. External Links: Link Cited by: SS1.
* P. F. Brown, J. Cocke, M. C. Lee, and W. J. Clark (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, pp. 7893-7905. External Links: Link Cited by: SS1.
* K. Clark, M. Luong, U. Khandelwal, C. D Manning, and Q. 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, pp. 5931-5937. External Links: Link Cited by: SS1.
* J. Devlin, M. Chang, K. Lee, and K. 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), Vol. , pp. 4171-4186. External Links: Document, ISSN 1574-4186 Cited by: SS1.
* J. Gehring, M. Auli, D. Grangier, D. Yarats, and Y. N. Dauphin (2017)Convolutional sequence to sequence learning. In Proceedings of the 34th International Conference on Machine Learning-Volume 70, pp. 1243-1252. External Links: Link Cited by: SS1.
* G. Hinton, O. Vinyals, and J. Dean (2015a)Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531. Cited by: SS1.
* E. Hinton, O. Vinyals, and J. Dean (2015b)Distilling the knowledge in a neural network. CoRRabs/1503.02531. External Links: Link Cited by: SS1.
* J. Johnson, M. Douze, and H. Jegou (2017)Billion-scale similarity search with gpus. CoRRabs/1702.08734. External Links: Link Cited by: SS1.
* U. Khandelwal, A. Fan, D. Jurafsky, L. Zettlemoyer, and M. Lewis (2021)Nearest neighbor machine translation. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021, pp. . External Links: Link Cited by: SS1.
* Y. Kim and A. 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, pp. 1317-1327. External Links: Link Cited by: SS1.
* P. Koehn, H. Hoang, A. Birch, C. Callison-Burch, M. Federico, N. Bertoldi, B. Cowan, W. Shen, C. Moran, R. Zens, C. Dyer, O. Bojar, A. Constantin, and E. 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.
* Luong and Manning (2015) 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_.
* Meng et al. (2021) Yuxian Meng, Xiaoya Li, Xiayu Zheng, Fei Wu, Xiaofei Sun, Tianwei Zhang, and Jiwei Li. 2021. Fast nearest neighbor machine translation. abs/2105.14528.
* Och (2003) 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.
* Ott et al. (2019) 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_.
* Post (2018) 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.
* Sennrich et al. (2016) 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.
* Sutskever et al. (2014) 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.
* Tan et al. (2019) Xu Tan, Yi Ren, Di He, Tao Qin, Zhou Zhao, and Tie-Yan 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.
* Vaswani et al. (2017) 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.
* Wang et al. (2021a) 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.
* Wang et al. (2021b) Shuhe Wang, Jiwei Li, Yuxian Meng, Rongbin Ouyang, Guoyin Wang, Xiaaya Li, Tianwei Zhang, and Shi Zong. 2021b. Faster nearest neighbor machine translation. _CoRR_, abs/2112.08152.
* Zhang et al. (2019) 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.
* Zheng et al. (2021) 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.

## Appendix A Experimental Setup

### Datasets

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.

### 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

\begin{table}
\begin{tabular}{c|c c c} \hline  & Train & Valid & Test \\ \hline 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 \\ \hline \end{tabular}
\end{table}
Table 5: The number of examples 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\).

## Appendix 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.

\begin{table}
\begin{tabular}{l|c|c} \hline Hyperparameters & IWSLT & Multi-domains \\ \hline 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 \\ \(\alpha\) in Equation 12 & 0.5 & 0.5 \\ Optimizer & Adam & Adam \\ -\(\beta_{1}\) & 0.9 & 0.9 \\ -\(\beta_{2}\) & 0.98 & 0.98 \\ -\(\epsilon\) & 1e-8 & 1e-8 \\ \hline \end{tabular}
\end{table}
Table 6: Hyperparameter settings for different datasets.