Title: VARIATIONAL DIFFUSION AUTOENCODERS WITH RANDOM WALK SAMPLING

Abstract: Variational inference (VI) methods and especially variational autoencoders (VAEs) specify scalable generative models that enjoy an intuitive connection to manifold learning -with many default priors the posterior/likelihood pair q(z|x)/p(x|z) can be viewed as an approximate homeomorphism (and its inverse) between the data manifold and a latent Euclidean space. However, these approximations are well-documented to become degenerate in training. Unless the subjective prior is carefully chosen, the topologies of the prior and data distributions often will not match. Conversely, diffusion maps (DM) automatically infer the data topology and enjoy a rigorous connection to manifold learning, but do not scale easily or provide the inverse homeomorphism. In this paper, we propose a) a principled measure for recognizing the mismatch between data and latent distributions and b) a method that combines the advantages of variational inference and diffusion maps to learn a homeomorphic generative model. The measure, the locally bi-Lipschitz property, is a sufficient condition for a homeomorphism and easy to compute and interpret. The method, the variational diffusion autoencoder (VDAE), is a novel generative algorithm that first infers the topology of the data distribution, then models a diffusion random walk over the data. To achieve efficient computation in VDAEs, we use stochastic versions of both variational inference and manifold learning optimization. We prove approximation theoretic results for the dimension dependence of VDAEs, and that locally isotropic sampling in the latent space results in a random walk over the reconstructed manifold. Finally, we demonstrate our method on various real and synthetic datasets, and show that it exhibits performance superior to other generative models.
Variational inference (VI, Jordan et al. (1999); Wainwright et al. (2008) ) is a machine learning method that combines Bayesian statistics and latent variable models to approximate some probability density p(x). VI assumes and exploits a latent variable structure in the assumed data generation process, that the observations x ∼ p(x) are conditionally distributed given unobserved latent vari-

1. INTRODUCTION:

Recent developments in generative models such as variational auto-encoders (VAEs, Kingma & Welling (2013) ) and generative adversarial networks (GANs, Goodfellow et al. (2014) ) have made it possible to sample remarkably realistic points from complex high dimensional distributions at low computational cost. While their methods are very different -one is derived from variational inference and the other from game theory -their ends both involve learning smooth mappings from a user-defined prior distribution to the modeled distribution.

These maps are closely tied to manifold learning when the prior is supported over a Euclidean space (e.g. Gaussian or uniform priors) and the data lie on a manifold (also known as the Manifold Hypothesis, see Narayanan & Mitter (2010) ; Fefferman et al. (2016) ). This is because manifolds themselves are defined by sets that have homeomorphisms to such spaces. Learning such maps is beneficial to any machine learning task, and may shed light on the success of VAEs and GANs in modeling complex distributions. Furthermore, the connection to manifold learning may explain why these generative models fail when they do. Known as posterior collapse in VAEs (Alemi et al., 2017; Zhao et al., 2017; He et al., 2019; Razavi et al., 2019) and mode collapse in GANs (Goodfellow, 2017) , both describe cases where the forward/reverse mapping to/from Euclidean space collapses large parts of the input to a single output. This violates the bijective requirement of the homeomorphic mapping. It also results in degenerate latent spaces and poor generative performance. A major cause of such failings is when Figure 1 : A diagram depicting one step of the diffusion process modeled by the variational diffusion autoencoder (VDAE). The diffusion and inverse diffusion maps ψ, ψ -1 , as well as the covariance C of the random walk on M Z , are all approximated by neural networks.

the geometries of the prior and target data do not agree. We explore this issue of prior mismatch and previous treatments of it in Section 3.

Given their connection to manifold learning, it is natural to look to classical approaches in the field for ways to improve VAEs. One of the most principled methods is spectral learning (Schölkopf et al., 1998; Roweis & Saul, 2000; Belkin & Niyogi, 2002) which involves describing data from a manifold X ⊂ M X by the eigenfunctions of a kernel on M X . We focus specifically on DMs, where Coifman & Lafon (2006) show that normalizations of the kernel approximate a very specific diffusion process, the heat kernel over M X . A crucial property of the heat kernel is that, like its physical analogue, it defines a diffusion process that has a uniform stationary distribution -in other words, drawing from this stationary distribution draws uniformly from the data manifold. Moreover, Jones et al. (2008) established another crucial property of DMs, namely that distances in local neighborhoods in the eigenfunction space are nearly isometric to corresponding geodesic distances on the manifold. However, despite its strong theoretical guarantees, DMs are poorly equipped for large scale generative modeling as they are not easily scalable and do not provide an inverse mapping from the intrinsic feature space.

In this paper we address issues in variational inference and manifold learning by combining ideas from both. Theory in manifold learning allows us to better recognize prior mismatch, whereas variational inference provides a method to learn the difficult to approximate inverse diffusion map.

