MCMC S HOULD M IX : L EARNING E NERGY -B ASED M ODEL WITH F LOW -B ASED B ACKBONE
Anonymous authors
Paper under double-blind review
A BSTRACT
Learning energy-based model (EBM) requires MCMC sampling of the learned model as an inner loop of the learning algorithm.
However, MCMC sampling of
EBMs in high-dimensional data space is generally not mixing, because the energy function, which is usually parametrized by deep network, is highly multi-modal in the data space.
This is a serious handicap for both theory and practice of EBMs.
In this paper we propose to learn EBM with a ﬂow-based model serving as a backbone, so that the EBM is a correction or an exponential tilting of the ﬂowbased model.
We show that the model has a particularly simple form in the space of the latent variables of the ﬂow-based model and MCMC sampling of the EBM in the latent space mixes well and traverses modes in the data space.
This enables proper sampling and learning of EBMs.
I NTRODUCTION
The energy-based model (EBM) (LeCun et al., 2006; Ngiam et al., 2011; Kim & Bengio, 2016;
Zhao et al., 2016; Xie et al., 2016; Gao et al., 2018; Kumar et al., 2019b; Nijkamp et al., 2019; Du &
Mordatch, 2019; Finn et al., 2016; Atchad ´ e et al., 2017; De Bortoli et al., 2021; Song & Ou, 2018) deﬁnes an unnormalized probability density function on the observed data such as images via an energy function, so that the density is proportional to the exponential of the negative energy.
Taking advantage of the approximation capacity of modern deep networks such as convolutional networks (ConvNet) (LeCun et al., 1998; Krizhevsky et al., 2012), recent papers (Xie et al., 2016; Gao et al., 2018; Kumar et al., 2019b; Nijkamp et al., 2019; Du & Mordatch, 2019) parametrize the energy function by ConvNet.
The ConvNet-EBM is highly expressive and the learned EBM can produce realistic synthesized examples.
The EBM can be learned by maximum likelihood estimation (MLE), which follows an “analysis by synthesis” scheme.
In the synthesis step, synthesized examples are generated by sampling from the current model.
In the analysis step, the model parameters are updated based on the statistical difference between the synthesized examples and the observed examples.
The synthesis step usually requires Markov chain Monte Carlo (MCMC) sampling, and gradient-based sampling such as
Langevin dynamics (Langevin, 1908) or Hamiltonian Monte Carlo (HMC) (Neal, 2011) can be conveniently implemented on the current deep learning platforms where gradients can be efﬁciently and automatically computed by back-propagation.
However, gradient-based MCMC sampling in the data space generally does not mix, which is a fundamental issue from statistical perspective.
The data distribution is typically highly multi-modal.
To approximate such a distribution, the density function or the energy function of the ConvNetEBM needs to be highly multi-modal as well.
When sampling from such a multi-modal density in the data space, gradient-based MCMC tends to get trapped in local modes with little chance to traverse the modes freely, rendering the MCMC non-mixing.
Without being able to generate fair examples from the model, the estimated gradient of the maximum likelihood learning can be highly biased, and the learned model parameters can be far from the unbiased estimator given by MLE.
Even if we can learn the model by other means without resorting to MCMC sampling, e.g., by noise contrastive estimation (NCE) (Gutmann & Hyv¨arinen, 2010; Gao et al., 2019; Wang & Ou, 2018) or by amortized sampling (Grathwohl et al., 2020), it is still necessary to be able to draw fair examples from the learned model for the purpose of model checking or downstream applications based on the learned model.
Accepting the fact that MCMC sampling is not mixing, contrastive divergence (Tieleman, 2008) initializes ﬁnite step MCMC from the observed examples, so that the learned model is admittedly biased from the MLE.
Du et al.
(2020) improves contrastive divergence by initializing MCMC from augmented samples.
Recently, Nijkamp et al.
(2019) proposes to initialize short-run MCMC from a ﬁxed noise distribution, and shows that even though the learned EBM is biased, the short-run
MCMC can be considered a valid model that can generate realistic examples.
This partially explains why EBM learning algorithm can synthesize high quality examples even though the MCMC does not mix.
However, the problem of non-mixing MCMC remains unsolved.
Without proper MCMC sampling, the theory and practice of learning EBMs is on a very shaky ground.
The primary goal of this paper is to address the problem of MCMC mixing, which is crucial for proper learning of
EBMs.
The quality of synthesis is a secondary issue in our work.
We believe that ﬁtting EBMs properly with mixing MCMC is crucial to downstream tasks that go beyond generating high-quality samples, such as out-of-distribution detection and feature learning.
We will investigate our model on those tasks in future work.
We propose to learn an EBM with a ﬂow-based model as a backbone model or base measure, so that the EBM is in the form of a correction, or an exponential tilting, of the ﬂow-based model.
Flowbased models have gained popularity in generative modeling (Dinh et al., 2014; 2016; Kingma &
Dhariwal, 2018; Grathwohl et al., 2018; Behrmann et al., 2018; Kumar et al., 2019a; Tran et al., 2019) and variational inference (Kingma & Welling, 2013; Rezende & Mohamed, 2015; Kingma et al., 2016; Kingma & Welling, 2014; Khemakhem et al., 2019).
Similar to the generator model (Kingma & Welling, 2013; Goodfellow et al., 2014), ﬂow-based model is based on a mapping from latent space to the data space.
However, unlike the generator model, the mapping in the ﬂow-based model is deterministic and one-one, with closed-form inversion and Jacobian that can be efﬁciently computed.
This leads to an explicit normalized density via change of variable .
However, to ensure tractable inversion and Jacobian, the mapping in the ﬂow-based model has to be a composition of a sequence of simple transformations of highly constrained forms.
In order to approximate a complex distribution, it is necessary to compose a large number of such transformations.
In our work, we propose to learn EBM by correcting a relatively simple ﬂow-based model with a relatively simple energy function parametrized by a free-form ConvNet.
We show that the resulting EBM has a particularly simple form in the space of the latent variables of the ﬂow-based model.
MCMC sampling of the EBM in the latent space, which is a simple special case of neural transport MCMC (Hoffman et al., 2019), mixes well and is able to traverse modes in the data space.
This enables proper sampling and learning of EBMs.
Our experiments demonstrate the efﬁcacy of learning EBM with ﬂow-based backbone, and the neural transport sampling of the learned EBM greatly mitigates the non-mixing problem of MCMC.
R ELATED WORK AND CONTRIBUTIONS
The following are research themes in generative modeling and MCMC sampling that are closely related to our work.
Neural transport MCMC .
Our work is inspired by neural transport sampling (Hoffman et al., 2019).
For an unnormalized target distribution, the neural transport sampler trains a ﬂow-based model as a variational approximation to the target distribution, and then samples the target distri- bution in the space of latent variables of the ﬂow-based model via change of variable.
In the latent space, the target distribution is close to the prior distribution of the latent variables of the ﬂow-based model, which is usually a unimodal Gaussian white noise distribution.
Consequently the target distribution in the latent space is close to be unimodal and is much more conducive to mixing and fast convergence of MCMC than sampling in the original space (Mangoubi & Smith, 2017).
Our work is a simpliﬁed special case of this idea, where we learn the EBM as a correction of a pre-trained ﬂow-based model, so that we do not need to train a separate ﬂow-based approximation to the EBM.
The energy function, which is a correction of the ﬂow-based model, does not need to reproduce the content of the ﬂow-based model, and thus can be kept relatively simple.
Moreover, in the latent space, the resulting EBM takes on a very simple form where the inversion and Jacobian in the ﬂow-based model disappear.
This may allow for free-form ﬂow-based model where inversion and Jacobian do not need to be in closed form (Grathwohl et al., 2018; Behrmann et al., 2018).
Energy-based corrections .
Our model is based on an energy-based correction or an exponential tilting of a more tractable model.
This idea has been explored in noise contrastive estimation (NCE) (Gutmann & Hyv¨arinen, 2010; Gao et al., 2019) and introspective neural networks (INN) (Tu, 2007; Jin et al., 2017; Lazarow et al., 2017), where the correction is obtained by discriminative learning.
Earlier works include Rosenfeld et al.
; Wang & Ou (2018).
Recently Xiao et al.
(2020) recruits an EBM to correct a variational autoencoder with MCMC-based learning methods.
Correcting or reﬁning a simpler and more tractable backbone model can be much easier than learning an EBM from scratch, because the EBM does not need to reproduce the knowledge learned by the backbone model.
It also allows easier sampling of EBM.
Latent space sampling .
Non-mixing MCMC sampling of an EBM is a clear call for latent variables to represent multiple modes of the original model distribution via explicit top-down mapping, so that the distribution of the latent variables is less multi-modal.
Earlier works in this direction include Bengio et al.
; Brock et al.
; Kumar et al.
(2019b).
In this paper, we choose to use ﬂow-based model for its simplicity, because the distribution in the data space can be translated into the distribution in the latent space by a simple change of variable, without requiring integrating out extra dimensions as in the generator model.
Proper learning of EBMs .
Wang & Ou (2017) studies the “proper” learning of EBMs in the modality of languages and recruits Gibbs sampling from the discrete distributions.
In contrast, our work concerns images in continuous space for which we sample by gradient-based MCMC.
Moreover, our work emphasizes the empirical evaluation of the mixing behavior of Markov chains.
Contributions . This paper tackles the problem of non-mixing MCMC for sampling from an EBM. We propose to learn EBM with a ﬂow-based backbone model. The resulting EBM in the latent space is of a simple form that is much more friendly to MCMC mixing.
Our work provides strong empirical evidence regarding the feasibility of mixing MCMC sampling in EBMs parametrized by modern ConvNet for the modality of images.
M ODEL AND LEARNING 3. F LOW - BASED MODEL
Let x be the input example, such as an image.
A ﬂow-based model is of the form
[Equation]
(1) where z is the latent vector of the same dimensionality as x , and q 0 is a known prior distribution such as Gaussian white noise distribution.
g α is a composition of a sequence of invertible transformations whose inversion and log-determinants of the Jacobians can be obtained in closed form.
As a result, these transformations are of highly constrained forms.
α denotes the parameters.
Let q α ( x ) be the probability density at x under the transformation x = g α ( z ) , then according to the change of variable,
[Equation]
where dz and dx are understood as the volumes of the inﬁnitesimal local neighborhoods around z and x respectively under the mapping x
= g α ( z ) .
Then for a given x , z = g − 1 α ( x ) , and
[Equation]
where the ratio between the volumes dz/dx is the absolute value of the determinant of the Jacobian.
Suppose we observe training examples ( x i , i = 1 , ..., n ) ∼ p data ( x ) , where p data is the data distribution, which is typically highly multi-modal.
We can learn α by MLE.
For large n , the MLE of α minimizes the Kullback-Leibler divergence D KL ( p data (cid:107) q α ) .
q α strives to cover most of the modes in p data , and the learned q α tends to be more dispersed than p data .
In order for q α to approximate p data closely, it is usually necessary for g to be a composition of a large number of transformations of highly constrained forms with closed-form inversions and Jacobians.
The learned mapping g α ( z ) transports the unimodal Gaussian white noise distribution to a highly multi-modal distribution q α in the data space as an approximation to the data distribution p data .
3. E NERGY - BASED MODEL
An energy-based model (EBM) is deﬁned as follows:
[Equation]
where q ( x ) is a reference measure, such as the uniform measure or a Gaussian white noise distribution as in Xie et al.
f θ is deﬁned by a bottom-up ConvNet whose parameters are denoted by θ .
The normalizing constant or the partition function Z ( θ ) =
(cid:82) exp( f θ ( x ))
q ( x ) dx = E q [exp( f θ ( x ))] is typically analytically intractable.
Suppose we observe training examples x i ∼ p data for i = 1 , ..., n .
For large n , the sample average over { x i } approximates the expectation with respect to p data .
For notational convenience, we treat the sample average and the expectation as the same.
The log-likelihood is
[Equation]
n (cid:88)
[Equation]
The derivative of the log-likelihood is
[Equation]
where x − i ∼ p θ ( x ) for i = 1 , ..., n are synthesized examples sampled from the current model p θ ( x ) .
The above equation leads to the “analysis by synthesis” learning algorithm.
At iteration t , let θ t be the current model parameters.
We generate x − i ∼ p θ t ( x ) for i = 1 , ..., n .
Then we update θ t +
=
θ t + η t L (cid:48) ( θ t ) , where η t is the learning rate.
To generate synthesized examples from p θ , we can use gradient-based MCMC sampling such as
Langevin dynamics (Langevin, 1908) or Hamiltonian Monte Carlo (HMC) (Neal, 2011), where ∇ x f θ ( x ) can be automatically computed.
Since p data is in general highly multi-modal, the learned p θ or f θ tends to be multi-modal as well.
As a result, gradient-based MCMC tends to get trapped in the local modes of f θ with little chance of mixing between the modes.
3. E NERGY - BASED MODEL WITH FLOW - BASED BACKBONE
Instead of using uniform or Gaussian white noise distribution for the reference distribution q ( x ) in the EBM in (4), we can use a relatively simple ﬂow-based model q α as the reference model.
q α can be pre-trained by MLE, and serves as the backbone of the model, so that the model is of the following form
[Equation]
which is almost the same as in (4) except that the reference distribution q ( x ) is a pre-trained ﬂowbased model q α ( x ) .
The resulting model p θ ( x ) is a correction or reﬁnement of q α , or an exponential tilting of q α ( x ) , and f θ ( x ) is a free-form ConvNet to parametrize the correction.
The overall negative energy is f θ ( x ) + log q α ( x ) .
In the latent space of z , let p ( z ) be the distribution of z under p θ ( x ) , then
[Equation]
Recall equation (2), q α ( x ) dx = q 0 ( z ) dz , we have
[Equation]
p ( z ) is an exponential tilting of the prior noise distribution q 0 ( z ) .
It is a very simple form that does not involve the Jacobian or inversion of g α ( z ) .
We can also apply the above exponential tilting and change of variable scheme to the generator model, i.e., using the generator model as the backbone model.
However, for the generator model, q ( x ) is not in closed form, and after exponential tilting, the marginal p ( z ) requires integral.
See
Appendix A.2 for details.
In comparison, ﬂow-based model is simpler and more explicit.
3. L EARNING BY H AMILTONIAN NEURAL TRANSPORT SAMPLING
Instead of sampling p θ ( x ) , we can sample p ( z ) in (9).
While q α ( x ) is multi-modal, q 0 ( z ) is unimodal.
Since f θ ( x ) is a correction of q α , p ( z ) is a correction of p 0 ( z ) , and can be much less multi-modal than p θ ( x ) in the data space.
After sampling z from p ( z ) , we can generate x = g α
( z ) .
The above MCMC sampling scheme is a special case of neutral transport MCMC proposed by
Hoffman et al.
for sampling from an EBM or the posterior distribution of a generative model.
The basic idea is to train a ﬂow-based model as a variational approximation to the target EBM, and sample the EBM in the latent space of the ﬂow-based model.
In our case, since p θ is a correction of q α , we can simply use q α directly as the approximate ﬂow-based model in neural transport sampler.
The extra beneﬁt is that the distribution p ( z ) is of an even simpler form than p θ ( x ) , because p ( z ) does not involve the inversion and Jacobian of g α .
As a result, we may use a ﬂow-based backbone model of a more free form such as one based on residual network (Behrmann et al., 2018), and we will leave this issue to future investigation.
We use HMC (Neal, 2011) to sample from p ( z ) .
We can then learn θ by MLE according to equation (6).
Algorithm 1 describes the details.
for t = 0 : T − 1 do
Update { z i,t } mi =1 by HMC with target distribution p ( z ) in equation (9) for K steps.
Push the z -space samples forward through g α to obtain synthesized examples { x − i } mi =1 . 3.
Draw observed training examples
Update θ
according to (
3. L EARNING BY NOISE CONTRASTIVE ESTIMATION
We may also learn the correction f
θ ( x ) discriminatively, as in noise contrastive estimation (NCE) (Gutmann & Hyv¨arinen, 2010) or introspective neural networks (INN) (Tu, 2007; Jin et al., 2017; Lazarow et al., 2017).
Let x
+ i , i = 1 , ..., n be the training examples, which are treated as positive examples, and let x − i , i = 1 , ..., n − be the examples generated from q α ( x ) , which are treated as negative examples.
For each batch, let ρ be the proportion of positive examples, and 1 − ρ the proportion of negative examples.
We have
[Equation]
where b = log (cid:104) ρ 1 − ρ (cid:105) − log Z ( θ ) is treated as a separate bias parameter.
Then we can estimate b and θ by ﬁtting a logistic regression on the positive and negative examples.
Note, that NCE is the discriminator side of GAN.
Similar to GAN, we can also improve the ﬂow-based model based on the value function of GAN.
This may further improve the NCE results.
E XPERIMENTS
In the subsequent empirical evaluations, we shall address the following questions:
(1) Is the mixing of HMC with neural transport, both qualitatively and quantitatively, apparent?
In the latent space, does smooth interpolation remain feasible?
(3) Does the exponential tilting with correction term f θ ( x ) improve the quality of synthesis?
In terms of ablation, what is the effect of amount of parameters α for ﬂow-based q α ?
(5) Is discriminative learning in the form of NCE an efﬁcient alternative learning method?
The primary concern of our work is the mixing of MCMC, which is addressed in (1) and (2).
The questions addressed in (3), (4), and (5) are secondary to this concern.
We use Glow (Kingma &
Dhariwal, 2018) as the ﬂow-based model q α .
We refer to Appendix A.3 and A.4 for details on the model architectures and training settings, respectively.
4. M IXING
In the following, we will recruit several diagnostics to quantitatively and qualitatively address the question of mixing MCMC.
We will ﬁrst evaluate the famous Gelman-Rubin statistic for Markov chains running in latent space and contrast those against chains in data space.
Then, we will evaluate auto-correlation of long-run Markov chains as a weaker measure of mixing.
Finally, we provide a visual inspection of long-run Markov chains in our model and compare those with a biased model known to not be amenable to mixing.
[Figure]
Gelman-Rubin.
The Gelman-Rubin statistic (Gelman et al., 1992; Brooks & Gelman, 1998) measures the convergence of Markov chains to the target distribution.
It is based on the notion that if multiple chains have converged, by deﬁnition, they should appear “similar” to one another, else, one or more chains have failed to converge.
Speciﬁcally, the diagnostic recruits an analysis of variance to access the difference between the between-chain and within-chain variances.
We refer to the Appendix A.6 for details.
Figure 2(a-b) depicts the histograms of ˆ R for m = 64 chains over n = 2 , 000 steps with a burn-in time of 400 steps, learned from SVHN dataset.
The mean ˆ R value is 1 . 13 , which we treat as approximative convergence to the target distribution.
We contrast this result with HMC sampling in data space with unfavorable diagnostics of mixing.
Auto-Correlation.
MCMC sampling leads to autocorrelated samples due to the inherent Markovian dependence structure.
The ∆ t (sample) auto-correlation is the correlation between samples ∆ t steps apart in time.
Figure 2(c-d) shows auto-correlation against increasing time lag ∆ t , learned from
SVHN dataset.
While the auto-correlation of HMC chains with neural transport vanishes within ∆ t = 200 steps, the over-damped Langevin sampler requires ∆ t > 1 , 000 steps.
The single longrun Markov chain behavior is consistent with the Gelman-Rubin statistic assessing multiple chains.
Visual Inspection.
Assume a Markov chain is run for a large numbers of steps n = 2 , 000 with a
Hamiltonian neural transport.
Then, we pull back Markov chains into data space and visualize the long run trajectory in Figures 4 and 10 with p θ learned on the SVHN ( 32 × 32 × 3 ) (Netzer et al., 2011) dataset and CelebA (64 × 64 × 3) (Liu et al., 2015) dataset, respectively.
Figure 3 contrasts the
Markov chain that samples the EBM learned with short-run MCMC (Nijkamp et al., 2019), which does not mix, against our method in which the pulled back Markov chains mix freely.
We observe the Markov chain is freely traversing between local modes, which we consider a weak indication of mixing MCMC.
In Appendix A.7 we report the FID scores computed on samples from a single long-run chain, which converge to the FID score computed on samples from multiple chains with a ﬁxed number of sampling steps (Table 1).
4. I NTERPOLATION
Interpolation allows us to appraise the smoothness of the latent space.
In particular, two samples zand z 2 are drawn from the prior distribution q 0 .
We may spherically interpolate between them in z -space and then push forward into data space to assess q α .
To evaluate the tilted model p θ ( z ) , we run a magnetized form of the over-damped Langevin equation for which we alter negative energy
U ( z ) =
f θ ( g α ( z ))
+ log q 0 ( z ) to U γ ( z )
=
U ( z ) − γ (cid:107) z − z ∗ (
cid:107)with magnetization constant γ (Hill et al., 2019).
Note, ddz (cid:107)
z (cid:107) 2 = z/ (cid:107)
z (cid:107) 2 , thus, the magnetization term introduces a vector ﬁeld pointing with uniform strength γ towards z ∗ .
The resulting Langevin equation is dz ( t ) =
(cid:16) ∆ U ( z ( t ))
+ γ z ( t ) − z ∗ (cid:107)
z ( t ) − z ∗ (cid:107) 2 (cid:17) dt + √ 2 dW ( t ) with Wiener process W ( t ) .
To ﬁnd a low energy path from z 1 towards z 2 , we set z ∗ =
z 2 , z
= z 1 and perform n = 1 , 000 steps of the discretized, magnetized Langevin equation with small γ .
Figure 5 depicts the low-energy path in data-space and energy U ( z ) over time.
The qualitatively smooth interpolation and narrow energy spectrum indicate that a Langevin dynamics in latent space (with small magnetization) is able to traverse two arbitrary local modes, thus, substantiating our claim that the underlying geometry is amenable to mixing.
4. S YNTHESIS
While the emphasis of our work is on the mixing of Markov chains, we do evaluate the quality of synthesis on four datasets which include MNIST ( 28 × 28 × 1 ) (LeCun et al., 2010), SVHN ( 32 × 32 × 3 ) (Netzer et al., 2011), CelebA ( 64 × 64 × 3 ) (Liu et al., 2015), and, CIFAR-10 ( 32 × 32 × 3 ) (Krizhevsky et al.).
[Figure]
The qualitative results are depicted in Figure 6 which contrast generated samples from Glow q α against long-run Markov chains by Hamiltonian neural transport.
Table 1 compares the Fr ´ echet Inception Distance (FID) (Heusel et al., 2017) with Inception v3 classiﬁer (Szegedy et al., 2016) on 50 , 000 generated examples.
Both, qualitatively and quantitatively speaking, we observe a signiﬁcant improvement in quality of synthesis with exponentially tilting of the reference distribution q α by the correction f θ .
[Table]
Note, the models in Table 1 and 3 have been trained based on the ofﬁcial reference implementations.
To ensure a fair comparison of learning Glow by MLE and our methods, we disabled learning of the spatial prior and use additive coupling layers for Glow.
[Table]
In Table 2, we show the FID scores for samples obtained from every 1,000 steps of a single longrun chain for a model learned on SVHN.
That is, the ﬁrst FID score is calculated over the ﬁrst 1,000 consecutive samples, the second FID score over the ﬁrst 2,000 consecutive samples, and so forth.
The FID score converges to our reported FID score with multiple sampling chains and a ﬁxed number of sampling steps (Table 1), which indicates that one can obtain fair samples of the model by sampling from a single very long-run HMC chain.
4. A BLATION
We investigate the inﬂuence of the number of parameters α of ﬂow-based q α on the quality of synthesis.
In Table 3, we show (1) the threshold of a “large” q α learned by MLE outperforming
NT with a “small” tilted q α , and, (2) the minimal size of q α which allows for the learning by our method.
Our method with a “medium” sized backbone signiﬁcantly outperforms the “largest” Glow baseline.
[Table]
[Table]
N OISE C ONTRASTIVE E STIMATION
Noise Contrastive Estimation (NCE) is a computationally efﬁcient alternative learning procedure which avoids MCMC sampling by re-casting the learning problem in the form of a logistic regression.
Hence, we wish to learn the correction f
θ with our NCE-EBM algorithm according to (10), while we still sample from the learned model with neural transport MCMC.
Table 1 compares the learned models with both learning methods.
The long-run Markov chains in the energy-based models learned by NCE are conducive to mixing and remain of high visual quality.
Figure 7 depicts samples from q α (left) and samples from p θ learned by our NCE algorithm for which sampling is performed using Hamiltonian neural transport (right) for CelebA.
[Figure]
ancestral sampling
[Figure]
Assume a Markov chain is run for a large numbers of steps with a Hamiltonian neural transport.
Then, we pull back Markov chain into data space and visualize the long run trajectory in Figure(see Appendix A.7).
The Markov chain synthesizes realistic images with high diversity.
This ﬁnding indicates the efﬁcacy of learning EBM with ﬂow-based backbone by NCE, while, after learning the model, we may draw samples by HMC with neural transport.
C ONCLUSION
This paper proposes to learn EBM as a correction or exponential tilting of a ﬂow-based model, so that neural transport MCMC sampling in the latent space of the ﬂow-based model can mix well and traverse the modes in the data space.
From a statistical perspective, the mixing of MCMC is a fundamental problem and is crucial for proper learning of the EBMs.
In future work, we will investigate and identify downstream tasks which signiﬁcant beneﬁt from mixing MCMC and properly learned models by our approach.
Moreover, we pave the way towards computationally efﬁcient learning by noise contrastive estimation while drawing samples from the learned model by neural transport HMC, which we believe is an important direction to explore in future work.
A A PPENDIX A. C HANGE OF VARIABLE
Under the invertible transformation x = g ( z ) , let p ( z ) be the density of z , and p ( x ) be the density of x .
Let D z be an inﬁnitesimal neighborhood around z , and let D x be an inﬁnitesimal neighborhood around x , so that g maps z to x , and maps D z to D x .
Then
[Equation]
Pr( D z ) = p ( z ) | D z | + o ( | D z | ) , and Pr( D x ) = p ( x ) | D x | + o ( | D x | ) , where | D z | and | D x | are the volumes of D z and D x respectively.
Thus we have
[Equation]
11) where we ignore o ( | D z | ) and o ( | D x | ) terms.
This is the meaning of
[Equation] [Equation]
where | D x | / | D z | or dx/dz is the determinant of the Jacobian of g .
Equation (12) is a convenient starting point for deriving densities under change of variable.
A. E NERGY - BASED CORRECTION AND CHANGE OF VARIABLE FOR GENERATOR MODEL
The generator model is of the form z ∼ N(0 , I d ) , and x = g α ( z ) + (cid:15) , (cid:15) ∼ N(0 , σ 2 I D ) , where D is the dimensionality of x , and d (cid:28) D is the dimensionality of the latent vector.
Unlike the ﬂow-based model, the marginal distribution of x involves intractable integral.
We shall study exponential tilting of generator model using the simple equation (12) for change of variable.
To that end, we let ˜ z
= ( z, (cid:15) ) , and let ˜ x
= ( z, x ) .
Then
[Equation] [Equation]
Let q
(˜ z ) be the Gaussian white noise distribution of ˜ z
under the generator model.
Let q α (˜ x ) be the distribution of ˜ x under the generator model.
Consider the change of variable between ˜ z and ˜ x .
In parallel to equation (12), we have
[Equation]
The marginal distribution q α
[Equation]
z
[Equation]
) dz , which is intractable.
Suppose we exponentially tilt q α (˜ x ) to
[Equation] [Equation]
Again this can be translated into the space of ˜ z
so that under p θ (˜ x
[Equation]
[Equation]
Combining equations (14), (15), and (16), we have
[Equation]
17) that is, under the tilted model p θ (˜ x ) ,
[Equation]
For p ( z, (cid:15) ) , the marginal distribution p ( z ) cannot be obtained in closed form, in particular, for gradient-based sampling, we need to compute
[Equation]
[Equation]
That is, there is an inner loop for approximating E p ( (cid:15) | z ) .
This is less convenient than the ﬂow-based model.
A. M ODEL ARCHITECTURES
For Glow model q α , we follow the setting of Kingma & Dhariwal (2018) with n bits x = 8 , flow permutation = 2 , flow coupling = 0 .
For the EBM model f θ , we use the following Conv-Net structure.
We use the following notation.
Convolutional operation conv ( n ) with n output feature maps and bias term.
We recruit LipSwish ( x ) =
Swish ( x ) / 1 .
(Chen et al.) nonlinearity where Swish ( x ) =
x ∗ sigmoid ( x ) (Ramachandran et al., 2017) as activation function .
We set n f ∈ { 32 , 64 } .
Speciﬁcally
[Table]
we set use the following hyper-parameters:
MNIST :
For Glow, n levels = 3 , depth = 8 , width = 128 .
For EBM, n f = 32 . 2.
SVHN : For Glow, n levels = 3 , depth = 8 , width = 128 .
For EBM, n f = 32 . 3.
CelebA : For Glow, n levels = 3 , depth = 16 , width = 256 .
For EBM, n f = 32 . 4.
CIFAR-10 :
For Glow, n levels = 3 , depth = 16 , width = 512 .
For EBM, n f = 32 .
[Table] [Table]
A. T RAINING
Data.
The training image dataset are resized and scaled to [ − 1 , 1] .
We use 60,000, 70,000, 30,000, 50,000 observed examples for MNIST (28 × 28 × 1) , SVHN (32 × 32 × 3) , CelebA (64 × 64 × 3) , and CIFAR-10 (32 × 32 × 3) , respectively.
Glow.
The parameters α of the ﬂow model q α are pre-trained following the conﬁguration and reference implementation provided in (Kingma & Dhariwal, 2018).
Note, for comparability the
Glow models have been trained without spatial prior and with additive coupling layers.
We refer to
Appendix A.3 for a speciﬁcation of the Glow model conﬁguration.
EBM.
The network parameters are initialized with Xavier (Glorot & Bengio, 2010) and optimized using Adam (Kingma & Ba, 2015) with ( β 1 , β 2 ) =
For NT-EBM, the learning rates used are 5e − 5 , 5e − 5 , 1e − 5 , 5e − 5 for MNIST, SVHN, CelebA, CIFAR-10, respectively and a batch-size of 64 examples.
For NCE-EBM, the learning rates used are 1e − 5 , 1e − 5 , 1e − 5 for MNIST,
SVHN, and CelebA, respectively, and a batch-size of 128 examples.
For NT-EBM, in training the maximum number of parameter θ updates was 40 , 000 .
For NCE-EBM, in training the maximum number of parameter θ updates was 80 , 000 .
HMC.
We run Hamiltonian Monte Carlo (HMC) with persistent chains (Tieleman, 2008) initialized from q α and 20 steps of MCMC and 3 leapfrog integrator steps per update of parameters of θ .
The initial discretization step-size 0 .
15 with a simple adaptive policy multiplicatively increases or decreasing the step-size of the inner kernel based on the value of the Metropolis-Hastings acceptance rate (Andrieu & Thoms, 2008).
The target acceptance-rate is set to 0 .
651 (Beskos et al., 2013).
Figure 8 depicts the MH acceptance-rate and adaptive step-size over time.
FID.
The Fr ´ echet Inception Distance (FID) (Heusel et al., 2017) with Inception v3 classiﬁer (Szegedy et al., 2016) was computed on 50 , 000 generated examples with 50 , 000 observed examples as reference.
A. S YNTHESIS
Figure 9 depicts samples from pre-trained ﬂow q α and samples from p θ learned by neural transport
MCMC for the dataset CIFAR-10 ( 32 × 32 × 3 ).
[Figure]
A. G ELMAN -R UBIN S TATISTIC
The Gelman-Rubin statistic (Gelman et al., 1992; Brooks & Gelman, 1998) measures the convergence of Markov chains to the target distribution.
It is based on the notion that if multiple chains have converged, by deﬁnition, they should appear “similar” to one another, else, one or more chains have failed to converge.
Speciﬁcally, the diagnostic recruits an analysis of variance to access the difference between the between-chain and within-chain variances.
Let p denote the target distribution with mean µ ∈ R and variance σ 2 < ∞ .
Gelman et al.
(1992) designs two estimators of σ 2 and compares the square root of their ratio to 1 .
Let
X =
{ X ij , i = 1 , . . .
, m, j = 1 , . .
. , n } denote m
Markov chains of length n .
Let s 2 w = 1 m (cid:80) mi
=1 (cid:104) 1 n − 1 (cid:80) nj =1 ( X ij − ¯
X i · ) 2 (cid:105) be the within-chain variance.
The quantity s 2 w underestimates σ 2 due to positive correlation in the Markov chain.
Let ˆ σ 2 = n − 1 n s
+ s 2 b n be a mixture of withinchain variance s 2 w and between-chain variance s 2 b = nm − 1 (cid:80) mj
=1 ( ¯ X i · − ¯ X ·· ) 2 .
The quantity ˆ σ 2 will overestimate σ 2 , if an over-dispersed initial distribution for the Markov chains was used (Gelman et al., 1992).
That is, s 2 w underestimates while ˆ σ 2 overestimates σ 2 .
Both estimators are consistent for σ 2 as n → ∞ (Vats & Knudson, 2018).
In light of this, the Gelman-Rubin statistic monitors convergence as the ratio ˆ R = (cid:113) ˆ σ 2 s 2 .
Hence, ˆ R measures the degree to which variance (of the means) between chains exceeds what one would expect if the chains were identically distributed.
A. V ERY L ONG M ARKOV C HAIN
In Figure 10, the model p θ is learned with NT-EBM on the CelebA (64 × 64 × 3) (Liu et al., 2015) dataset. In Figure 11, the model p θ is learned with NCE-EBM on the SVHN (3232 × 3) (Netzer et al., 2011) dataset.