A T ALE OF T WO F LOWS : C OOPERATIVE L EARNING OF L ANGEVIN F LOW AND N ORMALIZING F LOW T OWARD E NERGY -B ASED M ODEL
Anonymous authors
Paper under double-blind review
A BSTRACT
This paper studies the cooperative learning of two generative flow models, in which the two models are iteratively updated based on the jointly synthesized examples.
The first flow model is a normalizing flow that transforms an initial simple density into a target density by applying a sequence of invertible transformations.
The second flow model is a Langevin flow that runs finite steps of gradient-based
MCMC toward an energy-based model.
We start from proposing a generative framework that trains an energy-based model with a normalizing flow as an amortized sampler to rapidly initialize the MCMC chains of the energy-based model.
In each learning iteration, we generate synthesized examples by using a normalizing flow initialization followed by a short-run Langevin flow revision toward the current energy-based model.
After that, we treat the synthesized examples as fair samples from the energy-based model and update the model parameters with the maximum likelihood learning gradient, while the normalizing flow directly learns from the synthesized examples by maximizing the tractable likelihood.
We show that, under the short-run non-mixing MCMC scenario, the estimation of the energy-based model actually follows the perturbation of maximum likelihood, and both the short-run Langevin flow and normalizing flow form a two-flow generator that we call CoopFlow.
We provide an understating of the CoopFlow algorithm by information geometry and show that it is a valid generator because it converges to a moment matching estimator.
We demonstrate that the trained CoopFlow is capable of synthesizing realistic images, reconstructing images, and interpolating between images.
I NTRODUCTION AND M OTIVATION
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.
C ONTRIBUTIONS AND R ELATED W ORK
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.
(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 & Hyv ¨ arinen, 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.
T WO F LOW M ODELS 3. L ANGEVIN F LOW
Energy-based model Let x ∈ R D be the observed signal such as an image.
An energy-based model defines an unnormalized probability distribution of x as follows:
[Equation]
where f θ : R D
→ 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 ( θ ) =
(cid:82) 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 { x i , i = 1 , ..., n } from unknown data distribution p data ( x ) , the energy-based model in Eq.(1) can be trained from { x i } 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 ( θ ) = 1 n (cid:80) ni
=1 log p θ ( x i ) .
For a large n , maximizing L ( θ ) is equivalent to minimizing the
Kullback-Leibler (KL) divergence D KL ( p data || p θ ) .
The learning gradient is given by
[Equation]
where the expectations are approximated by averaging over the observed examples { x i } 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
[Equation]
where t indexes the Langevin time step, δ denotes the Langevin step size, ϵ t is a Brownian motion that explores different modes, and p 0 ( x ) is a uniform distribution that initializes the MCMC chains.
Langevin flow As T → ∞ and δ → 0 , x T 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 x T , which is the resulting distribution of x after T steps of Langevin update starting from x 0 ∼ p 0 ( x ) .
Due to fixed initial distribution p 0 ( x ) and fixed K and δ , the distribution ˜ p
θ ( x ) is well defined, which can be implicitly expressed by
[Equation]
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 D KL (˜ p θ ( x ) || p
θ ( x ))
= − entropy (˜ p
θ ( x ))
− E ˜ p
θ ( x ) [ f ( x )]
+ log Z .
The gradient ascent part of Langevin dynamics, x t +
= x t + δ 2 / 2 ∇ x f θ ( x t ) , 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), D KL (˜ p θ ( x ) || p θ ( x )) decreases to zero monotonically as T → ∞ .
3. N ORMALIZING F LOW
Let z ∈ R D be the latent vector of the same dimensionality as x .
A normalizing flow is of the form
[Equation]
where q 0 ( z ) is a known prior distribution such as Gaussian white noise distribution N (0 , I D ) , and g α : R D
→ R D is a mapping that consists of a sequence of L invertible transformations, i.e., g ( z ) =
g L ◦ · · · ◦ g 2 ◦ g 1 ( z ) , whose inversion z = g − 1 α ( x ) and log-determinants of the Jacobians can be computed in closed form.
α are the parameters of g α .
The mapping is used to transform a random variable z that follows a simple distribution q 0 into a flexible distribution.
Under the change-of-variables law, the resulting random variable x = g α ( z ) has a probability density q α ( x ) =
q 0 ( g − 1 α ( x )) | det( ∂g − 1 α ( x ) /∂x ) | .
Let h l = g l ( h l − 1 ) .
The successive transformations between x and z can be expressed as a flow z
g 1 ←→ h 1 g 2 ←→ h 2 · · · g L ←→ x , where we define z := h 0 and x := h L for succinctness.
Then the determinant becomes | det( ∂g − 1 α ( x ) /∂x )
| = (cid:81)
Ll =
| det( ∂h l − 1 /∂h l ) | .
The log-likelihood of a datapoint x can be easily computed by
[Equation]
[Equation]
With some smart designs of the sequence of transformations g α = { g l , l = 1 , ..., L } , the logdeterminant in Eq.(6) can be easily computed, then the flow-based model can be trained by maximizing the data log-likelihood L ( α )
=
(cid:80) ni log q α ( x i ) .
C OOP F LOW : C OOPERATIVE T RAINING OF T WO F LOWS 4. C OOP F LOW A LGORITHM
We have moved from trying to use a convergent Langevin dynamics to train a valid EBM.
Instead, we accept the fact that the short-run non-convergent MCMC is inevitable and more affordable in practice, and we treats a non-convergent short-run Langevin flow as a generator and propose to jointly train it with a normalizing flow as a rapid initializer for more efficient generation.
The resulting generator is called CoopFlow, which consists of both Langevin flow and normalizing flow.
Specifically, at each iteration, for i = 1 , ..., m , we first generate z i ∼ N (0 , I D ) , and then transform z i by a normalizing flow to obtain ˆ x i =
g α
( z i ) .
Next, starting from each ˆ x i , we run a Langevin flow (i.e., a finite number of Langevin steps toward an EBM p θ ( x ) ) to obtain ˜ x i .
˜
x
i are considered as the synthesized examples that are generated by the CoopFlow models.
We then update α of the normalizing flow by treating ˜ x i as training data, and update θ of the Langevin flow according to the learning gradient of the EBM, which is computed with the synthesized examples ˜ x i and the observed examples.
Algorithm 1 presents a description of the proposed CoopFlow algorithm.
The advantage of this training scheme is that we only need to minimally modify the existing code for the
MLE training of EBM p θ and normalizing flow q α .
The density of the CoopFlow π ( θ,α ) ( x ) is well defined, which can be implicitly expressed by
[Equation]
If we increase the length of Langevin flow, π ( θ,α ) will converge to EBM p θ ( x ) .
The network f θ ( x ) in Langevin flow is scalar valued and is of free form, whereas the network g α ( x ) in normalizing flow has high-dimensional output and is of a severely constrained form.
Thus Langevin flow can potentially provide a tighter fit to p data ( x ) than normalizing flow.
On the flip side, sampling from
Langevin flow requires MCMC, whereas normalizing flow can synthesize examples directly via ancestral sampling.
It is thus desirable to train these two flows, where normalizing flow serves as an approximate sampler to amortize the sampling cost of Langevin flow.
Meanwhile the normalizing flow can be updated by a temporal difference MCMC teaching scheme provided by Langevin flow.
Given maximizing log α
Given following the gradient ∇
4. U NDERSTANDING THE L EARNED T WO F LOWS
Convergence equations In the traditional contrastive divergence (CD) algorithm, MCMC chains are initialized with observed data so that the CD learning seeks to minimize D KL ( p data ( x ) ||
p θ ( x ))
− D
KL (( K θ p data )( x ) || p
θ ( x )) , where ( K θ p data )( x ) denotes the marginal distribution obtained by running the Markov transition K , which is specified by the Langevin flow, from the data distribution.
In the proposed CoopFlow algorithm, the learning of the EBM (or the Langevin flow model) follows a modified contrastive divergence, where the initial distribution of the Langevin flow is modified to be a normalizing flow q α .
Thus, at iteration t , the update of θ follows the gradient of D
KL ( p data || p θ ) − D KL ( K θ ( t )
q α ( t ) ||
p θ ) with respect to θ .
Compared to the traditional CD loss, the modified one replaces p data by q α in the second KL divergence term.
At iteration t , the update of the normalizing flow q α follows the gradient of D KL ( K θ ( t ) q α ( t )
||
q α ) with respect to α .
Let ( θ ∗ , α ∗ ) be a fixed point the learning algorithm converges to, then we have the following convergence equations
[Equation]
[Equation]
Ideal case analysis In the idealized scenario where the normalizing flow q α has infinite capacity and the Langevin sampling can mix and converge to the sampled EBM, Eq.(13) means that q α ∗ attempts to be the stationary distribution of K θ ∗ , which is p θ ∗ because K θ ∗ is the Markov transition kernel of a convergent MCMC sampled from p θ ∗ .
That is, min α D KL ( K θ ∗ q α ∗
|| q α )
= 0 , thus q α ∗ =
p θ ∗ .
With the second term becoming 0 and vanishing, Eq.(12) degrades to min θ D KL ( p data || p θ ) and thus θ ∗ is the maximum likelihood estimate.
On the other hand, the normalizing flow q α chases the
EBM p θ toward p data , thus α ∗ is also a maximum likelihood estimate.
Moment matching estimator In the practical scenario where the Langevin sampling is not mixing, the CoopFlow model π t =
K θ ( t ) q α ( t ) is an interpolation between the learned q α ( t ) and p θ ( t ) , and it converges to π ∗ =
K θ ∗ q α ∗ , which is an interpolation between q α ∗ and p θ ∗ .
π ∗ is the short-run Langevin flow starting from q α ∗ towards EBM p θ ∗ .
π ∗ is a legitimate generator because E p data [ ∇ θ f θ ∗ ( x )] =
E π ∗ [ ∇ θ f θ ∗ ( x )] at convergence.
That is, π ∗ leads to moment matching in the feature statistics ∇ θ f θ ∗ ( x ) .
In other words, π ∗ satisfies the above estimating equation.
Consider a simple EBM with f θ ( x ) = ⟨ θ, h ( x ) ⟩ , where h ( x ) is the feature statistics.
Since ∇ θ f θ ( x ) = h ( x ) , the MLE of the
EBM p θ MLE is a moment matching estimator due to E p data [ h ( x )] =
E p θ MLE [ h ( x )] .
The CoopFlow π ∗ also converges to a moment matching estimator, i.e., E p data [ h ( x )]
=
E π ∗ [ h ( x )] .
Figure 1 is an illustration of different model distributions that correspond to different parameters at convergence.
We first introduce three families of distributions: Ω = { p : E p [ h ( x )] =
E p data [ h ( x )] } , Θ = { p θ ( x ) = exp( ⟨ θ, h ( x ) ⟩ ) /Z
( θ ) , ∀ θ } , and A = { q α , ∀ α } , which are shown by the red, blue and green curves respectively in Figure 1.
Ω is the set of distributions that reproduce statistical property h ( x ) of the data distribution.
Obviously, p θ MLE , p data , and π ∗ =
K θ ∗ q α ∗ are included in Ω .
Θ is the set of EBMs with different values of θ , thus p θ MLE and p θ ∗ belong to Θ .
Because of the short-run Langevin flow, p θ ∗ is not a valid model that matches the data distribution in terms of h ( x ) , and thus p θ ∗ is not in Ω .
A is the set of normalizing flow models with different values of α , thus q α ∗ and q α MLE belong to A .
The yellow line shows the MCMC trajectory.
The solid segment of the yellow line, starting from q α ∗ to K θ ∗ q α ∗ , illustrates the short-run non-mixing MCMC that is initialized by the normalizing flow q α ∗ in A and arrives at π ∗ =
K θ ∗ q α ∗ in Ω .
The dotted segment of the yellow line, starting from π ∗ =
K θ ∗ q α ∗ in Ω to p θ ∗ in Θ , shows the potential long-run MCMC trajectory, which is not really realized because we stop short in MCMC.
If we increase the number of steps of short-run Langevin flow, KL ( π ∗ || p θ ∗ ) will be monotonically decreasing to 0.
Though π ∗ stops midway in the path toward p θ ∗ , π ∗ is still a valid generator because it is in Ω .
The EBM with MLE p θ MLE is the intersection between Θ and Ω .
It is the projection of p data onto Θ because θ MLE = arg min θ D KL ( p data || p θ ) .
The normalizing flow with MLE q α MLE is also a projection of p data onto A because α MLE = arg min α D KL ( p data || q α ) .
Generally, q α MLE is far from p data due to its restricted form network, whereas p θ MLE is very close to p data due to its scalar-valued free form network.
Because the short-run Langevin flow is not mixing, θ ∗ ̸ = θ MLE and α ∗ ̸ = α MLE .
Specifically, θ ∗ and α ∗ are perturbations of θ MLE and α MLE , respectively.
p θ MLE shown by an empty dot is not attainable unfortunately.
Note that Ω , as well as Θ and A , is high-dimensional manifold instead of 1D curve as depicted in the Figure, and π ∗ may be farther away from p data than p θ MLE is.
During learning, q α ( t +1) is the projection of K θ ( t ) q α ( t ) on A .
At convergence of the learning algorithm, there is an infinite looping, i.e., K θ ∗ q α ∗ (i.e., π ∗ ) lifts
q α ∗ off the ground A , and the projection drops
K θ ∗ q α ∗ back to q α ∗ .
E XPERIMENTS
In this section, we showcase the experiment results on various tasks.
We start from a toy example to illustrate the basic idea of the CoopFlow algorithm in Section 5.1.
Then we show the image generation results in section 5.2.
Section 5.3 demonstrates the learned CoopFlow can be useful for image reconstruction and inpainting, while Section 5.4 shows that the learned latent space is meaningful so that it can be used for interpolation.
5. T OY E XAMPLE S TUDY
We first demonstrate our idea using a two-dimensional toy example where data lies on a spiral.
We train 3 CoopFlow models with different lengths of Langevin flows.
As shown in Figure 2, the rightmost box shows the results obtained with 10,000 Langevin steps.
We can see in this situation, both the normalizing flow q α and the EBM model p θ can capture the underlying distribution perfectly.
This validates that, with a sufficiently long Langevin flow, the CoopFlow algorithm can learn both a valid EBM and a valid normalizing flow.
The leftmost green box represents the model trains with 100-step Langevin updates.
This is a more typical non-convergent MCMC setting.
In this case, neither q α nor p θ are perfect.
Instead, they need to work together and the Langevin updates process over EBM model actually works as an extra generator that modifies the initial proposal by the normalizing flow model.
We can see that the samples from q α are not perfect, but after the modification, the samples from π ( θ,α ) follow the underlying distribution well.
The middle green box shows the results obtained with 500-step Langevin steps.
It is in the middle of short-run and long-run.
p θ becomes better but still not perfect.
Using the toy example here, we want to make the points that short-run non-convergent MCMC situation might still exist even when we have a generator for initial proposal and using relatively long MCMC chain.
Besides, if we use the MCMC on EBM as a special generator, i.e. Langevin Flow, then together with the normalizing flow, we can get pretty good samples even with few MCMC updates.
5. I MAGE GENERATION
We test our model on 3 image datasets for image synthesis.
( i ) CIFAR-10 (Krizhevsky et al., 2009) is a dataset containing 50k training images and 10k testing images in 10 classes; ( ii ) SVHN (Netzer et al., 2011) is a dataset containing over 70k training images and over 20k testing images of house numbers; ( iii ) CelebA
(Liu et al., 2015) is a celebrity facial dataset containing over 200k imagess.
We downsample all the images to the resolution of 32 × 32 .
For our model, we show results of three different settings.
CoopFlow denotes the setting where we train a normalizing flow and a Langevin flow together from scratch and use 30 Langevin steps.
CoopFlow(long) denotes the setting that we increase the number of Langevin updates to 200.
In CoopFlow(Pre) setting, we first pretrained a normalizing flow from observed data, and then we train the CoopFlow with the parameters of the normalizing flow being initialized by the pretrained one.
We use a 30-step long Langevin flow in this setting.
For all the three settings, we slightly increase the step size of Langevin update at the test stage for better performance.
For more training details, please refer to Appendix.
We show both the qualitative results in Figure 4 and quantitative results in Table 1.
To calculate FID
Heusel et al.
(2017) scores, we generate 50,000 samples on each dataset.
From these results, we can see that our models outperform most of the baseline algorithms.
We get lower FID scores comparing to the purely normalizing flow models and the prior works that try to jointly train a normalizing flow with an EBM, e.g., Gao et al.
; Nijkamp et al.
(2020a).
We also achieve comparable results with the SOTA EBMs.
Among our different settings, we can see using a longer Langevin flow consistently improves the performance.
And using a pretrained normalizing flow can provide the lowest FID scores among all the three settings.
More results can be found in Appendix
[Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
5. I MAGE RECONSTRUCTION
We show that the learned CoopFlow model is able to reconstruct observed images.
We may consider the CoopFlow model π ( θ,α ) ( x ) to be a latent variable generative model: z ∼ p 0 ( z ); ˆ x = g α
( z ); x = F θ (ˆ x, e ) , where z denotes the latent variables, e denotes all the injected noises in the Langevin flow, and F θ denotes the mapping realized by a T -step Langevin flow that is actually a T -layer noise-injected residual network.
Since the Langevin flow is not mixing, x is dependent on ˆ x in the Langevin flow, thus also dependent on z .
The CoopFlow is a generator x = F θ ( g α ( z ) , e ) , so we can reconstruct any x by inferring the corresponding latent variables z using gradient descent on L ( z ) =
|| x − F θ ( g α
( z ) , e )
|| 2 , with z being initialized by p 0 .
However, g is an invertible transformation, so we can infer z by an efficient way, i.e., we first find ˆ x by gradient descent on L (ˆ x ) =
|| x − F θ (ˆ x, e ) || 2 , with ˆ x being initialized by ˆ x 0 =
g α ( z ) where z ∼ p 0 ( z ) and e being set to be 0, and then use z = g − 1 α (ˆ x ) to get the latent variables.
These two methods are equivalent, but the latter one is computationally efficient, since computing the gradient on the whole two-flow generator F θ ( g α ( z ) , e ) is difficult and time-consuming.
Let ˆ x ∗ =
arg min ˆ x L (ˆ x ) .
The reconstruction is given by
F θ (ˆ x ∗ ) .
The optimization is done using 200 steps of gradient descent over ˆ x .
The results are shown in Figure 3.
We do experiment on CIFAR10 test set.
The right column displays the original images x that need to be reconstructed.
The left column and the middle column display ˆ x ∗ and F θ (ˆ x ∗ ) , respectively.
We can see that our model can successfully reconstruct the observed images, which verifies that the CoopFlow with a non-mixing MCMC is a valid latent variable generator.
We further show that our model is also capable of doing image inpainting.
Similar to image reconstruction, given a masked observation x mask with a binary matrix M indicating the positions of the unmasked pixels, we optimize ˆ x to minimize the reconstruction error between F θ (ˆ x ) and x mask in the unmasked area, i.e., L (ˆ x ) =
||
M ⊙ ( x mask − F θ (ˆ x ))
|| 2 , where ⊙ is the element-wise multiplication operator.
ˆ x is still initialized by the normalizing flow.
We do experiment on Celeba training set. In Figure 5, each
[Figure]
5. I NTERPOLATION IN THE LATENT SPACE
The CoopFlow model is capable of doing interpolation in the latent space z .
Given an image x , we first find its corresponding ˆ x ∗ using the reconstruction method described in section 5.3.
We then infer z by the inversion of the normalizing flow z ∗ =
g − 1 α (ˆ x ∗ ) .
The results of interpolation in latent space experiments on Celeba (training set) are shown in Figure 6.
Comparing the first and second columns to the left (or to the right), We can see that our reconstructions of the original images are good.
Further, our interpolations are smooth and the generated images are meaningful.
C ONCLUSION
This paper studies an interesting problem of learning two types of deep flow models in the context of energy-based framework for image representation and generation.
One is the normalizing flow model that generates synthesized examples by transforming Gaussian noise examples through a sequence of invertible transformations, while the other is the Langevin flow that generates synthesized examples by running a non-mixing, non-convergent short-run MCMC toward an EBM.
We propose the CoopFlow algorithm to train the short-run Langevin flow model jointly with the normalizing flow serving as a rapid initializer in a cooperative manner.
The experiments show that the CoopFlow is a valid generative model that can be useful for image generation, reconstruction, and interpolation.
A A PPENDIX A. M ODEL A RCHITECTURE
We use the same model structure for all the experiments.
For the normalizing flow model, we use the Flow++ (Ho et al., 2019) structure designed for the CIFAR-10 dataset.
For the EBM model, we use the structure shown in Table 2.
3 × 3 Conv2d, s=1, pad=1, ch=128; Swish
Res
Block, ch=
Res
[Table]
=
[Table]
Res
[Table] [Table]
3 × 3 Conv2d, s=4, pad=0, ch=100; Swish
[Table]
A. E XPERIMENTAL S ETTING
We have tried three different settings.
In CoopFlow setting and CoopFlow(long) setting, we train both normalizing flow model and Langevin flow model from scratch.
The difference between them are the number of Langevin steps use.
The CoopFlow(Long) use a relatively long chain of 200 steps
MCMC updates.
Following Ho et al.
(2019), we use data-dependent parameter initipalization (Salimans & Kingma, 2016) for our normalizing flow model.
On the other hand, in the CoopFlow(Pre) setting, we first pretrained the normalizing flow model using true observed data.
Then we use this pretrained normalizing flow model as initialization and train it together with our Langevin flow model using Algorithm alg1.
To make this switching smoother, we include a warm-up phase for
Langevin flow model.
During this phase, instead of updating both the normalizing flow model and
Langevin flow model from the beginning, we first fix the weight of pretrained normalizing flow model and only update the weight of EBM.
After a certain number of epochs of updating, the
Langevin flow model itself has been able to roughly capture the underlying data distribuion.
Then we beginning to update both the normalizing flow and Langevin flow together.
This may prevent the
Langevin flow from proposing bad samples at the begining of training and affecting the pretrained normalizing flow model.
During training, we use 1e-4 as learning rate for both the normalizing flow and Langevin flow model, we use Adam (Kingma & Ba, 2014) as our optimizer.
We use βs =
(0 . 9 , 0 . 999) for normalizing flow model and βs =
(0 . 5 , 0 . 5) for Langevin flow.
We keep monitoring the FID scores during training and pick the best one over the training trajectory.
For data augmentation, we only include random horizontal flip for the CIFAR-10 dataset and do not apply other tricks.
Given we are actually using the short-run MCMC setting, we remove the noise term in Langevin updates so that the updates becomes gradient descent over the energy function.
The other hyper-parameter setting can be found in Table 3.
[Figure]
We further show the influence of Langevin step size and number of Langevin steps on the CIFAR10 dataset using the CoopFlow(Pre) model.
We first fix the Langevin step size to be 0.03 and try different number of Langevin updates from 10 to 50.
The results are shown in Table 4.
On the other hand, we show the influence of Langevin Step size in Table 5.
Here, we keep the number of Langevin updates to be 30 and we adjust the training Langevin step size.
We keep the testing step size and training step size to have a ratio of 43 .
We can see that our choice of 30 as number of Langevin steps and 0.03 as Langevin step size is reasonable.
Based on the comparison between
CoopFlow and CoopFlow(Long), further increase the number of Langevin steps to a relatively long chain may helps with the FID score, but given long MCMC chain can be computationally expensive, we keep to use 30 steps here.
A. A BLATION S TUDY
Here we compare the results of training normalizing flow and Langevin flow using or Cooperative algorithm with training them separately.
Training flow itself follows the original normalizing flow training algorithm
(we use the flow++ (Ho et al., 2019) structure here, which is the same as the one used in our CoopFlow model).
And for training Langevin flow individually, we use the short-run
MCMC proposed by Nijkamp et al.
with our own EBM structure.
Through our experiments, we find that removing the noise injection to both observed data and Langevin updates can give us better result.
We show generated samples of normalizing flow and short-run EBM on CIFARdataset in Figure 7 and FID scores in Table 6.
We can see that the CoopFlow outperforms both the normalizing flow and short-run EBM by a large margin.
[Figure]
A. M ORE G ENERATION R ESULTS
In section 5.2, we show the samples from CoopFlow(Pre).
Here, we show the samples generated by
CoopFlow in Figure 8 and samples generated by CoopFlow(Long) in Figure 9.
We can see that all the samples are meaningful.
On the other hand, we the initial samples from the normalizing flow trained under the CoopFlow setting in Figure 10.
Comparing Figure 4 and Figure 10, we can see that there is a gap between the normalizing flow samples and the final CoopFlow outputs.
The initial samples from normalizing flow are blurred and after modified by the Langevin flow, the samples become sharp and clear.
This agrees with our discussion in section 4.2.
The CoopFlow(Long) setting and CoopFlow(pre) setting give similar observations.
We show more image inpainting results in Figure 11.
In Figure 12, we show the FID score curve over training epochs.
The FID keep decreasing as the learning proceeds.
We can see the FID curves converge to a low value for both
CoopFlow setting and CoopFlow(pre) on CIFAR-10 dataset.
This is an evidence to show that our algorithm is a convergent algorithm.
[Figure] [Figure]
[Table]
[Figure]
[Figure] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table] [Table]
A. Q UANTITATIVE R ESULTS F OR T HE I MAGE R ECONSTRUCTION E XPERIMENT
Following Nijkamp et al.
(2019), we calculate the per pixel MSE on 1,000 examples in the leave-out testset of CIFAR-10 dataset.
We use a 200-step gradient descent to minimize the reconstruction loss.
We plot the MSE per update iteration in Figure 13 and we report the final per pixel MSE in Table 7.
For a baseline, we test the reconstruction of the same 1,000 images using the short-run EBM model which has the same EBM structure as our CoopFlow and uses 100 steps of MCMC.
We can see the
CoopFlow works better than the short-run EBM in this image reconstruction task.
[Table]
[Figure]
A. C OMPARISON OF MODEL SIZES
In Table 8, we present an comparison of different models according to their size and the FID scores they got.
Here we mainly compare those models that employ a normalizing flow as a generator, e.g.
EBM-FCE, NT-EBM, GLOW or EBM jointly trained with other generator structure, e.g., VAEBM.
We can see our CoopFlow achieves a good computation and performance balance.
On one hand, we have similar number of parameters as EBM-FCE but has a much lower FID.
Also, the Flow++ only represents the setting that we train the normalizing flow model with the same structure as the one used in our CoopFlow.
Comparing it to CoopFlow, we can see that introducing a Langevin Flow increases the FID a lot.
On the other hand, the VAEBM model achieves better FID than us, but they need to use a pretrained NVAE model (Vahdat & Kautz, 2020), which is much larger than the one we used here.
[Table]
A. A NALYSIS OF THE NUMBER OF L ANGEVIN STEPS
In this section, we compare our method with other baselines in generating examples with the same number of Langevin steps.
We compare our methods with two closely related baselines, (1) the single EBM with short-run MCMC starting from noise distribution (Nijkamp et al., 2019), and (ii)
Cooperative learning of an EBM and a generic generator (Xie et al., 2018a).
The numbers of MCMC steps and FID reported in Table 9 are collected from the original papers.
In Table 10, we also report the FID scores of different baseline methods using the different numbers of MCMC steps.
We can see that with the same number of Langevin steps, the CoopFlow can generate much more realistic image patterns than the other two baselines.
The results of the experiments show that the proposed
CoopFlow can use less number of MCMC steps to achieve better results.
A. M ORE C ONVERGENCE A NALYSIS
The CoopFlow algorithm simply involves two MLE learning algorithms: (i) the MLE learning of the EBM p θ and (ii) the MLE learning of the normalizing flow q α .
The convergence of each of the two learning algorithms has been well studied and verified in the existing literature, e.g., Younes
[Table] [Table]
; Xie et al. (2016); Kingma & Dhariwal (2018b).
That is, each of them has a fixed point.
The only interaction between these two algorithms in the proposed CoopFlow algorithm is that, in each iteration, they feed each other with their synthesized examples and use the cooperatively synthesized examples in their parameter update formulas.
To be specific, the normalizing flow uses its synthesized examples to initialize the MCMC of the EBM, and the EBM feed the normalizing flow with its synthesized examples as training examples.
The synthesized examples from the Langevin flow are considered the cooperatively synthesized examples by two models, and used to compute their learning gradients.
Unlike other amortized sampling methods (e.g., Han et al.
Grathwohl et al.
(2021)) that uses variational learning, the EBM and normalizing flow in our framework don’t back-propagate each other through the cooperatively synthesized examples.
They just feed each other with some input data for their own training algorithms.
That is, each learning algorithm will still converge and have a fixed point.
Now let us analyze the convergence of the whole CoopFlow algorithm that alternates two maximum likelihood learning algorithms.
We first analyze the convergence of the objective function at each learning step, and then we conclude the convergence of the whole algorithm.
The convergence of CD learning of EBM .
The learning objective of the EBM is to minimize the
KL distance between the EBM and the data distribution.
Since the MCMC of EBM in our case is initialized by the normalizing flow, it follows a modified contrastive divergence algorithm.
That is, at iteration t , it has the following objective,
[Equation]
No matter what kind of distribution is used to initialize the MCMC, it will has a fixed point when the learning gradient of θ equals to 0, i.e.,
[Equation]
The initialization of the MCMC only affects the location of the fixed point of the learning algorithm.
The convergence and the analysis of the fixed point of contrastive divergence has been studied by
Hinton (2002); Carreira-Perpi˜n´an & Hinton (2005).
The convergence of MLE learning of normalizing flow The normalizing flow’s objective is to learn to minimize the KL divergence between the normalizing flow and the Langevin flow because, in each learning iteration, the normalizing flow uses the EBM’s synthesized examples as training data.
At iteration t , it has the following objective
[Equation]
which is a convergent algorithm at each t .
The convergence has been studied by Younes (1999).
The convergence of CoopFlow
The CoopFlow alternates the above two learning algorithms.
The
EBM learning will reduce the KL distance between EBM and data, i.e., p θ → p data ; while the
MLE learning of normalizing flow will reduce the KL distance between normalizing flow and EBM, i.e., q α → p θ .
Therefore, the normalizing flow will chase the EBM toward the data distribution gradually.
Because the process p θ → p data will stop at a fixed point, therefore q α → p θ will also stop at a fixed point obviously.
Such a chasing game is a contraction algorithm, therefore the fixed point of CoopFlow exists.
Empirical evidence has also supported that our EM-like learning algorithm is a contraction.
If we use ( θ ∗ , α ∗ ) to denote the fixed point of the CoopFlow, according to the definition of a fixed point, ( θ ∗ , α ∗ ) will satisfy
[Equation]
The convergence of the cooperative learning framework (CoopNets) that integrates MLE algorithm of an EBM and MLE algorithm of a generic generator has been verified and discussed in Xie et al.
The CoopFlow that uses a normalizing flow instead of a generic generator has the same convergence property as that of the original CoopNets.
The major contribution in our paper is to start from the above fixed point equation to discuss where the fixed point will be in our algorithm when the MCMC of the EBM is non-mixing and non-convergent.
This will go beyond all the prior works about cooperative learning.
A. D ISCUSSION ON THE BIAS OF THE CONTRASTIVE DIVERGENCE OBJECTIVE FOR EBM Bias
in CD learning .
The EBM trained by contrastive divergence (CD) still follows the “analysis by synthesis” scheme to update the model parameters θ .
The learning gradient of θ is still computed by
[Equation] [Equation]
except that the synthesized examples ˜ x are sampled by a short-run MCMC initialized by the training data p data .
For example, the CD-1 is commonly used to train Restricted Boltzmann machines (Hinton, 2002), and “1” means the single-step MCMC.
We denote ˜ x ∼ K θ p data ( x ) , where K θ is the
MCMC transition kernel of T steps of MCMC that samples p θ , and ( K θ p data )( x ) is the distribution of x after T steps of MCMC update starting from q data .
Because ˜ x ∼ K θ p data ( x ) ̸
= p θ ( x ) , therefore, the estimated θ using the above equation is biasd (i.e., it is not an MLE solution, but it still a convergent algorithm with a fixed point.)
Hinton (2002) point out that the CD learning actually minimizes the following objective, which is a perturbation of the MLE,
[Equation]
where the first term is the MLE due to the use of Eq.(14), the second term is the bia due to the gap between the actually used sampling distribution K θ p data and the correct distribution p θ ( x ) we should sample from.
Only if D KL ( K θ p data || p
θ )
= 0 , the CD estimate will be unbiased.
However, the original CD doesn’t explicitly minimize the second term.
Recently, Du et al.
(2021) improves the CD algorithm by explicitly minimizing the second term that has been neglected in the original
CD to further reduce the bias.
We need to point out that, even though the CD is biased, it is still a convergent algorithm, whose convergence has been well studied in its original paper (Hinton, 2002).
“Biasd” doesn’t mean it is wrong.
“Biasd” is just because D KL ( K θ p data || p θ ) ̸ = 0 .
Bias in amortized inference Actually, many useful algorithms and models using amortized sampling or amortized inference are biased (with respect to MLE).
For example, variational autoencoder (VAE) (Kingma & Welling, 2013), whose objective is also a perturbation of MLE,
[Equation]
16) where q α is the distribution of the decoder/generator in VAE, and
π ϕ ( z | x ) is the encoder to approximate the true posterior q α
( z | x ) .
Even though VAE explicitly minimizes the second term (i.e., bias) in its loss, but in practice, this term can not be zero due to the mean-field design of the inference model π ϕ , the model α is still biased to MLE.
Even though our model is not related to VAE, we just want to mention that “bias” is not a bad thing, it is just a trade-off between computation efficiency and the exact MLE solution due to the usage of amortized sampling/inference.
Bias in CoopFlow Now let us go back to our CoopFlow algorithm.
The update of the EBM in our CoopFlow can be considered a modified CD algorithm, because its MCMC starts from the normalizing flow q α .
At iteration t , the update of the EBM follows the following objective
[Equation]
which is biased with respect to MLE if D KL ( K θ ( t )
q α ( t ) || p θ ) ̸
= 0 .
However, the CoopFlow algorithm has the other part that can handle the second term.
It is the MLE learning of the normalizing flow q α ( x ) , with ˜ x ∼ K θ ( t ) q α ( t ) as training examples.
That is, at iteration t , we fix K θ ( t ) q α ( t ) and minimize the following objective
[Equation]
Eq.(18) tries to reduce the KL distance between the flow and the EBM.
In an ideal case where capacities of all models are large enough and the MCMC K θ is mixing and convergent MCMC, minimizing Eq.(18) will lead to D KL ( K θ q α
||
q α ) → 0 , which means q α = p θ because q α has become the stationary distribution of K θ .
If q α = p θ = 0 , then the second term of Eq.(17) will be vanishing because D KL ( K θ q α || p θ )
=
D KL ( K θ p θ || p
θ )
= 0 , the Eq.(17) only has the first term left, which is the
MLE solution.
Since q α = p θ , thus α is also MLE solution.
Therefore, our estimates of θ and α in the CoopFlow are unbiased in the ideal case.
However, in practice, the conditions, such as large enough capacity of models and convergent
MCMC are very hard to get.
Eventually, the CoopFlow might converge to a fixed point but not the fixed point of MLE.
We have used information geometry to illustrate the fixed point in Figureof the main text, and analyze the biased estimator in the paragraph “moment matching estimator” of section 4.2.
That is, even though the CoopFlow is not MLE solution, but it is a moment matching estimator (MME), which is a perturbation of MLE shown in Figure 1.
MLE is shown as p θ MLE in
Figure 1, and the obtained MME solution is marked as a star π ∗ .
Again, bias doesn’t means wrong.
π θ,α with MME is a valid generator.
A. C OMPARISON BETWEEN C OOP F LOW AND SHORT - RUN EBM VIA INFORMATION GEOMETRY
Figure 14 illustrates the convergences of both CoopFlow and the EBM with a short-run MCMC starting from an initial noise distribution q 0 .
We define the short-run EBM in a more generic form (Xie et al., 2016)
[Equation]
(19) which is an exponential tilting of a known reference distribution q 0 ( x ) .
In general, the reference distribution can be either a Gaussian distribution or a uniform distribution.
When the reference distribution is a uniform distribution, q 0 can be removed in Eq.(19).
Since the initial distribution of the CoopFlow is a Gaussian distribution q 0 , which is actually the prior distribution of the normalizing flow.
For a convenient and fair comparison, we use Gaussian distribution as the reference distribution of the EBM in Eq.(19).
And both CoopFlow and the baseline short-run EBM use the same EBM defined in Eq.(19) in their frameworks.
We use ¯ p ¯ θ to denote the baseline short-run EBM.
There are three families of distributions: Ω = { p : E p
[ h ( x )]
=
E p data [ h ( x )] } , Θ = { p θ ( x ) = exp( ⟨ θ, h ( x ) ⟩ ) q 0 ( x ) /Z
( θ ) , ∀ θ } , and
A = { q α , ∀ α } , which are shown by the red, blue and green curves respectively in Figure 14.
Figure 14 is an extension of Figure 1 by adding the following elements.
(i) q 0 , which is the initial distribution for both CoopFlow and short-run EBM.
It belongs to Θ because it corresponds to θ = 0 .
q 0 is just a noise distribution thus it is far under the green curve.
That is, it is very far from q α ∗ because q α ∗ has been already a good approximation of p θ ∗ .
(ii) g α ∗ , which is the learned transformation of the normalizing flow q α , and is visualized as a mapping from q 0 to q α ∗ by a directed brown line segment.
(iii)
The MCMC trajectory of the baseline shortrun EBM ¯ p ¯ θ , which is shown by the yellow line on the right hand side of the blue curve.
The solid part of the yellow line, starting from q 0 to ¯ π ∗ =
K ¯ θ ∗ q 0 , shows the short-run non-mixing MCMC starting from the initial Gaussian distribution q 0 in Θ and arriving at ¯ π ∗ in Ω .
The dotted part of the yellow line is the potential long-run MCMC trajectory that is unrealized.
To compare the MCMC trajectories of the CoopFlow and the short-run EBM, we can find that CoopFlow has a much shorter
MCMC trajectory than that of the short-run EBM, since the normalizing flow amortizes the sampling workload for the Langevin flow in the CoopFlow model.