Diffusion Variational Autoencoders

1 Introduction

A large part of unsupervised learning is devoted to the extraction of meaningful latent factors that explain a certain data set. The terminology around Variational Autoencoders suggests that they are a good tool for this task.

A Variational Autoencoder (Kingma & Welling, 2014; Rezende et al., 2014) consists of a “latent space” $Z$

(often just Euclidean space), a probability measure

P_{Z}

Z

, an encoder from the data space

X

Z

and a decoder from

Z

X

. The term “latent space” suggests that an element in

z

captures semantic information about its decoded data point

x = D e c (z)

. But is this interpretation actually warranted? Especially considering that in practice, a Variational Autoencoder is often trained without any knowledge about an underlying generative process. The only available information is the dataset itself. Nonetheless, one may train a Variational Autoencoder with any “latent space”

Z

and purely based on this terminology, one could be tempted to think of elements in

Z

as latent variables and factors explaining the data.

Figure 1: MNIST trained on a sphere and on a torus

Certainly, there is a heuristic argument that gives a partial justification of this interpretation. The starting point and guiding principle is that of Occam’s razor, that such latent factors might arise from the construction of simple, low-complexity models that explain the data

(Portegies, 2018). Strict formalizations of Occam’s razor exist, through Kolmogorov complexity and inductive inference (Solomonoff, 1964; Schmidhuber, 1997), but are often computationally intractable. A more practical approach leads to the principle of minimum description length (Rissanen, 1978; Hinton & Zemel, 1994) and variational inference (Honkela & Valpola, 2004).

From a different angle, part of the interpretation as latent space is warranted by the loss function of the Variational Autoencoder, which stimulates a continuous dependence between the latent variables and the corresponding data points. Close-by points in data space should also be close-by in latent space. This suggests that a Variational Autoencoder could capture topological and geometrical properties of a data set.

However, a standard Variational Autoencoder (with a Euclidean latent space) is at times structurally incapable of accurately capturing topological properties of a data set. Take for example the case of a spinning object placed on a turntable and being recorded by a camera from a fixed position. The data set for this example is the collection of all frames. The true latent factor is the angle of the turntable. However, the space of angles is topologically and geometrically different from Euclidean space. In an extreme example, if we train a Variational Autoencoder with a one-dimensional latent space on the pictures from the object on the turntable, there will be pictures taken from almost the same angle ending up at completely different parts of the latent space.

This phenomenon has been called manifold mismatch (Davidson et al., 2018; Falorsi et al., 2018). To match the latent space with the data structure, Davidson et al. implemented spheres as latent spaces, whereas Falorsi et al. implemented the special orthogonal group $S O (3)$ .

As further examples of datasets with topologically nontrivial latent factors, we can think of many translations of the same periodic picture, where the translation is the latent variable, or many pictures of the same object which has been rotated arbitrarily. In these cases, there are still clear latent variables, but their topological and geometrical structure is neither that of Euclidean space nor that of a sphere, but rather that of a torus and that of the $S O (3)$ respectively.

To address the problem of manifold mismatch, we developed the Diffusion Variational Autoencoder ( $Δ$ VAE) which allows for an arbitrary manifold as a latent space. Our implementation includes a version of the reparametrization trick, and a fast approximation of the evaluation of the KL divergence in the loss.

We implemented $Δ$ VAEs with latent spaces of $d$ -dimensional spheres, a two-dimensional flat torus, an embedded torus in $R^{3}$ , the special orthogonal group $S O (3)$ and the real projective spaces ${R P}^{d}$ .

We trained a synthetic data set of translated images on a flat torus. Our results show that the $Δ$ VAE can capture the topological properties of the data. We further observed that the success rate of the $Δ$ VAE in capturing global topological properties depends on the weight of higher Fourier components in the image. In data sets with more pronounced lower Fourier components the $Δ$ VAE is more successful in capturing the topology.

Mainly as a proof of concept, we trained on MNIST using the manifolds mentioned above.

2 Related work

Our work originated out of the search for algorithms that find semantically meaningful latent factors of data. The use of VAEs and their extensions to this end has mostly taken place in the context of disentanglement of latent factors (Higgins et al., 2017, 2018; Burgess et al., 2018). Examples of extensions that aim at disentangling latent factors are the $β$ -VAE (Higgins et al., 2017), the factor-VAE (Kim & Mnih, 2018), the $β$ -TCVAE (Chen et al., 2018) and the DIP-VAE (Kumar et al., 2018).

