summaries: The paper considers Gaussian VAEs and their tendency to suffer from posterior collapse. In particular, the authors analyse the impact of the usually fixed covariance $\sigma_x$ of the decoder Gaussian on the learned encoder variance. They show that the former can be seen as a regulariser for the latter and therefore impacts the "smoothness" of the encoder. The authors hypothesize that a large value of $\sigma_x$ causes posterior collapse as a consequence.

summaries: This paper studies the Gaussian VAE and figures out that the decoder variance regularizes the VAE and affects the model smoothness, and an inappropriate estimation of this parameter would raise posterior collapse, which is supported by theoretical analysis and empirical demonstrations. Hence, this paper then proposes an ELBO with adaptive decoder variance to avoid oversmoothing the model. Overall, the idea is interesting and provides some new insights for our community. The major concerns regarding this paper are listed as below.

summaries: This paper analyses the "posterior collapse" phenomenon observed in training latent variable models (in particular Variational Auto-Encoders), and propose a new training objective to remedy the problem. The theoretical analysis of the authors suggests that the posterior collapse is induced by an inappropriate choice of variance in the decoder distribution. The new objective they propose, the AR-ELBO, jointly optimises this variance along with the usual network parameters. The authors demonstrate that their objective yields relatively good results on image modelling, compared to other standard VAE methods.

abstractText: Variational autoencoders (VAEs) often suffer from posterior collapse, which is a phenomenon that the learned latent space becomes uninformative. This is related to local optima of the objective function that are often introduced by a fixed hyperparameter resembling the data variance. We suggest that this variance parameter regularizes the VAE and affects its smoothness, which is the magnitude of its gradient. An inappropriate choice of this parameter causes oversmoothness and leads to posterior collapse. This is shown theoretically by analysis on the linear approximated objective function and empirically in general cases. We propose AR-ELBO, which stands for adaptively regularized ELBO (Evidence Lower BOund). It controls the strength of regularization by adapting the variance parameter, and thus avoids oversmoothing the model. Generation models trained by proposed objectives show improved Fréchet inception distance (FID) of images generated from the MNIST and CelebA datasets.

authors:
Alexander A Alemi
Ben Poole
Ian Fischer
Joshua V Dillon
Rif A Saurous
Kevin Murphy

title: Fixing a broken ELBO

venue: arXiv preprint arXiv:1711.00464,

authors:
Samuel R Bowman
Luke Vilnis
Oriol Vinyals
Andrew M Dai
Rafal Jozefowicz
Samy Bengio

title: Generating sentences from a continuous space

venue: arXiv preprint arXiv:1511.06349,

authors:
Xi Chen
Yan Duan
Rein Houthooft
John Schulman
Ilya Sutskever
Pieter Abbeel

title: InfoGAN: Interpretable representation learning by information maximizing generative adversarial nets

venue: In Proc. Advances in Neural Information Processing Systems (NeurIPS),

authors:
Bin Dai
David Wipf

title: Diagnosing and enhancing VAE models

venue: In Proc. International Conference on Learning Representation (ICLR),

authors:
Bin Dai
Ziyu Wang
David Wipf

title: The usual suspects? Reassessing blame for VAE posterior collapse

venue: In Proc. International Conference on Machine Learning (ICML),

authors:
Hao Fu
Chunyuan Li
Xiaodong Liu
Jianfeng Gao
Asli Celikyilmaz
Lawrence Carin

title: Cyclical annealing schedule: A simple approach to mitigating KL vanishing

authors:
Partha Ghosh
Mehdi SM Sajjadi
Antonio Vergari
Michael Black
Bernhard Schölkopf

title: From variational to deterministic autoencoders

authors:
Ian Goodfellow
Jean Pouget-Abadie
Mehdi Mirza
Bing Xu
David Warde-Farley
Sherjil Ozair
Aaron Courville
Yoshua Bengio

title: Generative adversarial nets

authors:
Ishaan Gulrajani
Faruk Ahmed
Martin Arjovsky
Vincent Dumoulin
Aaron C Courville

title: Improved training of Wasserstein GANs

authors:
Junxian He
Daniel Spokoyny
Graham Neubig
Taylor Berg-Kirkpatrick

title: Lagging inference networks and posterior collapse in variational autoencoders

authors:
Martin Heusel
Hubert Ramsauer
Thomas Unterthiner
Bernhard Nessler
Sepp Hochreiter

title: GANs trained by a two time-scale update rule converge to a local Nash equilibrium

authors:
Irina Higgins
Loic Matthey
Arka Pal
Christopher Burgess
Xavier Glorot
Matthew Botvinick
Shakir Mohamed
Alexander Lerchner

title: beta-VAE: Learning basic visual concepts with a constrained variational framework

authors:
Sergey Ioffe
Christian Szegedy

title: Batch normalization: Accelerating deep network training by reducing internal covariate shift

authors:
Michael I Jordan
Zoubin Ghahramani
Tommi S Jaakkola
Lawrence K Saul

title: An introduction to variational methods for graphical models

venue: Machine Learning,

authors:
Diederik P Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization

authors:
Diederik P Kingma
Max Welling

title: Auto-encoding variational Bayes

authors:
Yann LeCun
Léon Bottou
Yoshua Bengio
Patrick Haffner

title: Gradient-based learning applied to document recognition

venue: Proc. IEEE,

authors:
Ziwei Liu
Ping Luo
Xiaogang Wang
Xiaoou Tang

title: Deep learning face attributes in the wild

venue: In Proc. IEEE Conference on Computer Vision and Pattern Recognition (CVPR),

authors:
James Lucas
George Tucker
Roger B Grosse
Mohammad Norouzi

title: Don’t blame the ELBO! A linear VAE perspective on posterior collapse

authors:
Xuezhe Ma
Chunting Zhou
Eduard Hovy

title: MAE: Mutual posterior-divergence regularization for variational autoencoders

authors:
Laurens van der Maaten
Geoffrey Hinton

title: Visualizing data using t-SNE

venue: Journal of Machine Learning Research,

authors:
Alireza Makhzani
Jonathon Shlens
Navdeep Jaitly
Ian Goodfellow
Brendan Frey

title: Adversarial autoencoders

venue: In Proc. International Conference on Learning Representation (ICLR) Workshop,

authors:
Takeru Miyato
Toshiki Kataoka
Masanori Koyama
Yuichi Yoshida

title: Spectral normalization for generative adversarial networks

authors:
Rogan Morrow
Wei-Chen Chiu

title: Variational autoencoders with normalizing flow decoders

venue: arXiv preprint arXiv:2004.05617,

authors:
Ali Razavi
Aäron van den Oord
Ben Poole
Oriol Vinyals

title: Preventing posterior collapse with delta-VAEs

authors:
Ali Razavi
Aaron van den Oord
Oriol Vinyals

title: Generating diverse high-fidelity images with VQ-VAE-2

authors:
Shibani Santurkar
Dimitris Tsipras
Andrew Ilyas
Aleksander Madry

title: How does batch normalization help optimization

authors:
Huajie Shao
Shuochao Yao
Dachun Sun
Aston Zhang
Shengzhong Liu
Dongxin Liu
Jun Wang
Tarek Abdelzaher

title: Controllable variational autoencoder

authors:
Casper Kaae Sønderby
Tapani Raiko
Lars Maaløe
Søren Kaae Sønderby
Ole Winther

title: Ladder variational autoencoders

authors:
Ilya Tolstikhin
Olivier Bousquet
Sylvain Gelly
Bernhard Schoelkopf

title: Wasserstein autoencoders

authors:
Aäron van den Oord
Oriol Vinyals
Koray Kavukcuoglu

title: Neural discrete representation learning

authors:
Jiacheng Xu
Greg Durrett

title: Spherical latent spaces for stable variational autoencoders

authors:
Shengjia Zhao
Jiaming Song
Stefano Ermon

title: InfoVAE: Balancing learning and inference in variational autoencoders

venue: In Proc. AAAI Conference on Artificial Intelligence (AAAI),

