N O C ONDITIONAL M ODELS FOR ME : T RAINING J OINT EBM S ON M IXED C ONTINUOUS AND D ISCRETE D ATA
A BSTRACT
We propose energy-based models of the joint distribution of data and supervision.
While challenging to work with, this approach gives ﬂexibility when designing energy functions and easy parameterization for structured supervision.
Further, these models naturally allow training on partially observed data and predictions conditioned on any subset of the modeled variables.
We identify and address the main difﬁculty in working with these models: sampling from the joint distribution of data and supervision.
We build upon recent developments in discrete
MCMC sampling and apply them alongside continuous MCMC techniques developed for energy-based models.
We present experimental results demonstrating that our proposed approach can successfully train joint energy-based models on high-dimensional data with structured supervision capable of both accurate prediction and conditional sampling.
I NTRODUCTION
The ﬂexibility of Energy-Based Models (EBMs) make them an appealing approach for a number of generative modelling tasks.
The ease with which structure may be incorporated and their relationship to physical systems has long made them popular in the broader scientiﬁc community (Ingraham et al., 2019; Du et al., 2020c;
No ´ e et al., 2019).
After a period of reduced interest from the machine learning community, EBMs have regained popularity as an approach for generative modeling (Nijkamp et al., 2020a; Du & Mordatch, 2019) and have found utility in many applications.
One such application is Joint Energy-based Models (JEM) (Grathwohl et al., 2019).
These models re-purpose existing state-of-the-art classiﬁer architectures to deﬁne an energy-based model of the joint distribution p ( x, y ) between data and class labels, instead of the conditional distribution of labels given data p ( y | x ) as is typically used for classiﬁcation.
Moving from conditional to joint models gives improved calibration, out-of-distribution detection, and adversarial robustness while retaining high discriminative performance.
This approach has been extended to domains beyond images (He et al., 2021;
Hataya et al., 2021) and to new applications such as semi-supervised learning (Grathwohl et al., 2020; Zhao et al., 2020).
Grathwohl et al.
(2019) focus on models for image classiﬁcation where x is a continuous images and y is a discrete, 1-dimensional class label.
Naively training this model would require MCMC sampling from the joint distribution of x and y which is challenging because y is discrete.
Thankfully, the structure of the energy-function allows us to marginalize out the label y to compute an EBM for log p ( x ) and a normalized p ( y | x ) model.
We can then train to maximize the factorized likelihood log p ( x, y )
= log p ( x ) + log p ( y | x ) .
Techniques for training EBMs on continuous data are used to maximize the ﬁrst term and the second term is optimized to minimize cross-entropy.
In many situations we are interested in supervision which is more complicated than a 1-dimensional class label.
We may, for example, wish to predict a segmentation mask, or a text caption given an image.
Here the supervisory information is rich, discrete, high dimensional, and its variables have correlations which are important to model.
In settings such as these we may not be able to marginalize out the y or compute a normalized p ( y | x ) .
Thus, we will not be able to train in the same way as Grathwohl et al.
An alternative approach to train this model would be to maximize log p ( x, y ) directly.
This would require sampling from the joint directly which can be challenging given the mixed continuousdiscrete nature of the sampling problem.
Notably though, solving this sampling task would enable the training of models simultaneously capable of discriminative modeling, conditional sampling, and semi-supervised learning.
In this work we take a ﬁrst step to address these issues and demonstrate that joint EBMs can be trained on high-dimensional mixed continuous and and discrete data.
These models generate convincing conditional samples and make accurate predictions.
To accomplish this, we build upon recent work on MCMC sampling in discrete spaces (Grathwohl et al., 2021) and MCMC-based
EBM training (Nijkamp et al., 2020a; Du & Mordatch, 2019).
In the next sections we provide some background on EBMs and MCMC sampling.
We then present our approach to train and sample from joint continuous-discrete EBMs.
Finally, we present results training these models on some small but exemplary datasets.
B ACKGROUND
In the following sections we refer to our input data x ∈ X and supervision y ∈ Y .
We assume that x is continuous and y is discrete (either binary or categorical).
2. E NERGY -B ASED M ODELS
An energy-based model is any model which parameterizes a probability distribution as
[Equation]
where f θ : R D
→ R fully speciﬁes the model and Z ( θ ) = (cid:82) e f θ ( x ) d x is the normalizing constant.
While the ﬂexibility of EBMs makes them appealing, this ﬂexibility comes with the cost of making sampling and likelihood evaluation much more challenging than it is for more restricted model classes.
In fact, likelihoods cannot be exactly computed, or even lower-bounded.
EBMs are typically trained with gradient-based optimization using the following estimator for the gradient of the maximum likelihood objective:
[Equation]
Use of this estimator requires generating samples from p θ ( x ) .
Since exact sampling is intractable, we resort to using approximate samples generated with MCMC (Tieleman, 2008).
Fortunately, a host of techniques and tricks have been developed to make training with MCMC efﬁcient when using deep neural networks to deﬁne the energy-function (Du & Mordatch, 2019; Du et al., 2020b;
Nijkamp et al., 2019; 2020a; Xie et al., 2016).
The preferred sampling approach for continuous data is Langevin Dynamics (Welling & Teh, 2011) which updates samples with
[Equation]
where the step-size (cid:15) , and the temperature λ are hyperparameters.
Common values for image data are (cid:15) = 0 .
01 , λ = 1 20 , 000 .
The low temperature is necessary to generate samples quickly enough for efﬁcient training.
2. G IBBS -W ITH -G RADIENTS
Gradient-based MCMC methods such as Hamiltonian Monte-Carlo (Neal et al., 2011) or Langevin
Dynamics (Equation 3) have enabled the training of deep EBMs on continuous data.
Similar progress has not been made on discrete data because a lack of scalable MCMC methods for discrete distributions.
In the following we discuss sampling from p ( y ) =
e f ( y ) Z where y ∈ Y is discrete.
Recent work (Grathwohl et al., 2021) presents an MCMC sampler for discrete distributions which has enabled the training of deep EBMs on discrete data.
The sampler, Gibbs-With-Gradients, is based on the Metropolis-Hastings algorithm and deﬁnes a proposal distribution, q ( y (cid:48) | y ) , which makes local updates y (cid:48)
∈ H ( y ) (where H ( y ) is the Hamming-ball of size 1 around x ).
Given the current state of a sampling chain, y t , the sampler proposes an update by sampling y (cid:48) ∼ q ( y (cid:48) | y t ) .
This proposed update is accepted with probability min { exp( f ( y (cid:48) )
− f ( y ))
q ( y | y (cid:48) )
q ( y (cid:48) | y ) , 1 } .
If the sample is accepted we set y t +
=
y (cid:48) , otherwise we set y t
+1 = y t .
Recent theoretical work on MCMC (Zanella, 2020) shows that the following proposal:
[Equation]
is near optimal for MCMC in these settings.
Unfortunately, sampling from this proposal and computing likelihoods requires the computation of f ( y (cid:48) )
− f ( y ) for each y (cid:48) ∈ H ( y ) which requires
O ( D ) function evaluations for D -dimensional data.
Grathwohl et al.
(2021) noticed that discrete distributions are often implemented as continuous functions of real-valued inputs.
In these settings the discrete structure is created by restricting the domain to a discrete subset of the possible inputs.
A prime example is the Ising model log p ( y ) =
y T Wy + b T y − log Z .
The log-probability function can accept inputs in R D but the distribution is deﬁned on y ∈ { 0 , 1 } D .
In these settings we can utilize a Taylor-series approximation
[Equation]
Using this approximation, we can approximate the proposal of Equation 4 using only O (1) function evaluations.
Despite these approximations, Grathwohl et al.
(2021) demonstrate that this approach to sampling (when it can be applied) greatly outperforms other generic discrete MCMC samplers such as Gibbs sampling or the Hamming-Ball Sampler (Titsias & Yau, 2017).
J OINT E NERGY -B ASED M ODELS
We wish to train an EBM for the joint distribution of data x ∈ X and supervision y ∈ Y where the supervision may be high dimensional and structured.
Our models take the form
[Equation]
where f θ : X × Y → R is a ﬂexible parametric function.
Examples could include:
• f θ ( x, y )
=
neural net ( concat ( x, y ); θ ) • f θ ( x, y )
=
g θ ( x ) T h θ ( y ) where g θ : X → R D and h θ : Y → R D
Importantly, models of this form can represent correlations between the variables of y and can deﬁne distributions over highly-structured y .
For example, if y is permutation invariant or graph-structured, a Deep Set (Zaheer et al., 2017) or Graph Convolutional network (Zhang et al., 2019), respectively can be used to parameterize the energy function deﬁning a model which respects the invariance structure in the data.
Working with models of this form can be challenging because both conditionals p ( x | y ) , p ( y | x ) and both marginals p ( x ) , p ( y ) may be unnormalized.
Intriguingly though, if we are able to sample from the joint then we will easily be able to maximize likelihood of any marginal and make predictions of any chosen variables given any subset of the remaining variables.
Thus in this setting, many distinct problems that have been addressed in notably different ways are reduced to solving a single, albeit challenging, problem.
To make predictions of y given x we can use MCMC to draw samples from p ( y | x ) .
This distribution has the form
[Equation]
where C ( x ) is a constant that does not depend on y .
Thus, p ( y | x ) is an EBM with unnormalized log probability given by evaluating f θ ( x, y ) with x ﬁxed to the conditioning value.
Symmetrically, we can apply the same approach to sample from p ( x | y ) .
Given our joint model, we can sample from any conditional by ﬁxing the observed variables and evaluating f θ ( x, y ) (and its gradient).
A similar argument allows us to train on partially-observed data.
Assume we have observed a sample x (cid:48) with no corresponding y .
We can estimate the marginal likelihood’s gradient with
[Equation]
Thus, dealing with unnormalized joint models, making predictions, training on fully-observed data, training on partially-observed data, and drawing joint samples are roughly equivalent problems.
Of course, joint sampling from high-dimensional, unnormalized distributions is an incredibly difﬁcult task but recent advances in gradient-based sampling and EBMs (Welling & Teh, 2011; Grathwohl et al., 2021; Du et al., 2020c; Grathwohl et al., 2019) have demonstrated that accurate samples can be generated from large-scale EBMs.
T RAINING
In this section we outline our simple approach to sample from p θ ( x, y ) where x is continuous and y is discrete.
Deﬁning efﬁcient MCMC samplers for such distributons is still an open problem.
When y is held ﬁxed, we ﬁnd that p θ ( x | y ) is simply an EBM deﬁned on continuous x .
In this setting
Langevin Dynamics has been successfully applied to generate samples.
With x held ﬁxed, we ﬁnd p θ ( y | x ) is an EBM deﬁned on discrete variables and Gibbs-With-Gradients (Grathwohl et al., 2021) may be applied to generate samples.
To sample from p θ ( x, y ) we can apply block Metropolis-Within-Gibbs (MWG) using x and y each as their own block.
MWG works similarly to Gibbs sampling where we iteratively sample y t
+1 ∼ p θ ( y | x t ) and then x t +1 ∼ p θ ( x | y t +1 ) .
Instead of sampling exactly from the conditional (
which we cannot do since the conditional is unnormalized), we update our current sample y t using a Markov transition kernel applied to p ( y | x t ) and then update x t likewise.
In practice, we hold x t ﬁxed, then update y t using one step of Gibbs-With-Gradients to obtain y t
+1 .
We then hold y t +1 ﬁxed and apply one step of Langevin dynamics to sample x t +1 .
It is likely that more involved approaches (Zhou, 2019) could lead to improved performance or more efﬁcient sampling but we found this simple approach to work well for our applications.
In the following we will refer to the Markov transition kernel for 1 step of Langevin Dyanmics applied to p θ ( x | y t ) as T c ( x t +1 | x t , y t ) and we refer to the Markov transition kernel for 1 step of Gibbs-WithGradients applied to p θ ( y | x t ) as T d ( y t +1 | y t , x t ) .
Pseudo-code for our proposed joint sampler can be seen in Algorithm 1.
Algorithm 1 Joint Sampling
Input: EBM p θ ( x, y )
∝ e f θ ( x,y ) , initial distributions p 0 ( x ) , q 0 ( y ) , number of steps T
Output: Approximate samples x T , y T ∼ p θ ( x, y )
Initialize samples x 0 ∼ p 0 ( x ) , y 0 ∼ q
( y ) t
=while t < T do
Sample y t +1 ∼ T d ( y t +1 | y t , x t ) { Gibbs-With-Gradients }
Sample x t +1 ∼ T c ( x t +1 | x t , y t +1 ) { Langevin Dynamics } t = t +end while return x T , y T
E VALUATION
Prediction In classiﬁcation tasks, predictions are typically made by returning the conﬁguration of y which maximizes p θ ( y | x ) .
When y is high-dimensional and its dimensions are correlated, performing this maximization is not straightforward.
In these settings, we propose to use MCMC to instead sample from p θ ( y | x ) by repeatedly applying the transition kernel T d ( y t +1 | y t , x t ) .
We refer to this method of prediction as “sample-based”.
As a useful comparison, in some of our experiments y is sufﬁciently low-dimensional that it is tractable to marginalize over all conﬁgurations and obtain the exact y (cid:48)
=
arg max y p θ ( y | x ) .
We refer to this method of prediction as “energy-based”.
Conditional Sampling Conditional sampling can be achieved through repeated application of the
Markov transition kernel T c ( x t +1 | x t , y t ) .
Sometimes, we may condition on only a subset of the supervision.
In this case, we apply MWG as in Algorithm 1 with a modiﬁed Markov transition kernel
T d ( y t +1 | y t , x t ) .
This kernel updates the subset of supervision which has not been conditioned on and holds the subset of supervision that is being conditioned on constant.
E XPERIMENTS
We demonstrate that our approach is able to train joint models which can accurately predict the labels and generate conditional samples.
Throughout our experiments, we use models of the form
[Equation]
where y is a concatenation of one-hot vectors for each of the attributes, and h φ encodes data features.
The appendix includes details of the dataset preprocessing, architectures, additional samples, and density visualizations (where applicable) for all experiments.
6. A T OY M ODEL
First, we demonstrate our approach on a toy dataset generated with NumPy (Harris et al., 2020) consisting of eight gaussian modes on a circle.
The supervision consists of three binary attributes, each of which deﬁnes an even partition of the eight modes.
The label y has six dimensions as it is the concatenation of three 2-dimensional one-hot vectors.
In Figure 1, we plot samples conditioned on either conﬁguration of one of these attributes.
This model obtains > 99% predictive accuracy of the three attributes both with sample-based and energybased prediction.
6. C ONDITIONAL G ENERATIVE M ODELING ON MNIST
We scale our approach to jointly model the distribution of 28 × 28 MNIST handwritten digits and their 10 possible labels.
Figure 2 shows our model obtains both high-quality unconditional samples and high-quality class-conditional samples.
Additionally, our model achieves 93% sampling-based accuracy and 98% energy-based accuracy on the test set.
6. H IGH - DIMENSIONAL S UPERVISION ON D S PRITES
Last, we apply our method to the dSprites dataset (Matthey et al., 2017).
This dataset consists of 64 × 64 (downsampled to 32 × 32 ) binary images of several shapes at different sizes, orientations, and positions.
We restrict this dataset to the ellipse shape and orientations that are multiples ofdegrees.
The attributes are the horizontal and vertical coordinates of the shape within the image, each of which has 16 possibilities.
This results in a 32 dimensional one-hot attribute vector y with 256 possible conﬁgurations.
In Figure 3 we demonstrate our model can conditionally sample the horizontal and vertical coordinates.
In the left ﬁgure, each row conditions on a different horizontal coordinate.
As we move from the top row to the bottom row, samples are conditioned to be further to the right.
In the right ﬁgure, each column conditions on a different vertical coordinate.
As we move from the leftmost column to the rightmost column, samples are conditioned to be further to the bottom.
Our model achieves 89% sampling-based accuracy.
R ELATED W ORK
Structured prediction problems have been a key application of EBMs.
In this setting we wish to make predictions of highly structured y given inputs x .
To capture complex correlations, this is often phrased in an energy-minimization framework (Gygli et al., 2017; Belanger et al., 2017) which has many similarities to the MCMC sampling we use.
We believe our approach could be applied to these applications to add many of the beneﬁts reported in Grathwohl et al.
Next are works that explore the unique capabilities of EBMs for challenging tasks such as continual learning (Li et al., 2020) and compositional generation (Du et al., 2020a).
We believe the techniques and architectures presented in this work could extend the range of problems these ideas can be applied to.
C ONCLUSIONS AND N EXT S TEPS
In this work we developed an approach to model the joint distribution of data and supervision using
EBMs.
This approach allows us to train on data with limited or partial supervision, and allows us to obtain any of the conditional distributions.
Our work is enabled by recent advances in discrete and continuous sampling techniques for EBMs.
We demonstrated our approach simultaneously achieves accurate prediction and conditional sampling.
This work is a ﬁrst step, and in the future we hope to scale our approach to larger datasets and incorporate more complex structured supervision.
A J OINT E NERGY -B ASED M ODELS A. M AXIMUM L IKELIHOOD G RADIENT
We give a short proof of Equation 8 adapted from Nijkamp et al.
(2020b).
Assuming we are only given x (cid:48) [Equation] then we can see
[Equation]
differentiate under the integral)
[Equation]
where the last line is obtained by applying Equation 2 to the gradient of the joint log density ∇ θ log p θ ( x (cid:48) , y ) .
A. JEM
We give a brief background on JEM (Grathwohl et al., 2019) for the uninitiated reader as an informative contrast to our approach.
Typically, a K -class classiﬁcation task is solved by ﬁtting a model of p θ ( y | x ) .
This is usually done with a function f θ : X → R K where
[Equation]
and y ∈ { 1 , . .
, K } is the index representing the class-label.
Grathwohl et al.
(2019) notices that the same function f θ can be used as well to deﬁne an energy-based model of the joint as
[Equation]
With this model, the label can be marginalized out to reveal an EBM for the unconditional data distribution
[Equation]
and, when computing the conditional we ﬁnd
[Equation]
which is identical to the conditional distribution parameterized by standard classiﬁers.
JEM models are trained to maximize log p θ ( x ) + log p θ ( y | x ) as deﬁned in Equations 12 and 13.
The ﬁrst term is optimized using the gradient estimator of Equation 2.
B E XPERIMENTAL D ETAILS
We use the Adam optimizer with default parameters β 1 = 0 .
9 , β 2 = 0 .
999 (Kingma & Ba, 2014) throughout our experiments.
In Table 1, Table 2, and Table 3 we provide hyperparameters for each of our experiments.
The main hyperparameters we tuned were the step size (cid:15) and temperature λ for
Langevin dynamics.
Sampling During training, we use Persistent Contrastive Divergence (PCD) and a replay buffer with a standard size of 10000 and replacement rate of 0 . 05 throughout all of our experiments (Du &
Mordatch, 2019).
At test time, we generate samples from random noise and do not use the replay buffer.
Sampling Noise The initial distribution for Langevin dynamics is uniform with bounds equal to the per-dimension minimum and maximum values of one batch of data.
The initial distribution for
Gibbs-With-Gradients is set to be uniform over all possible conﬁgurations on dSprites, and with probabilities proportional to the number of occurrences in one batch of data for the toy data and
MNIST experiments.
Exponential Moving Average We keep an exponential moving average (EMA), as in Du et al.
(2020b), of the training parameters θ being optimized by applying the following update after each training iteration:
[Equation]
We initialize ˆ θ
= θ 0 and set µ = 0 . 99 in our experiments.
At test time we use the parameters ˆ θ .
[Table]
[Table]
[Table]
B. T OY D ATA B.1. A RCHITECTURE
The following architecture is implemented in PyTorch (Paszke et al., 2019).
Sequential( (0):Linear(in_features=8, out_features=256, bias=
True) (1): ELU(alpha=1.0) (2): Linear(in_features=256, out_features=256, bias=
True) (3): ELU(alpha=1.0) (4): Linear(in_features=256, out_features=256, bias=
True) (5): ELU(alpha=1.0) (6): Linear(in_features=256, out_features=1, bias=True) )
B. MNIST B.2. D ATA P ROCESSING
We applied standard dequantization
[Equation]
x =
[Equation]
255 + u , u ∼ U [0 , 1] followed by the standard logit transformation [Equation] [Equation] log (cid:18)
[Equation]
B.2. A RCHITECTURE
The following architecture is implemented in PyTorch (Paszke et al., 2019).
Sequential
( (0):Linear(in_features=794, out_features=1000, bias=
True) (1): ELU(alpha=1.0) (2): Linear(in_features=1000, out_features=500, bias=
True) (3): ELU(alpha=1.0) (4): Linear(in_features=500, out_features=500, bias=
True) (5): ELU(alpha=1.0) (6): Linear(in_features=500, out_features=250, bias=
True) (7): ELU(alpha=1.0) (8): Linear(in_features=250, out_features=250, bias=
True) (9): ELU(alpha=1.0) (10): Linear(in_features=250, out_features=250, bias=
True) (11): ELU(alpha=1.0) (12): Linear(in_features=250, out_features=1, bias=True) )
B . D S PRITES B.3. D ATA P ROCESSING
We downsampled images from their original 64 × 64 size to 32 × 32 .
We binned the horizontal and vertical attributes from 32 options each to 16 options each.
We applied a signiﬁcant amount of dequantization to the images with the following function
[Equation]
followed by the same logit transformation as was used for MNIST.
The test set was constructed by taking random samples from the dataset.
The ﬁnal processed dataset has 46,652 examples in the training set, and 2,500 examples in the test set.
B.3. A RCHITECTURE
We use the following CNN architecture for h φ ( x ) , implemented in PyTorch (Paszke et al., 2019).
(cnn): Sequential( (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) (1): SiLU(inplace=True) (2): Conv2d(16, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)) (3): SiLU(inplace=True) (4): Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)) (5): SiLU(inplace=True)
where SiLU is the Sigmoid Linear Unit, also known as the Swish activation function (Hendrycks &
Gimpel, 2016; Elfwing et al., 2018; Ramachandran et al., 2017).
The convolutional neural network has 256 ﬂattened output features which are concatenated to the 32 dimensional supervision and fed through the following MLP.
C A DDITIONAL S AMPLES