However, the examples in the introduction already show that in some situations, the topological structure of the latent space makes it practically impossible to disentangle latent factors. The latent factors are inherently, topologically entangled: in the case of a 3d rotation of an object, one cannot assign globally linearly independent angles of rotation.

Still, it is exactly global topological properties that we feel a VAE has a chance of capturing. What do we mean by this? One instance of ‘capturing’ topological structure is when the encoder and decoder of the VAE provide bijective, continuous maps between data and latent space, also called homeomorphic auto-encoding (Falorsi et al., 2018; de Haan & Falorsi, 2018). This can only be done when the latent space has a particular topological structure, for instance that of a particular manifold.

One of the main challenges when implementing a manifold as a latent space is the design of the reparametrization trick. In (Davidson et al., 2018), a VAE was implemented with a hyperspherical latent space. To our understanding, they implemented a reparametrization function which was discontinuous; see also Section 3.5 below.

If a manifold has the additional structure of a Lie group, this structure allows for a more straightforward implementation of the reparametrization trick (Falorsi et al., 2018). In our work, we do not assume the additional structure of a Lie group, but develop a reparametrization trick that works for general submanifolds of Euclidean space, and therefore by the Whitney (respectively Nash) embedding theorem, for general closed (Riemannian) manifolds.

The method that we use has similarities with the approach of Hamiltonian Variational Inference (Salimans et al., 2015). Moreover, the implementation of a manifold as a latent space can be seen as enabling a particular, informative, prior distribution. In that sense, our work relates to (Dilokthanakul et al., 2016; Tomczak & Welling, 2017). The prior distribution we implement is very degenerate, in that it is does not assign weight to points outside of the manifold.

There are also other ways to implement approximate Bayesian inference on Riemannian manifolds. For instance, Liu and Zhu adapted the Stein variational gradient method to enable training on a Riemannian manifold

(Liu & Zhu, 2017). However, their proposed method is rather expensive computationally.

The family of approximate posteriors that we implement is a direct generalization of the standard choice for a Euclidean VAE. Indeed, the Gaussian distributions are solutions to the heat equations, i.e. they are transition kernels of Brownian motion. One may want to increase the flexibility of the family of approximate posterior distributions, for instance by applying normalizing flows

(Rezende & Mohamed, 2015; Kingma et al., 2016; Gemici et al., 2016).

3 Methods

3.1 Variational autoencoders

A VAE has generally the following ingredients:

a latent space $Z$ ,
a prior probability distribution
$P_{Z}$ on $Z$ ,
a family of encoder distributions $Q_{Z}^{α}$ on $Z$ , parametrized by $α$ in a parameter space $A$ ,
a family of decoder distributions $P_{X}^{β}$ on data space $X$ , parametrized by $β$ in a parameter space $B$ ; in the usual setup, and in our paper, in fact $B$ corresponds to the data space, and refers to the mean of a Gaussian distribution with identity covariance,
an encoder neural network
$α$ which maps from data space $X$ to the parameter space $A$ ,
a decoder neural network $β$ which maps from latent space $Z$ to parameter space $B$ .

The weights of these neural networks are optimized as to minimize the negated evidence lower bound (ELBO)

The first term is called reconstruction error (RE); up to additive and multiplicative constants it equals the mean squared error (MSE). The second term is called the KL-loss.

In a very common implementation, both latent space $Z$ and data space $X$

are Euclidean, and the families of decoder and encoder distributions are multivariate Gaussian. The encoder and decoder networks then assign to a datapoint or a latent variable a mean and a variance respectively.

When we implement $Z$ as a Riemannian manifold, we need to find an appropriate prior distribution, for which we will choose the normalized Riemannian volume measure, a family of encoder distributions $Q_{Z}^{α}$ , for which we will take transition kernels of Brownian motion, and an encoder network mapping to the correct parameters.

3.2 Brownian motion on a Riemannian manifold

We will briefly discuss Brownian motion on a Riemannian manifold, recommending lecture notes by Hsu (Hsu, 2008) as a more extensive introduction.

In the paper, we always assume that $Z$ is a smooth Riemannian submanifold of Euclidean space, which is closed, i.e. it is compact and has no boundary.

There are many different, equivalent definitions of Brownian motion. We present here the definition that is closest to our eventual approximation and implementation.

Figure 2: Random walk on a (one-dimensional) submanifold $Z$ of $R^{2}$ , with time step $τ = 1$ .

