SlowGuess/NAACL_2022_OCR huggingface dataset:
Chenze Shao $^{1,2}$ , Xuanfu Wu $^{1,2}$ , Yang Feng $^{1,2*}$
{shaochenze18z，wuxuanfu20s，fengyang}@ict.ac.cn
Non-autoregressive neural machine translation (NAT) suffers from the multi-modality problem: the source sentence may have multiple correct translations, but the loss function is calculated only according to the reference sentence. Sequence-level knowledge distillation makes the target more deterministic by replacing the target with the output from an autoregressive model. However, the multi-modality problem in the distilled dataset is still nonnegligible. Furthermore, learning from a specific teacher limits the upper bound of the model capability, restricting the potential of NAT models. In this paper, we argue that one reference is not enough and propose diverse distillation with reference selection (DDRS) for NAT. Specifically, we first propose a method called SeedDiv for diverse machine translation, which enables us to generate a dataset containing multiple high-quality reference translations for each source sentence. During the training, we compare the NAT output with all references and select the one that best fits the NAT output to train the model. Experiments on widely-used machine translation benchmarks demonstrate the effectiveness of DDRS, which achieves 29.82 BLEU with only one decoding pass on WMT14 En-De, improving the state-of-the-art performance for NAT by over 1 BLEU.<sup>1</sup>
Non-autoregressive machine translation (Gu et al., 2018) has received increasing attention in the field of neural machine translation for the property of parallel decoding. Despite the significant speedup, NAT suffers from the performance degradation compared to autoregressive models (Bahdanau et al., 2015; Vaswani et al., 2017) due to the multimodality problem: the source sentence may have
multiple correct translations, but the loss is calculated only according to the reference sentence. The multi-modality problem will cause the inaccuracy of the loss function since NAT has no prior knowledge about the reference sentence during the generation, where the teacher forcing algorithm (Williams and Zipser, 1989) makes autoregressive models less affected by feeding the golden context.
How to overcome the multi-modality problem has been a central focus in recent efforts for improving NAT models (Shao et al., 2019, 2020, 2021; Ran et al., 2020; Sun and Yang, 2020; Ghazvininejad et al., 2020; Du et al., 2021). A standard approach is to use sequence-level knowledge distillation (Kim and Rush, 2016), which attacks the multimodality problem by replacing the target-side of the training set with the output from an autoregressive model. The distilled dataset is less complex and more deterministic (Zhou et al., 2020), which becomes a default configuration of NAT. However, the multi-modality problem in the distilled dataset is still nonnegligible (Zhou et al., 2020). Furthermore, the distillation requires NAT models to imitate the behavior of a specific autoregressive teacher, which limits the upper bound of the model capability and restricts the potential of developing stronger NAT models.
In this paper, we argue that one reference is not enough and propose diverse distillation with reference selection (DDRS) for NAT. Diverse distillation generates a dataset containing multiple reference translations for each source sentence, and reference selection finds the reference translation that best fits the model output for the training. As illustrated in Figure 1, diverse distillation provides candidate references "I must leave tomorrow" and "Tomorrow I must leave", and reference selection selects the former which fits better with the model output. More importantly, NAT with DDRS does not imitate the behavior of a specific teacher but learns selectively from multiple references, which
Source code: https://github.com/ictnlp/DDRS-NAT.
Human Language Technologies, pages 3779 - 3791
Figure 1: Illustration of diverse distillation and reference selection. Diverse distillation provides multiple references and reference selection selects the one that best fits the model output for the training.
improves the upper bound of the model capability and allows for developing stronger NAT models.
The object of diverse distillation is similar to the task of diverse machine translation, which aims to generate diverse translations with high translation quality (Li et al., 2016; Vijayakumar et al., 2018; Shen et al., 2019; Wu et al., 2020; Li et al., 2021). We propose a simple yet effective method called SeedDiv, which directly uses the randomness in model training controlled by random seeds to produce diverse reference translations without losing translation quality. For reference selection, we compare the model output with all references and select the one that best fits the model output, which can be efficiently conducted without extra neural computations. The model learns from all references indiscriminately in the beginning, and gradually focuses more on the selected reference that provides accurate training signals for the model. We also extend the reference selection approach to reinforcement learning, where we encourage the model to move towards the selected reference that gives the maximum reward to the model output.
We conduct experiments on widely-used machine translation benchmarks to demonstrate the effectiveness of our method. On the competitive task WMT14 En-De, DDRS achieves 27.60 BLEU with  $14.7 \times$  speedup and 28.33 BLEU with  $5.0 \times$  speedup, outperforming the autoregressive Transformer while maintaining considerable speedup. When using the larger version of Transformer, DDRS even achieves 29.82 BLEU with only one decoding pass, improving the state-of-the-art performance level for NAT by over 1 BLEU.
Gu et al. (2018) proposes non-autoregressive ma
chine translation to reduce the translation latency through parallel decoding. The vanilla-NAT models the translation probability from the source sentence  $\mathbf{x}$  to the target sentence  $\mathbf{y} = \{y_{1},\dots,y_{T}\}$  as:
$$
p (\boldsymbol {y} | \boldsymbol {x}, \theta) = \prod_ {t = 1} ^ {T} p _ {t} \left(y _ {t} | \boldsymbol {x}, \theta\right), \tag {1}
$$
where  $\theta$  is a set of model parameters and  $p_t(y_t|\boldsymbol{x},\theta)$  is the translation probability of word  $y_{t}$  in position  $t$ . The vanilla-NAT is trained to minimize the cross-entropy loss:
$$
\mathcal {L} _ {C E} (\theta) = - \sum_ {t = 1} ^ {T} \log \left(p _ {t} \left(y _ {t} \mid \boldsymbol {x}, \theta\right)\right). \tag {2}
$$
The vanilla-NAT has to know the target length before constructing the decoder inputs. The target length  $T$  is set as the reference length during the training and obtained from a length predictor during the inference. The target length cannot be changed dynamically during the inference, so it often requires generating multiple candidates with different lengths and re-scoring them to produce the final translation (Gu et al., 2018).
The length issue can be overcome by connectionist temporal classification (CTC, Graves et al., 2006). CTC-based models usually generate a long alignment containing repetitions and blank tokens. The alignment will be post-processed by a collapsing function  $\Gamma^{-1}$  to recover a normal sentence, which first collapses consecutive repeated tokens and then removes all blank tokens. CTC is capable of efficiently finding all alignments  $a$  which the reference sentence  $y$  can be recovered from, and marginalizing the log-likelihood with dynamic programming:
$$
\log p (\boldsymbol {y} | \boldsymbol {x}, \theta) = \log \sum_ {\boldsymbol {a} \in \Gamma (\boldsymbol {y})} p (\boldsymbol {a} | \boldsymbol {x}, \theta). \tag {3}
$$
Due to the superior performance and the flexibility of generating predictions with variable length, CTC is receiving increasing attention in non-autoregressive translation (Libovicky and Helcl, 2018; Kasner et al., 2020; Sahara et al., 2020; Gu and Kong, 2020; Zheng et al., 2021).
2.2 Sequence-Level Knowledge Distillation
Sequence-level Knowledge Distillation (SeqKD, Kim and Rush, 2016) is a widely used knowledge distillation method in NMT, which trains the student model to mimic teacher's actions at sequence-level. Given the student prediction  $p$  and the teacher prediction  $q$ , the distillation loss is:
$$
\begin{array}{l} \mathcal {L} _ {\text {S e q K D}} (\theta) = - \sum_ {\boldsymbol {y}} q (\boldsymbol {y} | \boldsymbol {x}) \log p (\boldsymbol {y} | \boldsymbol {x}, \theta) \tag {4} \\ \approx - \log p (\hat {\boldsymbol {y}} | \boldsymbol {x}, \theta), \\ \end{array}
$$
where  $\theta$  are parameters of the student model and  $\hat{\pmb{y}}$  is the output from running beam search with the teacher model. The teacher output  $\hat{\pmb{y}}$  is used to approximate the teacher distribution otherwise the distillation loss will be intractable.
The procedure of sequence-level knowledge distillation is: (1) train a teacher model, (2) run beam search over the training set with this model, (3) train the student model with cross-entropy on the source sentence and teacher translation pairs. The distilled dataset is less complex and more deterministic (Zhou et al., 2020), which helps to alleviate the multi-modality problem and becomes a default configuration in NAT models.
2.3 Diverse Machine Translation
The task of diverse machine translation requires to generate diverse translations and meanwhile maintain high translation quality. Assume the reference sentence is  $\mathbf{y}$  and we have multiple translations  $\{\mathbf{y}_1, \dots, \mathbf{y}_k\}$ , the translation quality is measured by the average reference BLEU (rfb):
$$
\operatorname {r f b} = \frac {1}{k} \sum_ {i = 1} ^ {k} \operatorname {B L E U} (\boldsymbol {y}, \boldsymbol {y} _ {i}), \tag {5}
$$
and the translation diversity is measured by the average pairwise BLEU (pwb):
$$
\operatorname {p w b} = \frac {1}{(k - 1) k} \sum_ {i = 1} ^ {k} \sum_ {j \neq i} \mathrm {B L E U} \left(\boldsymbol {y} _ {i}, \boldsymbol {y} _ {j}\right). \tag {6}
$$
Higher reference BLEU indicates better translation quality and lower pairwise BLEU indicates
better translation diversity. Generally speaking, there is a trade-off between quality and diversity. In existing methods, translation diversity has to be achieved at the cost of losing translation quality.
In this section, we first introduce the diverse distillation technique we use to generate multiple reference translations for each source sentence, and then apply reference selection to select the reference that best fits the model output for the training.
3.1 Diverse Distillation
The objective of diverse distillation is to obtain a dataset containing multiple high-quality references for each source sentence, which is similar to the task of diverse machine translation that aims to generate diverse translations with high translation quality. However, the translation diversity is achieved at a certain cost of translation quality in previous work, which is not desired in diverse distillation.
Using the randomness in model training, we propose a simple yet effective method called SeedDiv to achieve translation diversity without losing translation quality. Specifically, given the desired number of translations  $k$ , we directly set  $k$  different random seeds to train  $k$  translation models, where random seeds control the random factors during the model training such as parameter initialization, batch order, and dropout. During the decoding, each model translates the source sentence with beam search, which gives  $k$  different translations in total. Notably, SeedDiv does not sacrifice the translation quality to achieve diversity since random seeds do not affect the expected model performance.
We conduct the experiment on WMT14 En-De to evaluate the performance of SeedDiv. We use the base setting of Transformer and train the model for 150K steps. The detailed configuration is described in section 4.1. We also re-implement several existing methods with the same setting for comparison, including Beam Search, Diverse Beam Search (Vijayakumar et al., 2018), HardMoE (Shen et al., 2019), Head Sampling (Sun et al., 2020) and Concrete Dropout (Wu et al., 2020). We set the number of translations  $k = 3$ , and set the number of heads to be sampled as 3 for head sampling. We also implement a weaker version of our method SeedDiv-ES, which early stops the training process with only  $\frac{1}{k}$  of total training steps. We report the results of
Figure 2: Reference BLEU and pairwise BLEU scores of SeedDiv and other diverse translation methods on the test set of WMT14 En-De. We do not use compound split to keep consistency with previous work.
these methods in Figure 2.
It is surprising to see that SeedDiv achieves outstanding translation diversity besides the superior translation quality, which outperforms most methods on both translation quality and diversity. Only HardMoe has a better pairwise BLEU than SeedDiv, but its reference BLEU is much lower. The only concern is that SeedDiv requires a larger training cost to train multiple models, so we also use a weaker version SeedDiv-ES for comparison. Though the model performance is degraded due to the early stop, SeedDiv-ES still achieves a good trade-off between the translation quality and diversity, demonstrating the advantage of using the training randomness controlled random seeds to generate diverse translations. Therefore, we use SeedDiv as the technique for diverse distillation.
3.2 Reference Selection
3.2.1 Losses under Diverse Distillation
After diverse distillation, we obtain a dataset containing  $k$  reference sentences  $y_{1:k}$  for each source sentence  $x$ . Traditional data augmentation algorithms for NMT (Sennrich et al., 2016a; Zhang and Zong, 2016; Zhou and Keung, 2020; Nguyen et al., 2020) generally calculate cross-entropy losses on all data and use their summation to train the model:
$$
\mathcal {L} _ {\text {s u m}} (\theta) = - \frac {1}{k} \sum_ {i = 1} ^ {k} \log p (\boldsymbol {y} _ {i} | \boldsymbol {x}, \theta). \tag {7}
$$
However, this loss function is inaccurate for NAT due to the increase of data complexity. Sequence-level knowledge distillation works well on NAT by
reducing the complexity of target data (Zhou et al., 2020). In comparison, the target data generated by diverse distillation is relatively more complex. If NAT learns from the  $k$  references indiscriminately, it will not eventually converge to any one reference but generate a mixture of all references.
Using the multi-reference dataset, we propose to train NAT with reference selection to evaluate the model output with better accuracy. We compare the model output with all reference sentences and select the one with the maximum probability assigned by the model. We train the model with only the selected reference:
$$
\mathcal {L} _ {\max } (\theta) = - \log \max  _ {1 \leq i \leq k} p \left(\boldsymbol {y} _ {i} \mid \boldsymbol {x}, \theta\right). \tag {8}
$$
In this way, we do not fit the model to all references but only encourage it to generate the nearest reference, which is an easier but more suitable objective for the model. Besides, when the ability of autoregressive teacher is limited, the NAT model can learn to ignore bad references in the data and select the clean reference for the training, which makes the capability of NAT not limited by a specific autoregressive teacher.
In addition to minimizing all losses  $\mathcal{L}_{sum}(\theta)$  or the selected loss  $\mathcal{L}_{max}(\theta)$ , there is also an intermediate choice to assign different weights to reference sentences. We can optimize the log-likelihood of generating any reference sentence as follows:
$$
\mathcal {L} _ {m i d} (\theta) = - \log \sum_ {i = 1} ^ {k} p (\boldsymbol {y} _ {i} | \boldsymbol {x}, \theta). \tag {9}
$$
The gradient of Equation 9 is equivalent to assigning weight  $\frac{p(\boldsymbol{y}_i|\boldsymbol{x},\theta)}{\sum_i p(\boldsymbol{y}_i|\boldsymbol{x},\theta)}$  to the cross-entropy loss of each reference sentence  $\boldsymbol{y}_i$ . In this way, the model focuses more on suitable references but also assigns non-zero weights to other references.
We use a linear annealing schedule with two stages to train the NAT model. In the first stage, we begin with the summation  $\mathcal{L}_{sum}(\theta)$  and linearly anneal the loss to  $\mathcal{L}_{mid}(\theta)$ . Similarly, we linearly switch to the selected loss  $\mathcal{L}_{max}(\theta)$  in the second stage. We use  $t$  and  $T$  to denote the current time step and total training steps respectively, and use a constant  $\lambda$  to denote the length of the first stage. The loss function is:
$$
\mathcal {L} (\theta) = \left\{ \begin{array}{l} T _ {1} \mathcal {L} _ {\text {m i d}} (\theta) + (1 - T _ {1}) \mathcal {L} _ {\text {s u m}} (\theta), t \leq \lambda T \\ T _ {2} \mathcal {L} _ {\text {m a x}} (\theta) + (1 - T _ {2}) \mathcal {L} _ {\text {m i d}} (\theta), t > \lambda T \end{array} , \right. \tag {10}
$$
Models	Lmax(θ)	Lsum(θ)
For	Back	For	Back
AT	k×	1×	k×	k×
vanilla-NAT	k×	1×	k×	k×
CTC	1×	1×	1×	1×
Table 1: The calculation cost of  $\mathcal{L}_{max}(\theta)$  and  $\mathcal{L}_{sum}(\theta)$  for different models. 'For' and 'Back' indicate forward and backward propagations respectively.
where  $T_{1}$  and  $T_{2}$  are defined as:
$$
T _ {1} = \frac {t}{\lambda T}, \quad T _ {2} = \frac {t - \lambda T}{T - \lambda T}. \tag {11}
$$
In this way, the model learns from all references indiscriminately at the beginning, which serves as a pretraining stage that provides comprehensive knowledge to the model. As the training progresses, the model focuses more on the selected reference, which provides accurate training signals and gradually finetunes the model to the optimal state.
3.2.2 Efficient Calculation with CTC
To calculate the probability  $p(\pmb{y}|\pmb{x},\theta)$ , the vanilla-NAT must set the decoder length to the length of  $\pmb{y}$ . Therefore, calculating the probability of  $k$  reference sentences requires running the decoder for at most  $k$  times, which will greatly increase the training cost. Fortunately, for CTC-based NAT, the training cost is nearly the same since its decoder length is only determined by the source sentence. We only need to run the model once and calculate the probabilities of the  $k$  reference sentences with dynamic programming, which has a minor cost compared with forward and backward propagations. In Table 1, we show the calculation cost of  $\mathcal{L}_{max}(\theta)$  and  $\mathcal{L}_{sum}(\theta)$  for different models. We use CTC as the baseline model due to its superior performance and training efficiency.
3.2.3 Max-Reward Reinforcement Learning
Following Shao et al. (2019, 2021), we finetune the NAT model with the reinforcement learning objective (Williams, 1992; Ranzato et al., 2015):
$$
\mathcal {L} _ {r l} (\theta) = \underset {\boldsymbol {y}} {\mathbb {E}} [ \log p (\boldsymbol {y} | \boldsymbol {x}, \theta) \cdot r (\boldsymbol {y}) ], \tag {12}
$$
where  $r(\pmb{y})$  is the reward function and will be discussed later. The usual practice is to sample a sentence  $\pmb{y}$  from the distribution  $p(\pmb{y}|\pmb{x},\theta)$  to estimate the above equation. For CTC based NAT,  $p(\pmb{y}|\pmb{x},\theta)$  cannot be directly sampled, so we sample from
the equivalent distribution  $p(\boldsymbol{a}|\boldsymbol{x},\theta)$  instead. We recover the target sentence by the collapsing function  $\Gamma^{-1}$  and calculate its probability with dynamic programming to estimate the following equation:
$$
\mathcal {L} _ {r l} (\theta) = \underset {\boldsymbol {a}} {\mathbb {E}} [ \log p \left(\Gamma^ {- 1} (\boldsymbol {a}) \mid \boldsymbol {x}, \theta\right) \cdot r \left(\Gamma^ {- 1} (\boldsymbol {a})\right) ]. \tag {13}
$$
The reward function is usually evaluation metrics for machine translation (e.g., BLEU, GLEU), which evaluate the prediction by comparing it with the reference sentence. We use  $r(\pmb{y}_1,\pmb{y}_2)$  to denote the reward of prediction  $\pmb{y}_1$  when  $\pmb{y}_2$  is the reference. As we have  $k$  references  $\pmb{y}_{1:k}$ , we define our reward function to be the maximum reward:
$$
r (\boldsymbol {y}) = \max  _ {1 \leq i \leq k} r (\boldsymbol {y}, \boldsymbol {y} _ {i}). \tag {14}
$$
By optimizing the maximum reward, we encourage the model to move towards the selected reference, which is the closest to the model. Otherwise, rewards provided by other references may mislead the model to generate a mixture of all references.
Datasets We conduct experiments on major benchmark datasets for NAT: WMT14 English  $\leftrightarrow$  German (En  $\leftrightarrow$  De, 4.5M sentence pairs) and WMT16 English  $\leftrightarrow$  Romanian (En  $\leftrightarrow$  Ro, 0.6M sentence pairs). We also evaluate our approach on a large-scale dataset WMT14 English  $\rightarrow$  French (En  $\rightarrow$  Fr, 23.7M sentence pairs) and a small-scale dataset IWSLT14 German  $\rightarrow$  English (De  $\rightarrow$  En, 160K sentence pairs). The datasets are tokenized into subword units using a joint BPE model (Sennrich et al., 2016b). We use BLEU (Papineni et al., 2002) to evaluate the translation quality.
Hyperparameters We use 3 teachers for diverse distillation and set the seed to  $i$  when training the  $i$ -th teacher. We set the first stage length  $\lambda$  to  $2/3$ . We use sentence-level BLEU as the reward. We adopt Transformer-base (Vaswani et al., 2017) as our autoregressive baseline as well as the teacher model. The NAT model shares the same architecture as Transformer-base. We uniformly copy encoder outputs to construct decoder inputs, where the length of decoder inputs is  $3\times$  as long as the source length. All models are optimized with Adam (Kingma and Ba, 2014) with  $\beta = (0.9,0.98)$  and  $\epsilon = 10^{-8}$ , and each batch contains approximately 32K source words. On WMT14 En  $\leftrightarrow$  De and WMT14 En  $\rightarrow$  Fr,
Models	Iterations	Speedup	WMT14	WMT16
EN-DE	DE-EN	EN-RO	RO-EN
AT	Transformer (Vaswani et al., 2017) + distillation (k=3)	N	1.0×	27.51	31.52	34.39	33.76
N	1.0×	28.04	32.17	35.10	34.83
One-pass NAT	NAT-FT (Gu et al., 2018)	1	15.6×	17.69	21.47	27.29	29.06
CTC (Libovický and Helcl, 2018)	1	-	16.56	18.64	19.54	24.67
NAT-REG (Wang et al., 2019)	1	27.6×	20.65	24.77	-	-
Bag-of-ngrams (Shao et al., 2020)	1	10.7×	20.90	24.61	28.31	29.29
AXE (Ghazvininejad et al., 2020)	1	-	23.53	27.90	30.75	31.54
SNAT (Liu et al., 2021)	1	22.6×	24.64	28.42	32.87	32.21
GLAT (Qian et al., 2021)	1	15.3×	25.21	29.84	31.19	32.04
Seq-NAT (Shao et al., 2021)	1	15.6×	25.54	29.91	31.69	31.78
CNAT (Bao et al., 2021)	1	10.37×	25.56	29.36	-	-
Imputer (Saharia et al., 2020)	1	-	25.80	28.40	32.30	31.70
OAXE (Du et al., 2021)	1	-	26.10	30.20	32.40	33.30
AlgNART (Song et al., 2021)	1	13.4×	26.40	30.40	32.50	33.10
REDER (Zheng et al., 2021)	1	15.5×	26.70	30.68	33.10	33.23
CTC w/ DSLP&MT (Huang et al., 2021)	1	14.8×	27.02	31.61	34.17	34.60
Fully-NAT (Gu and Kong, 2020)	1	16.8×	27.20	31.39	33.71	34.16
REDER + beam20 + AT reranking	1	5.5×	27.36	31.10	33.60	34.03
Iterative NAT	iNAT (Lee et al., 2018)	10	2.0×	21.61	25.48	29.32	30.19
CMLM (Ghazvininejad et al., 2019)	10	-	27.03	30.53	33.08	33.31
RecoverSAT (Ran et al., 2020)	N/2	2.1×	27.11	31.67	32.92	33.19
LevT (Gu et al., 2019)	2.05	4.0×	27.27	-	-	33.26
DisCO (Kasai et al., 2020)	4.82	-	27.34	31.31	33.22	33.25
JM-NAT (Guo et al., 2020b)	10	-	27.69	32.24	33.52	33.72
RewriteNAT (Geng et al., 2021)	2.70	-	27.83	31.52	33.63	34.09
Imputer (Saharia et al., 2020)	8	-	28.20	31.80	34.40	34.10
Our work	CTC	1	14.7×	26.09	29.50	33.55	32.98
CTC + distillation (k=3)	1	14.7×	26.35	29.73	33.51	32.82
DDRS w/o RL	1	14.7×	27.18	30.91	34.42	34.31
DDRS	1	14.7×	27.60	31.48	34.60	34.65
DDRS + beam20 + 4-gram LM	1	5.0×	28.33	32.43	35.42	35.81
Table 2: Performance comparison between our models and existing methods. The speedup is measured on WMT14 En-De test set. N denotes the length of translation.  $k$  means ensemble distillation (Freitag et al., 2017) from an ensemble of  $k$  AT models. - means not reported.