# A GEOMETRIC PERSPECTIVE ON VARIATIONAL AU## TOENCODERS

**Anonymous authors**
Paper under double-blind review

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 datasets.
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 dataset combining both high
dimensional data and low sample sizes.

1 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 (Yang et al., 2019) or data augmentation (Chadebec et al., 2021).

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 to several limitations of the VAE
framework. First, the training of a VAE aims at maximizing the Evidence Lower BOund (ELBO)
which is only a 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 as 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 a given prior distribution
inducing over-regularization (Connor et al., 2021). Hence, trying to tackle those limitations through
richer posterior distributions (Salimans et al., 2015; Rezende & Mohamed, 2015) or better priors
(Tomczak & Welling, 2018) represents a 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, we 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 improves significantly the generation process from a vanilla VAE and makes it able to perform
as well as more advanced VAE models in terms of Frechet Inception Distance (Heusel et al.,
2017) and Precision and Recall (Sajjadi et al., 2019) scores on four benchmark datasets.

-  We also show that the propose method appears more robust to dataset size changes and
outperforms even more strongly peers when only smaller sample sizes are considered.


-----

-  We validate the method on complex neuroimaging data from OASIS (Marcus et al., 2007).

2 VARIATIONAL AUTOENCODERS

Considering that we are given x ∈X a set of data points deriving from an unknown distribution
_p(x), a VAE aims at inferring such a distribution with a parametric model_ _pθ, θ_ Θ 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
**z** _qprior(z)_ ; **x** _pθ(x_ **z),**
_∼_ _∼_ _|_

where qprior is a prior distribution over the latent variables often taken as a standard Gaussian and
_pθ(x_ **z) is referred to as the decoder and is most of the time taken as a parametric distribution the**
_|_
parameters of which are estimated using neural networks. Hence, the likelihood pθ writes:


_pθ(x_ **_z)q(z)dz ._** (1)
_|_


_pθ(x) =_


As this integral is most of the time intractable so is pθ(z **x), the posterior distribution. Hence,**
_|_
Variational Inference (Jordan et al., 1999) is used and a simple parametrized variational distribution
_qϕ(z_ **x) is introduced to approximate the posterior pθ(z** **x). qϕ(z** **x) is referred to as the encoder**
_|_ _|_ _|_
and, in the vanilla VAE, qϕ is chosen as a multivariate Gaussian whose parameters µϕ and Σϕ are
again given by neural networks. An unbiased estimate of the likelihood pθ(x) can then be derived
using importance sampling with qϕ(z **x) and the ELBO objective follows using Jensen’s inequality:**
_|_

log pθ(x) = log Ez _qϕ_ _pˆθ_ Ez _qϕ_ log ˆpθ = Ez _qϕ log pθ(x_ **z)** _DKL(qϕ(z_ **x)** _p(z))_
_∼_ _≥_ _∼_ _∼_ _|_ _−_ _|_ _∥_ (2)
    _ELBO_

The Evidence Lower BOund (ELBO) is now tractable since both| _p{z_ _θ(x_ **z) and qϕ(z** **x}** ) are
_|_ _|_
parametrized and so can be optimized with respect to the encoder and decoder parameters.

**Remark 1 In practice, pθ(x|z) is chosen depending on the modeling of the input data but is often**
_taken as a 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 is a reconstruction term given_
_by pθ(x_ **z) while the second one is a regularizer corresponding to the KL divergence between the**
_|_
_posterior and the prior. For instance, in the case of a multivariate Gaussian we have_

_LREC = ∥x −_ **_µθ(z)∥2[2][,]_** _LREG = DKL(qϕ(z|x)∥p(z)) ._ (3)

