A Geometric Perspective on Variational Autoencoders
Authors: Clément Chadebec
Stephanie Allassonniere
Venue: ICLR 2022 Submitted (2022)
Keywords: Variational Autoencoders
Riemannian geometry

Abstract: In this paper, we propose a geometrical interpretation of the Variational Autoencoder framework. We show that VAEs naturally unveil a Riemannian structure of the learned latent space. Moreover, we show that using these geometrical considerations can significantly improve the generation from the vanilla VAE which can now compete with more advanced VAE models on four benchmark data sets. In particular, we propose a new way to generate samples consisting in sampling from the uniform distribution deriving intrinsically from the Riemannian manifold learned by a VAE. We also stress the proposed method's robustness in the low data regime which is known as very challenging for deep generative models. Finally, we validate the method on a complex neuroimaging data set combining both high dimensional data and low sample sizes.

Under review as a conference paper at ICLR 2022
A GEOMETRIC PERSPECTIVE ON VARIATIONAL AU- TOENCODERS
Anonymous authors Paper under double-blind review
ABSTRACT
In this paper, we propose geometrical interpretation of the Variational Autoen- coder framework. We show that VAEs naturally unveil a Riemannian structure of the learned latent space. Moreover, we show that using these geometrical con- siderations can significantly improve the generation from the vanilla VAE which can now compete with more advanced VAE models on four benchmark datasets. In particular, we propose new way to generate samples consisting in sampling from the uniform distribution deriving intrinsically from the Riemannian manifold learned by VAE. We also stress the proposed method's robustness in the low data regime which is known as very challenging for deep generative models. Finally, complex neuroimaging dataset combining both high
we validate the method dimensional data and low sample sizes.
INTRODUCTION
Variational Autoencoders (VAE) (Kingma Welling 2014 Rezende et al., 2014) are powerful generative models that map complex input data into a much lower dimensional space referred to as the latent space while driving the latent variables to follow a given prior distribution. Their simplicity to use in practice has made them very attractive models to perform various tasks such as high-fidelity image generation (Razavi et al., 2020), speech modeling (Blaauw & Bonada 2016), clustering (Yngg et al. 2019) or data augmentation (Chadebec al.) (021).
Nonetheless, when taken in their simplest version, it was noted that these models produce most of the time blurry samples. This undesired behavior may be due several limitations of the VAE framework. First, the training of a VAE aims at maximizing the Evidence Lower BOund (ELBO) which is only lower bound on the true likelihood and so does not ensure that we are always actually improving the true objective (Burda et al.] 2016 Alemi et al., 2016 Higgins et al.) 2017 Cremer et al., 2018, Zhang et al. 2018). Second, the prior distribution used regularization and for sampling may be too simplistic (Dai & Wipf, 2018) leading to poor data generation and there exists no guarantee that the actual distribution of the latent code will match given prior distribution inducing over-regularization (Connor et al. 202T) Hence, trying to tackle those limitations through richer posterior distributions (Salimans et al. 2015) Rezende & Mohamed] (0155) or better priors Tomczak & Welling, 2018) represents major part of the proposed improvements over the past few years. However, the tractability of the ELBO constrains the choice in distributions and so finding a trade-off between model expressiveness and tractability remains crucial.
In this paper, we take a rather different approach and focus on the geometrical aspects a vanilla VAE is able to capture in its latent space. In particular, propose the following contributions:
We show that VAES unveil naturally a latent space with a structure that can be modeled as a Riemannian manifold through the learned covariance matrices in the posterior distributions.
We propose a natural sampling scheme consisting in sampling from a uniform distribution defined on the learned manifold and given by the metric. We show that this procedure im- proves significantly the generation process from vanilla VAE and makes it able to perform as well as more advanced VAE models in terms of Frechet Inception Distance (Heusel al.) 2017) and Precision and Recall (Sajjadi et al., 2019) scores four benchmark datasets.
We also show that the propose method appears more robust to dataset size changes and outperforms more strongly peers when only smaller sample sizes considered.

---

