MCMC Should Mix: Learning Energy-Based Model with Flow-Based Backbone

1 INTRODUCTION . 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 . 2 RELATED 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 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 . 3 MODEL AND LEARNING . 3.1 FLOW-BASED MODEL . 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 .

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