3 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 of Posterior
(VAMP) (Tomczak & Welling, 2018) as priors was proposed. In the same vein, hierarchical latent
variable models (Sønderby et al., 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 a given dataset. Acceptance/rejection sampling method was also proposed to try to improve the
expressiveness of the prior distribution (Bauer & Mnih, 2019). Some recent works linking energybased models (EBM) and VAEs (Xiao et al., 2020) 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 modelings as hypershpere (Davidson et al., 2018), torus (Falorsi et al., 2018) or Poincar´e disk (Mathieu
et al., 2019) or discrete latent representations (Razavi et al., 2020) were proposed. Other recent contributions proposed to see the latent space as a Riemannian manifold where the Riemannian metric
is given by the Jacobian of the generator function (Arvanitidis et al., 2018; Chen et al., 2018; Shao
et al., 2018). This metric was then used directly within the prior modeled by Brownian motions


-----

(Kalatzis et al., 2020). Others proposed to learn the metric directly from the data throughout training 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 of the learned latent space they either require the computation of a time
consuming model-dependent function, the Jacobian, or add further parameters to the model to learn
the metric or transport operators adding some computational burden to the method.

Arguing that VAEs are essentially autoencoders regularized with a Gaussian noise, Ghosh et al.
(2020) proposed another interesting interpretation of the VAE framework and showed that other
types of regularization may be of interest as well. Since the generation process from these autoencoders is no longer relying on the prior distribution, the authors proposed to use ex-post density
estimation by fitting simple distributions such a Gaussian mixture in the latent space. While this
paves the way for consideration of other ways to generate data, it mainly reduces the VAE framework to an autoencoder while we believe that it can also unveil interesting geometrical aspects.

Another widely discussed improvement of the model consists in trying to tweak the approximate
posterior in the ELBO so that it better matches the true posterior using MCMC methods (Salimans 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 on
the impact of both the intrinsic structure of the variance of the data (Rakowski & Lippert, 2021) and
the shape of covariance matrices in the posterior distributions (Rolinek et al., 2019) on disentanglement 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 to see their inverse as the
value of a Riemannian metric evaluated at the embedding points µi.

4 PROPOSED METHOD

In this section, we argue that a vanilla VAE shows naturally a Riemannian structure of the latent
space through the learned covariance matrices in the posterior distributions. We then propose a 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 in the rest of the paper.
A more detailed discussion on integration and probability densities on manifolds may be found in
Appendix A. A d-dimensional manifold M is a manifold which is locally homeomorphic to a ddimensional Euclidean space. If the manifold M is further connected and differentiable it possesses
a tangent space Tz at any z ∈M composed of the tangent vectors of the curves passing by z. If M
is equipped with a smooth inner product g = ⟨·|·⟩z defined on its tangent space Tz for any z ∈M
then M it is called a Riemannian manifold and g is the associated Riemannian metric. Since g is
an inner product, a local representation of g at any z ∈M is given by the positive definite matrix
**G(z). The notion of length of curves γ : R →M traveling in M can be defined as follows**


_γ˙_ (t)[⊤]G(γ(t))˙γ(t)dt .


_L(γ) =_


_γ˙_ (t) _γ˙_ (t) _γ(t)dt =_
_⟨_ _|_ _⟩_


Curves minimizing L are geodesics and a Riemannian distance between z1, z2 can be defined
_∈M_
distG(z1, z2) = inf s.t. **_z1 = γ(0), z2 = γ(1) ._** (4)
_γ_ _[L][(][γ][)]_

The manifold M is said to be geodesically complete if all geodesic curves can be extended to R. In
an Euclidean space, G reduces to the Id 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 Id. In such a case the induced_
_Riemannian distance is the well-known Mahalanobis distance which writes_


distΣ =


(z2 **_z1)[⊤]Σ(z2_** **_z1) ._**
_−_ _−_


-----

4.2 THE RIEMANNIAN GAUSSIAN DISTRIBUTION

The notion of measure and so of probability distribution can be extended to geodesically complete
Riemannian manifolds as well (Pennec, 2006). Given the Riemannian manifold M endowed with
the Riemannian metric G and a chart z, an infinitesimal volume element may be defined on each
tangent space Tz of the manifold M as follows

_dMz =_ det G(z)dz, (5)

with dz being the Lebesgue measure. Hence, a Riemannian Gaussian distribution onp _M can be_
defined and consists in using the Riemannian distance defined in Eq. 4 instead of the Euclidean one


riem(z _σ, µ) = [1]_
_N_ _|_ _C_ [exp] _−_ [dist][G]2[(]σ[z][,][ µ][)][2]




exp
_−_ [dist][G]2[(]σ[z][,][ µ][)][2]



_dMz,_ (6)


_, C =_


where dMz is the volume element defined by Eq. 5. Hence, the multivariate normal distribution is
only a specific case of the Riemannian distribution with σ = 1, defined on the manifold M = R[d]
endowed with the constant Riemannian metric G(z) = Σ[−][1], ∀z ∈M.

4.3 GEOMETRICAL INTERPRETATION OF THE VAE FRAMEWORK

Within the VAE framework, the variational distribution qϕ(z **x) is voluntarily chosen as a sim-**
_|_
ple multivariate Gaussian distribution defined on R[d] with d being the latent space dimension.
Hence, as explained in the previous section, given an input data point xi, the posterior qϕ(z **x) =**
_|_
(µ(xi), Σ(xi)) can also be seen as a Riemannian Gaussian distribution where the Riemannian
_N_
distance is simply the distance with respect to the metric tensor Σ[−][1](xi). Hence, the VAE framework can be seen as follows. As with an autoencoder, the VAE provides a code µ(xi) which is a
lower dimensional representation of an input data point xi. However, it also gives a tensor Σ[−][1](xi)
depending on xi which can be seen as the value of a Riemannian metric G at µ(xi) i.e.

**G(µ(xi)) = Σ[−][1](xi) .**

This metric is crucial since it impacts the notion of distance in the latent space now seen as the Riemannian manifold M = (R[d], G) and so changes the directions that are favored in the sampling from
the posterior distribution qϕ(z **x). Then, a sample z is drawn from a standard (i.e. σ = 1 in Eq. 6)**
_|_
Riemannian Gaussian distribution and fed to the decoder. As first approximation and since we only
have access to a finite number of metric tensors Σ[−][1](xi), the VAE model assumes that the metric is
locally constant close to µ(xi) 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. 3, that z gives a sample x _pθ(x_ **z) close to xi**
_∼_ _|_
when decoded. Since z has a probability density function imposing higher probability for samples
having the smallest Riemannian distance to µ, the VAE imposes in a way that latent variables that
are 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. 1. Noteworthy is that the latter
distance can be amended through the choice of the decoder pθ(x **z). This is a 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 0d 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 µ(xi) and a family of metric tensors (Gi = Σ[−][1](xi)) corresponding
to the value of a Riemannian metric defined locally on the latent space. Inspired from Hauberg et al.
(2012), we propose to build a smooth continuous Riemannian metric defined on the entire latent
space by performing the following interpolation:


_−_ [dist][Σ][−][1][(][x]ρ[i][2][)][(][z][,][ µ][i][)][2]


**Σ[−][1](xi)** _ωi(z) + λ_ **_Id,_** _ωi(z) = exp_
_·_ _·_
_i=1_

X


(7)


**G(z) =**


where distΣ−1(xi)(z, µi) = (z **_µi)[⊤]Σ[−][1](xi)(z_** **_µi) is the Riemannian distance between z_**
_−_ _−_
and µi with respect to the locally constant metric G(µ(xi)) = Σ[−][1](xi). Since the sum in Eq. 7


-----

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 k < N metric tensors[1] and increasing ρ to reduce
memory usage. We provide an ablation study on the impact of λ, the number of centroids k and
their choice along with a discussion on the choice for ρ in Appendix G. Then, we have:

**Proposition 1 The Riemannian manifold M = (R[d], G) is geodesically complete.**

Prop. 1 (proved in Appendix B) allows now to refer to probability densities on M. Rigorously, the
metric defined in Eq. 7 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. 4 needed, ii) the sampling from Eq. 6 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 µi (i.e. Σ[−]i [1][(][x][i][)][), the VAE training remains unchanged, stable and computationally reasonable]
since Riemannian Gaussians become multivariate Gaussians in qϕ(z **x). Noteworthy is the fact that,**
_|_
likewise (Ghosh et al., 2020), in 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. In 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 a space seen as a Riemannian manifold, we propose to exploit this strong property to enhance the generation procedure. A
natural way to sample from such a latent space would consist in sampling from the uniform distribution 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 ⊂M, the uniform distribution writes (Pennec, 2006)

**1** (z)

_p_ (z) = **[1][A][(][z][)]** _A_ _._
_A_ Vol(A) [=] **1A(z)dMz**

_M_

This density is taken with respect to d _z, the Riemannian measure but using Eq. 5 and a coordinateR_
_M_
system z allows to obtain a pdf now defined with respect to the Lebesgue measure:

_URiem(z) ∝_ det G(z) .

Since the Riemannian metric has a closed form expression given by Eq. 7, sampling from this

p

distribution is quite easy and may be performed using the HMC sampler (Neal, 2005) for instance.
Now we are able to sample from the intrinsic uniform distribution which is 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 1 where a vanilla
VAE was trained with a toy dataset composed of binary images of disks and rings of different size
and thickness (example inspired by Chadebec et al. (2021)). 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 as a (0, Id)) in such a case is far from being optimal since the sampling will
_N_
be performed regardless of the underlying distribution of the latent variables and so will create
irrelevant samples. To further illustrate this, we propose to 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 V (z) = ( det G(z))[−][1] all along the curve

_i.e. solving the minimization problem_

p


_V (γ(t))dt_ s.t. _γ(0) = z1, γ(1) = z2 ._ (8)


inf


1This may be performed with k-medoids algorithm for instance.


-----

(a)



10


10


_−1_

_−2_

_−3−3_ _−2_ _−1_

Circles
Rings

Affine

(a)

Ours

Affine

(b)

Ours


_−1_

_−2_

_−3−3_ _−2_ _−1_


_−1_

_−2_

_−3−3_ _−2_ _−1_


_−2_

|Col1|1|
|---|---|
|||
|Ours Affine Circles Rings||

|(b)|Col2|
|---|---|
|||
|Ours Affine Circles Rings||


Ours Affine Circles Rings


Ours Affine Circles Rings


Figure 1: Top: Visualization and interpolation in a 2D latent space learned by a vanilla VAE trained
with binary images of rings and disks. The log of the metric volume element det G(z) (also

proportional to the log of the density we propose to sample from) is represented in gray scale. Top

p