Under review as a conference paper at ICLR 2022
We validate the method complex neuroimaging data from OASIS Marcus et al.
2007
2 VARIATIONAL. AUTOENCODERS
Considering that we are given x X' set of data points deriving from unknown distribution p(æ), VAE aims at inferring such a distribution with a parametric model (po. 0 E } using a maximum likelihood estimator. A key assumption behind the VAE is to assume that the generation process also involves latent variables z living in a lower dimensional space such that the generative model writes
Sprior (2)
~~ (x(xx)
where Qpeior is a prior distribution over the latent variables often taken as standard Gaussian and po(x|z) referred to the decoder and most of the time taken as parametric distribution the parameters of which are estimated using neural networks. Hence, the likelihood pa writes:
po = Po (x) 2)g(z)dz.
(1)
As this integral is most of the time intractable so is pz(z))), the posterior distribution. Hence, Variational Inference (Jordan et al. 1999) is used and a simple parametrized variational distribution (z/x)) is introduced to approximate the posterior po(z|x). (z]x)) is referred to as the encoder and, in the vanilla VAE, do is chosen as a multivariate Gaussian whose parameters H& and Eg are again given by neural networks. An unbiased estimate of the likelihood p5 (x) can then be derived using importance sampling with qo (z/x) and the ELBO objective follows using Jensen's 's inequality:
(x()) log Ex-qg [Po] 2 Exzgs [log po] = Ex-qg log (x|z) DRI. (un( (xxx(()
(2)
ELBO
The Evidence Lower BOund (ELBO) is now tractable since both pa (a z) and (z(z)x) are parametrized and so can be optimized with respect to the encoder and decoder parameters.
In practice, po (x|z) chosen depending on the modeling of the input data but is often taken simple distribution (e.g multivariate Gaussian, Bernoulli ...). Hence, the ELBO can also be seen as a two terms objective (Ghosh et al. 2020). The first one a reconstruction term given by Po (x(x)) while the second one is regularizer corresponding to the KL divergence between the posterior and the prior. For instance, in the case of a multivariate Gaussian we have
Remark
LREC = ||x po(z)]13
LREG = (xx.(xo(z)x))lll
RELATED WORK
A natural way to improve the generation from VAEs consists in trying to use more complex priors (Hoffman & Johnson 2016) than the standard Gaussian distribution used in the initial version such that they better match the true distribution of the latent codes. For instance, using a Mixture of Gaussian (Nalisnick et al. 2016 Dilokthanakul et al., 2017) or a Variational Mixture Posterior (VAMP) (Tomezak& & Welling 2018) as priors was proposed. In the same vein, hierarchical latent variable models (Sønderby
2016, Klushyn et al. 2019 or prior learning (Chen et al. 2016 Aneja et al. 2020) have recently emerged and aimed at finding the best suited prior distribution for given dataset. Acceptanceerrjeccton sampling method was also proposed to try to improve the expressiveness of the prior distribution (Bauer & Mnih 2019). Some recent works linking energy- based models (EBM) and VAEs (Xiao et al.] (020) or modeling the prior as an EBM (Pang et al.) 2020 have demonstrated promising results and are also worth citing.
On the ground that the latent space must adapt to the data as well, geometry-aware latent space mod- elings as hypershpere (Davidson al.) 2018), torus (Falorsi al. 2018) or Poincaré disk (Mathieu) et al. 2019) or discrete latent representations (Razavi et al.) 2020) were proposed. Other recent con- the latent space Riemannian manifold where the Riemannian metric is given by the Jacobian of the generator function (Arvanitidis et al., 2018 Chen et al., 2018 Shao
tributions proposed
et al. 2018). This metric was then used directly within the prior modeled by Brownian motions
2

---

Under review as a conference paper at ICLR 2022
(Kalatzis et al. 2020 Others proposed to learn the metric directly from the data throughout train- ing thanks to geometry-aware normalizing flows (Chadebec et al. 2020) or learn the latent structure of the data using transport operators (Connor et al. 2021). While these geometry -based methods show interesting properties the learned latent space they either require the computation of a time consuming model-dependent function, the Jacobian, add further parameters to the model to learn the metric or transport operators adding some computational burden to the method.
Arguing that VABS are essentially autoencoders regularized with Gaussian noise, Ghosh et al. (2020) proposed another interesting interpretation of the VAE framework and showed that other types regularization may be of interest as well. Since the generation process from these autoen- coders is no longer relying on the prior distribution, the authors proposed to use ex-post density estimation by fitting simple distributions such Gaussian mixture the latent space. While this paves the way for consideration of other ways to generate data, it mainly reduces the VAE frame- work to an autoencoder while we believe that can also unveil interesting geometrical aspects.
Another widely discussed improvement of the model consists in trying to tweak the approximate posterior the ELBO so that it better matches the true posterior using MCMC methods (Sali- mans et al. 2015) or normalizing flows (Rezende & Mohamed] 2015). For instance, methods using Hamiltonian equations in the flows to target the true posterior (Caterini et al. 2018) were proposed.
Finally, while discussing the potential link between PCA and autoencoders some intuitions arose the impact of both the intrinsic structure of the variance of the data (Rakowski & Lippert (02T) and the shape of covariance matrices in the posterior distributions (Rolinek al. (2019) on disentangle- ment in the latent space. We also believe that these covariance matrices indeed play a crucial role in the modeling of the latent space but in this paper, we instead propose see their inverse as the value of a Riemannian metric evaluated at the embedding points (4;-
4 PROPOSED METHOD
In this section, we argue that vanilla VAE shows naturally a Riemannian structure of the latent space through the learned covariance matrices in the posterior distributions. We then propose new natural generation scheme guided by this estimated geometry and consisting in sampling from a uniform distribution deriving intrinsically from the learned Riemannian manifold.
4.1 A WORD ON RIEMANNIAN GEOMETRY
First, we briefly recall some basic elements of Riemannian geometry needed the rest of the paper. A more detailed discussion integration and probability densities manifolds may be found in Appendix A. A d-dimensional manifold M is a manifold which is locally homcomorphic to a d- dimensional Euclidean space. If the manifold M is further connected and differentiable it possesses a tangent space 7= at any z E .M composed of the tangent vectors of the curves passing by 2. If M equipped with smooth inner product g (-I:)= defined on its tangent space T- for any M then M it is called a Riemannian manifold and is the associated Riemannian metric. Since is an inner product, local representation of g any . is given by the positive definite matrix (z)). The notion of length of curves R + M traveling in A can be defined as follows
I(7) / (1(0))()))))) /
V1(0TG((0))0))).
Curves minimizing L are geodesics and a Riemannian distance between 21, E M can be defined 21 = 70), z2 = 7(1)
dist x2) = inf L(7) s.t.
(4)
The manifold M is said to be geodesically complete if all geodesic curves can be extended to R. In an Fuclidean space, G reduces to the Ia and the distance becomes the classic Euclidean one.
Remark 2 A simple extension of this Euclidean framework consists in assuming that the metric is given by a constant positive definite matrix Σ different from I In such a case the induced Riemannian distance the well-known Mahalanobis distance which writes
dist v (x2 x11 [<> 21).
3

---

Under review as a conference paper at ICLR 2022
4.2 THE RIEMANNIAN GAUSSIAN DISTRIBUTION
The notion of measure and so of probability distribution can be extended geodesically complete Riemannian manifolds as well (Pennec) 2006). Given the Riemannian manifold M endowed with the Riemannian metric G and chart 2, an infinitesimal volume element may be defined on each tangent space T² of the manifold M as follows
dM₃ / G(z)dz,
(5)
with dz being the Lebesgue measure. Hence, a Riemannian Gaussian distribution on M can be defined and consists in using the Riemannian distance defined in Eq. 4instedd of the Euclidean one
exp (- distG 11) R
d.dists d.iist 22)
Nriem (z|o, v)) = C
C
exp
dMz
(6)
20
2a
M
where d.M is the volume element defined by Eq. 5 Hence, the multivariate normal distribution is only a specific case of the Riemannian distribution with a 1, defined on the manifold M R
endowed with the constant Riemannian metric G(z) = Σ E- E..
4.3 GEOMETRICAL INTERPRETATION OF THE VAE FRAMEWORK
Within the VAE framework, the variational distribution (g(z)x) is voluntarily chosen as sim- ple multivariate Gaussian distribution defined on Rd with d being the latent space dimension. Hence, as explained in the previous section, given an input data point æ; the posterior go (z(z))) = ((((z)) E(x:)) can also be seen as a Riemannian Gaussian distribution where the Riemannian distance is simply the distance with respect to the metric tensor Σ (x,), Hence, the VAE frame- work can be seen follows. As with autoencoder, the VAE provides code μl(;;) which is lower dimensional representation of an input data point I, However, it also gives a tensor Σ (z,) depending æ; which can be seen the value of a Riemannian metric G at (xx;)) ie.
G(u(x:)) Σ Σ (x;)
This metric is crucial since it impacts the notion of distance in the latent space seen as the Ric- mannian manifold M = (TR G) and so changes the directions that are favored in the sampling from
the posterior distribution 2o (z(x))). Then, sample is drawn from standard (i.e. a in Eq. ( Riemannian Gaussian distribution and fed to the decoder. As first approximation and since we only have access a finite number of metric tensors Σ Σ-1xx,) , the VAE model assumes that the metric is locally constant close to µı(;;) and so the Riemannian distance reduces to the Mahalanobis distance in the posterior distribution. This simplified drastically the training process since now Riemannian distances have closed form and so are easily computable. Interestingly, the VAE framework will impose through the ELBO expression given in Eq, that z gives a sample x ~ po (x (||z) close to æ; when decoded. Since z has probability density function imposing higher probability for samples having the smallest Riemannian distance to μ. the VAE imposes in a way that latent variables that close in the latent space with respect to the metric G will also provide samples that are close in the data space X' in terms of L2 distance as noticed in Remark. T Noteworthy is that the latter distance can be amended through the choice of the decoder po (xlz). This is an interesting property since it allows the VAE to directly link the learned Riemannian distance in the latent space to the distance in the data space. The regularization term in Eq. 3 ensures that the covariance matrices do not collapse to Og and constraints the latent codes to remain close to the origin easing optimization. Finally, at the end of training, we have a lower dimensional representation of the training data given by the means of the posteriors u(æ,) and a family of metric tensors (G; = Σ E (xx,)) corresponding to the value of a Riemannian metric defined locally on the latent space. Inspired from hommHHuberrg et al. (2012), we propose to build smooth continuous Riemannian metric defined on the entire latent space by performing the following interpolation:
N
w1(2) = eexp (- dists- m.))
G(*) = [ (x,) wx(z) + A Ias
(7)
where distæ (<=; (2, m4) (z mi) (x;) ms) is the Riemannian distance between and f4 with respect to the locally constant metric G(u(xz)) Σ-²(x)). Since the sum in Eq.
4

---

Under review as a conference paper at ICLR 2022
is made on the total number of training samples N, the number of reference metric tensors can be decreased for huge datasets by selecting only h < N metric tensors and increasing to reduce memory usage. We provide an ablation study the impact of X, the number of centroids & and their choice along with a discussion on the choice for p in Appendix G. Then, we have:
Proposition 1 The Riemannian manifold M = (R=", G) is geodesically complete.
Prop. 1 (proved in Appendix B) allows now to refer probability densities M. Rigorously, the metric defined in Eq. should have been used during the training process. Nonetheless, this would have made the training longer and trickier since it would involve i) the computation of Riemannian distances that have no longer close form and so make the resolution of the optimization problem in Eq. [ needed. ii) the sampling from Eq. o which is not trivial and iii) the computation of the regularization term. Instead, by approximating the value of the metric during training by its value at f4 (i.e. (x,)), the VAE training remains unchanged, stable and computationally reasonable since Riemannian Gaussians become multivariate Gaussians in do (z)x). Noteworthy is the fact that, likewise (Ghosh et al.) 2020), our vision of the VAE, the prior distribution is only seen as a regularizer though the KL term and other latent space regularization schemes may have been also envisioned. the following, we keep the proposed vision and do not amend the training process.
4.4 GEOMETRY- AWARE SAMPLING
Assuming that the VAE has learned a latent representation of the data in space seen Rieman- nian manifold, we propose to exploit this strong property to enhance the generation procedure. A natural way sample from such latent space would consist in sampling from the uniform distribu- tion intrinsically defined on the learned manifold. Similar to the Gaussian distribution presented in Sec. 4.2 the notion of uniform distribution can indeed be extended to Riemannian manifolds. Given a bounded set A C M, the uniform distribution writes (Pennec) 2006
(z) Vol(AA)
1A(Z) [LA=))
pa(z)
.M
This density is taken with respect to d.Ms the Riemannian measure but using Eq. and coordinate system z allows to obtain pdf now defined with respect to the Lebesgue measure:
Ustie (z) α det G(z).
Since the Riemannian metric has a closed form expression given by Eq. 7 sampling from this
distribution is quite easy and may be performed using the HMC sampler (Neal (005 for instance. Now we are able to sample from the intrinsic uniform distribution which a natural way of exploring the estimated manifold and the sampling is guided by the geometry of the latent space. A discussion on practical outcomes can be found in Appendix. C.
4.5 ILLUSTRATION ON A TOY DATASET
The usefulness of such sampling procedure may be easily appended in Figure ! where vanilla VAE was trained with a toy dataset composed binary images of disks and rings of different size and thickness (example inspired by Chadebec et al. (02T)). On the left is presented the learned latent space along with the embedded training points given by the colored dots. The log of the metric volume element is given in gray scale. In this example, we clearly see a geometrical structure appearing since the disks and rings seem to wrap around each other. Obviously, sampling using the prior (taken (0, Ia)) in such a case is far from being optimal since the sampling will be performed regardless of the underlying distribution of the latent variables and so will create irrelevant samples. To further illustrate this, we propose interpolate between points in the latent space using different cost functions. Dashed lines represent affine interpolations while the solid ones show interpolation aiming at minimizing the potential (zz) = (ddet G(z)) all along the curve i.e. solving the minimization problem
inf 7
(()()) s.t. 7(0) 21, 711) 22
(8)
'This may he performed with A- k-medoids algorithm for instance.
5

---

Under review as a conference paper at ICLR 2022
Gadd Keat
Affine
O
(1)
Ours
A Afine
Ours
Figure I: Top: Visualization and interpolation in 2D latent space learned by vanilla VAE trained with binary images of rings and disks. The log of the metric volume element v det G(z) (also proportional to the log of the density we propose to sample from) is represented in gray scale. Top right: The Riemannian distance from starting point is presented with the color maps. The dashed lines are affine interpolations between two points in the latent space and the solid ones are obtained Bottom: Decoded samples along the interpolation curves.
by solving Eq.
Below are presented the decoded samples all along the interpolation curves. Thanks to those interpo- lations see that i) the latent space seems to really have specific geometrical structure since decoding all along the interpolation curves obtained by solving Eq.Bleadd to qualitatively satisfying results, ii) certain locations of the latent space must be avoided since sampling there will produce irrelevant samples (see red frames and corresponding red dashes). Using the proposed sampling scheme will allow to sample in the white areas and so ensure that the sampling remains close the data iee where information is available and so does not produce irrelevant images when decoded.
EXPERIMENTS
In this section, we conduct a comparison with other VAE models using other regularization schemes, more complex priors, richer posteriors, ex-post density estimation trying to take into account geometrical aspects including our method. In the following and to ensure a fair comparison, all the models share the same auto-encoding neural network architectures described in Appendix E.
5.1 GENERATION WITH BENCHMARK DATASETS
First, we compare the proposed sampling method to several VAE variants such as a Wasserstein Autoencoder (WAE) (Tolstikhin al. 2018), Regularized Autoencoders (Ghosh al. 2020 with either 1.2 decoder's parameters regulard larzzaton (RAE-L2), gradient penalty (RAE-GP), spectral nor- malization (RAE-SN) or simple L2 latent code regularization (RAE). a vamp-prios VAE (VAMP) (Tomezak & Welling 2018), a Hamiltonian VAE (HVAE) (Caterini et al.) 2018), a geometry-awant VAE (RHVAE) (Chadebec al. (2200) and Autoeneoder (AE). We elect these models since they use different ways to generate the data using either the prior or ex ex-post density estimation. For the latter, we use the approach of Ghosh et al. (2020) and fit a 10-component mixture of Gaussian in the latent space after training. The models
trained on MNIST (LeCun) 1998). SVHN (Netzer) et al. 2011), CIFAR 10 (Krizhevsky et al. 2009) and CELEBA (Liu et al. 2015) and we keep the model achieving the best validation loss. See Appendix for the comprehensive experimental set- up. Figure 2 shows a qualitative comparison between the resulting generated samples for MNIST and CELEBA, the same plots made available in Appendix D for SVHN and CIFAR 10. Interest-

---

