Year: 2021

Abstract:
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 flow-based model serving as a backbone, so that the EBM is a correction or an exponential tilting of the flowbased model. We show that the model has a particularly simple form in the space of the latent variables of the flow-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.

heading: 1 INTRODUCTION
text: 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; Atchadé et al., 2017; De Bortoli et al., 2021; Song & Ou, 2018) defines 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 efficiently 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 & Hyvärinen, 2010; Gao et al., 2019; Wang & Ou, 2018) or by amortized sampling (Kim & Bengio, 2016; Song & Ou, 2018; 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 finite 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 fixed 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 goal of this paper is to address the problem of MCMC mixing, which is important for proper learning of EBMs. The sub-par quality of synthesis of our approach is a concern, which we believe may be addressed with recent flow architectures (Durkan et al., 2019) and jointly updating the flow model in future work. We believe that fitting 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 flow-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 flow-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), flow-based model is based on a mapping from latent space to the data space. However, unlike the generator model, the mapping in the flow-based model is deterministic and one-one, with closed-form inversion and Jacobian that can be efficiently computed. This leads to an explicit normalized density via change of variable. However, to ensure tractable inversion and Jacobian, the mapping in the flow-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 flow-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 flow-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 efficacy of learning EBM with flow-based backbone, and the neural transport sampling of the learned EBM greatly mitigates the non-mixing problem of MCMC.

heading: 2 RELATED WORK AND CONTRIBUTIONS
text: 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 flow-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 flow-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 flow-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 simplified special case of this idea, where we learn the EBM as a correction of a pre-trained flow-based model, so that we do not need to train a separate flow-based approximation to the EBM. The energy function, which is a correction of the flow-based model, does not need to reproduce the content of the flow-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 flow-based model disappear. This may allow for free-form flow-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 & Hyvärinen, 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. (2001); Wang & Ou (2018). Recently Xiao et al. (2020) recruits an EBM to correct a variational autoencoder with MCMC-based learning methods. Correcting or refining 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.
Amortized 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. (2013); Kim & Bengio (2016); Dai et al. (2017); Song & Ou (2018); Brock et al. (2018); Xie et al. (2018); Han et al. (2019); Kumar et al. (2019b); Grathwohl et al. (2020). In this paper, we choose to use flow-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 comparison, 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 flow-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.

heading: 3 MODEL AND LEARNING

heading: 3.1 FLOW-BASED MODEL
text: Let x be the input example, such as an image. A flow-based model is of the form z ∼ q0(z), x = gα(z), (1)
where z is the latent vector of the same dimensionality as x, and q0 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, q0(z)dz = qα(x)dx, (2) where dz and dx are understood as the volumes of the infinitesimal local neighborhoods around z and x respectively under the mapping x = gα(z). Then for a given x, z = g−1α (x), and
qα(x) = q0(z)dz/dx = q0(g −1 α (x))|det(∂g−1α (x)/∂x)|, (3)
where the ratio between the volumes dz/dx is the absolute value of the determinant of the Jacobian.
Suppose we observe training examples (xi, i = 1, ..., n) ∼ pdata(x), where pdata 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 DKL(pdata‖qα). qα strives to cover most of the modes in pdata, and the learned qα tends to be more dispersed than pdata. In order for qα to approximate pdata 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 pdata.

