WiSE-ALE: Wide Sample Estimator for Approximate Latent Embedding

TL;DR

WiSE-ALE introduces a novel variational lower bound for VAEs that emphasizes aggregate dataset statistics, enabling better reconstruction quality and more flexible latent representations.

Contribution

It proposes an alternative lower bound for VAEs that focuses on aggregate posterior constraints, leading to improved reconstruction and representation flexibility.

Findings

Achieves state-of-the-art reconstruction quality.

Retains smooth, compact latent representations.

Offers an efficient learning algorithm, WiSE-ALE.

Abstract

Variational Auto-encoders (VAEs) have been very successful as methods for forming compressed latent representations of complex, often high-dimensional, data. In this paper, we derive an alternative variational lower bound from the one common in VAEs, which aims to minimize aggregate information loss. Using our lower bound as the objective function for an auto-encoder enables us to place a prior on the bulk statistics, corresponding to an aggregate posterior for the entire dataset, as opposed to a single sample posterior as in the original VAE. This alternative form of prior constraint allows individual posteriors more flexibility to preserve necessary information for good reconstruction quality. We further derive an analytic approximation to our lower bound, leading to an efficient learning algorithm - WiSE-ALE. Through various examples, we demonstrate that WiSE-ALE can reach excellent reconstruction quality in comparison to other state-of-the-art VAE models, while still retaining the ability to learn a smooth, compact representation.

Figures

Equations

$$p_{\theta }(\mathbf{x},\mathbf{z}|{\mathcal{D}}_N)$$

$$=p_{\theta }({\mathbf{x}}_1,{\mathbf{x}}_2,\cdots ,{\mathbf{x}}_N|{\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_N)p_{\theta }({\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_N|{\mathcal{D}}_N)$$

$$=\prod _{i=1}^N{p}_{\theta }({\mathbf{x}}_i|{\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_N)\prod _{i=1}^N{p}_{\theta }({\mathbf{z}}_i|{\mathcal{D}}_N)$$

$$=\prod _{i=1}^N{p}_{\theta }({\mathbf{x}}_i|{\mathbf{z}}_i)\prod _{i=1}^N{p}_{\theta }({\mathbf{z}}_i|{\mathcal{D}}_N)$$

$$\displaystyle=\prod_{i=1}^{N}\bigg{(}p_{\theta}({\mathbf{x}}_{i}|{\mathbf{z}}_{i})\;p_{\theta}({\mathbf{z}}_{i}|\mathcal{D}_{N})\bigg{)}.$$

$$p_{\theta }(\mathbf{x},\mathbf{z}|{\mathcal{D}}_N)$$

$$=p_{\theta }({\mathbf{x}}_1,{\mathbf{x}}_2,\cdots ,{\mathbf{x}}_N,|\mathbf{z})p_{\theta }(\mathbf{z}|{\mathcal{D}}_N)$$

$$=p_{\theta }(\mathbf{z}|{\mathcal{D}}_N)\prod _{i=1}^N{p}_{\theta }({\mathbf{x}}_i|\mathbf{z}),$$

$$p(\mathbf{z}|{\mathcal{D}}_N)=\int p_{\theta }(\mathbf{z}|{\mathbf{x}}_j)p({\mathbf{x}}_j|{\mathcal{D}}_N)\text{d}{\mathbf{x}}_j=\sum _{i=1}^N{p}_{\theta }(\mathbf{z}|{\mathbf{x}}_j={\mathbf{x}}^{(i)})P({\mathbf{x}}_j={\mathbf{x}}^{(i)}|{\mathcal{D}}_N)=\frac{1}{N}\sum _{i=1}^N{p}_{\theta }(\mathbf{z}|{\mathbf{x}}^{(i)}),$$

$$q(\mathbf{z}|{\mathcal{D}}_N)=\frac{1}{N}\sum _{i=1}^N{q}_{\varphi }(\mathbf{z}|{\mathbf{x}}^{(i)}).$$

$$\displaystyle D_{\mathrm{KL}}\big{[}\,q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})\|p_{\theta}({\mathbf{z}}|\mathcal{D}_{N})\big{]}\,=\,\mathbb{E}_{q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})}\Bigg{[}\log\frac{q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})}{p_{\theta}({\mathbf{z}}|\mathcal{D}_{N})}\Bigg{]}.$$

$$\displaystyle D_{\mathrm{KL}}\big{[}\,q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})\|p_{\theta}({\mathbf{z}}|\mathcal{D}_{N})\big{]}=\;\mathbb{E}_{q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})}\big{[}\log\,q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})-\log\,p_{\theta}(\mathcal{D}_{N}|{\mathbf{z}})-\log p({\mathbf{z}})\big{]}\,+\,\log\,p(\mathcal{D}_{N}).$$

$$\displaystyle\log\,p(\mathcal{D}_{N})\,-\,D_{\mathrm{KL}}\big{[}\,q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})\|p_{\phi}({\mathbf{z}}|\mathcal{D}_{N})\big{]}\;=\;\mathbb{E}_{q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})}\big{[}\log p_{\theta}(\mathcal{D}_{N}|{\mathbf{z}})\big{]}-D_{\mathrm{KL}}\big{[}\;q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})\|p({\mathbf{z}})\big{]}.$$