venue: (Fu et al.,

heading: 1 INTRODUCTION

text:
The variational autoencoder (VAE) framework (Kingma & Welling, 2014; Higgins et al., 2017; Zhao et al., 2019) is a popular approach to achieve generative modeling in the field of machine learning. In this framework, a model that approximates the true posterior of observation data, is learned by a joint training of encoder and decoder, which creates a stochastic mapping between the observation data and the learned deep latent space. The latent space is assumed to follow a prior distribution. The generation of a new data sample can be done by sampling the latent space and passing the sample through the decoder. It is common to assume that both the prior on the latent space and the posterior of the observation data follow a Gaussian distribution. This setup is also known as the Gaussian VAE. In this case, the variance of the decoder output is usually modeled as an isotropic matrix σ2xI with a scalar parameter σ2x ≥ 0. Furthermore, in order to deal with the intractable log-likelihood of the true posterior, the evidence lower bound (ELBO) (Jordan et al., 1999) is adopted as the objective function instead.
While VAE-based generative models are usually considered to be more stable and easier to train than generative adversarial networks (Goodfellow et al., 2014), they often suffer from the problem of posterior collapse (Bowman et al., 2015; Sønderby et al., 2016; Alemi et al., 2017; Xu & Durrett, 2018; He et al., 2019; Razavi et al., 2019a; Ma et al., 2019), in which the latent space has little information of the input data. The phenomenon is generally mentioned as “the posterior collapses to the prior in the latent space” (Razavi et al., 2019a). Recently, several works have suggested that the variance parameter σ2x in the ELBO is strongly related to posterior collapse. For example, Lucas et al. (2019) analyzed posterior collapse through the analysis on a linear VAE. It revealed that an inappropriate choice of σ2x will introduce sub-optimal local optima and cause posterior collapse. Moreover, Lucas et al. (2019) reveals that contrary to the popular belief, these local optima are not introduced by replacing the log-likelihood with the ELBO, but by an excessively large σ2x. On the other hand, it can be shown that fixing σ2x to an excessively small value leads to under-regularization of the decoder, which can cause overfitting. However, in most implementations of a Gaussian VAE, the variance parameter σ2x is a fixed constant regardless of the input data and is usually 1.0. In another work, Dai & Wipf (2019) proposed a two-stage VAE and treated σ2x as a training parameter. Besides the inappropriate choice of the variance parameter, posterior collapse can also induced by other causes. For example, Dai et al. (2020) found that, small nonlinear perturbation introduced in
the network architecture can also result into extra sub-optimal local minima. However, in this work we will keep our focus on the variance parameter.
We suggest that σ2x affects the strength of regulation over the gradient magnitude of the decoder. We call the expected gradient magnitude the smoothness throughout this paper. The smaller the gradient magnitude, the smoother the model. In particular, we would like to focus on the local smoothness of the model, which is the smoothness evaluated within the neighborhood of the encoded latent variable of the observation data. Thus, we begin with the following hypothesis:
Main Hypothesis. The value of σ2x controls the regularization strength of the smoothness of the decoder. Therefore, an excessively large σ2x causes oversmoothness, which results in posterior collapse.
Following the hypothesis, the estimation of σ2x should be related to properties of the approximated posterior of the latent space, such as its local smoothness. We will start with analyzing how σ2x regularizes the local smoothness of the stochastic decoder and then propose new objective functions that inherently determine σ2x via maximum likelihood estimation (MLE). This proposed objective function is named AR-ELBO (Adaptively Regularized ELBO), which controls the regularization strength via σ2x. Furthermore, several variations are derived for different parameterizations of variance parameters.
Our main contributions are listed as follows:
1. We show that our main hypothesis holds for linear approximated ELBO and empirically holds in the general case in Section 3. This also suggests that the variance parameter σ2x should be estimated from properties of the approximated posterior instead of being treated as a hyperparameter. 2. We propose the AR-ELBO, an ELBO-based objective function that adaptively regularizes the smoothness of the decoder by MLE of the variance parameter σ2x in Section 4. Variations of AR-ELBO for several variance parameterizations of posterior distributions are also derived. AR-ELBO prevents the model from the posterior collapse induced by oversmoothing and improves the quality of generation, which is shown in Section 5.
The organization of this paper is as follows. In Section 2, we propose a mathematical definition of posterior collapse in the form of mutual information. This also includes the conventional definition: “the posterior collapses to the prior in the latent space”. In Section 3, the theoretical analysis and empirical support of the main hypothesis are given. We perform an analysis showing that σ2x affects the smoothness of the decoder via variance parameters of the latent space learned by the encoder. In Section 4, we propose new AR-ELBO objective functions for various variance parameterizations of posterior distributions, which can relieve the decoder from being oversmoothed in the training and prevent posterior collapse. These objective functions no longer include any hyperparameters and can adaptively estimate the variance parameter σ2x from the observation data. It should be noted that if we adaptively determine σ2x with the proposed AR-ELBO, the strength of regularization of the decoder smoothness will gradually decrease as training progresses. In Section 5, we conduct an experiment on the MNIST and CelebA datasets, which shows that utilizing the proposed AR-ELBO with the standard Gaussian VAE can be competitive with many other variations of VAE models in most situations.
Throughout this paper, we use a, a and A for a scalar, a column vector and a matrix, and ln and log denote the natural logarithm and common logarithm, respectively. Our code is available from the following URL1.

heading: 2 POSTERIOR COLLAPSE IN GAUSSIAN VAE

text:
We begin with the standard formulation of the Gaussian VAE, which is the foundation of our research. A definition of posterior collapse is proposed by using mutual information (MI).
1URL hidden due to blind review.

heading: 2.1 GAUSSIAN VAE

text:
Consider a data space X ⊂ Rdx and a sample set {xi}Ni=1 ⊂ X , where xi ∼ pdata(x). The empirical distribution p̃data(x) on X can be evaluated by p̃data(x) = 1N ∑N n=1 δ(x − xn), where δ(·) denotes the Dirac delta function. In the standard VAE framework, a latent space Z ⊂ Rdz is learned and the sampled latent variables z ∈ Z are used to generate data samples x′ ∈ X . Let qφ(z|x) and pθ(x|z) denote the stochastic encoder and decoder, respectively. Trainable parameters of the two neural networks are denoted as φ and θ. The decoder generates data samples by pθ(x) := Ep(z)[pθ(x|z)], where p(z) is the prior distribution onZ . The encoder and decoder are jointly trained by minimizing the following objective function:
L = −Epdata(x) [ln pθ(x)] + Epdata(x)DKL (qφ(z|x) ‖ pθ(z|x)) + Epdata(x) [ln pdata(x)] = DKL (pdata(x) ‖ pθ(x)) + Epdata(x)DKL (qφ(z|x) ‖ pθ(z|x)) . (1)
This objective function was derived in Zhao et al. (2019), which represents everything in the form of Kullback–Leibler divergence, and is equivalent to ELBO maximization up to an additive constant.
In the context of the Gaussian VAE, the encoder and decoder are assumed to satisfy
qφ(z|x) = N (z|µφ(x),diag(σ2φ(x))) and pθ(x|z) = N (x|µθ(z), σ2xI). (2)
The prior p(z) is also assumed to be the Gaussian distribution as p(z) = N (z|0, I). Substituting (2) into (1) while omitting terms independent of θ and φ leads to the following objective:
J̃σ2x(θ, φ) = Ep̃data(x) [ 1
2σ2x Eqφ(z|x)[‖x− µθ(z)‖ 2 2] +DKL(qφ(z|x) ‖ p(z))
] , (3)
which can be interpreted as the sum of the expected values of the reconstruction loss and the regularization term. In the case of a Gaussian prior and posterior, the regularization term is equal to 1 2 ∑dz i=1(σ 2 φ,i(x) + µφ,i(x) 2 − log σ2φ,i(x)− 1).

heading: 2.2 POSTERIOR COLLAPSE

text:
Posterior collapse is a major problem, where the encoder learns to map inputs to the latent space while ignoring the data distribution. In this phenomenon, the MI between input data and reconstructed data through the encoder-decoder path is reduced because the latent space has less information about the data distribution. Here, we suggest the following definition of posterior collapse. Definition 1. Posterior collapse is defined as the MI I(x; x′) becoming nearly zero, where x′ := µθ(z) with z ∼ qφ(z|x).
In many works (Bowman et al., 2015; Sønderby et al., 2016; Alemi et al., 2017; He et al., 2019; Razavi et al., 2019a), the phenomenon denoted as posterior collapse has been mathematically represented as Epdata(x)DKL(qφ(z|x) ‖ p(z)) → 0, which we hereafter refer to as KL collapse (Xu & Durrett, 2018). However, posterior collapse is not always caused by the diminished KL divergence, i.e., posterior collapse with Epdata(x)DKL(qφ(z|x) ‖ p(z)) 6= 0 can occur. The proposed definition of posterior collapse includes KL collapse from the following theorem, which is proven in Appendix A. Theorem 2. I(x; x′)→ 0 as Epdata(x)DKL(qφ(z|x) ‖ p(z))→ 0 holds for any pθ(x|z).
In Appendix E, it is demonstrated that posterior collapse can happen even if the KL divergence is nonzero when the posterior variance is fixed in Z .
3 VARIANCE PARAMETER σ2x AND THE LOCAL SMOOTHNESS
In this section, we provide mathematical and empirical support of the main hypothesis. Throughout this section, we adopt the following parameterization for the encoder for simplicity: qφ,σ2z (z|x) = N (z|µφ(x), σ2zI), where the variance is parameterized as an isotropic matrix unlike the conventional VAE. A similar analysis on the conventional VAE can be found in Appendix B. It begins with showing that the choice of σ2x affects the convergence point of σ 2 z , which is the variance parameter of the latent space. Then, we show that σ2z acts as the weight of the gradient penalty, which is
implicitly included in (3). This supports the main hypothesis that the over-regulation imposed by a large σ2x via σ 2 z causes the oversmoothness of the decoder and leads to posterior collapse. It is also empirically supported by observing the tendencies of the convergence point of σ2z , the smoothness and the MI I(x,x′). Ultimately, these items of evidence motivated us to develop a method that adapts σ2x to prevent oversmoothing of the decoder.
3.1 REGULARIZATION EFFECT OF σ2x IN LINEAR APPROXIMATED ELBO
The effect of σ2x on the convergence point of the variance parameter σ 2 z can be observed from two extreme cases, σ2x → 0+ and σ2x → ∞. In the first case, J̃σ2x reduces to Ep̃data(x)Eqφ,σ2z (z|x)[‖x − µθ(z)‖ 2 2], and it becomes zero only if σ 2 z = 0. In the second case, J̃σ2x reduces to Ep̃data(x)DKL(qφ,σ2z (z|x) ‖ p(z)), and σ 2 z becomes 1 at the minimum point, from DKL(qφ,σ2z (z|x) ‖ p(z)) = dz 2 (σ 2 z − log σ2z − 1) + ‖µφ(x)‖22. This shows that a small σ2x makes σ2z converge to a value around 0, while a large σ2x makes σ 2 z converge to a value around 1.
If σ2z is sufficiently small, as the training progresses to a certain extent, the perturbed decoding process µθ(z + z) around z = µφ(x) with z ∼ N ( z|0, σ2zI) can be approximated as a linear function. The ELBO can be approximated as follows by using the linear approximation of µθ(·) and omitting terms independent of θ and φ:
J̃σ2x(θ, φ, σ 2 z) ≈
1
2σ2x Ep̃data(x)
[ ‖x−µθ(µφ(x))‖22 +σ 2 z ‖∇µθ(µφ(x))‖ 2 F +2σ 2 x ‖µφ(x)‖ 2 2 ] . (4)
In the approximation above, ‖ · ‖F is the Frobenius norm and σ2z is treated as a function parameter. The derivation of the above approximation can be found in Appendix B. Equation (4) decomposes the objective function into three terms: a reconstruction error term, gradient penalty term and L2 regularization term. As one can see from (4), σ2z regularizes the smoothness of the decoder by penalizing its gradient norm in training. Although the linear approximation above is derived for the simplified VAE parameterization, we also provide the linear approximation of the ELBO for the standard VAE parameterization (2) in Appendix B, where the second term in (4) becomes a weighted gradient penalty.
Summarizing the above observations shows that σ2x affects the smoothness via σ 2 z , while σ 2 z directly regularizes the smoothness. This means that if σ2x is excessively large, it will cause overregularization of the decoder and suppress I(z,x′)(≥ I(x,x′)), which finally leads to posterior collapse. To avoid such over-regularization, σ2x and σ 2 z should be determined appropriately. In addition, an experiment shows that posterior collapse can be triggered by directly manipulating σ2z , as discussed in Appendix E.

heading: 3.2 EMPIRICAL STUDY ON SMOOTHNESS OF DECODER IN THE GENERAL CASE

text:
Section 3.1 shows the impact of σ2x on the regularization of the decoder smoothness through the linear approximated objective function. To support the main hypothesis in the general case, an experiment on the MNIST dataset (LeCun et al., 1998) is conducted. Several criteria are accessed to provide evidence for the regularization effect of σ2x on the decoder smoothness and its consequential effect on MI I(x,x′). To confirm that σ2x affects the smoothness via σ2z , we conduct the experiment for two cases: stochastic encoding and deterministic encoding. While the stochastic encoder qφ,σ2z (z|x) is used in the former case, a VAE equipped with a deterministic encoder, i.e., σ2z is fixed to zero during the training, is investigated in the latter case. Observing the difference between the two cases provides empirical support for Section 3.1. To investigate the relation between σ2x and the smoothness of the decoder clearly, common generalization techniques such as batch normalization (Ioffe & Szegedy, 2015; Santurkar et al., 2018) and weight decay are not used.
Criteria We used several criteria to observe the impact of σ2x in the experiment, such as the reconstruction error (MSE), KL divergence value Ep̃data(x)DKL(qφ(z|x) ‖ p(z)) and the final converged value of σ2z . In addition, we also estimate the local smoothness of the decoder and the MI between x and x′, denoted as I(x,x′). To access the local smoothness, the expected local smoothness (ELS) is introduced, which is the lower bound of the Lipschitz constant of the decoder. Consider a sample that is decoded with perturbation µθ(µφ(x) + z), where the perturbation follows a zero-mean
Guassian distribution with variance s2z , z ∼ N ( z|0, s2zI). Let z and ′z be i.i.d. random variables. We define the expected gap ∆2(s2z) as
∆2(s2z) := Epdata(x)EN ( z|0,s2zI)N ( ′z|0,s2zI)[∆ 2(x, z, ′ z)] (5)
with ∆2(x, z, ′z) := ‖µθ(µφ(x) + z)− µθ(µφ(x) + ′z)‖22.
As s2z decereases, the ratio ∆ 2(s2z)/(2s 2 z) converges and becomes an indicator of Ep̃data(x)[‖∇µθ(µφ(x))‖2F ], which is regularized by σ2z in (4). Therefore, we can now define the ELS as
Ep̃data(x)[‖∇µθ(µφ(x))‖ 2 F ]. (6)
Further details of the ELS are described in Appendix C. As a reference, we estimate the upper bound of I(x,x′), which is I(x′; z), by Monte Carlo estimation. Results Table 1 summarizes the results for different σ2x. In the stochastic encoding case, a larger σ2x consistently leads to a larger σ2z . This results in a smaller expected gap, a smaller ELS and a lower upper bound of MI. This supports the main hypothesis that a larger σ2x makes the decoder smoother. In the case of σ2x = 1.0, all the criteria except MSE become nearly zero, where KL collapse and posterior collapse both occur due to the over-regularization of the smoothness of the latent space. On the other hand, in the deterministic encoding case, the ELS increases with σ2x. This is because σ2x does not directly regularize the decoder via the gradient penalty as in (4). As a result, the MI upper bound does not shrink to zero even if σ2x = 1.0, where posterior collapse occurs in the case of stochastic encoding. The difference in the results between the two cases clearly suggests that a large σ2x triggers the oversmoothness via σ 2 z , which is consistent with the discussion in Section 3.1. These results provide empirical support of the main hypothesis as well as the discussion in Section 3.1. Further details and examples of images are shown in Appendix D.
3.3 DIFFICULTY OF CHOOSING AN APPROPRIATE σ2x
According to previous sections, a large σ2x will cause oversmoothness. Therefore, we consider the case of fixing σ2x to a sufficiently small value to avoid the problem. We arrive at the following theorem, whose proof can be found in Appendix F:
Theorem 3. Consider the global optimum of Jσ2x(θ, φ, σ 2 z) w.r.t. a given σ 2 x. If σ 2 x → 0, then σ2z → 0.
In Theorem 3, Jσ2x(θ, φ, σ 2 z) is optimized on the basis of the true data distribution instead of the empirical data distribution. According to the theorem, σ2z converges to zero as σ 2 x approaches zero, which leads to zero gradient penalty for the decoder as the VAE training progresses. In practice, we have no access to pdata(x), but we have access to the empirical distribution p̃data(x). Theorem 3 is satisfied even when pdata(x) is replaced with p̃data(x). In this case, where σ2x is chosen to be small, the optimization process of J̃σ2x will fit pθ,σ2x(x) to the empirical distribution p̃data(x), which usually results in overfitting.
As shown above, choosing an appropriate σ2x that avoids both oversmoothness and overfitting is nontrivial. Moreover, it is likely that σ2x should be adapted depending on the status of training. Therefore, it is intuitive to adapt σ2x instead of fixing it, which will be described in the next section.

heading: 4 ADAPTIVELY REGULARIZED ELBO

text:
A modified ELBO-based objective function is proposed, which can be interpreted as an implicit update scheme that simultaneously updates σ2x and the rest of the parameters. We also derive corresponding objective functions for models with different variance parameterizations.
4.1 ELBO WITH ADAPTIVE σ2x
In this subsection, we newly optimize the VAE objective function (1) w.r.t. all the parameters including σ2x, which is usually fixed in existing implementations. Following the process of establishing (3) but keeping the terms related to σ2x, we arrive at:
J̃ (θ, φ, σ2x) = Ep̃data(x) [ 1
2σ2x Eqφ(z|x)
[ ‖x− µθ(z)‖22 ] +DKL (qφ(z|x) ‖ p(z)) ] + dx 2 lnσ2x. (7)
From the partial derivative of J̃ w.r.t. σ2x, the MLE of σ2x, denoted as σ̂2x, can be evaluated with the other parameters fixed. Then, the ordinary network parameters θ and φ can be updated by optimizing (7) with the variance σ2x fixed. This combination of MLE and the alternative update between (θ, φ) and σ2x guarantees that (i) if θ and φ are fixed, then there exists σ̂ 2 x such that J̃ (θ, φ, σ̂2x) ≤ J̃ (θ, φ, σ2x) and (ii) for the σ̂2x obtained in the previous step, there exist θ̂ and φ̂, such that J̃ (θ̂, φ̂, σ̂2x) ≤ J̃ (θ, φ, σ̂2x). In this respect, the convergence of the optimization is assured and the parameter σ2x is kept as the MLE during the whole training stage. This inspired us to develop a weight scheduling scheme for σ̂2x, leading to a modified ELBO-based objective function. Consider the trainable network parameters (θ, φ) and the variance parameter σ2x. The update of the objective J̃ (θ, φ, σ2x) is divided as
σ2(t+1)x = 1
dx Ep̃data(x)Eqφ(t) (z|x)
[ ‖x− µθ(t)(z)‖ 2 2 ] (8a)
θ(t+1), φ(t+1) = arg min θ,φ J̃ σ 2(t+1) x (θ, φ), (8b)
where t is the iteration index. The step updating (θ, φ) is the same as that in the standard VAE; the step updating σ2x in (8a) can be interpreted as determining an appropriate balance between the reconstruction error and the KL term in J̃σ2x(θ, φ). As the learning progresses, the parameter σ 2 x will decrease along with the MSE Ep̃data(x)Eqφ(z|x)[‖x−µθ(z)‖22], which is consistent with the discussion in Dai & Wipf (2019).
Proposed objective function (AR-ELBO) The update scheme above can be further simplified by substituting (8a) into (7), which converts J̃ (θ, φ, σ̂2x) into
J̃AR(θ, φ) = dx 2
lnEp̃data(x)Eqφ(z|x) [ ‖x− µθ(z)‖22 ] + Ep̃data(x)DKL (qφ(z|x) ‖ p(z)) , (9)
where all constant terms w.r.t. the parameters are omitted. Optimizing (9) also makes σ2x remain as the MLE during the VAE training. Moreover, (9) is equivalent to the standard Gaussian VAE plus weight balancing with (8b). This relieves the VAE from the problem of imbalance between the KL divergence term and the reconstruction loss. Also, as stated by Theorem 3, decreasing σ2x also decreases σ2z . This gradually relieves the regularization of the ELS (6), which can be observed from (4). However, this eventually diminishes the gradient penalty; therefore, we suggest using earlystopping and learning rate scheduling to deal with this situation, which can give the decoder both appropriate smoothness and generalization capability.

heading: 4.2 OBJECTIVES FOR VARIOUS PARAMETERIZATIONS

text:
In the standard VAE given by (2), the variance of the decoded distribution on X , denoted as Σx, is modeled as an identity matrix, i.e., Σx = σ2xI. In this case, σ 2 x is simply a scalar value and the reconstruction objective is the same as conventional MSE and is minimized as in (9). However, out of
curiosity, we would like to explore three other variance parameterizations in addition to (2) and derive corresponding reconstruction objectives for these cases, in which the reconstruction objectives are no longer equal to MSE.
In fact, the variance Σx can not only be parameterized by an isotropic/diagonal matrix but also be chosen to be independent or dependent on z. We denote these in Table 2 as Iso-I (IsotropicIndependent), Iso-D (Isotropic-Dependent), Diag-I (Diagonal-Independent) and Diag-D (DiagonalDependent). The first case, Iso-I, corresponds to the standard variance model Σx = σ2xI. For these parameterizations, the corresponding objectives may be summarized as
J̃AR(θ, φ, Σ̂x) = J̃rec(θ, φ, Σ̂x) + Ep̃data(x)DKL (q(z|x) ‖ p(z)) (10a)
J̃rec(θ, φ,Σx) = 1
2 Ep̃data(x)Eq(z|x)
[ trace ( Σ−1x (x− µθ(z))(x− µθ(z))> ) + ln |Σx| ] . (10b)
By evaluating the partial derivative of J̃rec w.r.t. Σx, i.e., using its MLE, the reconstruction loss J̃rec corresponding to each case can be derived. All the derivations can be found in Appendix G. The final reconstruction objectives with different parameterizations of Σx are listed in Table 2. It is interesting to note that the reconstruction error for each dimension in the data space has to be calculated separately for Diag-D; meanwhile, only MSE for the whole minibatch is needed in Iso-I. Considering the optimization stability in practical situations, we suggest adding a small constant, e.g., 10−6, before taking the logarithm except for in the case of Iso-I.
Although the proposed objective functions are capable of determining Σx appropriately, it should be noted that a gap still exists between the prior p(z) and the aggregated posterior qφ(z) obtained by the proposed methods. The cause can be observed from the reformulated (1) (see Appendix H):
L = DKL (pdata(x) ‖ pθ(x)) +DKL (pdata(x)qφ(z|x) ‖ qφ(z)pθ(x|z)) +DKL(qφ(z) ‖ p(z)) . (11)
The first two terms in (11) can eventually become dominant in the VAE training. As a consequence, generation through sampling latent variables from the prior can cause off-distribution samples to be generated. To overcome this prior–posterior mismatch, at least two approaches can be adopted: (i) conduct another posterior estimation after the ordinary VAE training (van den Oord et al., 2017; Razavi et al., 2019b; Dai & Wipf, 2019; Ghosh et al., 2020; Morrow & Chiu, 2020) or (ii) add another regularizing term to the objective function (Makhzani et al., 2015; Tolstikhin et al., 2018; Zhao et al., 2019). The former approach is adopted in our work since it is effective (Ghosh et al., 2020) and is applicable to any VAE variation.
To summarize, the proposed AR-ELBO regularizes the Gaussian VAE with appropriate weighting for the gradient penalty without the need for an extra hyperparameter. Moreover, the remaining mismatch between the prior and posterior is mitigated by an extra pass of posterior estimation. A detailed discussion comparing the proposed objective function with previous works can be found in Appendix I.

heading: 5 EXPERIMENTS

text:
We compare the proposed methods with the following models: VAE, RAE (Ghosh et al., 2020), WAE-MMD (Tolstikhin et al., 2018) and plain autoencoder (AE). The quality of generated images is evaluated using Fréchet Inception Distance (FID) (Heusel et al., 2017) on the MNIST and
CelebA (Liu et al., 2015) datasets with the default train/test split. Regarding the prior–posterior mismatch, three approaches are tested on all the models. The first approach is the conventional case, which samples the latent variables from the prior. The other two approaches are applied after the ordinary training. The second approach forms an aggregated posterior qφ(z) by a second-stage VAE (Dai & Wipf, 2019). The third approach uses a Gaussian mixture model (GMM) with 10 to 100 components (Ghosh et al., 2020) to fit the posterior. The baseline is the standard VAE, and two methods of choosing σ2x are tested: (i) σ 2 x fixed to 1.0 as in general implementations; and (ii) σ 2 x learned with (7) by an optimizer as an usual trainable parameter (Dai & Wipf, 2019). In RAE, its objective function is the sum of the reconstruction error, the regularization of the decoder and the L2 regularization of the latent space. In RAE-GP, gradient penalty is used as the regularization of the decoder. The objective function of RAE-GP is equivalent to (4) except for that the weighting parameters are determined manually. For WAE-MMD, inverse multi-quadratic kernels with seven scales are used as proposed by Tolstikhin et al. (2018). Note that both WAE and RAE have hyperparameters in their objective functions. In contrast, the other methods, as well as the proposed objective functions, include no hyperparameters. However, regarding the major factor that affects the smoothness of the decoder, the proposed methods and VAE control smoothness by regularizing terms in their objective functions, while WAE and AE rely only on the network architecture and generalization techniques.
The latent space dimensions for MNIST and CelebA were set to dz = 16 and 64, respectively, consistent with Ghosh et al. (2020). A common network architecture, which is adopted from Chen et al. (2016) and described in Appendix J, is used for all models. In Table 3, we report the evaluation result of each method as: (i) the MSE of the reconstructed test data and (ii) the FID of generated images. The proposed methods (Diag-I, Iso-D and Diag-D) except for Iso-I do not necessarily achieve the best MSE, since the MSE is no longer the reconstruction loss for them. In the case of sampling z ∈ Z from the prior, a low FID is achieved by WAE due to the relatively strong regularization of the aggregated posterior with MMD. The images generated by the proposed method achieve the best FID score on MNIST, and is competitive on CelebA. It should be noted that the learned σ2x values on MNIST and CelebA from Iso-I are 0.0056 and 0.0050, respectively, which are much smaller than 1.0. Examples of reconstructed and generated images are shown in Appendix K.
As shown in Table 3, different parameterizations of variances can affect the FID greatly. In order to clearly observe the advantage of estimating Σx by MLE rather than estimating it as an usual trainable parameter, we examined the two approaches on all the four parameterizations (Iso-I, IsoD, Diag-I and Diag-D): (i) solve (10b) using MLE as (9) (AR-ELBO); (ii) simply treat Σx as a
trainable parameter (Dai & Wipf, 2019). The comparative result can be obtained from the bottom eight lines of Table 3. It shows that applying AR-ELBO improves FID scores in most of the cases.
Furthermore, in order to examine the feasibility of these learned latent spaces, we also evaluated the FID scores of the images generated by interpolating two latent variables with ratios between [0, 1] with all the models above. The proposed method still shows the best performance in terms of FID on MNIST and is competitive on CelebA. This result suggests the feasibility of proposed method on downstream tasks. The detail and image samples can be found in Appendix L.

heading: 6 CONCLUSION

text: We analyzed the posterior collapse phenomenon on the Gaussian VAE and investigated how strongly the variance parameter impacts the local smoothness of the decoder. The relation between the variance parameter and the local smoothness is examined both theoretically and empirically. We proposed optimization schemes to regulate the local smoothness appropriately, which leads to the prevention of posterior collapse due to oversmoothness. The proposed AR-ELBO implicitly optimizes the variance parameter to avoid over-regularizing of the smoothness. In addition, we proposed several parameterizations (Iso-D, Diag-I, Diag-D) of posterior variances, which are the extensions of the conventional VAE (Iso-I). The corresponding AR-ELBOs for these parameterizations are also derived. Our experiments show that the Gaussian VAE equipped with the proposed objective functions is competitive with other state-of-the-art models in terms of FID for both generated and interpolated images. Moreover, the proposed method remains stable in the most complicated parameterization (Diag-D). In this work, the prior–posterior mismatch was covered by extra posterior estimation methods; however, we would like to seek a thorough solution for this in the future.

heading: A PROOF OF THEOREM 2

text:
Let x be the input sample. We denote its corresponding latent space vector as z and the reconstructed sample as x′. We have the following relation:
I(x; z) ≥ I(x; x′), (12)
which can be proved similarly to the proof of Lemma 5 in Appendix F. On the other hand, I(x; z) can be evaluated by using the definition of the MI as
I(x; z) = DKL(p̃data(x)qφ(z|x) ‖ p̃data(x)qφ(z)) = Ep̃data(x)Eqφ(z|x)[ln qφ(z|x)− ln qφ(z)] = Ep̃data(x)DKL(qφ(z|x) ‖ p(z))−DKL(qφ(z) ‖ p(z)) (13) ≤ Ep̃data(x)DKL(qφ(z|x) ‖ p(z)), (14)
where I(x; z), DKL(qφ(z|x) ‖ p(z)) and DKL(qφ(z) ‖ p(z)) are all non-negative. Inequalities (12) and (14) lead to the proof.
B LINEAR APPROXIMATION OF THE ELBO-BASED OBJECTIVE Jσ2x We start with parameterizing the encoder while following the assumption in (2). Given a sufficiently small perturbation with p( z) = N (z|0,diag(σ2φ(x))), the linear approximation of µθ(·) at µφ(x) can be represented as
µθ(µφ(x) + z) = µθ(µφ(x)) + Jµθ (µφ(x)) z, (15)
where Jµθ (µφ(x)) represents the Jacobian matrix of µθ(z) at z = µφ(x). Substituting (15) into (3) leads to
Eqφ(z|x)[‖x− µθ(z)‖ 2 2] = EN (z|0,diag(σ2φ(x))) [ ‖x− (µθ(µφ(x)) + Jµθ (µφ(x)) z)‖ 2 2 ] = ‖x− µθ(µφ(x))‖22 + EN (z|0,diag(σ2φ(x))) [ >z Jµθ (µφ(x)) >Jµθ (µφ(x)) z ]
+ EN (z|0,diag(σ2φ(x))) [ (x− µθ(µφ(x)))>Jµθ (µφ(x)) z ]︸ ︷︷ ︸ =0 , (16)
where the last term is zero under the assumption that the perturbation is sufficiently small. The expectation in the second right-hand-side term can be evaluated as
EN (z|0,diag(σ2φ(x))) [ >z Jµθ (µφ(x)) >Jµθ (µφ(x)) z ]
= trace ( EN (z|0,diag(σ2φ(x))) [ z > z ] Jµθ (µφ(x)) >Jµθ (µφ(x)) )
= trace ( diag(σ2φ(x))Jµθ (µφ(x)) >Jµθ (µφ(x)) )
= dx∑ i=1 dz∑ j=1 σ2φ,j(x)
( ∂µθ,i(z)
∂zj ∣∣∣∣ z=µφ(z) )2 , (17)
which can be interpreted as the gradient penalty for the decoder weighted by σ2φ(x). By substituting the above result into (3), its linear approximation can be obtained as
J̃σ2x(θ, φ) ≈ 1
2σ2x Ep̃data(x)
[ ‖x− µθ(µφ(x))‖22
+ dx∑ i=1 dz∑ j=1 σ2φ,j(x)
( ∂µθ,i(z)
∂zj ∣∣∣∣ z=µφ(x) )2 + 2σ2x‖µφ(x)‖22 ] . (18)
In the case of the simplified parameterization described in Section 3, the second right-hand-side term in (18) can be further reduced to σ2z‖∇µθ(µφ(x))‖2F . In the simplified case, the perturbation follows a multivariate i.i.d. Gaussian distribution, z ∼ N (z|0, σ2zI). Under this assumption, we have
Ep( z) [ >z Jµθ (µφ(x)) >Jµθ (µφ(x)) z ] = Ep( z) [ >z ( dz∑ i=1 λiui(x)ui(x) > ) z ]
= dz∑ i=1 λiui(x) >Ep( z)[ z > z ]ui(x)
= 2σ2z dz∑ i=1 λi, (19)
where λi is the ith eigenvalue of Jµθ (µφ(x))Jµθ (µφ(x)) >, which is a symmetrical positive definite matrix, and the corresponding eigenvectors are (ui(x))dzi=1. Following the simplified assumption, the second right-hand-side term in (16) now becomes EN (z|0,σ2zI) [ >z Jµθ (µφ(x)) >Jµθ (µφ(x)) z ] .
Combining (19) and the fact that ∑dz i=1 λi = trace(Jµθ (µφ(x))Jµθ (µφ(x)) >) =
‖∇µθ(µφ(x))‖22, we can finally obtain the following linear approximation for the simplified parameterization:
EN (z|0,σ2zI) [ >z Jµθ (µφ(x)) >Jµθ (µφ(x)) z ]
= EN (z|0,σ2zI) [ trace ( Jµθ (µφ(x)) >Jµθ (µφ(x)) )]
= σ2z dx∑ i=1 dz∑ j=1
( ∂µθ,i(z)
∂zj ∣∣∣∣ z=µφ(z) )2 = σ2z ‖∇µθ(µφ(x))‖ 2 F . (20)

heading: C EXPECTED LOCAL SMOOTHNESS OF DECODER

text:
Here, we describe the relation between the expected local smoothness Ep̃data(x)[‖∇µθ(µφ(x))‖2F ] and the expected gap ∆2(s2z). First, consider the relation
∆2(x, z, ′ z) := ‖µθ(µφ(x) + z)− µθ(µφ(x) + ′z)‖22 = Kθ(µφ(x), z, ′z)2‖ z − ′z‖22,
(21) with the perturbation z following the Gaussian distributionN ( z|0, s2zI). Applying the expectation operator to (21) leads to
∆2(x, z, ′ z) = Ep( z, ′z) [ Kθ(µφ(x), z, ′ z) 2‖ z − ′z‖22 ]
(22) ≤ Ep( z, ′z) [ Kθ(µφ(x), z, ′ z) 2 ] Ep( z, ′z) [ ‖ z − ′z‖22 ] (23) =: 2K2θ (µφ(x), s 2 z)dzs 2 z, (24)
where p( z, ′z) := N ( z|0, s2zI)N ( ′z|0, s2zI), z − ′z ∼ N ( z − ′z|0, 2s2zI) and K2θ (µφ(x)) := Ep( z, ′z) [ Kθ(µφ(x), z, ′ z) 2 ] . Note that in (23), we assume that K2θ (µφ(x)) is independent of z and ′z . Consider the case that the variance s 2 z is sufficiently small to approximate µθ(z) linearly around z = µφ(x), which is perturbed with variance s2z . In such a case, Kθ(µφ(x), z, ′ z) is independent of z and ′z , which fits the assumption in (23). Under this local linearity assumption, K2θ (µφ(x)) is bounded as
K2θ (µφ(x), s 2 z) ≤ K2θ , (25)
where Kθ denotes the Lipschitz constant of the decoder.
Following the assumption, K2θ (µφ(x), s 2 z) can be formulated by invoking (15) as
K2θ (µφ(x), s 2 z) =
Ep( z, ′z)[‖µθ(µφ(x) + z)− µθ(µφ(x) + ′ z)‖22]
Ep( z, ′z)[‖ z − ′z‖ 2 2]
= Ep( z, ′z)[( z − ′ z) >Jµθ (µφ(x)) >Jµθ (µφ(x))( z − ′z)] 2dzs2z
= trace
( Ep( z, ′z)[( z − ′ z)( z − ′z)>]Jµθ (µφ(x))>Jµθ (µφ(x)) ) 2dzs2z
= trace
( Jµθ (µφ(x)) >Jµθ (µφ(x)) )
dz . (26)
Applying the expectation operator to (26) leads to K2θ (s 2 z) := Ep̃data [ K2θ (µφ(x), s 2 z) ]
= Ep̃data(x)
[ trace ( Jµθ (µφ(x)) >Jµθ (µφ(x)) )]
dz
= Ep̃data
[ ‖∇µθ(µφ(x))‖2F ] dz . (27)
Finally, combining (24) and (27) yields the following connection between the expected gap and the expected local smoothness:
∆2(s2z) = 2Ep̃data [ ‖∇µθ(µφ(x))‖2F ] s2z. (28)

heading: D EXPERIMENTAL DETAILS FOR SECTION 3.2

text:
D.1 EXPERIMENTAL SETUP
In the experiment, the model is trained with the Adam optimizer with a learning rate of 10−3. The dimension of the latent space is set to 8. We run 200 epochs with a minibatch size of 64 for all σ2x. We use the following DNN architectures for the encoder and decoder, respectively:
x ∈ R28×28 → Conv64 → ReLU size of (64, 14, 14) → Conv128 → ReLU→ Reshape(128, 7, 7) size of (128, 7, 7) → Flatten→ FC1024 → ReLU → FC16,
z ∈ R16 → FC1024 → ReLU → FC128×7×7 → ReLU size of (128, 7, 7) → ConvT64 → ReLU size of (64, 14, 14) → ConvT1 → Sigmoid size of (1, 28, 28).
Here, FCk, Convk, ConvTk and ReLU denote the fully connected layer mapping to Rk, the convolutional layer mapping to k channels, the transpose convolutional layer mapping to k channels and the rectified linear units (ReLU), respectively. The 3-tuple (channels,height,width) in the right column represents the output shape of each layer. In all the Convk and ConvTk layers, 4× 4 convolutional filters are used with a common stride of (2, 2).
Regarding the evaluation of criteria, MSE and KL are evaluated on the training set because the aim of the experiment is to validate the relation between σ2z and the smoothness of the decoder. The upper bound of the MI is obtained by calculating
−Ep̃data(x)Eqφ,σ2z (z|x) [ lnEp̃data(x′) exp ( − 1 2σ2z ‖z− µφ(x′)‖22 )] − dx 2 (29)
for each minibatch and then taking their mean, where the batch size is 10, 000 for all the evaluations.
D.2 SAMPLES OF GENERATED IMAGES AND T-SNE VISUALIZATION OF LATENT SPACES
Figure 1 shows several images decoded from µφ(x) + z with z ∼ N ( z|0, s2zI) for the cases with σ2x = 1.0 and 0.1. Posterior collapse can be observed from these blurry images decoded from the stochastic encoding case with σ2x = 1.0. This is due to the removal of batch normalization, which makes σ2x = 1.0 become an inappropriate choice. However, if σ 2 x is determined or adapted appropriately such as by using the proposed method, posterior collapse will not happen. In the other settings, the tendency of how the image changes with the perturbation is similar, as shown in Table 1.
The latent spaces are also visualized via t-SNE (Maaten & Hinton, 2008) in Figure 2. The dots with different colors represent the latent vectors encoded from images of different labels (numbers), and the pink dots are the sampling points generated from the prior p(z). As mentioned earlier, to observe the effect of σ2x clearly, we remove batch normalization, which usually helps prevent posterior collapse to a certain extent. As a result, the latent space with σ2x = 1.0 completely collapses and qφ(z) approaches p(z) as shown in Figures 1 and 2(a). In this case, both KL collapse and posterior collapse occur.

heading: E FIXING THE POSTERIOR VARIANCE OF LATENT SPACE

text:
From the previous sections, we know that σ2x affects the smoothness via σ 2 x. However, it would be interesting to see what will happen if σ2z is fixed while σ 2 x is optimized. In this experiment, the variance parameter σ2z is fixed while σ 2 x is optimized with the AR-ELBO (9) under the parameterization in Section 3. The other settings remain the same as those in Section 3.2. We evaluate the numerical results for different σ2z with the criteria listed in Section 3.2. According to Table 5, the tendencies of the expected gap and ELS show that a large σ2z makes the decoder smoother, which is consistent with the discussion in Section 3.1. However, the tendency of the KL divergence is different from that in Section 3.2. Although a larger σ2z consistently leads to a smaller MI, and eventually the MI collapses to zero; the KL divergence still remains far from zero, which means that posterior collapse
can happen without KL collapse. This phenomenon can be visually confirmed by observing the t-SNE plot in Figure 3. The cause of this phenomenon can be roughly reasoned from the linear approximated ELBO (4), in which σ2z directly affects the gradient penalty and causes oversmoothness. It should be pointed out that the strength of L2 regularization in (4) is gradually decreased with decreasing σ2x; therefore, it does not dominate the whole objective function. As a result, the mean of the approximated posterior qφ(z) is far from the mean of the prior p(z) (which is 0), and therefore DKL(qφ(z) ‖ p(z)) in (13) does not diminish to zero.

heading: F PROOF OF THEOREM 3

text:
According to Theorem 4 in Dai & Wipf (2019), we know that lim σ2x→0 Epdata(x)Eqφ,σ2z (z|x) [ ‖x− µθ(z)‖22 ] = 0, (30)
which also leads to σ̂2x → 0 (σ2x → 0). Here, σ̂2x is estimated through MLE and is given by
σ̂2x = 1
dx Ep̃data(x)Eqφ,σ2z (z|x)
[ ‖x− µθ(z)‖22 ] . (31)
To prove Theorem 3, we need the following auxiliary theorem: Theorem 4. In the training stage of VAE, we have σ2z → 0 (σ̂2x → 0).
First, we state the two lemmas with proofs. Lemma 5. In a VAE, I(x,x′) ≤ I(z, ze) always holds, where ze is the encoded latent variable ze = µφ(x) with x ∼ pdata(x).
Proof. The data processing flow of the VAE is x → ze → z → x′; ze = µφ(x), z = ze + z , and x′ = µθ(z), where z ∼ N ( z|0, σ2zI). The MI I(x; z,x′) can be represented as
I(x; z,x′) = I(x; x′) + I(x; z|x′) (32) = I(x; z) + I(x; x′|z). (33)
Since x and x′ are conditionally independent on the given z, it follows that I(x; x′|z) = 0. From the non-negativity of MI, we have I(x; z) ≥ I(x; x′). Repeating the same procedure for I(x; ze, z) leads to the proof.
Lemma 6. The MI between x and x′ diverges to positive infinity as σ2x → 0, where x′ is obtained from x ∼ pdata(x) as x′ = µθ(µφ(x) + z).
Proof. A lower bound of I(x; x′) is
I(x; x′) = DKL(pdata(x)pθ,φ(x′|x) ‖ pdata(x)pθ,φ(x′)) = Epdata(x)Epθ,φ(x′|x) [ln pθ,φ(x
′|x)− ln pθ,φ(x′)] = H [pθ,φ(x′)]− Epdata(x)H [pθ,φ(x
′|x)] ≥ H [pθ,φ(x′)]− Epdata(x)H(σ̂ 2 xI), (34)
where pθ,φ(x′|x) := Eqφ(z|x)[pθ(x′|z)] and pθ,φ(x′) := Epdata(x)Eqφ(z|x)[pθ(x′|z)]. Here, we denote the differential entropy of the Gaussian with variance σ̂2xI as
H(σ̂2xI) := 1
2 ln(2πeσ̂2dxx ). (35)
Since σ̂2x → 0 as σ2x → 0, H[pθ,φ(x′)] → H[pdata(x)] and H(σ̂2x) → −∞ in the inequality of (34). Therefore, I(x; x′)→∞ as σ2x → 0.
Now we prove Theorem 3. The MI I(z; ze) satisfies I(z; ze) = DKL ( qφ(ze)qσ2z (z|ze) ‖ qφ(ze)qφ,σ2z (z) ) = Eqφ(ze)Eqσ2z (z|ze) [ ln qσ2z (z|ze)− ln qφ,σ2z (z)
] = H [ qφ,σ2z (z) ] − Eqφ(ze)H [ qσ2z (z|ze)
] ≤ H(Σφ,σ2z )−H(σ 2 zI)
= dz 2 ln
( det(Σφ,σ2z )
σ2z
) , (36)
where Σφ,σ2z denotes the variance of qφ,σ2z (z). Invoking Lemma 5 and (36) leads to
I(x; x′) ≤ dz 2 ln
( det(Σφ,σ2z )
σ2z
) . (37)
Now, consider σ2z 6→ 0 as L → 0. According to Lemma 6, it follows that det(Σφ,σ2z ) → +∞, which contradicts the fact that qφ,σ2z (z) → p(z). Thus, we must have σ 2 z → 0 as σ2x converges to zero.

heading: G DERIVATION OF PROPOSED OBJECTIVES

text:
Here, we derive the objectives listed in Table 2. Consider an arbitrary Σx without any condition. The MLE of Σx, Σ̂x, can be obtained by
Σ̂x = Ep̃data(x)Eqφ(z|x) [ (x− µθ(z))(x− µθ(z))> ] . (38)
From the partial derivative of J̃rec(θ, φ,Σx) w.r.t. Σx, we have
∂J̃rec(θ, φ,Σx) ∂Σx = 1 2
( Ep̃dataEqφ(z|x) [ (x− µθ(z))(x− µθ(z))> ] + Σ−1x ) . (39)
The MLE of Σ̂x and the objectives for the different parameterizations are described in the following.
G.1 ISO-I
For Iso-I, the MLE of σ2x can be given as
σ̂2x = 1
dx Ep̃data(x)Eqφ(z|x)
[ ‖x− µθ(z)‖22 ] . (40)
Substituting (40) into (7) leads to J̃AR(θ, φ, σ̂2x) = Ep̃data(x) [ 1
2σ̂2x Eqφ(z|x)[‖x− µθ(z)‖ 2 2] +DKL(qφ(z|x) ‖ p(z)) ] + dx 2 ln σ̂2x
= dx 2 + Ep̃data(x)DKL(qφ(z|x) ‖ p(z))
+ dx 2
lnEp̃data(x)Eqφ(z|x) [ ‖x− µθ(z)‖22 ] − dx
2 ln dx, (41)
where the first and fourth terms are constants and thus omitted in (9).
G.2 ISO-D
First, substitute Σx = σ2x(z)I into (10b): J̃rec(θ, φ,Σx) = Ep̃data(x)Eqφ(z|x) [ 1
2σ2x(z) ‖x− µθ(z)‖22 + dx 2 lnσ2x(z)
] . (42)
Also, we know that the MLE of σ2x(z) is
σ̂2x(z) = 1
dx ‖x− µθ(z)‖22 . (43)
Substituting (43) into (42) leads to the reconstruction objective of Iso-D: J̃rec(θ, φ, Σ̂x) = Ep̃data(x)Eqφ(z|x) [ 1
2σ̂2x(z) ‖x− µθ(z)‖22 + dx 2 ln σ̂2x(z)
] (44)
= dx 2 + dx 2 Ep̃data(x)Eqφ(z|x)
[ ln ‖x− µθ(z)‖22 ] − dx
2 ln dx. (45)
G.3 DIAG-I
First, substitute Σx = diag(σ2x) into (10b):
J̃rec(θ, φ,Σx) = Ep̃data(x) [ dx∑ i=1 1 2σ2x,i Eqφ(z|x) [ (xi − µθ,i(z))2 ]] + dx∑ i=1 1 2 lnσ2x,i. (46)
Also, we know that the MLE of σ2x,i is σ̂2x,i = Ep̃data(x)Eqφ(z|x) [ (xi − µθ,i(z))2 ] . (47)
Substituting (47) into (46) leads to the reconstruction objective for Diag-I:
J̃rec(θ, φ, Σ̂x) = Ep̃data(x) [ dx∑ i=1 1 2σ̂2x,i Eqφ(z|x) [ (xi − µθ,i(z))2 ]] + dx∑ i=1 1 2 ln σ̂2x,i (48)
= dx 2 + 1 2 dx∑ i=1 lnEp̃data(x)Eqφ(z|x) [ (xi − µθ,i(z))2 ] . (49)
G.4 DIAG-D
First, substitute Σx = diag(σ2x(z)) into (10b):
J̃rec(θ, φ,Σx) = Ep̃data(x)Eqφ(z|x) [ dx∑ i=1 ( 1 2σ2x,i(z) (xi − µθ,i(z))2 + 1 2 lnσ2x,i(z) )] . (50)
Also, we know that the MLE of σ2x,i(z) is
σ̂2x,i(z) = (xi − µθ,i(z)) 2 . (51)
Substituting (51) into (50) leads to the reconstruction objective for Diag-D:
J̃rec(θ, φ, Σ̂x) = Ep̃data(x)Eqφ(z|x) [ dx∑ i=1 ( 1 2σ̂2x,i(z) (xi − µθ,i(z))2 + 1 2 ln σ̂2x,i(z) )] (52)
= dx 2 + 1 2 dx∑ i=1 Ep̃data(x)Eqφ(z|x) [ ln (xi − µθ,i(z))2 ] . (53)

heading: H DERIVATION OF (11)

text:
The KL divergence terms of (1) can be represented as
Epdata(x)DKL(qφ(z|x) ‖ pθ(z|x)) = Epdata(x)qφ(z|x)[ln qφ(z|x)− ln pθ(z|x)] (54) = Epdata(x)qφ(z|x) [ ln pdata(x)qφ(z|x) p(z)pθ(x|z) ] and
Epdata(x)qφ(z|x) [ln qφ(z)− ln p(z)] = DKL(qφ(z) ‖ p(z)). (55) Substituting the two equations above into (1), then L can be reformulated into (11).

heading: I RELATED WORKS

text:
To the best of our knowledge, Lucas et al. (2019) were among the first to suggest that posterior collapse may be caused by a sub-optimal σ2x. In the past, one of the common approaches for dealing with posterior collapse was to anneal the weight of the KL term in the ELBO. The first such attempt was KL annealing (Bowman et al., 2015). Bowman et al. (2015) introduced a weighting coefficient on the KL term in the cost function during training. The weighting scheduling is determined in advance, e.g., the weight increases monotonically (Bowman et al., 2015; Sønderby et al., 2016) or changes cyclically (Fu et al., 2019) as the training progresses.
The weighting coefficient also appears in Higgins et al. (2017), and is interpreted as a hyperparameter that controls the information capacity of the latent space. The suggested value for such hyperparameter is larger than 1. This is analogous to setting σ2x larger than the MLE value σ̂ 2 x for (7) and (8a), which enforces a stronger smoothness in exchange of better latent space disentanglement. The differences between Higgins et al. (2017) and the proposed method are: (i) σ2x changes between every minibatch; and (ii) the estimation of σ2x is aimed to prevent the oversmoothness.
Shao et al. (2020) proposed ControlVAE, which combined control theory with the VAE, and applied PI/PID control to determine the weight on the KL term. Although applying control theory to the weighting of the KL term makes it possible to reflect the status of the optimization, ControlVAE needs extra hyperparameters to be tuned in advance. On the other hand, our method can be interpreted as automatic KL annealing that estimates σ2x through MLE without the need of tuning an extra hyperparameter.
Ghosh et al. (2020) interpreted the stochastic autoencoder with the reparameterization trick as noise injection process and proposed replacing such a mechanism with an explicit regularized autoencoder (RAE). RAE regularizes its decoder in several ways: L2 regularization, a gradient penalty (Gulrajani et al., 2017) and spectral normalization (Miyato et al., 2018). As discussed in Section 3.1, if σ2z is sufficiently small, the ELBO can also be approximately represented as a sum of three losses (4), which correspond to the terms included in the basic RAE objective function. The approximated objective function (4) can be obtained when RAE with a gradient penalty (RAE-GP) is used and tuned appropriately.
Dai & Wipf (2019) optimized σ2x using an optimizer, which is the most similar approach to our Iso-I model in that (7) is used as an objective function. Our proposed method provides simplified objective functions that enable the variance parameter to be optimized automatically and guarantee that σ2x decreases as the reconstruction loss decreases, enabling the gradient penalty to be gradually weakened.
Comparing our work with these works showed that the proposed AR-ELBO and its variations are also capable of regularizing the Gaussian VAE via weighting of the gradient penalty. In the standard parameterization such as (2), the second term in (4) is a weighted gradient penalty, which makes it possible to regularize each dimension of the latent space differently according to the property of the input data. In addition, the combination of the proposed objective functions together with standard Gaussian VAE allows implicit gradient regularization on the decoder with lower computational cost than that of explicitly adding the gradient penalty to the objective function, such as in RAE.

heading: J DETAILS OF EXPERIMENTAL SETUP IN SECTION 5

text:
In this experiment, the Adam optimizer (Kingma & Ba, 2015) is used and the maximum number of epochs is set to 100 for MNIST and 70 for CelebA. The learning rates are 0.001 for MNIST and 0.0002 for CelebA. A minibatch size of 64 is used. All the FID2 values are evaluated with 10, 000 generated samples.
For the posterior estimation by the second-stage VAE, we adopt the same networks for the encoder and decoder as those in Dai & Wipf (2019). For GMM fitting, we use the same settings as those in Ghosh et al. (2020). Experimental details including the network architectures for each dataset are described in the following.
J.1 MNIST
We construct the encoder and decoder for the MNIST dataset using the architecture in Chen et al. (2016). The encoder is constructed as
x ∈ R28×28 → Conv64 → ReLU size of (64, 14, 14) → Conv128 → ReLU→ Reshape(128, 7, 7) size of (128, 7, 7) → Flatten→ FC1024 → BatchNorm→ ReLU → FC16×2.
The decoder is constructed as
z ∈ R16 → FC1024 → BatchNorm→ ReLU → FC128×7×7 → BatchNorm→ ReLU size of (128, 7, 7) → ConvT64 → BatchNorm→ ReLU size of (64, 14, 14) → ConvT1 → Sigmoid size of (1, 28, 28).
In all the Convk layers and all the ConvTk layers except for the last, 5×5 convolutional filters with stride (2, 2) are used. The difference between this architecture and those used in Appendix D.1 is whether batch normalization is applied or not. Although in the original work of Chen et al. (2016), the discriminator used leaky ReLU (lReLU), we adopt ReLU for the encoder part, which improves the performance for all the models evenly.
J.2 CELEBA
The CelebA images are preprocessed with center cropping of 140 × 140, then resized to 64 × 64 as described in Tolstikhin et al. (2018) and Ghosh et al. (2020). It should be noted that the size of cropping differs among the previous works, and it markedly affects the FID score. We choose the above cropping size as is the largest among the related works and seems to be the most difficult case for image generation. Moreover, Tolstikhin et al. (2018) and Ghosh et al. (2020) also used this cropping size. Similarly to in the previous section, the encoder and decoder are constructed on the basis of the discriminator and generator for CelebA used in Chen et al. (2016). The encoder is
2We used the PyTorch version of the FID implementation from https://github.com/mseitzer/ pytorch-fid for all the models. However, the result may slightly differ from that obtained with the TensorFlow implementation https://github.com/bioinf-jku/TTUR.
constructed as x ∈ R64×64 → Conv128 → ReLU size of (128, 32, 32)
→ Conv256 → BatchNorm→ ReLU size of (256, 16, 16) → Conv512 → BatchNorm→ ReLU size of (512, 8, 8) → Conv1024 → BatchNorm→ ReLU size of (1024, 4, 4) → Flatten→ FC64×2.
The decoder is constructed as z ∈ R64 → FC8×8×1024
→ ConvT512 → ReLU size of (512, 16, 16) → ConvT256 → BatchNorm→ ReLU size of (256, 32, 32) → ConvT128 → BatchNorm→ ReLU size of (128, 64, 64) → ConvT3 size of (3, 64, 64).
In all the Convk layers and all the ConvTk layers except for the last, 5×5 convolutional filters with stride (2, 2) are used. We use ReLU instead of leaky ReLU due to the performance consideration described in the previous subsection. To fit the size of the input images in our experiment, one extra convolutional layer is added for the encoder and the channel size is twice as large as that in Chen et al. (2016),

heading: K EXAMPLES OF RECONSTRUCTED AND GENERATED IMAGES IN SECTION 5

text:
We show examples of reconstructed images and images generated by sampling the learned approximated posterior from the proposed method and other works in Figures 4 and 5.
L INTERPOLATION OF LATENT VARIABLES
This section aims to investigate the feasibility of the learned latent space of the methods mentioned in Section 5. If high quality images can be generated by interpolating the latent variables in a latent
space, the corresponding latent space is more likely to be feasible for other downstream tasks. Therefore, in this section, we evaluate the FID scores of images generated by latent variable interpolation for various models mentioned in section 5.
First, we choose 10,000 random pairs of images from both MNIST and CelebA datasets. The interpolation is done by applying spherical interpolation (Ghosh et al., 2020) in latent spaces and then generate the interpolated images with the decoders. In the end, we evaluate the FID of these interpolated images.
Furthermore, the experiment has been proceeded with two different setups of mixing ratios: (i) a fixed ratio of 0.5, i.e., the mid-point of two latent variables; and (ii) a uniformly distributed random ratio between [0, 1] for each image pair. The result is shown in Table 4, where the proposed method achieved the best score on MNIST and is competitive on CelebA. This suggests that a generative model with proper smoothness achieved via the proposed method is also feasible for other applications such as interpolation and possibly applicable for other semantic controls.

title: AR-ELBO: PREVENTING POSTERIOR COLLAPSE INDUCED BY OVERSMOOTHING IN GAUSSIAN VAE


summaries: This paper proposed a deep generative model based on the non i.i.d. VAE framework in an unsupervised version. The model which combines a mixture prior in the local latent space with global latent space has three advantages: First, the latent space can capture interpretable features. Second, the model performs domain alignment. Third, the model can discriminate among their global posterior representations. Although this paper has mild improvement on the basic VAE structure, the model displays a good interpretability power, and the setup of the latent variables are illustrated reasonably in the paper.

summaries: This article introduces a VAE-based method for separating local variation factors from global variation factors in the data in an unsupervised manner. It achieves so by designing a graphical model with a mix of example-local and batch-shared variables, and training it using the ELBO. The article provide an detailed experimental analysis on MNIST and CelebA, and experimental evidence that all parts of the model (notably the discrete d variable) are relevant.

summaries: This paper presents a novel deep generative model based on noni.i.d. variational autoencoders that captures global dependencies among observations in a fully unsupervised fashion. The proposed model combines a mixture model in the local or data-dependent space and a global Gaussian latent variable, which captures interpretable disentangled representations with no user-defined regularization in the evidence lower bound. The proposed model is being evaluated in two tasks: (1) disentanglement, and (2) domain alignment.

abstractText: We present a novel deep generative model based on non i.i.d. variational autoencoders that captures global dependencies among observations in a fully unsupervised fashion. In contrast to the recent semi-supervised alternatives for global modeling in deep generative models, our approach combines a mixture model in the local or data-dependent space and a global Gaussian latent variable, which lead us to obtain three particular insights. First, the induced latent global space captures interpretable disentangled representations with no user-defined regularization in the evidence lower bound (as in beta-VAE and its generalizations). Second, we show that the model performs domain alignment to find correlations and interpolate between different databases. Finally, we study the ability of the global space to discriminate between groups of observations with non-trivial underlying structures, such as face images with shared attributes or defined sequences of digits images.

authors:
Mathieu Aubry
Daniel Maturana
Alexei A Efros
Bryan C Russell
Josef Sivic

title: Seeing 3d chairs: exemplar part-based 2d-3d alignment using a large dataset of cad models

venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,

authors:
Solon Barocas
Moritz Hardt
Arvind Narayanan

title: Fairness in machine learning

venue: NIPS Tutorial,

authors:
Diane Bouchacourt
Ryota Tomioka
Sebastian Nowozin

title: Multi-level variational autoencoder: Learning disentangled representations from grouped observations

venue: In Thirty-Second AAAI Conference on Artificial Intelligence,

title: Products and convolutions of gaussian probability density functions

venue: Tina-Vision Memo,

authors:
Junyoung Chung
Kyle Kastner
Laurent Dinh
Kratarth Goel
Aaron C Courville
Yoshua Bengio

title: A recurrent latent variable model for sequential data

venue: In Advances in neural information processing systems,

authors:
Nat Dilokthanakul
Pedro AM Mediano
Marta Garnelo
Matthew CH Lee
Hugh Salimbeni
Kai Arulkumaran
Murray Shanahan

title: Deep unsupervised clustering with gaussian mixture variational autoencoders

venue: arXiv preprint arXiv:1611.02648,

authors:
Sanja Fidler
Sven Dickinson
Raquel Urtasun

title: 3d object detection and viewpoint estimation with a deformable 3d cuboid model

authors:
Prashnna Gyawali
Zhiyuan Li
Cameron Knight
Sandesh Ghimire
B Milan Horacek
John Sapp
Linwei Wang

title: Improving disentangled representation learning with the beta bernoulli process

venue: IEEE International Conference on Data Mining (ICDM),

authors:
Christina Heinze-Deml
Nicolai Meinshausen

title: Conditional variance penalties and domain shift robustness

venue: arXiv preprint arXiv:1710.11469,

authors:
Irina Higgins
Loic Matthey
Arka Pal
Christopher Burgess
Xavier Glorot
Matthew Botvinick
Shakir Mohamed
Alexander Lerchner

title: β-VAE: Learning basic visual concepts with a constrained variational framework

title: Group-based learning of disentangled representations with generalizability for novel contents

venue: In IJCAI,

authors:
Matthew J Johnson
David K Duvenaud
Alex Wiltschko
Ryan P Adams
Sandeep R Datta

title: Composing graphical models with neural networks for structured representations and fast inference

authors:
Diederik P Kingma
Max Welling

title: Auto-encoding variational bayes

venue: arXiv preprint arXiv:1312.6114,

authors:
Iryna Korshunova
Jonas Degrave
Ferenc Huszár
Yarin Gal
Arthur Gretton
Joni Dambre

title: Bruno: A deep recurrent model for exchangeable data

venue: In Advances in Neural Information Processing Systems,

authors:
Jonathan Krause
Michael Stark
Jia Deng
Li Fei-Fei

title: 3d object representations for fine-grained categorization

venue: In 4th International IEEE Workshop on 3D Representation and Recognition (3dRR-13), Sydney, Australia,

authors:
Emile Mathieu
Tom Rainforth
N Siddharth
Yee Whye Teh

title: Disentangling disentanglement in variational autoencoders

venue: In International Conference on Machine Learning,

authors:
Alfredo Nazabal
Pablo M Olmos
Zoubin Ghahramani
Isabel Valera

title: Handling incomplete heterogeneous data using vaes

venue: Pattern Recognition,

authors:
Pascal Paysan
Reinhard Knothe
Brian Amberg
Sami Romdhani
Thomas Vetter

title: A 3d face model for pose and illumination invariant face recognition

venue: In 2009 Sixth IEEE International Conference on Advanced Video and Signal Based Surveillance,

authors:
Rajesh Ranganath
Dustin Tran
David Blei

title: Hierarchical variational models

authors:
Da Tang
Dawen Liang
Tony Jebara
Nicholas Ruozzi

title: Correlated variational auto-encoders

venue: arXiv preprint arXiv:1905.05335,

authors:
Jakub Tomczak
Max Welling

title: Vae with a vampprior

venue: In International Conference on Artificial Intelligence and Statistics,

authors:
Matthew J Vowels
Necati Cihan Camgoz
Richard Bowden

title: Nestedvae: Isolating common factors via weak supervision

venue: In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,

authors:
Kai Xu
Akash Srivastava
Charles Sutton

title: Variational russian roulette for deep bayesian nonparametrics

authors:
Aron Yu
Kristen Grauman

title: Fine-grained visual comparisons with local learning

venue: In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,

heading: 1 INTRODUCTION

text:
Since its first proposal by Kingma & Welling (2013), Variational Autoencoders (VAEs) have evolved into a vast amount of variants. To name some representative examples, we can include VAEs with latent mixture models priors (Dilokthanakul et al. (2016)), adapted to model time-series (Chung et al. (2015)), trained via deep hierarchical variational families (Ranganath et al. (2016), Tomczak & Welling (2018)), or that naturally handle heterogeneous data types and missing data (Nazabal et al. (2020)).
The large majority of VAE-like models are designed over the assumption that data is i.i.d., which remains a valid strategy for simplifying the learning and inference processes in generative models with latent variables. A different modelling approach may drop the i.i.d. assumption with the goal of capturing a higher level of dependence between samples. Inferring such kind of higher level dependencies can directly improve current approaches to find interpretable disentangled generative models (Bouchacourt et al. (2018)), to perform domain alignment (Heinze-Deml & Meinshausen (2017)) or to ensure fairness and unbiased data (Barocas et al. (2017)).
The main contribution of this paper is to show that a deep probabilistic VAE non i.i.d. model with both local and global latent variable can capture meaningful and interpretable correlation among data points in a completely unsupervised fashion. Namely, weak supervision to group the data samples is not required. In the following we refer to our model as Unsupervised Global VAE (UG-VAE). We combine a clustering inducing mixture model prior in the local space, that helps to separate the fundamental data features that an i.i.d. VAE would separate, with a global latent variable that modulates the properties of such latent clusters depending on the observed samples, capturing fundamental and interpretable data features. We demonstrate such a result using both CelebA, MNIST and the 3D FACES dataset in Paysan et al. (2009). Furthermore, we show that the global latent space can explain common features in samples coming from two different databases without requiring any domain label for each sample, establishing a probabilistic unsupervised framework for domain alignment. Up to our knowledge, UG-VAE is the first VAE model in the literature that performs unsupervised domain alignment using global latent variables.
Finally, we demonstrate that, even when the model parameters have been trained using an unsupervised approach, the global latent space in UG-VAE can discriminate groups of samples with non-trivial structures, separating groups of people with black and blond hair in CelebA or series of
numbers in MNIST. In other words, if weak supervision is applied at test time, the posterior distribution of the global latent variable provides with an informative representation of the user defined groups of correlated data.

heading: 2 RELATED WORK

text:
Non i.i.d. deep generative models are getting recent attention but the literature is still scarse. First we find VAE models that implement non-parametric priors: in Gyawali et al. (2019) the authors make use of a global latent variable that induces a non-parametric Beta process prior, and more efficient variational mechanism for this kind of IBP prior are introduced in Xu et al. (2019). Second, both Tang et al. (2019) and Korshunova et al. (2018) proposed non i.i.d. exchangable models by including correlation information between datapoints via an undirected graph. Finally, some other works rely on simpler generative models (compared to these previous approaches), including global variables with fixed-complexity priors, typically a multi-variate Gaussian distribution, that aim at modelling the correlation between user-specified groups of correlated samples (e.g. images of the same class in MNIST, or faces of the same person). In Bouchacourt et al. (2018) or Hosoya (2019), authors apply weak supervision by grouping image samples by identity, and include in the probabilistic model a global latent variable for each of these groups, along with a local latent variable that models the distribution for each individual sample. Below we specify the two most relevant lines of research, in relation to our work.
VAEs with mixture priors. Several previous works have demonstrated that incorporating a mixture in the latent space leads to learn significantly better models. In Johnson et al. (2016) authors introduce a latent GMM prior with nonlinear observations, where the means are learned and remain invariant with the data. The GMVAE proposal by Dilokthanakul et al. (2016) aims at incorporating unsupervised clustering in deep generative models for increasing interpretability. In the VAMP VAE model Tomczak & Welling (2018), the authors define the prior as a mixture with components given by approximated variational posteriors, that are conditioned on learnable pseudo-inputs. This approach leads to an improved performance, avoiding typical local optima difficulties that might be related to irrelevant latent dimensions.
Semi-supervised deep models for grouped data. In contrast to the i.i.d. vanilla VAE model in Figure 1 (a), and its augmented version for unsupervised clustering, GMVAE, in Figure 1 (b), the graphical model of the Multi-Level Variational Autoencoder (ML-VAE) in Bouchacourt et al. (2018) is shown in Figure 1 (c), where G denotes the number of groups. ML-VAE includes a local Gaussian variable Si that encodes style-related information for each sample, and global Gaussian variable CG to model shared in a group of samples. For instance, they feed their algorithm with batches of face images from the same person, modeling content shared within the group that characterize a person. This approach leads to learning a disentangled representations at the group and observations level, in a content-style fashion. Nevertheless, the groups are user-specified, hence resulting in a semisupervised modelling approach. In Vowels et al. (2020) authors use weak supervision for pairing samples. They implement two outer VAEs with shared weights for the reconstruction, and a Nested
VAE that reconstructs latent representation off one to another, modelling correlations across pairs of samples. The graphical model for Nested VAE is depicted in Figure 1 (d).

heading: 3 UNSUPERVISED GLOBAL VAE

text: We present UG-VAE, a deep generative VAE framework for modeling non-i.i.d. data with global dependencies. It generalizes the ML-VAE graphical model in Figure 1 (c) to i) remove the group supervision, ii) include a clustering-inducing prior in the local space, and iii) propose a more structured variational family.

