Flow Contrastive Estimation of Energy-Based Models

Introduction

Recently, flow-based models (henceforth simply called flow models) have gained popularity as a type of deep generative model and for use in variational inference .

Flow models have two properties that set them apart from other types of deep generative models: (1) they allow for efficient evaluation of the density function, and (2) they allow for efficient sampling from the model. Efficient evaluation of the log-density allows flow models to be directly optimized towards the log-likelihood objective, unlike variational autoencoders (VAEs) , which are optimized towards a bound on the log-likelihood, and generative adversarial networks (GANs) . Auto-regressive models , on the other hand, are (in principle) inefficient to sample from, since synthesis requires computation that is proportional to the dimensionality of the data.

These properties of efficient density evaluation and efficient sampling are typically viewed as advantageous. However, they have a potential downside: these properties also acts as assumptions on the true data distribution that they are trying to model. By choosing a flow model, one is making the assumption that the true data distribution is one that is in principle simple to sample from, and is computationally efficient to normalize. In addition, flow models assume that the data is generated by a finite sequence of invertible functions. If these assumptions do not hold, flow-based models can result in a poor fit.

On the other end of the spectrum of deep generative models lies the family of energy-based models (EBMs) . Energy-based models define an unnormalized density that is the exponential of the negative energy function. The energy function is directly defined as a (learned) scalar function of the input, and is often parameterized by a neural network, such as a convolutional network . Evaluation of the density function for a given data point involves calculating a normalizing constant, which requires an intractable integral. Sampling from EBMs is expensive and requires approximation as well, such as computationally expensive Markov Chain Monte Carlo (MCMC) sampling. EBMs, therefore, do not make any of the two assumptions above: they do not assume that the density of data is easily normalized, and they do not assume efficient synthesis. Moreover, they do not constrain the data distribution by invertible functions.

Contrasting an EBM with a flow model, the former is on the side of representation where different layers represent features of different complexities, whereas the latter is on the side of learned computation, where each layer, or each transformation, is like a step in the computation. The EBM is like an objective function or a target distribution whereas the flow model is like a finite step iterative algorithm or a learned sampler. Borrowing language from reinforcement learning , the flow model is like an actor whereas the EBM is like a critic or an evaluator. The EBM can be simpler and more flexible in form than the flow model which is highly constrained, and thus the EBM may capture the modes of the data distribution more accurately than the flow model. In contrast, the flow model is capable of direct generation via ancestral sampling, which is sorely lacking in an EBM. It may thus be desirable to train the two models jointly, combining the tractability of flow model and the flexibility of EBM. This is the goal of this paper.

Our joint training method is inspired by the noise contrastive estimation (NCE) of , where an EBM is learned discriminatively by classifying the real data and the data generated by a noise model. In NCE, the noise model must have an explicit normalized density function. Moreover, it is desirable for the noise distribution to be close to the data distribution for accurate estimation of the EBM. However, the noise distribution can be far away from the data distribution. The flow model can potentially transform or transport the noise distribution to a distribution closer to the data distribution. With the advent of strong flow-based generative models , it is natural to recruit the flow model as the contrast distribution for noise contrastive estimation of the EBM.

However, even with the flow-based model pre-trained by maximum likelihood estimation (MLE) on the data distribution, it may still not be strong enough as a contrast distribution, in the sense that the synthesized examples generated by the pre-trained flow model may still be distinguished from the real examples by a classifier based on an EBM. Thus, we want the flow model to be a stronger contrast or a stronger training opponent for EBM. To achieve this goal, we can simply use the same objective function of NCE, which is the log-likelihood of the logistic regression for classification. While NCE updates the EBM by maximizing this objective function, we can also update the flow model by minimizing the same objective function to make the classification task harder for the EBM. Such update of flow model combines MLE and variational approximation, and helps correct the over-dispersion of MLE. If the EBM is close to the data distribution, this amounts to minimizing the Jensen-Shannon divergence (JSD) between the data distribution and the flow model. In this sense, the learning scheme relates closely to GANs . However, unlike GANs, which learns a generator model that defines an implicit probability density function via a low-dimensional latent vector, our method learns two probabilistic models with explicit probability densities (a normalized one and an unnormalized one).