$$\displaystyle\log p(\mathcal{D}_{N})\,\geq\,L^{\textrm{WiSE-ALE}}(\phi,\theta;\mathcal{D}_{N})\,=\,\underbrace{\mathbb{E}_{q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})}\big{[}\log p_{\theta}(\mathcal{D}_{N}|{\mathbf{z}})\big{]}}_{\text{ \leavevmode\hbox to11.75pt{\vbox to11.75pt{\pgfpicture\makeatletter\hbox{\hskip 5.87407pt\lower-5.87407pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{ }\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }{ {{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{ }{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{{}\pgfsys@moveto{5.67407pt}{0.0pt}\pgfsys@curveto{5.67407pt}{3.13374pt}{3.13374pt}{5.67407pt}{0.0pt}{5.67407pt}\pgfsys@curveto{-3.13374pt}{5.67407pt}{-5.67407pt}{3.13374pt}{-5.67407pt}{0.0pt}\pgfsys@curveto{-5.67407pt}{-3.13374pt}{-3.13374pt}{-5.67407pt}{0.0pt}{-5.67407pt}\pgfsys@curveto{3.13374pt}{-5.67407pt}{5.67407pt}{-3.13374pt}{5.67407pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@stroke\pgfsys@invoke{ } }{{{{}}\pgfsys@beginscope\pgfsys@invoke{ }\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-1.75pt}{-2.25555pt}\pgfsys@invoke{ }\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{ }\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{{1}} }}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} } \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}} {Reconstruction likelihood}}}\;-\;\underbrace{D_{\mathrm{KL}}\big{[}\,q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})\|p({\mathbf{z}})\big{]}}_{\text{ \leavevmode\hbox to11.75pt{\vbox to11.75pt{\pgfpicture\makeatletter\hbox{\hskip 5.87407pt\lower-5.87407pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{ }\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }{ {{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{ }{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{{}\pgfsys@moveto{5.67407pt}{0.0pt}\pgfsys@curveto{5.67407pt}{3.13374pt}{3.13374pt}{5.67407pt}{0.0pt}{5.67407pt}\pgfsys@curveto{-3.13374pt}{5.67407pt}{-5.67407pt}{3.13374pt}{-5.67407pt}{0.0pt}\pgfsys@curveto{-5.67407pt}{-3.13374pt}{-3.13374pt}{-5.67407pt}{0.0pt}{-5.67407pt}\pgfsys@curveto{3.13374pt}{-5.67407pt}{5.67407pt}{-3.13374pt}{5.67407pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@stroke\pgfsys@invoke{ } }{{{{}}\pgfsys@beginscope\pgfsys@invoke{ }\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-1.75pt}{-2.25555pt}\pgfsys@invoke{ }\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{ }\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{{2}} }}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} } \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}} {Prior constraint}}}.$$

$$\displaystyle\mathbb{E}_{q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})}\big{[}\log p_{\theta}(\mathcal{D}_{N}|{\mathbf{z}})\big{]}=\frac{1}{N}\sum_{i=1}^{N}\mathbb{E}_{q_{\phi}({\mathbf{z}}|{\bm{x}}^{(i)})}\big{[}\log p_{\theta}(\mathcal{D}_{N}|{\mathbf{z}})\big{]}.$$

$$\log p_{\theta }({\mathcal{D}}_N|\mathbf{z})=\log \prod _{j=1}^N{p}_{\theta }({\mathbf{x}}^{(j)}|\mathbf{z})=\sum _{j=1}^N\log p_{\theta }({\mathbf{x}}^{(j)}|\mathbf{z}).$$

$$\displaystyle\mathbb{E}_{q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})}[\log p_{\theta}(\mathcal{D}_{N}|{\mathbf{z}})]=\sum_{i=1}^{N}\mathbb{E}_{q_{\phi}({\mathbf{z}}|{\bm{x}}^{(i)})}\Bigg{[}\frac{1}{N}\sum_{j=1}^{N}\log p_{\theta}({\bm{x}}^{(j)}|{\mathbf{z}})\Bigg{]}.$$

$$\displaystyle\mathbb{E}_{q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})}[\log p_{\theta}(\mathcal{D}_{N}|{\mathbf{z}})]\leq\sum_{i=1}^{N}\mathbb{E}_{q_{\phi}({\mathbf{z}}|{\bm{x}}^{(i)})}\Bigg{[}\log\bigg{(}\frac{1}{N}\sum_{j=1}^{N}p_{\theta}({\bm{x}}^{(j)}|{\mathbf{z}})\bigg{)}\Bigg{]}.$$

$$\displaystyle\mathbb{E}_{q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})}[\log p_{\theta}(\mathcal{D}_{N}|{\mathbf{z}})]\leq\sum_{i=1}^{N}\mathbb{E}_{q_{\phi}({\mathbf{z}}|{\bm{x}}^{(i)})}\Big{[}\log p_{\theta}({\bm{x}}^{(i)}|{\mathbf{z}})\Big{]}-N\log N.$$