_right: The Riemannian distance from a 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
by solving Eq. 8. Bottom: Decoded samples along the interpolation curves.

Below are presented the decoded samples all along the interpolation curves. Thanks to those interpolations we can see that i) the latent space seems to really have a specific geometrical structure since
decoding all along the interpolation curves obtained by solving Eq. 8 leads 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 to the
data i.e. where information is available and so does not produce irrelevant images when decoded.

5 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 or 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 et al., 2018), Regularized Autoencoders (Ghosh et al., 2020) with
either L2 decoder’s parameters regularization (RAE-L2), gradient penalty (RAE-GP), spectral normalization (RAE-SN) or simple L2 latent code regularization (RAE), a vamp-prior VAE (VAMP)
(Tomczak & Welling, 2018), a Hamiltonian VAE (HVAE) (Caterini et al., 2018), a geometry-aware
VAE (RHVAE) (Chadebec et al., 2020) and an Autoencoder (AE). We elect these models since they
use different ways to generate the data using either the prior or 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 are 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 E for the comprehensive experimental setup. Figure 2 shows a qualitative comparison between the resulting generated samples for MNIST
and CELEBA, the same plots are made available in Appendix D for SVHN and CIFAR 10. Interest

-----

MNIST CELEBA

AE - N

VAE - N

WAE

VAMP

HVAE

RHVAE

AE - GMM

VAE - GMM

RAE

VAE - Ours

Figure 2: Generated samples with different models and generation processes. Generated samples
with RAE variants are also provided in Appendix D.

Gen. Near. train Near. rec. Gen. Near. train Near. rec. Gen. Near. train Near. rec. Gen. Near. train Near. rec.

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-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, we clearly see that using the proposed method represents a strong improvement
of the generation process from a vanilla VAE when compared to the samples coming from a normal
prior (second row). This insists on the fact that even the simplest VAE model actually contains a lot
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 a 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 be impossible. These observations are even more supported by Table 1 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 an interesting idea since it allows for a better expressiveness and latent
space prospecting. For instance, on MNIST the FID falls from 40.7 with the prior to 13.1 when
using a GMM. Nonetheless, with the proposed method we are 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 on the 3 other datasets. Impressively, in almost all cases, the proposed generation method


-----

Table 1: 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-post density estimation by fitting a 10-component mixture of Gaussian in the latent space.

|Model|MNIST (16) FID ↓ PRD ↑|SVHN (16) FID ↓ PRD ↑|CIFAR 10 (32) FID ↓ PRD ↑|Celeba (64) FID ↓ PRD ↑|
|---|---|---|---|---|
|AE - N (0, 1) WAE VAE - N (0, 1) VAMP HVAE RHVAE|46.41 0.86/0.77 20.71 0.93/0.88 40.70 0.83/0.75 34.02 0.83/0.88 15.54 0.97/0.95 36.51 0.73/0.28|119.65 0.54/0.37 49.07 0.80/0.85 83.55 0.69/0.55 91.98 0.55/0.63 98.05 0.64/0.68 121.69 0.55/0.41|196.50 0.05/0.17 132.99 0.24/0.52 162.58 0.10/0.32 198.14 0.05/0.11 201.70 0.13/0.21 167.41 0.12/0.22|64.64 0.29/0.42 54.56 0.57/0.55 64.13 0.27/0.39 73.87 0.09/0.10 52.00 0.38/0.58 55.12 0.45/0.56|
|AE - GMM RAE (GP) RAE (L2) RAE (SN) RAE VAE - GMM|9.60 0.95/0.90 9.44 0.97/0.98 9.89 0.97/0.98 11.22 0.97/0.98 11.23 0.98/0.98 13.13 0.95/0.92|54.21 0.82/0.83 61.43 0.79/0.78 58.32 0.82/0.79 95.64 0.53/0.63 66.20 0.76/0.80 52.32 0.82/0.85|130.28 0.35/0.58 120.32 0.34/0.58 123.25 0.33/0.54 114.59 0.32/0.53 118.25 0.35/0.57 138.25 0.29/0.53|56.07 0.32/0.48 59.41 0.28/0.49 54.45 0.35/0.55 55.04 0.36/0.56 53.29 0.36/0.58 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 or outperform peers both in terms of FID and PRD scores. Finally, we check
if 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
not 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 between 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 consideration 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 a
10-component GMM and 0.6s for prior based methods on a single GPU V100-16GB.


MNIST


CIFAR 10


300
200

100

50

20

10


200

150

100


|Col1|WAE - N (0, 1) RAE-GP - N (0, 1) RAE-L2 - N (0, 1) RAE-SN - N (0, 1) VAE - N (0, 1) WAE - GMM RAE-GP - GMM RAE-L2 - GMM RAE-SN - GMM VAE - GMM VAE - Ours|
|---|---|


0.1 0.5 1.0 5.0 0.1 0.5 1.0 5.0

_×10[4]_ _×10[4]_

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 progressively decrease the size of the training set to see the robustness of the different sampling methods
according to the number of samples. This experiment is 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, CIFAR10 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 ELBO


-----

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|F1 Balanced Accuracy AD CN|
|---|---|
|Original (unbalanced) Original (resampled)|66.2 ± 7.6 47.6 ± 15.8 87.3 ± 2.0 81.8 ± 2.6 72.1 ± 3.6 88.0 ± 2.3|
|AE - N (0, Id) WAE VAE - N (0, Id) VAMP HVAE RHVAE|50.0 ± 0.0 0.0 ± 0.0 84.1 ± 0.0 57.4 ± 9.7 21.0± 24.5 84.4 ± 2.3 51.8 ± 3.8 6.1 ± 11.8 84.6 ± 1.1 83.1 ± 2.6 70.4 ± 3.6 82.2 ± 4.7 56.3 ± 7.9 19.6 ± 21.7 85.4 ± 1.7 68.0 ± 10.9 47.0 ± 24.2 85.1 ± 3.3|
|AE - GMM RAE (GP) RAE (L2) RAE (SN) RAE VAE - GMM|82.4 ± 2.3 69.5 ± 3.1 82.0 ± 3.6 63.9 ± 9.8 46.5 ± 15.9 70.6 ± 19.6 74.1 ± 6.0 60.6 ± 9.5 82.1 ± 5.9 62.3 ± 8.9 37.8 ± 22.6 80.1 ± 7.9 69.3 ± 8.1 53.8 ± 12.9 80.0 ± 10.7 83.0 ± 3.6 71.4 ± 4.3 85.3 ± 3.0|
|VAE - Ours|85.4 ± 2.5 74.7 ± 3.5 87.3 ± 2.7|



on the validation set the size of which is set as 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 are presented in Appendix F. Again, the proposed sampling
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-post density estimation seems
to still enhance the generation capability of the model when the number of training samples remains
quite high (≥5k), this gain seems to worsen when the dataset size reduces as illustrated on CIFAR.

5.3 GENERATION WITH COMPLEX DATA

Finally, we also propose to stress the proposed generation procedure in a day-to-day scenario where
the limited data regime is 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 to the lack of a large test set, we 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 is trained on each label of the train set and used to
generate 2k new samples per class. Then a 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 2
shows the mean balanced accuracy and F1 scores across 20 runs. These metrics provide a 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 is 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.

