Speed Up Iterative Non-Autoregressive Transformers by Distilling Multiple Steps
Authors: Sajad Norouzi
Rasa Hosseinzadeh
Felipe Perez
Maksims Volkovs
Venue: Submitted to ICLR 2023 (2023)
Keywords: non-autoregressive machine translation
knowledge distillation

Abstract: The computational benefits of iterative non-autoregressive transformers decrease as the number of decoding steps increases. As a remedy, we introduce Distill Multiple Steps (DiMS), a simple yet effective distillation technique to decrease the number of required steps to reach a certain translation quality. The distilled model enjoys the computational benefits of early iterations while preserving the enhancements from several iterative steps. DiMS relies on two models namely student and teacher. The student is optimized to predict the output of the teacher after multiple decoding steps while the teacher follows the student via a slow-moving average. The moving average keeps the teacher's knowledge updated and enhances the quality of the labels provided by the teacher. During inference, the student is used for translation and no additional computation is added. We verify the effectiveness of DiMS on various models obtaining 7 and 12.9 BLEU points improvements on distilled and raw versions of WMT'14 De-En, respectively.

Under review as a conference paper at ICLR 2023
SPEED UP ITERATIVE NON-AUTOREGRESSIVE TRANS- FORMERS BY DISTILLING MULTIPLE STEPS
Anonymous authors Paper under double-blind review
ABSTRACT
The computational benefits of iterative non-autoregressive transformers decrease as the number of decoding steps increases. As a remedy, we introduce Distill Multiple Steps (DiMS), simple yet effective distillation technique to decrease the number of required steps to reach certain translation quality. The distilled model enjoys the computational benefits of early iterations while preserving the enhancements from several iterative steps. DiMS relies on two models namely student and teacher. The student is optimized to predict the output of the teacher after multiple decoding steps while the teacher follows the student via slow-moving average. The moving average keeps the teacher's knowledge updated and enhances the quality of the labels provided by the teacher. During inference, the student is used for translation and no additional computation is added. We verify the effectiveness of DiMS various models obtaining 7 and 12.9 BLEU points improvements on distilled and raw versions of WMT' 14 De-En, respectively.
INTRODUCTION
Neural machine translation models typically follow an autoregressive decoding strategy, generating the target sen- tence one token a time. This sequential nature makes the inference process slow and dependent on the output se- quence length. To address this limitation Gu et al. (2018) introduces the Auoreerereeeeeivvve Transformer (NAT).
How does the solar system function ?
Gradien
Predict
Teacher
EMA
Student
NAT generates the entire target sentence in parallel, re- ducing the latency by an order of magnitude. NAT can be considered as a member of a broader family of iterative non autoregresseevee Transformers (iNAT) (Lee et al., 2020; Stern et al., 2019; Ghazvininejad et al., 2019) where the number of decoding steps is fixed and independent of the sequence length. By tuning the number of decoding steps, one can control the trade-off between speed and quality. While iNATs can be considered as efficient alternatives to their autoregressive counterparts, Kasai et (2020b) shows that autoregressive models can be sped up without loss in accuracy by combining shallow decoders with deep encoders. This diminishes the computational advantage of iNATs and challenges their motivation. The focus of recent work has thus shifted to design single-step NAT models (Ghazvininejad et al., 2020a; Qian et al., 2021; Du et al., 2021).
mask
xN
How <mask> <masko solar <masko function
Figure 1:
DIMS training. The student is trained to match the predictions of the teacher after several iterative steps. Teacher is up- dated with an exponential moving average of the student.
In order to preserve the enhancements obtained by multiple decoding iterations of iNATs, we introduce Distill Multiple Steps (DiMS), distillation algorithm applicable to wide range of iterative models. Given a tre-trained iNAT, referred to as teacher, student aims to replicate the behavior of multiple iterative steps of the teacher with one decoding pass. This process resembles the well-known knowledge distillation framework (Hinton et al., 2015). However, instead of reducing the number of parameters, we aim decrease the number of decoding passes. The final model then enjoys the translation quality of multi-steps iNAT with the computational efficiency of single-step translation.

---

