We propose to train a non-autoregressive machine translation model to minimize the energy defined by a pretrained autoregressive model. In particular, we view our non-autoregressive translation system as an inference network (Tu and Gimpel, 2018) trained to minimize the autoregressive teacher energy. This contrasts with the popular approach of training a non-autoregressive model on a distilled corpus consisting of the beam-searched outputs of such a teacher model. Our approach, which we call ENGINE (ENerGy-based Inference NEtworks), achieves state-of-the-art non-autoregressive results on the IWSLT 2014 DE-EN and WMT 2016 RO-EN datasets, approaching the performance of autoregressive models. 1 1 Introduction The performance of non-autoregressive neural machine translation (NAT) systems, which predict tokens in the target language independently of each other conditioned on the source sentence, has been improving steadily in recent years (Lee et al., 2018; Ghazvininejad et al., 2019; Ma et al., 2019). One common ingredient in getting non-autoregressive systems to perform well is to train them on a corpus of distilled translations (Kim and Rush, 2016). This distilled corpus consists of source sentences paired with the translations produced by a pretrained autoregressive teacher system. As an alternative to training non-autoregressive translation systems on distilled corpora, we instead propose to train them to minimize the energy defined by a pretrained autoregressive teacher model. That is, we view non-autoregressive machine transWork partly done at Toyota Technological Institute at Chicago and the University of Chicago. lation systems as inference networks (Tu and Gimpel, 2018, 2019; Tu et al., 2019) trained to minimize the teacher's energy. This provides the non-autoregressive model with additional information related to the energy of the teacher, rather than just the approximate minimizers of the teacher's energy appearing in a distilled corpus. In order to train inference networks to minimize an energy function, the energy must be differentiable with respect to the inference network output. We describe several approaches for relaxing the autoregressive teacher's energy to make it amenable to minimization with an inference network, and compare them empirically. We experiment with two non-autoregressive inference network architectures, one based on bidirectional RNNs and the other based on the transformer model of Ghazvininejad et al. (2019). In experiments on the IWSLT 2014 DE-EN and WMT 2016 RO-EN datasets, we show that training to minimize the teacher's energy significantly outperforms training with distilled outputs. Our approach, which we call ENGINE (ENerGy-based Inference NEtworks), achieves state-of-the-art results for non-autoregressive translation on these datasets, approaching the results of the autoregressive teachers. Our hope is that ENGINE will enable energy-based models to be applied more broadly for non-autoregressive generation in the future. Non-autoregressive neural machine translation began with the work of Gu et al. (2018a), who found benefit from using knowledge distillation (Hin-ton et al., 2015), and in particular sequence-level distilled outputs (Kim and Rush, 2016). Subsequent work has narrowed the gap between non-autoregressive and autoregressive translation, including multi-iteration refinements (Lee et al., 2018; Ghazvininejad et al., 2019; Saharia et al., 2020; Kasai et al., 2020) and rescoring with autoregressive models (Kaiser et al., 2018; Wei et al., 2019; Ma et al., 2019; Sun et al., 2019). Ghazvininejad et al. (2020) and Saharia et al. (2020) proposed aligned cross entropy or latent alignment models and achieved the best results of all non-autoregressive models without refinement or rescoring. We propose training inference networks with autoregressive energies and outperform the best purely non-autoregressive methods. Another related approach trains an actor network to manipulate the hidden state of an autoregressive neural MT system (Gu et al., 2017; Chen et al., 2018; Zhou et al., 2020) in order to bias it toward outputs with better BLEU scores. This work modifies the original pretrained network rather than using it to define an energy for training an inference network. Energy-based models have had limited application in text generation due to the computational challenges involved in learning and inference in extremely large search spaces (Bakhtin et al., 2020). The use of inference networks to output approximate minimizers of a loss function is popular in variational inference (Kingma and Welling, 2013; Rezende et al., 2014), and, more recently, in structured prediction (Tu and Gimpel, 2018, 2019; Tu et al., 2019), including previously for neural MT (Gu et al., 2018b). Most neural machine translation (NMT) systems model the conditional distribution p  ( y | x ) of a target sequence y = (cid:104) y 1 , y 2 , ..., y T (cid:105) given a source sequence x = (cid:104) x 1 , x 2 , ..., x T s (cid:105) , where each y t comes from a vocabulary V , y T is (cid:104) eos (cid:105) , and y 0 is (cid:104) bos (cid:105) . It is common in NMT to define this conditional distribution using an autoregressive factorization (Sutskever et al., 2014; Bahdanau et al., 2015; Vaswani et al., 2017): log p  ( y | x ) = | y | (cid:88) t =1 log p  ( y t | y 0: t  1 , x ) This model can be viewed as an energy-based model (LeCun et al., 2006) by defining the energy function E  ( x , y ) =  log p  ( y | x ) . Given trained parameters  , test time inference seeks to find the translation for a given source sentence x with the lowest energy:  y = arg min y E  ( x , y ) . Finding the translation that minimizes the energy involves combinatorial search. In this paper, we train inference networks to perform this search approximately. The idea of this approach is to replace the test time combinatorial search typically employed in structured prediction with the output of a network trained to produce approximately optimal predictions (Tu and Gimpel, 2018, 2019). More formally, we define an inference network A  which maps an input x to a translation y and is trained with the goal that A  ( x )  arg min y E  ( x , y ) . Specifically, we train the inference network parameters  as follows (assuming  is pretrained and fixed): (cid:98)  = arg min  (cid:88) (cid:104) x , y (cid:105)D E  ( x , A  ( x )) (1) where D is a training set of sentence pairs. The network architecture of A  can be different from the architectures used in the energy function. In this paper, we combine an autoregressive energy function with a non-autoregressive inference network. By doing so, we seek to combine the effectiveness of the autoregressive energy with the fast inference speed of a non-autoregressive network. In order to allow for gradient-based optimization of the inference network parameters  , we now define a more general family of energy functions for NMT. First, we change the representation of the translation y in the energy, redefining y = (cid:104) y 0 , . . . , y | y | (cid:105) as a sequence of distributions over words instead of a sequence of words. In particular, we consider the generalized energy E  ( x , y ) = | y | (cid:88) t =1 e t ( x , y ) (2) where e t ( x , y ) =  y (cid:62) t log p  (  | y 0 , y 1 , . . . , y t  1 , x ) . (3) We use the  notation in p  (  | . . . ) above to indicate that we may need the full distribution over words. Note that by replacing the y t with one-hot distributions we recover the original energy. In order to train an inference network to minimize this energy, we simply need a network architecture that can produce a sequence of word distributions, which is satisfied by recent non-autoregressive NMT models (Ghazvininejad et al., Figure 1: The ENGINE framework trains a non-autoregressive inference network A  to produce translations with low energy under a pretrained autoregressive energy E . 2019). However, because the distributions involved in the original energy are one-hot, it may be advantageous for the inference network too to output distributions that are one-hot or approximately so. We will accordingly view inference networks as producing a sequence of T logit vectors z t  R |V| , and we will consider two operators O 1 and O 2 that will be used to map these z t logits into distributions for use in the energy. Figure 1 provides an overview of our approach, including this generalized energy function, the inference network, and the two operators O 1 and O 2 . We describe choices for these operators in the next section. We now consider ways of defining the two operators that govern the interface between the inference network and the energy function. As shown in Figure 1, we seek an operator O 1 to modulate the way that logits z t output by the inference network are fed to the decoder input slots in the energy function, and an operator O 2 to determine how the distribution p  (  | . . . ) is used to compute the log probability of a word in y . Explicitly, then, we O ( z )  O ( z )  z SX q  q  z STL onehot (arg max( z )) I SG onehot (arg max( q ))   q   z ST onehot (arg max( q ))  q  z GX  q   q   z Table 1: Let O ( z )   |V| 1 be the result of applying an O 1 or O 2 operation to logits z output by the inference network. The choices we consider for O 1 and O 2 , which we present generically for operator O and logit vector z , are shown in Table 1, and described in more detail below. Some of these O operations are not differentiable, and so the Jacobian matrix  O ( z )  z must be approximated during learning; we show the approximations we use in Table 1 as well. We consider five choices for each O : Jacobian approximation is necessary. (b) STL : straight-through logits. Here O ( z ) = onehot (arg max i z ) . O ( z )  z is approximated by the identity matrix I (see Bengio et al. (2013)). (c) SG : straight-through Gumbel-Softmax. Here O ( z ) = onehot (arg max i softmax( z + g )) , where g i is Gumbel noise. 2  O ( z )  z is approximated with  softmax( z + g )  z (Jang et al., 2016). (d) ST : straight-through. This setting is identical to SG with g = 0 (see Bengio et al. (2013)). (e) GX : Gumbel-Softmax. Here O ( z ) = softmax( z + g ) , where again g i is Gumbel noise; no Jacobian approximation is necessary. 2 g i =  log(  log( u i )) and u i  Uniform (0 , 1) . We evaluate our methods on two datasets: IWSLT14 German (DE)  English (EN) and WMT16 Romanian (RO)  English (EN). All data are tokenized and then segmented into subword units using byte-pair encoding (Sennrich et al., 2016). We use the data provided by Lee et al. (2018) for RO-EN. We consider two architectures for the pretrained autoregressive (AR) energy function. The first is an autoregressive sequence-to-sequence (seq2seq) model with attention (Luong et al., 2015). The encoder is a two-layer BiLSTM with 512 units in each direction, the decoder is a two-layer LSTM with 768 units, and the word embedding size is 512. The second is an autoregressive transformer model (Vaswani et al., 2017), where both the encoder and decoder have 6 layers, 8 attention heads per layer, model dimension 512, and hidden dimension 2048. We choose two different architectures: a BiLSTM tagger (a 2-layer BiLSTM followed by a fully-connected layer) and a conditional masked language model (CMLM; Ghazvininejad et al., 2019), a transformer with 6 layers per stack, 8 attention heads per layer, model dimension 512, and hidden dimension 2048. Both architectures require the target sequence length in advance; methods for handling length are discussed in Sec. 4.5. For baselines, we train these inference network architectures as non-autoregressive models using the standard per-position cross-entropy loss. For faster inference network training, we initialize inference networks with the baselines trained with cross-entropy loss in our experiments. The baseline CMLMs use the partial masking strategy described by Ghazvininejad et al. (2019). This involves using some masked input tokens and some provided input tokens during training. At test time, multiple iterations (refinement iterations) can be used for improved results (Ghazvininejad et al., 2019). Each iteration uses partially-masked input from the preceding iteration. We consider the use of multiple refinement iterations for both the CMLM baseline and the CMLM inference network. 3 4.4 Hyperparameters For inference network training, the batch size is 1024 tokens. We train with the Adam optimizer (Kingma and Ba, 2015). We tune the learning rate in { 5e  4 , 1e  4 , 5e  5 , 1e  5 , 5e  6 , 1e  6 } . For regularization, we use L2 weight decay with rate 0.01, and dropout with rate 0.1. We train all models for 30 epochs. For the baselines, we train the models with local cross entropy loss and do early stopping based on the BLEU score on the dev set. For the inference network, we train the model to minimize the energy (Eq. 1) and do early stopping based on the energy on the dev set. Non-autoregressive models often need a target sequence length in advance (Lee et al., 2018). We report results both with oracle lengths and with a simple method of predicting it. We follow Ghazvininejad et al. (2019) in predicting the length of the 3 The CMLM inference network is trained according to Eq. 1 with full masking (no partial masking like in the CMLM baseline). However, since the CMLM inference network is initialized using the CMLM baseline, which is trained using partial masking, the CMLM inference network is still compatible with refinement iterations at test time. translation using a representation of the source sequence from the encoder. The length loss is added to the cross-entropy loss for the target sequence. During decoding, we select the top k = 3 length candidates with the highest probabilities, decode with the different lengths in parallel, and return the translation with the highest average of log probabilities of its tokens. Effect of choices for O 1 and O 2 . Table 2 compares various choices for the operations O 1 and O 2 . For subsequent experiments, we choose the setting that feeds the whole distribution into the energy function ( O 1 = SX) and computes the loss with straight-through ( O 2 = ST). Using Gumbel noise in O 2 has only minimal effect, and rarely helps. Using ST instead also speeds up training by avoiding the noise sampling step. Training with distilled outputs vs. training with energy. We compared training non-autoregressive models using the references, distilled outputs, and as inference networks on both datasets. Table 5 in the Appendix shows the results when using BiLSTM inference networks and seq2seq AR energies. The inference networks improve over training with the references by 11.27 BLEU on DE-EN and 12.22 BLEU on RO-EN. In addition, inference networks consistently improve over non-autoregressive networks trained on the distilled outputs. Impact of refinement iterations. Ghazvininejad et al. (2019) show improvements with multiple refinement iterations. Table 3 shows refinement results of CMLM and ENGINE. Both improve with multiple iterations, though the improvement is much larger with CMLM. However, even with IWSLT14 WMT16 DE-EN RO-EN Autoregressive (Transformer) Greedy Decoding 33.00 33.33 Beam Search 34.11 34.07 Non-autoregressive Iterative Refinement (Lee et al., 2018) -25.73  NAT with Fertility (Gu et al., 2018a) -29.06  CTC (Libovicky and Helcl, 2018) -24.71  FlowSeq (Ma et al., 2019) 27.55  30.44  CMLM(Ghazvininejad et al., 2019) 28.25 28.20  Bag-of-ngrams-based loss (Shao et al., 2020) -29.29  AXE CMLM (Ghazvininejad et al., 2020) -31.54  Imputer-based model (Saharia et al., 2020) -31.7  ENGINE (ours) 31.99 33.16 Table 4: BLEU scores on two datasets for several non-autoregressive methods. Comparison to other NAT models. Table 4 shows 1-iteration results on two datasets. To the best of our knowledge, ENGINE achieves state-of-the-art NAT performance: 31.99 on IWSLT14 DE-EN and 33.16 on WMT16 RO-EN. In addition, ENGINE achieves comparable performance with the autoregressive NMT model. We proposed a new method to train non-autoregressive neural machine translation systems via minimizing pretrained energy functions with inference networks. In the future, we seek to expand upon energy-based translation using our method. We would like to thank Graham Neubig for helpful discussions and the reviewers for insightful comments. This research was supported in part by an Amazon Research Award to K. Gimpel.
objective method abstain result abstain abstain abstain objective abstain abstain abstain abstain method method result result method other other other objective other method other other other other other other method other objective method other method method method abstain method method method method abstain method method abstain method abstain other abstain abstain method abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain method abstain method method abstain abstain abstain method abstain method method abstain abstain abstain abstain method abstain method method method method method method abstain method method abstain abstain abstain abstain method abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain other abstain abstain result abstain objective method other other