Our contributions: 1) We introduce the locally bi-Lipschitz property, a sufficient condition for a homeomorphism, for measuring the stability of a mapping between latent and data distributions. 2) We introduce VDAEs, a class of variational autoencoders whose encoder-decoder feedforward pass approximates the diffusion process on the data manifold with respect to a user-defined kernel k. 3) We show that deep neural networks are capable of learning such diffusion processes, and 4) that networks approximating this process produce random walks that have certain desirable properties, including well defined transition and stationary distributions. 5) Finally, we demonstrate the utility of the VDAE framework on a set of real and synthetic datasets, and show that they have superior performance and satisfy the locally bi-Lipschitz property where GANs and VAEs do not. ables z. By modeling the conditional distribution, then marginalizing over z, as in

EQUATION

we obtain the model evidence, or likelihood that x could have instead been drawn from p θ (x).

Maximizing Eq. 1 leads to an algorithm for finding likely approximations of p(x). As the cost of computing this integral scales exponentially with the dimension of z, we instead maximize the evidence lower bound (ELBO):

EQUATION

where q(z|x) is usually an approximation of p θ (z|x). Optimizing the ELBO is sped up by taking stochastic gradients (Hoffman et al., 2013) , and further accelerated by learning a global function approximator q φ in an autoencoding structure (Kingma & Welling, 2013) .

Diffusion maps (DMs, Coifman & Lafon (2006) ) on the other hand, are a class of kernel methods that perform non-linear dimensionality reduction on a set of observations X ⊆ M X , where M X is the data manifold. Given a symmetric and positive kernel k, DM considers the induced random walk on the graph of X, where given x, y ∈ X, the transition probabilities p(y|x) = p(x, y) are row normalized versions of k(x, y). Moreover, the diffusion map ψ embeds the data X ∈ R m into the Euclidean space R D so that the diffusion distance is approximated by Euclidean distance. This is a powerful property, as it allows the arbitrarily complex random walk induced by k on M X to become an isotropic Gaussian random walk on ψ(M X ).

SpectralNet is an algorithm introduced by algorithm in Shaham et al. (2018b) to speed up the diffusion map. Until recently, the method ψ k could only be computed via the eigendecomposition of K. As a result, DMs were only be tractable for small datasets, or on larger datasets by combining landmark-based estimates and Nystrom approximation techniques. However, Shaham et al. (2018b) propose approximations of the function ψ itself in the case that the kernel k is symmetric.

In particular, we will leverage SpectralNet to enforce our diffusion embedding prior.

Locally bi-lipschitz coordinates by kernel eigenfunctions. (Jones et al. ( 2008)) analyzed the construction of local coordinates of Riemannian manifolds by Laplacian eigenfunctions and diffusion map coordinates. They establish, for all x ∈ X, the existence of some neighborhood U (x) and d spectral coordinates given U (x) that define a bi-Lipschitz mapping from U (x) to R d . With a smooth compact Riemannian manifold, U (x) can be chosen to be a geodesic ball with radius a constant multiple of the inradius (the radius of the largest possible ball around x without intersecting with the manifold boundary), where the constant is uniform for all x, but the indices of the d spectral coordinates as well as the local bi-Lipschitz constants may depend on x. Specifically, the Lipschitz constants involve inverse of the inradius at x multiplied again by some global constants. For completeness we give a simplified statement of the Jones et al. (2008) result in the supplementary material.

Using the compactness of the manifold, one can always cover the manifold with m many neighborhoods (geodesic balls) on which the bi-Lipschitz property in Jones et al. (2008) holds. As a result, there are a total of D spectral coordinates, D ≤ md (in practice D is much smaller than md, since the selected spectral coordinates in the proof of Jones et al. (2008) tend to be low-frequency ones, and thus the selection on different neighborhoods tend to overlap), such that on each of the m neighborhoods, there exists a subset of d spectral coordinates out of the D ones which are bi-Lipschitz on the neighborhood, and the Lipschitz constants can be bounded uniformly from below and above.

3. MOTIVATION AND RELATED WORK:

Our proposed measure and model is motivated by degenerate latent spaces and poor generative performance in a variational inference framework arising from prior mismatch: when the topologies of the data and prior distributions do not agree. In real world data, this is usually due to two factors: first, when the dimensionalities of the distributions do not match, and second, when the geometries do not match. It is easy to see that homeomorphisms between the distributions will not exist in either case: pointwise correspondences cannot be established, thus the bijective condition cannot be met. As a result, the model has poor generative performance -for each point not captured in the pointwise correspondence, the latent or generated distribution loses expressivity.

Though the default choice of Gaussian distribution for p(z) is mathematically elegant and computationally expedient, there are many datasets, real and synthetic, for which this distribution is ill-suited.