We will construct Brownian motion out of random walks on a manifold. We first fix a small time step $τ > 0$ . We will imagine a particle, jumping from point to point on the manifold after each time step, see also Fig. 2. It will start off at a point $z \in Z$ . We describe the first jump, after which the process just repeats. After time $τ$ , the particle makes a random jump $\sqrt{τ} ϵ_{1}$ from its current position, into the surrounding space, where $ϵ_{1}$ is distributed according to a radially symmetric distribution in $R^{n}$ with identity covariance matrix. The position of the particle after the jump, $z + \sqrt{τ} ϵ_{1}$ , will therefore in general not be on the manifold, so we project the particle back: The particle’s new position will be

z_{1} = P (z + \sqrt{τ} ϵ_{1})

where the closest-point-projection $P : R^{n} \to Z$ assigns to every point $x \in R^{n}$ the point in $Z$ that is closest to $x$ . After another time $τ > 0$ the particle makes a new, independent, jump $ϵ_{2}$ according to the same radially symmetric distribution, and its new position will be $z_{2} = P (P (z + \sqrt{τ} ϵ_{1}) + \sqrt{τ} ϵ_{2})$ . This process just repeats.

Key to this construction, and also to our implementation, is the projection map $P$ . It has nice properties, that follow from general theory of smooth manifolds. In particular, $P (x)$ smoothly depends on $x$ , as long as $x$ is not too far away from $Z$ .

This way, for $τ > 0$ fixed, we have constructed a random walk, a random path on the manifold. We can think of this path as a discretized version of Brownian motion. Let now $τ_{N}$ be a sequence converging to $0$ as $N \to \infty$ . For fixed $N \in N$ , we can construct a random walk with time step $τ_{N}$ , and get a random path $W^{N} : [0, \infty) \to Z$ .

The random paths $W^{N}$ converge as $N \to \infty$ to a random path $W$ (in distribution). This random path $W$ is called Brownian motion. The convergence statement can be made precise by for instance combining powerful, general results by (Jørgensen, 1975)

with standard facts from Riemannian geometry. But, because Riemannian manifolds are locally, i.e. when you zoom in far enough, very similar to Euclidean space, the convergence result essentially comes down to the central limit theorem and its upgraded version, Donsker’s invariance theorem.

In fact, $W$ can be interpreted as a Markov process, and even as a diffusion process. If $A$ is a subset of $Z$ , the probability that the Brownian motion $W (t)$ started at $z$ is in the set $A$ at time $t$ is measured by a probability measure $Q_{Z}^{t, z}$ applied to the set $A$ . We denote the density of this measure with respect to the standard Riemannian volume measure by $q_{Z} (t; z, \cdot)$ . The function $q_{Z}$ is sometimes referred to as the heat kernel.

Let us close this subsection with an alternative description of the function $q_{Z}$ . It is also characterized by the fact that for every function $u_{0} : Z \to R$

, the solution to the partial differential equation