$$\displaystyle D_{\mathrm{KL}}\big{[}\,q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})\|p({\mathbf{z}})\big{]}=\frac{1}{N}\sum_{i=1}^{N}\Big{(}\mathbb{E}_{q_{\phi}({\mathbf{z}}|{\bm{x}}^{(i)})}\big{[}\log q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})\big{]}-\mathbb{E}_{q_{\phi}({\mathbf{z}}|{\bm{x}}^{(i)})}\big{[}\log p({\mathbf{z}})\big{]}\Big{)}.$$

$$\displaystyle D_{\mathrm{KL}}\big{[}\,q_{\phi}({\mathbf{z}}$$

$$\displaystyle\,=\,\frac{1}{N}\sum_{i=1}^{N}\Big{(}\log\mathbb{E}_{q_{\phi}({\mathbf{z}}|{\bm{x}}^{(i)})}\big{[}q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})\big{]}\Big{)}\,-\,\frac{1}{N}\sum_{i=1}^{N}\Big{(}\mathbb{E}_{q_{\phi}({\mathbf{z}}|{\bm{x}}^{(i)})}\big{[}\log p({\mathbf{z}})\big{]}\Big{)}.$$

$${\text{KL}}_{\text{approx}}^{\text{UB}}$$

$$\text{mbox}whereA$$

$$B$$

$$C$$

$$\displaystyle L_{\textrm{approx}}^{\textrm{WiSE-ALE}}(\phi,\theta;\mathcal{D}_{N})\;=\;\sum_{i=1}^{N}\mathcal{L}(\phi,\theta|{\bm{x}}^{(i)})\;-\;KL\big{[}\,q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})\,||\,p({\mathbf{z}})\big{]},$$

$$L^{\text{AEVB}}(\varphi ,\theta ;{\mathcal{D}}_N)$$

$$\displaystyle L^{\textrm{WiSE-ALE}}(\phi,\theta;\mathcal{B}_{M})=\sum_{i=1}^{M}\Big{(}\mathbb{E}_{q_{\phi}({\mathbf{z}}|{\bm{x}}^{(i)})}\big{[}\log p_{\theta}({\bm{x}}^{(i)}|{\mathbf{z}})\big{]}\Big{)}\,-\,D_{\mathrm{KL}}\big{[}\;q_{\phi}({\mathbf{z}}|\mathcal{B}_{M})\|p({\mathbf{z}})\big{]}.$$

$$L^{\text{WiSE-ALE}}(\varphi ,\theta ;{\mathcal{D}}_N)\approx \frac{N}{M}{L}^{\text{WiSE-ALE}}(\varphi ,\theta ;{\mathcal{B}}_M).$$

Full text

Machine Learning, ICML

1 Introduction

Unsupervised learning is a central task in machine learning. Its objective can be informally described as learning a representation of some observed forms of information in a way that the representation summarizes the overall statistical regularities of the data (Barlow, 1989). Deep generative models are a popular choice for unsupervised learning, as they marry deep learning with probabilistic models to estimate a joint probability between high dimensional input variables ${\mathbf{x}}$ and unobserved latent variables ${\mathbf{z}}$. Early successes of deep generative models came from Restricted Boltzmann Machines (Hinton & Salakhutdinov, 2006) and Deep Boltzmann Machines (Salakhutdinov & Hinton, 2009), which aim to learn a compact representation of data. However, the fully stochastic nature of the network requires layer-by-layer pre-training using MCMC-based sampling algorithms, resulting in heavy computation cost.

Kingma & Welling (2013) consider the objective of optimizing the parameters in an auto-encoder network by deriving an analytic solution to a variational lower bound of the log likelihood of the data, leading to the Auto-Encoding Variational Bayes (AEVB) algorithm. They apply a reparameterization trick to maximally utilize deterministic mappings in the network, significantly simplifying the training procedure and reducing instability. Furthermore, a regularization term naturally occurs in their model, allowing a prior $p({\mathbf{z}})$ to be placed over every sample embedding $q({\mathbf{z}}|{\mathbf{x}})$. As a result, the learned representation becomes compact and smooth; see e.g. Fig. 1 where we learn a 2D embedding of MNIST digits using 4 different methods and visualize the aggregate posterior distribution of 64 random samples in the learnt 2D embedding space.

However, because the choice of the prior is often uninformative, the smoothness constraint imposed by this regularization term can cause information loss between the input samples and the latent embeddings, as shown by the merging of individual embedding distributions in Fig. 1(d) (especially in the outer areas away from zero code). Extreme effects of such behaviours can be noticed from $\beta$-VAE (Higgins et al., 2016), a derivative algorithm of AEVB which further increases the weighting on the regularizing term with the aim of learning an even smoother, disentangled representation of the data. As shown in Fig. 1(e), the individual embedding distributions are almost indistinguishable, leading to an overly severe information bottleneck which can cause high rates of distortion (Tishby et al., 1999). The other end of the spectrum can be indicated by Fig. 1(b), where perfect reconstruction can be achieved but the learnt embedding distributions appear to severely sharp, indicating a latent representation which is heavily non-smooth and likely to be unstable due to a small amount of noise.

In this paper, we propose WiSE-ALE (a wide sample estimator), which imposes a prior on the bulk statistics of a mini-batch of latent embeddings. Learning under our WiSE-ALE objective does not penalize individual embeddings lying away from the zero code, so long as the aggregate distribution (the average of all individual embedding distributions) does not violate the prior significantly. Hence, our approach mitigates the distortion caused by the current form of the prior constraint in the AEVB objective. Furthermore, the objective of our WiSE-ALE algorithm is derived by applying variational inference in a simple latent variable model (Section 2) and with further approximation, we derive an analytic form of the learning objective, resulting in efficient learning algorithm.

In general, the latent representation learned using our algorithm enjoys the following properties: 1) smoothness, as indicated in Fig. 1(d), the probability density for each individual embedding distribution decays smoothly from the peak value; 2) compactness, as individual embeddings tend to occupy a maximal local area in the latent space with minimal gaps in between; and 3) separation, indicated by the narrow, but clear borders between neighbouring embedding distributions as opposed to the merging seen in AEVB. In summary, our contributions are:

• •

An alternative variational lower bound to the data log likelihood is derived, allowing us to impose prior constraint on the bulk statistics of a mini-batch embedding distributions.

• •

Analytic approximations to the lower bound are derived, allowing efficient optimization without sampling-based methods and leading to our WiSE-ALE algorithm.

• •

Extensive analysis of our algorithm’s performance in comparison with three related VAE algorithms, namely AEVB, $\beta$-VAE and WAE (Tolstikhin et al., 2017).

In the rest of the paper, we first review directed graphical models in Section 2. We then derive our variational lower bound and its analytic approximations in Section 3. Related work is discussed in Section 4. Experiment results are analyzed in Section 5, leading to conclusions in Section 6.

2 Background: Latent Variable Models

Here we briefly review the latent variable model that allows variational inference through an auto-encoding task and highlight the difference between the latent variable model for our WiSE-ALE algorithm and that for the AEVB algorithm (Kingma & Welling, 2013).

Given $N$ observations of input samples ${\bm{x}}\in R^{d_{x}}$ denoted \mathcal{D}_{N}=\big{(}{\bm{x}}^{(1)},{\bm{x}}^{(2)},\cdots,{\bm{x}}^{(N)}\big{)}, we assume ${\mathbf{x}}$ is generated from a latent variable ${\mathbf{z}}\in R^{d_{z}}$ of a much lower dimension. Here we denote ${\mathbf{x}}$ and ${\mathbf{z}}$ as random variables, ${\bm{x}}^{(i)}$ or ${\bm{z}}^{(i)}$ as the $i$-th input or latent code sample (i.e. a vector), and ${\mathbf{x}}_{i}$ and ${\mathbf{z}}_{i}$ as the random variable for ${\bm{x}}^{(i)}$ and ${\bm{z}}^{(i)}$. As shown in Fig. 2(a), this generative process can be modelled by a simple directed graphical model (Jordan et al., 1999), which models the joint probability distribution $p_{\theta}({\mathbf{x}},{\mathbf{z}}|\mathcal{D}_{N})=p_{\theta}({\mathbf{x}}|{\mathbf{z}})p({\mathbf{z}}|\mathcal{D}_{N})=p_{\theta}({\mathbf{z}}|{\mathbf{x}})p({\mathbf{x}}|\mathcal{D}_{N})$ between ${\mathbf{x}}$ and ${\mathbf{z}}$, given the current observations $\mathcal{D}_{N}$. $p({\mathbf{z}}|\mathcal{D}_{N})$ is the latent distribution given $\mathcal{D}_{N}$, $p({\mathbf{x}}|\mathcal{D}_{N})$ is the data distribution for $\mathcal{D}_{N}$ and $p_{\theta}({\mathbf{x}}|{\mathbf{z}})$ and $p_{\theta}({\mathbf{z}}|{\mathbf{x}})$ denote the complex transformation from the latent to the input space and reverse, where the transformation mapping is parameterised by $\theta$. The learning task is to estimate the optimal set of ${\bm{\theta}}$ so that this latent variable model can explain the data $\mathcal{D}_{N}$ well.

As the inference of the latent variable ${\mathbf{z}}$ given ${\mathbf{x}}$ (i.e. $p_{\theta}({\mathbf{z}}|{\mathbf{x}})$) cannot be directly estimated because $p({\mathbf{x}}|\mathcal{D}_{N})$ is unknown, both AEVB (Fig. 2(b)) and our WiSE-ALE (Fig. 2(c)) resort to variational method to approximate the target distribution $p_{\theta}({\mathbf{z}}|{\mathbf{x}})$ by a proposal distribution $q_{\phi}({\mathbf{z}}|{\mathbf{x}})$ with the modified learning objective that both $\theta$ and $\phi$ are optimised so that the model can explain the data well and $q_{\phi}({\mathbf{z}}|{\mathbf{x}})$ approaches $p_{\theta}({\mathbf{z}}|{\mathbf{x}})$. The primary difference between the AEVB model and our WiSE-ALE model lies in how the joint probability $p_{\theta}({\mathbf{x}},{\mathbf{z}}|\mathcal{D}_{N})$ is modelled and specifically whether we assume an individual random variable for each latent code ${\bm{z}}^{(i)}$. The AEVB model assumes a pair of random variables $({\mathbf{x}}_{i},{\mathbf{z}}_{i})$ for each ${\bm{x}}^{(i)}$ and estimates the joint probability as

[TABLE]

The equality between Eq. 2 and Eq. 3 can only be made with the assumption that the generation process for each ${\mathbf{x}}_{i}$ is independent (first product in Eq. 3) and each ${\mathbf{z}}_{i}$ is also independent (second product in Eq. 3). Such interpretation of the joint probability leads to the latent variable model in Fig. 2(b) and the prior constraint (often taken as $\mathcal{N}(0,{\bm{I}})$ to encourage shrinkage when no data is observed) is imposed on every ${\mathbf{z}}_{i}$.

In contrast, our WiSE-ALE model takes a single random variable to estimate the latent distribution for the entire dataset $\mathcal{D}_{N}$. Hence, the joint probability in our model can be broken down as

[TABLE]

leading to the latent variable model illustrated in Fig. 2(c). The only assumption we make in our model is assuming the generative process of different input samples given the latent distribution of the current dataset as independent, which we consider as a sensible assumption. More significantly, we do not require independence between different ${\mathbf{z}}_{i}$ as opposed to the AEVB model, leading to a more flexible model. Furthermore, the prior constraint in our model is naturally imposed on the aggregate posterior $p({\mathbf{z}}|\mathcal{D}_{N})$ for the entire dataset, leading to more flexibility for each individual sample latent code to shape an embedding distribution to preserve a better quality of information about the corresponding input sample.

Neural networks can be used to parameterize $p_{\theta}({\mathbf{x}}_{i}|{\mathbf{z}}_{i})$ in the generative model and $q_{\phi}({\mathbf{z}}_{i}|{\mathbf{x}}_{i})$ in the inference model from the AEVB latent variable model or $p_{\theta}({\mathbf{x}}_{i}|{\mathbf{z}})$ and $q_{\phi}({\mathbf{z}}|{\mathbf{x}}_{i})$ correspondingly from our WiSE-ALE latent variable model. Both networks can be implemented through an auto-encoder network illustrated in Fig. 2(d).

3 Our Method

In this section, we first define the aggregate posterior distribution $p({\mathbf{z}}|\mathcal{D}_{N})$ which serves as a core concept in our WiSE-ALE proposal. We then derive a variational lower bound to the marginal log likelihood of the data $\log p(\mathcal{D}_{N})$ with the focus on the aggregate posterior distribution. Further, analytic approximation to the lower bound is derived, allowing efficient optimization of the model parameters and leading to our WiSE-ALE learning algorithm. Intuition of our proposal is also discussed.

3.1 Aggregate Posterior

Here we formally define the aggregate posterior distribution $p({\mathbf{z}}|\mathcal{D}_{N})$, i.e. the latent distribution given the entire dataset $\mathcal{D}_{N}$. Considering

[TABLE]

we have the aggregate posterior distribution for the entire dataset as the average of all the individual sample posteriors. The second equality in Eq. 9 is made approximating the integral through summation. The third equality is obtained following the conventional assumption in the VAE literature that each input sample, ${\bm{x}}^{(i)}$, is drawn from the data set $\mathcal{D}_{N}$ with equal probability, i.e. $P({\bm{x}}^{(i)}|\mathcal{D}_{N})=\frac{1}{N}$. Similarly, for the estimated aggregate posterior distribution $q({\mathbf{z}}|\mathcal{D}_{N})$, we have

[TABLE]

3.2 Alternative Variational Lower Bound (LB)

To carry out variational inference, we minimize the KL divergence between the estimated and the true aggregate posterior distributions $q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})$ and $p_{\theta}({\mathbf{z}}|\mathcal{D}_{N})$, i.e.

[TABLE]

Substituting $p_{\theta}({\mathbf{z}}|\mathcal{D}_{N})=\frac{p_{\theta}(\mathcal{D}_{N}|{\mathbf{z}})\,p({\mathbf{z}})}{p_{\theta}(\mathcal{D}_{N})}$ in Eq. 11 and breaking down the products and fractions inside the $\log$, we have

[TABLE]

Re-arranging the above equation, we have

[TABLE]

As D_{\mathrm{KL}}\big{[}q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})\|p_{\phi}({\mathbf{z}}|\mathcal{D}_{N})\big{]} is non-negative, we have obtained a variational lower bound $L^{\textrm{WiSE-ALE}}(\phi,\theta;\mathcal{D}_{N})$ to the marginal log likelihood of the data $\log p(\mathcal{D}_{N})$ as

[TABLE]

There are two terms in the derived lower bound:
1 a reconstruction likelihood term that indicates how likely the current dataset $\mathcal{D}_{N}$ are generated by the aggregate latent posterior distribution $q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})$ and
2 a prior constraint that penalizes severe deviation of the aggregate latent posterior distribution $q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})$ from the preferred prior $p({\mathbf{z}})$, acting naturally as a regularizer. By maximizing the lower bound $L^{\textrm{WiSE-ALE}}(\phi,\theta;\mathcal{D}_{N})$ defined in Eq. 12, we are approaching to $\log p(\mathcal{D}_{N})$ and, hence, obtaining a set of parameters $\theta$ and $\phi$ that find a natural balance between a good reconstruction likelihood (good explanation of the observed data) and a reasonable level of compliance to the prior assumption (achieving some preferable properties of the posterior distribution, such as smoothness and compactness).