It is well known that spherical distributions are superior for modeling directional data (Fisher et al., 1993; Mardia, 2014) , which can be found in fields as diverse as bioinformatics (Hamelryck et al., 2006) , geology (Peel et al., 2001) , material science (Krieger Lassen et al., 1994) , natural image processing (Bahlmann, 2006) , and simply preprocessed datasets1 . Additionally observe that no homeomorphism exists between R k and S 1 for any k. For data distributed on more complex manifolds, the literature is sparse due to the difficult nature of such study. However, the manifold hypothesis is well-known and studied (Narayanan & Mitter, 2010; Fefferman et al., 2016) .

Previous research on alleviating prior mismatch exists. Davidson et al. (2018) ; Xu & Durrett (2018) consider VAEs with the von-Mises Fisher prior, a geometrically hyperspherical prior. (Rey et al., 2019) further model arbitrarily complex manifolds as priors, but require explicit knowledge of the manifold (i.e. its projection map, scalar curvature, and volume). Finally, Tomczak & Welling (2017) consider mixtures of any pre-existing priors. But while these methods increase the expressivity of the priors available, they do not prescribe a method for choosing the prior itself. That responsibility still lies with the user.

Convserly, our method chooses the best prior automatically. To our knowledge, ours is the first to take a data-driven approach to prior selection. By using some data to inform the prior, we not only guarantee the existence of a homeomorphism between data and prior distributions, we explicitly define it by the learned diffusion map ψ.

4. METHOD:

In this section we propose VDAEs, a variational inference method that, given the data manifold M X , observations X ⊂ M X , and a kernel k, models the geometry of X by approximating a random walk over the latent diffusion manifold M Z := ψ(M X ). The model is trained by maximizing the local evidence: the evidence (i.e. log-likelihood) of each point given its random walk neighborhood.

Points are generated from the trained model by sampling from π, the stationary distribution of the resulting random walk.

Starting from some point x ∈ X, we can roughly describe one step of the walk as the composition of three functions: 1) the approximate diffusion map ψΘ : M X → M Z , 2) a sampling procedure from the learned diffusion process z ∼ q φ (z |x) = N ( ψΘ (x), Cφ ) on M Z , and 3) the learned inverse diffusion map ψs-1

θ : M Z → M X that produces x ∼ p(x |z ) = N ( ψ-1 θ (z ), cI)

where the constant c is user-defined and fixed.

We rely crucially on three advantages of our latent space ψΘ (X): a) that it is well-defined (given the first D eigenvalues of k are distinct), b) well-approximated (given SpectralNet) and c) that Euclidean distances in M Z approximate single-step random walk distances on M X (see Section 2 and Coifman & Lafon ( 2006)). Thus the transition probabilities induced by k can be approximated by Gaussian kernels2 in M Z .

Therefore, to model a diffusion random walk over M Z , we must learn the functions ψΘ , ψ-1 θ , Cφ that approximate the diffusion map, its inverse, and the covariance of the random walk on M Z , at all points z ∈ M Z . SpectralNet gives us ψΘ . To learn ψ-1 θ and Cφ , we use variational inference.

4.1. THE LOWER BOUND:

Formally, let us define U x := B d (x, δ) ∩ M X , where B d (x, δ) as the δ-ball around x with respect to d(•, •), the diffusion distance on M Z . For each x ∈ X, we define the local evidence of x as

EQUATION

where p(x |x)| Ux is the restriction of p(x |x) to U x . The resulting local evidence lower bound is:

log p θ (x |x) ≥ -D KL (q φ (z |x)||p θ (z |x))

divergence of random walk distributions

+ E z ∼q φ (z |x) log p θ (x |z ) neighborhood reconstruction error . (4)

Note that the neighborhood reconstruction error should be differentiated from the self reconstruction error that is in VAEs. Eq. 4 produces the empirical loss function:

EQUATION

where Here we discuss the algorithm for generating data points from p(x). Composing q φ (z |x)(≈ p θ (z |x)) with p θ (x |z ) gives us an approximation of p θ (x |x). Then the simple, parallelizable, and fast random walk based sampling procedure naturally arises: initialize with an arbitrary point on the manifold x 0 ∈ M X , then pick suitably large N and for n = 1, . . . , N draw x n ∼ p(x|x n-1 ).

z i = g φ,Θ (x, i ), i ∼ N (0, I). g φ,

Eventually, our diffusion random walk converges on its stationary distribution π. By Coifman & Lafon ( 2006), this is guaranteed to be the uniform distribution on the data manifold. See Section 6.2 for examples of points drawn from this procedure.

4.3. A PRACTICAL IMPLEMENTATION:

We now introduce a practical implementation VDAEs, considering the case where ψΘ (x), q φ (z |x) and p θ (x |z ) are neural network functions, as they are in VAEs and SpectralNet, respectively.

The neighborhood reconstruction error. Since q φ (z |x) models the neighborhood of ψΘ (x), we may sample q φ to obtain z (the neighbor of x in the latent space). This gives p θ (x |x) ≈ ψ-1 θ (q φ (z |x)), where ψ -1 exists due to the bi-Lipschitz property. We can efficiently approximate x ∈ M X by considering the closest embedded data point ψΘ (x) ∈ M Z to z = ψΘ (x ). This is because Euclidean distance on M Z approximates the diffusion distance on M X . In other words, x ∼ p θ (x |x) ≈ ψ-1 θ (q φ (z |x)) which we approximate empirically by

EQUATION

where A ⊆ X is the training batch.

On the other hand, the divergence of random walk distributions -D KL (q φ (z |x)||p θ (z |x)) can be modeled simply as the divergence of two Gaussian kernels defined on M Z . Though p θ (z |x) is intractable, the diffusion map ψ gives us the diffusion embedding Z, which is an approximation of the true distribution of p θ (z |x) in a neighborhood around z = ψ(x). We estimate the first and second moments of this distribution in R D by computing the local Mahalanobis distance of points in the neighborhood. Then, by minimizing the KL divergence between q φ (z |x) and the one implied by this Mahalanobis distance, we obtain the loss:

EQUATION

where

C φ (x) is a neural network function, Σ * (x) = Cov(B d (ψ(x), δ) ∩ Z)

is the covariance of the points in a neighborhood of z = ψ(x) ∈ Z, and α is a scaling parameter. Note that C φ (x) does not have to be diagonal, and in fact is most likely not. Combining Eqs. 6 and 7 we obtain Algorithm 1.

Now we consider the sampling procedure. Since we use neural networks to approximate q φ (z |x) and p θ (x |z ), the generation procedure is highly parallelizable. We empirically observe the random walk enjoys rapid mixing properties -it does not take many iterations of the random walk to sample from all of M Z4 . This leads to Algorithm 2.

Algorithm 2 VDAE sampling

X 0 ← Initialize with points X 0 ⊂ X t ← 0 while p(X 0 ) ≈ π do for x t ∈ X t do z t+1 ∼ p φ (z | ψΘ (x t ))

Take one step of the diffusion random walk

x t+1 ∼ p θ (x|z t+1 )

Map back into input space t ← t + 1

5. THEORY:

We theoretically prove that the desired inverse map ψ -1 from spectral coordinate codes back to the manifold can be approximated by a decoder network, where the network complexity is bounded by quantities related to the intrinsic geometry of the manifold. This section relies heavily on the known bi-Lipschitz property of DMs Jones et al. ( 2008), which we are approximating with the VDAE latent space without the need for regularization.

5.1. THEOREMS:

The theory for the capacity of the encoder to map M to the diffusion map space ψ(M) has already been considered in Shaham et al. (2018a) and Mishne et al. (2017) . We instead focus on the decoder, which requires a different treatment. The following theorem is proved in Appendix A.3, based upon the result in Jones et al. (2008) . Theorem 1. Let M X ⊂ R m be a smooth d-dimensional manifold, ψ(M X ) ⊂ R D be the diffusion map for D ≥ d large enough to have a subset of coordinates that are locally bi-Lipschitz. Let X = [X 1 , ..., X m ] be the set of all m extrinsic coordinates of the manifold. Then there exists a sparsely-connected ReLU network f N , with 4DC M X nodes in the first layer, 8dmN nodes in the second layer, and 2mN nodes in the third layer, and m nodes in the output layer, such that

X(ψ(x)) -f N (ψ(x)) L 2 (ψ(M X )) ≤ √ mC ψ / √ N (8)

where the norm is interpreted as

F 2 L 2 (ψ(M)) := F (ψ(x)) 2 2 dψ(x).

Here C ψ depends on how sparsely X(ψ(x)) Ui can be represented in terms of the ReLU wavelet frame on each neighborhood U i , and C M X on the curvature and dimension of the manifold M X .

Theorem 1 is complementary to the theorem in Shaham et al. (2018a) , which provides guarantees for the encoder, as Theorem 1 demonstrates a similar approximation theoretic argument for the decoder. The proof is built on two properties of ReLU neural networks: 1) their ability to split curved domains into small, almost Euclidean patches, 2) their ability to build differences of bump functions We also discuss the connections between the distribution at each point in diffusion map space, q φ (z|x), and the result of this distribution after being decoded through the decoder network f N (z) for z ∼ q φ (z|X). Similar to Singer & Coifman (2008) , we characterize the covariance matrix

Cov(f N (z)) := E z∈q φ (z|x) [f N (z)f N (z) T ].