The contributions of our paper are as follows. We explore a parameter estimation method that couples estimation of an EBM and a flow model using a shared objective function. It improves NCE with a flow-transformed noise distribution, and it modifies MLE of the flow model to approximate JSD minimization, and helps correct the over-dispersion of MLE. Experiments on 2D synthetic data show that the learned EBM achieves accurate density estimation with a much simpler network structure than the flow model. On real image datasets, we demonstrate a significant improvement on the synthesis quality of the flow model, and the effectiveness of unsupervised feature learning by the energy-based model. Furthermore, we show that the proposed method can be easily adapted to semi-supervised learning, achieving performance comparable to state-of-the-art semi-supervised methods.

Related work

For learning the energy-based model by MLE, the main difficulty lies in drawing fair samples from the current model. A prominent approximation of MLE is the contrastive divergence (CD) framework, requiring MCMC initialized from the data distribution. CD has been generalized to persistent CD , and has more recently been generalized to modified CD , adversarial CD with modern CNN structure. scale up sampling-based methods to large image datasets with white noise as the starting point of sampling. However, these sampling based methods may still have difficulty traversing different modes of the learned model, which may result in biased model, and may take a long time to converge. Another variant is to An advantage of noise contrastive estimation (NCE), and our adaptive version of it, is that it avoids MCMC sampling in estimation of the energy-based model, by turning the estimation problem into a classification problem.

Generalizing from , developed an introspective parameter estimation method, where the EBM is discriminatively learned and composed of a sequence of discriminative models obtained through the learning process. Another line of work is to estimate the parameters of EBM by score matching . connects GAN to the estimation of EBM.

NCE and it variants has gained popularity in natural language processing (NLP) . applied NCE to log-bilinear models and in NCE is applied to neural probabilistic language models. NCE shows effectiveness in typical NLP tasks such as word embeddings and order embeddings .

In the context of inverse reinforcement learning, proposes a guided policy search method, and connects it to GAN. Our method is closely related to this method, where the energy function can be viewed as the cost function, and the flow model can be viewed as the unrolled policy.

Learning method

Let $x$ be the input variable, such as an image. We use $p_{\theta}(x)$ to denote a model’s probability density function of $x$ with parameter $\theta$. The energy-based model (EBM) is defined as follows:

where $f_{\theta}(x)$ is defined by a bottom-up convolutional neural network whose parameters are denoted by $\theta$. The normalizing constant $Z(\theta)=\int\exp[f_{\theta}(x)]dx$ is intractable to compute exactly for high-dimensional $x$.

The energy-based model in eqn. 1 can be estimated from unlabeled data by maximum likelihood estimation (MLE). Suppose we observe training examples $\{x_{i},i=1,...,n\}$ from unknown true distribution $p_{\rm data}(x)$. We can view this dataset as forming empirical data distribution, and thus expectation with respect to $p_{\rm data}(x)$ can be approximated by averaging over the training examples. In MLE, we seek to maximize the log-likelihood function

Maximizing the log-likelihood function is equivalent to minimizing the Kullback-Leibler divergence ${\rm KL}(p_{\rm data}||p_{\theta})$ for large $n$. Its gradient can be written as:

which is the difference between the expectations of the gradient of $f_{\theta}(x)$ under $p_{\rm data}$ and $p_{\theta}$ respectively. The expectations can be approximated by averaging over the observed examples and synthesized samples generated from the current model $p_{\theta}(x)$ respectively. The difficulty lies in the fact that sampling from $p_{\theta}(x)$ requires MCMC such as Hamiltonian monte carlo or Langevin dynamics , which may take a long time to converge, especially on high dimensional and multi-modal space such as image space.

The MLE of $p_{\theta}(x)$ seeks to cover all the models of $p_{\rm data}(x)$. Given the flexibility of model form of $f_{\theta}(x)$, the MLE of $p_{\theta}(x)$ has the chance to approximate $p_{\rm data}(x)$ reasonably well.

1.2 Noise contrastive estimation

which transforms estimation of EBM into a classification problem.