Under review as a conference paper ICLR 2023
The proposed distillation can be repeated iteratively, where the end of each round the newly optimized student becomes the next teacher. While effective, iterative distillation is slow as it requires multiple rounds of training until convergence. Altematively, the teacher with an exponential moving average (EMA) of the student. This gradually transfers the new knowledge learned by the student to the teacher and can be viewed as a continuous variant of iterative distillation. Figure depicts the DiMS algorithm on high level.
propose updating the parameters of
We demonstrate the effectiveness of our approach on several public datasets by showing that DiMS obtains substantial improvements on single-step translation with gains of up to 7 BLEU points on the distilled training dataset, while the gains on raw datasets are even greater. Notably, we are able to surpass many leading NAT models designed specifically for single-step translation and we set new state the art on raw datasets. We further show that EMA considerably speeds up training and converges to comparable accuracy with iterative distillation in a fraction of epochs.
2 BACKGROUND
In this section, we lay out formal framework for iNATs. We use the setup of Conditional Masked Language Models (CMLM), the approach first introduced in Ghazvininejad et al. (2019) and sub- sequently adopted in many iNAT models (Ghazvininejad et al., 2020b; Kasai et al., 2020a; Saharia et al., 2020; Huang et al., 2021). The source sentence, target sentence, and target sequence length are denoted by x, y and N, respectively.
2.1 TRAINING
Given a partially masked reference sentence ý and the corresponding source context x, the model trained to reveal all the masked positions simultaneously (Ghazvininejad et al., 2019). From probabilistic perspective, this imposes a conditional independence assumption on the predicted tokens. Formally, the training loss
Ey~M(y)
Σ
-oog
1EE(S)
where M is distribution over all partially masked target sentences and & is function that returns the set of masked indices. The training objective above implicitly assumes access to the target sentence length. To resolve this issue, CMLM trains parametric model, length predictor, predict the output length.
2.2 INFERENCE
The inference begins by creating a template (() with Ñ masked tokens, where Ñ is the output of the length predictor. At iteration of the inference, the model predicts the translation r!() given (t-1) and x as inputs. Depending on the number of decoding iterations S, typically a linear unmasking policy is used where at each step N V/S tokens with the highest probability are revealed. This process repeated S times, resulting in a fully revealed sentence. Denoting the output probability of the model by Pe the t-th step of the inference can be formally written as:
(t) y = (t)
arg-topk
(t-1) {Po (1,0(x,,(01)) )((11)} jEC(S(C-1)) else y;
Note that multiple length candidates can be considered (e.g. 1) with the average token probability as a ranking criterion. This is similar to beam search in autoregressive models but applied to the output sequence length. It is referred to as length beam.
DISTILLATION OF ITERATIVE NON-AUTOREEEESSIVE TRANSFORMERS
Increasing the number of decoding steps diminishes the computational advantage of iNATs. Our ob- jective is to limit the number of decoding steps without degrading the performance. More specifically, we want to compress the performance gain from multiple iterative steps of a teacher into one decoding
2

---

Under review as a conference paper ICLR 2023
y
y
yr
Masked token
Step one revealed
Teacher step 1
Teacher step 2
Step two revealed
Reference token
1o, (5((), x, 2)
Figure 2: Two iterative steps of the teacher applied to a partially masked sentence.
pass of student. For instance, consider an iterative model (teacher) that uses eight decoding steps. By replicating four steps of the teacher with one decoding pass, two steps of the student would be sufficient reach similar performance. Ideally, one would hope for a single iteration of the student to replicate the top performance of the teacher, but the model's s limited expressiveness might prevent this.
The standard way of knowledge distillation would have the teacher generate soft labels for all intermediate iterations, and optimize the student to track the teacher's output with fewer steps, but doing such generation on-the-fly greatly increases the training cost. One can move this process to a pre-pree pre-processing phase, at the cost of large memory requirements. We propose to use partially masked reference sentences an approximation to the intermediate predictions of the teacher, which eliminates the need for several decoding passes or large memory capacity.
The distillation process starts by initializing the student and the teacher to the same pre-trained model with parameters ø i.e. 0₂ = 02 = & where O, and 0, denote the parameters of the student and teacher respectively. Then, the teacher processes partially masked sentence through n iterative steps with a linear unmasking policy. More precisely, i/n of the originally masked tokens are revealed up to step and after the final pass, no masked token remains. This similar the inference procedure explained in Section 2.2, but instead of starting from a fully masked sentence, it starts from partially masked one. The student optimized to match the teacher's soft labels and a temperature used control the smoothness of the labels. With enough capacity, the student is expected to imitate the behavior of n consecutive steps of the teacher with merely one decoding pass.
3.1 TRAINING LOss
We denote the output distribution after n iterative steps the partially masked sentence ý by To (y, (y,x, n.) where represents the parameters the model. The distillation loss can be described as: Eie((s) (L(,...P...) where P. To, (y,x, n2) P. Ie. (y, x, 1) and in subscript denotes the index the sentence. Note that the teacher's ss soft labels not come from the same decoding iteration i.e. whenever a token is revealed, the corresponding soft labels are fixed in pt. Thus, the student receives labels from various decoding steps of the teacher. Figure 2 depicts the process teachers follow to produce the labels for two iterative steps. From the student's point of view, the primary difference between DiMS and CMLM training (Section 2.1) is the use of soft labels generated by the teacher instead of the ground truth tokens.
To facilitate the distillation, we combine the KL-divergence with the Euclidean distance of the last layers' hidden states of the teacher and the student. This transfers the knowledge concealed within the hidden states that might not be discernible in soft labels. We refer to this as hidden state loss. Similar to the KL-divergence, the hidden state loss is computed over the masked indices. The fact that the hidden states have a positive impact on the distillation has also been observed in previous works (Romero al., 2014 4; Sanh al., 2019; Jiao et al., 2020).
To summarize, DIMS training loss has two terms: i) KL-divergence between distributions predicted by the teacher and the student. The Euclidean distance between the last hidden states of two

---

Under review as a conference paper ICLR 2023
Algorithm DiMS
Require:
Data set D, pre-trained model Ø, Hidden state loss factor X, teacher steps n, EMA momentum J4, learning rate
0g, 0g f -
Initialize teacher and student
while not converged do
(x, D
Sample data Sample masking
~ M(y)
P: To. X y,
Run the teacher for n iterative steps D Run the student for a single step D Compute the DIMS loss > Gradient based optimization of the student EMA Update of the teacher
P. + Io, (x
DDMS + Σ, (P.. + + - e
0. f Optimizer(( 20,, Vo, (D.MM.))))
02 - (1 10), + plf
end while
models. Denoting teacher's and student's last hidden state by e, and e,, DiMS loss can be written formally as:
DDMMS = Σ EL.. Xller, -ex,il?
where
To, (,,x,) and Io. (V, x.
The Syper-parameter À controls the contribution of hidden state loss. When the distillation is completed, the student is used for evaluation following the teacher's inference process explained in Section 2.2, thus the computation remains the same.
3.2 EMA UPDATE THE TEACHER
As the distillation progresses, the performance gap between multiple steps of the teacher and single-pass of the student shrinks, making the teacher's labels less informative. Two approaches can be considered sustain the usefulness of the teacher's labels: i) Increasing the number of teacher's iterative steps. ii) Restarting the distillation where the recently optimized student becomes the new
teacher and repeating this process several times, i.e. 0(m) 0!=-1) The former makes the training more expensive as the number of sequential steps grows and the latter requires repeated distillation rounds leading longer training time.
Instead, we propose updating the teacher with the student's recently learned knowledge. The idea is that the student's single-step output approaches the teacher's multi-step, the student's multi-step performance would improve as well, and it is in the interest of the distillation process to use the improved student as the new teacher. However, replacing the teacher directly with the student would hurt the training stability and can lead to pathological solution of mapping everything to constant vector. This degenerate solution shortcuts the LDMS loss by setting a global minimum zero. To alleviate this, update the teacher with
average of the student, which transfers the new knowledge learned by the student to the teacher a controlled manner. The updated teacher now provides better training target for the student, creating a positive feedback loop between the two models. Note that even with the slow-moving average, the degenerate solution is inevitable with long enough training. In practice,
www-exponential-moviiee
find that it can be appropriately controlled by learning rate tuning and early stopping. The teacher can also benefit from the ensembling effects of the EMA (Izmailov al., 2018). Figure depicts DIMS training conceptually.
Overall, DIMS introduces three new hyper-parameters: i) n, the number of iterative steps of the teacher. ii) X, the coefficient that controls the balance between KL-divergence and hidden state loss, and iii) the momentum of teacher's 's EMA update. Algorithm I contains the pseudo-code for DIMS training.
EXPERIMENTS
4.1 EXPERIMENTAL SETUP
We use the Fairseq library (Ott et al., 2019) for all the experiments and follow the default data splits. All models are transformers with encoder-decoder architecture, each having 6 layers and
4

---

Under review as a conference paper ICLR 2023
35
27.3
26.3
30.0
10.4
30.6
33.5
34
11.4
33.3
13.7
37.6
33
33.6
25
23.5
27
21.2
32.1
32.7
BLEU
10.9
25
OMLM
CHLM
29
CMLM CHLM
CHUM CHUM-DRRS
19
OHLM-DOM
OMLMC
23
22.8
CHUM-DIS
CHUNC
OMLIC
CHUMC
10.3
Step
Step
4) 1M 14 En-De
b) WMT De-En
WMT En-Ro
d) WMT 16 Ro-Fa
Figure 3: CMLM and CMLMC models distilled with DIMS on four WMT tasks. For each teacher we plot performance various decoding steps and contrast with single-step performance of the student.
512-dimensional hidden states. Adam optimizer with inverse squared root learning rate scheduler used along with mixed precision. EMA and hidden state loss are leveraged with two iterative steps of the teacher unless otherwise stated. We use early stopping based on single-step BLEU score on the validation set. The final model is the average of 5 best checkpoints. Dropout is disabled for the teacher and the student since empirical improvements are observed. We conduct experiments on both the raw and distilled dataset that is obtained from an autoregressive model (Gu et al., 2018). Training done with Tesla V100 GPUs (32 GB) and we report all the hypar-parameters in Section the appendix. The extra computational cost of distillation is fraction of original training. We report detailed comparison in Section F of the appendix.
4.2
MAIN RESULTS
Our main experiments are conducted on WMT' En-De and WMT' En-Ro datasets with two models: i) CMLM, a pivotal work in iNAT literature showing the effectiveness of conditional masked language models. ii) CMLMC, recent work improving CMLM by incorporating correction mechanism. The corresponding official repositories are used to train the teachers. Both models exploit a length predictor that is conditioned on the encoder's hidden states. To make the length predictor compatible with changes in the encoder, we keep the length predictor loss during distillation.
Figure contrasts the single-step BLEU score of students with teachers evaluated for various number of decoding steps. DiMS considerably improves the translation quality of the single-step inference, reducing or eliminating the gap with multi-step inference. For example, on the WMT" 14 De-En dataset, the single-step of CMLMC+DiMS matches the teacher's 8-step performance. We also compared our best single-step model with strong baselines in Table I showing the effectiveness of DIMS. Note that noisy parallel decoding (NPD) refers the usage of the autoregressive model that generated the distilled the data to rerank the NAT's predictions with different sentence lengths. NPD only applicable models with length predictor.
The performance of the leading iNATs is at best similar to the autoregressive model used for sequence level knowledge distillation. This limits the final performance of iNATs and makes training without distillation desirable (Huang et al., 1). The columns corresponding to raw dataset in Table I show that DIMS improves the raw performance by large margin even more than the corresponding distilled variant. For instance, DiMS gets more than 12 BLEU scores improvements on single-step evaluation of CMLMC.
For one decoding pass, when raw variants of CMLMC are distilled with DIMS the performance superior to training on the distilled dataset (without DiMS). This makes DiMS preferable to sequence-level knowledge distillation. Nevertheless, the best performance is obtained when the two distillation approaches are combined. On distilled datasets, DiMS either is competitive with state of the art or outperforms it. To the best our knowledge, DIMS achieves a new state of the art single step evaluation over raw datasets.
It is not completely clear why knowledge distillation works in general. But when it comes to DIMS, we hypothesize that the labels generated by the teacher make the task simpler for the student. In other words, it is difficult for the model to close the gap between its single step prediction and ground truth while distillation with maacher-generate labels reduces this gap. The importance of the gap between labels and the model capacity has also been observed before Mirzadeh et al. (2020).
5