The following theorem is proved in Appendix A.3. Theorem 2. Let f N be a neural network approximation to X as in Theorem 1, such that it approximates the extrinsic manifold coordinates. Let C ∈ R m×m be the covariance matrix

C = E z∈q φ (z|x) [f N (z)f N (z) T ]. Let q φ (z|x) ∼ N (ψ(x), Σ)

with small enough Σ that there exists a patch U z0 ⊂ M around z 0 satisfying the bi-Lipschitz property of Jones et al. ( 2008), and such that P r(z ∼ q φ (z|x) ∈ ψ(U z0 )) < . Then the number of eigenvalues of C greater than is at most d, and C = J z0 ΣJ T z0 + O( ) where J z0 is the m × D Jacobian matrix at z 0 .

Theorem 2 establishes the relationship between the covariance matrices used in the sampling procedure and their image under the decoder f N to approximate ψ -1 . Similar to Singer & Coifman (2008) , we are able to sample according to a multivariate normal distribution in the latent space. Thus, the resulting cloud in the data space is distorted (to first order) by the local Jacobian of the map f N . The key insight of Theorem 2 is from combining this idea with the observation of Jones et al. ( 2008) that ψ -1 depends locally on only d of the coordinates in the D dimensional latent space.

6 EXPERIMENTAL RESULTS

6.1. VIDEO OF ROTATING FIGURE:

We consider the problem of generating new frames from a video of rigid movement. We take 200 frames of a color video (each frame is 100 × 80 × 3) of a spinning bulldog Lederman & Talmon (2018) . Due to the spinning of figure and the fixed background, this creates a low-dimensional approximately circular manifold.

We compare our method to VAE, the Wasserstein GAN Gulrajani et al. (2017) (with a bi-lipchitz constraint on the critic), and the hyperspherical VAE Davidson et al. (2018) . For the VAE, we use a two dimensional Gaussian prior p θ (z), such that z ∼ N (0, I 2 ). The noise injected to the GAN is drawn from a two dimensional uniform distribution p θ (z), such that z i ∼ U (0, 1), i = 1, 2. For the spherical VAE, we use a latent dimension of D = 2, which highlights the dimension mismatch issue that occurs with a spherical prior. This is a benefit of VDAE, even if we choose D > d the latent embedding will still only be locally d dimensional. We use the same architecture for all networks which consists of one hidden layer with 512 neurons, activation function for all networks are tanh. In Fig. 2 , we present 300 generated samples, by displaying them on a scatter plot with coordinates corresponding to their latent dimensions z 1 and z 2 .

6.2. DATA GENERATION FROM UNIFORMLY SAMPLED MANIFOLDS:

In this series of experiments, we visualize the results of the sampling procedure in Algorithm 2 on three synthetic manifolds. As discussed in 4.2, we randomly select an initial seed point, then recursively sample from p θ (x |x) many times to simulate a random walk on the manifold. In the top row of Fig. 3 , we highlight the location of the initial seed point, take 20 steps of the random walk, and display the resulting generated points on three learned manifolds. Clearly after a large number of resampling iterations, the algorithm continues to generate points on the manifold, and the distribution of sampled points converges to a uniform stationary distribution on the manifold. Moreover, this stationary distribution is reached very quickly. In the bottom row of the same Fig. 3 , we show p θ (x |x) by sampling a large number of points sampled from the single seed point. As can be seen, a single step of p θ (x |x) covers a large part of the latent space. The architecture also uses one hidden layer of 512 neurons and tanh activations.

6.3. CLUSTER CONDITIONAL SAMPLING:

In this section, we deal with the problem of generating samples from data with multiple clusters in an unsupervised fashion (i.e. no a priori knowledge of the cluster structure). Clustered data creates a problem for many generative models, as the topology of the latent space (i.e. normal distribution) differs from the topology of the data space with multiple clusters.

In our first experiment, we show that our method is capable of generating new points from a particular cluster given an input point from that cluster. This generation is done in an unsupervised fashion, which is a different setting from the approach of conditional VAEs Sohn et al. (2015) that require training labels. We demonstrate this property on MNIST in Figure 4 , and show that the newly generated points after short diffusion time remain in the equivalent class to the seeded image. Here the architecture is a standard fully convolutional architecture. Details can be found in Appendix A.4. The problem of addressing difference in topologies between the latent space of a generative model and the output data has been acknowledged in recent works about rejection sampling (Azadi et al., 2018; Turner et al., 2018) . Rejection sampling of neural networks consists of generating a large collection of samples using a standard GAN, and then designing a probabilistic algorithm to decide in a post-hoc fashion whether the points were truly in the support of the data distribution p(x).

In the following experiment, we compare to the standard example in the generative model literature.

