A Tale of Two Flows: Cooperative Learning of Langevin Flow and Normalizing Flow Toward Energy-Based Model

1 INTRODUCTION AND MOTIVATION . Normalizing flows ( Dinh et al. , 2015 ; 2017 ; Kingma & Dhariwal , 2018a ) are a family of generative models that construct a complex distribution by transforming a simple probability density , such as Gaussian distribution , through a sequence of invertible and differentiable mappings . Due to the tractability of the exact log-likelihood and the efficiency of the inference and synthesis , normalizing flows have gained popularity in density estimation ( Kingma & Dhariwal , 2018a ; Ho et al. , 2019 ; Yang et al. , 2019 ; Prenger et al. , 2019 ; Kumar et al. , 2020 ) and variational inference ( Rezende & Mohamed , 2015 ; Kingma et al. , 2016 ) . However , for the sake of ensuring the favorable property of closed-form density evaluation , the normalizing flows typically require special designs of the sequence of transformations , which , in general , constrain the expressive power of the models . Energy-based models ( EBMs ) ( LeCun et al. , 2006 ) define an unnormalized probability density function of data , which is the exponential of the negative energy function . The energy function is directly defined on data domain , and assigns each input configuration with a scalar energy value , with lower energy values indicating more likely configurations . Recently , with the energy function parameterized by a modern deep network such as a ConvNet , the ConvNet-EBMs ( Xie et al. , 2016 ) have gained unprecedented success in modeling large-scale data sets ( Gao et al. , 2018 ; Nijkamp et al. , 2019 ; Du & Mordatch , 2019 ; Grathwohl et al. , 2020 ; Gao et al. , 2021a ; Zhao et al. , 2021 ) and exhibited stunning performance in synthesizing different modalities of data ( Xie et al. , 2017 ; 2018b ; 2021a ; Du et al. , 2020 ) . The parameters of the energy function can be trained by maximum likelihood . However , due to the intractable integral in computing the normalizing constant , the eval- uation of the gradient of the log-likelihood typically requires Markov chain Monte Carlo ( MCMC ) sampling ( Barbu & Zhu , 2020 ) , such as gradient-based Langevin dynamics ( Neal et al. , 2011 ) , to generate samples from the current model . However , the Langevin sampling on a highly multi-modal energy function , due to the use of deep network parameterization , is generally not mixing . When sampling from a density with a multimodal landscape , the Langevin dynamics , which follows the gradient information , is apt to get trapped by local modes of the density and is unlikely to jump out to explore other isolated modes . Relying on non-mixing MCMC samples , the estimated gradient of the likelihood can be biased , and the resulting learned energy-based model may become an invalid model , which is unable to approximate the data distribution as expected . Recently , Nijkamp et al . ( 2019 ) propose to train an EBM with a short-run non-convergent Langevin dynamics , and show that even though the energy function is invalid , the short-run MCMC can be treated as a valid flow model that generates realistic examples . This not only provides an explanation of why EBMs with non-convergent MCMC are still capable of synthesizing realistic examples , but also suggests a more practical computationally-affordable way to learn useful generative models under the existing energy-based learning frameworks . Although EBMs have been widely applied to different domains , learning short-run MCMC in the context of EBM is still underexplored . In this paper , we will accept the fact that MCMC sampling is not mixing in practice , and decide to give up the goal of training a valid EBM of data . Instead , we treat the short-run non-convergent Langevin dynamics , which shares parameters with the energy function , as a flow-like transformation that we call the Langevin flow because it can be considered as a noise-injected residual network . Even though implementing a short-run Langevin flow is surprisingly simple , which is nothing but a design of a bottom-up ConvNet for the energy function , it might still require a sufficiently large number of Langevin steps ( each Langevin step consists of one step of gradient descent plus a diffusion step ) to construct an effective Langevin flow , so that it can be expressive enough to represent and synthesize the data . Motivated by reducing the number of Langevin steps in the Langevin flow for computational efficiency , we propose the CoopFlow model that trains a Langevin flow jointly with a normalizing flow in a cooperative learning scheme , in which the normalizing flow model learns to serve as a rapid sampler to initialize the Langevin flow so that the Langevin flow can be shorter , while the Langevin flow teaches the normalizing flow by short-run MCMC transition toward the EBM so that the normalizing flow can accumulate the temporal difference in the transition to provide better initial samples . Besides , the motivation of using the normalizing flow instead of a generic generator , is because the cooperative learning framework need to incorporate the MLE algorithm of an EBM and the MLE algorithm of a generator . The MLE of a normalizing flow generator will be much more tractable than the MLE of any other generic generator . The latter might resort to either MCMC-based inference to evaluate the posterior distribution or another encoder network for variational inference . In the proposed CoopFlow , the Langevin flow can overcome expressivity limitation of the normalizing flow due to the invertibility constraint . Such a benefit also motivates us to study the CoopFlow framework . To be specific , we understand the dynamics of cooperative learning with short-run non-mixing MCMC by information geometry . We provide justifications that the CoopFlow trained in the context of EBM with non-mixing MCMC is a valid generator , because it converges to a moment matching estimator . Experiments , including image generation , image reconstruction , and latent space interpolation , are conducted to support our justification . 2 CONTRIBUTIONS AND RELATED WORK . Our paper studies amortized sampling for training a short-run non-mixing Langevin sampler toward an EBM . We propose a novel framework , the CoopFlow , which cooperatively trains a short-run Langevin flow as a valid generator and a normalizing flow as an amortized sampler for image representation . We provide both theoretical and empirical justifications for the proposed CoopFlow algorithm . The following are some closely related work . We will point out the differences between our work and the prior arts to further highlight the contributions and novelty of our paper . Learning short-run MCMC as a generator Recently , Nijkamp et al . ( 2019 ) propose to learn an EBM with short-run non-convergent MCMC that samples from the model , and they eventually keep the short-run MCMC as a valid generator and discard the biased EBM . Nijkamp et al . ( 2020b ) use short-run MCMC to sample the latent space of a top-down generative model in a variational learning framework . An et al . ( 2021 ) proposes to correct the bias of short-run MCMC inference by optimal transport in training latent variable model . Pang et al . ( 2020 ) adopts short-run MCMC to sample from both the EBM prior and the posterior of the latent variables . Our work studies learning a normalizing flow to amortize the sampling cost of a short-run non-mixing MCMC sampler ( i.e. , Langevin flow ) in data space , which makes a further step forward in this underexplored theme . Cooperative learning with MCMC teaching Our learning algorithm is inspired by the cooperative learning or the CoopNets ( Xie et al. , 2020 ) , in which an EBM and a top-down generator ( Han et al. , 2017 ) are jointly trained by jump-starting their maximum learning algorithms . Xie et al . ( 2021b ) replaces the generator in the original CoopNets by a VAE ( Kingma & Welling , 2013 ) for efficient inference . The proposed CoopFlow algorithm is different from the above prior work in the following two aspects . ( i ) In the idealized long-run mixing MCMC scenario , our algorithm is a cooperative learning framework that trains an unbiased EBM and a normalizing flow via MCMC teaching , where the updating of the normalizing flow with a tractable density is more efficient and less biased than updating a generic generator via variational inference as in Xie et al . ( 2021b ) or MCMC-based inference as in Xie et al . ( 2020 ) . ( ii ) Our paper has a novel emphasis on studying cooperative learning with short-run non-mixing MCMC , which is more practical and common in reality . Our algorithm actually trains a short-run Langevin flow and a normalizing flow together toward a biased EBM for image generation . We use information geometry to understand the learning dynamics and show that the learned two-flow generator ( i.e. , the CoopFlow ) is a valid generative model , even though the learned EBM is biased . Joint training of EBM and normalizing flow There are other two works studying EBM and normalizing flow together . To avoid MCMC , Gao et al . ( 2020 ) propose to train an EBM using noise contrastive estimation ( Gutmann & Hyvärinen , 2010 ) , where the noise distribution is a normalization flow . Nijkamp et al . ( 2020a ) propose to learn an EBM as an exponential tilting of a pretrained normalizing flow , so that neural transport MCMC sampling in the latent space of the normalizing flow can mix well . Our paper proposes to train an EBM and a normalizing flow via short-run MCMC teaching . More specially , we focus on short-run non-mixing MCMC and treat it as a valid flow-like model ( i.e. , short-run Langevin flow ) that is guided by the EBM . Disregarding the biased EBM , the resulting valid generator is the combination of the short-run Langevin flow and the normalizing flow , where the latter serves as a rapid initializer of the former . The form of this two-flow generator shares a similar fashion with the Stochastic Normalizing Flow ( Wu et al. , 2020 ) that consists of a sequence of deterministic invertible transformations and stochastic sampling blocks . 3 TWO FLOW MODELS . 3.1 LANGEVIN FLOW . Energy-based model Let x ∈ RD be the observed signal such as an image . An energy-based model defines an unnormalized probability distribution of x as follows : pθ ( x ) = 1 Z ( θ ) exp [ fθ ( x ) ] , ( 1 ) where fθ : RD → R is the negative energy function and defined by a bottom-up neural network whose parameters are denoted by θ . The normalizing constant or partition function Z ( θ ) =∫ exp [ fθ ( x ) ] dx is analytically intractable and difficult to compute due to high dimensionality of x . Maximum likelihood learning Suppose we observe unlabeled training examples { xi , i = 1 , ... , n } from unknown data distribution pdata ( x ) , the energy-based model in Eq . ( 1 ) can be trained from { xi } by Markov chain Monte Carlo ( MCMC ) -based maximum likelihood estimation , in which MCMC samples are drawn from the model pθ ( x ) to approximate the gradient of the log-likelihood function for updating the model parameters θ . Specifically , the log-likelihood is L ( θ ) = 1n ∑n i=1 log pθ ( xi ) . For a large n , maximizing L ( θ ) is equivalent to minimizing the Kullback-Leibler ( KL ) divergence DKL ( pdata||pθ ) . The learning gradient is given by L′ ( θ ) = Epdata [ ∇θfθ ( x ) ] − Epθ [ ∇θfθ ( x ) ] ≈ 1 n n∑ i=1 ∇θfθ ( xi ) − 1 n n∑ i=1 ∇θfθ ( x̃i ) , ( 2 ) where the expectations are approximated by averaging over the observed examples { xi } and the synthesized examples { x̃i } generated from the current model pθ ( x ) , respectively . Langevin dynamics as MCMC Generating synthesized examples from pθ ( x ) can be accomplished with a gradient-based MCMC such as Langevin dynamics , which is applied as follows xt+1 = xt + δ2 2 ∇xfθ ( xt ) + δϵt ; x0 ∼ p0 ( x ) , ϵt ∼ N ( 0 , ID ) , t = 1 , · · · , T , ( 3 ) where t indexes the Langevin time step , δ denotes the Langevin step size , ϵt is a Brownian motion that explores different modes , and p0 ( x ) is a uniform distribution that initializes the MCMC chains . Langevin flow As T → ∞ and δ → 0 , xT becomes an exact sample from pθ ( x ) under some regularity conditions . However , it is impractical to run infinite steps with infinitesimal step size to generate fair examples from the target distribution . Additionally , convergence of MCMC chains in many cases is hopeless because pθ ( x ) can be very complex and highly multi-modal , then the gradient-based Langevin dynamics has no way to escape from the local modes , so that different MCMC chains with different starting points are unable to mix . Let p̃θ ( x ) be the distribution of xT , which is the resulting distribution of x after T steps of Langevin update starting from x0 ∼ p0 ( x ) . Due to fixed initial distribution p0 ( x ) and fixed K and δ , the distribution p̃θ ( x ) is well defined , which can be implicitly expressed by p̃θ ( x ) = ( Kθp0 ) ( x ) = ∫ p0 ( z ) Kθ ( x|z ) dz , ( 4 ) where Kθ denotes the transition kernel of T steps of Langevin dynamics that samples pθ . Generally , p̃θ ( x ) is not necessarily equal to pθ ( x ) . p̃θ ( x ) is dependent on T and s , which are omitted in the notation for simplicity . The KL-divergence DKL ( p̃θ ( x ) ||pθ ( x ) ) = −entropy ( p̃θ ( x ) ) − Ep̃θ ( x ) [ f ( x ) ] + logZ . The gradient ascent part of Langevin dynamics , xt+1 = xt + δ2/2∇xfθ ( xt ) , increases the negative energy fθ ( x ) by moving x to the local modes of fθ ( x ) , and the diffusion part δϵt increases the entropy of p̃θ ( x ) by jumping out of the local modes . According to the law of thermodynamics ( Cover & Thomas , 2006 ) , DKL ( p̃θ ( x ) ||pθ ( x ) ) decreases to zero monotonically as T →∞ .

Summary: The paper proposes a generative model that consists of two components: a normalizing flow and an energy-based model with short-run MCMC as proposed in (Nijkamp, 2019). The whole procedure operates by pushing some simple distribution through the NF and then running the short-run chain on these samples. Both models are learned jointly as follows. The normalizing flow tries to capture the limiting distribution of the EBM by minimizing  $$\text{KL}(K_\theta q_\alpha \Vert q_\alpha),$$  where $q_\alpha$ is the density of the NF with parameters $\alpha$ and $K_\theta$ is the transition kernel of the short-run chain with parameters $\theta$. Simultaneously, the EBM is learned to capture the data by minimizing  $$\text{KL}(p_{\text{data}}\Vert p_\theta) - \text{KL}(K_\theta q_\alpha\Vert p_\theta).$$  In Section 4.2, the authors provide some analysis of the proposed procedure. Namely, the authors consider the perfect scenario (where the NF matches the EBM density and they both match data) and give some intuition for the case when the models are not expressive enough to perfectly match the data.  Finally, the authors provide an empirical study of the proposed procedure. They report metrics for image generation on CIFAR10, SVHN, Celeba (downsampled to 32x32); give several examples of image inpainting and latent space interpolation.