ENGINE : Energy-Based Inference Networks for Non-Autoregressive Machine Translation (2020)

Abstract: 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.

## 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 @xcite @xcite . one common ingredient in getting non-autoregressive systems to perform well is to train them on a corpus of distilled translations @xcite . 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 trans-lation systems as inference networks @xcite @xcite trained to minimize the teacher's energy. this provides the nonautoregressive 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 @xcite . 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. ## related work non-autoregressive neural machine translation began with the work of @xcite , who found benefit from using knowledge distillation @xcite , and in particular sequence-level distilled outputs @xcite . subsequent work has narrowed the gap between nonautoregressive and autoregressive translation, including multi-iteration refinements @xcite @xcite and rescoring with autoregressive models @xcite @xcite . @xcite and @xcite 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 @xcite @xcite 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 @xcite . the use of inference networks to output approximate minimizers of a loss function is popular in variational inference @xcite , and, more recently, in structured prediction @xcite @xcite , including previously for neural mt @xcite . ## energy-based inference networks for non-autoregressive nmt most neural machine translation (nmt) systems model the conditional distribution p θ (y | x) of a target sequence @xmath0, y 2 , ..., y t given a source sequence @xmath1, x 2 , ..., x ts , where each y t comes from a vocabulary v, y t is eos , and y 0 is bos . it is common in nmt to define this conditional distribution using an "autoregressive" factorization @xcite @xcite : this model can be viewed as an energy-based model @xcite by defining the energy function given trained parameters θ, test time inference seeks to find the translation for a given source sentence x with the lowest energy: ŷ = 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 @xcite . 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): ## results effect of choices for o 1 and o 2 . ## iterations, engine is comparable to cmlm on de-en and outperforms it on ro-en. comparison to other nat models. ## conclusion we proposed a new method to train nonautoregressive 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.