The data consists of nine bounded spherical densities with significant minimal separation, lying on a 5 × 5 grid. A standard GAN or VAE struggles to avoid generating points in the gaps between these densities, and thus requires the post-sampling rejection analysis. On the other hand, our model creates a latent space that separates each of these clusters into their own features and only generates points that exist in the neighborhood of training data. Figure 5 clearly shows that this results in significantly fewer points generated in the gaps between clusters, as well as eliminating the need to generate additional points that are not in final generated set. Our VDAE architecture here uses one hidden layer of 512 neurons and tanh activations. GAN and DRS-GAN architectures are as described in Azadi et al. (2018) . Here we describe a practical method for computing the local bi-Lipschitz property, then use it to evaluate several methods on the MNIST dataset. Let Z and X be metric spaces and f : Z → X.

We define, for each z ∈ Z and k ∈ N, the function bilip k (z):

bilip k (z) = min K s.t. 1 K ≤ d X (f (z), f (z )) d Z (z, z ) ≤ K, ∀ z ∈ U z,k ∩ Z

where Z := f -1 (X) is the latent embedding of our dataset X5 , d X and d Z are metrics on X and Z, and U z,k is the k-nearest neighborhood of z. Intuitively, increasing values of K can be thought of as an increasing tendency of the learned map to stretch or compress regions of space. By analyzing various statistics of the local bi-Lipschitz measure evaluated at all points of a latent space Z, we can gain insight into how well-behaved a homeomorphism f is. In Table 1 we report the mean and standard deviation, over 10 runs, of the local bi-Lipschitz property for several methods trained on the MNIST dataset.

The comparison is between the Wassertein GAN (WGAN), the VAE, the hyperspherical VAE (SVAE), and our method. We use standard architectures prescribed by their respective papers to train the methods. For our method we use a single 500 unit hidden layer network architecture with ReLU nonlinearities for both the encoder and decoder.

By constraining our latent space to be the diffusion embedding of the data, our method finds a mapping that automatically enjoys the homeomorphic properties of an ideal mapping, and this is reflected in the low values of the local bi-Lipschitz constant. Conversely, other methods do not consider the topology of the data in the prior distribution. This is especially appparent in the VAE and SVAE, which must generate from the entirety of the input distribution X since they minimize a reconstruction loss. Interestingly, the mode collapse tendency of GANs alleviate the pathology of the bi-Lipschitz constant by allowing the GAN to focus on a subset of the distribution -but this comes at the cost, of course, of collapsing to a few modes of the dataset. Our method is able to reconstruct the entirety of X while simultaneously maintaining a low local bi-Lipschitz constant.

A.3 PROOFS:

Proof of Theorem 1. The proof of Theorem 1 is actually a simple extension of the following theorem, Theorem 4, which needs to be proved for each individual extrinsic coordinate X k , hence the additional factor of m coming from the L2 norm of m functions.

Theorem 4. Let M ⊂ R m be a smooth d-dimensional manifold, ψ(M) ⊂ R D be the diffusion map for D ≥ d large enough to have a subset of coordinates that are locally bi-Lipschitz. Let one of the m extrinsic coordinates of the manifold be denoted X(ψ(x)) for x ∈ M. Then there exists a sparsely-connected ReLU network f N , with 4DC M nodes in the first layer, 8dN nodes in the second layer, and 2N nodes in the third layer, such that

X -f N L 2 (ψ(M)) ≤ C ψ √ N (A.7)

where C ψ depends on how sparsely X(ψ(x)) Ui can be represented in terms of the ReLU wavelet frame on each neighborhood U i , and C M on the curvature and dimension of the manifold M.

Proof of Theorem 4. The proof borrows from the main theorem of Shaham et al. (2018a) . We adopt this notation and summarize the changes in the proof here. For a full description of the theory and guarantees for neural networks on manifolds, see Shaham et al. (2018a) . Let C M be the number of neighborhoods U i = B(x i , δ) ∩ M needed to cover M such that ∀x, y ∈ U i , (1 -) x -y ≤ d M (x, y) ≤ (1 + ) x -y . Here, we choose δ = min(δ M , κ -1 ρ) where δ M is the largest δ that preserves locally Euclidean neighborhoods and κ -1 ρ is the smallest value from Jones et al. ( 2008) such that every neighborhood U i has a bi-Lipschitz set of diffusion coordinates.

Because of the locally bi-Lipschitz guarantee from Jones et al. ( 2008), we know for each U i there exists an equivalent neighborhood ψ(U i ) in the diffusion map space, where ψ(x) = [ψ i1 (x), ..., ψ i d (x)]. Note that the choice of these d coordinates depends on the neighborhood U i . Moreover, we know the Euclidean distance on ψ(U i ) is locally bi-Lipschitz w.r.t. d M (•, •) on U i .