The objective function connects to logistic regression in supervised learning in the following sense. Suppose for each training or generated examples we assign a binary class label $y$: $y=1$ if $x$ is from training dataset and $y=0$ if $x$ is generated from $q(x)$. In logistic regression, the posterior probabilities of classes given the data $x$ are estimated. As the data distribution $p_{\rm data}(x)$ is unknown, the class-conditional probability $p(\cdot|y=1)$ is modeled with $p_{\theta}(x)$. And $p(\cdot|y=0)$ is modeled by $q(x)$. Suppose we assume equal probabilities for the two class labels, i.e., $p(y=1)=p(y=0)=0.5$. Then we obtain the posterior probabilities:

The class-labels $y$ are Bernoulli-distributed, so that the log-likelihood of the parameter $\theta$ becomes

which is, up to a factor of $1/n$, an approximation of eqn. 4.

The choice of the noise distribution $q(x)$ is a design issue. Generally speaking, we expect $q(x)$ to satisfy the following: (1) analytically tractable expression of normalized density; (2) easy to draw samples from; (3) close to data distribution. In practice, (3) is important for learning a model over high dimensional data. If $q(x)$ is not close to the data distribution, the classification problem would be too easy and would not require $p_{\theta}$ to learn much about the modality of the data.

2 Flow-based model

where $q_{0}$ is a known noise distribution. $g_{\alpha}$ is a composition of a sequence of invertible transformations where the log-determinants of the Jacobians of the transformations can be explicitly obtained. $\alpha$ denotes the parameters. Let $q_{\alpha}(x)$ be the probability density of the model given a datapoint $x$ with parameter $\alpha$. Then under the change of variables $q_{\alpha}(x)$ can be expressed as

More specifically, suppose $g_{\alpha}$ is composed of a sequence of transformations $g_{\alpha}=g_{\alpha_{1}}\circ\cdots\circ g_{\alpha_{m}}$. The relation between $z$ and $x$ can be written as $z\leftrightarrow h_{1}\leftrightarrow\cdots\leftrightarrow h_{m-1}\leftrightarrow x$. And thus we have

where we define $z\coloneqq h_{0}$ and $x\coloneqq h_{m}$ for conciseness. With carefully designed transformations, as explored in flow-based methods, the determinant of the Jacobian matrix $(\partial h_{i-1}/\partial h_{i})$ can be incredibly simple to compute. The key idea is to choose transformations whose Jacobian is a triangle matrix, so that the determinant becomes

The following are the two scenarios for estimating $q_{\alpha}$:

(2) Variational approximation to an unnormalized target density $p$ , based on $\min_{\alpha}{\rm KL}(q_{\alpha}\|p)$, where

${\rm KL}(q_{\alpha}\|p)$ is the difference between energy and entropy, i.e., we want $q_{\alpha}$ to have low energy but high entropy. ${\rm KL}(q_{\alpha}\|p)$ can be calculated without inversion of $g_{\alpha}$.

When $q_{\alpha}$ appears on the right of KL-divergence, as in (1), it is forced to cover most of the modes of $p_{\rm data}$, When $q_{\alpha}$ appears on the left of KL-divergence, as in (2), it tends to chase the major modes of $p$ while ignoring the minor modes . As shown in the following section, our proposed method learns a flow model by combining (1) and (2).

3 Flow Contrastive Estimation

A natural improvement to NCE is to transform the noise so that the resulting distribution is closer to the data distribution. This is exactly what the flow model achieves. That is, a flow model transform a known noise distribution $q_{0}(z)$ by a composition of a sequence of invertible transformations $g_{\alpha}(\cdot)$. It also fulfills (1) and (2) of the requirements of NCE. However, in practice, we find that a pre-trained $q_{\alpha}(x)$, such as learned by MLE, is not strong enough for learning an EBM $p_{\theta}(x)$ because the synthesized data from the MLE of $q_{\alpha}(x)$ can still be easily distinguished from the real data by an EBM. Thus, we propose to iteratively train the EBM and flow model, in which case the flow model is adaptively adjusted to become a stronger contrast distribution or a stronger training opponent for EBM. This is achieved by a parameter estimation scheme similar to GAN, where $p_{\theta}(x)$ and $q_{\alpha}(x)$ play a minimax game with a unified value function: $\min_{\alpha}\max_{\theta}V(\theta,\alpha)$,