3.3 Approximation of the Proposed Lower Bound

To allow fast and efficient optimization of the model parameters $\theta$ and $\phi$, we derive analytic approximations for the two terms in our proposed lower bound (Eq. 12).

3.3.1 Approximation to Reconstruction Likelihood Term

To approximate
1 reconstruction likelihood term in Eq. 12, we first substitute the definition of the approximate aggregate posterior given in Eq. 10 in the expectation operation in \mathbb{E}_{q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})}\big{[}\log p_{\theta}(\mathcal{D}_{N}|{\mathbf{z}})\big{]}, i.e.

[TABLE]

Now we can decompose the $p_{\theta}(\mathcal{D}_{N}|{\mathbf{z}})$ as a product of individual sample likelihood, due to the conditional independence, i.e.

[TABLE]

Substituting this into Eq. 13, we have

[TABLE]

Eq. 15 can be used to evaluate the reconstruction likelihood for $\mathcal{D}_{N}$. However, learning directly with this reconstruction estimate does not lead to convergence in our experiments and the computation is quite costly, as for every evaluation of the reconstruction likelihood, we need to evaluate $N^{2}$ expectation operations. We choose to simplify the reconstruction likelihood further to be able to reach convergence during learning at the cost of losing the lower bound property of the objective function $L^{\textrm{WiSE-ALE}}(\phi,\theta;\mathcal{D}_{N})$. Firstly, we apply Jensen inequality to the term inside the expectation in Eq. 15, leading to an upper bound of the reconstruction likelihood term as

[TABLE]

Now $(N-1)$ sample-WiSE-ALE likelihood distributions in the summation inside the $\log$ can be dropped with the assumption that the $p_{\theta}({\bm{x}}^{(j)}|{\mathbf{z}})$ will only be non-zero if ${\mathbf{z}}$ is sampled from the posterior distribution of the same sample ${\bm{x}}^{(j)}$ at the encoder, i.e. $i=j$. Therefore, the approximation becomes

[TABLE]

Using the approximation of the reconstruction likelihood term given by Eq. 17 rather than Eq. 15, we are able to reach convergence efficiently during learning at the cost of the estimated objective no longer remaining a lower bound to $\log p(\mathcal{D}_{N})$. Details of deriving the above approximation are given in Appendix A.

3.3.2 Approximation to Prior Constraint Term

The
2 prior constraint term D_{\mathrm{KL}}\big{[}\,q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})\|p({\mathbf{z}})\big{]} in our objective function (Eq. 12) evaluates the KL divergence between the approximate aggregate posterior distribution $q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})$ and a zero-mean, unit-variance Gaussian distribution $p({\mathbf{z}})$. Here we assume that each sample-WiSE-ALE posterior distribution can be modelled by a factorial Gaussian distribution, i.e. q_{\phi}({\mathbf{z}}|{\bm{x}}^{(i)})=\prod_{k=1}^{d_{z}}\mathcal{N}\big{(}{\mathbf{z}}_{k}|\,\mu_{k}({\bm{x}}^{(i)}),\sigma_{k}^{2}({\bm{x}}^{(i)})\big{)}, where $k$ indicates the $k$-th dimension of the latent variable ${\mathbf{z}}$ and $\mu_{k}({\bm{x}}^{(i)})$ and $\sigma_{k}^{2}({\bm{x}}^{(i)})$ are the mean and variance of the $k$-th dimension embedding distribution for the input ${\bm{x}}^{(i)}$. Therefore, D_{\mathrm{KL}}\big{[}\,q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})\|p({\mathbf{z}})\big{]} computes the KL divergence between a mixture of Gaussians (as Eq. 10) and $\mathcal{N}(0,{\bm{I}})$. There is no analytical solution for such KL divergences. Hence, we derive an analytic upper bound allowing for efficient evaluation.

Firstly, we substitute $q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})=\frac{1}{N}\sum_{i=1}^{N}q_{\phi}({\mathbf{z}}|{\bm{x}}^{(i)})$ (Eq. 10) to D_{\mathrm{KL}}\big{[}\,q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})\|p({\mathbf{z}})\big{]}, giving

[TABLE]

Applying Jensen inequality, i.e. \mathbb{E}_{x}\big{[}\log f(x)\big{]}\leq\log\mathbb{E}_{x}\big{[}f(x)\big{]}, to the first term inside the summation in Eq. 18, we have

[TABLE]

Taking advantage of the Gaussian assumption for $q_{\phi}({\mathbf{z}}|{\bm{x}}^{(i)})$ and $p({\mathbf{z}})$, we can compute the expectations in Eq. 20 analytically with the result quoted below and the full derivation given in Appendix B.1.

[TABLE]

When the overall objective function $L^{\textrm{WiSE-ALE}}(\phi,\theta;\mathcal{D}_{N})$ in Eq. 12 is maximised, this upper bound approximation will approach the true KL divergence D_{\mathrm{KL}}\big{[}\,q_{\phi}({\mathbf{z}}|\mathcal{D}_{N})\|p({\mathbf{z}})\big{]}, which ensures that the prior constraint on the overall aggregate posterior distribution takes effects.

3.3.3 Overall Objective Functions