First, we note that as in Shaham et al. (2018a) , the first layer of a neural network is capable of using 4D units to select the subset of d coordinates ψ(x) from ψ(x) for x ∈ U i and zeroing out the other D-d coordinates with ReLU bump functions. Then we can define X( ψ(x)) = X(ψ(x)) on x ∈ U i . Now to apply the theorem from Shaham et al. (2018a) , we must establish that X Ui : ψ(U i ) → R can be written efficiently in terms of ReLU functions. Because of the manifold and diffusion metrics being bi-Lipschitz, we know at a minimum that ψ is invertible on ψ(U i ). Because of this invertibility, we will slightly abuse notation and refer to X(ψ(x)) = X(x), where this is understood to be the extrinsic coordinate of the manifold at the point x that cooresponds to ψ(x). we also know that ∀x, y ∈ where ∇X(z) is understood to be the gradient of X(z) at the point z ∈ M. This means X( ψ(x)) is a Lipschitz function w.r.t. ψ(x). Because X( ψ(x)) Lipschitz continuous, it can be approximated by step functions on a ball of radius 2 -to an error that is at most

max z∈U i ∇X(z) 1-

2 -. This means the maximum ReLU wavelet coefficient is less than

max z∈U i ∇X(z) 1- (2 -+ 2 -+1

). This fact, along with the fact that ψ(U i ) is compact, gives the fact that on ψ(U i ), set of ReLU wavelet coefficients is in 1 . And from Shaham et al. (2018a) , if on a local patch the function is expressible in terms of ReLU wavelet coefficients in 1 , then there is an approximation rate of 1 √ N for N ReLU wavelet terms.

:

Any dataset where the data points have been normalized to be unit length becomes a subset of a hypersphere.

2 Importantly, note that the choice of a Gaussian kernel in the latent space is not dependent on the choice of k. We have this invariance due to the aforementioned property of diffusion embeddings.

Though q depends on φ and Θ, we will use q φ := q φ,Θ to be consistent with existing VAE notation and to indicate that Θ is not learned by VI.

For all experiments in Section 6, the number of steps required to draw from π is less than 10.

For VAE, SVAE, and our method, these are the means of the posterior distributions. For GAN it is points drawn from the N (0, 1) prior.

Figures and Tables:

Figure 2: Reconstructed images from the rotating bulldog example plotted in the latent space of VDAE (left), Spherical VAE (SVAE, left-middle) and VAE (right-middle), and GAN (right) on each patch, which allows one to borrow approximation results from the theory of wavelets on spaces of homogeneous type. The proof also crucially uses the bi-Lipschitz property of the diffusion embedding Jones et al. (2008). The key insight of Theorem 1 is that, because of the bi-Lipschitz property, the coordinates of the manifold in the ambient space R m can be thought of as functions of the diffusion coordinates. We show that because each of coordinates function X i is a Lipschitz function, the ReLU wavelet coefficients of X i are necessarily 1 . This allows us to use the existing guarantees ofShaham et al. (2018a)  to complete the desired bound.

Figure 3: An example of distributions reconstructed from a random walk on M Z (via Algorithm 2), given a single seed point drawn from X. (Bottom): An example of a single burst p θ (x|z). The distributions are a loop (left), sphere (middle), and the Stanford bunny (right).

Figure 4: An example of cluster conditional sampling with our method, given a seed point (top left of each image grid). The DVAE is able to produce examples via the random walk that stay approximately within the cluster of the seed point, without any supervised knowledge of the cluster.

Figure 5: Comparison between GAN, DRS-GAN, and our samples on a 5 × 5 Gaussian grid. GAN and DRS-GAN samples taken from Azadi et al. (2018). Shown from left-right are Original, GAN, DRS-GAN, and our method.

U i , |X( ψ(x)) -X( ψ(y))| = |X(x) -

Θ is the deterministic, differentiable function, depending on ψΘ and Cφ , that generates q φ by the reparameterization trick 3(Kingma & Welling, 2013). Take one step of the diffusion random walkx ← arg min y∈A\{x} | ψΘ (y) -z | 2

Table content: <table><tr><td>Algorithm 1 VDAE training</td></tr><tr><td>Θ, φ, θ ← Initialize parameters</td></tr><tr><td>Obtain parameters Θ for the approximate diffusion map ψΘ by Shaham et al. (2018b)</td></tr><tr><td>while not converged do</td></tr></table>

Mean and standard deviation of the local bi-Lipschitz constant, local to k-nearest neighborhoods on MNIST where k = 3, 5, 10, 100, 1000. Lower is better.