heading: 3.1 GENERATIVE MODEL

text:
Figure 2 represents the generative graphical model of UG-VAE. A global variable β ∈ Rg induces shared features to generate a group of B samples X = {x1, ..., xB} ⊆ RD, and G is the number of groups we jointly use to amortize the learning of the model parameters. During amortized variational training, groups are simply random data mini-batches from the training dataset, being G the number of data mini-batches. We could certainly take B = N (the training set size) and hence G = 1, but this leads to less interpretable global latent space (too much data to correlate with a single global random variable), and a slow training process.
Conditioned to β, data samples are independent and distributed according to a Gaussian mixture local (one per data) latent variable Z = {z1, ..., zB} ⊆ Rd, and d = {d1, ..., dB} ⊆ {1, ...,K} are independent discrete categorical variables with uniform prior distributions. This prior, along with the conditional distribution p(zi|di,β), defines a Gaussian mixture latent space, which helps to infer similarities between samples from different batches (by assigning them to the same cluster), and thus, di plays a similar role than the semi-supervision included in Bouchacourt et al. (2018) by grouping. Our experimental results demonstrate that this level of structure in the local space is crucial to acquire interpretable information at the global space, and specially, if we fix di for all the samples within a batch, that the global variable β is able to tune different generative factors for each cluster.
The joint distribution for a single group is therefore defined by:
pθ(X,Z,d,β) = p(X|Z,β) p(Z|d,β) p(d) p(β) (1)
where the likelihood term of each sample is a Gaussian distribution, whose parameters are obtained from a concatenation of zi and β as input of a decoder network:
p(X|Z,β) = B∏ i=1 p(xi|zi,β) = B∏ i=1 N (µθx([zi,β]),Σθx([zi,β])) (2)
In contrast with Johnson et al. (2016), where the parameters of the clusters are learned but shared by all the observations, in UG-VAE, the parameters of each component are obtained with networks fed
with β. Thus, the prior of each local latent continuous variable is defined by a mixture of Gaussians, where di defines the component and β is the input of a NN that outputs its parameters:
p(Z|d,β) = B∏ i=1 p(zi|di,β) = B∏ i=1 N ( µ (di) θz (β),Σ (di) θz (β) ) , (3)
hence we trained as many NNs as discrete categories. This local space encodes samples in representative clusters to model local factors of variation. The prior of the discrete latent variable is defined as uniform:
p(d) = B∏ i=1 Cat(π) πk = 1/K (4)
and the prior over the continuous latent variable β follows an isotropic Gaussian, p(β) = N (0, I).