heading: 3.2 ENERGY-BASED MODEL
text: An energy-based model (EBM) is defined as follows:
pθ(x) = 1
Z(θ) exp(fθ(x))q(x), (4)
where q(x) is a reference measure, such as the uniform measure or a Gaussian white noise distribution as in Xie et al. (2016). fθ is defined by a bottom-up ConvNet whose parameters are denoted by θ. The normalizing constant or the partition function Z(θ) = ∫ exp(fθ(x))q(x)dx = Eq[exp(fθ(x))] is typically analytically intractable.
Suppose we observe training examples xi ∼ pdata for i = 1, ..., n. For large n, the sample average over {xi} approximates the expectation with respect to pdata. For notational convenience, we treat the sample average and the expectation as the same.
The log-likelihood is
L(θ) = 1
n n∑ i=1 log pθ(xi) . = Epdata [log pθ(x)]. (5)
The derivative of the log-likelihood is
L′(θ) = Epdata [∇θfθ(x)]− Epθ [∇θfθ(x)] . =
1
n n∑ i=1 ∇θfθ(xi)− 1 n n∑ i=1 ∇θfθ(x−i ), (6)
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+1 = θt + ηtL ′(θ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 ∇xfθ(x) can be automatically computed. Since pdata 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.

heading: 3.3 ENERGY-BASED MODEL WITH FLOW-BASED BACKBONE
text: 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 flow-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
pθ(x) = 1
Z(θ) exp(fθ(x))qα(x), (7)
which is almost the same as in (4) except that the reference distribution q(x) is a pre-trained flowbased model qα(x). The resulting model pθ(x) is a correction or refinement 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
p(z)dz = pθ(x)dx = 1
Z(θ) exp(fθ(x))qα(x)dx. (8)
Recall equation (2), qα(x)dx = q0(z)dz, we have
p(z) = 1
Z(θ) exp(fθ(gα(z)))q0(z). (9)
p(z) is an exponential tilting of the prior noise distribution q0(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, flow-based model is simpler and more explicit.

heading: 3.4 LEARNING BY HAMILTONIAN NEURAL TRANSPORT SAMPLING
text: Instead of sampling pθ(x), we can sample p(z) in (9). While qα(x) is multi-modal, q0(z) is unimodal. Since fθ(x) is a correction of qα, p(z) is a correction of p0(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. (2019) for sampling from an EBM or the posterior distribution of a generative model. The basic idea is to train a flow-based model as a variational approximation to the target EBM, and sample the EBM in the latent space of the flow-based model. In our case, since pθ is a correction of qα, we can simply use qα directly as the approximate flow-based model in neural transport sampler. The extra benefit 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 flow-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.
Algorithm 1: Learning the correction fθ of flow qα with Neural Transport (NT-EBM). input : Learning iterations T , learning rate η, batch size m, pre-trained parameters α, initial
parameters θ0, initial latent variables {zi,0}mi=1 ∼ q0(z), observed examples {xi}ni=1, number of MCMC steps K in each learning iteration.
output: Parameters {θT }. for t = 0 : T − 1 do
1. Update {zi,t}mi=1 by HMC with target distribution p(z) in equation (9) for K steps. 2. Push the z-space samples forward through gα to obtain synthesized examples {x−i }mi=1. 3. Draw observed training examples {xi}mi=1. 4. Update θ according to (6).

heading: 3.5 LEARNING BY NOISE CONTRASTIVE ESTIMATION
text: We may also learn the correction fθ(x) discriminatively, as in noise contrastive estimation (NCE) (Gutmann & Hyvärinen, 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
log [ P (+|x) P (−|x) ] = log [ ρ 1− ρ ] − logZ(θ) + fθ(x) = b+ fθ(x),
where b = log [
ρ 1−ρ
] − logZ(θ) is treated as a separate bias parameter. Then we can estimate b
and θ by fitting 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 flow-based model based on the value function of GAN. This may further improve the NCE results.

heading: 4 EXPERIMENTS
text: 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? (2) In the latent space, does smooth interpolation remain feasible? (3) Does the exponential tilting with correction term fθ(x) improve the quality of synthesis? (4) In terms of ablation, what is the effect of amount of parameters α for flow-based qα? (5) Is discriminative learning in the form of NCE an efficient 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 flow-based model qα. We refer to Appendix A.3 and A.4 for details on the model architectures and training settings, respectively.

heading: 4.1 MIXING
text: In the following, we will recruit several diagnostics to quantitatively and qualitatively address the question of mixing MCMC. We will first 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.
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 definition, they should appear “similar” to one another, else, one or more chains have failed to converge. Specifically, 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 fixed number of sampling steps (Table 1).

heading: 4.2 INTERPOLATION
text: Interpolation allows us to appraise the smoothness of the latent space. In particular, two samples z1 and z2 are drawn from the prior distribution q0. 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 q0(z) to Uγ(z) = U(z)− γ‖z − z∗‖2 with magnetization constant γ (Hill et al., 2019). Note, ddz‖z‖2 = z/‖z‖2, thus, the magnetization term introduces a vector field pointing with uniform strength γ towards z∗. The resulting Langevin equation is dz(t) = ( ∆U(z(t)) + γ z(t)−z ∗
‖z(t)−z∗‖2
) dt+ √
2dW (t) with Wiener process W (t). To find a low energy path from z1 towards z2, we set z∗ = z2, z = z1 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.

heading: 4.3 SYNTHESIS
text: 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.).
(a) Samples from flow qα by ancestral sampling. (b) Samples from pθ by neural transport. (c) Samples from flow qα by ancestral sampling. (d) Samples from pθ by neural transport.
Figure 6: Comparison of generated samples by ancestral sampling from flow qα and neural transport sampling via Hamiltonian Monte Carlo from pθ learned by NT-EBM. Left: SVHN (32 × 32 × 3). Right: CelebA (64× 64× 3).
The qualitative results are depicted in Figure 6 which contrast generated samples from Glow qα against Markov chains by Hamiltonian neural transport. Table 1 compares the Fréchet Inception Distance (FID) (Heusel et al., 2017) with Inception v3 (Szegedy et al., 2016) on 50, 000 generated examples. Both, qualitatively and quantitatively speaking, we observe a significant improvement in quality of synthesis with exponentially tilting of the reference distribution qα by the correction fθ. However, the overall quality of synthesis is relatively low in comparison with baselines (Miyato et al., 2018; Song & Ou, 2018) which do not involve inference of latent variables. We hope advances in flow architectures and jointly learning the flow model may address these issues in future work.
Note, the models in Table 1 and 3 have been trained based on the official 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.
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 first FID score is calculated over the first 1,000 consecutive samples, the second FID score over the first 2,000 consecutive samples, and so forth. The FID score converges to our reported FID score with multiple sampling chains and a fixed 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.

heading: 4.4 ABLATION
text: We investigate the influence of the number of parameters α of flow-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 significantly outperforms the “largest” Glow baseline.

heading: 4.5 NOISE CONTRASTIVE ESTIMATION
text: Noise Contrastive Estimation (NCE) is a computationally efficient 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.
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 11 (see Appendix A.7). The Markov chain synthesizes realistic images with high diversity. This finding indicates the efficacy of learning EBM with flow-based backbone by NCE, while, after learning the model, we may draw samples by HMC with neural transport.

heading: 5 CONCLUSION
text: This paper proposes to learn EBM as a correction or exponential tilting of a flow-based model, so that neural transport MCMC sampling in the latent space of the flow-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 significant benefit from mixing MCMC and properly learned models by our approach.
Moreover, we pave the way towards computationally efficient 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.

heading: A APPENDIX
text: A.1 CHANGE 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 Dz be an infinitesimal neighborhood around z, and let Dx be an infinitesimal neighborhood around x, so that g maps z to x, and maps Dz to Dx. Then
Pr(Dz) = Pr(Dx). (10) Pr(Dz) = p(z)|Dz| + o(|Dz|), and Pr(Dx) = p(x)|Dx| + o(|Dx|), where |Dz| and |Dx| are the volumes of Dz and Dx respectively. Thus we have
p(z)|Dz| = p(x)|Dx|, (11) where we ignore o(|Dz|) and o(|Dx|) terms. This is the meaning of
p(z)dz = p(x)dx, (12) where |Dx|/|Dz| 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.2 ENERGY-BASED CORRECTION AND CHANGE OF VARIABLE FOR GENERATOR MODEL
The generator model is of the form z ∼ N(0, Id), and x = gα(z) + , ∼ N(0, σ2ID), where D is the dimensionality of x, and d D is the dimensionality of the latent vector. Unlike the flow-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, ), and let x̃ = (z, x). Then
x̃ = (z, x) = Gα(z̃) = Gα(z, ) = (z, gα(z) + ). (13) Let q0(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
q0(z̃)dz̃ = qα(x̃)dx̃. (14) The marginal distribution qα(x) = ∫ qα(x̃)dz = ∫ qα(z, x)dz, which is intractable.
Suppose we exponentially tilt qα(x̃) to
pθ(x̃) = 1
Z(θ) exp(fθ(x̃))qα(x̃). (15)
Again this can be translated into the space of z̃ so that under pθ(x̃), p(z̃)dz̃ = pθ(x̃)dx̃. (16)
Combining equations (14), (15), and (16), we have
p(z̃) = 1
Z(θ) exp(fθ(Gα(z̃))q0(z̃), (17)
that is, under the tilted model pθ(x̃),
p(z, ) = 1
Z(θ) exp(fθ(z, gα(z) + ))q0(z, ). (18)
For p(z, ), the marginal distribution p(z) cannot be obtained in closed form, in particular, for gradient-based sampling, we need to compute
∇z log p(z) = 1
p(z)
∫ ∇zp(z, )d (19)
= ∫ [∇z log p(z, )] p(z, )
p(z) d (20)
= Ep( |z)[∇z log p(z, )]. (21) That is, there is an inner loop for approximating Ep( |z). This is less convenient than the flow-based model.
A.3 MODEL 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.1 (Chen et al.) nonlinearity where Swish(x) = x ∗ sigmoid(x) (Ramachandran et al., 2017) as activation function . We set nf ∈ {32, 64}. Specifically, we set use the following hyper-parameters:
A.4 TRAINING
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 flow model qα are pre-trained following the configuration 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 specification of the Glow model configuration.
EBM. The network parameters are initialized with Xavier (Glorot & Bengio, 2010) and optimized using Adam (Kingma & Ba, 2015) with (β1, β2) = (0.99, 0.999). 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 Fréchet Inception Distance (FID) (Heusel et al., 2017) with Inception v3 classifier (Szegedy et al., 2016) was computed on 50, 000 generated examples with 50, 000 observed examples as reference.
A.5 SYNTHESIS
Figure 9 depicts samples from pre-trained flow qα and samples from pθ learned by neural transport MCMC for the dataset CIFAR-10 (32× 32× 3).
A.6 GELMAN-RUBIN STATISTIC
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 definition, they should appear “similar” to one another, else, one or more chains have failed to converge. Specifically, 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 = {Xij , i = 1, . . . ,m, j = 1, . . . , n} denote m Markov chains of length n. Let s2w = 1 m ∑m i=1 [ 1 n−1 ∑n j=1(Xij − X̄i·)2 ] be the within-chain variance. The quantity s2w underestimates
σ2 due to positive correlation in the Markov chain. Let σ̂2 = n−1n s 2 + s2b n be a mixture of withinchain variance s2w and between-chain variance s 2 b = n m−1 ∑m j=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, s2w 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̂ = √ σ̂2
s2 . 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.7 VERY LONG MARKOV CHAIN
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.

Review summaries:
I find that this paper is a re-submission from NeurIPS2020, for which I acted as one of the reviewers. The content almost remains the same.  This paper studies the learning of a special class of EBMs, which is a correction or an exponential tilting of a flow-based model. An interesting observation that the resulting EBM in the latent space is of a simple form that is much more friendly to MCMC mixing. It is said that HMC sampling of the EBM in the latent space, which is a simple special case of neural transport HMC, mixes well and traverses modes in the data space. Regarding this main claim, the authors raise a number of scientific questions, which are validated by empirical evaluations.

This paper proposes a product of EBM and FLOW and samples in latent space to improve mixing. Such a model can be learned using both MLE or NCE. Empirically, the paper shows that a single MCMC chain of such a model can traverse through different modes. 

Inspired by NueTra [1], in this paper the authors propose a new approach to train deep energy-based models (EBMs). The motivation here is that EBMs are typically trained via contrastive divergence which requires MCMC sampling in a high-dimensional space and from  a mutli-modal distribution. As a result, most MCMC sampling methods do not mix well which results in a biased estimation of MLE. To address this problem, the authors propose to train an energy-based model with a (pre-trained) backbone flow; the EBM can be interpreted as a correction step to the flow model. The flow $q_\alpha(z)$ defines a latent space $z$ with a base distribution $q_0(z)$. As pointed out in [1], we can use the flow network to additionally yield the following distribution on $z$-space: $p(z) = p(x)\frac{\delta x}{\delta z}$. Similar to NeuTra [1], the authors propose to run HMC to sample from $p(z)$ which geometrically is a simpler target, and then feed $z$ to the flow to get $x$ again. The authors perform some experiments on both toy and some of the common vision datasets to show that unlike SGLD, this approach mixes way better and is able to traverse to different local models.   

The authors propose to learn energy-based models with a flow-based model as "backbone". This is to utilize Neural Transport and ultimately make the Markov Chains mix well in data space. This is an important problem.