SlowGuess/NAACL_2021_OCR huggingface dataset:
Shuhao Gu $^{1,2}$ , Yang Feng $^{1,2*}$ , Wanying Xie $^{1,3}$
<sup>1</sup> Key Laboratory of Intelligent Information Processing,
<sup>3</sup> Beijing Language and Culture University, China
{gushuhao19b,fengyang} @ict.ac.cn, xiewanying07@gmail.com
Domain Adaptation is widely used in practical applications of neural machine translation, which aims to achieve good performance on both general domain and in-domain data. However, the existing methods for domain adaptation usually suffer from catastrophic forgetting, large domain divergence, and model explosion. To address these three problems, we propose a method of "divide and conquer" which is based on the importance of neurons or parameters for the translation model. In this method, we first prune the model and only keep the important neurons or parameters, making them responsible for both general-domain and in-domain translation. Then we further train the pruned model supervised by the original whole model with knowledge distillation. Last we expand the model to the original size and fine-tune the added parameters for the in-domain translation. We conducted experiments on different language pairs and domains and the results show that our method can achieve significant improvements compared with several strong baselines.
Neural machine translation (NMT) models (Kalchbrenner and Blunsom, 2013; Cho et al., 2014; Sutskever et al., 2014; Bahdanau et al., 2015; Gehring et al., 2017; Vaswani et al., 2017) are data-driven and hence require large-scale training data to achieve good performance (Zhang et al., 2019a). In practical applications, NMT models usually need to produce translation for some specific domains with only a small quantity of in-domain data available, so domain adaptation is applied to address the problem. A typical domain adaptation scenario as discussed in Freitag and Al-Onaizan (2016) is that an NMT model have been trained with large-scale general-domain data and then is adapted to specific
domains, hoping the model can fit in-domain data well meanwhile the performance will not degrade too much on the general domain.
Towards this end, many researchers have made their attempts. The fine-tuning method (Luong and Manning, 2015) performs in-domain training based on the general-domain model by first training the model on general-domain data and then continuing to train on in-domain data. Despite its convenience for use and high-quality for in-domain translation, this method suffers from catastrophic forgetting which leads to poor performance in the previous domains. Regularization-based methods (Dakwale and Monz, 2017; Thompson et al., 2019; Barone et al., 2017; Khayrallah et al., 2018) instead introduce an additional loss to the original objective so that the translation model can trade off between general-domain and in-domain. This kind of methods usually has all the parameters shared by general-domain and in-domain, with the assumption that the optimal parameter spaces for all the domains will overlap with each other, and retaining these overlapped parameters can balance over all the domains. This assumption is feasible when the domains are similar, but when the divergence of the domains is large, it is not reasonable anymore. In contrast, the methods with domain-specific networks (Dakwale and Monz, 2017; Wang et al., 2019; Bapna and First, 2019; Gu et al., 2019) can be often (but not always) immune to domain divergence as it can capture domain-specific features. But unfortunately, as the number of domains increases, the parameters of this kind of methods will surge. Besides, the structure of these networks needs to be carefully designed and tuned, which prevents them from being used in many cases.
Given the above, we propose a method of domain adaptation that can not only deal with large domain divergence during domain transferring but also keep a stable model size even with multiple domains. Inspired by the analysis work on NMT (Bau
Corresponding author: Yang Feng.
Reproducible code: https://github.com/ictnlp/PTE-NMT.
Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 3942-3952 June 6-11, 2021. ©2021 Association for Computational Linguistics
et al., 2019; Voita et al., 2019; Gu and Feng, 2020), we find that only some important parameters in a well-trained NMT model play an important role when generating the translation and unimportant parameters can be erased without affecting the translation quality too much. According to these findings, we can preserve important parameters for general-domain translation, while tuning unimportant parameters for in-domain translation. To achieve this, we first train a model on the general domain and then shrink the model with neuron pruning or weight pruning methods, only retaining the important neurons/parameters. To ensure the model can still perform well on general-domain data, we adjust the model on in-domain data with knowledge distillation where the original whole model is used as the teacher and the pruned model as the student. Finally, we expand the model to the original size and fine-tune the added parameters on the in-domain data. Experimental results on different languages and domains show that our method can avoid catastrophic forgetting on general-domain data and achieve significant improvements over strong baselines on multiple in-domain data sets.
- We prove that the parameters that are unimportant for general-domain data can be utilized to improve in-domain translation quality.
- Our model can keep superior performance over baselines even when continually transferring to multiple domains.
- Our model can fit in the continual learning scenario where the data for the previous domains cannot be got anymore which is the common situation in practice.
2.1 The Transformer
In our work, we apply our method in the framework of TRANSFORMER (Vaswani et al., 2017) which will be briefly introduced here. However, we note that our method can also be combined with other NMT architectures. We denote the input sequence of symbols as  $\mathbf{x} = (x_{1},\ldots ,x_{J})$ , the ground-truth sequence as  $\mathbf{y}^{*} = (y_{1}^{*},\dots,y_{K^{*}}^{*})$  and the translation as  $\mathbf{y} = (y_{1},\dots,y_{K})$ .
The Encoder & Decoder The encoder is composed of  $N$  identical layers. Each layer has two sublayers. The first is a multi-head self-attention
sublayer and the second is a fully connected feedforward network. Both of the sublayers are followed by a residual connection operation and a layer normalization operation. The input sequence  $\mathbf{x}$  will be first converted to a sequence of vectors  $\mathbf{E}_x = [E_x[x_1];\ldots ;E_x[x_J]]$  where  $E_{x}[x_{j}]$  is the sum of word embedding and position embedding of the source word  $x_{j}$ . Then, this sequence of vectors will be fed into the encoder and the output of the  $N$ -th layer will be taken as source hidden states. and we denote it as  $\mathbf{H}$ . The decoder is also composed of  $N$  identical layers. In addition to the same kind of two sublayers in each encoder layer, the cross-attention sublayer is inserted between them, which performs multi-head attention over the output of the encoder. The final output of the  $N$ -th layer gives the target hidden states  $\mathbf{S} = [\mathbf{s}_1;\dots ;\mathbf{s}_{K*}]$ , where  $\mathbf{s}_k$  is the hidden states of  $y_{k}$ .
The Objective We can get the predicted probability of the  $k$ -th target word over the target vocabulary by performing a linear transformation and a softmax operation to the target hidden states:
$$
p \left(y _ {k} \mid \mathbf {y} _ {<   k}, \mathbf {x}\right) \propto \exp \left(\mathbf {W} _ {o} \mathbf {s} _ {k} + \mathbf {b} _ {o}\right), \tag {1}
$$
where  $\mathbf{W}_o\in \mathbb{R}^{d_{model}\times |V_t|}$  and  $|\mathrm{V}_t|$  are the size of target vocabulary. The model is optimized by minimizing a cross-entropy loss of the ground-truth sequence with teacher forcing training:
$$
\mathcal {L} (\theta) = - \frac {1}{K} \sum_ {k = 1} ^ {K} \log p \left(y _ {k} ^ {*} | \mathbf {y} _ {<   k}, \mathbf {x}; \theta\right), \tag {2}
$$
where  $K$  is the length of the target sentence and  $\theta$  denotes the model parameters.
Knowledge Distillation (KD) method (Hinton et al., 2015) is for distilling knowledge from a teacher network to a student network. Normally, the teacher network is considered to be with higher capability. A smaller student network can be trained to perform comparably or even better by mimicking the output distribution of the teacher network on the same data. This is usually done by minimizing the cross entropy between the two distributions:
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {K D}} (\theta , \theta_ {T}) = - \frac {1}{K} \sum_ {k = 1} ^ {K} q \left(\mathbf {y} _ {k} \mid \mathbf {y} _ {<   k}, \mathbf {x}; \theta_ {T}\right) \tag {3} \\ \times \log p (\mathbf {y} _ {k} | \mathbf {y} _ {<   k}, \mathbf {x}; \theta), \\ \end{array}
$$
where  $q$  denotes the output distribution of the teacher network and  $\theta$  and  $\theta_T$  denote the parameters of the student and teacher network, respectively.
The parameters of the teacher network usually keep fixed during the KD process.
The main idea of our method is that different neurons or parameters have different importance to the translation model and hence different roles in domain adaptation. Based on this, we distinguish them into important and unimportant ones and make important neurons or parameters compromise between domains while unimportant ones focus on in-domain. Specifically, our method involves the following steps shown in Figure 1. First, we train a model on the general domain and then evaluate the importance of different neurons or parameters. Then we erase the unimportant neurons or parameters and only keep the ones that are related to the general domain so that our method will not be subjected to domain divergence. Next, we further adjust our model under the framework of knowledge distillation (Hinton et al., 2015) on the in-domain with the unpruned model as the teacher and the pruned model as the student. In this way, the pruned model can regain some of its lost performance because of pruning. Finally, we expand the pruned model to the original size and fine-tune the added parameters for the in-domain.
3.1 Model Pruning
Model pruning aims to find a good subset of neurons and parameters of the general-domain model while maintaining the original performance as much as possible. Therefore, under the premise of retaining most of the model's capability, we want to remove those unimportant neurons or parameters to reduce the size of the whole model first. To achieve this, we adopt two pruning schemes. The first is neuron pruning, where we evaluate the importance of neurons directly and then prune unimportant neurons and relevant parameters. The second is weight pruning, where we evaluate and prune each parameter directly.
Neuron Pruning To evaluate the importance of each neuron, we adopt a criterion based on the Taylor expansion (Molchanov et al., 2017), where we directly approximate the change in loss when removing a particular neuron. Let  $h_i$  be the output produced from neuron  $i$  and  $H$  represents the set of other neurons. Assuming the independence of each neuron in the model, the change of loss when
removing a certain neuron can be represented as:
$$
\left| \Delta \mathcal {L} \left(h _ {i}\right) \right| = \left| \mathcal {L} \left(H, h _ {i} = 0\right) - \mathcal {L} \left(H, h _ {i}\right) \right|, \tag {4}
$$
where  $\mathcal{L}(H,h_i = 0)$  is the loss value if the neuron  $i$  is pruned and  $\mathcal{L}(H,h_i)$  is the loss if it is not pruned. For the function  $\mathcal{L}(H,h_i)$ , its Taylor expansion at point  $h_i = a$  is:
$$
\mathcal {L} (H, h _ {i}) = \sum_ {n = 0} ^ {N} \frac {\mathcal {L} ^ {n} (H , a)}{n !} \left(h _ {i} - a\right) ^ {n} + R _ {N} \left(h _ {i}\right), \tag {5}
$$
where  $\mathcal{L}^n (H,a)$  is the  $n$  -th derivative of  $\mathcal{L}(H,h_i)$  evaluated at point  $a$  and  $R_{N}(h_{i})$  is  $N$  -th remainder. Then, approximating  $\mathcal{L}(H,h_i = 0)$  with a firstorder Taylor polynomial where  $h_i$  equals zero:
$$
\mathcal {L} (H, h _ {i} = 0) = \mathcal {L} (H, h _ {i}) - \frac {\partial \mathcal {L} (H , h _ {i})}{\partial h _ {i}} h _ {i} - R _ {1} (h _ {i}). \tag {6}
$$
The remainder  $R_{1}$  can be represented in the form of Lagrange:
$$
R _ {1} \left(h _ {i}\right) = \frac {\partial^ {2} \mathcal {L} \left(H , h _ {i}\right)}{\partial^ {2} \delta h _ {i}} h _ {i} ^ {2}, \tag {7}
$$
where  $\delta \in (0,1)$ . Considering the use of ReLU activation function in the model, the first derivative of loss function tends to be constant, so the second order term tends to be zero in the end of training. Thus, we can ignore the remainder and get the importance evaluation function as follows:
$$
\Theta_ {\mathrm {T E}} \left(h _ {i}\right) = \left| \Delta \mathcal {L} \left(h _ {i}\right) \right| = \left| \frac {\partial \mathcal {L} \left(H , h _ {i}\right)}{\partial h _ {i}} h _ {i} \right|. \tag {8}
$$
In practice, we need to accumulate the product of the activation and the gradient of the objective function w.r.t to the activation, which is easily computed during back-propagation. Finally, the evaluation function is shown as:
$$
\Theta_ {\mathrm {T E}} \left(h _ {i} ^ {l}\right) = \frac {1}{T} \sum_ {t} \left| \frac {\delta \mathcal {L} \left(H , h _ {i} ^ {l}\right)}{\delta h _ {i} ^ {l}} h _ {i} ^ {l} \right|, \tag {9}
$$
where  $h_i^l$  is the activation value of the  $i$ -th neuron of  $l$ -th layer and  $T$  is the number of the training examples. The criterion is computed on the general-domain data and averaged over  $T$ . Finally, we prune a certain percentage of neurons and relevant parameters in each target layer based on this criterion.
Weight Pruning We adopt the magnitude-based weight pruning scheme (See et al., 2016), where
Figure 1: The whole training process of the proposed method.
the absolute value of each parameter in the target matrix is treated as the importance:
$$
\Theta_ {\mathrm {A V}} \left(\mathrm {w} _ {m n}\right) = \left| \mathrm {w} _ {m n} \right|, \mathrm {w} _ {m n} \in \mathbf {W}, \tag {10}
$$
where  $\mathbf{w}_{mn}$  denotes the  $m$ -th row and  $n$ -th column parameter of the weight matrix  $\mathbf{W}$ . The weight matrix  $\mathbf{W}$  represents different parts of the model, e.g., embedding layer, attention layer, output layer, etc. Finally, a certain percentage of parameters in each target parameter matrix are pruned.
3.2 Knowledge Distillation
Though only limited degradation will be brought in performance after removing the unimportant neurons or parameters, we want to further reduce this loss. To achieve this, we minimize the difference in the output distribution of the unpruned and pruned model. In this work, the general-domain model (parameters denoted as  $\theta_G^*$ ) acts as the teacher model and the pruned model (parameters denoted as  $\theta_G$ ) acts as the student model. So, the objective in this training phase is:
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {K D}} \left(\theta_ {G}, \theta_ {G} ^ {*}\right) = - \frac {1}{K} \sum_ {k = 1} ^ {K} q \left(\mathbf {y} _ {k} \mid \mathbf {y} _ {<   k}, \mathbf {x}; \theta_ {G} ^ {*}\right) \tag {11} \\ \times \log p (\mathbf {y} _ {k} | \mathbf {y} _ {<   k}, \mathbf {x}; \theta_ {G}). \\ \end{array}
$$
Considering that the general-domain data is not always available in some scenarios when adapting the model to new domains, e.g., continual learning, we adopt the word-level knowledge distillation method using the in-domain data. Because the teacher model is trained on general-domain, it can
still transfer the general-domain knowledge to the student model even with the in-domain data. We can fine-tune the pruned model on general-domain if the data is available which can simplify the training procedure. We have also tried the sentence-level knowledge distillation method, but the results are much worse. The parameters of the teacher model keep fixed during this training phase and the parameters of the pruned model are updated with this KD loss. After convergence, the parameters of the pruned model  $(\theta_{G})$  will be solely responsible for the general-domain and will also participate in the translation of in-domain data. These parameters will be kept fixed during the following training phase, so our model won't suffer catastrophic forgetting on the general-domain during the fine-tuning process.
3.3 Model Expansion
After getting the well-trained pruned model, we add new parameters (denoted as  $\theta_I$ ) to it, which expands the model to its original size. Then we fine-tune these newly added parameters with indomain data, which is supervised by the ground truth sequences. As we have indicated above, the parameters of the pruned model (denoted as  $\theta_G$ ), which are responsible for generating the general-domain translation, keep fixed during this training phase. The objective function is:
$$
\mathcal {L} \left(\theta_ {G}, \theta_ {I}\right) = - \frac {1}{K} \sum_ {k = 1} ^ {K} \log p \left(y _ {k} ^ {*} \mid \mathbf {y} _ {<   k}, \mathbf {x}; \theta_ {G}, \theta_ {I}\right). \tag {12}
$$
After convergence, the parameters of the pruned model  $(\theta_G)$  and new parameters  $(\theta_I)$  are combined together for generating the in-domain translation.
Chinese  $\rightarrow$  English. For this task, the general-domain data is from WMT 2017 Zh-En translation task that contains 23.97M sentence pairs. The data is mainly related to the News domain. The newsdev2017 and newstest2017 are chosen as the development and test set, respectively. We choose the parallel sentences with the domain label Thesis from the UM-Corpus (Tian et al., 2014) as our indomain data. This portion covers 15 journal topics in the research area. We filter out the duplicate sentences and then choose 75K, 1K, and 1K sentences randomly as our training, development, and test data, respectively. We tokenize and truecase the English sentences with Moses scripts. $^{1}$  For the Chinese data, we perform word segmentation by using Stanford Segmenter. $^{2}$
English  $\rightarrow$  French. For this task, the general-domain data is from the UN corpus of the WMT 2014 En-Fr translation task that contains 12.78M sentence pairs, which are mainly related to the News domain. We choose newstest2013 and newstest2014 as our development and test set, respectively. The in-domain data with 53K sentence pairs are from WMT 2019 biomedical translation task, and it is mainly related to the Biomedical domain. We choose 1K and 1K sentences randomly from the corpora as our development and test data, respectively. We tokenize and truecase the corpora.
English  $\rightarrow$  German. For this task, general-domain data is from the WMT16 En-De translation task which is mainly News texts. It contains about 4.5M sentence pairs. We choose the newstest2013 for validation and newstest2014 for test. For the indomain data, we use the parallel training data from the IWSLT 2015 which is mainly from the Spoken domain. It contains about 194K sentences. We choose the 2014test for validation and the 2015test for test. We tokenize and truecase the corpora.
Besides, integrating operations of 32K, 32K, and 30K are performed to learn BPE (Sennrich et al., 2016) on the general-domain data and then applied to both the general-domain and in-domain data. Then we filter out the sentences which are longer
than 128 sub-words. For the Zh-En translation task,  $44\mathrm{K}$  size of the Chinese dictionary and  $33\mathrm{K}$  size of the English dictionary are built based on the general-domain data. For the En-Fr and En-De tasks,  $32\mathrm{K}$  size of the dictionaries for the source and target languages are also built on the corresponding general-domain data.
4.2 Systems
We use the open-source toolkit called Fairseqpy (Ott et al., 2019) released by Facebook as our Transformer system. The contrast methods can be divided into two categories. The models of the first category are capacity-fixed while the second category are capacity-increased. The first category includes the following systems:
- General This baseline system is trained only with the general-domain training data.
- In This baseline system is trained only with the in-domain training data.
- Fine-tuning (Luong and Manning, 2015) This method just continues to train the general-domain model with the in-domain data.
- SeqKD (Kim and Rush, 2016) The in-domain source sentences are first translated by the general-domain model. Then the model is further trained with the combined pseudo and real data.
- Multi-objective Learning (MOL) (Dakwale and Monz, 2017) This method is based on the Finetuning method. Besides minimizing the loss between the ground truth words and the output distribution of the network, this method also minimizes the cross-entropy between the output distribution of the general-domain model and the network. The final objective is:
$$
\mathcal {L} _ {\mathrm {M O L}} (\theta) = \mathcal {L} (\theta) + \alpha \mathcal {L} _ {\mathrm {K D}} (\theta) \tag {13}
$$
where  $\alpha$  is the hyper-parameter which controls the contribution of the two parts. The bigger the value, the less degradation on the general-domain.
- Elastic Weight Consolidation (EWC) (Thompson et al., 2019) This method models the importance of the parameters with Fisher information matrix and puts more constrains on the important parameters to let them stay close to the original values during the fine-tuning process. The training objective is:
$$
\mathcal {L} _ {\mathrm {E W C}} (\theta) = \mathcal {L} (\theta) + \alpha \sum_ {i} F _ {i} \left(\theta_ {i} - \theta_ {i} ^ {G}\right) ^ {2} \tag {14}
$$
where  $i$  represents the  $i$ -th parameter and  $F_{i}$  is the modeled importance for the  $i$ -th parameter.
<sup>1</sup>http://www.statmt.org/moses/
2https://nlp.stanford.edu/
ID	System	Chinese-English	English-French	English-German
#Para.	Gen.	In.	Avg.	#Para.	Gen.	In.	Avg.	#Para.	Gen.	In.	Avg.
0	General	100.5M	23.26	9.97	16.62	93.3M	33.05	25.25	29.15	94.4M	26.22	29.53	27.88
1	In	0.87	3.64	2.26-14.36	3.88	7.87	5.88-23.27	12.63	24.77	18.70-9.18
2	Fine-tuning	13.78	17.05	15.42-1.20	14.03	33.24	23.64-5.51	23.37	32.63	28.00+0.12
3	SeqKD	14.02	13.45	13.74-2.88	20.37	27.93	24.15-5.00	24.50	30.20	27.35-0.53
4	MOL	18.44	13.86	16.15-0.47	26.58	28.05	27.32-1.83	25.42	32.09	28.76+0.88
5	EWC	17.45	15.72	16.59-0.03	25.29	32.24	28.77-0.38	25.62	32.17	28.90+1.02
6	Full Bias	100.6M	23.26	11.55	17.41+0.79	93.4M	33.05	26.47	9.76+0.61	94.5M	26.22	30.00	28.11+0.23
7	Adapter	101.3M	23.26	15.82	19.54+2.92	94.1M	33.05	29.61	31.33+2.18	95.3M	26.22	31.54	28.88+1.00
8	MLL	117.4M	23.30	16.04	19.67+3.05	109.7M	32.70	30.78	31.74+2.59	111.2M	26.08	32.07	29.08+1.20
9	PTE	0+NP	81.1M	18.35	9.70	14.03-2.45	76.3M	29.27	24.86	27.07-2.60	77.2M	24.99	26.23	25.61-2.27
10	9+KD	81.1M	22.62	9.99	16.31-0.17	76.3M	32.77	25.83	29.30-0.37	77.2M	26.04	27.71	26.88-1.00
11	10+FT	100.5M	22.62	15.94	19.28+2.81	93.3M	32.77	30.69	31.73+2.07	94.4M	26.04	32.57*	29.31+1.43
12	0+WP	70.4M	20.74	9.54	15.14-1.48	65.3M	29.65	25.03	27.34-1.81	66.1M	25.02	26.66	25.84-2.04
13	12+KD	70.4M	23.50	9.77	16.64+0.02	65.3M	32.64	25.98	29.31+0.16	66.1M	26.38	26.74	26.56-1.32
14	13+FT	100.5M	23.50	16.98**	20.24+3.62	93.3M	32.64	33.16**	32.90+3.75	94.4M	26.38	33.02**	29.70+1.82
Table 1: BLEU scores on three translation tasks. '#Para.' denotes the number of parameters of the whole model, 'Gen.' and 'In.' denote the BLEU on general-domain and in-domain, and 'Avg.' denotes the average BLEU of the two test sets. 'NP', 'WP', 'KD', and 'FT' represent neuron pruning, weight pruning, knowledge distillation, and fine-tuning, respectively. The numbers on the right of 'PTE' denote that this training phase is based on the previous corresponding models. After knowledge distillation, the parameters in the pruned model (system 10, 13) are fixed, so the general-domain BLEU is unchanged after fine-tuning (system 11, 14). * and ** mean the improvements over the MLL method is statistically significant ( $\rho < 0.05$  and  $\rho < 0.01$ , respectively). (Collins et al., 2005)
The second category includes the following three systems: