G UIDING E NERGY - BASED M ODELS VIA C ONTRASTIVE L ATENT V ARIABLES
A BSTRACT
An energy-based model (EBM) is a popular generative framework that offers both explicit density and architectural flexibility, but training them is difficult since it is often unstable and time-consuming.
In recent years, various training techniques have been developed, e.g. , better divergence measures or stabilization in MCMC sampling, but there often exists a large gap between EBMs and other generative frameworks like GANs in terms of generation quality.
In this paper, we propose a novel and effective framework for improving EBMs via contrastive representation learning (CRL).
To be specific, we consider representations learned by contrastive methods as the true underlying latent variable.
This contrastive latent variable could guide EBMs to understand the data structure better, so it can improve and accelerate EBM training significantly.
To enable the joint training of
EBM and CRL, we also design a new class of latent-variable EBMs for learning the joint density of data and the contrastive latent variable.
Our experimental results demonstrate that our scheme achieves lower FID scores, compared to prior-art EBM methods ( e.g. , additionally using variational autoencoders or diffusion techniques), even with significantly faster and more memory-efficient training.
We also show conditional and compositional generation abilities of our latent-variable EBMs as their additional benefits, even without explicit conditional training.
The code is available at https://github.com/hankook/CLEL .
I NTRODUCTION
Generative modeling is a fundamental machine learning task for learning complex high-dimensional data distributions p data ( x ) .
Among a number of generative frameworks, energy-based models (EBMs,
LeCun et al., 2006; Salakhutdinov et al., 2007), whose density is proportional to the exponential negative energy, i.e. , p θ ( x ) ∝ exp( − E θ ( x )) , have recently gained much attention due to their attractive properties.
For example, EBMs can naturally provide the explicit (unnormalized) density, unlike generative adversarial networks (GANs, Goodfellow et al., 2014).
Furthermore, they are much less restrictive in architectural designs than other explicit density models such as autoregressive (Oord et al., 2016b;a) and flow-based models (Rezende & Mohamed, 2015;
Dinh et al., 2017).
Hence,
EBMs have found wide applications, including image inpainting (Du & Mordatch, 2019), hybrid discriminative-generative models (Grathwohl et al., 2019; Yang & Ji, 2021), protein design (Ingraham et al., 2019; Du et al., 2020b), and text generation (Deng et al., 2020).
Despite the attractive properties, training EBMs has remained challenging; e.g. , it often suffers from the training instability due to the intractable sampling and the absence of the normalizing constant.
Recently, various techniques have been developed for improving the training stability and the quality of generated samples, for example, gradient clipping (Du & Mordatch, 2019), short
MCMC runs (Nijkamp et al., 2019), data augmentations in MCMC sampling (Du et al., 2021), and better divergence measures (Yu et al., 2020; 2021; Du et al., 2021).
To further improve EBMs, there are several recent attempts to incorporate other generative models into EBM training, e.g. , variational autoencoders (VAEs) (Xiao et al., 2021), flow models (Gao et al., 2020; Xie et al., 2022), or diffusion techniques (Gao et al., 2021).
However, they often require a high computational cost for training such an extra generative model, or there still exists a large gap between EBMs and state-of-the-art generative frameworks like GANs (Kang [Table] [Table] [Table]
[Table]
[Table]
/K/s/TDO/0syCLn9luVoh1ovBtX8/
[Table]
+EYzuAcWiBAwT08wGPwO/gb/KsFM+lTh8/
[Table] =</latexit> [Table] /eNjZeXRo8lJzHPNc5vo4AYNSZDgmQRKPC42gEolHydn7Rj86R21Enh3QosCJglkmUsGBLBV3fkUK6DRJq3kdGaH8Io4I51RNgaD2ehfqeR2HXm9N7F9oX+tv/4o883oUBy8oDptyS56DrA5srd03l6oFu9bWX8u7bpK40x0Ohsvwr4KgBV3Wxn6844TRNOelwoy4BGNOgmFBkwo0CS6x9qLSYAH8DGZ4YmEGCs2kWi609nuWmfppru3JyF+ylzMqUMYsVGKdzdBmU2vI/2pJojZaU/p6UomsKAkzvuqcltKn3G++y58KjZzkwgLgWtjhfX4KGjjZT/W86APa12n8ZDt9LlAD5fp5FYGeKZjXVXtfZxPZymZvz2452NzpVXAYDoLR4OWXUXdv1O57iz1hT1mfBewV22Mf2T4bM+5sO6Hz1nnnave7+8P9ubK6TpvzmK2F+/sv7ZPv4Q==</latexit>
[Table]
<
[Table]
sha1_
[Table]
=
[Table]
3eNtNH0utTxV4VUllBNnBe/js7I=">AAADCHicfVHditNAFJ7EvzX+bFcvvQnWQhUpSansIiws6IUI4ort7sKmhJPppDvsTBJmTpbWMS/g03gn3njhW/g2TtqIbVf3wHC+Od93fuZMUgiuMQh+Oe616zdu3tq67d25e+/+dmvnwZHOS0XZiOYiVycJaCZ4xkbIUbCTQjGQiWDHyfmrmj++YErzPBvivGBjCdOMp5wC2lDc+oFx+BzjfqS5jCTgGQVhhpXXWVyS1MyqmvKLOEI2QzMBhBX2oor7XmeN7P7hPlWf/xZ56nX+12n/5Uq1cN/Kumt5V00St9pBL1iYfxmEDWiTxg7jHacfTXJaSpYhFaD1aRgUODagkFPBKi8qNSuAnsOUnVqYgWR6bBaLrvyOjUz8NFf2ZOgvoqsZBqTWc5lYZT203uTq4D+5JJEbrTHdGxueFSWyjC47p6XwMffrb/QnXDGKYm4BUMXt8D49AwUU7Wd7XvSa2dcp9s52el8wBZirZyYCNZUwq0zjr5LxbCmz3rNbDjd3ehkc9XvhoPfiw6B9MGj2vUUekcekS0KySw7IG3JIRoQ6T5y3zkdn6H5xv7rf3O9Lqes0OQ/Jmrk/fwOulflc</latexit>
[Table]
<latexit sha1_base64="RCPxUuU4Jf4F8cbf0xtwtxqYlBI
[Table]
p ✓ ( x , z
[Table]
<latexit sha
[Table]
base
[Table]
/0N++3Mwsc3ExFrtYRs/USe+p9LTuebPzB3chv2RQf2Pee5mqbqrqZrb6VcoGm1uNncZsGHnDTIwtLRlHg0qhqbpBQlfdZuJgIU7MRgD9CHNgxKFSR6GgASaneERPlOlhl4p+NHuC0qgqj20Mfa4+D4yZN50RYVeIqWspZCxarMZi59qYZbkr1DB80Y+YF4RAPTJnHoaOAb4RP3DDZpwScKbKwIQzJd4gY8wxAfUb0HV0SNXuOH2nmD4ElGPw+ZMIYT5y8URGyXwVjHlzmJp11WVztad5o9vcMXd3nn3c3dpvJv3e0B5qj7SaZmrPtX3tjXakdTRSelWyS27JK38tfyt/L/+YQ4uFJOeBlhnl338AQNPIEw==</latexit>
Instead of utilizing extra expensive generative models, in this paper we ask whether EBMs can be improved by other unsupervised techniques of low cost.
To this end we are inspired by recent advances in unsupervised representation learning literature (Chen et al., 2020;
Grill et al., 2020;
He et al., 2021) especially by the fact that the discriminative representations can be obtained much easier than generative modeling.
Interestingly, such representations have been used to detect out-ofdistribution samples (Hendrycks et al., 2019a;b), so we expect that training EBMs can benefit from good representations.
In particular, we primarily focus on contrastive representation learning (Oord et al., 2018;
Chen et al., 2020; He et al., 2020) since it can learn instance discriminability, which has been shown to be effective in not only representation learning, but also training GANs (Jeong & Shin, 2021;
Kang et al., 2021) and out-of-distribution detection (Tack et al., 2020).
In this paper, we propose Contrastive Latent-guided Energy Learning (CLEL) , a simple yet effective framework for improving EBMs via contrastive representation learning (CRL).
Our CLEL consists of two components, which are illustrated in Figure 1.
• Contrastive latent encoder.
Our key idea is to consider representations learned by CRL as an underlying latent variable distribution p data ( z | x ) .
Specifically, we train an encoder h ϕ via CRL, and treat the encoded representation z := h ϕ ( x ) as the true latent variable given data x , i e. , z ∼ p data ( ·| x ) .
This latent variable could guide EBMs to understand the underlying data structure more quickly and accelerate training since the latent variable contains semantic information of the data thanks to CRL.
Here, we assume the latent variables are spherical,
i e. , ∥ z ∥ 2 = 1 , since recent CRL methods (He et al., 2020; Chen et al., 2020) use the cosine distance on the latent space.
• Spherical latent-variable EBM.
We introduce a new class of latent-variable EBMs p θ ( x , z ) for modeling the joint distribution p data ( x , z ) generated by the contrastive latent encoder.
Since the latent variables are spherical, we separate the output vector f :=
f θ ( x ) into its norm ∥ f ∥ 2 and direction f/ ∥ f ∥ 2 for modeling p θ ( x ) and p θ ( z | x ) , respectively.
We found that this separation technique reduces the conflict between p θ ( x ) and p θ ( z | x ) optimizations, which makes training stable.
In addition, we treat the latent variables drawn from our EBM, ˜ z ∼ p θ ( z ) , as additional negatives in CRL, which further improves our CLEL.
Namely, CRL guides EBM and vice versa.
We demonstrate the effectiveness of the proposed framework through extensive experiments.
For example, our EBM achieves 8.61 FID under unconditional CIFAR-10 generation, which is lower than those of existing EBM models.
Here, we remark that utilizing CRL into our EBM training increases training time by only 10% in our experiments ( e.g. , 38 → 41 GPU hours).
This enables us to achieve the lower FID score even with significantly less computational resources ( e.g. , we use single
RTX3090 GPU only) than the prior EBMs that utilize VAEs (Xiao et al., 2021) or diffusion-based recovery likelihood (Gao et al., 2021).
Furthermore, even without explicit conditional training, our latent-variable EBMs naturally can provide the latent-conditional density p θ ( x | z ) ; we verify its effectiveness under various applications: out-of-distribution (OOD) detection, conditional sampling, and compositional sampling.
For example, OOD detection using the conditional density shows superiority over various likelihood-based models.
Finally, we remark that our idea is not limited to contrastive representation learning and we show EBMs can be also improved by other representation learning methods like BYOL (Grill et al., 2020) or MAE (He et al., 2021)
(see Section 4.5).
P RELIMINARIES
In this work, we mainly consider unconditional generative modeling: given a set of i.i.d.
samples { x ( i ) } Ni =1 drawn from an unknown data distribution p data ( x ) , our goal is to learn a model distribution p θ ( x ) parameterized by θ to approximate the data distribution p data ( x ) .
To this end, we parameterize p θ ( x ) using energy-based models (EBMs), and incorporate contrastive representation learning (CRL) into EBMs for improving them.
We briefly describe the concepts of EBMs and CRL in Section 2.
and Section 2.2, respectively, and then introduce our framework in Section 3.
2. E NERGY - BASED MODELS
An energy-based model (EBM) is a probability distribution on R d x , defined as follows: for x ∈ R d x ,
[Equation]
where E θ ( x ) is the energy function parameterized by θ and Z θ denotes the normalizing constant, called the partition function .
An important application of EBMs is to find a parameter θ such that p θ is close to p data .
A popular method for finding such θ is to minimize Kullback–Leibler (KL) divergence between p data and p θ via gradient descent:
[Equation]
Since this gradient computation (3) is NP-hard in general (Jerrum & Sinclair, 1993), it is often approximated via Markov chain Monte Carlo (MCMC) methods.
In this work, we use the stochastic gradient Langevin dynamics (SGLD, Welling & Teh, 2011), a gradient-based MCMC method for approximate sampling.
Specifically, at the ( t + 1) -th iteration, SGLD updates the current sample ˜ x
( t ) to ˜ x ( t +1) using the following procedure:
[Equation]
where ε > 0 is some predefined constant, ˜ x
(0) denotes an initial state, and N denotes the multivariate normal distribution.
Here, it is known that the distribution of ˜ x ( T ) (weakly) converges to p θ with small enough ε and large enough T under various assumptions (Vollmer et al., 2016;
Raginsky et al., 2017; Xu et al., 2018;
Zou et al., 2021).
Latent-variable energy-based models.
EBMs can naturally incorporate a latent variable by specifying the joint density p θ ( x , z ) ∝ exp( − E θ ( x , z )) of observed data x and the latent variable z .
This class includes a number of EBMs: e.g. , deep Boltzmann machines (Salakhutdinov & Hinton, 2009) and conjugate EBMs (Wu et al., 2021).
Similar to standard EBMs, these latent-variable EBMs can be trained by minimizing KL divergence between p data ( x ) and p θ ( x ) as described in (3):
[Equation]
where E θ ( x ) is the marginal energy, i.e. , E
[Equation]
2. C ONTRASTIVE REPRESENTATION LEARNING
Generally speaking, contrastive learning aims to learn a meaningful representation by minimizing distance between similar ( i.e. , positive) samples, and maximizing distance between dissimilar ( i.e. , negative) samples on the representation space.
Formally, let h ϕ : R d x → R d z be a ϕ -parameterized encoder, ( x , x + ) and ( x , x − ) be positive and negative pairs, respectively.
Contrastive learning then maximizes sim ( h ϕ ( x ) , h ϕ ( x + )) and minimizes sim ( h ϕ ( x ) , h ϕ ( x − )) where sim ( · , · ) is a similarity metric defined on the representation space R d z .
Under the unsupervised setup, various methods for constructing positive and negative pairs have been proposed: e.g. , data augmentations (He et al., 2020; Chen et al., 2020; Tian et al., 2020), spatial or temporal co-occurrence
(Oord et al., 2018), and image channels (Tian et al., 2019).
In this work, we mainly focus on a popular contrastive learning framework, SimCLR (Chen et al., 2020), which constructs positive and negative pairs via various data augmentations such as cropping and color jittering.
Specifically, given a mini-batch B = { x ( i ) } ni =1 , SimCLR first constructs two augmented views { v ( i ) j :=
t ( i ) j
( x ( i ) ) } j ∈{ 1 , 2 } for each data sample x ( i ) via random augmentations t ( i ) j ∼ T .
Then, it considers ( v ( i ) 1 , v ( i ) 2 ) as a positive pair and ( v ( i ) 1 , v ( k ) 2 ) as a negative pair for all k ̸
= i .
The
SimCLR objective L SimCLR is defined as follows:
[Equation]
where sim ( u , v )
=
u
⊤ v / ∥ u ∥ 2 ∥ v ∥ 2 is the cosine similarity, τ is a hyperparameter for temperature scaling, and
L NT-Xent denotes the normalized temperature-scaled cross entropy (Chen et al., 2020).
M ETHOD
Recall that our goal is to learn an energy-based model (EBM) p θ ( x ) ∝ exp( − E θ ( x )) to approximate a complex underlying data distribution p data ( x ) .
In this work, we propose Contrastive Latent-guided
Energy Learning (CLEL), a simple yet effective framework for improving EBMs via contrastive representation learning.
Our key idea is that directly incorporating with semantically meaningful contexts of data could improve EBMs.
To this end, we consider the (random) representation z ∼ p data ( z | x ) of x , generated by contrastive learning, as the underlying latent variable.
Namely, we model the joint distribution p data ( x , z ) =
p data ( x ) p data ( z | x ) via a latent-variable EBM p θ ( x , z ) .
Our intuition on the benefit of modeling p data ( x , z ) is two-fold: (i) conditional generative modeling p data ( x | z ) given some good contexts (e.g., labels) of data is much easier than unconditional modeling p data ( x ) (Mirza & Osindero, 2014; Van den Oord et al., 2016;
Reed et al., 2016), and (ii) the mode collapse problem of generation can be resolved by predicting the contexts p data ( z | x ) (Odena et al., 2017; Bang & Shim, 2021).
The detailed implementations of p data ( z | x ) , called the contrastive latent encoder, and the latent-variable EBM are described in Section 3.1 and 3.2, respectively, while
Section 3.3 presents how to train them in detail.
Our overall framework is illustrated in Figure 1.
3. C ONTRASTIVE LATENT ENCODER
To construct a meaningful latent distribution p data ( z | x ) for improving EBMs, we use contrastive representation learning.
To be specific, we first train a latent encoder h ϕ : R d x → R d z , which is a deep neural network (DNN) parameterized by ϕ , using a variant of the SimCLR objective L SimCLR (7) (we describe its detail in Section 3.3) with a random augmentation distribution T .
Since our objective only measures the cosine similarity between distinct representations, one can consider the encoder h ϕ maps a randomly augmented sample to a unit vector .
We define the latent sampling procedure z ∼ p data ( z | x ) as follows:
[Equation]
3. S PHERICAL LATENT - VARIABLE ENERGY - BASED MODELS
We use a DNN f θ : R d x → R d z parameterized by θ for modeling p θ ( x , z ) .
Following that the latent variable z ∼ p data ( z | x ) is on the unit sphere, we utilize the directional information f θ ( x ) /
∥ f θ ( x ) ∥ for modeling p θ ( z | x ) , while the remaining information ∥
f θ ( x ) ∥ 2 is used for modeling p θ ( x ) .
We empirically found that this norm-direction separation stabilizes the latent-variable EBM training.
For better modeling p θ ( z | x ) , we additionally apply a directional projector g θ : S d z − 1 → S d z − 1 to f θ ( x ) /
∥ f θ ( x ) ∥ 2 , which is constructed by a two-layer MLP, followed by ℓ 2 normalization.
We found that it is useful for narrowing the gap between distributions of the direction f θ ( x ) /
∥ f θ ( x ) ∥ 2 and the uniformly-distributed latent variable p data ( z ) (see Section 4.5 and Appendix E for detailed discussion).
Overall, we define the joint energy E θ ( x , z ) as follows:
[Equation]
(11) where β ≥ 0 is a hyperparameter.
Note that the marginal energy E θ ( x ) only depends on ∥ f θ ( x ) ∥since (cid:82) S d − 1 exp( βg θ ( f θ ( x ) /
∥ f θ ( x ) ∥ 2 )
⊤ z ) d z is independent of x due to the symmetry.
Also, the norm-based design does not sacrifice the flexibility for energy modeling (see Appendix F for details).
3. T RAINING
Remark that Section 3.1 and 3.2 define p data ( x , z ) and p θ ( x , z ) , respectively.
We now describe how to train the contrastive latent encoder h ϕ and the spherical latent-variable EBM p θ via mini-batch stochastic optimization algorithms in detail (see Appendix A for the pseudo-code).
Let { x ( i ) } ni =1 be real samples randomly drawn from the training dataset.
We first generate n samples { ˜ x ( i ) } ni =1 ∼ p θ ( x ) using the current EBM via stochastic gradient Langevin dynamics (SGLD) (4).
Here, to reduce the computational complexity and improve the generation quality of SGLD, we use two techniques: a replay buffer to maintain Markov chains persistently (Du & Mordatch, 2019), and periodic data augmentation transitions to encourage exploration (Du et al., 2021).
We then draw latent variables from p data and p θ : z ( i ) ∼ p data ( z | x ( i ) ) and ˜ z ( i ) ∼ p θ ( z | ˜ x
( i ) ) for all i .
For the latter case, we simply use the mode of p θ ( z | x ( i ) ) instead of sampling, namely,
˜ z
( i ) :
=
g θ ( f θ ( x ) /
∥ f θ ( x ) ∥ 2 ) .
Let
B := { ( x ( i ) , z
( i ) ) }
ni =1 and ˜ B := { (˜ x ( i ) ,
˜ z ( i ) ) } ni =1 be real and generated mini-batches, respectively.
Under this setup, we define the objective L EBM for the EBM parameter θ as follows:
[Equation]
(12) where the first two terms correspond to the empirical average of D KL ( p data ∥
p θ ) 4 and α is a hyperparameter for energy regularization to prevent divergence, following Du & Mordatch (2019).
When training the latent encoder h ϕ via contrastive learning, we use the SimCLR
(Chen et al., 2020) loss
L SimCLR (7) with additional negative latent variables { ˜ z
( i ) } ni
=1 .
To be specific, we define the objective
L LE for the latent encoder parameter ϕ as follows:
[Equation]
where L NT-Xent is the normalized temperature-scaled cross entropy defined in (8), τ is a hyperparameter for temperature scaling, z ( i ) j :
=
h ϕ ( t ( i ) j
( x ( i ) )) , and { t ( i ) j } ∼ T are random augmentations.
We found that considering { ˜ z ( i ) } ni =1 as negative representations for contrastive learning increases the latent diversity, which further improves the generation quality in our CLEL framework.
To sum up, our CLEL jointly optimizes the latent encoder h ϕ and the latent-variable EBM ( f θ , g θ ) from scratch via the following optimization: min ϕ,θ E B , ˜ B [
[Equation]
EBM
[Equation] [Equation]
B
[Equation]
θ, α, β [Equation] + L LE ( B , ˜ B ; ϕ, τ )] .
After training, we only utilize our latent-variable EBM ( f θ , g θ ) when generating samples.
The latent encoder h ϕ is used only when extracting a representation of a specific sample during training.
Here
∇ θ (˜ ˜ z )
∇ θ E θ (
[Table]
[Table] [Table]
[Table]
E XPERIMENTS
We verify the effectiveness of our Contrastive Latent-guided Energy Learning (CLEL) framework under various scenarios: (a) unconditional generation (Section 4.1), (b) out-of-distribution detection (Section 4.2), (c) conditional sampling (Section 4.3)
[Table]
and (d) compositional sampling (Section 4.4).
All the architecture, training, and evaluation details are described in Appendidx B.
4. U NCONDITIONAL IMAGE GENERATION
An important application of EBMs is to generate images using the energy function E θ ( x ) .
To this end
[Table]
we train our CLEL framework on CIFAR-10 (Krizhevsky et al., 2009) and ImageNet 32 ×(Deng et al., 2009; Chrabaszcz et al., 2017) under the unsupervised setting.
We then generate 50k samples using SGLD and evaluate their qualities using Fréchet Inception Distance (FID) scores (Heusel et al., 2017; Seitzer, 2020).
The unconditionally generated samples are provided in Figure 2.
Table 1 and 2 show the FID scores of our CLEL and other generative models for unconditional generation on CIFAR-10 and ImageNet 32 × 32, respectively.
We first find that CLEL outperforms previous EBMs under both CIFAR-10 and ImageNet 32 × 32 datasets.
As shown in Table 3, our method can benefit from a multi-scale architecture as Du et al.
(2021) did, contrastive representation learning (CRL) with a larger batch, more channels at lower layers in our EBM f θ .
As a result, we achieve 8.61 FID on CIFAR-10, which is lower than that of the prior-art EBM based on diffusion recovery likelihood, EBM-Diffusion (Gao et al., 2021), even with 5 × faster and 4 × more memoryefficient training (when using the similar number of parameters for EBMs).
Then, we narrow the gap between EBMs and state-of-the-art frameworks like GANs without help from other generative models.
We think our CLEL can be further improved by incorporating an auxiliary generator (Arbel et al., 2021; Xiao et al., 2021) or diffusion (Gao et al., 2021), and we leave it for future work.
[Figure]
[Table]
4. O UT - OF - DISTRIBUTION DETECTION
EBMs can be also used for detecting out-of-distribution (OOD) samples.
For the OOD sample detection, previous EBM-based approaches often use the (marginal) unnormalized likelihood p θ ( x ) ∝ exp( − E θ ( x )) .
In contrast, our CLEL is capable of modeling the joint density p θ ( x , z ) ∝ E θ ( x , z ) .
Using this capability, we propose an energy-based OOD detection score: given x ,
[Equation]
We found that the second term in (14) helps to detect the semantic difference between in- and out-of-distribution samples.
Table 4 shows our CLEL’s superiority over other explicit density models in OOD detection, especially when OOD samples are drawn from different domains, e.g. , SVHN (Netzer et al., 2011) and Texture (Cimpoi et al., 2014) datasets.
4. C ONDITIONAL SAMPLING
One advantage of latent-variable EBMs is that they can offer the latent-conditional density p θ ( x | z ) ∝ exp( − E θ ( x , z )) .
Hence, our EBMs can enjoy the advantage even though CLEL does not explicitly train conditional models.
To verify this, we first test instance-conditional sampling : given a real sample x , we draw the underlying latent variable z ∼ p data ( z | x ) using our latent encoder h ϕ , and then perform SGLD sampling using our joint energy E θ ( x , z ) defined in (10).
We here use our
CIFAR-10 model.
As shown in Figure 3a, the instance-conditionally generated samples contain similar information ( e.g. , color, shape, and background) to the given instance.
[Figure]
car bird cat deer dog frog horse ship trunk Class
[Figure]
Predictions
C o nd i t i o n e d C l a ss
[Figure]
[Figure]
This successful result motivates us to extend the sampling procedure: given a set of instances { x ( i ) } , can we generate samples that contain the shared information in { x ( i ) } ?
To this end, we first draw latent variables
z ( i ) ∼ p data ( ·| x ( i ) ) for all i , and then aggregate them by summation and normalization:
¯ z := (cid:80)
i
z ( i ) /
∥ (cid:80)
i
z
( i ) ∥ 2 .
To demonstrate that samples generated from p θ ( x | ¯ z ) contains the shared information in { x ( i ) } , we collect the set of instances { x ( i ) y } for each label y in CIFAR-10, and check whether ˜ x y ∼ p θ ( ·| ¯ z y ) has the same label y .
Figure 3b shows the class-conditionally generated samples { ˜ x y } and Figure 3c presents the confusion matrix of predictions for { ˜ x y } computed by an external classifier c .
Formally, each ( i, j ) -th entry is equal to P ˜ x
i ( c (˜ x i ) = j ) .
We found that ˜ x y is likely to be predicted as the label y , except the case when y is dog: the generated dog images sometimes look like a semantically similar class, cat.
These results verify that our EBM can generate samples conditioning on a instance or class label, even without explicit conditional training.
4. C OMPOSITIONALITY VIA LATENT VARIABLES
An intriguing property of EBMs is compositionality (Du et al., 2020a): given two EBMs E ( x | c 1 ) and
E ( x | c 2 ) that are conditional energies on concepts c 1 and c 2 , respectively, one can construct a new energy conditioning on both concepts: p θ ( x | c 1 and c 2 ) ∝ exp( − E ( x | c 1 ) − E ( x | c 2 )) .
As shown in Section 4.3, our CLEL implicitly learns E ( x | z ) , and a latent variable z can be considered as a concept, e.g. , instance or class.
Hence, in this section, we test compositionality of our model.
To this end, we additionally train our CLEL in CelebA 64 × 64 (Liu et al., 2015).
For compositional sampling, we first acquire three attribute vectors ¯ z a for a ∈ A :=
{ Young , Female , Smiling } as we did in Section 4.3, then generate samples from a composition of conditional energies as follows:
[Equation]
where sim ( · , · ) is the cosine similarity.
Figure 4a and 4b show the generated samples conditioning on multiple attributes and their attribute prediction results computed by an external classifier, respectively.
They verify our compositionality qualitatively and quantitatively.
For example, almost generated faces conditioned by { Young , Female } look young and female (see the third row in Figure 4.)
[Table]
[Table] [Table]
4. A BLATION STUDY
Component analysis.
To verify the importance of our CLEL’s components, we conduct ablation experiments with training a smaller ResNet (He et al., 2016) in CIFAR-10 (Krizhevsky et al., 2009) for 50k training iterations.
Then, we evaluate the quality of energy functions using FID and OOD detection scores.
Here, we use SVHN (Netzer et al., 2011) as the OOD dataset.
Table 5 demonstrates the effectiveness of CLEL’s components.
First, we observe that learning p θ ( z | x ) to approximate p data ( z | x ) plays a crucial role for improving generation (see (a) vs. (b)).
In addition, using generated latent variables ˜
z ∼ p θ ( · ) as negatives for contrastive learning further improves not only generation, but also OOD detection performance (see (b) vs. (c)).
We also empirically found that using an additional projection head is critical; without projection g θ ( i.e. , (d))
[Table]
our EBM failed to approximate p data ( x )
[Table]
but an additional projection head ( i
[Table]
e. , (c) or (e)) makes learning feasible.
Hence, we use a 2-layer MLP (c) in all experiments since it is better than a simple linear function (e).
We also test various β ∈ { 0 . 1 , 0 . 01 , 0 .
001 } under this evaluation setup (see Table 6) and find β = 0 .
01 is the best.
Compatibility with other self-supervised representation learning methods.
While we have mainly focused on utilizing contrastive representation learning (CRL), our framework CLEL is not limited to CRL for learning the latent encoder h ϕ .
To verify this compatibility, we replace SimCLR with other self-supervised representation learning (SSRL) methods, BYOL (Grill et al., 2020) and MAE (He et al., 2021).
See Appendix C for implementation details.
Note that these methods have several advantages compared to SimCLR: e.g., BYOL does not require negative pairs, and MAE does not require heavy data augmentations.
Table 7 implies that any SSRL methods can be used to improve
EBMs under our framework, where the CRL method, SimCLR (Chen et al., 2020), is the best.
R ELATED WORKS
Energy-based models (EBMs) can offer an explicit density and are less restrictive in architecture design, but training them has been challenging.
For example, it often suffers from the training instability due to the time-consuming and unstable MCMC sampling procedure ( e.g. , a large number of SGLD steps).
To reduce the computational complexity and improve the quality of generated samples, various techniques have been proposed: a replay buffer (Du & Mordatch, 2019), short-run
MCMC (Nijkamp et al., 2019), augmentation-based MCMC transitions (Du et al., 2021).
Recently, researchers have also attempted to incorporate other generative frameworks into EBM training e.g. , adversarial training (Kumar et al., 2019; Arbel et al., 2021; Grathwohl et al., 2021), flow-based models (Gao et al., 2020; Nijkamp et al., 2022; Xie et al., 2022), variational autoencoders (Xiao et al., 2021), and diffusion techniques (Gao et al., 2021).
Another direction is on developing better divergence measures, e.g. , f -divergence (Yu et al., 2020), pseudo-spherical scoring rule (Yu et al., 2021), and improved contrastive divergence (Du et al., 2021).
Compared to the recent advances in the EBM literature, we have focused on an orthogonal research direction that investigates how to incorporate discriminative representations, especially of contrastive learning, into training EBMs.
C ONCLUSION
The early advances in deep learning was initiated from the pioneering energy-based model (EBM) works, e.g., restricted and
deep Boltzman machines (Salakhutdinov et al., 2007; Salakhutdinov &
Hinton, 2009), however the recent accomplishments rather rely on other generative frameworks such as diffusion models (Ho et al., 2020; Song et al., 2021).
To narrow the gap, in this paper, we suggest to utilize discriminative representations for improving EBMs, and achieve significant improvements.
We hope that our work would shed light again on the potential of EBMs, and would guide many further research directions for EBMs.
E THICS STATEMENT
In recent years, generative models have been successful in synthesizing diverse high-fidelity images.
However, high-quality image generation techniques have threats to be abused for unethical purposes, e g., creating sexual photos of an individual.
This significant threats call us researchers for the future efforts on developing techniques to detect such misuses.
Namely, we should learn the data distribution p ( x ) more directly.
In this respect, learning explicit density models like VAEs and EBMs could be an effective solution.
As we show the superior performance in both image generation and out-of-distribution detection, we believe that energy-based models, especially with discriminative representations, would be an important research direction for reliable generative modeling
R EPRODUCIBILITY STATEMENT
We provide all the details to reproduce our experimental results in Appendix B. Our code is available at https://github.com/hankook/CLEL .
In our experiments, we mainly use NVIDIA GTX3090 GPUs.
A CKNOWLEDGMENTS AND DISCLOSURE OF FUNDING
This work was mainly supported by Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No.2021-0-02068, Artificial
Intelligence Innovation Hub; No.2019-0-00075, Artificial Intelligence Graduate School Program (KAIST)).
This work was partly supported by KAIST-NAVER Hypercreative AI Center.
A T RAINING PROCEDURE OF CLEL
Require: a latent-variable EBM ( f θ , g θ ) , a latent encoder h ϕ , an augmentation distribution T , hyperparameters α, β, τ > 0 , and the stop-gradient operation sg ( · ) .
1: for # training iterations do 2: //
Construct batches B and ˜ B 3:
Sample { x ( i ) } ni
=1 ∼ p data ( x ) 4:
Sample { ˜ x ( i ) } ni =1 ∼ p θ ( x ) using stochastic gradient Langevin dynamics (SGLD) 5:
z ( i ) ← sg (cid:0) h ϕ ( t ( i ) ( x ( i ) )) /
∥ h ϕ ( t ( i ) ( x ( i ) ))
∥ 2 (cid:1) , t ( i ) ∼ T 6: ˜
z ( i ) ← sg (cid:0) g θ
( f θ (˜ x ( i ) ) /
∥ f θ (˜ x ( i ) )
∥ 2 ) (cid:1) 7: // Compute the EBM loss, L EBM 8:
L EBM ← 1 n (cid:80) ni
=1 E θ ( x ( i ) , z
( i ) )
− E θ (˜ x ( i ) )
+ α · ( E θ ( x ( i ) )
2 + E θ (˜ x ( i ) ) 2 ) 9: // Compute the encoder loss, L LE 10:
z ( i )
j ← h ϕ ( t ( i ) j
( x ( i ) )) , t ( i ) j ∼ T 11:
L LE ← 12 n (cid:80) ni
=1 (cid:80)
j
=1 , 2 L NT-Xent (cid:16) z
( i ) j , z
( i )− j , { z ( k ) l } k ̸
= i,l ∈{ 1 , 2 } ∪ {
˜ z
( i ) } ni =1 ; τ (cid:17) 12:
Update θ and ϕ to minimize L EBM + L LE 13: end for
B T RAINING DETAILS
Architectures.
For the spherical latent-variable energy-based model (EBM)
f θ , we use the 8-block
ResNet (He et al., 2016) architectures following Du & Mordatch (2019).
The details of the (a) small, (b) base, and (c) large ResNets are described in Table 8.
We append a 2-layer MLP with a output dimension of 128 to the ResNet, i
[Table]
e. , f θ : R 3 × 32 × 32 → R 128 .
Note that we use the small model for ablation experiments in Section 4.5.
To stabilize training, we apply spectral normalization (Miyato et al., 2018) to all convolutional layers.
For the projection g θ , we use a 2-layer MLP with a output dimension of 128, the leaky-ReLU activation, and no bias, i
[Table]
e. , g θ ( u ) =
W 2 σ ( W 1 u ) ∈
R 128 .
For the latent encoder h ϕ , we simpy use the CIFAR variant of ResNet-18 (He et al., 2016), followed by a 2-layer MLP with a output dimension of 128. [Table]
[Table]
Training.
For the EBM parameter θ , we use Adam optimizer (Kingma & Ba, 2015) with β 1 = 0 , β 2 = 0 . 999 , and a learning rate of 10 − 4 .
We use the linear learning rate warmup for the first 2k training iterations.
For the encoder parameter ϕ , we use SGD optimizer with a learning rate of 3 × 10 − 2 , a weight decay of 5 × 10 − 4 , and a momentum of 0 .
9 as described in Chen & He (2020).
For all experiments, we train our models up to 100k iterations with a batch size of 64, unless otherwise stated.
For data augmentation T , we follow Chen et al.
(2020), i.e. , T includes random cropping, flipping, color jittering, and color dropping.
For hyperparameters, we use α = 1 following Du &
Mordatch (2019), and β = 0 .
01 (see Section 4.5 for β -sensitivity experiments).
For our large model, we use a large batch size of 256 only for learning the contrastive encoder h ϕ .
After training, we utilize exponential moving average (EMA) models for evaluation.
SGLD sampling.
For each training iteration, we use 60 SGLD steps with a step size of 100 for sampling ˜ x ∼ p θ .
Following Du et al.
(2021), we apply a random augmentation t ∼ T for every 60 steps.
We also use a replay buffer with a size of 10000 and a resampling rate of 0.1% for maintaining diverse samples (Du & Mordatch, 2019).
For evaluation, we run 600 and 1200 SGLD steps from uniform noises for our base and large models, respectively.
C I MPLEMENTATION DETAILS FOR BYOL AND MAE
We here provide implementation details for replacing SimCLR (Chen et al., 2020) with BYOL (Grill et al., 2020) and MAE (He et al., 2021) under our CLEL framework as shown in Section 4.5.
BYOL.
Since BYOL also learns its representations on the unit sphere, the method can be directly incorporated with our CLEL framework.
MAE.
Since MAE’s representations do not lie on the unit sphere, we incorporate MAE into our
CLEL framework by the following procedure:
Pretrain a MAE framework and remove its MAE decoder.
To this end, we simply use a publicly-available checkpoint of the ViT-tiny architecture.
Freeze the MAE encoder parameters and construct a learnable 2-layer MLP on the top of the encoder.
Train only the MLP via contrastive representation learning without data augmentations using our objective (13) for the latent encoder.
D T RAINING STABILITY WITH NORM - DIRECTION SEPARATION
[Table]
hd+
[Table] + [Table] </latexit>
<latexit sha
[Table]
="
[Table] [Table] [Table] [Table]
f ✓
[Table] [Table] [Table]
/X2</latexit>
[Figure]
At the early stage of our research, we first tested a multi-head architecture for modeling p ( x ) and p ( z | x ) .
To be specific, E ( x )
=
g ′ θ
( f θ ( x )) ∈ R and E ( z | x ) =
z ⊤ g
θ
( f ( x )) where f is a shared backbone and g and g ′ are separate 2-layer MLPs, as shown in Figure 5a.
We found that, with this choice, learning p ( x ) causes a mode collapse for f θ ( x ) because all samples should be aligned with a specific direction.
In contrast, learning p ( z | x ) encourages f θ ( x ) to be diverse due to contrastive learning.
Namely, modeling p ( x ) and p ( z | x ) with the multi-head architecture makes some conflict during optimization.
We empirically observe that the multi-head architecture is unstable in EBM training, as shown in Figure 5c.
To remove such a conflict, we design the norm-direction separation for modeling p(x) and p(z|x) simultaneously (i.e., Figure 5b), which leads to training stability, as shown in Figure 5c.
E T HE ROLE OF DIRECTIONAL PROJECTOR
[Figure]
Our directional projector g θ is designed for narrowing the gap between the EBM feature direction f θ ( x ) /
∥ f θ ( x ) ∥ 2 and the “uniformly-distributed” latent variable p data ( z ) (i.e., h θ ( x ) /
∥ h θ ( x ) ∥ 2 ).
Specifically, the contrastive latent variable p data ( z ) is known to be uniformly distributed (Wang &
Isola, 2020), but we observed that it is difficult to optimize the feature direction f θ ( x ) / ∥ f θ ( x ) ∥ 2 to be uniform along with learning our norm-based EBM p θ ( x ) ∝ exp( −∥ f θ ( x ) ∥ 22 ) at the same time.
As empirical supports, we analyze the cosine similarity distributions using f θ ( x ) , g θ ( f θ ( x )) , and h ϕ ( x ) features on CIFAR-10, as shown in Figure 6.
This figure shows that f θ ( x ) tends to learn similar directions (see Figure 6a) while h ϕ ( x ) tends to be uniformly distributed (see Figure 6c).
Hence, it is necessary to employ a projection between them.
We found that our projector g θ successfully narrows the gap as shown in Figure 6b, which significantly improves training EBMs.
F F LEXIBILITY OF NORM - BASED ENERGY FUNCTION
Our norm-based energy parametrization does not sacrifice the flexibility compared to the vanilla parametrization of EBMs.
We here show that any vanilla EBM p 1 ( x ) ∝ exp( f 1 ( x )) , f 1 ( x ) ∈ R , can be formulated by a norm-based EBM
p 2 ( x ) ∝ exp( ∥ f 2 ( x ) ∥ 22 ) , f 2 ( x ) ∈ R d , on a compact input space
X (e.g., an image x lies on the continuous pixel space X =
[0 , 255] HWC ).
Let b := min x ∈X f 1 ( x ) be the minimum value of f 1 .
Then, p 3 ( x ) ∝ exp( f 1 ( x ) − b ) is identical to p 1 due to the normalizing constant.
Furthermore, p 3 can be formulated as a special case of the norm-based EBM p 2 : for example, if the first component of f 2 is the same as (cid:112) f 1 ( x ) − b and other components are zero, then p 1 and p 2 model the exactly same distribution.
Therefore, energy-modeling with our norm-based design is not much different from that with the vanilla form.