LaDDer: Latent Data Distribution Modelling with a Generative Prior

In this paper, we show that the performance of a learnt generative model is closely related to the model's ability to accurately represent the inferred **latent data distribution**, i.e. its topology and structural properties. We propose LaDDer to achieve accurate modelling of the latent data distribution in a variational autoencoder framework and to facilitate better representation learning. The central idea of LaDDer is a meta-embedding concept, which uses multiple VAE models to learn an embedding of the embeddings, forming a ladder of encodings. We use a non-parametric mixture as the hyper prior for the innermost VAE and learn all the parameters in a unified variational framework. From extensive experiments, we show that our LaDDer model is able to accurately estimate complex latent distribution and results in improvement in the representation quality. We also propose a novel latent space interpolation method that utilises the derived data distribution. The code and demos are available at <https://github.com/lin-shuyu/ladder-latent-data-distribution-modelling>.

Introduction {#sec:intro}
============

Variational autoencoders (VAEs) @VAE-1 [@VAE-2] are probabilistic latent variable models that aim to learn rich representations from large amounts of data in an unsupervised manner. A trained VAE consists of a generative decoder that generates a data sample from a latent code and a variational encoder that maps a data sample to an approximate posterior distribution over latent variables. Amongst different types of generative models including GANs @GAN, flow models @NICE [@realNVP] and autoregressive models @Nade, VAEs have been favoured for their training stability, strong theoretical grounding and ability to learn well-structured latent representations. In general, the quality of a learned generative model is dependent on two factors: 1) the quality of the inferred data distribution, which defines the topology and structural properties of the **latent space**; 2) the ability to generate good quality samples in the **data space**.

Thanks to all the preferable properties, VAEs have been widely applied in many computer vision applications, including image synthesis @image-synthesis-eg1 [@image-synthesis-eg2; @image-synthesis-eg3], human motion modelling @human-motion-modelling1 [@human-motion-modelling2] and 3D reconstruction @3D-recons-eg2 [@codeslam; @czarnowski2020deepfactors]. However, VAEs still struggle when modelling complex data such as images. Specifically, compared to other models, the generated images of VAEs can be somewhat blurry, lack of fine details and have limited diversity. An entire research field is dedicated to improving VAE learning quality. In this paper, we focus on a key component of VAE learning objective that involves restricting the learned latent space representation to a chosen prior probability. As we cannot observe the true latent distribution from which the data were generated, it is impossible to choose a perfect prior distribution a priori. Therefore, many works, which adopt a simple and inflexible prior in the form of a unit Gaussian distribution, often suffer from an over-regularised latent representation, as the model tries to get the encoder to shoehorn the data to fit the simple prior, sacrificing the quality of generated images. This suggests that using a more flexible prior distribution, or treating the prior as a parameterised model, can reduce or eliminate the over-regularisation issue and lead to improvement in the representation quality. Following these insights, we propose LaDDer, a method that allows us to accurately model the prior distribution in a VAE framework. The central idea of our method is a meta-embedding concept which, in short, derives a latent embedding of a latent embedding. Specifically as shown in Fig. \[fig:system\_overview\], our approach consists of multiple VAE models each acting on its predecessor's latent representation and forming a ladder of encodings. We use a non-parametric mixture as the hyper prior for the innermost VAE. The hyper prior together with the intermediate encoders forms a **generative prior** for the outermost VAE. We learn the parameters of all the VAE networks, along with the non-parametric mixture, in a unified variational inference framework. From extensive experiments, we show that our LaDDer model is able to accurately estimate complex **latent data distribution** and results in improvement in the representation quality. We also propose a **novel latent space interpolation** method and show how to best utilise the derived latent distribution in further tasks.

![**LaDDer's key concept**. LaDDer adopts a generative prior, which consists of a mixture hyper prior and a series of VAEs each acting on its predecessor's latent encodings.<span data-label="fig:system_overview"></span>](images/overview-sketch.pdf){width="90.00000%"}

Background {#sec:background}
==========

Here we introduce VAE models and explain why the prior distribution is crucial in producing good learning outcomes. Given a dataset of N observations $\mathcal{D}_N = \{\vx_1, \cdots, \vx_N\}$, VAEs assume all data samples $\vx_i$ are generated from a low-dimensional latent space $\rvz$ under a latent variable model $p_{\theta}(\rvx, \rvz) = p_{\theta}(\rvx | \rvz)p(\rvz)$, where $p_{\theta}(\rvx | \rvz)$ denotes the generative model (decoder) parameterised by $\theta$. VAEs learn the model parameters $\theta$ by maximising the marginal log likelihood for all data points in $\mathcal{D}_N$, i.e. $\operatorname*{arg\,max}_{\theta} {\mathbb{E}}_{p_{\mathcal{D}}(\rvx)} [\log p_{\theta}(\rvx)]$, where $p_{\mathcal{D}}(\rvx)$ is the empirical data distribution and $p_{\theta}(\rvx)=\int p_{\theta}(\rvx|\rvz)p(\rvz) \textrm{d}\rvz$. However, directly evaluating the marginal log likelihood is often not feasible, as integration over the network $p_{\theta}(\rvx|\rvz)$ is not trivial. To obtain an analytical learning objective, VAEs @VAE-1 [@VAE-2] use variational inference and derive an evidence lower bound (ELBO) $\mathcal{L}(\rvx; \theta, \phi)$ to the marginal log likelihood, i.e. $
    {\mathbb{E}}_{p_{\mathcal{D}}(\rvx)}
    [\log p_{\theta}(\rvx)]  
    \,\geq \, \mathcal{L}(\rvx; \theta, \phi) \label{eq:VAE-ELBO-obj1}
$, as shown in Eq (\[eq:VAE-ELBO-obj2\]): $$\begin{aligned}
    \mathcal{L}(\rvx; \theta, \phi) & \delequal \,
    {\mathbb{E}}_{p_{\mathcal{D}}(\rvx)}\big[
    {\mathbb{E}}_{q_{\phi}(\rvz|\rvx)}[\log p_{\theta}(\rvx|\rvz)]
    \;-\;
    {D_{\mathrm{KL}}}[\,q_{\phi}(\rvz|\rvx) \Vert p(\rvz)]
    \big] \label{eq:VAE-ELBO-obj2} \\
        \delequal  \, {\mathbb{E}}_{p_{\mathcal{D}}(\rvx)}&\big[ 
    \underbrace{{\mathbb{E}}_{q_{\phi}(\rvz|\rvx)}[\log p_{\theta}(\rvx|\rvz)]}
    _\text{\small \textcircled{1} \textrm{\scriptsize{reconstruction likelihood}}}
    \;-\;
    \underbrace{{\mathbb{E}}_{q_{\phi}(\rvz|\rvx)}[\log q_{\phi}(\rvz|\rvx)]}
    _\text{\small \textcircled{2} \textrm{\scriptsize{posterior entropy}}}
    \;+
    \underbrace{{\mathbb{E}}_{q_{\phi}(\rvz|\rvx)}[\log p(\rvz)]}
    _\text{\small \textcircled{3} \textrm{\scriptsize{cross-entropy wrt prior}}} \big].
    \label{eq:VAE-ELBO-obj3} \end{aligned}$$ By breaking down the KL divergence term in Eq (\[eq:VAE-ELBO-obj2\]) using the definition of KL, we obtain an ELBO expression in Eq (\[eq:VAE-ELBO-obj3\]). Here we see that the ELBO contains three terms: a reconstruction likelihood that encourages good reconstruction through the auto-encoding process, a negative posterior entropy that favours $q_{\phi}(\rvz|\rvx)$ with large variances and finally a cross-entropy between the posterior $q_{\phi}(\rvz|\rvx)$ and the prior $p(\rvz)$ that regularises the posteriors to comply with the target prior distribution. We can re-arrange into a cross-entropy between the aggregate posterior $q_{\phi}(\rvz)$ and the prior $p(\rvz)$, as shown below: $$\begin{gathered}
 \label{eq:rearrange-cross-entropy}
    {\mathbb{E}}_{p_{\mathcal{D}}(\rvx)}{\mathbb{E}}_{q_{\phi}(\rvz|\rvx)}[\log p(\rvz)]
    = {\mathbb{E}}_{q_{\phi}(\rvz)} [\log p(\rvz)], \;\;
    \textrm{where} \;\; 
    q_{\phi}(\rvz) 
    =
    {\mathbb{E}}_{p_{\mathcal{D}}(\rvx)} [q_{\phi}(\rvz|\rvx)]. \end{gathered}$$ This re-arrangement reveals that term in the ELBO loss encourages the prior $p(\rvz)$ and the inferred latent data distribution $q_{\phi}(\rvz)$ to match each other. If an overly limiting prior is used, then the learnt data distribution will diverge from the true data distribution. Details of the re-arrangement are given in Supplementary Materials (SM) A1.

Related work
------------

Many approaches have been proposed to improve VAE's modelling ability for rich, high-dimensional data. For example, @expressive-post-1 [@expressive-post-2; @expressive-post-3; @expressive-post-4] propose to use more flexible sample posteriors $q_{\phi}(\rvz|\vx)$ to increase the encoder's expressive power and improve the model's performance in all three loss terms of Eq (\[eq:VAE-ELBO-obj3\]). @IWAE [@factorVAE; @beta-VAE] focus on deriving alternative learning objectives in order to produce representations with preferred qualities, such as disentanglement.

We take a different approach from these lines of research. In this work, we focus on modelling the prior distribution $p(\rvz)$ accurately, which then results in improved learning performance and facilitates further tasks, such as latent space interpolation. Several works have considered using more flexible distributions than a unit Gaussian as priors, such as a stick-breaking prior @SBP-prior, a Chinese Restaurant Process prior @CRP-prior or a Gaussian mixture (GM) prior @deep-clustering [@VampPrior; @SVAE; @alpha-VAE]. Those methods often have limited performance for complex datasets where high-dimensional latent space is needed to facilitate the generative modelling. A closely related work is @hierarchical-prior, where a generative model is used to parameterise the prior. Our method differs from @hierarchical-prior, as we realise that the introduction of a generative prior alone does not guarantee accurate modelling of the latent data distribution. Thus, we introduce a GM hyper prior to complete the modelling task. In addition, we optimise all the model parameters in a coherent lower bound objective, whereas @hierarchical-prior has to set up a constrained optimisation to replace the ELBO objective in order to learn the model parameters.

Another line of research in VAEs is devoted to improving the generation quality, leading to the impressive image quality shown in VQ-VAE2 @VQ-VAE2. A major difference between VQ-VAE2 and our work lies in the different approaches taken to model the prior distribution. VQ-VAE2 learns an auto-regressive prior (using a pixelCNN model @PixelCNN) in a post-hoc, 2nd stage (after the autoencoder is trained). Our model trains both autoencoder and prior modules jointly under the same objective. Further, VQ-VAE2 employs a 2D latent representation, where a single feature vector corresponds to a local patch of the generated image. It is extremely difficult to manipulate such a 2D representation to generate globally consistent changes. In contrast, we adopt a global code to represent the whole image. We can easily traverse across the derived data manifold and generate smoothly changing data samples.

Our Method {#sec:proposal}
==========

As shown in Fig. \[fig:system\_overview\], we propose a generative prior, which includes an additional VAE model to project the original data VAE's encodings to an even lower dimensional space and a hyper prior for this prior VAE. We parameterise the hyper prior as a Gaussian mixture model (GMM). This design allows us to accurately estimate the latent data distribution, as the generative prior is sufficiently flexible to fit any arbitrarily complex distribution. At the same time, the optimisation of our LaDDer model can be easily integrated into the VAE ELBO objective, which we will demonstrate in this section. In the end, we also demonstrate how to use the latent data distribution derived from our LaDDer learning to facilitate a latent space interpolation task.

The VAE Unit in Our Generative Prior {#sec:generative-prior}
------------------------------------

We first introduce the VAE unit in our generative prior. Similar to the VAE for data samples introduced in Section \[sec:background\], the prior VAE is also formulated by a latent variable model $p_{\alpha}(\rvz, \rvt) = p_{\alpha}(\rvz|\rvt) p(\rvt)$, which governs the generation of latent encodings $\vz_i$ through 1) a hyper prior $p(\rvt)$ that resides in an even lower dimensional space and 2) an encoding decoder $p_{\alpha}(\rvz|\rvt)$ which is parameterised by a neural network with parameters $\alpha$. To optimise the prior VAE, we can introduce a variational encoder $q_{\beta}(\rvt|\rvz)$ parameterised by $\beta$ and learn both $\alpha$ and $\beta$ by maximising an ELBO objective $\mathcal{L}(\rvz; \alpha, \beta)$ similar to Eq (\[eq:VAE-ELBO-obj3\]) for this prior model, i.e. $ \label{eq:inner-VAE-ELBO}
    {\mathbb{E}}_{q_{\phi}(\rvz|\rvx)} [\log p_{\alpha}(\rvz)]  
    \geq 
    \mathcal{L}(\rvz; \alpha, \beta) 
$, where: $$\begin{aligned}
    \mathcal{L}(\rvz; \alpha, \beta) &\delequal
    {\mathbb{E}}_{q_{\phi}(\rvz|\rvx)}\big[
    {\mathbb{E}}_{q_{\beta}(\rvt|\rvz)}[\log p_{\alpha}(\rvz|\rvt)]
    -{\mathbb{E}}_{q_{\beta}(\rvt|\rvz)}[\log q_{\beta}(\rvt|\rvz)]
    +{\mathbb{E}}_{q_{\beta}(\rvt|\rvz)}[\log p(\rvt)] 
    \big]. 
    \label{eq:inner-VAE-ELBO2}\end{aligned}$$

Notice that $\mathcal{L}(\rvz; \alpha, \beta)$ is a lower bound to the likelihood ${\mathbb{E}}_{q_{\phi}(\rvz|\rvx)} [\log p_{\alpha}(\rvz)]$, which is equivalent to the cross-entropy term in Eq (\[eq:VAE-ELBO-obj3\]). This connection allows us to integrate the learning objective of this prior VAE, i.e. $\mathcal{L}(\rvz; \alpha, \beta)$, into the ELBO for the original data VAE, i.e. $\mathcal{L}(\rvx; \theta, \phi)$, and obtain a new lower bound $\mathcal{L}'(\rvx; \theta, \phi, \alpha, \beta)$ to the data ELBO $\mathcal{L}(\rvx; \theta, \phi)$, i.e. $ 
    \mathcal{L}(\rvx; \theta, \phi)
    \geq \mathcal{L}'(\rvx; \theta, \phi, \alpha, \beta) 
$, where: $$\begin{aligned}
    \mathcal{L}'(\rvx; \theta, \phi, \alpha, \beta)=
    {\mathbb{E}}_{p_{\mathcal{D}}(\rvx)} \big[
    {\mathbb{E}}_{q_{\phi}(\rvz|\rvx)}[\log p_{\theta}(\rvx|\rvz)]
    -{\mathbb{E}}_{q_{\phi}(\rvz|\rvx)}[\log q_{\phi}(\rvz|\rvx)]
    +\mathcal{L}(\rvz; \alpha, \beta)
    \big].
    \label{eq:prior-lower-bound-to-ELBO}\end{aligned}$$

Variational Gaussian Mixture Model for the Hyper Prior {#sec:hyper-prior}
------------------------------------------------------

The prior VAE unit defined in Section \[sec:generative-prior\] introduces a hyper prior $p(\rvt)$. From Eq (\[eq:rearrange-cross-entropy\]), we know that the optimal hyper prior should be matched to the aggregate hyper posterior, i.e. $p(\rvt) \approx q_{\beta, \phi}(\rvt)$. To facilitate the matching, we parameterise $p(\rvt)$ with a Gaussian mixture model (GMM) of M components (M $\ll$ N): $
    p(\rvt)
    = 
    \sum_{m=1}^{M}w_m \mathcal{N}(\rvt; \,\mu_m, \,\Sigma_m),
    $ where $w_m$ is the weight for each Gaussian mixture ($w_m > 0$ and $\sum_m w_m =1$) and $\mu_m$ and $\Sigma_m$ are the mean and covariance matrix for the $m$-th Gaussian mixture. To fit $p(\rvt)$ to $q_{\beta, \phi}(\rvt)$, we resort to variational inference techniques introduced in @VGMM. Here we give a high-level sketch of the algorithm and we refer interested readers to @VGMM for further mathematical details. Firstly, we define the generative model of the GMM by introducing the following distributions to the GMM parameters $(\vw, \mu, \Sigma)$ and all N prior encoding samples $\vt_i$ (the encodings of $\vz_i$): $ \label{eq:generative-distribution1}
    w_m \sim \mathcal{B}eta(1, \alpha_0), \;\; 
    \mu_m \sim \mathcal{N}(0, \mathcal{I}), \;\;
    \Sigma_m \sim \mathcal{W}(d_t, \mathcal{I}),\;\;
    k_i \sim \textrm{Cat}(\vw),\;\;
    \vt_i \sim \mathcal{N}(\mu_{k_i}, \Sigma_{k_i}),
$ where $k_i$ indicates the choice of mixture components for $i$-th sample $\vt_i$. Secondly, we introduce the following variational distributions under the mean-field assumption for all the model variables $W=(\vw, \mu, \Sigma, k_i)$: $ \label{eq:variational-distribution}
    w_m \sim \mathcal{B}eta(\gamma_{m,1}, \gamma_{m,2}), \;\; 
    \mu_m \sim \mathcal{N}(\vv_m, \mathcal{I}), \;\;
    \Sigma_m \sim \mathcal{W}(a_m, \mathcal{B}_m), \;\;
    k_i \sim \textrm{Discrete}(\vr_i),
$ where $\xi = (\gamma_{m,1}, \gamma_{m,2}, \vv_m, a_m, \mathcal{B}_m, \vr_i)$ for $m=1,\cdots,M$ and $i=1,\cdots,N$ denotes variational variables and needs to be optimised. Thirdly, a variational bound on the log likelihood of $p(\rvt)$ is introduced as the learning objective to optimise $\xi$, as shown below: $$\begin{aligned}
 \label{eq:VGMM-bound}
    \log p(\rvt | \alpha_0, M) 
    \geq
    \mathcal{L}(\rvt; \xi)
    =
    {\mathbb{E}}_{q_{\xi}(W)} [\log p(W, \rvt| \alpha_0, M)]
    -{\mathbb{E}}_{q_{\xi}(W)} [\log q_{\xi}(W)],\end{aligned}$$ where $p(W, \rvt| \alpha_0, M)$ denotes the product of all the generative distributions introduced in step 1 and $q_{\xi}(W)$ denotes the product of all the variational distributions introduced in step 2.

Now substituting $\mathcal{L}(\rvt; \xi)$ in Eq (\[eq:VGMM-bound\]) into Eq (\[eq:inner-VAE-ELBO2\]) to replace $\log p(\rvt)$, we obtain a lower bound $\mathcal{L}'(\rvz; \alpha, \beta, \xi)$ to the prior ELBO $\mathcal{L}(\rvz; \alpha, \beta)$, i.e. $\mathcal{L}(\rvz; \alpha, \beta) \geq \mathcal{L}'(\rvz; \alpha, \beta, \xi)$, as $$\begin{aligned}
 \label{eq:LB-prior-ELBO}
    \mathcal{L}'(\rvz; \alpha, \beta, \xi) 
    =
    {\mathbb{E}}_{q_{\phi}(\rvz|\rvx)}\big[
    {\mathbb{E}}_{q_{\beta}(\rvt|\rvz)}[\log p_{\alpha}(\rvz|\rvt)]
    -{\mathbb{E}}_{q_{\beta}(\rvt|\rvz)}[\log q_{\beta}(\rvt|\rvz)]
    + {\mathbb{E}}_{q_{\beta}(\rvt|\rvz)} [\mathcal{L}(\rvt; \xi)]
    \big]. \end{aligned}$$ Now substituting $\mathcal{L}'(\rvz; \alpha, \beta, \xi)$ into the lower bound $\mathcal{L}'(\rvx; \theta, \phi, \alpha, \beta)$ in Eq (\[eq:prior-lower-bound-to-ELBO\]) to replace $\mathcal{L}(\rvz; \alpha, \beta)$, we obtain a final lower bound $\mathcal{L}''(\rvx; \Theta)$ to the data likelihood ${\mathbb{E}}_{p_{\mathcal{D}}(\rvx)} [\log p_{\theta}(\rvx)]$, i.e. $ 
    {\mathbb{E}}_{p_{\mathcal{D}}(\rvx)} [\log p_{\theta}(\rvx)]  
    \geq 
    \mathcal{L}''(\rvx; \Theta)
$, where $\Theta = [\theta, \phi, \alpha, \beta, \xi]$ denotes all the parameters in the data VAE, prior VAE and GMM hyper prior, and $$\begin{aligned}
\label{final-LB}
    \mathcal{L}''(\rvx; \Theta)=
    {\mathbb{E}}_{p_{\mathcal{D}}(\rvx)} \big[
    {\mathbb{E}}_{q_{\phi}(\rvz|\rvx)}[\log p_{\theta}(\rvx|\rvz)]
    -{\mathbb{E}}_{q_{\phi}(\rvz|\rvx)}[\log q_{\phi}(\rvz|\rvx)]
    +\mathcal{L}'(\rvz; \alpha, \beta, \xi)
    \big].\end{aligned}$$

Block Coordinate Ascent to Optimise Model Parameters {#sec:BCA-optimisation}
----------------------------------------------------

With the new lower bound in Eq (\[final-LB\]), we are ready to introduce our block coordinate ascent algorithm which optimises all the model parameters $\Theta = [\theta, \phi, \alpha, \beta, \xi]$. Notice that the GMM hyper prior's parameter $\xi$ are only involved in the term $\mathcal{L}(\rvt; \xi)$ from Eq (\[eq:VGMM-bound\]). We can update $\xi$ by maximising $\mathcal{L}(\rvt; \xi)$ alone. With the newly updated $\xi$, we can then update the prior VAE's parameter $\alpha, \beta$, which only present in the prior VAE's ELBO $\mathcal{L}'(\rvz; \alpha, \beta | \xi)$ from Eq (\[eq:LB-prior-ELBO\]). Finally, with $\xi, \alpha, \beta$ updated, we can now update the data VAE's parameters $\theta, \phi$ by maximising $\mathcal{L}''(\rvx; \theta, \phi | \alpha, \beta, \xi)$ from Eq (\[final-LB\]) with $\alpha, \beta, \xi$ being fixed. An algorithmic illustration for this optimisation procedure is given in SM B. Notice that we initialise the GMM hyper prior as an uninformative unit Gaussian to begin the optimisation.

Shortest Likely Path to Traverse the Data Manifold {#sec:SLP}
--------------------------------------------------

In Section \[sec:generative-prior\]-\[sec:BCA-optimisation\], we introduce a generative prior and an optimisation scheme to facilitate accurate modelling of the latent data distribution. Now we demonstrate how to use the derived latent distribution in a latent space traversal task, where a path along the learnt data manifold needs to be inferred to interpolate between two data samples. Such interpolation has been used commonly in previous works @factorVAE [@Glow; @UnsupervisedRL; @AAR] to illustrate the smoothness of the learnt data manifold. Traditionally, the traversal is done by a linear interpolation between the encodings of a pair of query images, which we refer as shortest path (SP) interpolation. This method ignores the nonlinear topology of the data manifold in the latent space and often results in poor interpolated images, which are generated faraway from the data manifold.

To achieve better interpolation results, we propose to formulate the traversal task as an optimisation problem which aims to find an optimal path that is shortest while remaining close to the data manifold. We refer to our interpolation scheme as shortest likelihood path (SLP) interpolation and design an optimisation objective as follows: $$\begin{aligned}
 \label{eq:SLP-obj}
    \mathcal{O}_{\textrm{SLP}} 
    \;\;=\;\;
    \underbrace{\textrm{L}_{\textrm{path}}(\vs)}
    _\text{\textrm{\scriptsize{path length}}}
    \;\;+\;\;
    \underbrace{\textrm{std}_{\textrm{step length}}(\vs)}
    _\text{\textrm{\scriptsize{equal step length}}}
    \;\;-\;\;
    \underbrace{\log p(\rvt=\vs),}
    _\text{\textrm{\scriptsize{path likelihood}}}\end{aligned}$$ where $\vs=[s_1, \cdots, s_J]^T$ denotes a list of J steps along the traversal path and each $s_j$ denotes a single encoding that can be projected back to the data space as an interpolated image. As can be seen, our $\mathcal{O}_{\textrm{SLP}}$ objective contains three terms: 1) $\textrm{L}_{\textrm{path}}(\vs)$ minimises the current path length, 2) $\textrm{std}_{\textrm{step length}}(\vs)$ requires steps to be evenly distributed along the path and 3) $\log p(\rvt=\vs)$ ensures all the interpolated encodings to be generated from the inferred data distribution $p(\rvt)$. To find an optimal path $\vs$, we minimise $\mathcal{O}_{\textrm{SLP}}$ wrt $\vs$ by a standard optimisation scheme, such as AdamOptimiser @AdamOpt that is used in our experiments. Notice that our SLP interpolation can be generally applied to any representation learning algorithm, as long as the likelihood of the latent data distribution $p(\rvz)$ can be easily evaluated.

Experiments and Results {#sec:result}
=======================

We carry out extensive experiments on MNIST @MNIST, Fashion MNIST @fashion-mnist and CelebA @CelebA-Dataset datasets to evaluate our LaDDer model. For all datasets, images are treated as real-valued data and we use a Laplace distribution to model the decoder $p_{\theta}(\rvx|\rvz)$. The corresponding reconstruction likelihood (term in Eq (\[eq:VAE-ELBO-obj3\])) is derived in SM C, following @alpha-VAE. We compare our method to 4 other approaches, including the original VAE with a normal prior @VAE-1 [@VAE-2], VAE with a GMM prior @deep-clustering [@SVAE; @alpha-VAE], VAE with a hierarchical prior @hierarchical-prior where a generative prior with a normal hyper prior is used, and VampPrior @VampPrior where the prior is modelled as an average encoding of a set of inferred pseudo-inputs. More results, details of data pre-processing and model architectures are given in SM D-F.

Our Method Better Estimates the Latent Data Distribution
--------------------------------------------------------

We first show that our LaDDer model can achieve a better modelling of the latent data distribution, i.e. more accurate estimation of the aggregate posterior $q_{\phi}(\rvz)$. Fig. \[fig:mnist-2d-latent-space\] visualises the aggregate posterior $q_{\phi}(\rvz)$ and the optimised priors from 5 different approaches, where VAE models are trained for 2D latent space on MNIST dataset. The true data distribution shown in Fig. \[fig:mnist-2d-latent-space\]a contains complex low density regions, which correspond to natural boundaries between different object classes. An inflexible prior model, such as a unit Gaussian (Fig. \[fig:mnist-2d-latent-space\]b), over-represents such low density regions. The hierarchical prior (Fig. \[fig:mnist-2d-latent-space\]c) also fails to model $q_{\phi}(\rvz)$ accurately, because the prior VAE is overly regularised by its inflexible hyper prior. The other three methods (Fig. \[fig:mnist-2d-latent-space\]d-f), where the prior $p(\rvz)$ is parameterised with mixture models, obtain better fitting. The difference across these three methods lies in the number of mixture components. Both the GMM prior and the VampPrior have a finite number of mixture components, whereas our LaDDer prior contains an infinite number of mixture components, as our prior is estimated by integrating over all $t$-values: $p_{\alpha}(\rvz)=\int p_{\alpha}(\rvz|\rvt)p(\rvt)\,\textrm{d}\rvt$. As a result, our method is able to smoothly fit any arbitrarily complex $q_{\phi}(\rvz)$ with no artificial boundaries as the ones introduced in VampPrior.

\[t\] ![image](images/2d-mnist-latent-space_retrained.pdf){width="95.00000%"}

\[t\] ![image](images/celeba-generation-comparison.pdf){width="\textwidth"}

Having a prior that can better fit $q_{\phi}(\rvz)$ also leads to improved generation quality. Fig. \[fig:generation-quality-comparison\] gives examples of both reconstructed and generated samples from the 5 VAE models trained under the CelebA dataset. All 5 methods can produce good reconstruction, but the generation quality varies significantly. This indicates that the bad generation quality commonly reported in VAE models does not merely result from the autoencoder architecture, but is also a result of the prior $p(\rvz)$ failing to represent the topology of the learned latent data distribution. Generating from the prior that does not represent the true data distribution unsurprisingly produce unrealistic samples. Our proposed generative prior addresses this issue and produces generation quality that is almost as good as the model's reconstruction quality. The improved generation quality is supported by the FID score @FID1, which is a quantitative measure to evaluate visual quality of image samples, shown in Table \[tab:celeba-FID\].

<span>|B|A|A|A|A|A|</span> & Normal & VampPrior & GMM & Hierarchical & Ours\
Generation & 250.3 $\pm$ 1.5 & 182.9 $\pm$ 1.2 & 140.4 $\pm$ 0.9 & 143.8 $\pm$ 1.0 & **132.7** $\pm$ 0.8\
Reconstruction & 100.3 $\pm$ 0.6 & 95.5 $\pm$ 0.7 & 102.3 $\pm$ 0.5 & 99.7 $\pm$ 0.6 & **95.3** $\pm$ 0.5\
Difference & 150.0 $\pm$ 2.1 & 87.4 $\pm$ 1.9 & 38.1 $\pm$ 1.4 & 44.1 $\pm$ 1.6 & **37.4** $\pm$ 1.3\

\[tab:celeba-FID\]

We also evaluate the cross-entropy $E_{q_{\phi}(\rvz)}[\log p(\rvz)]$ ( in Eq (\[eq:VAE-ELBO-obj3\])) between the aggregate posterior $q_{\phi}(\rvz)$ and the prior $p(\rvz)$ to measure the level of matching between the two distributions. Higher $E_{q_{\phi}(\rvz)}[\log p(\rvz)]$ indicates a better fit of the prior $p(\rvz)$ to the learnt data distribution $q_{\phi}(\rvz)$. As shown in Table \[tab:cross-entropy\], GMM priors can achieve very good fitting for relatively simple datasets, such as MNIST. However, when the dataset becomes more complex and the required latent dimension increases, such as fashion-MNIST and CelebA datasets, our generative prior, which is more flexible to fit complex distributions, obtains the highest cross-entropy and achieve the best fit.

<span>|D|C|C|C|C|C|</span> Dataset & Normal & VampPrior & GMM & Hierarchical & Ours\
MNIST & -19.0 $\pm$ 0.1 & -16.2 $\pm$ 0.1 & **-8.4** $\pm$ 0.9 & -17.5 $\pm$ 0.1 & -12.3 $\pm$ 0.1\
fashion-MNIST & -52.5 $\pm$ 0.6 & -60.8 $\pm$ 1.0 & -45.2 $\pm$ 0.8 & -63.8 $\pm$ 0.9 & **-32.8** $\pm$ 0.7\
CelebA & -351.4 $\pm$ 4.3 & -380.7 $\pm$ 3.4 & -290.6 $\pm$ 54.8 & -358.0 $\pm$ 7.3 & **-71.9** $\pm$ 2.1\

\[tab:cross-entropy\]

Overall Generative Modelling Performance Improves
-------------------------------------------------

A more flexible prior $p(\rvz)$ also leads to better generative modelling performance, in terms of better reconstruction quality and higher ELBO objective. This effect is clearly shown in Table \[tab:ELBO-recons-error\], where our method produces the highest ELBO value and the lowest reconstruction error across all datasets. All models are trained under the same conditions and the evaluation is repeated for 5 times to estimate the variance of the results.

<span>|C|E|B|B|B|B|B|</span> Metric & Dataset & Normal & VampPrior & GMM & Hierarchical & Ours\
& MNIST & 1528.6 $\pm$ 21.3 & 1476.3 $\pm$ 19.4 & 1558.2 $\pm$ 24.4 & 1534.9 $\pm$ 19.8 & **1562.5** $\pm$ 25.1\
& fashion & 1078.4 $\pm$ 23.4 & 1069.6 $\pm$ 22.1 & 1133.3 $\pm$ 23.9 & 1152.9 $\pm$ 24.9 & **1174.0** $\pm$ 29.7\
& CelebA & 55564.8 $\pm$ 1751.3 & 56463.4 $\pm$ 1683.9 & 54162.5 $\pm$ 1716.5 & 56176.2 $\pm$ 1763.9 & **58596.1** $\pm$ 1815.3\
Per pixel & MNIST & 2.3 $\pm$ 0.06 & 2.5 $\pm$ 0.06 & 2.3 $\pm$ 0.05 & 2.4 $\pm$ 0.06 & **2.2** $\pm$ 0.06\
recons. error & fashion & 4.3 $\pm$ 0.12 & 4.1 $\pm$ 0.11 & 4.0 $\pm$ 0.11 & 3.7 $\pm$ 0.11 & **3.6** $\pm$ 0.11\
($\times 0.01$) & CelebA & 5.8 $\pm$ 0.21 & 5.9 $\pm$ 0.25 & 5.9 $\pm$ 0.20 & 5.7 $\pm$ 0.17 & **5.5** $\pm$ 0.18\

\[tab:ELBO-recons-error\]

Furthermore, our LaDDer model produces an embedding scheme that better preserves the semantics of the data. To see this, we visualise 30k latent encodings of the MNIST dataset in Fig. \[fig:mnist-clustering\], where a VAE of 2D latent space is trained with a GMM prior, VampPrior and our generative prior respectively. All encodings are coloured by their class labels. Our method clearly gives better clustering results, where encodings of different classes are better separated. Furthermore, the mixture components estimated in our method are better aligned with class labels, whereas the mixtures from VampPrior and GMM prior have significant overlaps and do not have a consistent correspondence to specific class labels.

![**The mixtures in our hyper prior naturally aligns with MNIST class labels**. We visualise encodings coloured by the digit labels for (a) VampPrior, (b) GMM prior and (c) our generative prior. Our method produces better clustering and the mixture components inferred in our method are naturally aligned with different classes.<span data-label="fig:mnist-clustering"></span>](images/clusters.pdf){width="90.00000%"}

Our Shortest Likely Path Traversal Gives Better Interpolation
-------------------------------------------------------------

In Section \[sec:SLP\], we formulate the latent space interpolation between a pair of images as a shortest likely path (SLP) optimisation task, which favours the paths that go through high density regions of the inferred latent data distribution $p(\rvz)$ or $p(\rvt)$. Here we demonstrate our SLP method informed by the learnt latent distribution outperforms the conventional linear shortest path (SP) interpolation.

![**Latent space interpolation for MNIST**. (a) The 3 pairs of MNIST images for the interpolation task. (b-c) The traversed paths produced by shortest likely path (SLP) and shortest path (SP) are visualised in our very low-dimensional $t$-space. (d-f) The interpolated images for both interpolation methods using the data distribution estimated by our method, a GMM prior and a normal prior.<span data-label="fig:mnist-traversal"></span>](images/mnist-2d-traversal-ours.pdf){width="95.00000%"}

Fig. \[fig:mnist-traversal\] illustrates the interpolation between 3 pairs of MNIST images through the data manifolds estimated by our method, a GMM prior and a normal prior. Notice that our LaDDer model can produce the same quality samples from a much lower dimension latent space compared to the VAE models with the GMM and the normal prior (2D vs 16D). The super low latent dimension allows us to visualise the learnt data distribution $p(\rvt)$ and the traversed path in Fig. \[fig:mnist-traversal\]b-c. Our SLP interpolation results in paths that only step on regions with high likelihood of $p(\rvt)$, whereas SP interpolation ignores the topology of $p(\rvt)$ and often lands on the low density regions which do not correspond to realistic data samples, hence producing unrealistic samples (bottom rows in Fig. \[fig:mnist-traversal\]d-f). Notice that neither SLP nor SP gives a good interpolation for the normal prior model. This is because the normal prior poorly represents the latent data distribution and even if an encoding has a high likelihood wrt the normal prior, it does not correspond to a realistic data sample. This reinforces the importance of obtaining an accurate modelling of the inferred data distribution in being able to utilise it for later tasks.

Fig. \[fig:celeba-traversal\] illustrates the interpolation between 2 pairs of CelebA images over the data manifold produced by our method, where the SLP optimisation takes place in the 32D latent space of the prior VAE. Here, we plot the different objectives in our SLP optimisation (blue line) for each example in (c-f). The overall objective ($\mathcal{O}_{\textrm{SLP}}$ in Eq (\[eq:SLP-obj\])) smoothly converges in all examples. It is clear that our SLP solution trade-offs path length and allows different step lengths for obtaining high likelihood over the traversed path. As a result, our SLP method only produces realistic images which smoothly transform into the target face, whereas the SP method produces unrealistic faces along the traversed path.

![**Latent space interpolation for CelebA**. (a) The pair of images to be interpolated. (b) The sequence of interpolated images along the data manifold estimated by our LaDDer model, using two different interpolation methods. (c-f) Optimisation of different objectives in the our SLP objective. More examples are given in Supplementary Materials F4.<span data-label="fig:celeba-traversal"></span>](images/celeba-traversal-2-examples.pdf){width="90.00000%"}

Conclusion and Future Work {#sec:conclusion}
==========================

In this paper, we recognise the importance of adopting a sufficiently flexible prior in a VAE model to facilitate accurate modelling of the inferred latent data distribution. We propose LaDDer, which consists of multiple VAE models each acting on its predecessor's latent representation and a non-parametric mixture as the hyper prior for the innermost VAE. From extensive experiments, we show that our method is able to accurately model latent data distribution of complex data. We also demonstrate how to use the derived latent distribution to facilitate further tasks, such as producing better interpolation along the derived data manifold. We believe that LaDDer can be helpful in estimating the data distribution for many challenging datasets and the derived data distribution can be useful for a wide range of applications. We will continue to explore along these directions.

Acknowledgements {#acknowledgements .unnumbered}
================

We thank Stephen Roberts and Niki Trigoni for engaging discussions and supporting the research behind this project. Shuyu Lin is supported by the EPSRC Centre for Doctoral Training, EP/L015897/1, and the China Scholarship Council. Ronald Clark is supported by an Imperial College Research Fellowship. We thank the reviewers for helpful comments.

{'source': 'arXiv_out/2009'}