---

Under review as a conference paper ICLR 2023
Model
1M
WMT'16
En-De
De-En
En-Ro
Ro-En
Raw
Dist.
Raw
Dist.
Raw
Dist.
Raw
Dist.
CTC (Libovický Helel (2018))
17.4
19.8
19.9
24.7
CMLM (Ghazvininejal el (2019))
10.6
18.1
21.8
21.2
27.3
28.2
SMART (Ghazvininejad al. (20206))
18.6
23.8
CMLMC (Huang et al. (2021))
11.7
19.6
16.4
23.6
21.4
28.2
21.8
29.0
Aux. Reg. (Wang el (2019))
20.7
24.8
baggoo-ngramn (Shao et al. (2020))
20.9
24.6
28.3
29.3
Hint-based Loss (Shao et (2020))
21.1
25.2
Bigram CRF (Sun et al. (2019)
23.4
27.2
AXE (Ghuzvininejad al. (2020a))
20.4
23.5
24.9
27.9
30.5
30.8
31.4
31.5
EM+ODD (Sun & Yang (2020)
24.5
27.9
ENGINE (Tu et al. (2020))
28.1
28.2
Imputer (Saharia al. (2020))
15.6
25.8
28.4
32.3
31.7
OAXE (Du et al. (2021))
22.4
26.1
26.8
30.2
32.4
33.3
FullyNAT + GLu & Kong (2020))
21.8
27.2
31.4
33.7
34.2
Flowseq NPD (m=30) (Maetal, (2019))
21.15
25.3
26.04
30.7
31.7
32.2
32.5
32.8
GLAT NPD (m=7) (Qian et al. (2021)
26.5
31.0
32.9
33.5
CMLMC DIMS
23.2
26.7
29.3
30.8
31.2
33.0
32.7
33.6
CMLMC + DIMS + NPD (m=7)
23.7
27.3
29.9
31.6
32.2
33.7
33.6
34.5
Table I: Comparison of the single-step test set BLEU score with previously published works.
4.3 RESULTS ALIGNMENT BASED MODEL
To show the versatility of DIMS, we conduct experiment on alingment-based models leveraging Connectionist Temporal Classification (CTC) (Graves et al., 2006) objective. The main difference between alignment- based models (Libovický & Helcl, 2018) and CMLM family is the method used for modeling the output length. In an alignment-basee model, output length is set to a large enough number (e.g. twice the length of the source sentence). By introducing special blank token, the model has the ability to decide the sentence length since blanks are removed from the final prediction. Training such models requires target which is what CTC objective does efficiently via dynamic programming. For example if the target sentence "ABC" and the output length is set four, then CTC considers " "_ABC", "A_BC", "AB_C", "ABC_" as valid target sentences where denotes the blank token.
objective that considers all the alignments mapping to the same
Imputer (Saharia et al., 2020) among few iterative models in this family. There is no official implementation of Imputer available online, therefore we implement a version ourselves (denoted with t) 1. We compare our implementation with the original paper Section A the appendix and show that two models have similar performance. Table 2 summarizes the results of the base DiMS applied to Imputer for both directions of the WMT' English- Germmnn dataset. DiMS outperforms its teacher which shows that the applicability of DIMS is not limited to models trained with the cross-entropy objective. Further details of Imputer training and distillation are explained in Section of the appendix.
Method
WMT' 14 En-De
WMT' 14 De-En
Imputer
28.96 26.43
25.87
Imputer' DIMS
29.75
Table Sinlge-step test set BLEU score for Imputer models trained on WMT' English-German.
Based imputer-pytoreeh
the following
implementation:
httpps//ggtaab.coo/zooiaaall

---

Under review as a conference paper ICLR 2023
35
30
Method
WMT De-En
25
CMLM
22.77
CMLM + U-DiMS
29,45
15
CMLM + DIMS
29.74
10
CMLMC
23.63
1eacher-
Teacher-10
CMLMC + U-DiMS
30,52
Student-1
CMLMC DIMS
30.81
33
Average Sentence Length
Table 3: Single-step test set BLEU score for models trained with U-DiMS.
Figure 4: Test set BLEU score on 1M De- De-En splitted based on the target sentence length for CMLM teacher and student.
4.4 UNSUPERVISED DIMS
In previous sections, we assume access to a parallel dataset and feed a partially masked reference sentence to both student and teacher. One can use the teacher to generate synthetic target sentences during the distillation. This relaxes the dependence on the references and enables using monolingual datasets for distillation. As usual, there is a trade-off between computation and sample quality i.e. using more decoding passes leads better data while increasing the computational requirements. We refer to this unsupervised distillation variant as U-DiMS. Note that unsupervised only refers to the distillation, and for training the teacher still require access to a parallel dataset. The only distinction between U-DiMS and DIMS is the usage of synthetic data generated by the teacher and the remaining parts are untouched. We run U-DiMS on WMT' 14 De-En for CMLM and CMLMC using two iterative steps to generate the synthetic samples. Table 3 shows the effectiveness of U-DiMS, obtaining a similar performance to DIMS.
4.5 ABLATION STUDIES
We conduct all of the ablation studies on CMLM over WMT' 16 En-Ro smaller than WMT' 14 and the validation set is used for evaluation.
4.5.1 HIDDEN STATE LOSS
To investigate the effects of hidden state loss, we conduct an ablation study in this section. The first block in Table 4 includes BLEU scores for the base DiMS model with and without this term. The single-step performance of the distilled model is improved over 2 BLEU points by leveraging this loss. This supports the fact that the hidden states contain extra information that is not available in soft labels. The exact value of is selected based on a grid search reported in Section of the appendix.
4.5.2 EMA
In order to establish the computational advantages of the slow-moving average, we compare it with running the base variant for 9 iterative rounds. Figure 5 demonstrates that the EMA variant is able match the iterative distillation with far fewer updates (almost equal to one round of the distillation).
We observed that it essential to move the teacher toward the student slowly. For example, when S 0.9, the collapse to degenerate solution (explained in Section 3.2) occurs before the end of the first epoch. We plot the validation curve for various values of fu Section B of the appendix showing the importance of the slow-moving average.
4.5.3 TEACHER DECODING STEPS
One hyper- parameter in DIMS algorithm is the number of teacher's decoding steps. In order to investigate the effect of this hyper-parameter, we set it to 2, 4, and 8 while turning EMA on and off. The two bottom blocks of Table 4 include the results of this ablation. Although running the teacher for 4 decoding steps shows superior performance without EMA, as soon as we turn it on the gap disappears. This shows that EMA can gradually improve the teacher and remove the need for several
7

---

Under review as a conference paper ICLR 2023
Method
1-Step BLEU
M
CMLM
25.77
31.4
CMLM + DIMS
30.85
31.2
CMLM + DIMS Hidden.
28.69
CMLM + DIMS (T=4)
31.04
CMLM + DIMS (T=8)
30.97
30.0
Roand
Round
CMLM + DIMS EMA
31.63
30.6
Round
CMLM + DIMS (T=4) + EMA
31.52
Round
DMA
CMLM + DiMS (T=8) EMA
31.36
40
90
140
190
240
Epoch
Table 4: BLEU score on WMT 16 En-Ro val- idation set with beam length set to one as its done for early stopping. T stands for the num- ber teacher decoding steps and it is set to two if not specified.
Figure 5: Validation set BLEU score on WMT" 16 En-Ro for iterative distillation and EMA model. Each round of iterative distillation is shown with a unique color and the end of the round is noted by black dot. The number of steps differs various rounds as we use early stopping.
iterative steps. Thus, we find no reason to set this distillation's computational cost.
larger than as only increases
4.6 ANALYSIS
We study the effect of target sentence lengths on DIMS performance. The test is divided into five equally-sized buckets based on the target length. The BLEU scores are reported for each bucket in Figure 4. The main benefit of the iterative model manifested by large sentences. The reason might be the fact that longer sentences require context and modeling it becomes challenging with the conditional independence assumption in NAT. It is clear in Figure 4, that the performance is improved in every bucket. This improvement is most visible in the bucket with the highest average sentence length. This is because the fact that the same bucket has the largest gap between the teacher's single and multi-step evaluation.
We combine the length predictor objective with ours to account for changes in the encoder's parame- ters. Interestingly enough, DIMS improves the performance of the length predictor depicted in Figure 6. This shows that the encoder benefits from the distillation as well enhancing the length prediction accuracy.
Student
Student
Teacher
Teacher
0.0
0.1
0.2
0.3
0.4
0.5
0.6
0.0
0.1
0.2
0.3
0.4
0.5
WMT'
CMLM on De-En
Figure 6: Comparison of teachers and students in predicting the target length. Top-1 means predicting the target exactly correct and Top-3 and Top-5 means being incorrect by I and 2 offsets, respectively.
Table shows an example from the WMT' 14 De-En dataset. The improvements in samples evident by comparing the predictions of the teacher and the student with the target sentence. We provide more qualitative examples in Section of the appendix.

---

Under review as a conference paper at ICLR 2023
Target
The antibodies hunt down any nicotine molecules in the bloodstream, neutralising them before they reached the brain, preventing smoker from getting a nicotine hit.
Teacher
The antibodies hunt nicotine molecules molechlood neutralize them before reach brain smoker not experience high nicotine.
Student
The antibodies hunt the nicotine molecules in the blood and neutralize them before they reach the brain, so smoker does not experience nicotine high.
Table 5: qualitative example from WMT' 14 De-En along with teacher and student's 's predictions on CMLMC.
RELATED WORKS
Many techniques have been proposed for iterative non-autoregreessive machine translation. Earlier attempts include denoising autoencoder (Lee al., 2018) and insertion-deletion (Stern al., 2019; Gu et al., 2019). More recently. Ghazvininejad et al. (2019) introduced the Mask-Predict improving the performance of iNAT's by employing a conditional masked language model. CMLMC (Huang et al., 20 and SMART (Ghazvininejad et al., 2020b) improve CMLM by incorporating correction mechanism. DisCo (Kasai et al., 2020b) is another variant conditioning each token on an arbitrary subset of the other tokens. DiMS is entangled with the progress in this domain as it requires a pre-trained iterative teacher.
The ordering assumption in cross-entropy can make the NAT training challenging, therefore Ghazvininejad et al. (2020a) propose aligned cross-entropy (AXE), an objective that considers the best monotonic alignment between the target and the model's predictions. Du et al. (2021) relaxes the monotonic assumption and introduces Order Agnostic Cross-Entropp (OAXE). CTC (Libovický & Helel, 2018) is similar alignment-based objective that fixes the model output length and considers various alignments leading to the same target. Imputer (Saharia et al., 2020) extends CTC to benefit from iterative refinements.
GLAT (Qian et al., 2021) shows that the optimization challenges of iNATs can be mitigated by introducing a curriculum learning focusing on sentences with only few masked tokens in the early stages of the training and gradually increasing the masking ratio. ENGINE (Tu et al., 2020) assumes access to pre-trained autoregressive model and optimizes a NAT model maximize the likelihood under the probability distribution defined by the pre-trained model.
Salimans & Ho (2021) applies a distillation technique similar to DIMS on generative models to decrease the number of required steps for generating high-quality images. In contrast to DiMS, the distillation is applied progressively. DIMS eliminates the need for progressive distillation by updating the teacher with EMA. Lastly, the proposed EMA has some resemblance to self self-supervised learning techniques (Grill et al., 2020; Caron et al., 2021; He et al., 2020) where two models are updated, one through gradient-based optimization and the other one through EMA. Despite this similarity, the motivations are quite different. In self-supervised learning, EMA is proposed as a technique to remove large negative sets whereas here EMA enhances the quality of the labels generated by the teacher.
6 CONCLUSION
We introduce DiMS, an efficient distillation algorithm that enhances the translation quality of a pre-trained iterative model, especially with single-step translation. This done by replicating the model's multi-step behavior through one decoding pass. The distillation can be repeated to achieve greater gains, but this increases the training time noticeably. We show that the same benefits are obtainable by setting the teacher time comparable to one round of the distillation. Experiments over raw and distilled datasets on four translation tasks with various models for supervised and unsupervised variants validate the effectiveness and versatility of DIMS.
moving average of the student while keeping the training
Potential directions for future works include: i) The same family of iterative models have been applied to automatic speech recognition, thus DiMS is applicable to this domain. ii) One can combine a pyramid of techniques introduced for iNATs obtain a strong iterative model and make it computationally efficient via DiMS. iii) Large monolingual sets can be used to distill models with U-DiMS.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties self-supervised vision transforerrs In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 9650-6660, 2021. (Cited on page 9.)
Cunxiao Du, Zhaopeng Tu, and Jing Jiang. Order-agnostic cross entropy for non-autoregressive machine translation. International Conference on Machine Learning, pp. 2849-2859. PMLR, 2021. (Cited pages 1, 6, and 9.)
Marjan Ghazvininejad, Omer Levy, Yinhan Liu, and Luke Zettlemoyer. Mask-predict: Parallel decoding of conditional masked language models. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp. 6112-6121, 2019. (Cited on pages 1, 2, 6, 9, and 13.)
Marjan Ghazvininejad, Vladimir Karpukhin, Luke Zettlemoyer, and Omer Levy. Aligned cross entropy for non von-autoregressive machine translation. In International Conference on Machine Learning, pp. 2515-3523. PMLR, 2020a. (Cited on pages 1, 6, and 9.)
Marjan Ghazvininejad, Omer Levy, and Luke Zettlemoyer. Semi-autoregressive training improves mask-predict decoding. arXiv preprint arXiv:2001.087: 2020b. (Cited on pages 2, 6, and 9.)
Alex Graves, Santiago Femández, Faustino Gomez, and Jürgen Schmidhuber. Connectionist temporal classification: labelling unsegmented sequence data with recurrent neural networks. In Proceedings of the 23rd international conference Machine learning, pp. 369-376, 2006. (Cited on page 6.)
Bean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self- supervupe learning. Advances in Neural Information Processing Systems, 33:21271-21288 2020. (Cited on page 9.)
Jiatao Gu and Xiang Kong. Fully non-eutorreeessvve neural machine translation: Tricks of the trade. arXiv preprint arX1v:2012./5833, 2020. (Cited on page 6.)
Jiatao Gu, James Bradbury, Caiming Xiong, Victor OK Li, and Richard Socher. Non-autoregressive
neural machine translation. In International Conference on Learning Representations, 2018. (Cited
on pages and 5.)
Jiatao Gu, Changhan Wang, and Junbo Zhao. Levenshtein transformer. Advances in Neural Informa- tion Processing Systems, 32, 2019. (Cited on page 9.)
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 9729-9738, 2020. (Cited on page 9.)
Geoffrey Hinton, Oriol Vinyals, Jeff Dean, et al. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.025:: 2(7), 2015. (Cited on page 1.)
Xiao Shi Huang, Felipe Pérez, and Maksims Volkovs. Improving non oon-autoregressvve translation models without distillation. In International Conference on Learning Representations, 2021. (Cited on pages 2, 5, 6,9 9, and 13.)
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Av- eraging weights leads to wider optima and better generalization. arXiv preprint arXiv: 800.055077 2018. (Cited page 4.)
Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. Tinybert: Distilling bert for natural language understanding. In Findings of the Association for Computational Linguistics: EMNLP 2020, pp. 4163-4174, 2020. (Cited on page 3.)
Jungo Kasai, James Cross, Marjan Ghazvininejad, and Jiatao Gu. Non-autoregressive machine trans- lation with disentangled context transformer. In International Conference on Machine Learning, pp. 1144-1155. PMLR, 2020a. (Cited on page 2.)
10