heading: 3.2 INFERENCE MODEL

text:
The graphical model of the proposed variational family is shown in Figure 2(b):
qφ(Z,d,β|X) = q(Z|X) q(d|Z)q(β|X,Z) (5)
where we employ an encoder network that maps the input data into the local latent posterior distribution, which is defined as a Gaussian:
q(Z|X) = B∏ i=1 q(zi|xi) = B∏ i=1 N (µφz (xi),Σφz (xi)) (6)
Given the posterior distribution of z, the categorical posterior distribution of di is parametrized by a NN that takes zi as input
q(d|Z) = B∏ i=1 q(di|zi) = B∏ i=1 Cat(πφd(zi)) (7)
The approximate posterior distribution of the global variable β is computed as a product of local contributions per datapoint. This strategy, as demonstrated by Bouchacourt et al. (2018), outperforms other approaches like, for example, a mixture of local contributions, as it allows to accumulate group evidence. For each sample, a NN encodes xi and the Categorical parameters πφd(zi) in a local Gaussian:
q(β|X,Z) = N (µβ ,Σβ) = B∏ i=1 N ( µφβ ([xi, πφd(zi)]),Σφβ ([xi, πφd(zi)]) ) (8)
If we denote by µi and Σi the parameters obtained by networks µφβ and Σφβ , respectively, the parameters of the global Gaussian distribution are given, following Bromiley (2003), by:
Λβ = Σ −1 β = B∑ i=1 Λi
µβ = (Λβ) −1 B∑ i=1 Λiµi
(9)
where Λβ = Σ−1β is defined as the precision matrix, which we model as a diagonal matrix.