Under review as a conference paper at ICLR 2022
MNIST
CELEBA
AE-
Q
7
9
3
5
1
4
VAB-N
0
3
6
g
2 8
7
Q
1
0
3
WAE
VAMP
8
3
9
6
0
4
5 7
HVAE
9
0
8
4
0
6
RHVAE
4
8
/
2
9
6 4 3
9
AE-GMM
3
0
6
5
6
8
VAB GMM
2
2
0
5
8
5
RAE
2
6
4
8
4
9
VAB Ours
5
/
7
2
3
8
4
Figure 2: Generated samples with different models and generation processes. with RAE variants are also provided in Appendix D.
Generated samples
Gen.
Near, train Near. rec.
Gon.
Near, train Near,
Gen.
Near. train Near,
Gen.
Near, train Near, roc.
5
5
5
9
9
9
3
3
3
2
2
2
pe
Figure 3: Nearest train image (near. train) and nearest image in all reconstructions of train images (near. recon) to the generated one (Gen.) with the proposed method. Note: the nearest reconstruction may be different from the reconstruction of the nearest train image.
ingly, using the non- non-prior based methods seems to produce qualitatively better samples (rows 7 to end). Nonetheless, the resulting samples seem even sharper when the sampling takes into account geometrical aspect of the latent space as we propose (last row). Additionally, even though the exact same model is used, clearly see that using the proposed method represents strong improvement of the generation process from a vanilla VAE when compared to the samples coming from a normal the simplest VAE model actually contains lot
prior (second row). This insists on the fact that of information in its latent space but the limited expressiveness of the prior impedes to access to it.
Hence, using more complex prior such as the VAMP may be tempting idea. However, one must keep in mind that the ELBO objective in Eq. 2 must remain tractable and so using more expressive priors may impossible. These observations are even more supported by Table where we report the Frechet Inception Distance (FID) and the precision and recall (PRD) score against the test set to assess the sampling quality and diversity. Again, fitting a mixture of Gaussian (GMM) in the latent space appears to be interesting idea since it allows for better expressiveness and latent space prospecting. For instance, on MNIST the FID falls from 40.7 with the prior to 13.1 when using GMM. Nonetheless, with the proposed method we able to make it even smaller (8.5) and PRD scores higher without changing the model and performing post processing. This can also be observed the other datasets. Impressively, in almost all cases, the proposed generation method

---

Under review as a conference paper at ICLR 2022
FID (lower is better) and PRD score (higher is better) for different models and datasets. In the first section are presented results using the prior distribution while in the second one, we use ex- ex-post density estimation by fitting a 10-component mixture of Gaussian in the latent space.
Table 1:
MNIST (16)
SVHN (16)
CIFAR 10 (32)
Celeba (64)
Model
FID J
PRD +
FID L
PRD f
FID L
PRD t
FID
PRD t
AE (O, I)
46.41
0.86/0.77
119.65
0.54/0.37
196.50
0.05/0.17
64.64
0.29/0.42
WAE
20.71
0,93/0.88
49.07
0,80/0.85
132.99
0.24/0.52
54.56
0.57/0.55
VAB- (O,-N((,)) I)
40.70
0,83/0.75
83.55
0.69/0.55
162.58
0.10/0.32
64.13
0.27/0.39
VAMP
34.02
0.83/0.88
91.98
0,55/0.63
198, 14
0.05/0.11
73.87
0.09/0.10
HVAE
15.54
0,97/0.95
98.
0.64/0.68
201.70
0.13/0.21
52.00
0.38/0.58
RHVAE
36.51
0,73/0.28
121. .699
0,55/0.41
167.41
0.12/0.22
55.12
0,45/0.56
AE GMM
9.60
0,9570.90
54.21
0.82/0.83
130.28
0.35/0.58
56.07
0.32/0.48
RAE (GP)
9.44
0,97/0.98
61.43
0.79/0.78
120.32
0.34/0.58
59.41
0.28/0.49
RAE (L2)
9.89
0,97/0.98
58.32
0.82/0.79
123.25
0.33/0.54
54.45
0.35/0.55
RAE (SN)
11.22
0.97/0.98
95.64
0.53/0.63
114.59
0.32/0.53
55.04
0.36/0.56
RAE
11.23
0.98/0.98
66.20
0.76/0.80
118.25
0.35/0.57
53,29
0.36/0.58
VAE- GMM
13.13
0.95/0.92
52.32
0.82/0.85
138.25
0.29/0.53
55,50
0.37/0.49
VAE Ours
8,53
0,98/0.97
46,99
0.84/0.85
93.53
0.71/0.68
48.71
0.44/0.62
can either compete outperform peers both in terms of FID and PRD scores. Finally, we check the proposed method does not overfit the training data and is able to produce diverse samples by showing the nearest neighbor in the train set and the nearest image in all the reconstructions of the train images to a generated image in Figure 3 This experiment shows that the generated samples are only resampled train images and that the sampling prospects quite well the manifold. To support even more this claim we provide in Appendix G an analysis in a case where only two centroids are selected in the metric. This also shows that the generated samples are not only an interpolation be- tween the k selected centroids since some generated images contain attributes that are not present in the images of the decoded centroids. The outcome of such an experiment is that using post training latent space processing such as ex-post density estimation or adding some geometrical considera- tion to the model allows to strongly improve the sampling without adding more complexity to the model. Generating 1000 samples on CELEBA takes approx. 5.5 min for our method vs. 4 min for compontttt GMM and 0.6s for prior based methods on single GPU V100-16GB.
MNIST
CIFAR 10
300
WAE-XIB
KAB-GP-AT.
200
RAB-12-NOLD D
100
VAE X/D. D)
50
150
PAE-GMM
KAB-GP-CMO
RAP-12 GMM
20
RAR-OU-GMA GMM
VAE GMM
10
100
0.1
0.5
1.0
5.0
0.1
0.5
1.0
5.0
x
X10*
Figure 4: Evolution of the FID score according to the number of training samples.
5.2
INVESTIGATING GENERATION ROBUSTNESS IN Low DATA REGIME
We perform a comparison using the same models and datasets as before but we decide to progres- sively decrease the size of the training set to see the robustness of the different sampling methods according to the number of samples. This experiment rarely performed in most generative models related papers even though it is well known that such a context may be very challenging for these models. Nonetheless, it appears to us very important since in day-to-day applications collecting such large databases may reveal costly if not impossible (think of medicine for instance). Hence, we consider MNIST, CIFARIO and SVHN and use either the full dataset size, 10k, 5k or 1k training samples. For each experiment, the best retained model is again the one achieving the best BLBO
8

---

Under review as a conference paper at ICLR 2022
Table 2:
Classification results averaged on 20 independent runs. For the generative models, the classifier is trained on 2K generated samples per class.
Generation method
Balanced Accuracy
AD
CN
Original (unbalanced)
66.2 7.6
47.6 = 15.8
87.3
Original (resampled)
81,8 1 2.6
72.1 3.6
88.0 1500::0
AE-N [o Fa)
50.0 0.0
0.0 0.0
84.1 0.0
WAB
5774
21.04 24.5
84.4
(0,
5188 3.8
6.1 118
84.6 1.1
VAMP
83.1 8 2.6
7:
82.2 4.7
HVAE
56.3 7.9
1966
85.4
RHVAE
68.0 10.9
47.0=
85.1 3.3
AE- GMM
8,
69.3 3.1
82.0 3.6
RAE (GP)
63
46.5 = 15,9
70.6 + 19.6
RAE (L2)
6.0
60.6 9.5
82.1 5.9
RAF
62.3
378 22.6
80.1
7.9
RAE
69.3 8.1
53.8 = 12,9
80.0 + 10.7
83.0
4.
VAE Ours
85.4
7..7
87.3 1.
on the validation set the size of which is set 20% of the train size. See Appendix E for further details about experiments set-up. Then, we report the evolution of the FID against the test set in Figure 4 Results obtained on SVHN method appears quite robust to the dataset size since it outperforms the other models' FID even when the number of training samples is smaller. This is made possible thanks to the proposed metric that allows to avoid regions of the latent space having poor information. Finally, our study shows that although using more complex generation procedures such as ex- ex-post density estimation seems to still enhance the generation capability of the model when the number of training samples remains quite high (25k), this gain seems to worsen when the dataset size reduces illustrated on CIFAR.
presented in Appendix F. Again, the proposed sampling
5.3 GENERATION WITH COMPLEX DATA
Finally, we also propose to stress the proposed generation procedure in day-to-day scenario where the limited data regime more than common. To stress the model in such condition, we consider the publicly available OASIS database composed of 416 MRI of patients, 100 of whom were diagnosed with Alzheimer disease (AD). Since both FID and PRD scores are not relevant and reliable in such low data regime due the lack of large test set, propose to assess quantitatively the generation quality with a data augmentation task. Hence, we split the dataset into a train set (70%), a validation set (10%) and a test set (20%). Each model trained on each label of the train set and used to generate 2k new samples per class. Then simple CNN classifier is trained on i) the original train set and ii) the 4k generated samples from the generative models and tested on the test set. Table P shows the mean balanced accuracy and FI scores across 20 runs. These metrics provide good way to assess i) if the generative model can generate data that are not too far from the test set and add information to the data that is relevant for classification and ii) allows to quantify the amount of overfitting. The proposed method the only one to be able to outperform the original (unbalanced) data both in terms of balanced accuracy and F1 scores for both labels meaning that generated samples are relevant to the classifier. This is also the sign of a good generalization power since the classifier achieves classification results that outperform the one observed on the original data.
CONCLUSION
In this paper, we provided geometric understanding of the latent space learned by VAE and showed that it can actually be seen as a Riemannian manifold. Then, we proposed a new natural generation process consisting in sampling from the intrinsic uniform distribution defined on this learned manifold. It showed to be competitive with more advanced versions of the VAEs using either more complex priors, ex ex-post density estimation, normalizing flows or other regularization schemes. Interestingly, the proposed method revealed good robustness properties in complex settings such as high dimensional data or low sample sizes. Future work would consist in trying to use this method to perform data augmentation in those challenging contexts and compare its reliability for such task with state of the art augmentation methods.

---

Under review as a conference paper at ICLR 2022
REPRODUCIBILITY STATEMENT
In order to make the method and the proposed experiments reproducible, we provide in Appendix the complete experimental set-up and in Appendix C pseudo-code algorithms detailing the imple- mentation from a practical point of view. We also provide an implementation in the supplementary material.
10

---