---

Under review as a conference paper at ICLR 2023
Jungo Kasai, Nikolaos Pappas, Hao Peng, James Cross, and Noah Smith. Deep encoder, shallow decoder: Reevaluating non-autoregressive machine translation. In International Conference on Learning Representations, 2020b. (Cited on pages and 9.)
Jason Lee, Elman Mansimov, and Kyunghyun Cho. Deterministic non-autoregressive neural sequence modeling by iterative refinement. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pp. 1173-1182, 2018. (Cited page 9.)
Jason Lee, Raphael Shu, and Kyunghyun Cho. Iterative refinement in the continuous space for non- autoregressive neural machine translation. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 1006-1015, 2020. (Cited on page 1.)
Jindřich Libovický and Jindřích Helel. End-to-end non-autoregressive neural machine translation with connectionist temporal classification. In 2018 Conference on Empirical Methods in Natural Language Processing, pp. 3016-3021. Association for Computational Linguisties, 2018. (Cited on pages 6 and 9.)
Xuezhe Ma, Chunting Zhou, Xian Li, Graham Neubig, and Eduard Hovy. Flowseq: Non- autoregressive conditional sequence generation with generative flow. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp. 4282 4292, 2019. (Cited on page 6.)
Seyed Iman Mirzadeh, Mehrdad Farajtabar, Ang Li, Nir Levine, Akihiro Matsukawa, and Hassan Ghasemzadeh. Improved knowledge distillation via teacher assistant. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 5191-5198, 2020. (Cited on page 5.)
Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. arXiv preprint arXiv:1904.00 00038, 2019. (Cited on page 4.)
Lihua Qian, Hao Zhou, Yu Bao, Mingxuan Wang. Lin Qiu, Weinan Zhang, Yong Yu, and Lei Li.
Glancing transformer for non-autoregressive neural machine translation. In Proceedings of the 59th Anmual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume I: Long Papers), pp. 1993-2003, 2021.
(Cited on pages 1, 6, and 9.)
Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and Yoshua Bengio. Fitnets: Hints for thin deep nets. arXiv preprint arXiv:1412.65555 2014. (Cited page 3.)
Chitwan Saharia, William Chan, Saurabh Saxena, and Mohammad Norouzi. Non-autoregressive machine translation with latent alignments. In EMNLP, 2020. (Cited on pages 2, 6, 9, 13, and 14.)
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusson models. In International Conference Learning Representations, 2021. (Cited page 9.)
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv: 1910.000 2019. (Cited page 3.)
Chenze Shao, Jinchao Zhang, Yang Feng, Fandong Meng, and Jie Zhou. Minimizing the bag-of- ngrams difference for autrerrveeeee neural machine translation. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 1 198-205, 2020. (Cited on page 6.)
Mitchell Stern, William Chan, Jamie Kiros, and Jakob Uszkoreit. Insertion transformer: Flexible sequence generation via insertion operations. In International Conference on Machine Learning, pp. 5976-5955 PMLR, 2019. (Cited on pages and 9.)
Zhiqing Sun and Yiming Yang. An em approach to non-autoregressive conditional sequence genera- tion. In International Conference on Machine Learning, pp. 9249-925 PMLR, 2020. (Cited on page 6.)
11