heading: 3.3 EVIDENCE LOWER BOUND

text:
Overall, the evidence lower bound reads as follows:
L(θ, φ; X,Z,d,β) = Eq(β) [Li(θ, φ; xi, zi,d,β)]− Eq(d) [DKL (q(β|X,Z)‖p(β))] (10)
The resulting ELBO is an expansion of the ELBO for a standard GMVAE with a new regularizer for the global variable. As the reader may appreciate, the ELBO for UG-VAE does not include extra
hyperparameters to enforce disentanglement, like other previous works as β-VAE, and thus, no extra validation is needed apart from the parameters of the networks architecture, the number of clusters and the latent dimensions. We denote by Li each local contribution to the ELBO:
Li(θ, φ; xi, zi,d,β) = Eq(di,zi) [log p(xi|zi, di,β)] −Eq(di) [DKL (q(zi|xi)‖p(zi|di,β))]−DKL (q(di|zi)‖p(di)))
(11)
The first part of equation 10 is an expectation over the global approximate posterior of the so-called local ELBO. This local ELBO differs from the vanilla ELBO proposed by Kingma & Welling (2013) in the regularizer for the discrete variable di, which is composed by the typical reconstruction term of each sample and two KL regularizers: one for zi, expected over di, and the other over di. The second part in equation 10 is a regularizer on the global posterior. The expectations over the discrete variable di are tractable and thus, analytically marginalized.
In contrast with GMVAE (Figure 1 (b)), in UG-VAE, β is shared by a group of observations, therefore the parameters of the mixture are the same for all the samples in a batch. In this manner, within each optimization step, the encoder q(β|X,Z) only learns from the global information obtained from the product of Gaussian contributions of every observation, with the aim at configuring the mixture to improve the representation of each datapoint in the batch, by means of p(Z|β,X) and p(X|Z,β). Hence, the control of the mixture is performed by using global information. In contrast with ML-VAE (whose encoder q(CG|X) is also global, but the model does not include a mixture), in UG-VAE, the β encoder incorporates information about which component each observation belongs to, as the weights of the mixture inferred by q(d|Z) are used to obtain q(β|X,Z). Thus, while each cluster will represent different local features, moving β will affect all the clusters. In other words, modifying β will have some effect in each local cluster. As the training progresses, the encoder q(β|X,Z) learns which information emerging from each batch of data allows to move the cluster in a way that the ELBO increases.