The objective function can be interpreted from the following perspectives:

(2) Minimization of Jensen-Shannon divergence for the flow model. If $p_{\theta}(x)$ is close to the data distribution, then the update of $\alpha$ is approximately minimizing the Jensen-Shannon divergence between the flow model $q_{\alpha}$ and data distribution $p_{\rm data}$:

(3) Connection with GAN. Our parameter estimation scheme is closely related to GAN. In GAN, the discriminator $D$ and generator $G$ play a minimax game: $\min_{G}\max_{D}V(G,D)$,

The discriminator $D(x)$ is learning the probability ratio $p_{\rm data}(x)/(p_{\rm data}(x)+p_{G}(x))$, which is about the difference between $p_{\rm data}$ and $p_{G}$ . In the end, if the generator $G$ learns to perfectly replicate $p_{\rm data}$, then the discriminator $D$ ends up with a random guess. However, in our method, the ratio is explicitly modeled by $p_{\theta}$ and $q_{\alpha}$. $p_{\theta}$ must contain all the learned knowledge in $q_{\alpha}$, in addition to the difference between $p_{\rm data}$ and $q_{\alpha}$. In the end, we learn two explicit probability distributions $p_{\theta}$ and $q_{\alpha}$ as approximations to $p_{\rm data}$.

Henceforth we simply refer to the proposed method as flow constrastive estimation, or FCE.

4 Semi-supervised learning

A class-conditional energy-based model can be transformed into a discriminative model in the following sense. Suppose there are $K$ categories $k=1,...,K$, and the model learns a distinct density $p_{\theta_{k}}(x)$ for each $k$. The networks $f_{\theta_{k}}(x)$ for $k=1,...,K$ may share common lower layers, but with different top layers. Let $\rho_{k}$ be the prior probability of category $k$, for $k=1,...,K$. Then the posterior probability for classifying $x$ to the category $k$ is a softmax multi-class classifier

where $b_{k}=\log(\rho_{k})-\log Z(\theta_{k})$.

Given this correspondence, we can modify FCE to do semi-supervised learning. Specifically, assume $\{(x_{i},y_{i}),i=1,...,m\}$ are observed examples with labels known, and $\{x_{i},i=m+1,...,m+n\}$ are observed unlabeled examples. For each category $k$, we can assume that class-conditional EBM is in the form

where $f_{\theta_{k}}(x)$ share all the weights except for the top layer. And we assume equal prior probability for each category. Let $\theta$ denotes all the parameters from class-conditional EBMs $\{\theta_{k},k=1,...,K\}$. For labeled examples, we can maximize the conditional posterior probability of label $y$, given $x$ and the fact that $x$ is an observed example (instead of a generated example from $q_{\alpha}$). By Bayes rule, this leads to maximizing the following objective function over $\theta$:

which is similar to a classifier in the form.

For unlabeled examples, the probability can be defined by an unconditional EBM, which is in the form of a mixture model:

Together with the generated examples from $q_{\alpha}(x)$, we can define the same value function $V(\theta,\alpha)$ as eqn. 12 for the unlabeled examples. The joint estimation algorithm alternate the following two steps: (1) update $\theta$ by $\max_{\theta}L_{\rm label}(\theta)+V(\theta,\alpha)$; (2) update $\alpha$ by $\min_{\alpha}V(\theta,\alpha)$. Due to the flexibility of EBM, $f_{\theta_{k}}(x)$ can be defined by any existing state-of-the-art network structures designed for semi-supervised learning.

Experiments

For FCE, we adaptively adjust the numbers of updates for EBM and Glow: we first update EBM for a few iterations until the classification accuracy is above $0.5$, and then we update Glow until the classification accuracy is below $0.5$. We use Adam with learning rate $\alpha=0.0003$ for the EBM and Adamax with learning rate $\alpha=0.00001$ for the Glow model. Code and more results can be found at http://www.stat.ucla.edu/~ruiqigao/fce/main.html