---

Under review as a conference paper at ICLR 2023
Zhiqing Sun, Zhuohan Li, Haoqing Wang. Di He, Zi Lin, and Zhihong Deng. Fast structured decoding for sequence models. Advances in Neural Information Processing Systems, 32, 2019. (Cited on page 6.)
Lifu Tu, Richard Yuanzhe Pang, Sam Wiseman, and Kevin Gimpel. Engine: Energy- based inference networks for con-autoregressive machine translation. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 2819-2826, 2020. (Cited on pages 6 and 9.)
Yiren Wang, Fei Tian, Di He, Tao Qin, ChengXiang Zhai, and Tie- Yan Liu. -on-attoreeressvve machine translation with auxiliary regularization. In AAAI, 2019. (Cited page 6.)
12


Reviews:
[{"review_id": "XjxYhjJ48w3", "reviewer": "Reviewer_mVeK", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Iterative NAT improves translation performance with multiple decoding steps, at the cost of sacrificing decoding speed. In response to this problem, this paper proposes to leverage distillation technique to imitate the behaviors of multiple steps with a single decoding step. To this end, the authors maintain two NAT models, namely teacher and student. \n\nThe student model learns to predict the multiple-step output of the teacher with two additional training objectives for the imitation learning: 1) a KL-divergence loss of the probability distributions of target tokens between the teacher and the student; 2) an Euclidean distance between the last layer representations of the two models.\n\nAs the training processes, the performance of student model is improved (for both single- and multiple- decoding steps), which can further enhance the teacher model. Accordingly, the teacher model is updated by the student via a slow-exponential-moving average.\n\nExperimental results show that the proposed approach consistently improves performance over CMLMC (Huang et al., 2022) across language pairs on both distilled and raw data.", "summary_of_the_review": "This paper focuses on improving the performance of single-step NAT (on raw data), which is a promising and important direction to make NAT be independent model. The proposed Distill Multiple Steps (DiMS) method is effective across datasets and model architectures.\n\nHowever, the method introduces several hyper-parameters, which require a grid-search for individual dataset (e.g. Section 4.5). This method also relies on an pretrained teacher model, which prevents NAT from learning from scratch.", "strength_and_weaknesses": "Pros:\n1. An promising direction to improve fully NAT, which can maintain the advantage of decoding speed;\n\n2. An interesting way to leverage the performance of multiple-step decoding via distillation.\n\nCons:\n\n1. The approach is relatively complicated. It introduces 3 additional hyper-parameters (decoding steps of teacher model, interpolation weights of KL and Euclidean losses, and ), which potentially threat the robustness of the approach. Are these hyper-parameters the same across models (e.g. CMLMC and Imputer) and datasets (e.g. En-De and En-Ro).\n\n2. The methods to distill multiple steps are practical but not innovative enough. The KL and Euclidean losses are straightforward, and widely-used in imitation learning. \n\n3. The distillation pipeline still relies on an external teacher model (e.g. the pertained NAT model in this work), which prevents the model from learning from scratch (e.g. DAT (Huang et al., ICML 2022)).\n\n\nSuggested Revisions:\n1. The ultimate goal of NAT is  to become an independent translation model that achieves comparable with or even better performance than the AT models. Now the majority of NAT models are base setting (i.e., 6 layers of 512 dimensions) trained on medium-scale data (e.g., 4.5M WMT14 En-De), which still lags behind AT models (e.g., big or deep models trained on 20M WMT17 En-De data). It would be appreciated if the authors can validate the effectiveness of NAT models on big settings and/or larger datasets.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "6531MdqLpss", "reviewer": "Reviewer_oFCe", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes to use progressive distillation for distilling non-autoregressive models, and update the teacher based on the exponential moving average of the student. The experiments assess the idea on the machine translation task and exhibit some improvement when comparing with single-step results.", "summary_of_the_review": "While the idea of using progressive distillation for reducing the inference iterations of non-autoregressive models seems appealing, the resulting quality is not convincing and therefore the work provides limited contribution.", "strength_and_weaknesses": "Strength: reducing the number of iterations for the non-autoregressive approach is a critical topic for the field, and the use of progressive distillation is a reasonable choice.\nWeakness: the final results are worse than without applying the progressive distillation, and therefore making the contribution limited.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "1-1ZL4SDu1", "reviewer": "Reviewer_ABHG", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This article proposed a distill multiple steps method to decrease the number of required steps to speed up iterative non-autoregressive transformer.", "summary_of_the_review": "Please see review in the previous sections.", "strength_and_weaknesses": "Strengths:\n\n1. The training of the distilled student model obtained improved training efficiency with a certain translation quality. \n\n2. The proposed approach gained improvement on the CMLM on the WMT14 EN-De and WMT16 Ro-EN benchmarks.\n\nWeaknesses:\n\n1. For many comparison methods in Table 1, the reported results appeared to be lower than the original paper (e.g., the CMLM model).\n\n2. If the training time of the teacher model is also calculated, the efficiency of the final student model may not necessarily be improved.\n\n3. As shown in Figure 3, the BLEU scores of the proposed method were inferior to those of the baseline CMLM and CMLMC models.\n\n4. There were too many hyper-parameters (e.g., n, lambda, and p) in the proposed DiMS.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject