Energy-Based Reranking: Improving Neural Machine Translation Using Energy-Based Models (2021)

Abstract: The discrepancy between maximum likelihood estimation (MLE) and task measures such as BLEU score has been studied before for autoregressive neural machine translation (NMT) and resulted in alternative training algorithms (Ranzato et al., 2016; Norouzi et al., 2016; Shen et al., 2016; Wu et al., 2018). However, MLE training remains the de facto approach for autoregressive NMT because of its computational efficiency and stability. Despite this mismatch between the training objective and task measure, we notice that the samples drawn from an MLE-based trained NMT support the desired distribution – there are samples with much higher BLEU score comparing to the beam decoding output. To benefit from this observation, we train an energy-based model to mimic the behavior of the task measure (i.e., the energy-based model assigns lower energy to samples with higher BLEU score), which is resulted in a re-ranking algorithm based on the samples drawn from NMT: energy-based re-ranking (EBR). We use both marginal energy models (over target sentence) and joint energy models (over both source and target sentences). Our EBR with the joint energy model consistently improves the performance of the Transformer-based NMT: +3.7 BLEU points on IWSLT’14 German-English, +3.37 BELU points on Sinhala-English, +1.4 BLEU points on WMT’16 English-German tasks.

## introduction autoregressive models are widely used for neural machine translation (nmt) @xcite @xcite . the autoregressive factorization provides a tractable likelihood computation as well as efficient sampling. the former results in the effective maximum likelihood estimation (mle) for training the * amirmohammad rooshenas is the corresponding author. interestingly, we observe that samples drawn from an nmt model trained using mle may have higher quality (measured with bleu) than the outputs of beam search. in particular, we draw 100 target samples for each source sentence from an nmt model trained using mle on the iwslt'14 german-english task, and observe that an oracle ranker -i.e. argmax y∼p nmt (y|x) bleu(., y * ), where (x, y * ) is the pair of source and gold target sentence -achieves the high score of 67.54, while the beam decoding achieves 33.87. we also look at the distribution of the spearman rank correlation coefficient of the drawn samples with respect to the log probability score of the baseline nmt (basenmt). figure 1 shows that there is no strong correlation between the bleu score ranking of samples and the log probability score ranking for the majority of source sentences; thus, maximum a priori (map) decoding is incapable of finding the desired output. in parallel to our study, @xcite also report that the mismatch regarding mle training of autoregressive models is attributable to the distribution of the probability mass rather than the parameter estimation, resulting in a poor map decoding. instead of looking for an alternate algorithm for parameter estimation, these results motivate us to explore training a parametric approximation of the metric, here bleu score: ω θ (y, x) ≈ bleu(y, y * ). therefore the decoding becomes: we use energy-based models (ebms) to parameterize ω θ (y, x). ebms @xcite are general parametric models that assign a scalar energy value to each configuration of input variables, thus defining an unnormalized probability distribution. although computing the partition function is intractable for general ebms, we only require the relative energy of the sampled sentences from the basenmt model, thus canceling out the normalization constant. in this paper we use two different energy-based models: marginal energy model (marginal-ebm) defined only over target sentences and joint energy model (joint-ebm) defined over both source and target sentences. figure 1 also shows the correlation coefficient of the energy ranking and bleu score using both marginal-ebm and joint-ebm. the shift in the coefficient distribution suggests that decoding based on energy scores results in better bleu scores compared to decoding based on the log probability scores of the basenmt model. also we observe that joint-ebm works better than using marginal-ebm as joint-ebm better captures the correlation of source and target sentences, while marginal-ebm is not directly conditioned on the source sentence. in this paper, we describe how to train ebms foot_0 to achieve the desired ranking. our energy ranker consistently improves the performance of transformerbased nmt on german-english, romanian-english and italian-english tasks from iwslt'14, the french-english task from iwslt'17, german-english task from wmt'14, and english-german task from wmt'16, as well as the low-resource sinhala-english and nepali-english tasks described in the flores dataset @xcite . ## energy-based reranking using ebm e θ to reweight the samples from an nmt defines a new probability distribution over the output sentences (see @xcite ): to sample from p θ (y|x), we sample k sentences from p nmt (y|x) using multinomial sampling from locally normalized distributions over the output and reweight the samples based on the energy network exp( -e θ (y,x) t ). then we resample two sentences, y 1 and y 2 , from the renormalized set, which defines a conditional distribution: p i (y|x) = exp(-e θ (y,x)/t ) k exp(-e θ (y k ,x)/t ) (a similar sampling approach has been used in @xcite ). now we train the energy model such that the ranking of y 1 and y 2 with respect to the energy model is consistent with their ranking with respect to the task metric, bleu score. in general, we assume y h is the sentence with the higher bleu score and y l is the sentence with with the lower bleu score. therefore, the training objective of e θ (y, x) becomes: where ξ(y i , x i ) is the margin violation and α is the margin weight. algorithm 1 outlines the whole training procedure. if we define the energy only over sentences of the target language, e θ (y), we can share the energymodel among multiple language pairs with the same target language. in this case we have to, first, sample the language l from our language set and then sample a sentence pair from the selected language training set d l . the probability of selecting a language is proportional to the number of sentences in its training set. algorithm 1 rank-based training of ebm pnmt(y|x) ← pretrained nmt e θ (y, x) ← energy based models for target sentences repeat l ← 0. for batch size do sample (xi, yi) from d yi ← collect k samples from pnmt(.|xi) p i (y) ← exp(-e θ (y,x)/t ) in this paper, we use bert @xcite to parameterize both e θ (y, x) and e θ (y). section 4.3 and 4.4 discuss the construction of e θ in detail. discriminative re-ranking was first introduced by @xcite for improving the performance of machine translation (mt). they have trained a linear separator using the perceptron learning algorithm to distinguish the top r translations from the rest of the translations in the n-best possible outputs. the features for the discriminator are extracted from both source and target sentences. @xcite combine the score of mt and the linear model using more complex syntactical features to re-rank the target sentences. here, we rely on the features learned by bert, and given the high capacity of the energy model, we train the energy model to respect the ranking of every pair of samples. @xcite describe using language model (lm) to improve the performance of nmt using shallow and deep fusion. shallow models combine the marginal probability of predicting each word in nmt and lm: log p nmt (y i |@xmath0, while deep fusion concatenates the hidden states of two models before predicting each word and uses parallel data to fine-tune the weights. similar to deep fusion, @xcite feed the unnormalized output of lm to the decoder of nmt. @xcite jointly train the lm and nmt using monolingual target-side data and parallel data, respectively. @xcite augment the parallel training data with monolingual data with the target language and back-translation. re-ranking with lm has also been explored by @xcite , where they decode the output based on log p(y|x) + λ 1 log p(x|y) + λ 2 log p(y), where p(y|x) is the direct model provided by nmt, p(x|y) is computed via back-translation and p(y) is an lm. our approach differs from the previous methods that use lms for re-ranking as we train our energy-based model to be consistent with the task measure instead of using pre-trained lms. in our experiments, we only explore the effect of using the direct model plus lm, nevertheless, backtranslation can also be added into our model for further improvement. recently, @xcite use masked language models (mlm) such as bert to score hypotheses from nmt. @xcite describe the score of a mlm as pseudo-log-likelihood score (pll). to calculate pll score of a sentence, each token w i in the sentence is sequentially masked, which allows the calculation of log p(w i |w \i ) from the output of the mlm. the normalized pseudolog-probability of the sentence is the average of logprobability of the masked words given the rest of the words in the sentence: in parallel to our work, @xcite proposes using two different bert models as an encoder of the source language (x-bert) and a decoder of the target language (y-bert). @xcite add an extra trainable encoder-decoder adaption module followed by a feed-forward module to each layer of the decoder and a feed-forward module to each layer of the encoder. (please see @xcite for more detail on the architecture.) for fine-tuning xy-bert for translation tasks, @xcite al. (2020) 2020 ) report a significant improvement over prior non-autoregressive models and superior performance comparing to autoregressive methods on iwslt'14 german-english task. their finding is consistent with our improvement using the pretained bert model. however, our joint-ebm model is a different way of using bert for translation, which does not require separate bert models for source and target language. please see section 4.9 for a detailed comparison. finally, other works also discuss using bert to improve the performance of nmt. @xcite describe initializing the embedding or the whole encoder with bert's parameters. @xcite use an attention model to incorporate the output of bert into encoder and decoder of nmt. in our approach, we use bert as an external energy-based ranker. ## analysis in this section, we study the sentence preference of marginal-ebr created by the energy ranking. ## conclusion and future work we introduce energy-based re-ranking (ebr) to improve the performance of autoregressive neural machine translation. despite its superior performance, ebr suffers from high latency because of its dependency on sampling from an autoregressive model. directly sampling from the underlying ebm can speed up the inference, which is our future direction in order to benefit from the power of energy-based models for machine translation.