Table content: <table><tr><td>Method</td><td>k=3</td><td>5</td><td>10</td><td>100</td><td>1000</td></tr><tr><td>GAN</td><td colspan="2">2.24 ± 0.35 2.32 ± 0.29</td><td>2.5 ± 0.34</td><td>2.8 ± 0.42</td><td>3.11 ± 0.44</td></tr><tr><td>VAE</td><td colspan="4">8.63 ± 2.53 8.93 ± 2.55 9.17 ± 2.59 9.51 ± 2.64</td><td>9.78 ± 2.6</td></tr><tr><td>SVAE</td><td colspan="4">47.2 ± 19.1 50.3 ± 19.0 52.6 ± 19.1 55.2 ± 19.5</td><td>58.5 ± 19.4</td></tr><tr><td>VDAE</td><td colspan="5">1.92 ± 0.5 2.05 ± 0.49 2.22 ± 0.50 2.64 ± 0.52 2.87 ± 0.516</td></tr><tr><td colspan="5">6.4 EMPIRICAL EVALUATION OF THE LOCAL BI-LIPSCHITZ MEASURE</td><td/></tr></table>

Peer review (iclr2020):
The paper proposes a new generative model for unsupervised learning, based on a diffusion random walk principle inspired by the manifold learning literature. The basic idea is to (probabilistically) map points to a latent space, perform a random walk in that space, and then map back to the original space again. Learning of the suitable maps is achieved by casting the problem in a variational inference framework. The paper is generally well-written, and clearly states out its goals and motivation. Sections 1 - 3 in particular give a nice overview of the broader context of the paper, and its aim of borrowing ideas from manifold learning and variational autoencoders. The particular aim on using concepts from manifold learning to avoid mode collapse - corresponding to the underlying homeomorphism losing its bijectivity - is in particular intriguing. The method itself is intuitive at a high level, although I did have some difficulty with Section 4: - in 4.1, one begins by considering the local evidence. This requires drawing a point from U_x, which is defined to a be set. I presume this means one draws uniformly from this set? - Eqn 3 does not apparently have the same structure as Eqn 2. In particular, the first term in (2) is a function of x, but for (3) it is not a function of x*. How does conditioning affect the ELBO? - I was not sure how to interpret the statement that p?(x*|x) ? ?^{-1}(q(z*|x)). Do you mean the distribution is strongly concentrated around this value? Note also an extra *]* in the latter. - Eqn 5 should presumably be an equality? Also, it was not clear what the d in |.|_d^2 means, and why one does not use ||.||^2. - At a higher level, given that x* ~ U_x originally, why do we now draw x* ~ p(x*|x)? - Arriving at 4.2, it was not clear what *The sampling procedure* refers to, i.e., which of the steps in 4.1 it is seeking to specify or augment. It would be useful to clearly lay out the objective function that is being optimised, and how this section fits into that. - In 4.3, it seemed as if the discussion of the neighbourhood reconstruction error would be better placed in 4.1 itself. It appears to be a justification of the already-derived Eqn 5. - Algorithm 2, is there a need to introduce Z_t? It is a bit confusing that, e.g., Z_1 is first written to in iteration 1 by g(Z_0, ?), and then by ?(X_1) in the second iteration. The authors also claim a contribution to be the identification of a principled measure to identify mismatch between latent and data distributions. This *bi-Lipschitz* property is only introduced in Sec 5.1, and the discussion is not too approachable to someone unfamiliar with the area. In particular: - it is not clear how precisely the discussion in this section relates to the VDAE algorithm described in the previous section. - precisely what quantity we are to compute so as to verify this condition remains elusive. The abstract and introduction made me expect that the property is practically verifiable, but it was not clear from this section whether it is so. - the conclusion or key takeaway of this subsection was unclear. I gather that Jones et al. established the existence of a neighbourhood wherein one can define a bi-Lipschitz mapping to R^d for suitable d. But how does this relate to latent and data space mismatch? The experiments show that the proposed method can generate meaningful samples for synthetic manifold data, as well as on the MNIST dataset. I would*ve preferred more discussion of the results in Sec 4.1. I also was hoping for a clearer illustration of mode-collapse problems on standard benchmarks for GANs, with comparison of results to, e.g., those of Wasserstein-GANs (beyond those in Sec 6.1) or other proposals that aim to mitigate mode collapse. Minor comments: - Fig 1, the text in the middle panel is hard to read in black and white. - SpectralNet is mentioned a few times but never formally introduced. - notationally, Section 4 is a little heavy. I would suggest considering to omit the subscript ?*s in the function ? and its inverse. - when mentioning the *reparametrisation trick*, please provide a citation. - what is mathbb{X} in Theorem 2? - *completementary* -> *complementary* - * rejection sampling Azadi et al. (2018); Turner et al. (2018).* -> * rejection sampling (Azadi et al. 2018; Turner et al. 2018).* - some form of Conclusion would be appropriate.