YOUR AUTOREGRESSIVE GENERATIVE MODEL CAN BE BETTER IF YOU TREAT IT AS AN ENERGY - BASED ONE Anonymous authors
Paper under double-blind review
A BSTRACT
Autoregressive generative models are commonly used, especially for those tasks involving sequential data.
They have, however, been plagued by a slew of inherent ﬂaws due to the intrinsic characteristics of chain-style conditional modeling (e.g., exposure bias or lack of long-range coherence), severely limiting their ability to model distributions properly.
In this paper, we propose a unique method for training the autoregressive generative model that takes advantage of a well-designed energy-based learning objective.
We show that our method is capable of alleviating the exposure bias problem and increase temporal coherence by imposing a constraint which ﬁts joint distributions at each time step.
Besides, unlike former energy-based models, we estimate energy scores based on the underlying autoregressive network itself, which does not require any extra network.
Finally, thanks to importance sampling, we can train the entire model efﬁciently without requiring an MCMC process.
Extensive empirical results, covering benchmarks like language modeling, neural machine translation, and image generation, demonstrate the effectiveness of the proposed approach.
I NTRODUCTION
By factorizing the joint distribution into the product of a series of conditional distributions, autoregressive generative models (abbr.
ARGMs)
(Vaswani et al., 2017; Dai et al., 2019; van den Oord et al., 2016a;b; Salimans et al., 2017; Chen et al., 2018) simplify the difﬁcult challenge of modeling high-dimensional joint distributions.
They can be trained efﬁciently via maximum likelihood and generate samples of exceptional quality, making this technique popular for modeling distributions, especially for sequential data.
Nonetheless, despite their potency and ﬂexibility, ARGMs still have inherent weaknesses due to the intrinsic characteristics of chain-style conditional modeling.
For example, ARGMs usually suffer from a discrepancy of the input context distributions between the training and inference stages, which causes consequent error propagation ( i.e. , Exposure Bias (Ranzato et al., 2016; Bengio et al., 2015)).
Besides, due to the nature of greedy selection of beam search approximations, the decoded results from ARGMs may also lack in long-range coherence.
We consider one approach by which ARGMs could be adapted to reduce these concerns.
Earlier work, both heuristic and theoretical, has already been proposed with those goals.
For instance, the exposure bias problem of ARGMs can be alleviated to some extent with scheduled sampling (Bengio et al., 2015; Mihaylova & Martins, 2019), by mixing input contexts from both real data and autoregressive generation, during the training stage.
However, this scheme suffers from an over-correcting problem (Zhang et al., 2019).
In addition, at the inference stage, beam search makes it possible to choose more diverse candidates, improving the quality of generated sequences.
Nevertheless, this results in only marginal improvements in temporal coherence, since
ARGMs can only leverage previous decoded contexts without consideration of the whole sequence information.
Moreover, setting aside the difﬁculty in training them, energy-based models (EBMs) have demonstrated their effectiveness in modeling high-dimensional distributions in a variety of machine learning applications (Zhao et al., 2017; Arbel et al., 2021; Gao et al., 2021), without requiring the transformation of the target distribution into a product of conditional distributions.
As a result, several studies (Deng et al., 2020; Bakhtin et al., 2021; Durkan & Nash, 2019) attempt to combine
EBMs with ARGMs, expecting to beneﬁt from the strengths of both approaches.
However, though some positive results were obtained, the existing works preferred a two-stage optimization, which ﬁrst obtains a well-trained ARGM and then trains an additional EBM based on it.
Such an optimiza- tion strategy does not enable ARGMs to beneﬁt from the properties of EBM in modeling the joint distribution in a temporally more coherent way.
In this paper, we present a novel design for seamlessly integrating E nergy - based models into
A uto R egressive M odels (E-ARM).
Our training is based on an energy-based learning objective, which forces ARGMs training to ﬁt the joint distribution along with the conditional one at each time step.
Thanks to our well-designed energy function, the two involved models can share a single base network without additional parameters, that is, the base network not only serves as a generator that provides fake data to facilitate the training of EBMs like previous works (Che et al., 2020; Xiao et al., 2021; Durkan & Nash, 2019;
Deng et al., 2020), but also plays the role of modeling the energy surface.
This property makes it easy to plug E-ARM into the training of any autoregressive generative models.
Intuitively, the exposure bias in ARGMs is caused by the fact that the model is trained on real data rather than data generated by the model.
On the other hand, in the EBM’s optimization process for modeling joint densities, the negative phase of wake-sleep algorithms (Hinton, 2002;
Kim &
Bengio, 2016) requires sampling data from the EBM itself.
Along with the fact that our method combines the EBM and the ARGM seamlessly as a whole, E-ARM can reduce the discrepancy between input data of the training and inference stage, which mitigates the exposure bias problem of the ARGM.
On top of it, unlike ARGMs, which factor the joint distribution into a product of conditional distributions, EBMs are able to model the joint distribution directly and score each input at the sequence level instead of at the token level, which makes them capable of modeling longrange coherence.
Additionally, in order to optimize the proposed energy-based learning objective efﬁciently via gradient-based wake-sleep algorithms (Kim & Bengio, 2016), we present a way to estimate the negative phase gradient (which is a necessary component in the gradient-based wakesleep algorithms) through those samples generated with the autoregressive view instead of the EBM view, which would require an expensive Markov Chain Monte Carlo (MCMC) process.
This allows us to sidestep extremely time-consuming MCMCs, thus accelerating training.
In summary, the following contributions are made with this paper: i)
We introduce a novel scheme,
E-ARM, to integrate the EBM view into autoregressive generative models seamlessly; ii) we attempt to reduce the intrinsic problems of autoregressive models, such as exposure bias and weak temporal coherence, by optimizing an energy-based learning objective, which uses samples autoregressively generated; iii)
We demonstrate how to efﬁciently optimize our model constructed from a single network, using wake-sleep algorithms without MCMC; iv)
In a number of applications, such as language modeling, neural machine translation, and image generation, our model can achieve better results in comparison with relevant baselines.
B ACKGROUND 2. E NERGY - BASED M ODELS
Energy-based models (LeCun et al., 2006) can express any probability p ( x ) for x ∈ R K as
[Equation]
where E θ : R D → R denotes an energy function which aims to map a D -dimensional datapoint to a scalar, and Z ( θ ) = (cid:80) x exp( − E θ ( x )) denotes the normalizing constant, also known as the partition function.
Any function can be used as an energy function to represent an EBM as long as it can generate a single scalar given some input x and the normalizing constant is ﬁnite 1 .
Wake-sleep algorithms are commonly used to optimize EBMs (Hinton, 2002; Kim & Bengio, 2016; Grathwohl et al., 2020) via gradient-based approximate maximum likelihood.
Speciﬁcally, the gradient of the log-likelihood, which needs to be maximized, with respect to θ can be expressed as
[Equation]
The ﬁrst term in the right hand side of Eq.
2 is the negative phase term while the second term is called the positive phase term.
MCMC methods have been used (Hinton, 2002; Welling & Teh, 2011a) to approximately sample from p θ ( x ) , for estimating the negative phase term.
2. M ODELING DISTRIBUTIONS AUTOREGRESSIVELY
Autoregressive generative models (ARGM) 2 can decompose any joint distribution p ( x ) into a product of conditional distributions using the product rule of probability by ordering those random variables within the joint distribution and characterizing each random variable given all variables preceding it in that order.
Formally, we use x <k to denote the vector variable covering all random variables before the time step k and use x k denote the random variable at time step k .
Then we have
[Equation]
In general, modeling distributions autoregressively has achieved remarkable accomplishments in numerous areas (Vaswani et al., 2017; Radford et al., 2019; van den Oord et al., 2016c;b; Salimans et al., 2017)
thanks to its ability to avoid the challenging target of modeling joint high-dimensional distributions directly.
We primarily focus on autoregressive language models in this paper, but we also conduct experiments on image generation to validate the generality of our method.
2. E XPOSURE BIAS AND INCOHERENCE PROBLEMS IN AUTOREGRESSIVE MODELS
In the discussion about the defects of sequential autoregressive generative models, the exposure bias problem (Bengio et al., 2015; Ranzato et al., 2016) is an important issue, which greatly affects the model’s deployment performance.
During the training stage, the autoregressive model is always conditioned on ground truth token sequences.
In generation stage, however, the model has to rely on its own previously generated tokens to predict the next token, when the model is deployed.
If an incorrect token is selected, this error can be ampliﬁed in following steps because the next prediction will be made using an unusual input (one unlike those in the training set).
Besides, out of the consideration of efﬁciency, autoregressive decoding usually selects the most probable token at each time step, given the ones previously selected.
Such a scheme assumes the largest joint probability of the whole sequence can be achieved by separately choosing the most probable next token (given its previous context) over all time steps, which is only the local optimum.
Correspondingly, the chosen sequence can not always be the model’s optimum result.
I NTEGRATE EBM S INTO AUTOREGRESSIVE MODELS SEAMLESSLY
For a long time, as a result of compromises for improving training stability and efﬁciency ( e.g. , modeling a joint distribution by decomposing it and using a teacher-forcing training strategy), conventional autoregressive generative models have suffered from ﬂaws such as the exposure bias and the lack of long-range coherence.
To tackle these issues, we attempt to seamlessly integrate E nergybased models into A uto R egressive M odels (E-ARM), which can be regarded as a variant of autoregressive generative models blending with an energy-based learning objective.
Given a joint sequential distribution, E-ARM also addresses it autoregressively, that is, tackling tokens step by step under a speciﬁc order.
However, what differs from conventional ARGMs is that we attempt to model both the conditional and the joint distributions simultaneously at each time step.
In this way, E-ARM can model distributions conveniently in an autoregressive manner while avoiding those potential problems brought by ARGMs.
Formally, given a sequence of random variables ( x 1 , x 2 , . . .
, x K ) with length K , we introduce a parametric autoregressive model q θ ( x <k ) =
(cid:81)
k − 1 l =1 q θ ( x l | x <l ) ( k denotes the time step) with parameters θ .
Particularly, we deﬁne ˜
q θ ( x <k ) =
(cid:81) k − 1 l
=
m q θ ( x l | x <l ) (cid:81) m − 1 n =1 q ( x n | x <n ) , which means only those conditional distributions q θ ( x l | x <l ) of the most recent k − m time steps are involved in the current update of parameters θ while those distant conditional distributions q ( x n | x <n ) are treated as ﬁxed (The rationale behind such a design will be elaborated in Sec.4).
Then, we deﬁne p θ ( x k , x <k ) as a product of the autoregressive model and an EBM as follows,
[Equation]
where the energy function φ θ ( x k , x <k ) is deﬁned as the x k ’s negative corresponding component of the base network’s output logit with the input preﬁx context x <k
= ( x 1 , x 2 , . . .
, x k − 1 ) ( e.g. , given a sequence “This is Friday.”
and assuming the corresponding index of the token “Friday” in the vocabulary is i , then the value of − φ θ ( “Friday”, “This is” ) is the i -th component of the output logit, which is the straight input tensor of the ﬁnal softmax layer), and the normalization term Z θ =
E x (cid:48) <
k ∼ ˜
q θ ( x <k )
[ (cid:80) x
k e − φ θ ( x k , x (cid:48) <k ) ] .
Our primary goal is to make the distribution q θ ( x k | x <k ) to approach the real conditional p d ( x k | x <k ) whilst maintaining p θ ( x k , x <k ) as close to the real joint p d ( x k , x <k ) as possible at each time step, which can be achieved by minimizing the Kullback-Leibler (KL) divergence between the distributions,
[Equation]
) where λ adjusts the ratio between the two objectives.
In Eq.
5, the ﬁrst objective at each time step k can be easily optimized by cross entropy while the second objective is optimized in the sense of EBMs by wake-sleep algorithms (Hinton et al., 1995; Kim & Bengio, 2016), which minimizes the objective by descending the following gradient of θ according to Eq.
[Equation]
where we have E θ ( x k , x <k )
=
φ θ ( x k , x <k ) − log ˜
q θ ( x <k ) .
Optimization via Eq.
2 or 6 involves sampling data from the model and can thus lead to the discovery of non-data-like samples, whose likelihood is then explicitly reduced by the energy function.
E-ARM is therefore not plagued by the exposure bias problem.
Besides, because we model the joint distribution throughout the training process, E-ARM can assess the entire sequence as a whole and generate more coherent data using energy sampling (Deng et al., 2020).
OPTIMIZATION
In this section, we present how to efﬁciently optimize E-ARM.
To begin with, we optimize the ﬁrst objective in Eq.
5 as with conventional autoregressive models by reducing the per time-step crossentropy loss.
As for the second objective, we resort to descend the estimated gradient as shown in
Eq. 6.
Thanks to the importance sampling technique and our well-deﬁned energy function, we now show that an improved version of Eq.
6 has a simple and symmetric form that can be easily estimated whilst not requiring an expensive MCMC.
Speciﬁcally, by replacing E θ ( x k , x <k ) with the speciﬁc form φ θ ( x k , x <k ) − log ˜ q θ ( x <k ) , the gradient w.r.t.
θ in the positive phase of Eq.
6 can be written into
[Equation]
Similarly, we can get the negative phase gradient as
[Equation]
The ﬁrst term − E x <k ∼ p d [ ∂∂θ log ˜ q θ ( x <k )] in Eq. 7 is equivalent to the log-likelihood gradient of ˜ q
θ ( x <k ) , which means improvements in this direction will be automatically taken care of as a result of steps arising from the gradient of the ﬁrst KL-divergence in Eq. 5, albeit at the expense of changing the weight given to the second vs. the ﬁrst KL, λ .
Besides, because the estimation of the expectation operator over the data distribution p d is easy, and the score φ θ ( x k , x <k ) can be acquired simply accessing the output logit of ARGM (see the deﬁnition of φ θ in Sec. 3), the second term can likewise be readily estimated and optimized.
As a result, the positive phase optimization is both feasible and efﬁcient.
The negative phase gradient estimation, on the other hand, is more involved.
In Eq. 8, sampling data from p θ is required for estimating the expectation E p θ , whereas p θ is a parametric joint probability involving an energy-based unnormalized probability estimator that may require time-consuming
MCMC methods to generate data.
However, thanks to importance sampling, we can substitute the troublesome computation of the expectation over the distribution p θ with the expectation over the distribution q θ , which can generate samples autoregressively without MCMC.
Formally, the negative phase gradient E x k , x <k ∼ p θ [ ∂∂θ E θ ( x k , x <k )] is equivalent to the following formulation (See the detailed derivation in Appendix A),
[Equation]
where w ( x <k ) =
[Equation]
According to Eq. 9, all the estimated expectations only need sampling from the autoregressive model rather than the joint distribution, and the reweighing weight w in Eq.
10 also does not involve expectation computation over distribution p θ .
Generally, producing data from an autoregressive model is a very simple ancerstral sampling process, as compared with sampling straight from an
EBM, which needs MCMC approaches (Durkan & Nash, 2019).
On account of that, the optimization process can be much more efﬁcient.
Besides, the term E x <k ∼ ˜ q
θ
( x <k )
[ w ( x <k ) ∂∂θ log ˜
q θ ( x <k )] in Eq.
9 is equivalent to a re-weighted version of the gradient of q θ ’s information entropy with respect to θ .
This term can be optimized similarly to the teacher-forcing training of autoregressive model with the “teacher” sequence generated autoregressively by the model itself.
Actually, the scheduled sampling methods (Bengio et al., 2015; Ranzato et al., 2016; Mihaylova & Martins, 2019) are similar to this term but without the reweighting factor.
Furthermore, it is worth noting that for a sequence with total length K , since we add a constraint to ﬁt the joint distribution p θ at each time step k , Eq. 9 actually has K counterparts with different time steps.
If we use the q θ ( x <k ) directly instead of ˜ q θ ( x <k ) in the Eq.
4 to deﬁne p θ ( x k , x <k ) , due to the fact that the distribution q θ ( x <k ) modeled by an autoregressive model can be naturally broken up into pieces, simply summing up these K gradients results in the term
K (cid:88) k
[Equation] [Equation]
where l indicates the speciﬁc index of the current token in the entire sequence.
As a result, earlier time steps (smaller l ) will get stronger training signals (larger K + 1 − l , indicating more gradient terms), giving rise to imbalanced training for different time steps.
To solve this, we introduce ˜ q θ ( x <k ) as (cid:81) k − 1 l = m q θ ( x l | x <l ) (cid:81) m − 1 n =1 q ( x n | x <n ) to deﬁne p θ ( x k , x <k ) shown in Sec. 3, allowing gradients only back propagate through conditional distributions w.r.t.
a few recent tokens 4 .
This explains our proposal of using ˜ q θ ( x <k ) to deﬁne p θ ( x k , x <k ) .
Ultimately, combining Eq.
7 and Eq. 9 , at each time step k , we can optimize p θ ( x k , x <k ) via descending the estimated gradient of θ as follows,
[Equation]
From Eq.
12, we can see that the only difference between two phases is that in the negative phase, the expectation over ˜ q θ have a reweighing weight w for each sample.
The reweighing weight w in
Eq.
10 and Eq.
12 can be further reﬁned (see the derivation in Appendix B) and we can observe that
[Equation]
where µ ( x <k ) =
p θ ( x <k ) ˜
q θ ( x <k ) indicates the possibility of which distribution the preﬁx context x <k is most likely to come from, the distribution p θ or the distribution ˜ q θ .
Correspondingly, w ( x <k ) reﬂects the context x <k ’s relative magnitude of µ ( x <k ) compared with the average among all potential contexts—the larger the value of w ( x <k ) , the more likely the context x <k in the data space coming from p θ , which is modeled by the product of autoregressive models and EBMs.
During training, those input sequences with contexts more likely under p θ than q θ will be assigned larger weights w while others will be assigned smaller weights w .
In general, E-ARM ought to be viewed as a new learning pattern for autoregressive models that ensures our base autoregressive network stays close to the real distribution p d .
We found that training from scratch with the energy-based learning objective of in Eq.12 alone did not work well. The reason is that at the initial stage of the training process, what we have is just a randomly initialized autoregressive network which outputs sequences with random values given any context.
This indicates disjoint supports between the real sequence’s distribution p d and distribution p θ modeled by ARGMs. If we only use the energy-based learning objective of Eq.
12, the whole gradient E p d ( x ) [ ∂∂θ log p θ ( x )] in Eq.2 would be 0 due to disjoint supports between p d and p θ .
As a result, in order to make the optimization more feasible, we must maintain the cross-entropy loss low throughout training and pre-train as a pure ARGM for a few epochs before introducing the E-ARM objective.
Actually, the starting epoch of E-ARM is a hyper-parameter, and we discuss it in the Sec. 5.2.
Following the excellent work of Deng et al.
(2020); Bakhtin et al.
, we also adopt Top-K energy re-sampling in the inference stage, which means that in the generative process, we ﬁrst gather multiple candidate sequences generated autoregressively, and then re-sample from them based on their energy scores estimated by the network’s logit at the last time step where the entire sequence has been processed.
Since we employ the EBM to model the joint distribution at each time step, such a re-sampling strategy can mitigate the undesirable impact of the greedy selection of one token at a time, and we found this variation to increase the coherence of generated samples.
E XPERIMENTS
To empirically corroborate the effectiveness of E-ARM and show its broad applicability, we conduct extensive experiments covering three machine learning applications, which are neural machine translation (NMT), language modeling, and image generation.
In this section, we will introduce the three corresponding experimental setups, followed by an analysis of the obtained results.
We will release the source code once upon acceptance.
5. A PPLICATION TO N EURAL M ACHINE T RANSLATION
E-ARM is ﬁrst evaluated in the context of neural machine translation (NMT), which is a conditional generation task and is important in the natural language processing (NLP) ﬁeld.
We ﬁrst analyze
E-ARM on the IWSLT14 dataset, which includes six different language pairs ( { German, Spanish,
Italian } → English and English → { German, Spanish, Italian } ).
In addition, we test E-ARM on the
WMT16 (English → German) benchmark to make sure we evaluating E-ARM on a larger dataset.
Hereafter we abbreviate English, German, Spanish, Italian as “En”, “De”, “Es”, “It”.
The weight λ in Eq.
5 is set as 0.05 for all translation tasks.
We use one size of transformer (“Base-IWSLT”) for the IWSLT14 benchmark and two sizes of transformer (“Base-WMT”, “Large-WMT”) for the
WMT16 benchmark 5 .
Scheduled Sampling is carried out following Mihaylova & Martins (2019).
The results of IWSLT14 tasks are shown in Table 1.
We test not only the pure performance of
E-ARM but also the compatibility with other techniques.
Speciﬁcally, we can observe that (1) without any particular engineering, E-ARM outperforms the base autoregressive translation model
[Table]
trained with cross-entropy singly by 0.62 (31.92 → 32.54) BLEU points in average, especially on three translation pairs—38.61 → 39.56 on Spanish-to-English, 30.48 → 31.13 on Italian-toEnglish, 35.42 → 36.07 on English-to-Spanish.
(2) E-ARM is compatible with other techniques like scheduled sampling, which can help alleviate the exposure bias problem to some extent.
They are not mutually exclusive and can work together to further improve the performance of the base
ARGM.
However, since scheduled sampling can reduce exposure bias and beam search can somewhat alleviate the ﬂaws caused by greedy selection at each time step, the performance gain of
E-ARM when all these tactics are combined is only 0.34 (34.21 → 34.55), which is lower than the 0.62 (31.92 → 32.54) obtained when the model is purely trained without these other techniques.
[Table]
Additionally, Table 2 shows the performance of E-ARM on the WMT
English → German task.
For two different model sizes, enabling label smoothing (L.S.) improves model performance by 0.52 and 0.35, respectively.
The performance of the base transformer model further increases to 28.36 BLEU points when scheduled sampling (S
[Table]
S.) is used, while the larger model improves to 29.23 points.
E-ARM paired with label smoothing and scheduled sampling yields the highest scores of 28.62 and 29.44, respectively.
Overall, our training strategy outperforms
ARGM’s vanilla teacher-forcing training and can have uniformly favorable impacts across different models and dataset sizes.
5. A PPLICATION TO L ANGUAGE M ODELING
[Table]
To further demonstrate E-ARM’s consistency in reducing ﬂaws of autoregressive generative models, we also conduct language modeling experiments.
The WikiTextdataset (Merity et al., 2017), which is the largest word-level language modeling benchmark with long-term dependency, was chosen as the testbed.
It comprises 103 million training tokens from 28 thousand articles, with an average length of 3.6 thousand tokens per article, which allows model to evaluate the ability of modeling long-term dependency.
Two network structures are mainly tested, which are Transformer-Base (Vaswani et al., 2017) and Transformer-XL (Dai et al., 2019)
(Tr-Base and
Tr-XL for short respectively hereafter).
[Table]
The ﬁnal results are reported in Table 3.
We can see from the results that E-ARM outperforms baselines with clear margins for different types of models.
Speciﬁcally, the Transformer-Base improves performance by 0.67 PPL points (from 30.56 to 29.89), while the Transformer-XL improves model by 0.
PPL points (from 24.20 to 23.81).
Our strategy does not change the structure of the base network nor introduces any additional module or learnable parameters, therefore we can conclude that the performance boost is solely from the introduced energybased learning objective.
In addition, we study the effect of hyper-parameter settings on the performance of language modeling, which can be seen in Table 4.
From this, we may deduce that starting E-ARM training at the 15-th epoch yields the best results, whereas starting earlier or later yields a performance decline.
It is reasonable because, if E-ARM was introduced too early, the autoregressive model may not have been optimized well at that moment.
As a result, generative quality would be terrible, and make energy-based training unstable.
On the other hand, the underlying autoregressive model can be modiﬁed only marginally if E-ARM is introduced when the ARGM training is virtually complete.
Besides, from the vertical perspective which presents the impact of different λ , we can observe that the best λ in
Eq. 5 is 0.05.
The ﬁrst line of the table indicates the baseline of training the autoregressive model with pure cross-entropy loss.
5. A PPLICATION TO I MAGE G ENERATION
In order to illustrate the effectiveness and generality of our method in processing different modality tasks, we further show the results of applying E-ARM to image generation in this section.
We apply
E-ARM to Pixel-CNN (Van Oord et al., 2016) and its variant Gated Pixel-CNN (Oord et al., 2016).
Experiments are carried out on the MNIST and CIFAR-10 datasets.
[Table]
Table 5 summarizes the quantitative results measured by per-pixel negative log-likelihood (NLL), while Figure 1 depicts some of the generated samples.
We can see that with the help of our E-ARM, both the Pixel-CNN and the Gated Pixel-CNN can obtain improvements in all datasets (0.17 → 0. and 3.14 → 3.07 for Pixel-CNN on MNIST and CIFAR10 respectively and 0.14 → 0.12 and 3.
→ 2.97 for Gated Pixel-CNN on MNIST and CIFAR10 respectively).
This is further evidence in favour of the energy-based learning objective for improving autoregressive models.
R ELATED W ORKS 6. A UTOREGRESSIVE GENERATIVE MODELS
Modeling high-dimensional data distributions directly is usually a rather challenging task due to “the curse of dimensionality” (Bellman, 1954).
One alternative method is to sequentialize the random variables and then factorize the joint probability distribution into the product of conditionals based on the sequence structure, which is exactly the core idea of autoregressive generative models (ARGMs).
ARGMs have been very successful, in particular for sequential data.
For example, ARGMs have been widely used in language modeling (Vaswani et al., 2017;
Dai et al., 2019; Radford et al., 2019), audio synthesis (van den Oord et al., 2016a), and even image generation (van den Oord et al., 2016c;b; Salimans et al., 2017).
The advantages of ARGMs are however balanced by issues of (1) exposure bias (Ranzato et al., 2016; Bengio et al., 2015; Song et al., 2020), due to the discrepancy in input context distributions between the training and inference stages, and (2) weak long-range coherence, due to the inherent greedy selection of one token at a time without look-ahead.
6. E NERGY - BASED MODELS
In the ﬁeld of generative modeling, energy-based models (EBMs) have been widely used (Zhao et al., 2017; Arbel et al., 2021; Gao et al., 2021).
The primary idea behind EBMs is to decompose the dependencies between variables ( e.g. images and labels) through different terms of an energy function, assigning low energies to proper conﬁgurations found in the dataset, while assigning high energies to incorrect or unseen ones (LeCun et al., 2006).
Due to the challenge of sampling from EBMs, training EBMs by wake-sleep algorithms (Hinton, 2002; Kim & Bengio, 2016; Grathwohl et al., 2021), which require expensive MCMC approaches, has been notoriously difﬁcult, especially on high-dimensional data like images or texts.
Stochastic
Gradient Langevin Dynamics (SGLD)
(Welling & Teh, 2011a) is a frequently used gradient-based
MCMC approach that injects noise into parameter updates and anneals the step size during the course of training, and which has been adopted in numerous prior works (Nijkamp et al., 2019; Du & Mordatch, 2019; Grathwohl et al., 2020).
However, these gradient-based MCMC methods require enormous extra computing overheads and are not applicable when the input is discrete like for text sequences (Deng et al., 2020).
As a result, a variety of recent works attempt to explore the strategy of training an EBM without
MCMC.
In particular, Bakhtin et al.
; Xu et al.
(2021a) ; Gao et al. (2020) optimize the EBMs by using noise contrastive estimation (NCE) (Gutmann & Hyv ¨ arinen, 2010; Ma & Collins, 2018).
Durkan & Nash (2019) estimate the intractable normalization component by utilizing ARGMs and importance sampling.
Che et al.
; Wang et al.
(2021) skirt the challenge of collecting data in the high-dimensional data space by producing data in the lower-dimensional feature space, which improves sampling efﬁciency.
C ONCLUSIONS AND FUTURE WORK
In this paper, we propose a novel method dubbed E-ARM to integrate energy-based models into autoregressive generative models seamlessly, with an energy-based training objective that exploits an underlying autoregressive model.
This is achieved by deﬁning the energy function from the output logits of the base autoregressive network, to model the unnormalized joint distribution of the subsequence up to each time step.
We also found ways to improve training of E-ARM using importance sampling, avoiding the requirement of MCMC for the energy-based training.
Experimental results on two language tasks and one vision task demonstrate the effectiveness of E-ARM to alleviate exposure bias and incoherence problems of ARGMs.
In the future, we expect to extend E-ARM on other sequential generation tasks ( e.g. text summarization, audio generation), and incorporate the proposed methodology into other advanced autoregressive architectures.
A T HE DERIVATION OF THE NEGATIVE PHASE GRADIENT
In this section, we show the detailed derivation of Eq. 9.
Formally, as shown in Sec. 3, given an autoregressive model q θ ( x <k ) = (cid:81) k − 1 l =1 q θ ( x l | x <l ) ( k denotes the time step) with parameters θ , we deﬁne a product of the autoregressive model and an EBM as follows
[Equation]
where ˜ q θ ( x <k )
= (cid:81) k − 1 l
=
m q θ ( x l | x <l ) (cid:81) m − 1 n =1 q ( x n | x <n ) .
Under such deﬁnition, only those conditional distributions q θ ( x l | x <l ) of the most recent k − m time steps are involved in the current update of parameters θ while those distant conditional distributions q ( x n | x <n ) are treated as ﬁxed.
We have explained the rationale and intuition in Sec.4.
Z θ is the normalization term and equal to
E x
[Equation]
<k ∼ ˜
q θ ( x <k )
[ (cid:80) x
k e − φ θ ( x k , x (cid:48) <k ) ] .
The optimization of p θ ( x k , x <k ) includes two phases, and the gradient w.r.t θ of negative phase is
[Equation]
Next, we will show the speciﬁc derivation of these two terms in Eq.
so that the entire Eq.
15 can be transformed into Eq. 9.
A. T HE DERIVATION OF THE FIRST TERM
The ﬁrst term E x <k ∼ p θ [ ∂∂θ log ˜
q θ ( x <k )] can be processed as follows
[Equation]
where we have w ( x <k )
= (cid:80) [Equation] e
[Equation]
φ
[Equation]
xk [Equation] x <k
[Equation]
k ∼
[Equation]
x
[Equation] [Equation]
xk
[Equation]
<k ) ]
because
[Equation]
A. T HE DERIVATION OF THE SECOND TERM
Then, we tackle the second term E x k , x <k ∼ p θ [ ∂∂θ φ θ ( x k , x <k )] as follows
[Equation]
where w ( x <k ) is also equal to (cid:80)
xk e − φ ( xk, x <k )
Z θ .
Combining Eq.
16 and Eq.
[Equation]
we can obtain an equivalent form of the gradient of the negative phase without any expectation over p θ as
[Equation]
where w ( x <k
[Equation]
B THE FURTHER REFINEMENT OF w
The reweighing weight w can be further deduced as
[Equation]
where µ ( x <k ) is deﬁned as p θ ( x
<k ) ˜
q θ ( x <k ) .
C E XPERIMENTAL S ETTINGS
In this section, we introduce the speciﬁc setup of different benchmarks in Table 6.
We uniformly use
Adam optimizer.
The training will be stopped once the model has not obtained better performance for 20 epochs on the validation set.
For translation tasks, the length of generated fake sentences, which is used for the computing of negative phase in Eq. 12, is dependent on the source sequence whilst for language modeling tasks, we ﬁx the length of generated fake sentences as 50 during training.
As for the model structures of the image generation task, we use the ofﬁcial structure reported by PixelCNN (van den Oord et al., 2016c) and Gated PixelCNN (van den Oord et al., 2016b) without modiﬁcation.
The source code will be released once upon acceptance.
We use the same batch of samples generated autoregressively to approximate both the expectations in Eq.
and weight w ( i.e. , shared), which does not need to sample twice.
The number of samples in a batch is dynamic while the maximum number of the total tokens in a batch are ﬁxed (4096 in our experiments).
If the length of sequences in a batch is 32, then it includes 4096 / 32 =samples in total.
It is a common strategy in language generation tasks, and has been used in many frameworks(e.g. Fairseq (Ott et al., 2019)). We generate samples autoregressively as many as the number of sequences in the current batch at each update iteration. [Table]
D M ORE E XPERIMENTAL A NALYSIS D. E FFECT ON I NCOHERENCE
In order to validate the effectiveness of our E-ARM for ameliorating the long-range coherence of generations, we undertake an experiment to assess the model’s performance under different test sets with varying sentence lengths.
We divided the test set of IWSLT14 (German → English, Italian →
English, Spanish → English) translation dataset into three subsets ([0, 25], [25, 50], and [50, ∞ )) based on the target sentence lengths.
Then, we incrementally applied scheduled sampling technique and our E-ARM above the base transformer network, and tested their performances on these three subsets.
Generally, the subset of samples with longer target sentences ([50, ∞ )) should have been more affected by the long-range incoherence problem (lower BLEU score).
In practice, we uniformly applied label smoothing and beam searching (with 5 beams) strategy for all experiments in
Table 7.
Speciﬁcally, Table 7 shows that the base translation model improved performance for all three test sets with varying target sentence lengths after using the scheduled sampling technique, especially for the two sets [0, 25) and [25, 50) which had relatively short target sentence lengths ( e.g. On
German to English task, 38.20 - 37.72 = +0.48 points and 33.76 - 33.24 = + 0.52 points for [0, 25) and [25, 50) test sets respectively).
We consider that this performance boost was achieved through alleviating the exposure bias problem, since scheduled sampling approaches (Ranzato et al., 2016;
Zhang et al., 2019; Mihaylova & Martins, 2019) have been veriﬁed in mitigating the exposure bias problem.
Besides, after applying our E-ARM together with the scheduled sampling technique, the
[Table]
base model can further obtain additional performance gain.
Speciﬁcally, the improvement on the longer sentence is more evident, since model can obtain large improvements on the [50, ∞ ) ( e.g.
On German to English task, 31.43 - 31.08 = +0.35 points for [50, ∞ ) test sets) than short sets [0, 25] and [25, 50] ( e.g. On German to English task, 38.37 - 38.20 = +0.17 points and 33.92 - 33.
= + 0.16 points for [0, 25) and [25, 50) test sets respectively).
This phenomenon indicates that our
E-ARM can resolve the incoherence problem to some extent.
D. E FFECT ON E XPOSURE B IAS
[Table]
We follow the analytic experiments in the work (Zhang et al., 2019) to show that our E-ARM is capable of alleviating the exposure bias problem.
Speciﬁcally, we randomly select 1K pairs from the training data for each translation pair and use the trained autoregressive model which applied
E-ARM (Label Smoothing with smoothing factor 0.1 is applied during training while scheduled sampling is not used) to decode the source sentences, and then count the ground truth words whose probabilities in the predicted distributions produced by our E-ARM are greater than those produced by the baseline and denote the number as N .
The ratio of N to the total number of words tested is calculated.
The detailed results are shown in Table 8.
We ﬁnd that the results on all different tasks are greater than 50%, which demonstrate the ability of our E-ARM in solving exposure bias problem.
D. A NALYSIS TO MODEL ’ S C ONVERGENCE
In this section, We will investigate the convergence of our E-ARM.
To begin, we ﬁrst train a base
Transformer model (“Tr-Base” architecture shown in Table 6) on the IWSLT14 Spanish to English training set for baseline and E-ARM model respectively, and then record the training loss and test loss (in cross entropy) at the end of each epoch.
The loss curves are plotted in the Figure 2.
From
Figure 2, we can see that (1) at the start of the training, our E-ARM converges slightly faster than the baseline.
As the training process progresses, the cross entropy of the baseline on the training set will gradually decrease, with a faster rate than E-ARM.
On the other hand, the test loss curve of the baseline will fall at initially and then slowly rose after 50 epochs while E-ARM always remains stable convergence.
This phenomenon also shows that our E-ARM model can effectively prevents over-ﬁtting and produce better generalization.
[Figure]
D. A NALYSIS TO T OP -K RE - SAMPLING
[Table]
Top-K energy re-sampling in the inference stage is introduced by Bakhtin et al.
(2021), which collects many candidate sequences generated autoregressively in the inference stage and then resamples from them depending on their energy scores estimated by the network.
To measure the contribution of the Top-K energy re-sampling in our method, we conduct ablation study to verify it by selecting different K = { 0, 5, 10 } .
The results are shown in Table 9 by using BLEU score.
From Table
[Table]
we observe that the beneﬁts brought by Top-K sampling is minor (K= { 5, 10 } ), when compared with model without Top-K sampling (K=0).
These results also indicate that the performance improvements of our E-ARM are mainly from our joint-training, rather than Top-K energy re-sampling
. D. E VALUATION WITH OTHER METRICS
[Table]
To further evaluate the effectiveness of the our proposed E-ARM
[Table]
we also evaluate our method by using other metrics, such as ROUGE Lin (2004) and METEOR Banerjee & Lavie (2005) for neural machine translation.
The results are shown in Table 10.
In Table 10, the improvements of E-ARM in different metrics is consistent with the conclusion of Table 1, which further prove the effectiveness of our E-ARM model.
D. E FFICIENCY STUDY
Our E-ARM has the advantage of being able to optimize an energy-based learning target using maximum log-likelihood, without the usage of MCMC procedures.
The requirement to sample data from the autoregressive model at each update step, on the other hand, remains a possible element that could slow down the training process.
Nonetheless, the extra overheads are still acceptable when compared to sampling data using MCMC algorithms.
The reasons are provided in below:
Assuming that the forward processes of the Transformer, with a length n sentence as the input, have the time cost τ .
We tested the time cost of gradients back-propagation for Transformer on Tesla
V100 GPU.
We found that the time cost of the backward process is approximately twice as the forward process, which is marked as 2 τ .
Therefore, the time cost of one step update is approximate 3 τ .
Autoregressively generating a sequence of length n by Transformer necessitates n feedforward processes, as each predicted token must use all previously created tokens as input.
One fact is that we simply need to use the previously produced k tokens as the input at each time step k , and gradient back-propagation is not required during the generation.
[Table]
[Table]
As a result, the time cost of generating a fake sentence for energy-based training in Eq.12 is actually 1 n τ +n + · · · + n − 1 n τ = n − 1 2 τ .
Considering the IWSLT
German to English translation task, which has training samples with length 20 in average, the time cost of generating fake data in each iteration is roundly 9.5 τ .
Furthermore, the generated fake sentence will be fed into the transformer and included in the overall loss computation, resulting in an extra forward and backward procedure apart from the update of original input.
Thus, the total time cost of our E-ARM’s one update is 15.5 τ , which is about 5.2 times as great as vanilla training.
Table 11 shows the time cost of training a 12 layer transformer with 100 iterations.
The time cost of our E-ARM roughly coincides the extra time cost as we analyzed above.
For long sequence tasks like image generation and language modeling, which usually have sequences consisting of hundreds of tokens, we randomly truncate a continuous sequence with length 50 for energy-based training in Eq.12.
When it comes to the MCMC sampling, one problematic issue is that for sequential data like text, the intrinsic discrete property prevents it from applying MCMC in the data space, which forces us to apply it in the latent feature space.
Here, we take the SGLD (Welling & Teh, 2011b) algorithm for example.
Assuming that we apply the SGLD at the ﬁrst layer of the network, then the time cost of one SGLD iteration is about 3 τ either.
Since the SGLD process requires k iterations to reach convergence, the total time cost of one update of our E-ARM with MCMC process is (3 k + 6) τ .
In practice, the k is usually set as 100 for stable training Grathwohl et al.
(2020), which results in the time cost being (3 × 100+6) τ 3 τ = 102 times as large as the vanilla training.
For short-run SGLD, which takes k as 20 with a sacriﬁce of performance, it still leads to the time cost being 22 times as large as the vanilla training.
D. C ASES S TUDIES
To better understand the advantages of our method in correcting error tokens
[Table]
we also prepare some translation cases in IWSLT14 German → English, as shown in Table
[Table]
E M ORE D ISCUSSION OF RELATED WORKS
The seminal idea of combing a generative model and an energy-based model has been explored by a plethora of great works (Pang et al., 2020; Durkan & Nash, 2019; Xie et al., 2019; 2020; Xiao et al., 2021; Bakhtin et al., 2021).
Our E-ARM can be considered as a member of this family of models in general, but it has a different mechanism and goal than the others.
In particular, Pang et al.
(2020) aimed to learn an energy-based model (EBM) in the latent space of a generator model, so that the EBM can act as a prior model on the generator model’s top-down network.
They believe that the energy-based correction of the prior noise distribution will beneﬁt the subsequent generator’s generating process.
Furthermore, Xie et al.
(2019) attempted to learn the conditional distribution of a high-dimensional output given an input by combining the efforts of a fast thinking initializer, which generates the output and a latent vector, and a slow thinking solver, which learns an objective function in the form of a conditional energy function, so that the output can be generated by optimizing the objective function, or more rigorously by sampling from the conditional energy-based model.
A similar work is GAMs (Parshakova et al., 2019a;b; Khalifa et al., 2021), which combine an autoregressive component with a log-linear component, allowing the use of global a priori features to compensate for lack of data.
Moreover, VAEBM, a symbiotic composition of a variational auto-encoder and an EBM, was proposed by (Xiao et al., 2021).
It can use a state-of-the-art VAE to capture the general mode structure of the data distribution while relying on its EBM component to explicitly eliminate non-data-like regions from the model and reﬁne the generation samples.
In addition, Bakhtin et al.
(2021) designed a novel mechanism to train an unnormalized energy-based models for modeling joint sequence by working in the residual of a pretrained locally normalized language model and training using noise contrastive estimation.
All of the above models require an additional network to learn the energy scores, which prevents the base autoregressive model from beneﬁting from EBM’s properties in modeling the joint distribution in a more temporally coherent manner.
In contrast, by carefully constructing an energy-based learning objective and its corresponding optimization procedure, we are able to smoothly integrate energy surface learning into autoregressive networks that do not require additional learnable parameters.
Rather than proposing a new generative model, our method is more likely to a novel training pattern for training a better autoregressive model.
Recently, instead of constructing an autoregressive model in the data space,
Xu et al.
(2021b) have proposed a unique way which uses autoregressive models in the latent space followed by a decoder which decodes the autoregressively generated latent feature into the original data space.
They attempt to learn a structured representation space where dimensions are ordered based on importance and trade off the sample quality for computational efﬁciency by truncating the dimensions of latent generations.
Their work is orthogonal to ours.
We think the combination between our E-ARM with anytime sampling is also a valuable work which is worth exploration in the future.