Combining results from Section 3.3.1 and 3.3.2, we obtain an analytic approximation $L_{\textrm{approx}}^{\textrm{WiSE-ALE}}(\phi,\theta;\mathcal{D}_{N})$ for the variational lower bound $L^{\textrm{WiSE-ALE}}(\phi,\theta;\mathcal{D}_{N})$ defined in Eq. 12, as shown below:

[TABLE]

where we use \mathcal{L}\big{(}\phi,\theta\,|\,{\bm{x}}^{(i)}\big{)} to denote the sample-WiSE-ALE reconstruction likelihood \mathbb{E}_{q_{\phi}({\mathbf{z}}|{\bm{x}}^{(i)})}\Big{[}\log p_{\theta}({\bm{x}}^{(i)}|{\mathbf{z}})\Big{]} given by Eq. 17 and the KL divergence term is estimated through $\text{KL}_{\text{approx}}^{\text{UB}}$ defined in Eq. 21. Optimizing $L_{\textrm{approx}}^{\textrm{WiSE-ALE}}(\phi,\theta;\mathcal{D}_{N})$ w.r.t the model parameters $\phi$ and $\theta$, we are able to learn a model that naturally balances between a good embedding of the observed data and some preferred properties of the latent embedding distributions, such as smoothness and compactness.

3.4 Comparison of AEVB and WiSE-ALE Learning Objective Functions

Comparing the objective function in our WiSE-ALE algorithm and that proposed in AEVB algorithm (Kingma & Welling, 2013) stated below,

[TABLE]

we notice that the difference lies in the form of prior constraint and the difference is illustrated in Fig. 3. AEVB learning algorithm imposes the prior constraint on every sample embedding and any deviation away from the zero code or the unit variance (e.g. variance of a sample posterior becomes less than 1, as the model becomes more certain about a specific input sample embedding) will incur penalty. In contrast, our WiSE-ALE learning objective imposes the prior constraint on the aggregate posterior distribution, i.e. the average of all the sample embeddings. Such prior constraint will allow more flexibility for each sample posterior to settle at a mean and variance value in favour for good reconstruction quality, while preventing too large mean values (acting as a regulariser) or too small variance values (ensuring smoothness of the learnt latent representation).

To investigate the different behaviours of the two prior constraints more concretely, we consider only two embedding distributions $q({\mathbf{z}}|{\bm{x}}^{(1)})$ and $q({\mathbf{z}}|{\bm{x}}^{(2)})$ (red dashed lines) in a 1D latent space, as shown in Fig. 4. The mean values of the two embedding distributions are fixed to make the analysis simple and their variances are allowed to change. When the variances of the two embedding distributions are large, such as Fig. 4(a), $q({\mathbf{z}}|{\bm{x}}^{(1)})$ and $q({\mathbf{z}}|{\bm{x}}^{(2)})$ have a large area of overlap and it is difficult to distinguish the input samples ${\bm{x}}^{(1)}$ and ${\bm{x}}^{(2)}$ in the latent space. On the other hand, when the two embedding distributions have small variances, such as Fig. 4(c), there is clear separation between ${\bm{x}}^{(1)}$ and ${\bm{x}}^{(2)}$ in the latent space, indicating the embedding only introduces a small level of information loss. Overall, the prior constraint in the AEVB objective favours the embedding distributions much closer to the uninformative $\mathcal{N}(0,{\bm{I}})$ prior, leading to large area of overlap between the individual posteriors, whereas our WiSE-ALE objective allows a wide range of acceptable embedding mean and variance, which will then offer more flexibility in the learnt posteriors to maintain a good reconstruction quality.

3.5 Efficient Mini-Batch Estimator

So far our derivation has been for the entire dataset $\mathcal{D}_{N}$. Given a small subset $\mathcal{B}_{M}$ with $M$ samples randomly drawn from $\mathcal{D}_{N}$, we can obtain a variational lower bound for a mini-batch as:

[TABLE]

When $\mathcal{B}_{M}$ is reasonably large, then $L^{\textrm{WiSE-ALE}}(\phi,\theta;\mathcal{B}_{M})$ becomes an good approximation of $L^{\textrm{WiSE-ALE}}(\phi,\theta;\mathcal{D}_{N})$ through

[TABLE]

Given the expressions for the objective functions derived in Section 3.3, we can compute the gradient for an approximation to the lower bound of a mini-batch $\mathcal{B}_{M}$ and apply stochastic gradient ascent algorithm to iteratively optimize the parameters $\phi$ and $\theta$. We can thus apply our WiSE-ALE algorithm efficiently to a mini-batch and learn a meaningful internal representation of the entire dataset. Algorithmically, WiSE-ALE is similar to AEVB, save for an alternate objective function as per Section 3.3.3. The procedural details of the algorithm are presented in Appendix C.

4 Related Work

Bengio et al. (2013) proposes that a learned representation of data should exhibit some general features, such as smoothness, sparsity and simplicity. These attributes are general, however, and are not tailored to any specific downstream tasks. Requirements from Bayesian decision making (see e.g. (Lacoste-Julien et al., 2011; Cobb et al., 2018)) adds consideration of a target task and proposes latent distribution approximations which optimise the performance over a particular task, as well as conforming to more general properties. The AEVB algorithm (Kingma & Welling, 2013) learns the latent posterior distribution under a reconstruction task, while simultaneously satisfying the prior, thus ensuring that the representation is smooth and compact. However, the prior form of the AEVB algorithm imposes significant influence on the solution space (as discussed in Section 3.4), and leads to a sacrifice of reconstruction quality. Our WiSE-ALE algorithm, however, prioritises the reconstruction task yet still enables globally desirable properties.