Figure 1 demonstrates the results of FCE on several 2D distributions, where FCE starts from a randomly initialized Glow. The learned EBM can fit multi-modal distributions accurately, and forms a better fit than Glow learned by either FCE or MLE. Notably, the EBM is defined by a much simpler network structure than Glow: for Glow we use $10$ affine coupling layers, which amount to $30$ fully-connected layers, while the energy-based model is defined by a $4$-layer fully-connected network with the same width as Glow. Another interesting finding is that the EBM can fit the distributions well, even if the flow model is not a perfect contrastive distribution.

For the distribution depicted in the first row of Figure 1, which is a mixture of eight Gaussian distributions, we can compare the estimated densities by the learned models with the ground truth densities. Figure 2 shows the mean squared error of the estimated log-density over numbers of training iterations of EBMs. We show the results of FCE either starting from a randomly initialized Glow (’rand’) or a Glow model pre-trained by MLE (’trained’), and compare with NCE with a Gaussian noise distribution. FCE starting from a randomly initialized Glow converges in fewer iterations. Both settings of FCE achieve a lower error rate than NCE.

2 Learning on real image datasets

We conduct experiments on the Street View House Numbers (SVHN) , CIFAR-10 and CelebA datasets. We resized the CelebA images to $32\times 32$ pixels, and used $20,000$ images as a test set. We initialize FCE with a pre-trained Glow model, trained by MLE, for the sake of efficiency. We again emphasize the simplicity of the EBM model structure compared to Glow. See Appendix for detailed model architectures. For Glow, depth per level is set as $8$, $16$, $32$ for SVHN, CelebA and CIFAR-10 respectively. Figure 3 depicts synthesized examples from learned Glow models. To evaluate the fidelity of synthesized examples, Table 1 summarizes the Fréchet Inception Distance (FID) of the synthesized examples computed with the Inception V3 classifier. The fidelity is significantly improved compared to Glow trained by MLE (see Appendix for qualitative comparisons), and is competitive to the other generative models. In Table 2, we report the average negative log-likelihood (bits per dimension) on the testing sets. The log-likelihood of the learned EBM is based on the estimated normalizing constant (i.e., a parameter of the model) and should be taken with a grain of salt. For the learned Glow model, the log-likelihood of the Glow model estimated with FCE is slightly lower than the log-likelihood of the Glow model trained with MLE.

3 Unsupervised feature learning

To further explore the EBM learned with FCE, we perform unsupervised feature learning with features from a learned EBM. Specifically, we first conduct FCE on the entire training set of SVHN in an unsupervised way. Then, we extract the top layer feature maps from the learned EBM, and train a linear classifier on top of the extracted features using only a subset of the training images and their corresponding labels. Figure 4 shows the classification accuracy as a function of the number of labeled examples. Meanwhile, we compare our method with a supervised model with the same model structure as the EBM, and is trained only on the same subset of labeled examples each time. We observe that FCE outperforms the supervised model when the number of labeled examples is small (less than $2000$).

Next we try to combine features from multiple layers together. Specifically, following the same procedure outlined in , the features from the top three convolutional layers are max pooled and concatenated to form a $14,336$-dimensional vector of feature. A regularized L2-SVM is then trained on these features with a subset of training examples and the corresponding labels. Table 3 summarizes the results of using $1,000$, $2,000$ and $4,000$ labeled examples from the training set. At the top part of the table, we compare with methods that estimate an EBM or a discriminative model coupled with a generator network. At the middle part of the table, we compare with methods that learn an EBM with contrastive divergence (CD) and modified versions of CD. For fair comparison, we use the same model structure for the EBMs or discriminative models used in all the methods. The results indicate that FCE outperforms these methods in terms of the effectiveness of learned features.

4 Semi-supervised learning

