Conjugate Energy-Based Models
Abstract
We propose conjugate energy-based models (EBMs), a class of deep latent-variable models with a tractable posterior.
Conjugate EBMs have similar use cases as variational autoencoders, in the sense that they learn an unsupervised mapping between data and latent variables.
However these models omit a generator, which allows them to learn more ﬂexible notions of similarity between data points.
Our experiments demonstrate that conjugate
EBMs achieve competitive results in terms of image modelling, predictive power of latent space, and out-of-distribution detection on a variety of datasets.
1. Introduction
Deep generative models approximate a data distribution by combining a prior over latent variables with a neural generator, which maps latent variables to points on a data manifold.
While it is common to evaluate these models in terms of their ability to generate realistic examples, an arguably more important use case of these models is that they learn representations in absence of supervision.
To be useful in downstream tasks, these representations should encode some set of “semantically meaningful” features rather than “nuisance variables” that are unlikely to have predictive power.
Guiding a model towards a semantically meaningful representation requires some form of inductive bias.
A large body of work on variational autoencoders (VAEs, Kingma and
Welling (2013); Rezende et al.
(2014)) has explored the use of priors as inductive biases.
Relatively mild biases in the form of conditional independence are common in the literature on disentangled representations (Higgins et al., 2016;
Kim and Mnih, 2018; Chen et al., 2018; Esmaeili et al., 2019).
More generally, recent work has employed priors that reﬂect structure of the underlying data to represent objects in an image (Eslami et al., 2016;
Lin et al., 2020b; Engelcke et al., 2019; Crawford and Pineau, 2019b), or moving objects in video (Crawford and Pineau, 2019a; Kosiorek et al., 2018; Wu et al., 2020; Lin et al., 2020a).
Despite steady progress, work on disentangled representations and structured VAEs still predominantly considers synthetic data sets.
To train a VAE we minimize a reconstruction loss, which treats all pixels in an image equally.
For complex natural scenes, learning a
[Table]
[Table] [Table]
[Figure]
CEBM model that can produce pixel-perfect reconstructions poses fundamental challenges, given the combinatorial explosion of possible inputs.
This is not only a problem from the perspective of generation, but also from the perspective of the learned representation; a VAE must encode all factors of variation that give rise to large deviations in pixel
, regardless of whether these factors are semantically meaningful (e.g. presence and locations of objects) or not (e.g. shadows of objects in the background of the image).
In this paper, we consider energy-based models (EBMs) with latent variables as an alternative VAEs for learning representations in an unsupervised manner.
The general idea of using EBMs for this purpose is by no means new; it has a long history in the context of restricted Boltzmann machines (RBMs) and related models (Smolensky, 1986; Welling et al., 2004; Hinton et al., 2006).
Our motivation with the present work is to design a class of energy-based models that retain the desirable features of VAEs while addressing what we see as one of their main weaknesses: We would like incorporate inductive biases, but model the data at an intermediate level of representation that does not necessarily encode all features of an image at the pixel level.
Concretely, we propose Conjugate EBMs (CEBMs), a class of models in which the energy function deﬁnes a neural exponential family.
While the normalizer of this family is intractable, we can compute its posterior in closed form when we pair the likelihood with an appropriate conjugate bias term in the energy function.
As a result, the neural suﬃcient statistics in a CEBM fully determine both the marginal likelihood and the encoder, hereby side-stepping the need for a generator.
In our experiments, we evaluate the representations learned by CEBM using class labels as a proxy for the primary factors of variation in a dataset.
We show that CEBMs learn a notion of similarity that aligns more closely with class labels in terms of the nearest neighbors in latent space (Figure 1).
Moreover, we show that the representations learned by CEBMs (in an unsupervised manner) can achieve a competitive performance in classiﬁcation and out-of-distribution detection tasks.
2. Background 2.1. Energy-Based Models
An EBM (LeCun et al., 2006) deﬁnes a density p θ ( x ) =
exp {− E θ ( x ) } /Z
θ in terms of an energy function E θ ( x ) : R D −→ R with parameters θ , which maps inputs x ∈ R D to a scalar value.
The density p θ ( x ) can only be evaluated up to an unknown constant of proportionality since the normalizing constant Z θ =
(cid:82) dx exp {− E θ ( x ) } is typically intractable.
To train an
EBM, we approximate the gradient of the expected log-likelihood
[Equation]
Estimating the gradient of the log normalizer ∇ θ log Z θ requires samples from the model x (cid:48) ∼ p θ ( x (cid:48) ), which in turn requires approximate inference since this density is intractable.
A commonly used method is Stochastic Gradient Langevin Dynamics (SGLD, Welling and
Teh (2011)), which initializes a sample x (cid:48) 0 ∼ p 0 ( x (cid:48) )
and then performs a sequence of gradient updates with additional injected noise (cid:15)
[Equation]
Recent work has shown that EBMs with convolutional energy functions can accurately model distributions over images, in the sense that SGLD produces realistic samples (Nijkamp et al., 2019a,b; Du and Mordatch, 2019; Xie et al., 2016).
2.2. Conjugate Exponential Families
An exponential family is a set of distributions whose probability density function or probability mass function can be expressed in the form
[Equation]
where h ( · ) is base measure, η is a vector of natural parameters , t ( · ) is a vector of suﬃcient statistics, and A ( · ) is a log normalizer .
If a likelihood belongs to an exponential family, then there exists a conjugate prior with the form p ( η | λ )
=
h 0 ( η ) exp (cid:8) t 0 ( η ) (cid:62) λ − A 0 ( λ ) (cid:9) .
A prior is conjugate to a likelihood when its vector of suﬃcient statistics comprises the natural parameters and the log-normalizer of the likelihood t 0 ( η ) = (cid:2) η, − A ( η ) (cid:3) , λ = (cid:2)
λ 1 , λ 2 (cid:3) .
The convenient property of conjugate exponential families is that both the marginal likelihood p ( x | λ ) and the posterior p ( η | x, λ ) are tractable.
The reason is that the joint probability has the form
[Equation]
If we substitute ˜ λ
1 = λ 1 + t ( x ) and ˜ λ
2 = λ 2 + 1, we can equivalently factorize this joint as
[Equation]
This shows that the posterior is in the same exponential family as the prior, and that we can express the marginal likelihood using the log normalizer A 0 ( · )
[Equation]
| x
[Equation]
λ )
[Equation]
p
[Equation]
η
[Equation]
3. Conjugate Energy-Based Models
In this paper we are interested in learning a probabilistic model that deﬁnes a joint distribution p θ ( x, z ) over high-dimensional data x ∈ R D and a lower-dimensional set of latent variables z ∈ R K .
Because most work on deep generative models has focused on images, we will restrict ourselves to this data modality.
We propose to consider models of the form
[Equation]
In the energy function, η : R K → R H decodes latent variables to a vector of natural parameters in an intermediate space of dimension H .
The function t θ : R D → R H plays the role of an encoder, which maps data to suﬃcient statistics in the same space as η ( z ).
The function b : R K → R serves as an inductive bias that plays a role analogous to the prior.
We will consider a bias b ( z ) = log p ( z | λ ) in form of a tractable exponential family
[Equation]
We can then express the energy function as
[Equation]
This form of the energy function has a very convenient property: It corresponds to a model p θ ( x, z ) in which the posterior p θ ( z | x ) is tractable.
To see this, we can make a substitution ˜ λ
= λ + t θ ( x ) analogous to the one in Equation 3, which allows us to express the energy as
[Equation]
We now see that we can factorize the he corresponding density p θ ( x, z ) as
[Equation]
which yields a posterior and a marginal that are analogous to the distributions in Equation
[Equation]
This posterior is conjugate, in the sense that it is in the same tractable family as the bias.
In addition to having a tractable posterior, CEBMs have the convenient property that the marginal likelihood p θ ( x | λ ) itself can be expressed as an energy-based model that is deﬁned in terms of the log normalizer A ( · ) and the encoder network t θ ( x ).
This means that we can train CEBMs using SGLD in the same way as other EBMs.
CEBMs diﬀer from VAEs in that they lack a generator network.
Instead, the density is fully speciﬁed by the encoder network t θ ( x ), which deﬁnes a notion of agreement ( λ + t θ ( x )) (cid:62) η ( z ) between data and latent variables in an intermediate feature space.
In VAEs, by contrast we would assume that the suﬃcient statistics t ( x ) are known, and learn a generator η θ ( z ) to compute a notion of agreement in data space (see Appendix B).
Note that when z is a categorical variable, the suﬃcient statistics t θ ( x ) deﬁne a classiﬁer, and CEBMs are equivalent to the models that are considered in recent work on EBMs for classiﬁcation (Grathwohl et al., 2019; Liu and Abbeel, 2020).
[Figure]
CIFAR
[Figure] [Figure]
[Figure] [Figure]
100 full  100 full
[Figure]
36 GMM-
[Figure]
{+IGEBM               49 CEBM        19     70 CEBMM 67  
+}
4. Inductive Biases
CEBMs have a property that is somewhat counter-intuitive.
While the posterior p θ ( z | x, λ ) is tractable, the prior p θ ( z ) is in general not tractable.
In particular, although the bias b θ ( z ) is the logarithm of a tractable exponential family, it is not the case that p θ ( z ) (cid:54) =
exp { b θ ( z ) } .
Rather the prior p θ ( z ) has the form,
[Equation]
In principle the bias a CEBM can take the form of any exponential family distribution.
Since products of exponential families are also in the exponential family, this covers a broad range of possible biases.
In this paper, we will constrain ourselves to a Spherical Gaussian and a Mixture of Gaussians.
We provide derivations for both cases in Appendix A.
5. Experiments
Our experiments evaluate to what extent CEBMs can learn representations that encode meaningful factors of variation, whilst discarding details about the input that we would consider noise.
We train using both inductive biases and will refer to them as CEBM (Spherical Gaussian) and CEBMM (Mixutre of Gaussian) for the rest of Section.
See
Appendix C and D for architecture and training details.
5.1. Samples and Latent Space
We begin with a qualitative evaluation by visualizing samples.
Figures 2 shows samples from CEBMs with uniform noise initialization an 500 SGLD steps.
To assess to what extent the representation in CEBMs aligns with classes in each dataset, we look at the agreement between the label for each data point and the label of its nearest neighbor in the latent space.
Figure 1 shows that the distance in pixel space is a poor measure of similarity in this dataset, whereas proximity in the latent space is more likely to agree with class labels in both VAEs and CEBMs.
Quantitatively, in Appendix F, we observe a stronger alignment between classes and the latent representation in CEBMs, which is reﬂected in higher numbers on the diagonal of the matrix.
5.2. Classiﬁcation
To evaluate performance in settings where few labels are available, we train a logistic classiﬁer using 1 , 10 , 100 examples per class, as well as the full training dataset.
We compare
CEBMs against the IGEBM (Du and Mordatch, 2019), a standard VAE with the spherical Gaussian prior, and the GMM-VAE (Tomczak and Welling, 2018) where the prior is a mixture of Gaussians.
As discussed in Section 2, IGEBM does not have an explicit representation.
In order to compare against IGEBM, we remove the last layer (which outputs the energy) and use the resulting intermediate representation as the latent code.
We report the classiﬁcation accuracy on the test set in Table 1.
We can see that that
CEBMs overall achieve a higher accuracy compared to VAEs in particular for CIFAR-10 and
SVHN where the pixel distance is not good measure for similarity.
Moreover, we observe that CEBMs outperform IGEBM which suggest that the inductive biases in CEBMs can lead to increased performance in downstream tasks.
5.3. Out-of-Distribution Detection
EBMs have formed the basis for encouraging results in out-of-distribution (OOD) detection (Du and Mordatch, 2019; Grathwohl et al., 2019).
In Appendix G, we report the area under the receiver-operator curve (AUROC) using two score functions: log p θ ( x ) and a gradient-based function proposed by Grathwohl et al.
CEBMs results for OOD detection in most cases improve upon VAE and IGEBM baselines.
6. Conclusion
We introduced CEBMs; a new family of energy-based models that deﬁne a joint energy function over both the data and latent variables.
The joint distribution factorizes into a tractable posterior and a marginal likelihood, imposing an inductive bias on the latent space.
This factorization allows us to directly optimize the marginal likelihood of the data, while at the same time imposing an inductive bias on the latent space.
Experimental results for this class of models are encouraging; we observe a closer agreement between unsupervised representations and class labels, which translates into improvements in downstream classiﬁcation tasks.
Acknowledgments
We would like to thank our reviewers for their thoughtful comments.
This work was supported by the Intel Corporation, the 3M Corporation, NSF award 1835309, startup funds from Northeastern University, the Air Force Research Laboratory (AFRL), and DARPA.
Appendix A. Derivation for Two Cases of Inductive Biases 1. Spherical Gaussian.
As a bias that is analogous to the standard prior in VAEs, we consider a spherical Gaussian with ﬁxed hyperparameters ( µ, σ ) = (0 , 1) for each dimension of z ∈ R K ,
[Equation]
Each term has suﬃcient statistics η ( z k )
=
( z k , z 2 k ) and natural parameters
[Equation]
The marginal likelihood of the CEBM is then
[Equation]
where ˜ λ
k = λ + t θ,k ( x ) and the log normalizer is
[Equation]
2. Mixture of Gaussians.
In our experiments we will consider datasets that are normally used for classiﬁcation.
These datasets, by design, exhibit multimodal structure that we would like to see reﬂected in the learned representation.
As an inductive bias that is amenable to uncovering this structure, we will consider a bias in the form of a mixture of
L Gaussians,
[Equation]
Here z ∈ R K is a vector of features and y ∈ { 1 , . . .
, L } is a categorical assignment variable.
The bias for each component l is a spherical Gaussian with hyperparameters λ l,k for each dimension k .
Again using the notation ˜ λ
l,k = λ l,k + t θ,l,k ( x ) to refer to the posterior parameters, then we obtain an energy
[Equation]
We can then deﬁne a joint probability over data x and the assignment y in terms the log normalizer A ( · ),
[Equation]
which then allows us to compute the marginal
[Equation]
We optimize this marginal with respect hyperaparameters λ as well as the weights θ .
Appendix B. Relationship between CEBMs and VAEs B.1. Variational Autoencoders
Variational autoencoders are a widely used class of deep generative models Kingma and
Welling (2013); Rezende et al.
A VAE deﬁnes a joint distribution p θ ( x, z ) over data x
and latent variables z ; it combines an unstructured prior (e.g. a spherical Gaussian) with a likelihood that is parameterized by an expressive neural network, often referred to as a decoder.
An inference model, also known as an encoder, approximates the posterior p θ ( z | x ) by mapping each data point x onto latent variables z .
These models are trained by maximizing the stochastic evidence lower bound (ELBO) deﬁned as
[Equation]
When the q φ ( z | x ) is reparameterizable, we can compute Monte Carlo estimates of the gradient of this objective using pathwise derivatives.
Non-reparameterizable cases, such as models with discrete variables, require likelihood-ratio estimators Williams (1992).
Despite their successes, VAEs have limitations.
By maximizing the ELBO, we favor encoder-decoder pair that perfectly reconstruct their input, so there is nothing preventing the VAE from mapping similar inputs to similar encoding, even when they might be semantically diﬀerent.
Likewise, examples that might be very dissimilar in pixel space because of noise or benign transformations might end up with very diﬀerent latent representations.
B.2. VAEs as Energy-based Models
We can interpret the generative model in a VAE as a model with an energy function
[Equation]
In this setting, η θ ( z ) is the generator network that maps low-dimensional latent variables to a high-dimensional vector of natural parameters.
The function t ( x ) is a known mapping of data to the suﬃcient statistics of a Gaussian or Bernoulli likelihood.
The bias b θ ( z ) contains the terms in the log density log p θ ( x, z ) that only depend on z .
In a VAE there are two such terms.
The ﬁrst is the log prior log p θ ( z ).
The second is the log normalizer of A ( η θ ( z )) for the likelihood log p θ ( x | z ), which can be computed in closed form because Gaussian and Bernoulli distributions are in the exponential family.
Combining these terms yields an energy function for the bias,
[Equation]
[Equation]
In other words, in a VAE we use known suﬃcient statistics t ( x ), and train a generator to learn the natural parameters
η θ ( z ).
In a CEBM, by contrast, we assume known natural parameters η ( z ) and train an encoder to learn the suﬃcient statistics t θ ( x ).
Appendix C. Model Architectures
CEBMs employ an encoder network t θ ( x ) in the form of 4-layer CNN (which is proposed by Nijkamp et al. (2019a)), followed by an MLP output layer.
For IGEBMs, we add one extra MLP as its ﬁnal layer which outputs a scalar value.
VAEs use the same encoder network as the CEBMs, and use a decoder network in form of an MLP followed by 4-layer
CNN.
Table Table C, Table Table C
[Table]
and Table Table C show the architectures used for CEBM,
VAE
[Table]
and IGEBM, respectively.
[Table] [Table]
[Table]
Appendix D. Training Details of Persistent Contrastive Divergence Optimization.
In CEBMs and VAEs, we choose the dimension of latent variables to be 128.
For CEBMS, We found that the optimization becomes diﬃcult with smaller dimensions.
We L2 regularize energy magnitudes (proposed by Du and Mordatch (2019))
[Table]
where the coeﬃcient of the L2 regularization term is 0.1.
We empirically found that the training would become unstable without this regularization.
We train our models using 60 SGLD steps where we initialize samples from the replay buﬀer with 0.95 probability, and initialize from uniform noise with 0.05 probability.
We train all the models with 90k gradient steps, batch size 128, Adam optimizer with learning rate 1e-4.
When doing PCD, we used a reply buﬀer of size 5000.
We set the α in the SGLD teps to be 0.075.
Similar to Du and Mordatch (2019), we found it useful to add some noise to the image before encoding.
In our experiments, we used Gaussian noise with σ 2 = 0 .
For the mixture models (CEBMM and GMM-VAE), we used 50 mixtures.
Training Stability
As observed in previous work Du and Mordatch (2019)
; Grathwohl et al.
(2019), training EBMs can be a challenging task that often requires a thorough hyperparameter search.
We found that the choices of activation function, learning rate, number of SGLD steps, and regularization will all aﬀect training stability.
Models regularly diverge during training, and it is diﬃcult to perform diagnostics given that log p θ ( x ) cannot be computed.
As suggested by Nijkamp et al. (2019a), we found checking the diﬀerence in energy between data and model samples to be helpful for verifying stable training.
We also note that in general, we observed a trade-oﬀ between sample quality and the predictive power of latent variables in our experiments.
We leave investigation of the source of this trade-oﬀ to future work, but we suspect that this is because SGLD is having more diﬃculty to converge when the latent space is more disjoint.
Appendix E. Persistent Constrastive Divergence
− θ
Appendix F. Confusion Matrices on 1-NN Classiﬁcation
We perform 1-nearest-neighbor classiﬁcation task for MNIST, Fashion-MNIST, SVHN, CIFAR10.
We compute the L2 distance in the latent space of VAE, IGEBM and CEBM
[Table]
and also in pixel space.
We visualize the confusion matrices
[Table]
[Figure]
[Figure] [Table] [Figure]
[Figure] [Figure]
[Figure]
[Table]
[Figure]
CEBM
[Table]
[Figure]
Appendix G. Out-of-Distribution Detection Table
[Figure]
[Table] [Table]
[Table]
[Table] [Table]