Under review as a conference paper at ICLR 2022
REFERENCES
Alexander A Alemi, Jan Fischer, Joshua V Dillon, and Kevin Murphy, Deep variational information bottleneck. arXiv preprint arXiv: www:::012400000 2016.
Jyoti Aneja, Alexander Schwing, Jan Kautz, and Arash Vahdat. NCP- VAE: Variational autoencoders with noise contrastive priors. arXi::2010.02917 [cs, stat), 2020.
Georgios Arvanitidis, Lars Kai Hansen, and Sören Hauberg. Latent space oddity: On the curvature of deep generative models. In 6th International Conference Learning Representations, ICLR 2018, 2018
Matthias Bauer and Andriy Mnih. Resampled priors for variational autoencoders. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 66-75. PMLR, 2019.
Merlijn Blaauw and Jordi Bonada. Modeling and transforming speech using variational autoen- coders. Morgan N, editor. Interspeech 2016; 2016 Sep 8-12; San Francisco, CA.[place un- known]: ISCA; 2016. p. 1770-4., 2016. Publisher: International Speech Communication As- sociation (ISCA).
Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance weighted autocncoders. 550,00019 [cs, stat], 2016.
Manfredo Perdigao do Carmo. Riemannian Geometry. Birkhäuser, 1992.
Anthony Caterini, Arnaud Doucet, and Dino Sejdinovic. Hamiltonian variational auto-encoder. In Advances in Neural Information Processing Systems, pp. 8167-8177, 2018.
Clément Chadebec, Clément Mantoux, and Stéphanie Allassonnière. Geometry-aware hamiltonian variational auto-encoder. arXiv:2010. 18888 2020.
Clément Chadebec, Elina Thibeau-Sutre, Ninon Burgos, and Stéphanie Allassonnière. Data Aug- mentation in High Dimensional Low Sample Size Setting Using a Geometry-Basee Variational Autoencoder. arXiv preprint arXiv:2 2021.
Nutan Chen, Alexej Klushyn, Richard Kurle, Xueyan Jiang, Justin Bayer, and Patrick Smagt. Met- rics for deep generative models. In International Conference on Artificial Intelligence and Statis- fics, pp. 1540-1550. PMLR, 2018.
Xi Chen, Diederik P Kingma, Tim Salimans, Yan Duan, Prafulla Dhariwal, John Schulman, Ilya Sutskever, and Pieter Abbeel. Variational lossy autoencoder. arXiv preprint arXiv:1611.0233. 2016.
Marissa Connor, Gregory Canal, and Christopher Rozell. Variational autoencoder with learned latent structure. In International Conference on Artificial Intelligence and Statistics, pp. 2359-2367, PMLR, 2021.
Chris Cremer, Xuechen Li, and David Duvenaud. Inference suboptimality in variational autoen- coders. International Conference on Machine Learning, pp. 1078-1086. PMIR, 2018
Bin Dai and David Wipf. Diagnosing and Enhancing VAE Models. In International Conference on Learning Representations, 2018.
Tim R Davidson, Luca Falorsi, Nicola De Cao, Thomas Kipf, and Jakub M Tomezak. Hyperspher- ical variational auto-encoders. In 34th Conference on Uncertainty in Artificial Intelligence 2018, UAI 2018, pp. 856-865. Association For Uncertainty in Artificial Intelligence (AUAI), 2018.
Nat Dilokthanakul, Pedro A. M. Mediano, Marta Gamelo, Matthew C. H. Lee, Hugh Salimbeni, Kai Arulkumaran, and Murray Shanahan. Deep unsupervised clustering with gaussian mixture variational autoencoders. urXiv:1611.02660 [cs, stat), 2017.
Simon Duane, Anthony D Kennedy, Brian J Pendleton, and Duncan Roweth. Hybrid monte carlo. Physics Letters B, 195(2):216-222, 1987.
11

---

Under review as a conference paper at ICLR 2022
Luca Falorsi, Pim de Haan, Tim R. Davidson, Nicola De Cao, Maurice Weiler, Patrick Forré, and Taco S. Cohen. Explorations homeomorphic variational auto-encoding. 1.X.v:8877.04688 [cs, stat], 2018.
Partha Ghosh, Mehdi SM Sajjadi, Antonio Vergari, Michael Black, and Bernhard Schölkopf. From variational to deterministic autoencoders. In 8th International Conference on Learning Represen- tations, ICLR 2020, 2020.
Mark Girolami and Ben Calderhead. Riemann manifold langevin and hamiltonian monte carlo methods. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 73(2): 123-214, 2011.
Søren Hauberg, Oren Freifeld, and Michael Black. ric Learning. In F. Pereira, C. J. C. Burges, L. Bottou, and K. Q. Weinberger (eds.), Advances in Neural Information Processing Systems, volume 25. Curran Asso- ciates, Inc., 2012. URL https:://procceednggs Paper...
Geometric take on Met-
caasscc//pe cc/paperr2012/file
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge local nash equilibrium. In Advances in Neural Information Processing Systems, 2017.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-VAE: Learning basic visual concepts with constrained variational framework. ICLR, 2(5))6, 2017.
Matthew D Hoffman and Matthew J Johnson. Elbo surgery: yet another way carve up the varia- tional evidence lower bound. In Workshop Advances in Approximate Bayesian Inference, NIPS, volume pp. 2, 2016.
Michael Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul. An introduction to variational methods for graphical models. Machine Learning, 37(2):183-233, 1999.
Dimitrios Kalatzis, David Eklund, Georgios Arvanitidis, and Soren Hauberg. Variational autoen- coders with riemannian brownian motion priors. In International Conference on Machine Learn- ing. pp. 5053-5066. PMLR, 2020.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.69555 2014.
Diederik P. Kingma and Max Welling, Auto-encoding variational bayes. arXiv:1312.6114 [cs, stat], 2014.
Alexej Klushyn, Nutan Chen, Richard Kurle, and Botond Cscke. Learning Hicrarchical Priors in VAEs. Advances in neural information processing systems, pp. 10, 2019.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Yann LeCun. The MNIST database handwritten digits. 1998.
Jun S Liu. Monte Carlo strategies in scientific computing. Springer Science Business Media, 2008.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
Maxime Louis Computational and statistical methods for trajectory analysis in Riemannian geometry setting. PhD Thesis, Sorbonnes universités, 2019.
Daniel S. Marcus, Tracy H. Wang. Jamie Parker, John G. Csemannsky John C. Morris, and Randy L. Buckner. Open access series of imaging studies (OASIS): Fross-sectional MRI data in young, middle aged, nondemented, and demented older adults. Journal of Cognitive Neuroscience, 19 498-449815577 2007.
12

---

Under review as a conference paper at ICLR 2022
Emile Mathieu, Charline Le Lan, Chris J Maddison, Ryota Tomioka, and Yee Whye Teh. Contin- uous hierarchical representations with poincaré variational auto-encoders. In Advances in neural information processing systems, pp. 12565-12576, 2019.
Eric Nalisnick, Lars Hertel, and Padhraic Smyth. mixtures. In NIPS Workshop on Bayesian Deep Learning, volume 2, pp. 131, 2016.
Approximate inference for deep latent gaussian
Radford M Neal. Hamiltonian importance sampling. In talk presented at the Banff International Research Station (BIRS) workshop on Mathematical Issues in Molecular Dynamics, 2005.
Radford M Neal and others. MCMC using hamiltonian dynamics. Handbook of Markov Chain
Monte Carlo, 2(11):2, 2011.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng. Reading digits in natural images with unsupervised feature learning. 2011.
Bo Pang, Tian Han, Erik Nijkamp, Song-Chun Zhu, and Ying Nian Wu. Learning latent space energy-based prior model. Advances in Neural Information Processing Systems, 33, 2020.
Xavier Pennec.
Intrinsic statistics on riemannian manifolds: Basic tools for geometric measure- ments. Journal Mathematical Imaging and Vision, 25(1):127-154, 2006. ISSN 0924-9907, 1573-7683. doi: 00.1007/s18851-000 6228-44
Alexander Rakowski and Christoph Lippert. Disentanglement and local directions of variance. In Joint European Conference on Machine Learning and Knowledge Discovery Databases, pp- 19-34. Springer, 2021.
Ali Razavi, Aaron van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq-vae-2. Advances in Neural Information Processing Systems, 2020.
Danilo Rezende and Shakir Mohamed, Variational inference with normalizing flows. In Interna- tional Conference on Machine Learning, pp. 1530-1538. PMLR, 2015.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and ap- proximate inference in deep generative models. In International conference on machine learning, pp. 1278-1286 PMLR, 2014.
Michal Rolinek, Dominik Zietlow, and Georg Martius. Variational autoencoders pursue pea direc- tions (by accident). In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12406- 12415, 2019.
MSM Sajjadi, O Bachem, M Lucic, O Bousquet, and S Gelly. Assessing generative models via precision and recall. 32nd Conference on Neural Information Processing Systems (NeurIPS 2018), pp. $228-5237,200
Tim Salimans. Diederik Kingma, and Max Welling. Markov chain monte carlo and variational inference: Bridging the gap. In International Conference on Machine Learning, pp. 1218-1226, 2015.
Hang Shao, Abhishek Kumar, and P. Thomas Fletcher. The riemannian geometry of deep generative models. In 2018 IEEE/CVF Conference Computer Vision and Pattern Recognition Workshops (CVPRW), pp. 2288-428. IEEE, 2018. ISBN 978-1-5386-6100-0. doi: 10.1109//sPR.2211 00071.
Casper Kaae Sønderby, Tapani Raiko, Lars Maaløe, Søren Kaae Sønderby, and Ole Winther. Ladder variational autoencoder. In 29th Annual Conference on Neural Information Processing Systems (NIPS 2016), 2016.
I Tolstikkhin, O Bousquet, S Gelly, and B Schölkopf. Wasserstein auto-encoders. In 6th International Conference on Learning Representations (ICLR 2018), 2018.
Jakub Tomezak and Max Welling. Vae with a vampprior. In International Conference on Artificial Intelligence and Statistics, pp. 1214-1223. PMLR, 2018.
13

---

Under review as a conference paper at ICLR 2022
Zhisheng Xiao, Karsten Kreis, Jan Kautz, and Arash Vahdat. Vaebm: A symbiosis between varia- tional autoencoders and energy-based models. In International Conference on Learning Repre- sentations, 2020.
Linxiao Yang, Ngai-Man Cheung, Jiaying Li, and Jun Fang. Deep clustering by gaussian mixture variational autoencoders with graph embedding. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 6440-6444 2019.
Cheng Zhang, Judith Bütepage, Hedvig Kjellström, and Stephan Mandt. Advances in variational inference. IEEE Transactions on Pattern Analysis and Machine Intelligence, 41(8):2008-200 2018.
14

---

Under review as a conference paper at ICLR 2022
A FURTHER ELEMENTS ON RIEMANNIAN GEOMETRY
In the field of differential geometry, a Riemannian manifold M can be defined as a connected and differentiable manifold endowed with a Riemannian metric g. The metric g is smooth inner product : p + <-I-), defined on each tangent space TyM of the manifold with p E M.. A chart (or coordinate system) (U,z) is homeomorphism mapping an open set U of the manifold to an open set V of an Euclidean space. The manifold M is further called a d dd-dimensional manifold if for each chart V Rd. This means that there exists neighborhood U of each point p E M such that U is homeomorphic to R.. Given p U, chart Φ (x'....,,
the tangent space M. Hence, the metric of a Riemannian manifold can be locally represented in the chart @ as positive definite matrix (p) (.s.))..ss..ja af 8zo )pooss.sa for each point p of the manifold. That for v, w E T.M and E M, the inner product writes (u)w))p = u Gp))..
....., x²) induces a basis a........ 3xd
0
Two ways of apprehending manifolds exist. The first one is the extrinsic view and it assumes that the manifold is embedded within a higher dimensional Euclidean space. A simple example is the 2-dimensional sphere S seen subspace of R The second one which is adopted in this paper the intrinsic view. In the intrinsic view the manifold is studied using its underlying structure and so the length of curve Y + .M traveling in the manifold cannot be interpreted using the Euclidean distance but requires to use the metric defined onto the manifold itself. Let ž1, E M be two points of the manifold, and y be a curve traveling in M parametrized by / [0, 1] such that 7(0)= 21 and 7(1) = 22. Then, the length of y is given by
[(7) /
/ (ay
1500)hoodd
Curves that are minimizer(s) of this criteria are called geodesic curves. A distance dist on the manifold M can then be derived and writes:
dist(21, 22) min ((7) s.t. 7(0) 21, 21, 7(1) = 22
The manifold M is said to be geodesically complete if all geodesic curves can be extended to R. For any p E M, the exponential map at p maps vector of the tangent space T,MM point of the manifold ğ E M such that the geodesic starting at with initial velocity v reached ğ at time 1.
+ -M T
Expp
- (x), (1)
where ") (1) means 70) = and 7(0) v. Saying that the manifold M is geodesically complete means that the exponential maps is defined on the entire tangent space T.M for each element E M.
In Pennec (0066), the author discusses the statistical framework that can be developed on geodesi- cally complete manifolds using an intrinsic point of view. In particular, given a positively oriented Riemannian manifold M and a chart Φ (II,--. ,T,), a volume form Vol, can be defined the d-form:
dVol, (ee dx A... A drd
This represents an infinitesimal volume element on each tangent space and so measure on the manifold M
(M) del g(0-(((())))))) dr It follows that we are able to integrate functions / M 3 U + R on a given chart (U,
/ [())aM))) = / f(o
dx dxd
d(U)
This notion then be extended to the whole manifold M using partition of unity. In particular, such a property allows us to define probability distributions whose density is defined with respect to the measure the manifold. We recall such definition from Pennec (0066) below
15

---

Under review as a conference paper at ICLR 2022
Definition Let ((M) be the Borel o-algebra of M. The random point p has probability density function Pp if:
/
AA B(M), P(p X)
pl)dM(p)
x
f p(p)dM(p)) M
and
Finally, given a chart Φ (2 .n) defined on the whole manifold M and a random point p on M, the point <(p) is random point whose density P may be written with respect to the Lebesgue measure as such (Pennee 2006):
(2(z) = (z)) det (z))
(10)
16