heading: 4 EXPERIMENTS

text: In this section we demonstrate the ability of the UG-VAE model to infer global factors of variation that are common among samples, even when coming from different datasets. In all cases, we have not validated in depth all the networks used, we have merely rely on encoder/decoder networks proposed in state-of-the-art VAE papers such as Kingma & Welling (2013), Bouchacourt et al. (2018) or Higgins et al. (2016). Our results must be hence regarded as a proof of concept about the flexibility and representation power of UG-VAE, rather than fine-tuned results for each case. Hence there is room for improvement in all cases. Details about network architecture and training parameters are provided in the supplementary material.

heading: 4.1 UNSUPERVISED LEARNING OF GLOBAL FACTORS

text:
In this section we first asses the interpretability of the global disentanglement features inferred by UG-VAE over both CelebA and MNIST. In Figure 3 we show samples of the generative model as we explore both the global and local latent spaces. We perform a linear interpolation with the aim at exploring the hypersphere centered at the mean of the distribution and with radius σi for each dimension i. To maximize the variation range across every dimension, we move diagonally through the latent space. Rows correspond to an interpolation on the global β between [−1, 1] on every dimension (p(β) follows a standard Gaussian). As the local p(z|d,β) (equation 3) depends on d and β, if we denote µz = µ (d) z (β), the local interpolation goes from [µz0 − 3, µz1 − 3, ...µzd − 3] to [µz0 + 3, µz1 + 3, ..., µzd + 3]. The range of ±3 for the local interpolation is determined to cover the variances Σ(d)z (β) that we observe upon training the model for MNIST and CelebA. The every image in Figure 3 correspond to samples from a different cluster (fixed values of d), in order to facilitate the interpretability of the information captured at both local and global levels. By using this set up, we demonstrate that the global information tuned by β is different and clearly interpretable inside each cluster.
The total number of clusters is set to K = 20 for CelebA and K = 10 for MNIST. Three of these components are presented in Figure 3. We can observe that each row (each value of β) induces a
shared generative factor, while z is in charge of variations inside this common feature. For instance, in CelebA (top), features like skin color, presence of beard or face contrast are encoded by the global variable, while local variations like hair style or light direction are controlled by the local variable. In a simple dataset like MNIST (bottom), results show that handwriting global features as cursive style, contrast or thickness are encoded by β, while the local z defines the shape of the digit. The characterization of whether these generative factors are local/global is based on an interpretation of the effect that varying z and β provokes in each image within a batch, and in the whole batch of images, respectively. In the supplementary material, we reproduce the same figures for the all the clusters, in which we can appreciate that there is a significant fraction of clusters with visually interpretable global/local features.
We stress here again the fact that the UG-VAE training is fully unsupervised: data batches during training are completely randomly chosen from the training dataset, with no structured correlation whatsoever. Unlike other approaches for disentanglement, see Higgins et al. (2016) or Mathieu et al. (2019), variational training in UG-VAE does not come with additional ELBO hyperparameters that need to be tuned to find a proper balance among terms in the ELBO.
One of the main contributions in the design of UG-VAE is the fact that, unless we include a clustering mixture prior in the local space controlled by the global variable β, unsupervised learning of global factors is non-informative. To illustrate such a result, in Figure 4 we reproduce the results in Figure 3 but for a probabilistic model in which the discrete local variable d is not included. Namely, we use the ML-VAE in Figure 2(c) but we trained it with random data batches. In this case, the local space is uni-modal given β and we show interpolated values between -1 to 1. Note that the disentanglement effect of variations in both β and z is mild and hard to interpret.

