SlowGuess/NAACL_2022_OCR huggingface dataset:
Zhixian Yang and Renliang Sun and Xiaojun Wan
yangzhixian@stu.pku.edu.cn
sunrenliangpku@gmail.com
wanxiaojun@oku.edu.cn
k-nearest-neighbor machine translation ( $k$  NNT), proposed by Khandelwal et al. (2021), has achieved many state-of-the-art results in machine translation tasks. Although effective,  $k$  NNT-MT requires conducting  $k$  NNT 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$  NNT search forward to the preprocessing phase, and then introduce  $k$  Nearest Neighbor Knowledge Distillation ( $k$  NNT-KD) that trains the base NMT model to directly learn the knowledge of  $k$  NNT. Distilling knowledge retrieved by  $k$  NNT 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$  NNT-MT, while maintaining the same training and decoding speed as the standard NMT model.
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 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  $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 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 (Zhang et al., 2019), 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  $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 (kNN-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  $p_{\mathrm{kNN}}^T$
1Our code is available at https://github.com/FadedCosine/kNN-KD
Human Language Technologies, pages 5546 - 5556
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 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 (Aharoni and Goldberg, 2020) and achieving 2.56 BLEU improvement over vanilla kNN-MT on average.
- 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).
The goal of the standard NMT model is to learn the conditional probability  $p_{\mathrm{MT}}(\mathbf{y} \mid \mathbf{x})$  for translating a sentence  $\mathbf{x} = \{x_1, \dots, x_m\}$  in source language to a sentence  $\mathbf{y} = \{y_1, \dots, y_n\}$  in target language. Translation is usually performed in a autoregressive manner, and its probability can be factored as  $p_{\mathrm{MT}}(\mathbf{y} \mid \mathbf{x}) = \prod_{i=1}^{n} p(y_i \mid \mathbf{x}, \mathbf{y}_{<i})$ . When predicting  $i$ -th token in the target sentence given  $(\mathbf{x}, \mathbf{y}_{<i})$  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
$\nu$  as follows:
$$
p _ {\mathrm {M T}} \left(y _ {i} \mid \mathbf {x}, \mathbf {y} _ {<   i}\right) = \frac {\exp \left(\mathbf {o} _ {y _ {i}} ^ {\top} \mathbf {h} _ {i - 1}\right)}{\sum_ {w \in \mathcal {V}} \exp \left(\mathbf {o} _ {w} ^ {\top} \mathbf {h} _ {i - 1}\right)}, \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} = \{y_1^{\star},\dots ,y_n^{\star}\}$ , and for each  $y_{i}^{\star}$  in the training set, the CE loss is usually used for optimizing NMT models:
$$
\mathcal {L} _ {\mathrm {C E}} = - \sum_ {y _ {i} \in \mathcal {V}} \mathbb {1} _ {y _ {i} = y _ {i} ^ {*}} \log p _ {\mathrm {M T}} \left(y _ {i} \mid \mathbf {x}, \mathbf {y} _ {<   i} ^ {\star}\right), \tag {2}
$$
where  $\mathbb{1}$  is the indicator function, and the groundtruth target sequence  $\mathbf{y}_{<i}^{\star}$  is used in the conditions of  $p_{\mathrm{MT}}$  due to the teacher forcing technique.
2.2 Nearest Neighbor Machine Translation
$k\mathrm{NN}$  -MT applies the nearest neighbor retrieval mechanism to the decoding phase of a NMT model, which allows the model direct access to a largescale datastore for better inference. Specifically,  $k\mathrm{NN}$  -MT includes two following steps:
Datastore Building Given a bilingual sentence pair in the training set  $(\mathbf{x},\mathbf{y}^{\star})\in (\mathcal{X},\mathcal{Y}^{\star})$  , kNNMT first constructs a datastore  $\mathcal{D}$  as follows:
$$
(\mathcal {K}, \mathcal {V}) = \bigcup_ {(\mathbf {x}, \mathbf {y} ^ {*}) \in (\mathcal {X}, \mathcal {Y} ^ {\star})} \left\{\left(f \left(\mathbf {x}, \mathbf {y} _ {<   i} ^ {\star}\right), y _ {i} ^ {\star}\right), \forall y _ {i} ^ {\star} \in \mathbf {y} ^ {\star} \right\}, \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\mathrm{NN}$  MT aims to interpolate the base NMT model's probability in Equation 1 with a  $k\mathrm{NN}$  model. At each decoding step  $i$ ,  $k\mathrm{NN}$  MT maps the current translation context to a representation  $f(\mathbf{x}, \mathbf{y}_{<i})$ , which is used to query the datastore for  $k$  nearest neighbors according to the  $l_2$  distances. Denote the retrieved neighbors as  $\mathcal{N}^i = \{(\mathbf{k}_j, v_j), j \in \{1, 2, \dots, k\}\}$ , and then a  $k\mathrm{NN}$  distribution over vocabulary  $\mathcal{V}$  can be computed as:
$$
\begin{array}{l} p _ {\mathrm {k N N}} \left(y _ {i} \mid \mathbf {x}, \mathbf {y} _ {<   i}\right) \propto \\ \sum_ {\left(\mathbf {k} _ {j}, v _ {j}\right) \in \mathcal {N} ^ {i}} \mathbb {1} _ {y _ {i} = v _ {j}} \exp \left(\frac {- d \left(\mathbf {k} _ {j} , f \left(\mathbf {x} , \mathbf {y} <   i\right)\right)}{\tau}\right), \tag {4} \\ \end{array}
$$
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 of
Figure 1: Illustration of kNN-KD. In the preprocessing phase, we finish the datastore building in Step 1, and conduct kNN search in advance in Step 2. These two steps can be done offline before training and inference. During training, we compute the kNN 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.
$p_{\mathrm{MT}}(y_i \mid \mathbf{x}, \mathbf{y}_{<i})$  and  $p_{\mathrm{kNN}}(y_i \mid \mathbf{x}, \mathbf{y}_{<i})$  with a tunable weight  $\lambda$ :
$$
\begin{array}{l} p \left(y _ {i} \mid \mathbf {x}, \mathbf {y} _ {<   i}\right) = (1 - \lambda) p _ {\mathrm {M T}} \left(y _ {i} \mid \mathbf {x}, \mathbf {y} _ {<   i}\right) \tag {5} \\ + \lambda p _ {\mathrm {k N N}} (y _ {i} | \mathbf {x}, \mathbf {y} _ {<   i}). \\ \end{array}
$$
Note that each decoding step of each beam requires a  $k\mathrm{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\mathrm{NN}$ -MT to be deployed on real-world applications.
2.3 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 ^ {\mathrm {S}} \left(y _ {i} \mid \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 ^ {\mathrm {T}} \left(y _ {i} \mid \mathbf {x}, \mathbf {y} _ {<   i}\right) = \frac {\exp \left(\mathbf {z} _ {y _ {i}} / \tau\right)}{\sum_ {w \in \nu} \exp \left(\mathbf {z} _ {w} / \tau\right)}. \tag {7}
$$
When  $\tau \to \infty$ ,  $p^{\mathrm{T}}$  degenerates into a uniform distribution, and when  $\tau \to 0$ ,  $p^{\mathrm{T}}$  becomes an one-hot vector. Specifically, KD defines the objective as:
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {K D}} = - \sum_ {y _ {i} \in \mathcal {V}} p ^ {\mathrm {T}} \left(y _ {i} \mid \mathbf {x}, \mathbf {y} _ {<   i} ^ {\star}\right) \tag {8} \\ \times \log p ^ {\mathrm {S}} \left(y _ {i} \mid \mathbf {x}, \mathbf {y} _ {<   i} ^ {\star}\right). \\ \end{array}
$$
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 {C E}} + \alpha \mathcal {L} _ {\mathrm {K D}}, \tag {9}
$$
where  $\alpha$  is a weight to balance two losses.
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.
3.1 Nearest Neighbor Knowledge Distillation
When we apply vanilla  $k\mathrm{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\mathrm{NN}$ -MT is prohibitively slow. This is mainly because vanilla  $k\mathrm{NN}$ -MT has to conduct a  $k\mathrm{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\mathrm{NN}$ -KD can be described as follows:
Step 1: Datastore Building We build the datastore  $\mathcal{D}$  in the same way as vanilla kNN-MT (Khandelwal et al., 2021) which has been described in Section 2.2, so we omit it here.
Step 2:  $k\mathbf{NN}$  Search in Advance For all the translation contexts  $(\mathbf{x},\mathbf{y}_{< i}^{\star})$  in the training set, we conduct a  $k\mathbf{NN}$  search using  $f(\mathbf{x},\mathbf{y}_{< i}^{\star})$  as a query to search through the datastore  $\mathcal{D}$  built in Step 1, and then we obtain the retrieved results  $\mathcal{N}^i = \{(\mathbf{k}_j,v_j),j\in \{1,2,\ldots ,k\} \}$ . Note that we are performing  $k\mathbf{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\mathbf{NN}$  search in advance, we extend the target sentence in the training set from  $\mathbf{y}^{\star} = \{y_1^{\star},\dots ,y_n^{\star}\}$  to  $\mathbf{y}^{\star} = \{(y_{1}^{\star},\mathcal{K}^{1}),\dots ,(y_{n}^{\star},\mathcal{K}^{n})\}$ , where  $\mathcal{K}^i = \{(d(\mathbf{k}_j,f(\mathbf{x},\mathbf{y}_{< i}^{\star})),v_j),j\in \{1,2,\ldots ,k\}\}$ .
Step 3:  $k\mathbf{NN}$  as a Teacher In the training phase, a  $k\mathbf{NN}$  distribution can be formulated as:
$$
\begin{array}{l} p _ {\mathrm {k N N}} ^ {\mathrm {T}} \left(y _ {i} \mid \mathbf {x}, \mathbf {y} _ {<   i} ^ {\star}\right) \propto \\ \sum_ {\left(d _ {j}, v _ {j}\right) \in \mathcal {K} ^ {i}} \mathbb {1} _ {y _ {i} = v _ {j}} \exp \left(\frac {- d _ {j}}{\tau}\right), \tag {10} \\ \end{array}
$$
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{array}{l} \mathcal {L} _ {\mathrm {k N N} - \mathrm {K D}} = - \sum_ {y _ {i} \in \mathcal {V}} p _ {\mathrm {k N N}} ^ {\mathrm {T}} \left(y _ {i} \mid \mathbf {x}, \mathbf {y} _ {<   i} ^ {\star}\right) \tag {11} \\ \times \log p _ {\mathrm {M T}} \left(y _ {i} \mid \mathbf {x}, \mathbf {y} _ {<   i} ^ {\star}\right). \\ \end{array}
$$
And the final training objective in Equation 9 can be rewritten as:
$$
\mathcal {L} = (1 - \alpha) \mathcal {L} _ {\mathrm {C E}} + \alpha \mathcal {L} _ {\mathrm {k N N} - \mathrm {K D}}, \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.
In this section, we show that our proposed  $k\mathrm{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}, y_i \in \mathcal{V}$  is:
$$
\begin{array}{l} \frac {\partial \mathcal {L}}{\partial \mathbf {z} _ {y _ {i}}} = (1 - \alpha) \left(p \left(y _ {i}\right) - \mathbb {1} _ {y _ {i} = y _ {i} ^ {*}}\right) + \alpha \left(p \left(y _ {i}\right) - p ^ {\mathrm {T}} \left(y _ {i}\right)\right) \\ = \left\{ \begin{array}{l l} p \left(y _ {i}\right) - \alpha p ^ {\mathrm {T}} \left(y _ {i}\right), & \text {i f} y _ {i} \neq y _ {i} ^ {\star} \text {a n d} y _ {i} \in \mathcal {K} ^ {i} \\ p \left(y _ {i}\right), & \text {i f} y _ {i} \neq y _ {i} ^ {\star} \text {a n d} y _ {i} \notin \mathcal {K} ^ {i} \\ p \left(y _ {i}\right) - \left(1 - \alpha + \alpha p ^ {\mathrm {T}} \left(y _ {i}\right)\right), & \text {i f} y _ {i} = y _ {i} ^ {\star} \end{array} \right. \tag {13} \\ \end{array}
$$
where we abbreviate  $p_{\mathrm{MT}}(y_i \mid \mathbf{x}, \mathbf{y}_{<i}^{\star})$  to  $p(y_i)$  and  $p_{\mathrm{kNN}}^{\mathrm{T}}(y_i \mid \mathbf{x}, \mathbf{y}_{<i}^{\star})$  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 non-ground-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 kNN 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 kNN 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  $(1 - \alpha + \alpha p^{\mathrm{T}}(y_{i}))$ . 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 kNN-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 Dank für ihren hilfreichen Vorschlag | Thanks for your", its groundtruth target token is "useful", while "helpful" is also reasonable for this translation. Assuming that we have conducted the  $k\mathrm{NN}$  search with  $k = 3$  in advance as shown in Figure 1, and set  $\tau$  to 1, we can then compute the  $k\mathrm{NN}$  teacher distribution as:
$$
p ^ {T} \left(y _ {4}\right) = \left\{ \begin{array}{l l} 0. 3 7 8, & \text {i f} y _ {4} \text {i s}" u s e f u l" \\ 0. 6 2 2, & \text {i f} y _ {4} \text {i s}" h e l p f u l" \\ 0, & \text {o t h e r w i s e} \end{array} \right. \tag {14}
$$
According to Equation 13, the gradient norms are  $|p(\text{"helpful")} - 0.622\alpha |$  for "helpful", and  $|p(\text{"useful")} - (1 - 0.622\alpha)|$  for "useful". Therefore, our kNN-KD can train the model to learn
from the  $k\mathrm{NN}$  model to increase the probability of "helpful" that is reasonable but not ground-truth, thus addressing the overcorrection issue.
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 (Aharoni and Goldberg, 2020) (De-En, 733k training samples). For IWSLT'14 De-En, we follow the preprocessing steps provided by fairseq² (Ott et al., 2019) to split the data, and process the text into bytepair encoding (BPE) (Sennrich et al., 2016). For IWSLT'15 En-Vi, we use the pre-processed dataset³ provided by Luong and Manning (2015). We use tst2012 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 dataset⁴ provided by Zheng et al. (2021), and consider domains including Koran, Medical, and Law in our experiments.
4.2 Competitive Models
The proposed  $k\mathrm{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\mathrm{NN}$ -KD with vanilla  $k\mathrm{NN}$ -MT and some typical KD methods:
- Word-KD (Hinton et al., 2015b). As described in Section 2.3, Word-KD is the standard KD that distills knowledge equally for each word.
- Seq-KD (Kim and Rush, 2016). 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
Datasets	|D|	k	τ
IWSLT'14 De-En	3,949,106	64	100
IWSLT'15 En-Vi	3,581,500	64	100
Koran	524,374	16	100
Medical	6,903,141	4	10
Law	19,062,738	4	10
Table 1: Hyper-parameter settings for different datasets.
data, and the training objective can be formulated as:
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {S e q - K D}} = - \sum_ {i = 1} ^ {n} \sum_ {y _ {i} \in \mathcal {V}} \mathbb {1} _ {y _ {i} = \hat {y} _ {i}} \tag {15} \\ \times \log p _ {\mathrm {M T}} \left(y _ {i} \mid \mathbf {x}, \hat {\mathbf {y}} _ {<   j}\right), \\ \end{array}
$$
where  $\hat{\mathbf{y}}$  is the target sequence generated by teacher model, and  $n$  is the length of it.
- BERT-KD (Chen et al., 2020). This method distills knowledge learned in BERT (Devlin et al., 2019) to the student NMT model to improve translation quality.
- Selective-KD (Wang et al., 2021a). This work finds that some of the teacher's knowledge will hurt the effect of KD, and then address this issue by introducing Selective-KD to select suitable samples for distillation.
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. Other details of the experimental setup can be seen in Appendix A.
Model Configuration We choose Transformer (Vaswani et al., 2017) as our base NMT model. For IWSLT'14 De-En and IWSLT'15 EnVi, we use transformer_iwslt_de_en configuration, which has 6 layers in both encoder and decoder, embedding size 512, feed-forward size 1,024 and attention heads 4. For multi-domains translation datasets, we follow Khandelwal et al. (2021) to adopt transformer_wmt19_de_en configuration, which has 6 layers in both encoder and decoder, embedding size 1,024, feed-forward size 8, 192 and attention heads 8.
Preprocessing Details When building the datastores, we use the context vectors input to the final output layer as keys in the datastore  $\mathcal{D}$ . For
$^{2}$ https://github.com/pytorch/fairseq/blob/main/examples/translation/prepare-iwslt14.sh
3https://nlp.stanford.edu/projects/nmt/
4https://github.com/zhengxxn/ adaptive-knn-mt