---

Under review as a conference paper at ICLR 2022
B PROOF OF PROP. 1
We adapt the proof in Louis (2019) and Chadebec et al. (2021) to our specific metric. We will show that given the manifold M E2 and the Riemannian metric whose local representation given by Eq. 7 any geodesic curve :)a, b[-> M is actually extensible to R. Let us consider a geodesic curve y such that y cannot be extended to There exists / =|a, bf such that / the maximum definition domain of %. We will show that with such an assumption we will end up with contradiction. We recall the shape of the Riemannian metric:
N G(z) = ΣE wi(z) X Iai
Since Σ; are positive definite matrices we have z E > 0, = E M - (0}. We further have (x) > V E .M since the manifold is geodesically complete. Let to Ela, bl we therefore have for any & EJa, b(.
A 1+(t) s A (1(())12 + Σ E(z)) (()) (:(r(t))
E(x;)
S 1000 = (l5(4)) ( (12()
where the last equality comes for the constant speed of geodesic curves (Carmo 1992). Hence we have:
(l:(1) - 7 [I((to)llx(to)
14-1p 4ol
VA
This means that for any Ela, bf the geodesic curve 7 remains within compact set. show now that the curve can actually be extended. Let us define the sequence In -> b. Since the geodesic
curves have constant speed the set 1 {(I:: )(x)))meN is compact. Moreover, using Cauchy- Lipchitz theorem, we can find > 0 such that for any n E N, the geodesic y can be extended to It, tn =[. Since, in can be close as b as desired can assure that the curve definition domain can be extended to Ja, + [
17

---

Under review as a conference paper at ICLR 2022
C THE GENERATION PROCESS ALGORITHM- IMPLEMENTATION D DETAILS
In this appendix, we provide pseudo-code algorithms explaining how to build the metric from trained VAE and how use the proposed sampling process. Noteworthy is the fact that do amend the training process of the vanilla VAE which remains pretty simple and stable.
C.I BUILDING THE METRIC
In this section, we explain how to build the proposed Riemannian metric. For the sake of clarity, we recall the expression of the metric in Eq. below
G(2)
Σ m (ax) wi(z) X Ia.
where
(a) 14:)²
(z u:) (x;) - mi)
wi(2) exp
=exp
exp
p²
Algorithm Building the Metric from a Trained Model
Input: A trained VAE model m, the training dataset X, À
for 2; E X do
p4, Σ, m(x))
> Retrieve training embeddings and covariance matrices
end for
Select & centroids c; in the μ4;
e.g. with k-medoids
Get corresponding covariance matrices E;
max min ||c, cyll2
b Set p the max distance between two closest neighbors
juft Build the metric using Eq.
N
G(z) = Σ Σ
wi(z) + - Ia
1=1
Return
D Return G as a function
As is standard VAE implementations, we assume that the covariance matrices E, given by the VAE are diagonal and that the encoder outputs mean vector and the log of the diagonal coefficients. In the implementation, the exponential is then applied to recover the Σ; so that no singular matrix arises.
C.2 SAMPLING PROCESS
Further to the description performed in the paper, we provide here detailed algorithm stating the main steps of the generation process.
C.2.1 THE HMC SAMPLER
In the sampling process propose to rely the Hamiltonian Monte Carlo sampler sample from the Riemanian uniform distribution. In a nutshell, the HMC sampler aims at sampling from a target distribution Ptarget(2) (2) with Rd using Hamiltonian dynamies. The main idea behind such sampler to introduce an auxiliary random variable v ~ AO Ta) independent from z and mimic the behavior of a particle having z (resp. v) as location (resp. velocity). The Hamiltonian of the particle then writes
11 (1, v) = ((z) + K(v).
where (2)) is the potential energy of such particle and K K(v) is its kinetic energy both given by
U(x) = - log Ptarget (z), K(v) 2 T
18

---

Under review as a conference paper at ICLR 2022
The following Hamilton's s equations govern the evolution in time of the particle.
(M(=.0) V V
(11)
0H (=.0)
log Ptarget (z)
8z
In order to integrate these equations, recourse to the leapfrog integrator is needed and consists in applying my times the following equations.
v(t + ) =v(l) + 2 Vx log Prarget (=()),
x = x(t) + Eur v(t 50 )
(12)
v(l Eir)
=v(1 + ) fu V= log Prarget (z(1 + Eir))
where Eir is called the leapfrog step size. This algorithm produces a proposal (z, 8) that is accepted with probability a where
min (1, (1(
This procedure is then repeated to create ergodic Markov chain (= converging the distribution
Ptarget (Neal & others 2011 Duane et al. 1987 Liu 2008 Girolami & Calderhead 2011
C.3 THE PROPOSED ALGORITHM
our setting the target density is given by the density of the Riemannian uniform distribution which writes with respect to Lebesgue measure as follows
p(z)
URsem (2) = C v det Giz).
(13)
The log density follows
log p(z) 2
log det G(z) log C
In such case, the Hamiltonian writes
1 Hzz,v) = =z)gp(()) + 200 D.v.
and Hamilton's 's equations become
(H((..)
de
(H(=.e) Us.
0 logg(z)
-ttr(( ( (z)
(G(=) 8z.
Os.
Since the covariance matrices are supposed to be diagonal is standard in VAE implementations, the computation of the inverse metric is straightforward. Moreover, since Gzz) is smooth and has closed form, it can be differentiated with respect given in Eq. 12 can be used and the acceptance ratio a is easy to compute. Noteworthy is the fact that the normalizing constant C' is never needed since it vanishes in the gradient computation and simplifies in the acceptance ratio a. We provide pseudo-code of the proposed sampling procedure in Alg. 2 typical choice in the sampler's 's hyper-parameters used the paper is N 100, ni = 10 and 0.01. The initialization of the chain can be done either randomly or points that belong to the manifold (i.e. the centroids e; or μ4).
pretty easily. Now, the leapfrog integrator
19

---

Under review as a conference paper at ICLR 2022
Algorithm 2 Proposed Sampling Process
Input: The metric function G, hyper-parameters of the HMC sampler (chain length N, number of leapfrog steps nir, leapfrog step size ěhF)
Initialization: z
b Initialize the chain
for i 1 N do
Id)
> Draw velocity
Ho - Hzz
Compute the starting Hamiltonian
Zo <-
for
my do ú f- EK (z,(zz
+ = + Elr
Leapfrog step Eq. 12
V.H(ž,0)
v ut D
end for
H H22
Compute the ending Hamiltonian
Accept ž with probability min min(1, ex, H))
if Accepted then
Z
else
20
end if
end for
Return z
20

---

Under review as a conference paper at ICLR 2022
D OTHER GENERATION
D.1 SOME FURTHER SAMPLES ON CELEBA AND MNIST
In this section, we provide some further generated samples using the proposed method. Figure and Figure 6 again support the fact that the method is able to generate sharp and diverse samples. We also add the other variants of the RAE model Figure
2
7
1
9
t
9
3
4
7
9
4
9
2
3
2
3
5
4
2 l
6 5
4
3
1
4
/
6
7
8 4
4
0
7
6
4
9
5
9
/
5
7
5
7
3 1
5
4
6
7
7 8 6 1
3 9
2
1
6
7
3
5
/
3
3
0
5
2
3
4
6
7 9
4
9
3
5
O
9
9
9
0
8
1
4
8
9
2
4
5
7
1
Figure 5: 100 samples with the proposed method MNIST dataset.
Figure 6: 100 samples with the proposed method on Celeba dataset.
21

---

Under review as a conference paper at ICLR 2022
MNIST
CELEBA
AE -
Q
7
3
5
1
4
VAE -
0
3
6
g
2
7 3
Q 1
1
6
WAE
VAMP
8
3
9
0
4
5
HVAE
9
0
3
4
0
6
7
RHVAE
4
2
9
6 4 3 7
9
AE-GMM
7
0
6
5
6
8
PAE-GME
2
2
0
5
8
5
E
5
/
?
8
0
RAB (GP)
5
2
6
7 1
g
?
RAE (L2)
7
7
5
/
0
2
RAB (SN)
RAE
2
6
4
4
9
VAE Ours
5
/ 7
2
3
8
4
Figure 7: Generated samples with different models and generation processes.
22

---

Under review as a conference paper at ICLR 2022
D.2 CIFAR AND SVHN
In this appendix, we gather the resulting samplings from the different considered models for SVHN and CIFAR 10.
SVHN
CIFAR 10
0
316
5
3
5
AE-A N
33
41
5
VAB-
0
9
3
5
0
2
WAE
2
9
6
7
8
6
0
5
VAMP
HVAE
11
11 8
S
4
2
RHVAE
3
2
6
AE-GMM
9
2
29
2
6
2
VAE GMM
6
8
866
9
3
2
8
RAB (GP)
40
1203
8
RAB (L2)
71
6
5
3
RAB (SN)
3
8
4
0
510
RAB
3
2 3
2
7
VAE Ours
Figure 8: Generated samples with different models and generation processes.
32
53
3
2
6
6
6
6
Figure 9: Closest element in the training set (Near.) to the generated one (Gen.) with the proposed method.
23

---

Under review as a conference paper at ICLR 2022
E EXPERIMENTAL SET-UP
The RAES, VAES and AEs are trained for 100 epochs for SVHN, M and Celeba and 200 on CIFARIO. Each time we use the official train and test split of the data. For MNIST and SVHN, 10k samples out of the train set are reserved for validation and 40k for CIFARIO. As to Celeba, we use the official validation set for validation. The model that is kept at the end of training is the one achieving the best validation loss. All the models are trained with a batch size of 100 and starting learning rate of 1e-3 (but CIFAR where the learning rate set to 5c-4) with an Adam optimizer Kingmm& & Ba 2014). We also use scheduler decreasing the learning rate by half if the validation loss stops increasing for 5 epochs. For the experiments on the sensitivity to the training set size, we keep the same set-up. For each dataset we ensure that the validation set is 1/5th the size of the train set but for CIFAR where we select the best model on the train set. The neural networks architectures can be found in Table B] and are inspired by Ghosh al. (020)). The metrics (FID and PRD scores) are computed with 10000 samples against the test set (for Celeba we selected only the 10000 first samples of the official test set). The factor p
p max min ||e, cylla ensure some JøL
smoothness of the manifold. For models coming from peers, we use the parameters provided by the authors when available.
For the data augmentation task, the generative models are trained on each class for 1000 epochs with batch size of 100 and starting learning rate of 1c-4. Again scheduler is used and the learning rate is cut by half if the loss does not improve for 20 epochs. All the models have the autoencoding architecture described in Table 3 As the classifier, it is trained with a batch size of 200 for 50 epochs with a starting learning rate of le le-4 4 and Adam optimizer. A scheduler reducing the learning rate by half every 5 epochs if the validation loss does not improve again used. The best kept model is the on achieving the best balanced accuracy on the validation set. Its neural network architecture may be found in Table E MRIs are only pre-processed such that the maximum value of voxel is and the minimum O for each data point.
Table 3: Neural networks used for the encoder and decoders of VAEs in the benchmarks
MNIST [CIFARIO)
SVHN 32, LINEAR RELU
CELEBA
OASIS
EXCOREE
(133 32, 32)
64)
208, 176)
CONv(128, (4, 4), STRIDE=2)
CONv(I28, (5. S), STRIDE=)
CoNv664, (5. S), STRIDE=2) RELU
LAYER
BATCH NORMALIZATOON RELU
BATCH NORMALIZATION
RELU
CoNv (4, 4), STRIDE=2)
(IN0AR(500) RELU
CoNv(256, (5. 5), STRDE=2)
CONV(128, (5. 5) STRIB=2) RELU
LAYER
BATCH NORMALIZATION
BATCH NORMALIZATION
RELU
RELU
CoNv(512, (4, 4), STRIDE=2)
CoNv(512, (5, S), STRIDE=2)
CoNW256, (5. 5). STRIDE=2) RELU
LAYER
BATCH NORMALIZATION
LINEAR[5O0, 16*)
BATCH NORMALIZATION
RELU
RELU
CONY M. 4), STRIDE=2) BATCH NORMALIZATION RELU
CoNv (3. S), STRIDE=2) BATCH NORMALIZATION RELU
CON 3), STRIDE=2) RELU
LAYER 4
LAYER
4006 16*)
L168AR 64*)
CONS 3), STRIDE=2) RELU
LAYER
L195 16*)
DECODER
(32))
(64)
LAYER
HESHAPPR1D244 8,
LINEAR(65530)
LINEAR(500) RELU
RESHAPPR X, 8)
LINEAR(63536)
LINEAR(6S530) RESHAPIC 1024, & 8)
CONVES12, (1. STRIDE=2]
LINEAR (TBO0)
CONVISSI2, (3, STRIDE=2)
(512, 3), ST
LAYER
BATCH NORMALIZATIN
RELU
BATCH NORMALIZATION RELU
RELU
RELU
CONVT2256, (4, 4), STRIDE=2]
3072) RESHAPEI3, 32, 32) SIGMOID
CONVT2256, (), STRIDE=2)
CONVT1256, SJ. STRIDE=))
LAYER
BATCH NORMALIZATION
BATCH NORMALIZATION
RELU
RELU
RELU
CONVTL3. (4. 41. STRIDE=1) BATCH NORMALIZATION
CONVT (5, 5), STRIDE=2) BATCH NORMALIZATION
CONVTU125 STRIDE=2)
LAYER
RELU
SIGMOID
RELU
CoNvIT3. 5J, STRIDE=]) BATCH NOIMMALIZATION SIGMOID
CONV STRIDE=2) RILU
LAYER
ItI. (3. STRIDE=
LAYER
RELu
MNIST images are re-scaled to 32x32 images with a 0 padding.
24

---

Under review as a conference paper at ICLR 2022
Table 4: Neural Network used for the classifier Sec. 5.3
OASIS CLASSIFIER
208, 176)
CONVIS. (3, STRIDE=1)
BATCH NORMALIZATION
LAYER 1
LEAKYRELU
MAXPOOL(2, STRIDE=2)
CoNy(16, 3), STRME=l)
BATCH NORMALIZATION
LEAKYRFLU
MAXPDOL(2. STRIDE=2)
CONv(32, (3. 3J, STRIDE=2)
BATCH NORMALIZATION
LAYER 3
LEAKYRELU
MAXPDOL STRIDE=2)
CoNv(64, 3). STRIDE=2)
BATCH NORMALIPATION
LAYER
LEAKYRELU
MANPOOL(2. STRIDE=2)
L256 100)
LAYER
LINEAR 100, 2)
SOFTMAX
25

