D i MS : Distilling Multiple Steps of Iterative Non-Autoregressive Transformers for Machine Translation (2023)

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.8 and 12.9 BLEU points improvements in single-step translation accuracy on distilled and raw versions of WMT’14 De-En.Full code for this work is available here: https://github.com/layer6ai-labs/DiMS

## introduction neural machine translation models typically follow an autoregressive decoding strategy, generating the target sentence one token at a time. this sequential nature makes the inference process slow and dependent on the output sequence length. to address this limitation @xcite introduces the non-autoregressive transformer (nat). nat generates the entire target sentence in parallel, reducing the latency by an order of magnitude. nat can be considered as a member of a broader family of iterative non-autoregressive transformers (inat) @xcite @xcite 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, @xcite 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 been shifted to design single-step nat models @xcite @xcite . in order to preserve the enhancements obtained by multiple decoding iterations of inats, we introduce distill multiple steps (dims), a distillation algorithm applicable to a wide range of iterative models. given a pre-trained inat, referred to as teacher, a student aims to replicate the behavior of multiple iterative steps of the teacher with one decoding pass. this process resembles the wellknown knowledge distillation framework @xcite . however, instead of reducing the number of parameters, we aim to 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. the proposed distillation can be repeated iteratively, where at 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. alternatively, we propose updating the parameters of 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 1 depicts the dims algorithm. we demonstrate the effectiveness of our approach on several public datasets by showing that dims obtains substantial improvements on singlestep translation with gains of up to 7.8 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. we further show that ema considerably speeds up training and converges to a comparable accuracy with iterative distillation in a fraction of epochs. ## background in this section, we lay out a formal framework for inats. we use the setup of conditional masked language models (cmlm). cmlm first introduced in @xcite and subsequently adopted in many inat models @xcite @xcite . the source sentence, target sentence, and target sequence length are denoted by x, y and n , respectively. ## distillation of iterative non-autoregressive transformers increasing the number of decoding steps typically improves accuracy, but diminishes the computational advantage of inats. our objective is to reduce the number of decoding steps without degrading the performance. more specifically, we want to condense the translation quality of multiple steps of a teacher into one decoding pass of a 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 to reach a similar performance. 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. this process can be moved to a pre-processing phase, at the cost of large memory requirement. we propose to use partially masked reference sentences as 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. θ @xmath0. then, the teacher processes a 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 i and after the final pass, no masked token remains. this is similar to the inference procedure outlined in section 2.2, but instead of starting from a fully masked sentence, it starts from a partially masked one. the student is optimized to match the teacher's soft labels and a temperature is used to 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 one decoding pass. ## related works many techniques have been proposed for iterative non-autoregressive machine translation. earlier attempts include denoising autoencoder @xcite and insertion-deletion @xcite . more recently, @xcite introduced the mask-predict improving the performance of inats by employing a conditional masked language model. cmlmc @xcite and smart @xcite improve cmlm by incorporating a correction mechanism. disco @xcite ) 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 position constraint in cross-entropy can make the nat training challenging, therefore @xcite propose aligned crossentropy (axe), an objective that considers the best monotonic alignment between the target and the model's predictions. @xcite relaxes the monotonic assumption and introduces order agnostic cross-entropy (oaxe). ctc @xcite is a similar alignment-based objective that fixes the model output length and considers various alignments leading to the same target. imputer @xcite extends ctc to benefit from iterative refinements. glat @xcite shows that the optimization challenges of inats can be mitigated by introducing a curriculum learning focusing on sentences with only a few masked tokens in the early stages of the training and gradually increasing the masking ratio. engine @xcite assumes access to a pre-trained autoregressive model and optimizes a nat model to maximize the likelihood under the probability distribution defined by the pre-trained model. @xcite 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-supervised learning techniques @xcite @xcite 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 selfsupervised 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. ## discussion it is not completely clear why knowledge distillation works in general @xcite . 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 teacher-generated labels reduces this gap. the importance of the gap between labels and the model capacity has also been observed before @xcite . ## conclusion we introduce dims, an effective distillation algorithm that enhances the single-step translation quality of a pre-trained iterative model. this is 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 as a moving average of the student while keeping the training time comparable to one round of the distillation. experiments over raw and distilled datasets on four translation tasks for supervised and unsupervised variants validate the effectiveness and versatility of dims. 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 to obtain a strong iterative model and make it computationally efficient via dims. iii) large monolingual target the antibodies hunt down any nicotine molecules in the bloodstream , neutralising them before they reached the brain , preventing a smoker from getting a nicotine hit . teacher the antibodies hunt the nicotine molecules molecblood neutralize them before reach brain a smoker not experience high nicotine . student the antibodies hunt the nicotine molecules in the blood and neutralize them before they reach the brain , so a smoker does not experience a nicotine high .