heading: 4.2 DOMAIN ALIGNMENT

text:
In this section, we evaluate the UG-VAE performance in an unsupervised domain alignment setup. During training, the model is fed with data batches that include random samples coming from two different datasets. In particular, we train our model with a mixed dataset between CelebA and 3D FACES Paysan et al. (2009), a dataset of 3D scanned faces, with a proportion of 50% samples from each dataset inside each batch.
Upon training with random batches, in Figure 5, we perform the following experiment using domain supervision to create test data batches. We create two batches containing only images from CelebA and 3D FACES. Let β1 and β2 be the mean global posterior computed using (8) associated for each batch. For two particular images in these two batches, let z1 and z2 be the mean local posterior of these two images, computed using (3). Figure 5 (a) shows samples of the UG-VAE model when we linearly interpolate between β1 and β2 (rows) and between z1 and z2 (columns)1. Certainly β is capturing the domain knowledge. For fixed z, e.g. z1 in the first column, the interpolation between β1 and β2 is transferring the CelebA image into the 3D FACES domain (note that background is turning white, and the image is rotated to get a 3D effect). Alternatively, for fixed β, e.g. β1 in the first row, interpolating between z1 and z2 modifies the first image into one that keeps the domain but resembles features of the image in the second domain, as face rotation.
In Figure 5(b) we show the 2D t-SNE plot of the posterior distribution of β for batches that are random mixtures between datasets (grey points), batches that contain only CelebA faces (blue squares), and batches that contain only 3D faces (green triangles). We also add the corresponding points of the β1 and β2 interpolation in Figure 5(a). In Figure 5(c), we reproduce the experiment in (a) but interpolating between two images and values of β that correspond to the same domain (brown interpolation line in Figure 5(b)). As expected, the interpolation of β in this case does not change the domain, which suggests that the domain structure in the global space is smooth, and that the interpolation along the local space z modifies image features to translate one image into the other. In Figure 6 experiments with more datasets are included. When mixing the 3DCars dataset (Fidler et al. (2012)) with the 3D Chairs dataset (Aubry et al. (2014)), we find that certain correlations between cars and chairs are captured. In Figure 6 (a), interpolating between a racing car and an office desk chair leads to a white car in the first domain (top right) and in a couch (bottom left). In Figure 6 (b), when using the 3D Cars along with the Cars Dataset (Krause et al. (2013)), rotations in the cars are induced.
Finally, in the supplementary material we show that, as expected, the rich structured captured by UG-VAE illustrated in Figure 5 is lost when we do not include the clustering effect in the local space, i.e. if we use ML-VAE with unsupervised random data batches, and all the transition between domains is performed within the local space.
1Note that since both β and z are deterministically interpolated, the discrete variable d plays no role to sample from the model.