In section 3.4 we show that FCE can be generalized to perform semi-supervised learning. We emphasize that for semi-supervised learning, FCE not only learns a classification boundary or a posterior label distribution $p(y|x)$. Instead, the algorithm ends up with $K$ estimated probabilistic distributions $p(x|y=k),k=1,...K$ for observed examples belonging to $K$ categories. Figure 5 illustrates this point by showing the learning process on a 2D example, where the data distribution consists of two twisted spirals belonging to two categories. Seven labeled points are provided for each category. As the training goes, the unconditional EBM $p_{\theta}(x)$ learns to capture all the modes of the data distribution, which is in the form of a mixture of class-conditional EBMs $p_{\theta_{1}}(x)$ and $p_{\theta_{2}}(x)$. Meanwhile, by maximizing the objective function $L_{\rm label}(\theta)$ (eqn. 17), $p_{\theta}(x)$ is forced to project the learned modes into different spaces, resulting in two well-separated class-conditional EBMs. As shown in Figure 5, within a single mode of one category, the EBM tends to learn a smoothly connected cluster, which is often what we desire in semi-supervised learning.

Then we test the proposed method on a dataset of real images. Following the setting in , we use two types of CNN structures (‘Conv-small’and ‘Conv-large’) for EBMs, which are commonly used in state-of-the-art semi-supervised learning methods. See Appendix for detailed model structures. We start FCE from a pre-trained Glow model. Before the joint training starts, EBMs are firstly trained for $50,000$ iterations with the Glow model fixed. In practice, this helps EBMs keep pace with the pre-trained Glow model, and equips EBMs with reasonable classification ability. We report the performance at this stage as ‘FCE-init’. Also, since virtual adversarial training (VAT) has been demonstrated as an effective regularization method for semi-supervised learning, we consider adopting it as an additional loss for learning the EBMs. More specifically, the loss is defined as the robustness of the conditional label distribution around each input data point against local purturbation. ‘FCE + VAT’ indicates the training with VAT.

Table 4 summarizes the results of semi-supervised learning on SVHN dataset. We report the mean error rates and standard deviations over three runs. All the methods listed in the table belong to the family of semi-supervised learning methods. Our method achieve competitive performance to these state-of-the-art methods. ‘FCE + VAT’ results show that the effectiveness of FCE does not overlap much with existing semi-supervised method, and thus they can be combined to further boost the performance.

Conclusion

This paper explores joint training of an energy-based model with a flow-based model, by combining the representational flexibility of the energy-based model and the computational tractability of the flow-based model. We may consider the learned energy-based model as the learned representation, while the learned flow-based model as the learned computation. This method can be considered as an adaptive version of noise contrastive estimation where the noise is transformed by a flow model to make its distribution closer to the data distribution and to make it a stronger contrast to the energy-based model. Meanwhile, the flow-based model is updated adaptively through the learning process, under the same adversarial value function.

In future work, we intend to generalize the joint training method by combining the energy-based model with other normalized probabilistic models, such as auto-regressive models. We also intend to explore other joint training methods such as those based on adversarial contrastive divergence or divergence triangle .

The work is partially supported by DARPA XAI project N66001-17-2-4029 and ARO project W911NF1810296. We thank Pavel Sountsov, Alex Alemi, Matthew D. Hoffman and Srinivas Vasudevan for their helpful discussions.

References

Appendix A Model architectures

Table 5 summarizes the EBM architectures used in unsupervised learning (subsections 4.1-4.3). The slope of all leaky ReLU (lReLU) functions are set to $0.2$. For semi-supervised learning from a 2D example (subsection 4.4), we use the same EBM structure as the one used in unsupervised learning from 2D examples, except that for the top fully connect layer, we change the number of output channels to $2$, to model EBMs of two categories respectively. Table 6 summarizes the EBM architectures used in semi-supervised learning from SVHN (subsection 4.4). After each convolutional layer, a weight normalization layer and a leaky ReLU layer is added. The slope of leaky ReLU functions is set to $0.2$. A weight normalization layer is added after the top fully connected layer.

For Glow model, we follow the setting of . The architecture has multi-scales with levels $L$. Within each level, there are $K$ flow blocks. Each block has three convolutional layers (or fully-connected layers) with a width of $W$ channels. After the first two layers, a ReLU activation is added. Table 7 summarizes the hyperparameters for different datasets.

Appendix B Synthesis comparison

In figures 6, 7 and 8, we display the synthesized examples from Glow trained by MLE and our FCE.