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: This paper develops a new generative model, called CoopFlow, by combining the energy-based model (EBM), normalizing flow (NF) and Langevin dynamics (LD). The CoopFlow updates all the parts in an cooperative manner: Parameters of NF are updated by using samples modified by LD, LD is performed based on the potential function provided by the EBM, and the gradient of log-likelihood of EBM are estimated according to samples generated by NF + LD. Here, LD can overcome the limitation of the model capacity of NF, and NF + LD can provide a more accurate estimation of the gradient of EBM loss. After learning, a powerful sampler and an accurate EBM can be obtained.