---

Under review as a conference paper at ICLR 2022
F DATASET SIZE SENSIBILITY ON SVHN
In Figure 10 we show the same plot for SVHN as Sec. 5.2 Again the proposed method appears be part of the most robust generation procedures to dataset size changes.
SVHN
WAE-X(0, 1]
300
RAE-GP-MO0 3)
AAE-L2-NV0,1
RAE-SN RAE-SN-N(0 -V(0,
200
VAE-N/0,
WAE CCM
150
RAF-GP- AM-GP-GMM
RAF-12- GMM
100
RAE-SN CMM
VAE-CMM
VAE Ours
0.1
0.5
1.0
5.0
x10
Figure 10: FID score evolution according to the number training samples.
26

---

Under review as a conference paper at ICLR 2022
G ABLATION STUDY
G.1 INFLUENCE OF THE NUMBER OF CENTROIDS IN THE METRIC
In order to assess the influence of the number of centroids and their choice in the metric in Eq. 7 we show in Figure [I the evolution of the FID according to the number of centroids in the metric (left) and the variation of FID according the choice the centroids. As expected, choosing small number of centroids will increase the value of the FID since it reduces the variability of the generated samples that will remain close to the centroids. Nonetheless, as soon as the number of centroids is higher than 1000 the FID score is either competitive or better than peers and continues decreasing as the number of centroids increases.
100
mnist
mnist
90
celeha
celeba
80
60
70
60-
50
40
30-
10
2
100
1000
10000
1
5
Number centroids the metric
Centroids choice
Figure 11: Left: FID score evolution according the number of centroids in the metric (Eq. 7 Right: The FID variation with respect to the choice in centroids. We generate 10000 samples by selecting each time different centroids (A: 1000).
To assess the variability of the generated samples, we propose to analyze some generated samples when only 2 centroids considered. In Figure [2] we display on the left the decoded centroids along with the closest image to these decoded centroids in the train set. On the right are presented some generated samples. We place these samples in the top row if they are closer the first decoded centroid and in the bottom row otherwise. Interestingly, even with a small number of centroids the proposed sampling scheme is able to access to a relatively good diversity samples. These samples are not simply resampled train images or simple interpolation between selected centroids some of the generated samples have attributes such as glasses that are not present in the images of decoded centroids.
Decoded centroid Nearest train image
Generated samples
Figure 12: Variability of the generated samples when only two centroids are considered in the metric. Left: The image obtained by decoding the centroids. Middle: The nearest image the train set to the decoded centroids. Right: Some generated samples. Each generated sample is assigned to the closest decoded centroid (top row for the first centroid and bottom row for the second one).
27

---

Under review as a conference paper at ICLR 2022
G.2 INFLUENCE OF À IN THE METRIC
In this section, we also assess the influence of the regularization factor À in Eq. on on the resulting sampling. To do so, we generate 10k samples using the proposed method both MNIST and Celeba datasets for values of À E [le- le- 1e-2, 1e-1, 11]. Then, we compute the FID against the test set. Each time, we consider influence of À remains limited. In practice, À is mainly here to avoid pathological cases such as the metric collapsing to far from the centroids. In the implementation, a typical choice for À is le-2.
1000 centroids the metric. As shown in Figure 13 the
mnist
70-
celeba
60-
40
30
20
10-6
10-5
10
10
10
10
100
Value
Figure 13: FID score evolution according to the value of X in the metric (Eq 7
G.3 THE CHOICE OF
In the experiments presented, the smoothing factor p in Eq. is set to the value of the maximum distance between two closest centroids = max min |cj cill2. This choice motivated by the jyt
fact that wanted to build smooth metric and ensure some smoothness of the manifold while trying to interpolate faithfully between the metric tensors G; = E_- In particular, a small value of p would have allowed disconnected regions and the sampling may have not prospected well the learned manifold and would have only become a resampling of the centroids. On the other hand, μ4;
setting high value for p would have biased the interpolation and the value of the metric As a result, ((s)) might have been very different from the one observed Σ; since the other μ. would have had strong influence its value. The proposed value for p appeared to work well in practice.
28

---

Under review as a conference paper at ICLR 2022
H OTHER CLASSIFICATION METRICS
In Table are presented some further classification results for each considered model while gener- ated samples using each generation procedure made available in Figure [4] On OASIS database the proposed method appears to produce visually the sharpest samples. This better generation per- formance also supported by the classification metries provided in Table Table
Table 5: Classification results averaged 20 independent runs. For the VAEs, the classifier is trained on 2K generated samples class.
Precision
Recall
Generation method
Balanced Accuracy
AD
CN
AD
CN
Original * Original (resampled)
66.2 + 7.6
74.7 + 8.4
80.3 4,0
35.7 = 16.3
95.7 + 1.5
81.8 +: 2.6
67.0 + 5.3
91.4 + 1.8
18.5
85.1 4.2
AE-1 .V
50.0 + 0.0
0.0
0.0
0.0
100.0
WAE
57.4 + 9.7
48.51
7.17
19.3
95.4
VAE- N
51.8 3.8
47.
7.74
3.7
9.7
VAMP
83.1 2.6
5.2
9.1
94.8
7..5
HVAE
56.3
48.7
3.8
13.9 17.6
98.6
RHVAE
68.0 10.9
56.1
8.50
46.7
89.2
AE - GMM
8.34
4.98.1499
96.8
5.6
6.2
RAE (GP)
6.8
18.5
8.62
00.9228.6
67.0 24,9
RAE (L.2)
7.0
57.8
8.2
70.0
78.3 11.7
RAE (SN)
6.9
24.9
6.6
41.7
82.9 16.4
RAE
6.1
56.2 13.5
8.22
60.0 24.0
78.5 17.5
VAE- GMM
83.0
6..7
9..9
9.5
77.9 5.9
VAE Ours
2.5
5.30
2.2
9..4
80.3 5.1
uunbalanced
29

---

Under review as a conference paper at ICLR 2022
OASIS
Train
VAE N
WAE
VAMP
HVAE
RHVAE
-GGM
RAE (GP)
RAE (L2)
RAE (SN)
RAE
VAE-Ours
Figure 14: Generated samples with different models and generation processes.
30

---

Under review as a conference paper at ICLR 2022
I LINK BETWEEN THE RIEMANNIAN VAE AND VANILLA VAE
We assume as (Ghosh et al. 2020) that VAE is essentially an autoencoder regularized with noise. Hence, a Riemannian-basee noise would be informed by the intrinsic geometry of the latent space. Indeed, the main assumption behind the Riemannian VAE would consist in assuming that given a set of data æ E X C R there exists a lower dimensional space, namely the latent space, that has apparently no reason to be Euclidean in which live the latent variables. In such a framework, it would be assumed that this space is the d-dimensional Riemannian manifold M = (R4, G) where G is an unknown Riemannian metric. Now the goal of the Riemannian VAE would be the same regularized autoencoder that is to learn a smooth representation of the data within a much lower dimensional space here seen as the Riemannian manifold M.
VAE could also be seen as a regularized autoencoder but the
To to so and similarly to autoencoder models, it would be assumed that there exist Ep : IR + M a parametrized encoding function mapping the input data onto the manifold eœ(x) - μa E M and dg .M- RDD parametrized decoding function that maps back the latent codes the data space. In such a case, the main objective would be to find o and 0 such that the reconstruction loss is minimized
min FREC = min (ax, do(e)(z))) EEX, 6.8
(14)
6.8
where I is a function measuring the distance between the input data and the reconstructions and is chosen depending on the problem and the data (e.g. mean square error, binary cross entropy...) In order to learn smooth latent space meaning that small variations in the latent space do not change completely the output of the decoder, the decoder also regularized using a Riemannian Gaussian noise. That is imposing that
dx(z) RJ x,
2
NG Vicm (2((1,))
where
NCG riem (z|a,, (1) = C exp
dista (2, 11)3
C=
exp (-
dista 11)2
dMz
(15)
20
20
M
and the reconstruction loss in Eq. 14 would become
min PRCC min (ar, do(z)),
E X,
~~ NC (2))
riem
Hence, no longer decode the embedding μ but rather that is obtained with the Riemannian Gaussian distribution centered on μ.. Since the metric G is unknown, a Riemannian VAE would aim learning the metric directly from the data. Thus, the encoder function would output an embedding of an input data point but also the value of the Riemannian metric at the embedding point i.e. ((a). Since, we would only have access to a finite number of metric tensors, a smooth metric G could be built using Eq. 7 Now, at least theoretically, we would be able to compute the geodesic distance involved in [q. As of now, the manifold is not regularized and so pathological cases such as the metric collapsing to O may arise. To avoid such a behavior, some smoothness conditions could be applied on the manifold by imposing for instance that the Riemannian Gaussian distribution not too far from standard Gaussian. This would prevent the metric from collapsing to 0 and ensure that the latent codes remain close to the origin as well. However, other regularization schemes could have be envisioned as well. Since we would be working in an ambient-lik manifold, there exists global chart z and so the density of the Riemannian Gaussian distribution can be written with respect to the Lebesgue measure dz in Rdd |Pennec 2006)
dist (z, 11)
NG riem (z|0, 1))
C exp
20 V det G(z)
dist (2, 11)
(16)
C
exp
v det G(z) dz,
2x
Ra
Hence, the regularization term set as the KL divergence between the Riemannian Gaussian distribu- tion and the standard normal would follow
LRDG DKL (NG riem (z|a, (O)
/ PNo (oss) (ocll/NNa...nn log(pw(o..a))) dz.
rmam
31

---

Under review as a conference paper at ICLR 2022
The final objective a Riemannian VAE would try minimize would then write
C = LREC + fREG.
Unfortunately, this framework cannot be used practice for at least two reasons. The first one is the sampling from the Riemaniann distribution in Eq. [1] which is far from being trivial. MCMC methods could have been envisioned to sample from such a distribution but this would have impeded backpropagation since this framework would not amenable the reparametrization trick (Sali- mans et al., 2015 Caterini et al. 2018). Second, the regularization term would involve computing the density of the Riemannian Gaussian distribution which explicitly involves the computation of the Riemannian distance and so the resolution of the optimization problem in Eq. 9 Since, this framework not usable in practice, it can be assumed that the value of the metric during training can be approximated by its value at μe (i.e. G(u)). With this approximation, Riemannian Gaussians become multivariate Gaussians and all the terms become computable. We find back the vanilla VAE framework if we further consider G( E(x) where E(x) is the covariance matrix given by the encoder vanilla VAE. We indeed have
1
disty H1)
riem (z/o, 1))
C
exp
20
v E-I E-1
disty (2, p1)2
(17)
C
exp
Vdet W-1 dz
20
Ra
where disty- (z. (2 (b)
(z n.). If further σ 1 have
(s- 14) E 1x 11)
exp
No Viem (z)o,,1)
(18)
(x-1e)
(x=
exp Rd
2
Assuming as is standard in the VAF framework, that E is diagonal makes the computation of LREG easy and we retrieve the training of vanilla VAE model. Indeed, as explained Remark. [ the log of the conditional distribution pa reduces to the reconstruction loss LREC and the KL between the variational posterior qs (z/z) and the prior taken as a standard Gaussian gives CREG.
32


Reviews:
[{"review_id": "kdYesHMSIk", "reviewer": "Reviewer_UafD", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "#### Strengths\n- The stated contributions of the paper are substantiated by the experiments.\n- The proposed work is competitive and often superior to related methods on several datasets.\n\n#### Weaknesses\n- The use of non-standard English makes reading somewhat difficult.\n- I would have liked to see some discussion of the cost incurred from HMC sampling\n- I would have liked to see some discussion of the relationship of the proposed approach to [1-2].\n\n[1] Rakowski, Alexander, and Christoph Lippert. \"Disentanglement and Local Directions of Variance.\" Joint European Conference on Machine Learning and Knowledge Discovery in Databases. Springer, Cham, 2021.\n\n[2] Rolinek, Michal, Dominik Zietlow, and Georg Martius. \"Variational autoencoders pursue pca directions (by accident).\" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2019.\n\n#### Possible typos:\n- “...using either other regularization…” → using other regularization (p.6)\n- “...4 well known databases” → “four well known datasets” (p.6)\n- “visually speaking” → “qualitatively” (p.6)\n- “...same models and data sets…” → same models and datasets (p.8)\n- “... quite robust to the data set” → quite robust to the dataset (p.6)\n- \"According to us, these metrics\" --> These metrics (p.9)\n- \"A d-dimensional manifold... d-dimensional Euclidean\" --> A $d$-dimensional manifold... $d$-dimensional Euclidean (p.3)\n- There are parts of figure 1 that aren’t fully described in the caption (e.g., part (a) and part (b) in top/bottom)\n", "summary_of_the_paper": "The authors propose an interpretation of the VAE latent space as a Riemannian manifold and leverage this structure to design an improved sampling strategy. To construct a Riemannian manifold, the authors interpret the covariance of the posterior of each latent as a local Riemannian metric. The local metrics are interpolated to produce the latent manifold. To generate samples, the authors use HMC on the manifold specific uniform distribution constructed using local metrics.", "summary_of_the_review": "The authors introduce a novel sampling procedure for the classical VAE that produces quantitatively and qualitatively strong results. The sampling approach leverages a novel theoretically-based interpretation of the VAE latent space structure. The paper itself is intuition building (e.g., toy example) but difficult to read in some places. It also appears the authors may have neglected a related body of work.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "erkAcgTpMMW", "reviewer": "Reviewer_v2S3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n- The proposed approach, i.e., using the Riemannian geometry to formulate distributions in the VAE, is interesting.\n- The paper proposes a general perspective on VAEs through a geometrical perspective. It generalizes some papers on VAEs with various latent manifolds (e.g., hyperspherical VAE, VAEs with Poincare disk).\n\nWeaknesses:\n- Remark 1 is not necessarily true for two reasons. First, this is not true that $p_{\\theta}(\\mathbf{x}|\\mathbf{z})$ is often taken as a multivariate Gaussian. The conditional distribution must be chosen according to the problem at hand. If data is binary, the Bernoulli distribution is taken; if data is continuous, we can pick from a plethora of distributions. Second, a very common choice to model RGB images is a discretized logistic distribution rather than a Gaussian distribution.\n- The example in Sec. 4.5 is a bit vague. First, the setup is not well explained. Second, Figure 1 is not properly explained: there are three times (a), and one (b), but they are neither referred to in the caption nor in the text. Moreover, the font used for the top (a) and top (b) are so small that it is hard to read.\n- The paper explains the idea, however, it lacks a more concrete statement of the proposed approach. The authors mention that they use Eq. (7) as the Riemannian metric and the HMC sampler for sampling, however, it is not necessarily obvious what they precisely use later on in the experiments. It would be highly beneficial to provide a specific instantiation of their approach, e.g., how the metric is calculated, what is the final objective, what is the sampling procedure, what is the training algorithm. At the moment, the paper presents the ideas at a rather general level.\n\n", "summary_of_the_paper": "The paper proposes to generalize the vanilla VAE by formulating the conditional distribution and the marginal distribution using the Riemannian geometry. The authors proposed a generalized perspective on various VAEs proposed so far. The paper is rather easy to follow, however, it is unclear what precisely is used in the experiments and how the proposed approach works.", "summary_of_the_review": "I like the paper and the idea of using Riemannian geometry in VAEs. However, the paper presents the idea at a pretty general level and it is hard to see the details that matter from the computational perspective. I doubt it would be easy and straightforward to code the proposed approach up. This is my main problem with the paper, namely, the reproducibility. And the follow-up problem is about the computational burden of the proposed approach. Since it is unclear how to sample and how to calculate various components of the VAE, it is hard to assess how “heavy” the method is and whether it could be scaled up. As we see nowadays, models that are easy to scale up (e.g., hierarchical VAEs like NVAE or diffusion-based models) allow achieving SOTA performance.\n\n=== UPDATE ===\nI appreciate the rebuttal. Overall, I still doubt whether the paper is a strong contribution. Nevertheless, the new additions (especially Appendix C) make the paper better. Therefore, I decided to increase my score to 6.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "I2glAczsmRc", "reviewer": "Reviewer_BQpw", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This work proposes a novel interpretation of variational encoding by the prism of differential geometry.  In particular the data $x_i$ provides two main elements, the “position” $\\mu(x_i)$ and instead of interpreting the estimated $\\Sigma(x_i)$ as the variance of the normal law i.e of p(z|x) it is interpreted as the value of the Riemanian metric at $\\mu(x_i)$. In that perspective, the authors propose to sample not from an isotropic Gaussian of $\\mathbb{R}^d$, but rather using uniformly from the learned manifold.\nThe experiments are interesting and well-conducted.\n\nRemarks and Question:\n1. On the interpretation of the Geometric interpretation of the VAE. Section 4.3 is a little confusing. Indeed, the initial assumption that the inverse variance corresponds to the local metric tensor is acceptable (the lower the variance the less the sample is allowed to move from the mean). However, the rest of the section is in my opinion confusing. Indeed, because the authors seem to suggest that the original-VAE is in fact an approximation of a Riemannian VAE, can the authors link the Riemannian VAE penalties and how they are approximated in the original VAE framework? For clarity, I highly recommend including such discussion in the final version.\n2. Can the authors comment on the link between their work and the desire of the VAE-community to rather enrich the expressivity of p(z|x) using normalizing flows or Hamiltonian scheme [1, 2]\n3. I can’t relate eq.7 to the interpretation of the inverse variance as a local Riemannian metric (i.e $G(\\mu(x_i) = \\Sigma^{-1}(x_i)$). Indeed as of eq.7 $G(\\mu(x_i)) \\neq \\Sigma^{-1}$, depending on how $\\mu(x_i)$ is close to the other $\\mu_i$.\n4. In the same line can the author comment on the need for the regularization term $+ \\lambda I_d$ ? How to choose $\\lambda$ in practice ?\n\nOn the Experiments:\n1. What is the influence of the number of selected points in the sampling based on 7). For example can the author provide a FID curve somehow like Fig.4 with varying number of “support vectors” $(\\mu_i, \\Sigma^{-1}_i)$ ? Indeed, the more complex the latent space geometry, the more “support-vectors” it may require for an acceptable decoding.\n2. Because the sampling method in the latent space is dependent on the chosen support vectors, did the authors observe any variance in the quality of the samples depending on the choice of the K-centroids (something like variance of the FID at a defined number of support vectors)?\n3. On the Decoding processes : One drawback of such sampling scheme (due to the choice of the centroids) is the fact that in the decoding processes, the decoder simply provides an interpolation between the decoded K-support vectors. Did the authors observe such an issue in practice ?\n\nI would be glad to increase my score if the major concerns are answered.\n\n[1]: Hamiltonian Variational Auto-Encoder,  Anthony L. Caterini, Arnaud Doucet, Dino Sejdinovic  \n\n[2]: Variational Inference with Normalizing Flows, Danilo Jimenez Rezende, Shakir Mohamed \n", "summary_of_the_paper": "The paper proposes to analyse the VAE framework using differential geometry tools. More specifically, they propose to impose a (Riemannian) metric in the latent space of a VAE. Such a Riemanian metric then enables data generation based on the (Riemanian) Uniform distribution on the learned latent space, improving the quality of the generated samples.\n", "summary_of_the_review": "The geometric interpretation of the VAE is interesting. Moreover, the author's proposition for resampling seems sound in the light of their proposition and the experiments are well conducted. Yet, the theoretical link between the VAE and the Riemannian distance based VAE requires a more thorough discussion.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Fq6WBPTaDpl", "reviewer": "Reviewer_yrKH", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "\nThe idea is to equip the latent space with a Riemannian structure, the measure of which depends on the inverse of covariance of the posterior.\nBefore sampling, the distribution needs to be approximated using MCMC.\nThe results are ok.\n\nPros: \n* The article is well written and smoothly walks the reader through the mathematics behind variational autoencoders (VAE) and Riemannian geometry. \n* The idea is clear and straightforward. The theory is robust. However, I am doubtful of the practical outcomes. \n\nCons: \n* The analyses present several shortcomings.\n* The baselines are chosen accordingly to the theoretical line of the paper. That is to improve the samples' quality without touching (too much) the loss nor the architecture.\nThat is fair, although I was expecting some Energy-based approaches such as VAEBM (Xiao et al. 2020) or Pang et al. (which is in the references but not cited)\n* The authors overemphasize sometimes. \n** \"Now that the VAE has learned and capture the intrinsic geometry of the data within the latent space seen as a Riemannian manifold,\" \n**    Why is the variance-based measure necessarily the one that captures the \"intrinsic geometry of the data\"?\n    Besides, what if the data is made of several connected components? Such a situation is not allowed by the theory developed here.\n**  VAE+GMM is almost as good as the proposed method in Table 2.\n", "summary_of_the_paper": "It is a theoretical paper that discusses the Riemannian structure of the latent space of a VAE.\nThe associated measure on a point \\mu(x) derives from \\Sigma^-1(x), which are the parameters of the posterior p(z|x) or simply the outputs of the encoder.\nThe measure allows to define a Riemannian normal distribution on that space/manifold, to compute geodesics but also to \"better\" sample.", "summary_of_the_review": "The document is good, easy to read, and fits within the scope of the conference.\nThe Riemannian point of view on VAE surely brings interesting insights that might lay the foundation of other works. \nNevertheless, I have doubts about the importance of the contributions and their impact on the field. \nMoreover, the analysis is pretty basic. Maybe because there is too much to say/study for a conference paper.\n\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "2IayOvK_0rJ", "reviewer": "Reviewer_NpUX", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper provides a way of incorporating geometry into the VAE setup for generating samples that belong to this manifold. While geometry-aware sampling strategies are available in prior literature, the authors claim that these processes are time-consuming and complex. To better position this statement, I think the runtime and complexity of the proposed method need to be discussed and substantiated.\n\nI assume the authors do not change the training process and are only proposing a better sample generation strategy - it is not clear from the current text if this is the case and if so, it is unclear what training procedure they are resorting to. The authors review Riemannian geometry concepts that are relevant to the method proposed and it is easy to follow through this part of the paper.\nThe toy example shows the advantage of using the proposed sample generation procedure in comparison to vanilla VAE sample generation - the proposed method generates relevant samples that maintain the simple geometric structure as in the input data.\nExperiments show improvements over RAE’s, VAMP, and RHVAE in terms of FID score and classification with a VAE generated set of data. The performance of the proposed method in a low data regime looks interesting and would be interesting to have a discussion on why the proposed method is able to achieve this. \n\nMy major concern is the variability of the samples generated. I wonder if the samples generated in real dataset experiments are biased towards resampling available input data. If \\rho variable used for the calculation of the metric in equation (7) is very large then the sampling process is biased towards just picking the closest input data point. I think the authors need to better support the claim on the variability of generated samples - interpolation scenarios in real images, random samples (not clear if the generated samples shown are random or cherry-picked), or even reconstructions for a given input (this would matter if the proposed method is modifying the training process as well). From the current set of visualizations and results, it is not clear if the images generated are new or just resampled input data.\n\nOther concerns (in no particular order) are listed below.\n- There is no clarity on the modification of the loss function (if any) used to train the VAE model. I assume the model is learning a mean vector and a covariance matrix for each input image - but are there any additional constraints to ensure that the matrix learned is a valid covariance matrix? More critically, what if the covariance matrix is not invertible? How are samples generated during training?\n\n- There is no discussion on hyperparameters (\\lambda, \\rho) and the sensitivity of the experiments to these parameters. The supplementary material says rho is assigned by a max-min distance formulation - is this based on the entire dataset or a subset? Any reason why this is the right choice? \n\n- In the OASIS dataset experiment, the generated dataset is made class balanced and evaluated. But the “original” classifier is trained with an unbalanced set - Is it possible to evaluate by resampling or reweighting the original training set to make them class balanced? I understand this introduces bias in the classifier but so does the experimental setup proposed. \n\n- Any reason for dropping the RHVAE evaluation in Table 2?\n\n- I'm not sure FID comparison with methods using prior is a fair setting - I believe it gives methods with “ex-post density estimate” an unfair advantage in reducing the gap between the posterior and the prior. This is by no means specific to this paper but I would be interested in knowing the author’s perspective on this.\n\nThe ideas and experiments by [1,2] are relevant to the proposed method and the current paper can be improved by incorporating and differentiating with respect to these recent works.\n\n[1] Kalatzis, Dimitrios, et al. \"Variational Autoencoders with Riemannian Brownian Motion Priors.\" International Conference on Machine Learning. PMLR, 2020.\n\n[2] Connor, Marissa, Gregory Canal, and Christopher Rozell. \"Variational Autoencoder with Learned Latent Structure.\" International Conference on Artificial Intelligence and Statistics. PMLR, 2021.", "summary_of_the_paper": "The paper introduces a geometric perspective on VAEs trained with a Gaussian mixture prior (standard multivariate or learned mixture of Gaussian): the covariance matrix encoded for a given input image in a VAE corresponds to a Riemannian metric (inverse covariance) around the local neighborhood of the data point encoded.\nThe authors suggest sample generation in VAE  should be based on this covariance-induced distance (Mahalanobis). To obtain a single continuous metric at the end of the training, the authors suggest an interpolation based on available data and then generating data using a Monte Carlo (HMC) sampler based on this interpolated metric. \nExperiments using samples generated using this method shows improvements in FID and precision-recall measures in MNIST, SVHN, CIFAR10, and CelebA datasets. The authors claim that the proposed sample generation produces good samples even with datasets of small size.\n", "summary_of_the_review": "Overall, I believe the paper needs revision to better support the claims in the paper. The idea of incorporating geometry in the VAE generation procedure is interesting and has potential. I believe the authors need to improve clarity on experimental setup and evaluation, and better support the variability of samples generated. The current form of paper has a large room for improvement and I hope the authors are able to incorporate the concerns raised and better position their work.\n\n----\nUpdate: I appreciate the authors for answering the concerns raised and am increasing my recommendation to 6. The additional experiments and details on the algorithm in the appendix improve the clarity of the paper. The proposed sampling strategy shows improvement (FID metric) in conjunction with a simple VAE model. It might be interesting to see if the proposed sampling can be used to improve other VAE models. The experimental setup with OASIS is not as compelling (my suggestion of resampling is only a simple strategy and I believe there are better class-balancing classification methods), but highlights the shortcoming of using earlier VAE methods in this setup.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes to use covariance of the approximate posterior to induce a metric on the latent space of the VAE and use it to sample from the Riemannian manifold learned by a VAE. Experiments on MNIST and CelebA show the method outperforms vanilla VAE in terms of sample quality (FID and PR scores). It is also shown to work better than baseline VAE models on a medical imaging classification task. While the reviewers have acknowledged the contributions of the paper, the novelty in the contributions and their importance/impact was seen to be rather limited. The main concern from the reviewers is -- while the paper is mainly based on the use of inverse covariance as the metric for manifold, it doesn't give a reasonable theoretical justification on it is a sensible metric that captures the intrinsic geometry of data. Authors in their response justify it as -- since the covariance matrices are learned from the data and favor through the posterior sampling some direction in the latent space, it is a natural choice as metric. This is not very convincing. A more technical justification for this will certainly make the paper more convincing. I suggest the authors to look at "Kumar, Abhishek, and Ben Poole. "On Implicit Regularization in $ β $-VAEs." International Conference on Machine Learning. PMLR, 2020" which theoretically connects inverse covariance and the Riemannian metric in Sec 5.2, and see if it can be adapted in their context.
Decision: Reject