6 CONCLUSION

In this paper, we provided a geometric understanding of the latent space learned by a 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-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 a
task with state of the art augmentation methods.


-----

REPRODUCIBILITY STATEMENT

In order to make the method and the proposed experiments reproducible, we provide in Appendix
E the complete experimental set-up and in Appendix C pseudo-code algorithms detailing the implementation from a practical point of view. We also provide an implementation in the supplementary
material.


-----

REFERENCES

Alexander A Alemi, Ian Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information
bottleneck. arXiv preprint arXiv:1612.00410, 2016.

Jyoti Aneja, Alexander Schwing, Jan Kautz, and Arash Vahdat. NCP-VAE: Variational autoencoders
with noise contrastive priors. arXiv:2010.02917 [cs, stat], 2020.

Georgios Arvanitidis, Lars Kai Hansen, and S¨oren Hauberg. Latent space oddity: On the curvature
of deep generative models. In 6th International Conference on 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 autoencoders. _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 autoencoders.
_arXiv:1509.00519 [cs, stat], 2016._

Manfredo Perdigao do Carmo. Riemannian Geometry. Birkh¨auser, 1992.

Anthony L Caterini, Arnaud Doucet, and Dino Sejdinovic. Hamiltonian variational auto-encoder.
In Advances in Neural Information Processing Systems, pp. 8167–8177, 2018.

Cl´ement Chadebec, Cl´ement Mantoux, and St´ephanie Allassonni`ere. Geometry-aware hamiltonian
variational auto-encoder. arXiv:2010.11518, 2020.

Cl´ement Chadebec, Elina Thibeau-Sutre, Ninon Burgos, and St´ephanie Allassonni`ere. Data Augmentation in High Dimensional Low Sample Size Setting Using a Geometry-Based Variational
Autoencoder. arXiv preprint arXiv:2105.00026, 2021.