{\begin{matrix} \partial_{t} u = \frac{1}{2} Δ u & on (0, \infty) \times Z u (t = 0, \cdot) = u_{0} & on Z \end{matrix}

is given by

u (t, z) = \int_{Z} u_{0} (y) q_{Z} (t; z, y) d y .

3.3 Riemannian manifold as latent space

A $Δ$ VAE is a VAE with a Riemannian submanifold of Euclidean space as a latent space, and the transition probability measures of Brownian motion $Q_{Z}^{t, z}$

as a parametric family of encoder distributions. We propose the uniform distribution for

P_{Z}

, which is the normalized standard measure on a Riemannian manifold (although the choice of prior distribution could easily be generalized).

As in the standard VAE, we then implement functions $z : X \to Z$ and $t : X \to (0, \infty)$ as neural networks.

3.4 Elbo

We optimize the weights in the network, aiming to minimize the average loss for the loss function

- E_{z \sim Q_{Z}^{t (x), z (x)}} [log p_{X}^{β (z)} (x)] + D_{K L} (Q_{Z}^{t (x), z (x)} ∥ P_{Z}) .

The first integral can often only be approached by sampling, and in that case it is often advantageous to perform a change of variables, commonly known as the reparametrization trick (Kingma & Welling, 2014).

3.5 Reparametrization trick

The reparametrization trick requires a space $Γ$ of noise variables, distributed according to $P_{Γ}$ , and a function $g : Γ \times (0, \infty) \times Z$ such that $g (γ, t, z)$ is approximated according to $Q_{Z}^{t, z}$ .

The following example illustrates that finding such a function $g$ can be a nontrivial task, even for a relatively simple case where $Z = S^{2}$ , the two-dimensional unit sphere in $R^{3}$

. In an attempt to construct a random variable distributed according to

Q_{S^{2}}^{t, z}

, we may first sample a random tangent vector

γ

at the south pole according to an appropriate distribution, and then walk along the great circle in that direction for a distance according to the length of the vector. In other words, the new point is obtained from applying the exponential map at the south pole to the random sample. Next, we compose with a rotation of the sphere that brings the south pole to the point

z

on the sphere. More precisely, we select a map of rotations

T : S^{2} \to S O (3)

such that every

y \in S^{2}

is the image of the south pole under the rotation

T (y)

. Then, a proposed reparametrization

g

is given by

g (γ, z) = T (z) [{exp}_{S P} (γ)] .

In some sense, this reparametrization works: One could find a distribution $P_{Γ}$ on the tangent space at the south pole such that the distribution of $g (γ, z)$ is $Q_{S^{2}}^{t, z}$ . However, for topological reasons, in particular by the hairy ball theorem, it is impossible to find a continuous map $T$ with this property.

3.6 Approximate reparametrization by random walk

Instead, we construct an approximate reparametrization map by approximating Brownian motion by a random walk, similar to how we defined it in this paper. Starting from a point $z$ on the manifold, we set a random step in ambient space $R^{n}$ . We then project back to the manifold and repeat: we take a new step and project back to the manifold. In total, we take $N$ steps, see Fig. 2.

We define the function $g : E^{N} \times (0, \infty) \times Z \to Z$ by

\begin{matrix} g (ϵ_{1}, \dots, ϵ_{N}, t, z) = P (\dots P (P (z + \sqrt{\frac{t}{N}} ϵ_{1}) + \sqrt{\frac{t}{N}} ϵ_{2}) \dots + \sqrt{\frac{t}{N}} ϵ_{N}) . \end{matrix}

If we take $ϵ_{1}, \dots, ϵ_{N}$ as i.i.d. random variables, distributed according to a radially symmetric distribution, then $y = g (ϵ_{1}, \dots, ϵ_{N}, z)$ is approximately distributed as a random variable with density $q_{Z} (t; z, \cdot)$ . This approximation is very accurate for small times $t$ , even for small values of $N$ , if we take $ϵ_{1}, \dots, ϵ_{N}$ approximately Gaussian. The observation that for small times, the diffusion kernel $q_{Z}$ is approximately Gaussian, is also very helpful in approximating the KL term in the loss.

3.7 Approximation of the KL-divergence

Unlike the standard VAE, or the hyperspherical VAE with the Von-Mises distribution, the KL-term cannot be computed exactly for the $Δ$ VAE. There are several techniques one could use to get, nonetheless, a good approximation of the KL divergence. We have implemented an asymptotic approximation, which we will describe first.

Asymptotic approximation

We can use short-term asymptotics, i.e. a parametrix expansion, of the heat kernel on Riemannian manifolds to obtain asymptotic expansions of the entropy.

For a $d$ -dimensional sphere of radius $R$ , the scalar curvature $S c$ equals $d (d - 1) / R^{2}$ . By using a parametrix expansion of the heat kernel, see for instance (Zhao & Song, 2017), we find that

where $r$ is the geodesic distance between $y$ and $z$ .

We then derive, with $Z = S^{2}$ , the following asymptotic behavior for the KL divergence

\begin{matrix} D_{K L} (Q^{t, z} ∥ P_{Z}) & = \int_{Z} q_{Z} (t; z, y) log q_{Z} (t; z, y) d y + log V o l (Z) = - \frac{d}{2} log (2 π t) - \frac{d}{2} + log V o l (Z) + \frac{d + 4}{24} S c t + o (t) . \end{matrix}

We conjecture that this asymptotic formula holds for general $d$ -dimensional Riemannian manifolds. In our implementation, we restrict $t$ so that it cannot become too large, thus ensuring a certain accuracy of the asymptotic expansion.

Figure 3: Latent space representation of MNIST for $S O (3) ≅ {R P}^{3}$ (top-left), ${R P}^{2}$ (bottom-left), flat torus (top-right), and $R^{3}$ (bottom-right). The flat torus is represented as a square with periodic boundary conditions. The projective spaces are represented by a $3$ - and $2$ -dimensional ball respectively, for which every point on the boundary is identified with its reflection through the center. The effect of this identification can be seen, since the same digits that map close to a point on the boundary also map close to the reflected point.

Manifold	LL	ELBO	KL	MSE ( $\times 10^{- 2}$ )
$S^{2}$	-738.76 $\pm$ 0.08	-739.25 $\pm$ 0.09	5.16 $\pm$ 0.13	3.48 $\pm$ 0.05
Embedded Torus	-738.58 $\pm$ 0.08	-740.97 $\pm$ 0.10	6.57 $\pm$ 0.01	3.55 $\pm$ 0.02
Flat Torus	-738.97 $\pm$ 0.08	-741.37 $\pm$ 0.11	6.42 $\pm$ 0.00	3.70 $\pm$ 0.02
$R P^{3}$	-738.81 $\pm$ 0.02	-739.32 $\pm$ 0.03	5.65 $\pm$ 0.08	3.37 $\pm$ 0.02
$R P^{2}$	-740.17 $\pm$ 0.35	-741.19 $\pm$ 0.53	3.15 $\pm$ 0.56	4.49 $\pm$ 0.02
$R^{3}$	-738.27 $\pm$ 0.03	-738.85 $\pm$ 0.04	5.74 $\pm$ 0.05	3.23 $\pm$ 0.02
$R^{2}$	-738.83 $\pm$ 0.08	-739.35 $\pm$ 0.11	4.95 $\pm$ 0.04	3.56 $\pm$ 0.04

Table 1: Numerical results for

Δ

VAEs trained on (non-binarized) MNIST. The values indicate mean and standard deviation over

10

runs. The columns represent the (data-averaged) log-likelihood estimate (LL), Evidence Lower Bound (ELBO), KL-divergence (KL) and mean squared error (MSE).

Manifold	LL	ELBO	KL	MSE ( $\times 10^{- 2}$ )
$S^{2}$	-3779.75 $\pm$ 2.25	-3825.11 $\pm$ 4.21	3.45 $\pm$ 0.01	2.82 $\pm$ 0.21
Embedded Torus	-3787.34 $\pm$ 11.8	-3809.33 $\pm$ 23.0	11.2 $\pm$ 1.20	1.85 $\pm$ 1.18
Flat Torus	-3773.99 $\pm$ 1.70	-3813.00 $\pm$ 5.16	6.42 $\pm$ 0.00	0.90 $\pm$ 0.25
$R P^{3}$	-3774.61 $\pm$ 0.53	-3821.25 $\pm$ 2.44	3.70 $\pm$ 0.00	2.62 $\pm$ 0.12
$R P^{2}$	-3789.13 $\pm$ 8.36	-3850.19 $\pm$ 33.7	3.92 $\pm$ 1.97	4.02 $\pm$ 1.74
$R^{3}$	-3779.67 $\pm$ 4.06	-3783.03 $\pm$ 5.71	9.73 $\pm$ 0.45	0.46 $\pm$ 0.26
$R^{2}$	-3785.64 $\pm$ 5.26	-3789.60 $\pm$ 6.19	8.26 $\pm$ 0.53	0.85 $\pm$ 0.28

Table 2: Numerical results for

Δ

VAEs trained on a simple picture consisting of the lowest non-trivial Fourier components. The values indicate mean and standard deviation over

10

runs. The columns represent the (data-averaged) log-likelihood estimate (LL), Evidence Lower Bound (ELBO), KL-divergence (KL) and mean squared error (MSE).

Numerical approximation and integration

For some manifolds such as spheres or flat tori, exact solutions to the heat kernel are available, usually in terms of infinite series and special functions. For other, small-dimensional, Riemannian manifolds, the heat kernel may be accurately computed numerically. In both cases, the integration in the KL-term may be performed numerically. If the dimensionality of the underlying manifold gets to large, one may have to resort to Monte Carlo approximation of the integral.

4 Experiments

We have implemented $Δ$ VAEs with latent spaces of $d$ -dimensional spheres, a flat two-dimensional torus, a torus embedded in $R^{3}$ , the $S O (3)$ and real projective spaces ${R P}^{d}$ .

For all our experiments we used multi-layer perceptrons for the encoder and decoder with three and two hidden layers respectively. Recall that the encoder needs to produce both a point

z

on the manifold and a time

t

for the transition kernel. These functions share all layers, except for the final step where we project, with the projection map

P

, from the last hidden layer to the manifold to get

z

, and use an output layer with a

tanh

activation function to obtain

t

The encoder and decoder are connected by a sampling layer, in which we approximate sampling from the transition kernel of Brownian motion according to the reparametrization trick described in Section 3.6.

4.1 $Δ$ VAEs for MNIST

We then trained $Δ$ VAEs on MNIST. We show the manifolds as latent space with encoded MNIST digits in Figs. 1 and 3. When MNIST is trained on different latent spaces, different adjacency structures between digits may become apparent, providing topological information.

The $S O (3)$ is isometric to a scaling of the ${R P}^{3}$ (with natural choices of Riemannian metrics). Although we have implemented an embedding and projection map for this embedding for the $S O (3)$

directly based on an SVD decomposition to find the nearest orthogonal matrix, training on the

{R P}^{3}

with the following trick was faster and we only present these results.

For training the projective spaces, we used an additional trick, where instead of embedding ${R P}^{d}$ in a Euclidean space, we embed $S^{d}$ in $R^{d + 1}$ , and make the decoder neural network even by construction (i.e. the decoder applied to a point $s$ on the sphere equals the decoder applied to a point $- s$ ). Then, an encoder and decoder to and from the ${R P}^{3}$ are defined implicitly. However, it must be noted that this setup does not allow for a homeomorphic encoding (because ${R P}^{d}$ does not embed in $S^{d}$ ).

The numerically computed ELBO, reconstruction error, KL-divergence and mean-squared-error are shown in Table 1 together with the estimated log-likelihood for a test dataset of MNIST.

Estimation of log-likelihood

For the evaluation of the proposed methods we have estimated the log-likelihood of the test dataset according to the importance sampling presented in (Burda et al., 2016). The approximate log-likelihood of datapoint $x$ is calculated by sampling $L$ latent variables $z^{(1)}, \dots, z^{(L)}$ according to the approximate posterior $Q_{Z}^{t (x), z (x)}$ . The estimated log-likelihood for datapoint $x$ is given by

log p_{X}^{β} (x) \approx log ⎛ ⎜ ⎝ \frac{1}{L} L \sum l = 1 \frac{p_{X}^{β (z^{(l)})} (x) p_{Z} (z^{(l)})}{q_{Z} (t (x); z (x), z^{(l)})} ⎞ ⎟ ⎠ .

The log-likelihood estimates presented in Table 1 are obtained with $L = 100$ samples for each datapoint, averaged over all datapoints.

4.2 Translations of periodic pictures

To test whether a $Δ$ VAE can capture topological properties, we trained it on synthetic datasets consisting of translations of the same periodic picture.

Our input pictures were discretized to 64 $\times$ 64 pixels, but to ease presentation, we discuss them below as if they were continuous. Note that with the choice of multi-layer perceptrons as encoder and decoder networks, the network has no information about which pixels are contiguous.

Simple picture

To illustrate the idea, we start with a very simple picture $p i c : [- π, π]^{2} \to R$ , which only consists of the lowest Fourier components in each direction

p i c (θ, ϕ) = cos (θ) + cos (ϕ) .

By translating the picture, i.e. by considering different phases, we obtain a submanifold of the space of all pictures. When we interpret this space as $L^{2} ([- π, π]^{2})$ , we get an embedding of the flat torus which is isometric up to a scaling. In other words, if we train the $Δ$ VAE with a flat torus as latent space, we essentially try to find the identity map.

Figs. 4 and 4 illustrate that for this simple case, the $Δ$ VAE with a flat torus as latent space indeed captures the translation of the picture as a latent variable. The fact that Fig. 4 is practically a reflection and translation of the legend in Fig. 5, shows that there is an almost isometric correspondence between the translation of the original picture and the encoded latent variable.

Although the figure does not provide a proof that the embedding is homeomorphic, it is in principle possible to give a guarantee that the encoder map is surjective (more precisely, that it has degree $1$ or $- 1$ ).

Figure 5: Colors used to color encoded pictures in latent space. The horizontal direction represents horizontal translation, the vertical direction represents vertical translation of the original, periodic picture. The boundary conditions are periodic.

We contrast this to when we train the same dataset on a $Δ$ VAE with a sphere as a latent space. Fig. 4 displays typical results, showing that large parts of the sphere are not covered.

We present further numerical results in Table 2. We note that of all the closed manifolds, the MSE is by far the lowest for the flat torus. The KL divergence is lower for the positively-curved spaces $S^{2}$ , ${R P}^{2}$ and ${R P}^{3}$ , but the MSE for those is about a factor three higher. For $R^{3}$ as latent space the MSE is even lower, but it is hard to compare because the KL term is not directly comparable, and it is the combination of the two that is optimized.

More complicated pictures

As an exploratory analysis of the extent to which the $Δ$ VAE can capture the topological structure, we trained on fixed datasets, each consisting of translations of a fixed random picture. These random pictures, viewed as functions $p i c : [- π, π]^{2} \to R$ , are created by randomly drawing complex Fourier coefficients $a_{k, ℓ}$ from a Gaussian distribution:

p i c (θ, ϕ) = R e ⎛ ⎝ N \sum k, ℓ = - N γ^{| k | + | ℓ | - 1} a_{k, ℓ} exp (i k θ + i ℓ ϕ) ⎞ ⎠

where $γ \in (0, 1]$ is a discounting factor. The $Δ$ VAE with a flat torus as latent space is also capable of capturing the translations when the picture is generate this way, as long as the higher Fourier components carry not too much weight, see Fig. 6.

Figure 6: Latent space and reconstruction for a $Δ$ VAE with a flat torus as latent space, trained on a picture consisting of several Fourier components.

Too complicated pictures

Generically, we see that when there is too much weight on the higher Fourier components in the picture ( $N ≳ 10$ with $γ ≳ 0.3$ ), the $Δ$ VAE is no longer capable of capturing the translations, see Fig. 7.

We only tested a very simple setup. It is likely that convolutional neural networks would significantly improve the performance in capturing the translations, since implicit information is added to the network about the underlying geometry of the pictures.

Initial experiments with pretraining the network on simple Fourier images also show an increased success rate of capturing the topological structure.

Figure 7: Flat torus latent space and reconstruction for a $Δ$ VAE trained on a picture with too much weight on higher Fourier components for the $Δ$ VAE to capture the topological structure.

Interesting patterns

For moderate weights on Fourier components, the data manifold is often mapped to the latent space in very interesting ways, see Fig. 8. The patterns that occur in this manner are very structured, suggesting that it may be possible to find underlying mathematical reasons for this structure. Moreover, it gives some indication that we may understand how the network arrives at such patterns, and how we may control them.

5 Conclusion

We developed and implemented Diffusion Variational Autoencoders, which allow for arbitrary manifolds as a latent space ¹¹1https://github.com/LuisArmandoPerez/DiffusionVAE. Our original motivation was to investigate to which extent VAEs find semantically meaningful latent variables, and more specifically, whether they can capture topological and geometrical structure in datasets. By allowing for an arbitrary manifold as a latent space, $Δ$ VAEs can remove obstructions to capturing such structure.

Indeed, our experiments with translations of periodic images show that a simple implementation of a $Δ$ VAE with a flat torus as latent space is capable of capturing topological properties.

However, we have also observed that when we use periodic images with significant high Fourier components, it is more challenging for the $Δ$ VAE to capture the topological structure. We will investigate further whether this can be resolved by better network design, by pretraining, different loss functions, or even more intensive extensions of the VAE algorithm.

Our exploratory analysis shows that we can ask well-defined questions about whether a VAE can capture topological structure, and that such questions are amenable to a statistical approach. Moreover, the patterns observed in latent space in Fig. 6 and Fig. 8, suggest that it may be possible to capture in mathematical theorems what the image of a data manifold in latent space would look like, when the networks are (almost) minimizing the loss. Combining these different perspectives, we may eventually develop more understanding on how to develop algorithms that capture semantically meaningful latent variables.

References

Burda et al. (2016) Burda, Y., Grosse, R., and Salakhutdinov, R. Importance weighted autoencoders. ICLR, 2016.
Burgess et al. (2018) Burgess, C. P., Higgins, I., Pal, A., Matthey, L., Watters, N., Desjardins, G., and Lerchner, A. Understanding disentangling in $β$ -VAE. arXiv preprint, arXiv:1804.03599, 2018.
Chen et al. (2018) Chen, T. Q., Li, X., Grosse, R. B., and Duvenaud, D. K. Isolating sources of disentanglement in Variational Autoencoders. In Advances in Neural Information Processing Systems 31. 2018.
Davidson et al. (2018) Davidson, T. R., Falorsi, L., De Cao, N., Kipf, T., and Tomczak, J. M. Hyperspherical variational auto-encoders. In 34th Conference on Uncertainty in Artificial Intelligence.
2018.
de Haan & Falorsi (2018) de Haan, P. and Falorsi, L. Topological constraints on homeomorphic auto-encoding. arXiv preprint, arXiv:1812.10783, 2018.
Dilokthanakul et al. (2016) Dilokthanakul, N., Mediano, P. A. M., Garnelo, M., Lee, M. C. H., Salimbeni, H., Arulkumaran, K., and Shanahan, M. Deep unsupervised clustering with Gaussian mixture Variational Autoencoders. arXiv preprint, arXiv:1611.02648, 2016.
Falorsi et al. (2018) Falorsi, L., de Haan, P., Davidson, T. R., De Cao, N., Weiler, M., Forré, P., and Cohen, T. S. Explorations in homeomorphic variational auto-encoding. arXiv preprint, arXiv:1807.04689, 2018.
Gemici et al. (2016) Gemici, M. C., Rezende, D., and Mohamed, S. Normalizing flows on Riemannian manifolds. arXiv preprint, arXiv:1611.02304, 2016.
Higgins et al. (2017) Higgins, I., Matthey, L., Pal, A., Burgess, C., Glorot, X., Botvinick, M., Mohamed, S., and Lerchner, A. $β$ -vae: Learning basic visual concepts with a constrained variational framework. ICLR, 2017.
Higgins et al. (2018) Higgins, I., Amos, D., Pfau, D., Racaniere, S., Matthey, L., Rezende, D., and Lerchner, A. Towards a definition of disentangled representations. arXiv preprint, arXiv:1812.02230, 2018.
Hinton & Zemel (1994) Hinton, G. E. and Zemel, R. S. Autoencoders, minimum description length and helmholtz free energy. In Advances in neural information processing systems, 1994.
Honkela & Valpola (2004) Honkela, A. and Valpola, H. Variational learning and bits-back coding: An information-theoretic view to Bayesian learning. IEEE Transactions on Neural Networks, 15(4):800–810, 7 2004. ISSN 1045-9227.
Hsu (2008) Hsu, E. P. A brief introduction to Brownian motion on a Riemannian manifold. Lecture notes, 2008.
Jørgensen (1975) Jørgensen, E. The central limit problem for geodesic random walks. Zeitschrift für Wahrscheinlichkeitstheorie und verwandte Gebiete, 32(1-2):1–64, 1975. ISSN 00443719.
Kim & Mnih (2018) Kim, H. and Mnih, A. Disentangling by factorising. arXiv preprint, arXiv:1802.05983, 2018.
Kingma & Welling (2014) Kingma, D. P. and Welling, M. Auto-Encoding Variational Bayes. In ICLR, 2014.
Kingma et al. (2016) Kingma, D. P., Salimans, T., Jozefowicz, R., Chen, X., Sutskever, I., and Welling, M. Improved Variational Inference with Inverse Autoregressive Flow. In Advances in Neural Information Processing Systems 29. 2016.
Kumar et al. (2018) Kumar, A., Sattigeri, P., and Balakrishnan, A. Variational inference of distentangled latent concepts from unlabeled observations. ICLR, 2018.
Liu & Zhu (2017) Liu, C. and Zhu, J. Riemannian Stein variational gradient descent for Bayesian inference. arXiv preprint, arXiv:1711.11216, 2017.
Portegies (2018) Portegies, J. W. Ergo learning. Nieuw Archief voor Wiskunde, 5(3):199–205, 2018.
Rezende & Mohamed (2015) Rezende, D. and Mohamed, S. Variational inference with normalizing flows. In Proceedings of the 32nd International Conference on Machine Learning, 2015.
Rezende et al. (2014) Rezende, D. J., Mohamed, S., and Wierstra, D. Stochastic backpropagation and approximate inference in deep generative models. arXiv preprint, arXiv:1401.4082, 2014.
Rissanen (1978) Rissanen, J. Modeling by shortest data description. Automatica, 14(5):465–471, 9 1978. ISSN 0005-1098.
Salimans et al. (2015) Salimans, T., Kingma, D., and Welling, M. Markov Chain Monte Carlo and Variational Inference: Bridging the gap. In Proceedings of the 32nd International Conference on Machine Learning, 2015.
Schmidhuber (1997) Schmidhuber, J. Discovering neural nets with low Kolmogorov complexity and high generalization capability. Neural Networks, 10(5):857 – 873, 1997. ISSN 0893-6080.
Solomonoff (1964) Solomonoff, R. A formal theory of inductive inference. Part I. Information and Control, 7(1):1–22, 3 1964. ISSN 0019-9958.
Tomczak & Welling (2017) Tomczak, J. M. and Welling, M. VAE with a VampPrior. arXiv preprint, arXiv:1705.07120, 2017.
Zhao & Song (2017) Zhao, C. and Song, J. S. Exact heat kernel on a hypersphere and its applications in kernel SVM. Frontiers in Applied Mathematics and Statistics, 4:1, 1 2017. ISSN 2297-4687.

Diffusion Variational Autoencoders

Authors

Variational Autoencoder with Learned Latent Structure

Disentanglement with Hyperspherical Latent Spaces using Diffusion Variational Autoencoders

Robust Feature Disentanglement in Imaging Data via Joint Invariant Variational Autoencoders: from Cards to Atoms

Encoded Prior Sliced Wasserstein AutoEncoder for learning latent manifold representations

Use of Student's t-Distribution for the Latent Layer in a Coupled Variational Autoencoder

Dirichlet Graph Variational Autoencoder

Quantifying the Effects of Enforcing Disentanglement on Variational Autoencoders

1 Introduction

2 Related work