H IERARCHICAL B AYES A UTOENCODERS
Anonymous authors
Paper under double-blind review
A BSTRACT
Autoencoders are powerful generative models for complex data, such as images.
However, standard models like the variational autoencoder (VAE) typically have unimodal Gaussian decoders, which cannot effectively represent the possible semantic variations in the space of images.
To address this problem, we present a new probabilistic generative model called the Hierarchical Bayes Autoencoder (HBAE) .
The HBAE contains a multimodal decoder in the form of an energybased model (EBM), instead of the commonly adopted unimodal Gaussian distribution.
The HBAE can be trained using variational inference, similar to a VAE, to recover latent codes conditioned on inputs.
For the decoder, we use an adversarial approximation where a conditional generator is trained to match the EBM distribution.
During inference time, the HBAE consists of two sampling steps: ﬁrst a latent code for the input is sampled, and then this code is passed to the conditional generator to output a stochastic reconstruction.
The HBAE is also capable of modeling sets, by inferring a latent code for a set of examples, and sampling set members through the multimodal decoder.
In both single image and set cases, the decoder generates plausible variations consistent with the input data, and generates realistic unconditional samples.
To the best our knowledge, Set-HBAE is the ﬁrst model that is able to generate complex image sets.
I NTRODUCTION
Autoencoders are a crucial generative modeling architecture for many applications involving complex data, including image (e.g., Razavi et al., 2019) and scene synthesis (Eslami et al., 2018), image compression (Minnen et al., 2018), video prediction (Lee et al., 2018), and model-based reinforcement learning (Ha & Schmidhuber, 2018).
The variational approach exempliﬁed by the
VAE (Kingma & Welling, 2013) has the beneﬁt of tractable inference and uncertainty estimation while also allowing reasonable control over bottleneck representation and reconstruction quality.
While a number of variants have been proposed to incorporate expressive priors for the VAE latent distribution, such as mixtures of Gaussians (Dilokthanakul et al., 2016) and discrete formulations (Rolfe, 2016), the VAE decoder is very limited given that data dimensions are decoded as conditionally independent Gaussian distributions.
VAE samples therefore tend to be blurry and exhibit little meaningful variation around the predicted mean.
On the other hand, complex multimodal decoder distributions can be learned by alternative generative models, including GANs (Goodfellow et al., 2014) and autoregressive models (e.g., pixelCNN) (van den Oord et al., 2016).
Conditional formulations of these models (Mirza & Osindero, 2014;
Van den Oord et al., 2016) allow controlling the sampling behavior to be class or embedding speciﬁc.
Typically a known class label or metric embedding is given when training these conditional models.
Conditioning on a class label or embedding allows, for instance in the case of face images, to generate a range of different images of a given person identity varying in pose, lighting, expression, etc., by conditioning on the latent embedding learned by a face recognition network (e.g., Van den Oord et al., 2016).
These models can be difﬁcult to train to ensure that the decoder pays attention to the conditioning variable.
Encoder-decoder architectures such as VAE-GAN (Larsen et al., 2015) and BiGAN (Donahue et al., 2016) improve sample and reconstruction sharpness through their use of adversarial losses.
However, none of these models incorporate multimodal decoders into their formulations, and therefore do not lead to diverse sampling behavior.
In this work we introduce the Hierarchical Bayes Autoencoder (HBAE).
The model incorporates a multimodal decoder, in the form of an EBM, which we show under mild assumptions can be trained similarly to a conditional GAN, where the conditioning is the latent code of a VAE.
We show that the model is capable of learning to generate crisp and varied samples.
Empirically, the stochastic reconstructions produced by the decoder retain semantic features of the respective input examples as veriﬁed by distance in Inception feature space, while unconditionally generated samples are diverse and cover the data manifold, as veriﬁed with FID scores (Heusel et al., 2017).
We also introduce a simple extension called the Set-HBAE where each data point is a collection of examples (e.g., set of images).
In this scenario, the traditional VAE is no longer applicable, because of the inherent one to many mapping between the latent code and the set elements.
We show that
Set-HBAE is able to be successfully applied to such a setting, and learns to reconstruct input sets in a semantically meaningful way, as well as generate new sets unconditionally by sampling from the latent code prior.
Furthermore, we show that the Set-HBAE works well as a method for few-shot classiﬁcation on ShapeNet (Chang et al., 2015) relative to a supervised baseline, especially in the small data regime.
To summarize, the contributions of this work are:
A novel probabilistic generative model trained using variational inference, with a ﬂexible learned multimodal decoder in the form of an EBM.
Empirical results showing that the multimodal decoder is able to produce stochastic reconstructions varying around the local data manifold of examples, and diverse unconditional samples.
An extension of the HBAE formulation to model sets of inputs, allowing a ﬁner degree of semantic control over sampling behavior, and a new method for few-shot classiﬁcation of sets.
M ETHODOLOGY 2. VAE S
VAEs (Kingma & Welling, 2013) represent a family of deep directed graphical models.
Given an input x ∈ X , a VAE formulates a joint distribution of x and a latent variable z ∈ R l as
[Equation]
where θ denotes the set of trainable parameters; p θ ( z ) and p θ ( x | z ) denote the prior and decoder distribution, respectively.
As a common design choice, p θ ( z ) is usually set as a simple low dimensional distribution (e.g., an Isotropic Gaussian) which can be easily sampled from; p θ ( x | z ) , referred to as the decoder, is parameterized as a Gaussian distribution with a learned mean µ ( z ; θ ) and an identity covariance matrix.
This popular instantiation of VAEs offers a simple objective function which corresponds to a regularized autoencoder where reconstructions are penalized with (cid:96) 2 loss.
We argue that such a unimodal decoder is not optimal for at least two scenarios.
Firstly, in the case where the dimensionality of z is much lower than the inherent dimensionality of the input x , a unimodal decoder becomes the bottleneck of learning.
In an extreme case, if the latent space is limited to contain only k bits, e.g., a categorical VAE as in Jang et al.
(2016) with 2 k latent units, with k being a small quantity, the model is forced to assign multiple inputs to any given latent code z .
This makes having a multimodal decoder essential to effectively ﬁt the data distribution.
Secondly, for settings where the input does not have a vector representation, such as set inputs (Zaheer et al., 2017), it is not clear how the unimodal decoder can be applied.
2. H IERARCHICAL B AYES A UTOENCODER
We now present the Hierarchical Bayes Autoencoder (HBAE), which addresses the drawbacks of the
VAE as discussed previously.
The key insight is to design a multimodal decoder that can ﬂexibily model the implicit (and explicit in the set case) one-to-many mapping from a latent code to target samples.
To do so, we replace the decoder distribution of a VAE with an EBM-based decoder LeCun et al.
Formally, we start by using the joint distribution in Equation 1.
We then deﬁne p θ ( x | z ) =
e − E ( x | z ; θ ) (cid:82) x e − E ( x | z ; θ ) d x , with E ( x | z ; θ ) being the energy function.
From the graphical model perspective, this formulation removes the conditional independence assumption between the dimensions of x given z , and allows for modeling rich interactions between them instead.
This enables the decoder to assign high probability to multiple possible outputs, if needed (refer to Figure 2 for an illustration).
Note that the EBM formulation of the decoder is a strict generalization of Gaussian and other exponential family energy-based models; see LeCun et al.
(2006) for more explanations.
During training, we similarly apply a variational distribution q φ ( z | x ) as in a VAE, which leads to the following variational objective:
[Equation]
where the ﬁrst row is the standard variational lower bound
Directly minimizing Equation 2 as an objective function is difﬁcult as it involves the logarithm of a partition function.
As a solution, we turn to recent connections between deep EBMs and GANs (Zhao et al., 2016; LeCun et al., 2006; Wang & Liu, 2016; Zhai et al., 2016).
As a brief recap, following Zhai et al.
(2016), for an EBM with energy E ( x ; θ ) , we can write its negative log likelihood
[Equation]
NLL) as: E x
[Equation]
p data ( x ) [ E ( x ; θ )]
+ log[
[Equation]
e
[Equation]
x
[Equation]
θ
[Equation]
d x ] , which can be further developed as:
[Equation]
where q ψ ( x ) is an auxiliary distribution with parameters ψ , with H ( q ψ ) denoting its entropy.
We then let D ( x ; θ ) =
− E ( x ; θ ) and q ψ ( x ) =
(cid:82)
z
[Equation]
p ψ ( z (cid:48) )
δ ( G ( z (cid:48) ; ψ ))
d z (cid:48) (i.e., an implicit distribution deﬁned by a generator G and a known prior p ψ ( z (cid:48) ) ), which leads to an optimization procedure as follows:
[Equation]
In practice, we rely on implicit entropy regularizers such as batch normalization (Ioffe & Szegedy, 2015) in place of the entropy term H ( q ψ ) , which gives us a learning objective:
[Equation]
which can be essentially implemented as a GAN, in particular a WGAN (Arjovsky et al., 2017), with D being the discriminator and G being the generator.
We then accordingly apply the same derivation and change of variables to Equation 2, and obtain a practical objective for HBAE:
[Equation]
Here both the discriminator D and generator G are conditional on z , which is a sample from the encoder q φ .
The additional noise distribution p ψ ( z (cid:48) | z ) can also be conditional, but we set it to a parameter free Gaussian distribution in practice.
Furthermore, we can also extend the objective above in a similar way to other VAE variants such as Higgins et al.
; Zhao et al.
Tolstikhin et al.
(2017) where the KL divergence term is generalized to a weighted distance measure between the aggregated posterior distribution q φ ( z ) =
E x ∼ p data ( x ) q φ ( z | x ) and the prior p θ ( z ) , which gives rise to our ﬁnal objective:
[Equation] [Equation]
where D denotes a distance between distributions, which we choose the MMD distance as in Tolstikhin et al.
(2017), β is a regularization factor.
The training procedure of HBAE is similar to that of a conditional GAN, except that the conditioning variable is learned jointly.
To be concrete, given a ﬁxed encoder q φ and discriminator D , the conditional generator G is updated in the inner loop to increase the conditional discriminator’s score on
G ’s output.
In the outer loop, the encoder q φ and discriminator D are jointly updated adversarially against G .
During inference, we can sample from the generator as with other generative models.
However, in a HBAE, sampling follows a two step hierarchy.
One ﬁrst sample a latent variable z , either from the encoder distribution q φ ( z | x ) , or from the prior p θ ( z ) .
One then samples from a secondary z (cid:48) ∼ p ψ ( z (cid:48) | z ) to generate an example.
This sampling procedure resembles the inference process of a Hierarchical Bayes model, and hence we name it Hierarchical Bayes Autoencoder.
2. S ET -HBAE
Given the HBAE formulation, it is also easy to extend the same generative modeling setting to cases where inputs are sets of images (e.g., different views of the same object).
We refer to this variant as the Set-HBAE.
In order to do so, we explicitly express the input in the set format x = { x 1 , x 2 , ..., x k } where k denotes the size of the set.
By construction, we let the elements within a set x to be conditionally independent given the set latent code z (see Figure 1).
We then accordingly modify the HBAE objective Equation 7 as:
[Equation]
where | x | denotes the size of the set.
The encoder q φ in this case takes form of a set encoder.
Here we adopt the DeepSet (Zaheer et al., 2017) mechanism to average pool the codes of elements within the set; the form of the generator remains unchanged.
During inference time, the hierarchical sampling procedure naturally corresponds to sampling a set code z ﬁrst, followed by sampling an element code z (cid:48) .
In this way, the Set-HBAE is able to perform both stochastic set reconstruction, and well as unconditional set generation.
R ELATED W ORK
Autoencoding generative models .
The HBAE is similar to a handful of existing works which involve an encoder, decoder and a discriminator.
Notably, BiGAN (Donahue et al., 2016) trains a two channel discriminator that tries to separate the two joint distributions p ( z ) p ( x | z ) and p ( x ) p
( z | x ) , which involves an encoder and decoder, respectively.
BiGAN differs from the HBAE in that the mapping between z and x is unimodal.
Also, in a BiGAN, the encoder and decoder work jointly against the discriminator, whereas in the HBAE, the encoder and the discriminator work jointly against the generator.
The VAE-GAN (Larsen et al., 2015) is a VAE/GAN hybrid model that consists of the three modules.
In addition to the deocder being unimodal, a VAE-GAN deploys an unconditional discriminator, which contrasts with the conditional discriminator used in HBAE.
Conditional GANs .
A conditional GAN (cGAN) (Mirza & Osindero, 2014; Miyato & Koyama, 2018) modiﬁes the generator and discriminator of a GAN to be conditional on class labels.
The
HBAE can be considered a generalization of cGANs, as it learns the conditioning representation jointly with the generator and discriminator.
To this end, it is possible to come up with a semisupervised learning version of HBAE where a small portion of input data have observed labels, which are jointly trained with unlabelled data.
Refer to Figure 2 for a comparison of the HBAE with the cGAN, BiGAN and VAE-GAN architectures.
PixelCNN-based multimodal decoders .
Another line of work that employs a multimodal encoder is represented by PixelVAE (Gulrajani et al., 2016), which uses a PixelCNN-based decoder (Van den
Oord et al., 2016).
While similar in spirit to our work, PixelCNN suffers from capturing long term dependency of images, and is also expensive during inference time.
We see HBAE as providing an alternative approach to this line of work, which leverages the power of state of the art GANs for decoding.
VAEs with large latent codes .
Another design philosophy of VAEs is represented by VQ-VAE (Razavi et al., 2019), where a large latent code is used for each example, with preserved spatial dimensions.
VQ-VAE2 then learns a powerful prior in the form of a PixelCNN, which is able to generate good quality samples at inference time.
We note that in this scenario, the use of a unimodal decoder is well justiﬁed where the latent space is large enough to allow learning a oneto-one mapping between a latent code and an input.
The HBAE on the other hand, works best in the setup where a compact latent code is used, or in applications where explicit multimodality is desired, such as with sets as inputs.
Generative Modeling of sets .
Generative modeling of sets is a challenging task.
Hierachical Bayes models have previously been widely applied to document modeling, which can be considered as a set of words, represented by the LDA model (Blei et al., 2003).
However, due to the discrete nature of words, LDA only needs to rely on a hierarchy of simple distributions such as the Dirichlet and multinomial distributions.
Recently, there has been a a body of work dedicated to point cloud generative modeling, where sets are composed of 3D points in Euclidean space (Li et al., 2018; Achlioptas et al., 2017).
However, sets considered in our work contain much richer structure than point clouds, which contain very little structure per element.
It is thus unclear how methods developed from the point cloud or document modeling lines of work generalize to more complex data.
E XPERIMENTS 4. A RCHITECTURE
The HBAE deﬁned in Equation 7 has three trainable modules: an encoder which corresponds to q φ ( z | x ) 1 , a conditional discriminator D ( x | z ; θ ) , and a conditional generator G ( z (cid:48) | z ; ψ ) .
Our architecture design largely mimics that of a conditional GAN, in particular the projection-based discriminator proposed in Miyato & Koyama (2018).
Here the latent code, z , takes place of the label, y , in a conditional GAN, and is jointly learned with other modules instead of being given.
We follow a simple design choice for the architectures, where both the encoder and discriminator consist of stacks of residual blocks with Spectral Normalization (Miyato et al., 2018), and the generator consists of stacks of residual blocks with Batch Normalization (Ioffe & Szegedy, 2015).
Refer to Figure 3 for more details.
4. T RAINING PROTOCOL
In all the experiments, we follow the training protocols of standard GANs as in Donahue et al.
(2016), where the encoder and discriminator are trained jointly against the generator, as shown in Equation 7.
We use Adam with a learning rate 2e-4 for the generator , and learning rate 4e4 for the encoder and discriminator.
The moment terms are set as {} for both optimizers.
In practice, one failure mode we have found by directly following Equation 7 is that the encoder can collapse to a constant mapping which makes HBVE regress to an unconditional GAN.
As a remedy, we put an additional reconstruction regularization term on the encoder’s parameters
[Equation]
φ ( z | x ) (cid:107) x − G ( 0 | z ) (cid:107) 1 (for the Set-HBAE, this corresponds to mapping z to all elements within the set: λ E x ∼ p data ( x ) E z ∼ q φ ( z | x ) 1 | x | (cid:80) x i ∈ x (cid:107) x
i − G ( 0 | z ) (cid:107) 1 ).
Here λ is the regularization strength which we set as 1 .
We note that this regularization term is not essential for
HBAE to train, but in most cases helps with the stability.
4. S INGLE IMAGE INPUTS
In the ﬁrst set of experiments, we test the HBAE in the standard setting where inputs are single images.
We use the CelebA Dataset (Liu et al., 2015) with resolution 64 × 64 .
In Figure 4, we show some of the qualitative results.
We see that the HBAE is able to reconstruct input examples in a semantically meaningful way by preserving certain aspects of the input face, while outputting variations that exhibit rich and changing details, such as adding or removing glasses, varying expressions subtly, or changing hairstyles.
When sampled unconditionally, the HBAE is also able to produce sharp examples.
As control experiments, we have accordingly trained a GAN and VAE using the same architectures (encoder and generator), and reported the FID score of the unconditional generations, and reconstructions (when applicable), in Table 1.
We have also included the reconstruction
MSE of the VAE and HBAE evaluated in the pixel space and Inception V3’s feature space.
Refer to the caption for detailed discussions.
Additionally, we show that the HBAE can be used to sample images in a Markov chain, exploring the data manifold around an example as a random walk.
The Markov chain starts from a real example, and recursively passes it through the encoder and the generator.
We show two such chains in Figure 5, starting from the same test examples, both of which exhibit smooth transitions from one sample to the next, without burning into a single mode.
Interestingly, the two chains also lead to different sampling paths, due to the multimodality of the generator.
This behavior is not possible with regular
VAE or GAN formulations.
4. S ET OF IMAGES AS INPUT
In the second set of experiments, we change the inputs to be a set of images.
The ﬁrst dataset we use is the ShapeNet dataset from Chang et al.
(2015), which consists of 3d objects with projected 2d views.
We use a subset which contains 13 popular categories, namely airplane, bench, cabinet, car, chair, monitor, lamp, speaker, gun, couch, table, phone, ship .
The training set consists of 32,
[Table] [Table]
unique objects and 788,088 images in total.
The second dataset we use is the VGGFace2 dataset from Cao et al.
(2018), which is a face dataset containing 9,131 subjects and 3.31M images.
For both datasets, we randomly construct ﬁxed-size sets of 8 elements.
For ShapeNet, this is done by randomly selecting an object and 8 of its views; for VGGFace2, we similarly select one identity and 8 of her/his faces.
A mini-batch is constructed by selecting N (which we set to 32) objects/subjects, which amounts to Nx 8 images.
Example set inputs together with their stochastic reconstructions are shown in Figure 6.
The qualitative results indicate that the samples are diverse, yet are consistent with each other and tend to resemble the semantic traits of their respective input sets (i.e., same object in different views for the ShapeNet case and same identity with varying attributes for the
VGGFace2 case).
As a further quantitative measure, we also conducted few shot classiﬁcation experiments on
ShapeNet with the learned features.
We compare the features learned from our Set-HBAE with a control VAE, as well as with randomly initialized features, and with supervised training results, all with the same encoder architecture.
For all settings, we trained a linear classiﬁer with crossentropy loss, without any explicit data augmentation or regularization (except for the supervised learning setting where the encoder is jointly trained with the linear classiﬁer).
We vary the number of training images from 64 to 2048, and report the results in Table 2.
We see that in the few shot classiﬁcation domain, the learned features derived from the Set-HBAE largely outperform the supervised learning approach as well as the VAE results, all of which far exceed random features.
Finally, we show the Set-HBAE’s ability to generate unconditional samples.
The generated results are shown in Figure 7 for both ShapeNet and VGGFace2.
Notably, the Set-HBAE is able to generate diverse sets, where elements within the generated sets are also diverse but consistent in ways that preserves object shape or face attributes.
To the best of our knowledge, this is the ﬁrst generative model that is able to generate complex sets with richly varying yet coherent structure.
C ONCLUSION
In this paper, we presented the HBAE, a novel probabilistic generative model with an EBM-based multimodal decoder.
In our implementation, we jointly trained the encoder with variational inference and the EBM decoder with a powerful conditional GAN.
We demonstrated that the decoder is capable of generating diverse conditional and unconditional samples.
Then we introduced the
Set-HBAE, a novel generative model for sets of examples, generalizing the HBAE framework.
The
Set-HBAE takes sets as input and can sample plausible variations that are consistent with the conditioning set, but generalize beyond the speciﬁc input instances.
In addition, we show that the
Set-HBAE performs well on a ShapeNet Chang et al.
(2015) few-shot classiﬁcation task, showing the advantage of leveraging generative models in the small data regime.
The ability to sample plausible variations conditioned on an input example or set of examples is a key building block for enabling predictive models, for which we see future applications in areas such as video (Lee et al., 2018), 3D scene generation (Eslami et al., 2018), and model-based reinforcement learning (Ha &
Schmidhuber, 2018).
A A PPENDIX