Nutan Chen, Alexej Klushyn, Richard Kurle, Xueyan Jiang, Justin Bayer, and Patrick Smagt. Metrics for deep generative models. In International Conference on Artificial Intelligence and Statis_tics, 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.02731,
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 autoencoders. In International Conference on Machine Learning, pp. 1078–1086. PMLR, 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 Tomczak. Hyperspherical 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 Garnelo, Matthew C. H. Lee, Hugh Salimbeni,
Kai Arulkumaran, and Murray Shanahan. Deep unsupervised clustering with gaussian mixture
variational autoencoders. arXiv:1611.02648 [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._


-----

Luca Falorsi, Pim de Haan, Tim R. Davidson, Nicola De Cao, Maurice Weiler, Patrick Forr´e, and
Taco S. Cohen. Explorations in homeomorphic variational auto-encoding. arXiv:1807.04689 [cs,
_stat], 2018._

Partha Ghosh, Mehdi SM Sajjadi, Antonio Vergari, Michael Black, and Bernhard Sch¨olkopf. 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. A Geometric take on Metric 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://proceedings.neurips.cc/paper/2012/file/](https://proceedings.neurips.cc/paper/2012/file/ec5aa0b7846082a2415f0902f0da88f2-Paper.pdf)
[ec5aa0b7846082a2415f0902f0da88f2-Paper.pdf.](https://proceedings.neurips.cc/paper/2012/file/ec5aa0b7846082a2415f0902f0da88f2-Paper.pdf)

Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter.
Gans trained by a two time-scale update rule converge to a 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 a
constrained variational framework. ICLR, 2(5):6, 2017.

Matthew D Hoffman and Matthew J Johnson. Elbo surgery: yet another way to carve up the variational evidence lower bound. In Workshop in Advances in Approximate Bayesian Inference, NIPS,
volume 1, pp. 2, 2016.

Michael I 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 autoencoders 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.6980, 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 Cseke. Learning Hierarchical 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 of 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 a Riemannian_
_geometry setting. PhD Thesis, Sorbonnes universit´es, 2019._

Daniel S. Marcus, Tracy H. Wang, Jamie Parker, John G. Csernansky, John C. Morris, and Randy L.
Buckner. Open access series of imaging studies (OASIS): Cross-sectional MRI data in young,
middle aged, nondemented, and demented older adults. Journal of Cognitive Neuroscience, 19
(9):1498–1507, 2007.


-----

Emile Mathieu, Charline Le Lan, Chris J Maddison, Ryota Tomioka, and Yee Whye Teh. Continuous hierarchical representations with poincar´e variational auto-encoders. In Advances in neural
_information processing systems, pp. 12565–12576, 2019._

Eric Nalisnick, Lars Hertel, and Padhraic Smyth. Approximate inference for deep latent gaussian
mixtures. In NIPS Workshop on Bayesian Deep Learning, volume 2, pp. 131, 2016.

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 measurements. Journal of Mathematical Imaging and Vision, 25(1):127–154, 2006. ISSN 0924-9907,
1573-7683. doi: 10.1007/s10851-006-6228-4.

Alexander Rakowski and Christoph Lippert. Disentanglement and local directions of variance. In
_Joint European Conference on Machine Learning and Knowledge Discovery in 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 approximate 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 pca directions (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. In 32nd Conference on Neural Information Processing Systems (NeurIPS
_2018), pp. 5228–5237, 2019._

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 on Computer Vision and Pattern Recognition Workshops
_(CVPRW), pp. 428–4288. IEEE, 2018. ISBN 978-1-5386-6100-0. doi: 10.1109/CVPRW.2018._
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 Tolstikhin, O Bousquet, S Gelly, and B Sch¨olkopf. Wasserstein auto-encoders. In 6th International
_Conference on Learning Representations (ICLR 2018), 2018._

Jakub Tomczak and Max Welling. Vae with a vampprior. In International Conference on Artificial
_Intelligence and Statistics, pp. 1214–1223. PMLR, 2018._


-----

Zhisheng Xiao, Karsten Kreis, Jan Kautz, and Arash Vahdat. Vaebm: A symbiosis between variational 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–6449, 2019._

Cheng Zhang, Judith B¨utepage, Hedvig Kjellstr¨om, and Stephan Mandt. Advances in variational
inference. IEEE Transactions on Pattern Analysis and Machine Intelligence, 41(8):2008–2026,
2018.


-----

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 a smooth inner
product g : p →⟨·|·⟩p defined on each tangent space TpM of the manifold with p ∈M. A chart (or
coordinate system) (U, x) is a 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−dimensional manifold if for each
charthomeomorphic to V ⊂ **R[d]. This means that there exists a neighborhood R[d]. Given p** _U_, a chart ϕ : (x[1], . . ., x[d]) induces a basis U of each point p∂x∂ ∈M[1][, . . .,] such that∂x∂ _[d]_ _U is_
_∈_ _p_ [on the]

tangent space Tp . Hence, the metric of a Riemannian manifold can be locally represented in the 
_M_
chart ϕ as a positive definite matrix G(p) = (gi,j)p,0≤i,j≤d = (⟨ _∂x[∂]_ _[i][ |][ ∂]∂x[j][ ⟩][p][)][0][≤][i,j][≤][d][ for each point][ p]_

of the manifold. That is for v, w ∈ _TpM and p ∈M, the inner product writes ⟨u|w⟩p = u[⊤]G(p)w._

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 [2] seen as a subspace of R[3]. The second one which is adopted in this paper
is the intrinsic view. In the intrinsic view the manifold is studied using its underlying structure and
so the length of a curve γ : R →M traveling in the manifold cannot be interpreted using the
Euclidean distance but requires to use the metric defined onto the manifold itself. Let z1, z2
be two points of the manifold, and γ be a curve traveling in M parametrized by t ∈ [0, 1] such that ∈M
_γ(0) = z1 and γ(1) = z2. Then, the length of γ is given by_


_L(γ) =_


_γ˙_ (t) _γ(t)dt =_
_∥_ _∥_


_γ˙_ (t) _γ˙_ (t) _γ(t)dt_
_⟨_ _|_ _⟩_


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(z1, z2) = min (γ) s.t. _γ(0) = z1, γ(1) = z2_ (9)
_γ_ _L_

The manifold M is said to be geodesically complete if all geodesic curves can be extended to R.
For any p, the exponential map at p maps a vector v of the tangent space Tp to a point of
_∈M_ _M_
the manifold ˜p ∈M such that the geodesic starting at p with initial velocity v reached ˜p at time 1.


_Tp_
_M_ _→M_
_v_ Expp(v) = γ(p,v)(1) _[,]_
_→_


Expp :


where γ(p,v)(1) means γ(0) = p and ˙γ(0) = v. Saying that the manifold is geodesically
_M_
_complete means that the exponential maps is defined on the entire tangent space Tp_ for each
_M_
element p ∈M.

In Pennec (2006), 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 ϕ = (x1, · · ·, xd), a volume form dVolg can be defined as the
_d-form:_
dVolg = det(gi,j) dx[1] _∧· · · ∧_ _dx[d]_

This represents an infinitesimal volume element on each tangent space and so a measure on theq
manifold M
_dM(ϕ[−][1](x)) =_ det g(ϕ[−][1](x)) dx

It follows that we are able to integrate functions fp : _U_ R on a given chart (U, ϕ).
_M ⊃_ _→_


det g(ϕ[−][1](x)) dx[1] _· · · dx[d]_


_f_ (ϕ[−][1](x))


_f_ (p)dM(p) =

_ϕ(ZU_ )


This notion can 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 on the manifold. We recall such definition from Pennec (2006) below


-----

**Definition 1 Let B(M) be the Borel σ-algebra of M. The random point p has a probability density**
_function ρp if:_


_∀X ∈B(M), P(p ∈X_ ) =

_and_
Z


_ρ(p)dM(p)_

Z

_X_

_ρ(p)dM(p) = 1_


Finally, given a chart ϕ = (z1, _, zn) defined on the whole manifold_ and a random point p
_· · ·_ _M_
on M, the point z = ϕ(p) is a random point whose density ρ[′]z [may be written with respect to the]
Lebesgue measure as such (Pennec, 2006):


_ρ[′]z[(][z][) =][ ρ][p][(][ϕ][−][1][(][z][))]_


det g(ϕ[−][1](z)) (10)


-----

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 = R[d] and the Riemannian metric whose local representation is given by
Eq. 7, any geodesic curve γ :]a, b[→M is actually extensible to R. Let us consider a geodesic curve
_γ such that γ cannot be extended to R. There exists I =]a, b[ such that I is the maximum definition_
domain of γ. We will show that with such an assumption we will end up with a contradiction. We
recall the shape of the Riemannian metric:


**Σ[−][1](xi) · ωi(z) + λ · Id,**
_i=1_

X


**G(z) =**


Since Σi are positive definite matrices we have z[⊤]Σiz > 0, ∀ **_z ∈M −{0}. We further have_**
_ωi(z) > 0,_ **_z_** since the manifold is geodesically complete. Let t0 ]a, b[ we therefore have
for any t ∈] ∀a, b[ ∈M. _∈_


_γ˙_ (t)[⊤]Σ(xi)[−][1]γ˙ (t) _ωi(γ(t))_
_·_
_i=1_

X


_λ · ∥γ˙_ (t)∥2[2] _[≤]_ _[λ][ · ∥]γ[˙]_ (t)∥2[2] [+]


_≤∥γ˙_ (t)∥γ[2](t) [=][ ∥]γ[˙] (t0)∥γ[2](t0) _[,]_

where the last equality comes for the constant speed of geodesic curves (Carmo, 1992). Hence we
have:

_γ(t0)_ _γ(t0)_
_γ(t)_ _γ(t0)_ 2 _∥_ _t_ _t0_ _._
_∥_ _−_ _∥_ _≤_ _[∥]_ [˙] _√λ_ _· |_ _−_ _|_

This means that for any t ∈]a, b[ the geodesic curve γ remains within a compact set. We show now
that the curve can actually be extended. Let us define the sequence tn
_−n−→∞−−→_ _[b][. Since the geodesic]_

curves have a constant speed the set I = {(tn, ˙γ(tn)}n∈N is compact. Moreover, using CauchyLipchitz theorem, we can find ε > 0 such that for any n ∈ N, the geodesic γ can be extended to
]domain can be extended totn − _ε, tn + ε[. Since, tn ] can be as close asa, b +_ 2[ε] [[][.] _b as desired we can assure that the curve definition_


-----

C THE GENERATION PROCESS ALGORITHM - IMPLEMENTATION DETAILS

In this appendix, we provide pseudo-code algorithms explaining how to build the metric from a
trained VAE and how to use the proposed sampling process. Noteworthy is the fact that we do not
amend the training process of the vanilla VAE which remains pretty simple and stable.

C.1 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. 7 below


**Σ[−][1](xi) · ωi(z) + λ · Id,**
_i=1_

X


**G(z) =**


where


_−_ [dist][Σ][−][1][(][x]ρ[i][2][)][(][z][,][ µ][i][)][2]


_−_ [(][z][ −] **_[µ][i][)][⊤][Σ][−]ρ[1][2][(][x][i][)(][z][ −]_** **_[µ][i][)]_**


_ωi(z) = exp_


= exp


**Algorithm 1 Building the Metric from a Trained Model**

**Input: A trained VAE model m, the training dataset X**, λ
**for xµii ∈X, Σi = do m(xi)** _▷_ Retrieve training embeddings and covariance matrices

**end for**
Select k centroids ci in the µi _▷_ e.g. with k-medoids

Get corresponding covariance matrices Σi
_ρ_ max min _▷_ Set ρ to the max distance between two closest neighbors
_←_ _i_ _j=i_
_̸_ _[∥][c][i][ −]_ **_[c][j][∥][2]_**

Build the metric using Eq. 7


**Σ[−]i** [1] _ωi(z) + λ_ **_Id_**

_·_ _·_
_i=1_

X


**G(z) =**


**Return G** _▷_ Return G as a function

As is standard in VAE implementations, we assume that the covariance matrices Σi given by the
VAE are diagonal and that the encoder outputs a mean vector and the log of the diagonal coefficients.
In the implementation, the exponential is then applied to recover the Σi so that no singular matrix
arises.

C.2 SAMPLING PROCESS

Further to the description performed in the paper, we provide here a detailed algorithm stating the
main steps of the generation process.

C.2.1 THE HMC SAMPLER

In the sampling process we propose to rely on the Hamiltonian Monte Carlo sampler to sample from
the Riemanian uniform distribution. In a nutshell, the HMC sampler aims at sampling from a target
distribution ptarget(z) with z ∈ R[d] using Hamiltonian dynamics. The main idea behind such a
sampler is to introduce an auxiliary random variable v (0, Id) independent from z and mimic
_∼N_
the behavior of a particle having z (resp. v) as location (resp. velocity). The Hamiltonian of the
particle then writes
_H(z, v) = U_ (z) + K(v),
where U (z) is the potential energy of such a particle and K(v) is its kinetic energy both given by

_U_ (z) = log ptarget(z), _K(v) = [1]_
_−_ 2 **_[v][⊤][v]_**


-----

The following Hamilton’s equations govern the evolution in time of the particle.
( _∂H∂H∂∂((vzzz,,vv))_ == **_v,_** **_z log ptarget(z) ._** (11)

_−∇_

In order to integrate these equations, recourse to the leapfrog integrator is needed and consists in
applying nlf times the following equations.

**_v(t +_** _[ε]2[lf]_ [)] = v(t) + _[ε]2[lf]_

**_z(t + εlf_** ) = z(t) + εlf _[· ∇] v([z]t +[ log][ε][ p]2[lf]_ [target][)][,] [(][z][(][t][))][,] (12)

 **_v(t + εlf_** ) = v(t + _[ε]2[lf]_ [) +] · _[ ε]2[lf]_

_[· ∇][z][ log][ p][target][(][z][(][t][ +][ ε][lf]_ [))][,]

where εlf is called the leapfrog step size. This algorithm produces a proposal (z, **_v) that is accepted_**



with probability α where
e e

_α = min_ 1, exp _H(z, v) −_ _H(z,_ **_v)_** _._
  

This procedure is then repeated to create an ergodic Markov chain (z[n]) converging to the distribution

e e

_ptarget (Neal & others, 2011; Duane et al., 1987; Liu, 2008; Girolami & Calderhead, 2011)._

C.3 THE PROPOSED ALGORITHM

In 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) =_ Riem(z) = [1]
_U_ _C_


det G(z) . (13)


The log density follows

log p(z) = [1]

2 [log det][ G][(][z][)][ −] [log][ C,]

In such a case, the Hamiltonian writes


_H(z, v) =_ log p(z) + [1]
_−_ 2 **_[v][⊤][v][,]_**

and Hamilton’s equations become
_∂H∂(vz,v)_ = **_v,_**
( _∂H∂(zzi,v)_ = _∂zi_ = 2 [tr] **G[−][1](z)** _[∂][G]∂z[(]i[z][)]_

_−_ _[∂]_ [log][ p][(][z][)] _−_ [1]




Since the covariance matrices are supposed to be diagonal as is standard in VAE implementations,
the computation of the inverse metric is straightforward. Moreover, since G(z) is smooth and has
a closed form, it can be differentiated with respect to z pretty easily. Now, the leapfrog integrator
given in Eq. 12 can be used and the acceptance ratio α 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 α. We provide a pseudo-code of the proposed sampling procedure
in Alg. 2. A typical choice in the sampler’s hyper-parameters used in the paper is N = 100, nlf = 10
and εlf = 0.01. The initialization of the chain can be done either randomly or on points that belong
to the manifold (i.e. the centroids ci or µi).


-----

**Algorithm 2 Proposed Sampling Process**

**Input: The metric function G, hyper-parameters of the HMC sampler (chain length N**, number
of leapfrog steps nlf, leapfrog step size εlf )
**Initialization: z** _▷_ Initialize the chain

**for i = 1 →** _N do_

**_v_** (0, Id) _▷_ Draw a velocity
_∼N_

_H0_ _H(z, v)_ _▷_ Compute the starting Hamiltonian

**_z0_** _←z_
**for ← k = 1 ←** _nlf do_

**_v¯_** **_v_** 2

**_z ←_** **_z + − ε[ε]lf[lf]_** _[· ∇] ¯v_ **_[z][H][(][z][,][ v][)]_** _▷_ Leapfrog step Eq. 12

**_v ←_** **_v¯_** 2 _·_ **_z, ¯v)_**

**_v ←_** **_v −_** _[ε][lf]_ _[· ∇][z][H][(][e]_

e ←

**_z_** **_z_**

e ←

**end for**

e

_H_ _H(z,_ **_v)_** _▷_ Compute the ending Hamiltonian
_←_ e

Accept **_z with probability α = min_** 1, exp(H0 _H)_
_−_

**if Acceptede** **then e**  

**_z_** ez
_←_

**else**

**_z_** **_z0_**
_←_ e

**end if**

**end for**
**Return z**


-----

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 5 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 in Figure 7.

Figure 5: 100 samples with the proposed method on MNIST dataset.

Figure 6: 100 samples with the proposed method on Celeba dataset.


-----

MNIST CELEBA

AE - N

VAE - N

WAE

VAMP

HVAE

RHVAE

AE - GMM

VAE - GMM

RAE (GP)

RAE (L2)

RAE (SN)

RAE

VAE - Ours

Figure 7: Generated samples with different models and generation processes.


-----

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

AE - N

VAE - N

WAE

VAMP

HVAE

RHVAE

AE - GMM

VAE - GMM

RAE (GP)

RAE (L2)

RAE (SN)

RAE

VAE - Ours

Figure 8: Generated samples with different models and generation processes.

Figure 9: Closest element in the training set (Near.) to the generated one (Gen.) with the proposed
method.


-----

E EXPERIMENTAL SET-UP

The RAEs, VAEs and AEs are trained for 100 epochs for SVHN, MNIST[2] and Celeba and 200 on
CIFAR10. 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 CIFAR10. 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 is set to 5e−4) with an Adam optimizer
(Kingma & Ba, 2014). We also use a 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/5[th] 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 3 and are inspired by Ghosh et al. (2020). 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 ρ is set to ρ = max min
_i_ _j≠_ _i_ _[∥][c][i][ −]_ **_[c][j][∥][2][ to ensure some]_**

_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
a batch size of 100 and a starting learning rate of 1e−4. Again a 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 to the classifier, it is trained with a batch size of 200 for 50
epochs with a starting learning rate of 1e−4 and Adam optimizer. A scheduler reducing the learning
rate by half every 5 epochs if the validation loss does not improve is 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 4. MRIs are only pre-processed such that the maximum value of a voxel is 1
and the minimum 0 for each data point.

Table 3: Neural networks used for the encoder and decoders of VAEs in the benchmarks

MNIST [CIFAR10] SVHN CELEBA OASIS

ENCODER (1[3], 32, 32) (3, 32, 32) (3, 64, 64) (1, 208, 176)

CONV(128, (4, 4), STRIDE=2) LINEAR(1000) CONV(128, (5, 5), STRIDE=2) CONV(64, (5, 5), STRIDE=2)
LAYER 1 BATCH NORMALIZATION RELU BATCH NORMALIZATION RELU
RELU RELU

CONV(256, (4, 4), STRIDE=2) LINEAR(500) CONV(256, (5, 5), STRIDE=2) CONV(128, (5, 5), STRIDE=2)
LAYER 2 BATCH NORMALIZATION RELU BATCH NORMALIZATION RELU
RELU RELU

CONV(512, (4, 4), STRIDE=2) CONV(512, (5, 5), STRIDE=2) CONV(256, (5, 5), STRIDE=2)
LAYER 3 BATCH NORMALIZATION LINEAR(500, 16*) BATCH NORMALIZATION RELU
RELU RELU

CONV(1024, (4, 4), STRIDE=2) CONV(1024, (5, 5), STRIDE=2) CONV(512, (5, 5), STRIDE=2)
LAYER 4 BATCH NORMALIZATION -  BATCH NORMALIZATION RELU
RELU RELU

CONV(1024, (5, 5), STRIDE=2)
LAYER 5 LINEAR(4096, 16*) -  LINEAR(16384, 64*)
RELU

LAYER 6 -  -  -  LINEAR(4096, 16*)

DECODER (16 [32]) (16) (64) (16)

LINEAR(65536) LINEAR(500) LINEAR(65536) LINEAR(65536)
LAYER 1
RESHAPE(1024, 8, 8) RELU RESHAPE(1024, 8, 8) RESHAPE(1024, 8, 8)

CONVT(512, (4, 4), STRIDE=2) LINEAR (1000) CONVT(512, (5, 5), STRIDE=2) CONVT(512, (5, 5), STRIDE=(3, 2))
LAYER 2 BATCH NORMALIZATION RELU BATCH NORMALIZATION RELU
RELU RELU

CONVT(256, (4, 4), STRIDE=2) LINEAR(3072) CONVT(256, (5, 5), STRIDE=2) CONVT(256, (5, 5), STRIDE=2)
LAYER 3 BATCH NORMALIZATION RESHAPE(3, 32, 32) BATCH NORMALIZATION RELU
RELU SIGMOID RELU

CONVT(3, (4, 4), STRIDE=1) CONVT(128, (5, 5), STRIDE=2) CONVT(128, (5, 5), STRIDE=2)
LAYER 4 BATCH NORMALIZATION -  BATCH NORMALIZATION RELU
SIGMOID RELU

CONVT(3, (5, 5), STRIDE=1) CONVT(64, (5, 5), STRIDE=2)
LAYER 5 -  -  BATCH NORMALIZATION RELU
SIGMOID

CONVT(1, (5, 5), STRIDE=1)
LAYER 6 -  -  - 
RELU

2MNIST images are re-scaled to 32x32 images with a 0 padding.


-----

Table 4: Neural Network used for the classifier in Sec. 5.3

OASIS CLASSIFIER

INPUT SHAPE (1, 208, 176)

CONV(8, (3, 3), STRIDE=1)
BATCH NORMALIZATION

LAYER 1

LEAKYRELU
MAXPOOL(2, STRIDE=2)

CONV(16, (3, 3), STRIDE=1)
BATCH NORMALIZATION

LAYER 2

LEAKYRELU
MAXPOOL(2, STRIDE=2)

CONV(32, (3, 3), STRIDE=2)
BATCH NORMALIZATION

LAYER 3

LEAKYRELU
MAXPOOL(2, STRIDE=2)

CONV(64, (3, 3), STRIDE=2)
BATCH NORMALIZATION

LAYER 4

LEAKYRELU
MAXPOOL(2, STRIDE=2)

LINEAR(256, 100)
LAYER 5
RELU

LINEAR(100, 2)
LAYER 6
SOFTMAX


-----

F DATASET SIZE SENSIBILITY ON SVHN

In Figure 10, we show the same plot for SVHN as in Sec. 5.2. Again the proposed method appears
to be part of the most robust generation procedures to dataset size changes.

SVHN

WAE - N (0, 1)
RAE-GP - N (0, 1)
RAE-L2 - N (0, 1)
RAE-SN - N (0, 1)
VAE - N (0, 1)
WAE - GMM
RAE-GP - GMM
RAE-L2 - GMM
RAE-SN - GMM
VAE - GMM
VAE - Ours


0.1 0.5 1.0 5.0

_×10[4]_

Figure 10: FID score evolution according to the number of training samples.


-----

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 11 the evolution of the FID according to the number of centroids in the metric
(left) and the variation of FID according to the choice in the centroids. As expected, choosing a
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.


110
100

90
80
70
60
50
40
30
20
10


80

70

60

50

40

30

20

10


10 100 1000 10000

mnist
celeba

Number of centroids in the metric


mnist
celeba


Centroids choice


Figure 11: Left: FID score evolution according to 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 (k = 1000).

To assess the variability of the generated samples, we propose to analyze some generated samples
when only 2 centroids are considered. In Figure 12, 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 to 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 of samples. These samples
are not simply resampled train images or a simple interpolation between selected centroids as some
of the generated samples have attributes such as glasses that are not present in the images of the
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 in 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).


-----

G.2 INFLUENCE OF λ IN THE METRIC

In this section, we also assess the influence of the regularization factor λ in Eq. 7 on the resulting
sampling. To do so, we generate 10k samples using the proposed method on both MNIST and
Celeba datasets for values of λ ∈ [1e[−][6], 1e[−][4], 1e[−][2], 1e[−][1], 1]. Then, we compute the FID against
the test set. Each time, we consider k = 1000 centroids in the metric. As shown in Figure 13, the
influence of λ remains limited. In practice, λ is mainly here to avoid pathological cases such as the
metric collapsing to 0 far from the centroids. In the implementation, a typical choice for λ is 1e[−][2].


80

70

60

50

40

30

20

10


10[−][6] 10[−][5] 10[−][4] 10[−][3] 10[−][2] 10[−][1] 10[0]

mnist
celeba

Value of λ


Figure 13: FID score evolution according to the value of λ in the metric (Eq.7).

G.3 THE CHOICE OF ρ


In the experiments presented, the smoothing factor ρ in Eq. 7 is set to the value of the maximum
distance between two closest centroids ρ = max min
_i_ _j≠_ _i_ _[∥][c][j][ −]_ **_[c][i][∥][2][. This choice is motivated by the]_**

fact that we wanted to build a smooth metric and so ensure some smoothness of the manifold while
trying to interpolate faithfully between the metric tensors Gi = Σ[−]i [1][. In particular, a small value]
of ρ 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,
setting a high value for ρ would have biased the interpolation and the value of the metric at a µi.
As a result, G(µi) might have been very different from the one observed Σ[−]i [1] since the other µj
would have had a strong influence on its value. The proposed value for ρ appeared to work well in
practice.


-----

H OTHER CLASSIFICATION METRICS

In Table 5 are presented some further classification results for each considered model while generated samples using each generation procedure are made available in Figure 14. On OASIS database
the proposed method appears to produce visually the sharpest samples. This better generation performance is also supported by the classification metrics provided in Table 2 and Table 5.

Table 5: Classification results averaged on 20 independent runs. For the VAEs, the classifier is
trained on 2K generated samples per class.

*unbalanced

|Generation method|Precision Recall 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 78.5 5.2 85.1 4.2 ± ± ± ± ±|
|AE - N WAE VAE - N VAMP HVAE RHVAE|50.0 0.0 0.0 0.0 72.6 0.0 0.0 0.0 100.0 0.0 ± ± ± ± ± 57.4 9.7 48.5 42.8 76.7 6.1 19.3 27.5 95.4 9.3 ± ± ± ± ± 51.8 3.8 38.0 47.3 73.4 1.7 3.7 7.8 99.8 0.7 ± ± ± ± ± 83.1 2.6 56.3 5.2 97.5 2.1 94.8 4.7 71.5 7.4 ± ± ± ± ± 56.3 7.9 48.7 41.7 75.5 3.8 13.9 17.6 98.6 2.2 ± ± ± ± ± 68.0 10.9 56.1 25.3 83.0 7.5 46.7 30.2 89.2 10.6 ± ± ± ± ±|
|AE - GMM RAE (GP) RAE (L2) RAE (SN) RAE VAE - GMM|82.4 2.3 55.8 4.9 96.8 2.4 93.3 5.6 71.5 6.2 ± ± ± ± ± 63.9 9.8 45.3 18.5 84.2 8.6 60.9 28.6 67.0 24.9 ± ± ± ± ± 74.1 6.0 57.8 10.1 88.3 5.2 70.0 18.7 78.3 11.7 ± ± ± ± ± 62.3 8.9 43.1 24.9 80.6 6.6 41.7 30.1 82.9 16.4 ± ± ± ± ± 69.3 8.1 56.2 13.5 85.2 6.2 60.0 24.0 78.5 17.5 ± ± ± ± ± 83.0 3.6 60.7 5.4 94.9 3.7 88.0 9.5 77.9 5.9 ± ± ± ± ±|
|VAE - Ours|85.4 2.5 64.0 5.3 95.8 2.2 90.4 5.6 80.3 5.1 ± ± ± ± ±|


-----

OASIS

Train

VAE - N

WAE

VAMP

HVAE

RHVAE

VAE - GMM

RAE (GP)

RAE (L2)

RAE (SN)

RAE

VAE - Ours

Figure 14: Generated samples with different models and generation processes.


-----

I LINK BETWEEN THE RIEMANNIAN VAE AND VANILLA VAE

We assume as in (Ghosh et al., 2020) that a VAE is essentially an autoencoder regularized with
noise. Hence, a Riemannian-based VAE could also be seen as a regularized autoencoder but the
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 x ∈X ⊂ R[D]
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 = (R[d], G) where G is an unknown Riemannian
metric. Now the goal of the Riemannian VAE would be the same as a 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.

To to so and similarly to autoencoder models, it would be assumed that there exist eϕ : R[D] _→M_
a parametrized encoding function mapping the input data onto the manifold eϕ(x) = µ and
_∈M_
_dθ : M →_ R[D] a parametrized decoding function that maps back the latent codes to the data space.
In such a case, the main objective would be to find ϕ and θ such that the reconstruction loss is
minimized
min _x_ _,_ (14)
_ϕ,θ_ _[L][REC][ = min]ϕ,θ_ _[l][(][x][, d][θ][(][e][ϕ][(][x][))][,]_ _∈X_

where l 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 a smooth latent space meaning that small variations in the latent space do not change
completely the output of the decoder, the decoder is also regularized using a Riemannian Gaussian
noise. That is imposing that
_dϕ(z) ≈_ **_x,_** **_z ∼Nriem[G]_** [(][z][|][µ][, σ][)][,]
where

riem[G] [(][z][|][σ,][ µ][) = 1] _, C =_ exp _d_ **_z,_** (15)
_N_ _C_ [exp] _−_ [dist][G]2[(]σ[z][,][ µ][)][2] _−_ [dist][G]2[(]σ[z][,][ µ][)][2] _M_

  Z  

_M_

and the reconstruction loss in Eq. 14 would become
min _x_ _,_ **_z_** riem[G] [(][z][|][e][ϕ][(][x][)][, σ][)][,]
_ϕ,θ_ _[L][REC][ = min]ϕ,θ_ _[l][(][x][, d][θ][(][z][))][,]_ _∈X_ _∼N_

Hence, we no longer decode the embedding µ but rather z that is obtained with the Riemannian
Gaussian distribution centered on µ. Since the metric G is unknown, a Riemannian VAE would aim
at 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._**
**G(µ). 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 Eq. 15. As of now, the manifold is not regularized and so pathological cases such
as the metric collapsing to 0 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 is
not too far from a 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-like manifold, there exists a
global chart z and so the density of the Riemannian Gaussian distribution can be written with respect
to the Lebesgue measure dz in R[d] (Pennec, 2006)

riem[G] [(][z][|][σ,][ µ][) = 1] det G(z)
_N_ _C_ [exp] _−_ [dist][G]2[(]σ[z][,][ µ][)][2]

 p (16)

_C =_ exp det G(z) dz,

_−_ [dist][G]2[(]σ[z][,][ µ][)][2]
RZ[d]  p

Hence, the regularization term set as the KL divergence between the Riemannian Gaussian distribution and the standard normal would follow

_LREG = DKL_ _Nriem[G]_ [(][z][|][σ,][ µ][)][∥N] [(0][, I][d][)]
 

= riem[(][σ,][µ][)] log(p **Griem[(][σ,][µ][)][)][ −]** [log(][p][N][ (0][,I][d][)][)] _dz ._

_N_
R[d][ p][N][ G]

Z  


-----

The final objective a Riemannian VAE would try to minimize would then write

_L = LREC + LREG ._

Unfortunately, this framework cannot be used in practice for at least two reasons. The first one is
the sampling from the Riemaniann distribution in Eq. 15 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 be amenable to the reparametrization trick (Salimans 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 is not usable in practice, it can be assumed that the value of the metric during training
can be approximated by its value at µ (i.e. G(µ)). 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(µ) = Σ(x)[−][1] where Σ(x) is the covariance matrix given by
the encoder of a vanilla VAE. We indeed have


_√det Σ[−][1]_
√det Σ[−][1] _dz,_ (17)




riem[Σ][−][1] [(][z][|][σ,][ µ][) = 1]
_N_ _C_ [exp] _−_ [dist][Σ][−]2[1]σ[(][z][,][ µ][)][2]




exp
_−_ [dist][Σ][−]2[1]σ[(][z][,][ µ][)][2]
RZ[d] 


_C =_


where distΣ−1 (z, µ) =


(z − **_µ)[⊤]Σ[−][1](z −_** **_µ). If we further set σ = 1 we have_**


exp 2
_Nriem[Σ][−][1]_ [(][z][|][σ,][ µ][) =]  _−_ [(][z][−][µ][)][⊤][Σ][−][1][(][z][−][µ][)]  _._ (18)

exp 2 _dz_
_−_ [(][z][−][µ][)][⊤][Σ][−][1][(][z][−][µ][)]
R[d]
 

R

Assuming as is standard in the VAE framework, that Σ is diagonal makes the computation of REG
_L_
easy and we retrieve the training of a vanilla VAE model. Indeed, as explained in Remark. 1, the log
of the conditional distribution pθ reduces to the reconstruction loss LREC and the KL between the
variational posterior qϕ(z _x) and the prior taken as a standard Gaussian gives_ REG.
_|_ _L_


-----