WiSE-ALE is, however, not the only algorithm that considers an alternate prior form to mitigate its impact on reconstruction quality. The Gaussian Mixture VAE (Dilokthanakul et al., 2016) uses a Gaussian mixture model to parameterise $p({\mathbf{z}})$, encouraging more flexible sample posteriors. The Adversarial Auto-Encoder (Makhzani et al., 2016) matches the aggregate posterior over the latent variables with a prior distribution through adversarial training. The WAE (Tolstikhin et al., 2017) minimises a penalised form of the Wasserstein distance between the aggregate posterior distribution and the prior, claiming a generalisation of the AAE algorithm under the theory of optimal transport (Villani, 2008). More recently, the Sinkhorn Auto-Encoder (Patrini et al., 2018) builds a formal analysis of auto-encoders using an optimal transport based prior and uses the Sinkhorn algorithm as an alternative to estimate the Wasserstein distance in WAE.

Our work differs from these in two main aspects. Firstly, our objective function can be evaluated analytically, leading to an efficient optimization process. In many of the above work, the optimization involves adversarial training and some hyper-parameter tuning, which leading to less efficient learning and slow or even no convergence. Secondly, our WiSE-ALE algorithm naturally finds a balance between good reconstruction quality and preferred latent representation properties, such as smoothness and compactness, as shown in Fig. 1(c). In contrast, some other work sacrifice the properties of smoothness and compactness severely for improved reconstruction quality, as shown in Fig.  1(b). Many works (Bloesch et al., 2018; Clark et al., 2018) have indicated that those properties of the learnt latent representation are essential for tasks that require optimisation over the latent space.

5 Experiments

We evaluate our WiSE-ALE algorithm in comparison with AEVB, $\beta$-VAE and WAE on the following 3 datasets. The implementation details for all experiments are given in Appendix E.

Sine Wave. We generated 200,000 sine waves with small random noise: $x(t)=A\,\text{sin}(2\pi ft+\varphi)+\epsilon$, each containing 256 samples, with independently sampled frequency $f\sim\text{Unif}(0,20\text{Hz})$, phase angle $\varphi\sim\text{Unif}(0,2\pi)$ and amplitude $A\sim\text{Unif}(0,2)$. 2. 2.

MNIST (LeCun, 1998). 70,000 $28\times 28$ binary images that contain hand-written digits. 3. 3.

CelebA (Liu et al., 2015). 202,599 RGB images of aligned celebrity faces of $218\times 178$ are cropped to square images of $178\times 178$ and resized to $64\times 64$.

5.1 Reconstruction Quality

Throughout all experiments, our method has shown consistently superior reconstruction quality compared to AEVB, $\beta$-VAE and WAE. Fig. 5 offers a graphical comparison across the reconstructed samples given by different methods for the sine wave and CelebA datasets. For the sine wave dataset, our WiSE-ALE algorithms achieves almost perfect reconstruction, whereas AEVB and $\beta$-VAE often struggle with low-frequency signals and have difficulty predicting the amplitude correctly. For the CelebA dataset, our WiSE-ALE manages to predict much sharper human faces, whereas the AEVB predictions are often blurry and personal characteristics are often ignored. WAE reaches a similar level of reconstruction quality to ours in some images, but it sometimes struggles with discovering the right elevation and azimuth angles, as shown in the second to the right column in Fig. 5(b).

5.2 Properties of the Learnt Representation Space

We understand that a good latent representation should not only reconstruct well, but also preserve some preferable qualities, such as smoothness, compactness and possibly meaningful interpretation of the original data. Fig. 1 indicates that our WiSE-ALE automatically learns a latent representation that finds a good tradeoff between minimizing the information loss and maintaining a smooth and compact aggregate posterior distribution. Furthermore, as shown in Fig. 6, we compare the ELBO values given by AEVB, $\beta$-VAE and our WiSE-ALE over training for the Sine dataset. Our WiSE-ALE manages to report the highest ELBO with a significantly lower reconstruction error and a fairly good performance in the KL divergence loss. This indicates that our WiSE-ALE is able to learn an overall good quality representation that is closest to the true latent distribution which gives rise to the data observation.

6 Conclusion and Future Work

In this paper, we derive a variational lower bound to the data log likelihood, which allows us to impose a prior constraint on the bulk statistics of the aggregate posterior distribution for the entire dataset. Using an analytic approximation to this lower bound as our learning objective, we propose WiSE-ALE algorithm. We have demonstrated its ability to achieve excellent reconstruction quality, as well as forming a smooth, compact and meaningful latent representation.

In the future, we are planning to analyse the error introduced in our approximation to our proposed lower bound. Further, we would like to investigate the potential to use the evaluation of the reconstruction likelihood term given by Eq. 15 as our learning objective, which will keep the lower bound property of our proposal and guarantee that our proposed posterior approaches the true posterior through optimisation.