heading: 4.3 UG-VAE REPRESENTATION OF STRUCTURED NON-TRIVIAL DATA BATCHES

text:
In the previous subsection, we showed that the UG-VAE global space is able to separate certain structure in the data batches (e.g. data domain) even though during training batches did not present such an explicit correlation. Using UG-VAE trained over CelebA with unsupervised random batches of 128 images as a running example, in this section we want to further demonstrate this result.
In Figure 7 we show the t-SNE 2D projection of structured batches using the posterior β distribution in (8) over CelebA test images. In Figure 7(a), we display the distribution of batches containing only men and women, while in Figure 7(b) the distribution of batches containing people with black or blond hair. In both cases we show the distribution of randomly constructed batches as the ones in the training set. To some extend, in both cases we obtain separable distributions among the different kinds of batches. A quantitive evaluation can be found in Table 1, in which we have used samples from the β distribution to train a supervised classifier to differentiate between different types of batches. When random batches are not taken as a class, the separability is evident. When random batches are included, it is expected that the classifier struggles to differentiate between a batch that contains 90% of male images and a batch that only contain male images, hence the drop in accuracy for the multi-case problem.
An extension with similar results and figures for another interpretation of global information capturing are exposed in the supplementary material, using structured grouped batches in MNIST dataset. In this experiment, the groups are digits that belong to certain mathematical series, including even numbers, odd numbers, Fibonacci series and prime numbers, and we prove that the model is able to discriminate among their global posterior representations.

heading: 5 CONCLUSION

text: In this paper we have presented UG-VAE, an unsupervised generative probabilistic model able to capture both local data features and global features among batches of data samples. Unlike similar approaches in the literature, by combining a structured clustering prior in the local latent space with a global latent space with Gaussian prior and a more structured variational family, we have demonstrated that interpretable group features can be inferred from the global latent space in a completely unsupervised fashion. Model training does not require artificial manipulation of the ELBO term to force latent interpretability, which makes UG-VAE stand out w.r.t. most of the current disentanglement approaches using VAEs. The ability of UG-VAE to infer diverse features from the training set is further demonstrated in a domain alignment setup, where we show that the global space allows interpolation between domains, and also by showing that images in correlated batches of data, related by non-trivial features such as hair color or gender in CelebA, define identifiable structures in the posterior global latent space distribution.