Score-based pullback Riemannian geometry

## Abstract

Abstract Data-driven Riemannian geometry has emerged as a powerful tool for interpretable representation learning, offering improved efficiency in downstream tasks. Moving forward, it is crucial to balance cheap manifold mappings with efficient training algorithms. In this work, we integrate concepts from pullback Riemannian geometry and generative models to propose a framework for data-driven Riemannian geometry that is scalable in both geometry and learning: score-based pullback Riemannian geometry. Focusing on unimodal distributions as a first step, we propose a score-based Riemannian structure with closed-form geodesics that pass through the data probability density. With this structure, we construct a Riemannian autoencoder (RAE) with error bounds for discovering the correct data manifold dimension. This framework can naturally be used with anisotropic normalizing flows by adopting isometry regularization during training. Through numerical experiments on various datasets, we demonstrate that our framework not only produces high-quality geodesics through the data support, but also reliably estimates the intrinsic dimension of the data manifold and provides a global chart of the manifold, even in high-dimensional ambient spaces.

## 1 Introduction

Data often reside near low-dimensional non-linear manifolds as illustrated in Figure 1. This manifold assumption (Fefferman et al., ) has been popular since the early work on non-linear dimension reduction (Belkin & Niyogi, ; Coifman & Lafon, ; Roweis & Saul, ; Sammon, ; Tenenbaum et al., ). Learning this non-linear structure, or representation learning, from data has proven to be highly successful (DeMers & Cottrell, ; Kingma & Welling, ) and continues to be a recurring theme in modern machine learning approaches and downstream applications (Chow et al., ; Gomari et al., ; Ternes et al., ; Vahdat & Kautz, ; Zhong et al., ).

Recent advances in data-driven Riemannian geometry have demonstrated its suitability for learning representations.
In this context, these representations are elements residing in a learned geodesic subspace of the data space, governed by a non-trivial Riemannian structure(^1^11rather than the standard $\ell^{2}$-inner product) across the entire ambient space (Arvanitidis et al., ; Diepeveen, ; Hauberg et al., ; Peltonen et al., ; Scarvelis & Solomon, ; Sorrenson et al., ; Sun et al., ).
Among these contributions, it is worth highlighting that Sorrenson et al. () are the first and only ones so far to use information from the full data distribution obtained though generative models (Dinh et al., ; Song et al., ), even though this seems a natural approach given recent studies such as Sakamoto et al. (); Stanczuk et al. (). A possible explanation for the limited use of generative models in constructing Riemannian geometry could lie in challenges regarding *scalability of the manifold mappings*. Indeed, even though the generative models can be trained efficiently, Sorrenson et al. () also mention themselves that it can be numerically challenging to work with their induced Riemannian geometry.

Figure: (a) Hemisphere (2,3)
Refer to caption: extracted/5896858/results/visualisations/RAE/projections/hemisphere_2_3/riemannian_autoencoder.jpg

If the manifold mapping scalability challenges were to be overcome, the combined power of Riemannian geometry and state of the art generative modelling could have profound implications on how to handle data in general.
Indeed, beyond typical data analysis tasks such as computing distances, means, and interpolations/extrapolations of data points as illustrated in Figures 2(a), 2(b), 2(c) and 2(d), a data-driven Riemannian structure also offers greater potential for representation learning and downstream applications. For instance, many advanced data processing methods, from Principal Component Analysis (PCA) to score and flow-matching, have Riemannian counterparts (Fletcher et al. (); Diepeveen et al. () and Chen & Lipman (); Huang et al. ()) that have proven beneficial by improving upon full black box methods in terms of interpretability (Diepeveen, ) or Euclidean counterparts in terms of efficiency (Kapusniak et al., ).
Here it is worth highlighting that scalability of manifold mappings was completely circumvented by Diepeveen () by using pullback geometry.
However, here learning a suitable (and stable) pullback geometry suffers from challenges regarding *scalability of the training algorithm*, contrary to the approach by Sorrenson et al. ().

Motivated by the above, this work aims to address the following question: How to strike a good balance between scalability of training a data-driven Riemannian structure and of evaluating its corresponding manifold mappings?

Figure: (a) Barycentre
Refer to caption: extracted/5896858/results/banana-distribution/barycentre.png

### 1.1 Contributions

In this paper, we take first steps towards striking such a balance and propose a score-based pullback Riemannian metric assuming a relatively simple but generally applicable family of probability densities, which we show to result in both scalable manifold mappings and scalable learning algorithms. We emphasize that we do not directly aim to find the perfect balance between the two types of scalability. Instead we start from a setting which has many nice properties, but will allow for generalization to multimodal densities, which we reserve for future work.

Specifically, we consider a family of unimodal probability densities whose negative log-likelihoods are compositions of strongly convex functions and diffeomorphisms. As this work is an attempt to bridge between the geometric data analysis community and the generative modeling community, we break down the contributions in two ways.
Theoretically,

- •
We propose a score-based pullback Riemannian metric such that manifold mappings respect the data distribution as illustrated in Figures 2(a), 2(b), 2(c) and 2(d).
- •
We demonstrate that this density-based Riemannian structure naturally leads to a Riemannian autoencoder(^2^22in the sense of Diepeveen ()) and provide error bounds on the expected reconstruction error, which allows for approximation of the data manifold as illustrated in Figure 1.
- •
We introduce a learning scheme based on adaptations of normalizing flows to find the density to be integrated into the Riemannian framework, which is tested on several synthetic data sets.

Practically, this work showcases how two simple adaptations to the normalizing flows framework enable data-driven Riemannian geometry. This significantly expands the potential for downstream applications compared to the unadapted framework.

### 1.2 Outline

After introducing notation in Section 2, Section 3 considers a family of probability distributions, from which we obtain suitable geometry, and Section 4 showcases how one can subsequently construct Riemannian Autoencoders with theoretical guarantees. From these observations Section 5 discusses the natural limitations of standard normalizing flows and how to change the parametrisation and training for downstream application in a Riemannian geometric setting. Section 6 showcases several use cases of data-driven Riemannian structure on several data sets. Finally, we summarize our findings in Section 7.

## 2 Notation

Here we present some basic notations from differential and Riemannian geometry, see Boothby (); Carmo (); Lee (); Sakai () for details.

Let $\mathcal{M}$ be a smooth manifold. We write $C^{\infty}(\mathcal{M})$ for the space of smooth functions over $\mathcal{M}$. The *tangent space* at $\mathbf{p}\in\mathcal{M}$, which is defined as the space of all *derivations* at $\mathbf{p}$, is denoted by $\mathcal{T}_{\mathbf{p}}\mathcal{M}$ and for *tangent vectors* we write $\Xi_{\mathbf{p}}\in\mathcal{T}_{\mathbf{p}}\mathcal{M}$. For the *tangent bundle* we write $\mathcal{T}\mathcal{M}$ and smooth vector fields, which are defined as *smooth sections* of the tangent bundle, are written as $\mathscr{X}(\mathcal{M})\subset\mathcal{T}\mathcal{M}$.

A smooth manifold $\mathcal{M}$ becomes a *Riemannian manifold* if it is equipped with a smoothly varying *metric tensor field* $(\,\cdot\,,\,\cdot\,)\colon\mathscr{X}(\mathcal{M})\times\mathscr{X}(\mathcal{
M})\to C^{\infty}(\mathcal{M})$. This tensor field induces a *(Riemannian) metric* $d_{\mathcal{M}}\colon\mathcal{M}\times\mathcal{M}\to\mathbb{R}$. The metric tensor can also be used to construct a unique affine connection, the *Levi-Civita connection*, that is denoted by $\nabla_{(\,\cdot\,)}(\,\cdot\,):\mathscr{X}(\mathcal{M})\times\mathscr{X}(
\mathcal{M})\to\mathscr{X}(\mathcal{M})$.
This connection is in turn the cornerstone of a myriad of manifold mappings.
One is the notion of a *geodesic*, which for two points $\mathbf{p},\mathbf{q}\in\mathcal{M}$ is defined as a curve $\gamma_{\mathbf{p},\mathbf{q}}\colon[0,1]\to\mathcal{M}$ with minimal length that connects $\mathbf{p}$ with $\mathbf{q}$.
Another closely related notion to geodesics is the curve $t\mapsto\gamma_{\mathbf{p},\Xi_{\mathbf{p}}}(t)$ for a geodesic starting from $\mathbf{p}\in\mathcal{M}$ with velocity $\dot{\gamma}_{\mathbf{p},\Xi_{\mathbf{p}}}(0)=\Xi_{\mathbf{p}}\in\mathcal{T}_{
\mathbf{p}}\mathcal{M}$. This can be used to define the *exponential map* $\exp_{\mathbf{p}}\colon\mathcal{D}_{\mathbf{p}}\to\mathcal{M}$ as

$$ $\exp_{\mathbf{p}}(\Xi_{\mathbf{p}}):=\gamma_{\mathbf{p},\Xi_{\mathbf{p}}}(1) \quad\text{where $\mathcal{D}_{\mathbf{p}}\subset\mathcal{T}_{\mathbf{p}} \mathcal{M}$ is the set on which $\gamma_{\mathbf{p},\Xi_{\mathbf{p}}}(1)$ is defined.}$ (1) $$

Furthermore, the *logarithmic map* $\log_{\mathbf{p}}\colon\exp(\mathcal{D}^{\prime}_{\mathbf{p}})\to\mathcal{D}^{
\prime}_{\mathbf{p}}$ is defined as the inverse of $\exp_{\mathbf{p}}$, so it is well-defined on $\mathcal{D}^{\prime}_{\mathbf{p}}\subset\mathcal{D}_{\mathbf{p}}$ where $\exp_{\mathbf{p}}$ is a diffeomorphism.

Finally, if $(\mathcal{M},(\cdot,\cdot))$ is a $d$-dimensional Riemannian manifold, $\mathcal{N}$ is a $d$-dimensional smooth manifold and $\phi:\mathcal{N}\to\mathcal{M}$ is a diffeomorphism, the *pullback metric*

$$ $(\Xi,\Phi)^{\phi}:=(D_{(\cdot)}\phi[\Xi_{(\cdot)}],D_{(\cdot)}\phi[\Phi_{( \cdot)}])_{\phi(\cdot)},$ (2) $$

where $D_{\mathbf{p}}\phi:\mathcal{T}_{\mathbf{p}}\mathcal{N}\to\mathcal{T}_{\phi(
\mathbf{p})}\mathcal{M}$ denotes the differential of $\phi$,
defines a Riemannian structure on $\mathcal{N}$, which we denote by $(\mathcal{N},(\cdot,\cdot)^{\phi})$.
Pullback metrics literally pull back all geometric information from the Riemannian structure on $\mathcal{M}$. In particular, closed-form manifold mappings on $(\mathcal{M},(\cdot,\cdot))$ yield under mild assumptions closed-form manifold mappings on $(\mathcal{N},(\cdot,\cdot)^{\phi})$.
Throughout the rest of the paper pullback mappings will be denoted similarly to eq. 2 with the diffeomorphism $\phi$ as a superscript, i.e., we write $d^{\phi}_{\mathcal{N}}(\mathbf{p},\mathbf{q})$, $\gamma^{\phi}_{\mathbf{p},\mathbf{q}}$, $\exp^{\phi}_{\mathbf{p}}(\Xi_{\mathbf{p}})$ and $\log^{\phi}_{\mathbf{p}}\mathbf{q}$
for $\mathbf{p},\mathbf{q}\in\mathcal{N}$ and $\Xi_{\mathbf{p}}\in\mathcal{T}_{\mathbf{p}}\mathcal{N}$.

## 3 Riemannian geometry from unimodal probability densities

We remind the reader that the ultimate goal of data-driven Riemannian geometry on $\mathbb{R}^{d}$ is to construct a Riemannian structure such that geodesics always pass through the support of data probability densities. In this section we will focus on constructing Riemannian geometry that does just that from unimodal densities $p:\mathbb{R}^{d}\to\mathbb{R}$ of the form

$$ $p(\mathbf{x})\propto e^{-\psi(\varphi(\mathbf{x}))}$ (3) $$

where $\psi:\mathbb{R}^{d}\to\mathbb{R}$ is a smooth strongly convex function and $\varphi:\mathbb{R}^{d}\to\mathbb{R}^{d}$ is a diffeomorphism, e.g., such as the density in Figure 2(^3^33Here, $\psi(\mathbf{x}):=2\mathbf{x}_{1}^{2}+\frac{1}{8}\mathbf{x}_{2}^{2}$ and $\varphi(\mathbf{x}):=(\mathbf{x}_{1}-\frac{1}{9}\mathbf{x}_{2}^{2},\mathbf{x}_ {2})$.). In particular, we will consider pullback Riemannian structures of the form

$$ $(\Xi,\Phi)^{\nabla\psi\circ\varphi}_{\mathbf{x}}:=(D_{\mathbf{x}}\nabla\psi \circ\varphi[\Xi],D_{\mathbf{x}}\nabla\psi\circ\varphi[\Phi])_{2},$ (4) $$

which are related to the Riemannian structure obtained from the *score function* $\nabla\log(p(\cdot)):\mathbb{R}^{d}\to\mathbb{R}^{d}$ if $\varphi$ is close to a linear $\ell^{2}$-isometry on the data support, i.e., $D_{\mathbf{x}}\varphi$ is an orthogonal operator:

$$ $(D_{\mathbf{x}}\nabla\log(p(\cdot))[\Xi],D_{\mathbf{x}}\nabla\log(p(\cdot))[ \Phi])_{2}=(D_{\mathbf{x}}\nabla(\psi\circ\varphi)[\Xi],D_{\mathbf{x}}\nabla( \psi\circ\varphi)[\Phi])_{2}\\ =(D_{\mathbf{x}}((D_{(\cdot)}\varphi)^{\top}\circ\nabla\psi\circ\varphi)[\Xi], D_{\mathbf{x}}((D_{(\cdot)}\varphi)^{\top}\circ\nabla\psi\circ\varphi)[\Phi])_ {2}\\ \approx(D_{\mathbf{x}}\nabla\psi\circ\varphi[\Xi],D_{\mathbf{x}}\nabla\psi \circ\varphi[\Phi])_{2}=(\Xi,\Phi)^{\nabla\psi\circ\varphi}_{\mathbf{x}}.$ (5) $$

For that reason, we call such an approach to data-driven Riemannian geometry: *score-based pullback Riemannian geometry*. Since we find ourselves in a pullback setting(^4^44This is generally not true when using the score itself for probability densities of the form eq. 3.), this allows to construct pullback geometry with closed-form manifold mappings.

What remains to be shown is that such geodesics and other manifold mappings pass through the data support (like in Figures 2(a), 2(b), 2(c) and 2(d)).
The following result, which is a direct application of (Diepeveen, , Prop. 2.1) and (Diepeveen, , Cor. 3.6.1), gives us closed-form expressions of several important manifold mappings under $(\cdot,\cdot)^{\nabla\psi\circ\varphi}$ and makes a connection with $(\cdot,\cdot)^{\varphi}$ if we choose

$$ $\psi(\mathbf{x})=\frac{1}{2}\mathbf{x}^{\top}\mathbf{A}^{-1}\mathbf{x},$ (6) $$

where $\mathbf{A}\in\mathbb{R}^{d\times d}$ is symmetric positive definite.

This special case highlights why, in general, we obtain geodesics and manifold mappings that pass through the data support. For instance, in the scenario depicted in Figure 2(b), where the correct form eq. 3 is used, geodesics are computed by first reversing the effect of the diffeomorphism – transforming the data distribution to resemble a Gaussian, then drawing straight lines between the morphed data points, and finally applying the diffeomorphism again. This approach results in geodesics that traverse regions of higher likelihood between the endpoints, due to the strong convexity of the quadratic function, which aligns perfectly with our objectives.

###### Proposition 1 .

Let $\varphi:\mathbb{R}^{d}\to\mathbb{R}^{d}$ be a smooth diffeomorphism and let $\psi:\mathbb{R}^{d}\to\mathbb{R}$ be a smooth strongly convex function, whose Fenchel conjugate is denoted by $\psi^{\star}:\mathbb{R}^{d}\to\mathbb{R}$. Next, consider the $\ell^{2}$-pullback manifolds $(\mathbb{R}^{d},(\cdot,\cdot)^{\nabla\psi\circ\varphi})$ and $(\mathbb{R}^{d},(\cdot,\cdot)^{\varphi})$ defined through metric tensor fields

$$ $(\Xi,\Phi)^{\nabla\psi\circ\varphi}_{\mathbf{x}}:=(D_{\mathbf{x}}\nabla\psi \circ\varphi[\Xi],D_{\mathbf{x}}\nabla\psi\circ\varphi[\Phi])_{2},\quad\text{ and}\quad(\Xi,\Phi)^{\varphi}_{\mathbf{x}}:=(D_{\mathbf{x}}\varphi[\Xi],D_{ \mathbf{x}}\varphi[\Phi])_{2}.$ (7) $$

Then,

- (i)
length-minimising geodesics $\gamma^{\nabla\psi\circ\varphi}_{\mathbf{x},\mathbf{y}}:[0,1]\to\mathbb{R}^{d}$ on $(\mathbb{R}^{d},(\cdot,\cdot)^{\nabla\psi\circ\varphi})$ are given by
$\gamma^{\nabla\psi\circ\varphi}_{\mathbf{x},\mathbf{y}}(t)=(\varphi^{-1}\circ
\nabla\psi^{*})((1-t)(\nabla\psi\circ\varphi)(\mathbf{x})+t(\nabla\psi\circ
\varphi)(\mathbf{y})).$
(8)
In addition, if $\psi$ is of the form eq. 6
$\gamma^{\nabla\psi\circ\varphi}_{\mathbf{x},\mathbf{y}}(t)=\gamma^{\varphi}_{
\mathbf{x},\mathbf{y}}(t)=\varphi^{-1}((1-t)\varphi(\mathbf{x})+t\varphi(
\mathbf{y})).$
(9)
- (ii)
the logarithmic map $\log^{\nabla\psi\circ\varphi}_{\mathbf{x}}(\cdot):\mathbb{R}^{d}\to\mathcal{T}
_{\mathbf{x}}\mathbb{R}^{d}$ on $(\mathbb{R}^{d},(\cdot,\cdot)^{\nabla\psi\circ\varphi})$ is given by
$\log^{\nabla\psi\circ\varphi}_{\mathbf{x}}\mathbf{y}=D_{\varphi(\mathbf{x})}
\varphi^{-1}[D_{(\nabla\psi\circ\varphi)(\mathbf{x})}\nabla\psi^{*}[(\nabla
\psi\circ\varphi)(\mathbf{y})-(\nabla\psi\circ\varphi)(\mathbf{x})]].$
(10)
In addition, if $\psi$ is of the form eq. 6
$\log^{\nabla\psi\circ\varphi}_{\mathbf{x}}\mathbf{y}=\log^{\varphi}_{\mathbf{x
}}\mathbf{y}=D_{\varphi(\mathbf{x})}\varphi^{-1}[\varphi(\mathbf{y})-\varphi(
\mathbf{x})].$
(11)
- (iii)
the exponential map $\exp^{\nabla\psi\circ\varphi}_{\mathbf{x}}(\cdot):\mathcal{T}_{\mathbf{x}}
\mathbb{R}^{d}\to\mathbb{R}^{d}$ on $(\mathbb{R}^{d},(\cdot,\cdot)^{\nabla\psi\circ\varphi})$ is given by
$\exp^{\nabla\psi\circ\varphi}_{\mathbf{x}}(\Xi_{\mathbf{x}})=(\varphi^{-1}
\circ\nabla\psi^{*})((\nabla\psi\circ\varphi)(\mathbf{x})+D_{\varphi(\mathbf{x
})}\nabla\psi[D_{\mathbf{x}}\varphi[\Xi_{\mathbf{x}}]]).$
(12)
In addition, if $\psi$ is of the form eq. 6
$\exp^{\nabla\psi\circ\varphi}_{\mathbf{x}}(\Xi_{\mathbf{x}})=\exp^{\varphi}_{
\mathbf{x}}(\Xi_{\mathbf{x}})=\varphi^{-1}(\varphi(\mathbf{x})+D_{\mathbf{x}}
\varphi[\Xi_{\mathbf{x}}]).$
(13)
- (iv)
the distance $d^{\nabla\psi\circ\varphi}_{\mathbb{R}^{d}}:\mathbb{R}^{d}\times\mathbb{R}^{d}
\to\mathbb{R}$ on $(\mathbb{R}^{d},(\cdot,\cdot)^{\nabla\psi\circ\varphi})$ is given by
$d^{\nabla\psi\circ\varphi}_{\mathbb{R}^{d}}(\mathbf{x},\mathbf{y})=\|(\nabla
\psi\circ\varphi)(\mathbf{x})-(\nabla\psi\circ\varphi)(\mathbf{y})\|_{2}.$
(14)
In addition, if $\psi$ is of the form eq. 6
$d^{\nabla\psi\circ\varphi}_{\mathbb{R}^{d}}(\mathbf{x},\mathbf{y})=\|\varphi(
\mathbf{x})-\varphi(\mathbf{y})\|_{\mathbf{A}^{-2}}:=\|\mathbf{A}^{-1}(\varphi
(\mathbf{x})-\varphi(\mathbf{y}))\|_{2}.$
(15)
- (v)
the Riemannian barycentre $\mathbf{x}^{*}\in\mathbb{R}^{d}$ of the data set $\{\mathbf{x}^{i}\}_{i=1}^{N}$ on $(\mathbb{R}^{d},(\cdot,\cdot)^{\nabla\psi\circ\varphi})$ is given by
$\mathbf{x}^{*}:=\operatorname*{arg\,min}_{\mathbf{x}\in\mathbb{R}^{d}}\Bigl{\{
}\frac{1}{2N}\sum_{i=1}^{N}d_{\mathbb{R}^{d}}^{\nabla\psi\circ\varphi}(\mathbf
{x},\mathbf{x}^{i})^{2}\Bigr{\}}=(\varphi^{-1}\circ\nabla\psi^{*})\Bigl{(}
\frac{1}{N}\sum_{i=1}^{N}\nabla\psi(\varphi(\mathbf{x}^{i}))\Bigr{)}.$
(16)
In addition, if $\psi$ is of the form eq. 6
$\mathbf{x}^{*}:=\operatorname*{arg\,min}_{\mathbf{x}\in\mathbb{R}^{d}}\Bigl{\{
}\frac{1}{2N}\sum_{i=1}^{N}d_{\mathbb{R}^{d}}^{\varphi}(\mathbf{x},\mathbf{x}^
{i})^{2}\Bigr{\}}=\varphi^{-1}\Bigl{(}\frac{1}{N}\sum_{i=1}^{N}\varphi(\mathbf
{x}^{i})\Bigr{)}.$
(17)

###### Remark 1 .

We note that $\ell^{2}$-stability of geodesics and the barycentre are inherited by (Diepeveen, , Thms. 3.4&3.8), if we have (approximate) local $\ell^{2}$-isometry of $\varphi$ on the data distribution.

###### Proposition 1 .

###### Remark 1 .

## 4 Riemannian autoencoder from unimodal probability densities

The connection between $(\cdot,\cdot)^{\nabla\psi\circ\varphi}$ and $(\cdot,\cdot)^{\varphi}$ begs the question what $\psi$ could still be used for if it is of the form eq. 6.
We note that this case comes down to having a data probability density that is a deformed Gaussian distribution. In the case of a regular (non-deformed) Gaussian, one can compress the data generated by it through projecting them onto a low rank approximation of the covariance matrix such that only the directions with highest variance are taken into account. This is the basic idea behind PCA. In the following we will generalize this idea to the Riemannian setting and observe that this amounts to constructing a *Riemannian autoencoder* (RAE) (Diepeveen, ), whose error we can bound by picking the dimension of the autoencoder in a clever way, reminiscent of the classical PCA error bound.

Concretely, we assume that we have a unimodal density of the form eq. 3 with a quadratic strongly convex function $\psi(\mathbf{x}):=\frac{1}{2}\mathbf{x}^{\top}\mathbf{A}^{-1}\mathbf{x}$ for some diagonal matrix $\mathbf{A}:=\operatorname*{diag}(\mathbf{a}_{1},\ldots\mathbf{a}_{d})$ with positive entries(^5^55Note that this is not restrictive as for a general symmetric positive definite matrix $\mathbf{A}$ the eigenvalues can be used as diagonal entries and the orthonormal matrices can be concatenated with the diffeomorphism.). Next, we define an indexing $u_{w}\in[d]:=\{1,\ldots,d\}$ for $w=1,\ldots,d$ such that

$$ $\mathbf{a}_{u_{1}}\geq\ldots\geq\mathbf{a}_{u_{d}},$ (18) $$

and consider a threshold $\varepsilon\in[0,1]$. We then consider $d_{\varepsilon}\in[d]$ defined as the integer that satisfies

$$ $d_{\varepsilon}:=\left\{\begin{matrix}\min\Bigl{\{}d^{\prime}\in[d-1]\;\Bigl{| }\;\sum_{w=d^{\prime}+1}^{d}\mathbf{a}_{u_{w}}\leq\varepsilon\sum_{u=1}^{d} \mathbf{a}_{u}\Bigr{\}},&\text{if }\mathbf{a}_{u_{d}}\leq\varepsilon\sum_{u=1} ^{d}\mathbf{a}_{u},\\ d,&\text{otherwise.}\\ \end{matrix}\right..$ (19) $$

Finally, we define the mapping $E_{\varepsilon}:\mathbb{R}^{d}\to\mathbb{R}^{d_{\varepsilon}}$ coordinate-wise as

$$ $E_{\varepsilon}(\mathbf{x})_{w}:=(\log^{\varphi}_{\varphi^{-1}(\mathbf{0})} \mathbf{x},D_{\mathbf{0}}\varphi^{-1}[\mathbf{e}^{u_{w}}])^{\varphi}_{\varphi^ {-1}(\mathbf{0})}\overset{\text{\lx@cref{creftype~refnum}{eq:thm-log-remetrize d-special-psi}}}{=}(\varphi(\mathbf{x}),\mathbf{e}^{u_{w}})_{2},\quad w=1, \ldots,d_{\varepsilon},$ (20) $$

and define $D_{\varepsilon}:\mathbb{R}^{d_{\varepsilon}}\to\mathbb{R}^{d}$ as

$$ $D_{\varepsilon}(\mathbf{p}):=\exp_{\varphi^{-1}(\mathbf{0})}^{\varphi}\Bigl{(} \sum_{w=1}^{d_{\varepsilon}}\mathbf{p}_{w}D_{\mathbf{0}}\varphi^{-1}[\mathbf{e }^{u_{w}}]\Bigr{)}\overset{\text{\lx@cref{creftype~refnum}{eq:thm-exp-remetriz ed-special-psi}}}{=}\varphi^{-1}\Bigl{(}\sum_{w=1}^{d_{\varepsilon}}\mathbf{p} _{w}\mathbf{e}^{u_{w}}\Bigr{)},$ (21) $$

which generate a Riemannian autoencoder and the set $D_{\varepsilon}(\mathbb{R}^{d_{\varepsilon}})\subset\mathbb{R}^{d}$ as an approximate data manifold as in the scenario in Figure 1.

As hinted above, this Riemannian autoencoder comes with an error bound on the expected approximation error, which is fully determined by the diffeomorphism’s deviation from isometry around the data manifold. For the proof, we refer the reader to Appendix A.

###### Theorem 1 .

Let $\varphi:\mathbb{R}^{d}\to\mathbb{R}^{d}$ be a smooth diffeomorphism and let $\psi:\mathbb{R}^{d}\to\mathbb{R}$ be a quadratic function of the form eq. 6 with diagonal $\mathbf{A}\in\mathbb{R}^{d\times d}$. Furthermore, let $p:\mathbb{R}^{d}\to\mathbb{R}$ be the corresponding probability density of the form eq. 3. Finally, consider $\varepsilon\in[0,1]$ and the mappings $E_{\varepsilon}:\mathbb{R}^{d}\to\mathbb{R}^{d_{\varepsilon}}$ and $D_{\varepsilon}:\mathbb{R}^{d_{\varepsilon}}\to\mathbb{R}^{d}$ in eqs. 20 and 21 with $d_{\varepsilon}\in[d]$ as in eq. 19.

Then,

$$ $\mathbb{E}_{\mathbf{X}\sim p}[\|D_{\varepsilon}(E_{\varepsilon}(\mathbf{X}))- \mathbf{X}\|_{2}^{2}]\leq\varepsilon\inf_{\beta\in[0,\frac{1}{2})}\Bigl{\{} \frac{B_{\beta,\varphi^{-1}}C_{\beta,\varphi}C_{\beta,\varphi^{-1}}}{1-2\beta} \Bigl{(}\frac{1+\beta}{1-2\beta}\Bigr{)}^{\frac{d}{2}}\Bigr{\}}\sum_{i=1}^{d} \mathbf{a}_{i}+o(\varepsilon),$ (22) $$

where

$$ $B_{\beta,\varphi^{-1}}:=\sup_{\mathbf{x}\in\mathbb{R}^{d}}\{\|D_{\varphi( \mathbf{x})}\varphi^{-1}\|_{2}^{2}e^{-\frac{\beta}{2}\varphi(\mathbf{x})^{\top }\mathbf{A}^{-1}\varphi(\mathbf{x})}\},$ (23) $$

$$ $C_{\beta,\varphi^{-1}}:=\sup_{\mathbf{x}\in\mathbb{R}^{d}}\{|\det(D_{\varphi( \mathbf{x})}\varphi^{-1})|e^{-\frac{\beta}{2}\varphi(\mathbf{x})^{\top}\mathbf {A}^{-1}\varphi(\mathbf{x})}\},$ (24) $$

and

$$ $C_{\beta,\varphi}:=\sup_{\mathbf{x}\in\mathbb{R}^{d}}\{|\det(D_{\mathbf{x}} \varphi)|e^{-\frac{\beta}{2}\varphi(\mathbf{x})^{\top}\mathbf{A}^{-1}\varphi( \mathbf{x})}\}.$ (25) $$

###### Remark 2 .

Note that the RAE latent space is interpretable as it is $\ell^{2}$-isometric to the data manifold if $\varphi$ is an approximate $\ell^{2}$-isometry on the data manifold. In other words, latent representations being close by or far away correspond to similar behaviour in data space, which is not the case for a VAE (Kingma & Welling, ).

###### Theorem 1 .

###### Remark 2 .

## 5 Learning unimodal probability densities

Naturally, we want to learn probability densities of the form eq. 3, which can then directly be inserted into the proposed score-based pullback Riemannian geometry framework. In this section we will consider how to adapt normalizing flow (NF) (Dinh et al., ) training to a setting that is more suitable for our purposes(^6^66We note that the choice for adapting the normalizing flow training scheme rather than using diffusion model training schemes is due to more robust results through the former.). In particular, we will consider how training a normalizing flow density $p:\mathbb{R}^{d}\to\mathbb{R}$ given by

$$ $p(\mathbf{x}):=\frac{1}{C_{\psi}}e^{-\psi(\varphi(\mathbf{x}))}|\det(D_{ \mathbf{x}}\varphi)|,$ (26) $$

where $C_{\psi}>0$ is a normalisation constant that only depends on the strongly convex function $\psi$, yields our target distribution eq. 3.

From Sections 3 and 4 we have seen that ideally the strongly convex function $\psi:\mathbb{R}^{d}\to\mathbb{R}$ corresponds to a Gaussian with a parameterised diagonal covariance matrix $\mathbf{A}\in\mathbb{R}^{d\times d}$, resulting in more parameters than in standard normalizing flows, whereas the diffeomorphism $\varphi:\mathbb{R}^{d}\to\mathbb{R}^{d}$ is regularized to be an isometry. In particular, $\mathbf{A}$ ideally allows for learnable anisotropy rather than having a fixed isotropic identity matrix. The main reason is that through anisotropy we can construct a Riemannian autoencoder (RAE), since it is known which dimensions are most important. Moreover, the diffeomorphism should be $\ell^{2}$-isometric, unlike standard normalizing flows which are typically non-volume preserving, enabling stability (Remark 1) and a practically useful and interpretable RAE (Theorems 1 and 2). In addition, $\ell^{2}$-isometry (on the data support) implies volume-preservation, which means that $|\det(D_{\mathbf{x}}\varphi)|\approx 1$ so that eq. 26 reduces to the target distribution eq. 3.

This leads to learning the density through minimizing the following adapted normalizing flow loss

$$ $\mathcal{L}(\theta_{1},\theta_{2}):=\mathbb{E}_{\mathbf{X}\sim p_{\text{data}} }\left[-\log p_{\theta_{1},\theta_{2}}(\mathbf{X})\right]\\ +\lambda_{\text{vol}}\mathbb{E}_{\mathbf{X}\sim p_{\text{data}}}\left[\log(| \det\bigl{(}D_{\mathbf{X}}\varphi_{\theta_{2}}\bigr{)}|)^{2}\right]+\lambda_{ \text{iso}}\mathbb{E}_{\mathbf{X}\sim p_{\text{data}}}\left[\|(D_{\mathbf{x}} \varphi_{\theta_{2}})^{\top}D_{\mathbf{X}}\varphi_{\theta_{2}}-\mathbf{I}_{d} \|_{F}^{2}\right]$ (27) $$

where $\lambda_{\text{vol}},\lambda_{\text{iso}}>0$ and the negative log likelihood term reduces to

$$ $\mathbb{E}_{\mathbf{X}\sim p_{\text{data}}}\left[-\log p_{\theta_{1},\theta_{2 }}(\mathbf{X})\right]=\frac{1}{2}\mathbb{E}_{\mathbf{X}\sim p_{\text{data}}} \left[\varphi_{\theta_{2}}(\mathbf{X})^{\top}\mathbf{A}_{\theta_{1}}^{-1} \varphi_{\theta_{2}}(\mathbf{X})\right]\\ -\mathbb{E}_{\mathbf{X}\sim p_{\text{data}}}\left[\log(|\det\bigl{(}D_{\mathbf {X}}\varphi_{\theta_{2}}\bigr{)}|)\right]+\frac{1}{2}\operatorname*{tr}( \mathbf{A}_{\theta_{1}})+\frac{d}{2}\log(2\pi),$ (28) $$

where $\mathbf{A}_{\theta_{1}}$ is a diagonal matrix and $\varphi_{\theta_{2}}$ is a normalizing flow with affine coupling layers(^7^77We note that the choice for affine coupling layers rather than using more expressive diffeomorphisms such as rational quadratic flows Durkan et al. () is due to our need for high regularity for stable manifold mappings (Remark 1) and an interpretable RAE (Remark 2), which has empirically shown to be more challenging to achieve for more expressive flows as higher-order derivatives of $\varphi$ will blow up the higher-order error terms in theorem 1.) (Dinh et al., ).

## 6 Experiments

We conducted two sets of experiments to evaluate the proposed scheme from Section 5 to learn suitable pullback Riemannian geometry. The first set investigates whether our adaptation of the standard normalizing flow (NF) training paradigm leads to more accurate and stable manifold mappings, as measured by the geodesic and variation errors. The second set assesses the capability of our method to generate a robust Riemannian autoencoder (RAE).

For all experiments in this section, detailed training configurations are provided in Appendix D.

### 6.1 Manifold mappings

As discussed in Diepeveen (), the quality of learned manifold mappings is determined by two key metrics: the *geodesic error* and the *variation error*. The geodesic error measures the average deviation form the ground truth geodesics implied by the ground truth pullback metric, while the variation error evaluates the stability of geodesics under small perturbations. We define these error metrics for the evaluation of pullback geometries in Appendix C.

Our approach introduces two key modifications to the normalizing flow (NF) training framework:

- 1.
Anisotropic Base Distribution: We parameterize the diagonal elements of the covariance matrix $\mathbf{A}_{\theta_{1}}$, introducing anisotropy into the base distribution.
- 2.
$\ell^{2}$-Isometry Regularization: We regularize the flow $\varphi_{\theta_{2}}$ to be approximately $\ell^{2}$-isometric.

To assess the effectiveness of these modifications in learning more accurate and robust manifold mappings, we compare our method against three baselines:

- (1)
*Normalizing Flow (NF)*: Uses an NF with a standard isotropic Gaussian base distribution $\mathcal{N}(\mathbf{0},\mathbf{I}_{d})$ and no isometry regularization of the flow.
- (2)
*Anisotropic Normalizing Flow*: Uses an NF with the same parameterization of the diagonal covariance matrix in the base distribution as in our method, but without regularization of the flow.
- (3)
*Isometric Normalizing Flow*: Uses an NF with an isotropic Gaussian base distribution $\mathcal{N}(\mathbf{0},\mathbf{I}_{d})$ and regularizes the flow to be approximately $\ell^{2}$-isometric.

We conduct experiments on three datasets, illustrated in Figure 6 in Section B.1: the *Single Banana Dataset*, the *Squeezed Single Banana Dataset*, and the *River Dataset*. Detailed descriptions of the construction and characteristics of these datasets are provided in Section B.1.

Table 1 presents the geodesic and variation errors for each method across the three datasets and Figure 3 visually compares the geodesics computed using each method on the river dataset. Our method consistently achieves significantly lower errors compared to the baselines, indicating more accurate and stable manifold mappings.

Introducing anisotropy in the base distribution without enforcing isometry in the flow offers no significant improvement over the standard flow. On the other hand, regularizing the flow to be approximately isometric without incorporating anisotropy in the base distribution results in underfitting, leading to noticeably worse performance than the standard flow. Our results demonstrate that the combination of anisotropy in the base distribution with isometry regularization (our method) yields the most accurate and stable manifold mappings, as evidenced by consistently lower geodesic and variation errors.

**Table 1: Comparison of evaluation metrics for different methods across three datasets. Best-performing results for each metric are highlighted in bold. Values are reported as mean (std). The proposed method performs best in all metrics on each data set.**
| Metric | Our Method | NF | Anisotropic NF | Isometric NF |
| --- | --- | --- | --- | --- |
| Single Banana Dataset |  |  |  |  |
| Geodesic Error | 0.0315 (0.0268) | 0.0406 (0.0288) | 0.0431 (0.0305) | 0.0817 (0.1063) |
| Variation Error | 0.0625 (0.0337) | 0.0638 (0.0352) | 0.0639 (0.0354) | 0.0639 (0.0355) |
| Squeezed Single Banana Dataset |  |  |  |  |
| Geodesic Error | 0.0180 (0.0226) | 0.0524 (0.0805) | 0.0505 (0.0787) | 0.1967 (0.2457) |
| Variation Error | 0.0631 (0.0326) | 0.0663 (0.0353) | 0.0661 (0.0350) | 0.0669 (0.0361) |
| River Dataset |  |  |  |  |
| Geodesic Error | 0.1691 (0.0978) | 0.2369 (0.1216) | 0.2561 (0.1338) | 0.3859 (0.2568) |
| Variation Error | 0.0763 (0.0486) | 0.1064 (0.0807) | 0.1113 (0.0863) | 0.0636 (0.0333) |

Figure: (a) Ground Truth
Refer to caption: extracted/5896858/results/visualisations/geodesics/gt.png

### 6.2 Riemannian autoencoder

To evaluate the capacity of our method to learn a Riemannian autoencoder, we conducted experiments on two synthetic datasets across various combinations of intrinsic dimension $d^{\prime}$ and ambient dimension $d$:

- •
Hemisphere($d^{\prime}$, $d$): Samples are drawn from the upper hemisphere of a $d^{\prime}$-dimensional unit sphere and embedded in an $d$-dimensional ambient space via a random isometric mapping.
- •
Sinusoid($d^{\prime}$, $d$): This dataset is generated by applying sinusoidal transformations to $d^{\prime}$-dimensional latent variables, resulting in a complex, nonlinear manifold embedded in $d$ dimensions.

For a detailed description of these datasets, refer to Section B.2.

#### 6.2.1 1D and 2D manifolds

In Figures 1 and 4, we present the data manifold approximations by our Riemannian autoencoder for four low-dimensional manifolds.: Hemisphere(2,3), Sinusoid(1,3), Sinusoid(2,3) and Sinusoid(1,100). In appendix E, we detail the process used to create the data manifold approximations for these experiments. In our experiments, we set $\epsilon=0.01$, which resulted in $d_{\epsilon}=d^{\prime}$ for all cases, accurately capturing the intrinsic dimension of each manifold and producing accurate global charts.

Figure: (a) Dims (5, 59, 92)
Refer to caption: extracted/5896858/results/visualisations/RAE/projections/sinusoid_1_100/more_transparent/5_59_92.jpg

#### 6.2.2 Higher-dimensional manifolds

To evaluate the scalability of our method to higher-dimensional manifolds, we conducted additional experiments on the Hemisphere(5,20) and Sinusoid(5,20) datasets.

Our theory suggests that the learned variances indicate the importance of each latent dimension: higher variances signal more important dimensions for reconstructing the manifold, while dimensions with vanishing variances are considered insignificant and are disregarded when constructing the Riemannian autoencoder. To test the model’s ability to correctly identify important and unimportant latent dimensions, we report the average $\ell^{2}$ reconstruction error for each dataset as a function of the number of latent dimensions used. In the reconstruction error plots (see figs. 5(b) and 5(d)), we report three variance-based orders for adding latent dimensions: decreasing variance order (blue line), increasing variance order (green line), and random order (red line).

For the Hemisphere(5,20) dataset, the model identified five non-vanishing variances (see fig. 5(a)), perfectly capturing the intrinsic dimension of the manifold. This is reflected in the blue curve in fig. 5(b), where the first five latent dimensions, corresponding to the largest variances, are sufficient to reduce the reconstruction error almost to zero. In contrast, the green curve illustrates that the remaining ambient dimensions do not encode useful information about the manifold. The red curve demonstrates improvement only when an important latent dimension is included.

For the more challenging Sinusoid(5,20) dataset, our method still performs very well, though not as perfectly as for the Hemisphere dataset. The first six most important latent dimensions explain approximately $97\%$ of the variance, increasing to over $99\%$ with the seventh dimension (see fig. 5(c)). This is reflected in the blue curve in fig. 5(d), where the first six latent dimensions reduce the reconstruction error to near zero, and the addition of the seventh dimension brings the error effectively to zero. The slight discrepancy between our results and the ground truth likely arises from increased optimization difficulty, as the normalizing flow must learn a more intricate distribution while maintaining approximate isometry. We believe that with deeper architectures and more careful tuning of the optimization loss, the model will converge to the correct intrinsic dimensionality of five. Currently, it predicts six dimensions at a threshold of $\epsilon=0.05$ and seven at $\epsilon=0.01$, slightly overestimating due to the manifold’s complexity.

Figure: (a) Learned variances in decreasing order.
Refer to caption: extracted/5896858/results/visualisations/RAE/reconstruction/hemisphere_5_20/diagonal_values_normal_scale.png

## 7 Conclusions

In this work we have taken a first step towards a practical Riemannian geometry framework for generative modelling, striking a balance between scalability of training a data-driven Riemannian structure and of evaluating its corresponding manifold mappings. We have considered a family of unimodal probability densities whose negative log-likelihoods are compositions of strongly convex functions and diffeomorphisms, and sought to learn them. We have shown that once these unimodal densities have been learned, the proposed score-based pullback geometry gives us closed-form geodesics that pass through the data probability density and a Riemannian autoencoder with error bounds that can be used to estimate the dimension of the data manifold. Finally, to learn the distribution we have proposed an adaptation to normalizing flow training. Through numerical experiments, we have shown that these modifications are crucial for extracting geometric information, and that our framework not only generates high-quality geodesics across the data support, but also accurately estimates the intrinsic dimension of the approximate data manifold while constructing a global chart, even in high-dimensional ambient spaces. Current challenges of the method lie in balancing the expressivity of the network architecture, e.g., through additional layers or more expressive architectures, and satisfying approximate $\ell^{2}$-isometry on the data support. For future work we aim to overcome these challenges, extending the method to multimodal distributions, while making it scalable for higher-dimensional data sets.

## Appendix A Proof of theorem 1

##### Auxiliary lemma

###### Lemma 1 .

Let $\varphi:\mathbb{R}^{d}\to\mathbb{R}^{d}$ be a smooth diffeomorphism and let $\psi:\mathbb{R}^{d}\to\mathbb{R}$ be a quadratic function of the form eq. 6 with diagonal $\mathbf{A}\in\mathbb{R}^{d\times d}$. Furthermore, let $p:\mathbb{R}^{d}\to\mathbb{R}$ be the corresponding probability density of the form eq. 3. Finally, consider $\varepsilon\in[0,1]$ and the mappings $E_{\varepsilon}:\mathbb{R}^{d}\to\mathbb{R}^{d_{\varepsilon}}$ and $D_{\varepsilon}:\mathbb{R}^{d_{\varepsilon}}\to\mathbb{R}^{d}$ in eqs. 20 and 21 with $d_{\varepsilon}\in[d]$ as in eq. 19.

Then, for any $\alpha\in[0,1)$ and any $\beta\in[0,1-\alpha)$

$$ $\mathbb{E}_{\mathbf{X}\sim p}[d_{\mathbb{R}^{d}}^{\varphi}(D_{\varepsilon}(E_{ \varepsilon}(\mathbf{X})),\mathbf{X})^{2}e^{\frac{\alpha}{2}\varphi(\mathbf{X} )^{\top}\mathbf{A}^{-1}\varphi(\mathbf{X})}]\leq\varepsilon\frac{C_{\beta, \varphi}C_{\beta,\varphi^{-1}}}{1-\alpha-\beta}\Bigl{(}\frac{1+\beta}{1-\alpha -\beta}\Bigr{)}^{\frac{d}{2}}\sum_{i=1}^{d}\mathbf{a}_{i},$ (29) $$

where

$$ $C_{\beta,\varphi^{-1}}:=\sup_{\mathbf{x}\in\mathbb{R}^{d}}\{|\det(D_{\varphi( \mathbf{x})}\varphi^{-1})|e^{-\frac{\beta}{2}\varphi(\mathbf{x})^{\top}\mathbf {A}^{-1}\varphi(\mathbf{x})}\},$ (30) $$

and

$$ $C_{\beta,\varphi}:=\sup_{\mathbf{x}\in\mathbb{R}^{d}}\{|\det(D_{\mathbf{x}} \varphi)|e^{-\frac{\beta}{2}\varphi(\mathbf{x})^{\top}\mathbf{A}^{-1}\varphi( \mathbf{x})}\}.$ (31) $$

###### Proof.

We need to distinct two cases: (i) $d_{\varepsilon}=d$ and (ii) $1\leq d_{\varepsilon}<d$

(i) If $d_{\varepsilon}=d$ we have that $D_{\varepsilon}(E_{\varepsilon}(\mathbf{x}))=\mathbf{x}$ for any $\mathbf{x}\in\mathbb{R}^{d}$. In other words

$$ $\mathbb{E}_{\mathbf{X}\sim p}[d_{\mathbb{R}^{d}}^{\varphi}(D_{\varepsilon}(E_{ \varepsilon}(\mathbf{X})),\mathbf{X})^{2}e^{\frac{\alpha}{2}\varphi(\mathbf{X} )^{\top}\mathbf{A}^{-1}\varphi(\mathbf{X})}]=0\leq\varepsilon\frac{C_{\beta, \varphi}C_{\beta,\varphi^{-1}}}{1-\alpha-\beta}\Bigl{(}\frac{1+\beta}{1-\alpha -\beta}\Bigr{)}^{\frac{d}{2}}\sum_{i=1}^{d}\mathbf{a}_{i}.$ (32) $$

(ii) Next, we consider the case $1\leq d_{\varepsilon}<d$.
First, notice that we can rewrite

$$ $\|\varphi(D_{\varepsilon}(E_{\varepsilon}(\mathbf{x})))-\varphi(\mathbf{x})\|_ {2}^{2}\overset{\text{\lx@cref{creftypeplural~refnum}{eq:rae-encoder} and~ \lx@cref{refnum}{eq:rae-decoder}}}{=}\|\sum_{k=1}^{d_{\varepsilon}}(\varphi( \mathbf{x}),\mathbf{e}^{i_{k}})_{2}\mathbf{e}^{i_{k}}-\varphi(\mathbf{x})\|_{2 }^{2}=\|\sum_{k=d_{\varepsilon}+1}^{d}(\varphi(\mathbf{x}),\mathbf{e}^{i_{k}}) _{2}\mathbf{e}^{i_{k}}\|_{2}^{2}\\ \overset{\text{orthogonality}}{=}\sum_{k=d_{\varepsilon}+1}^{d}\|(\varphi( \mathbf{x}),\mathbf{e}^{i_{k}})_{2}\mathbf{e}^{i_{k}}\|_{2}^{2}=\sum_{k=d_{ \varepsilon}+1}^{d}(\varphi(\mathbf{x}),\mathbf{e}^{i_{k}})_{2}^{2}=\sum_{k=d_ {\varepsilon}+1}^{d}\varphi(\mathbf{x})_{i_{k}}^{2}.$ (33) $$

Moreover, we define

$$ $C:=\int_{\mathbb{R}^{d}}e^{-\frac{1}{2}\varphi(\mathbf{x})^{\top}\mathbf{A}^{- 1}\varphi(\mathbf{x})}\mathrm{d}\mathbf{x}.$ (34) $$

Then,

$$ $\mathbb{E}_{\mathbf{X}\sim p}[d_{\mathbb{R}^{d}}^{\varphi}(D_{\varepsilon}(E_{ \varepsilon}(\mathbf{X})),\mathbf{X})^{2}e^{\frac{\alpha}{2}\varphi(\mathbf{X} )^{\top}\mathbf{A}^{-1}\varphi(\mathbf{X})}]=\frac{\int_{\mathbb{R}^{d}}\| \varphi(D_{\varepsilon}(E_{\varepsilon}(\mathbf{x})))-\varphi(\mathbf{x})\|_{2 }^{2}e^{-(\frac{1}{2}-\frac{\alpha}{2})\varphi(\mathbf{x})^{\top}\mathbf{A}^{- 1}\varphi(\mathbf{x})}\mathrm{d}\mathbf{x}}{\int_{\mathbb{R}^{d}}e^{-\frac{1}{ 2}\varphi(\mathbf{x})^{\top}\mathbf{A}^{-1}\varphi(\mathbf{x})}\mathrm{d} \mathbf{x}}\\ \overset{\text{\lx@cref{creftype~refnum}{eq:RAEnormalizationConstant}}}{=} \frac{1}{C}\int_{\mathbb{R}^{d}}\|\varphi(D_{\varepsilon}(E_{\varepsilon}( \mathbf{x})))-\varphi(\mathbf{x})\|_{2}^{2}e^{-(\frac{1}{2}-\frac{\alpha}{2}) \varphi(\mathbf{x})^{\top}\mathbf{A}^{-1}\varphi(\mathbf{x})}\mathrm{d}\mathbf {x}\\ \overset{\text{\lx@cref{creftype~refnum}{eq:rewrite-RAEx-x}}}{=}\frac{1}{C} \int_{\mathbb{R}^{d}}\sum_{k=d_{\varepsilon}+1}^{d}\varphi(\mathbf{x})_{i_{k}} ^{2}e^{-(\frac{1}{2}-\frac{\alpha}{2})\varphi(\mathbf{x})^{\top}\mathbf{A}^{-1 }\varphi(\mathbf{x})}\mathrm{d}\mathbf{x}=\frac{1}{C}\sum_{k=d_{\varepsilon}+1 }^{d}\int_{\mathbb{R}^{d}}\varphi(\mathbf{x})_{i_{k}}^{2}e^{-(\frac{1}{2}- \frac{\alpha}{2})\varphi(\mathbf{x})^{\top}\mathbf{A}^{-1}\varphi(\mathbf{x})} \mathrm{d}\mathbf{x}\\ \overset{\mathbf{x}=\varphi^{-1}(\mathbf{y})}{=}\frac{1}{C}\sum_{k=d_{ \varepsilon}+1}^{d}\int_{\mathbb{R}^{d}}\mathbf{y}_{i_{k}}^{2}e^{-(\frac{1}{2} -\frac{\alpha}{2})\mathbf{y}^{\top}\mathbf{A}^{-1}\mathbf{y}}|\det(D_{\mathbf{ y}}\varphi^{-1})|\mathrm{d}\mathbf{y}\\ =\frac{1}{C}\sum_{k=d_{\varepsilon}+1}^{d}\int_{\mathbb{R}^{d}}\mathbf{y}_{i_{ k}}^{2}e^{-(\frac{1}{2}-\frac{\alpha}{2}-\frac{\beta}{2})\mathbf{y}^{\top} \mathbf{A}^{-1}\mathbf{y}}|\det(D_{\mathbf{y}}\varphi^{-1})|e^{-\frac{\beta}{2 }\mathbf{y}^{\top}\mathbf{A}^{-1}\mathbf{y}}\mathrm{d}\mathbf{y}\\ \leq\frac{\sup_{\mathbf{y}\in\mathbb{R}^{d}}\{|\det(D_{\mathbf{y}}\varphi^{-1} )|e^{-\frac{\beta}{2}\mathbf{y}^{\top}\mathbf{A}^{-1}\mathbf{y}}\}}{C}\sum_{k= d_{\varepsilon}+1}^{d}\int_{\mathbb{R}^{d}}\mathbf{y}_{i_{k}}^{2}e^{-(\frac{1} {2}-\frac{\alpha}{2}-\frac{\beta}{2})\mathbf{y}^{\top}\mathbf{A}^{-1}\mathbf{y }}\mathrm{d}\mathbf{y}\\ \overset{\text{\lx@cref{creftype~refnum}{eq:RAEinvdiffeoRegConstant}}}{=}\frac {C_{\beta,\varphi}}{C}\sum_{k=d_{\varepsilon}+1}^{d}\int_{\mathbb{R}^{d}} \mathbf{y}_{i_{k}}^{2}e^{-(\frac{1}{2}-\frac{\alpha}{2}-\frac{\beta}{2}) \mathbf{y}^{\top}\mathbf{A}^{-1}\mathbf{y}}\mathrm{d}\mathbf{y}=\frac{C_{\beta ,\varphi}}{C}\sum_{k=d_{\varepsilon}+1}^{d}\int_{\mathbb{R}^{d}}\mathbf{y}_{i_ {k}}^{2}e^{-(\frac{1}{2}-\frac{\alpha}{2}-\frac{\beta}{2})\sum_{j=1}^{d}\frac{ \mathbf{y}_{j}^{2}}{\mathbf{a}_{j}}}\mathrm{d}\mathbf{y}\\ =\frac{C_{\beta,\varphi}}{C}\sum_{k=d_{\varepsilon}+1}^{d}\int_{\mathbb{R}} \mathbf{y}_{i_{k}}^{2}e^{-(\frac{1}{2}-\frac{\alpha}{2}-\frac{\beta}{2})\frac{ \mathbf{y}^{2}}{\mathbf{a}_{i_{k}}}}\mathrm{d}\mathbf{y}_{i_{k}}\int_{\mathbb{ R}^{d-1}}e^{-(\frac{1}{2}-\frac{\alpha}{2}-\frac{\beta}{2})\sum_{j\neq i_{k}}^ {d}\frac{\mathbf{y}_{j}^{2}}{\mathbf{a}_{j}}}\mathrm{d}\mathbf{y}_{1}\ldots \mathrm{d}\mathbf{y}_{i_{k}-1}\mathrm{d}\mathbf{y}_{i_{k}+1}\ldots\mathrm{d} \mathbf{y}_{d}\\ =\frac{C_{\beta,\varphi}}{C}\sum_{k=d_{\varepsilon}+1}^{d}\frac{\mathbf{a}_{i_ {k}}}{(1-\alpha-\beta)}\int_{\mathbb{R}}e^{-(\frac{1}{2}-\frac{\alpha}{2}- \frac{\beta}{2})\frac{\mathbf{y}^{2}}{\mathbf{a}_{i_{k}}}}\mathrm{d}\mathbf{y} _{i_{k}}\int_{\mathbb{R}^{d-1}}e^{-(\frac{1}{2}-\frac{\alpha}{2}-\frac{\beta}{ 2})\sum_{j\neq i_{k}}^{d}\frac{\mathbf{y}_{j}^{2}}{\mathbf{a}_{j}}}\mathrm{d} \mathbf{y}_{1}\ldots\mathrm{d}\mathbf{y}_{i_{k}-1}\mathrm{d}\mathbf{y}_{i_{k}+ 1}\ldots\mathrm{d}\mathbf{y}_{d}\\ =\frac{C_{\beta,\varphi}}{C}\sum_{k=d_{\varepsilon}+1}^{d}\frac{\mathbf{a}_{i_ {k}}}{(1-\alpha-\beta)}\int_{\mathbb{R}^{d}}e^{-(\frac{1}{2}-\frac{\alpha}{2}- \frac{\beta}{2})\mathbf{y}^{\top}\mathbf{A}^{-1}\mathbf{y}}\mathrm{d}\mathbf{y }\\ =\frac{C_{\beta,\varphi}}{C}\sum_{k=d_{\varepsilon}+1}^{d}\frac{\mathbf{a}_{i_ {k}}}{(1-\alpha-\beta)}\Bigl{(}\frac{1+\beta}{1-\alpha-\beta}\Bigr{)}^{\frac{d }{2}}\int_{\mathbb{R}^{d}}e^{-(\frac{1}{2}+\frac{\beta}{2})\mathbf{y}^{\top} \mathbf{A}^{-1}\mathbf{y}}\mathrm{d}\mathbf{y}\\ \overset{\mathbf{y}=\varphi(\mathbf{x})}{=}\frac{C_{\beta,\varphi}}{C}\sum_{k= d_{\varepsilon}+1}^{d}\frac{\mathbf{a}_{i_{k}}}{(1-\alpha-\beta)}\Bigl{(}\frac {1+\beta}{1-\alpha-\beta}\Bigr{)}^{\frac{d}{2}}\int_{\mathbb{R}^{d}}e^{-(\frac {1}{2}+\frac{\beta}{2})\varphi(\mathbf{x})^{\top}\mathbf{A}^{-1}\varphi( \mathbf{x})}|\det(D_{\mathbf{x}}\varphi)|\mathrm{d}\mathbf{x}\\ =\frac{C_{\beta,\varphi}}{C}\sum_{k=d_{\varepsilon}+1}^{d}\frac{\mathbf{a}_{i_ {k}}}{(1-\alpha-\beta)}\Bigl{(}\frac{1+\beta}{1-\alpha-\beta}\Bigr{)}^{\frac{d }{2}}\int_{\mathbb{R}^{d}}e^{-\frac{1}{2}\varphi(\mathbf{x})^{\top}\mathbf{A}^ {-1}\varphi(\mathbf{x})}|\det(D_{\mathbf{x}}\varphi)|e^{-\frac{\beta}{2} \varphi(\mathbf{x})^{\top}\mathbf{A}^{-1}\varphi(\mathbf{x})}\mathrm{d}\mathbf {x}\\ \leq\frac{C_{\beta,\varphi}\sup_{\mathbf{x}\in\mathbb{R}^{d}}\{|\det(D_{ \mathbf{x}}\varphi)|e^{-\frac{\beta}{2}\varphi(\mathbf{x})^{\top}\mathbf{A}^{- 1}\varphi(\mathbf{x})}\}}{C}\sum_{k=d_{\varepsilon}+1}^{d}\frac{\mathbf{a}_{i_ {k}}}{(1-\alpha-\beta)}\Bigl{(}\frac{1+\beta}{1-\alpha-\beta}\Bigr{)}^{\frac{d }{2}}\int_{\mathbb{R}^{d}}e^{-\frac{1}{2}\varphi(\mathbf{x})^{\top}\mathbf{A}^ {-1}\varphi(\mathbf{x})}\mathrm{d}\mathbf{x}\\ \overset{\text{\lx@cref{creftype~refnum}{eq:RAEdiffeoRegConstant}}}{=}\frac{C_ {\beta,\varphi}C_{\beta,\varphi^{-1}}}{C}\sum_{k=d_{\varepsilon}+1}^{d}\frac{ \mathbf{a}_{i_{k}}}{(1-\alpha-\beta)}\Bigl{(}\frac{1+\beta}{1-\alpha-\beta} \Bigr{)}^{\frac{d}{2}}\int_{\mathbb{R}^{d}}e^{-\frac{1}{2}\varphi(\mathbf{x})^ {\top}\mathbf{A}^{-1}\varphi(\mathbf{x})}\mathrm{d}\mathbf{x}\\ \overset{\text{\lx@cref{creftype~refnum}{eq:RAEnormalizationConstant}}}{=} \frac{C_{\beta,\varphi}C_{\beta,\varphi^{-1}}}{1-\alpha-\beta}\Bigl{(}\frac{1+ \beta}{1-\alpha-\beta}\Bigr{)}^{\frac{d}{2}}\sum_{k=d_{\varepsilon}+1}^{d} \mathbf{a}_{i_{k}}\\ \overset{\text{\lx@cref{creftype~refnum}{eq:dimind-epsilon}}}{\leq}\varepsilon \frac{C_{\beta,\varphi}C_{\beta,\varphi^{-1}}}{1-\alpha-\beta}\Bigl{(}\frac{1+ \beta}{1-\alpha-\beta}\Bigr{)}^{\frac{d}{2}}\sum_{i=1}^{d}\mathbf{a}_{i}.$ (35) $$

∎

###### Lemma 1 .

###### Proof.

##### Proof of the theorem

###### Proof of theorem 1 .

First, consider the Taylor approximation

$$ $\varphi^{-1}(\varphi(\mathbf{y}))-\varphi^{-1}(\varphi(\mathbf{y}))=D_{\varphi (\mathbf{x})}\varphi^{-1}[\varphi(\mathbf{y})-\varphi(\mathbf{x})]+\mathcal{O} (\|\varphi(\mathbf{y})-\varphi(\mathbf{x})\|_{2}^{2})\\ =D_{\varphi(\mathbf{x})}\varphi^{-1}[\varphi(\mathbf{y})-\varphi(\mathbf{x})]+ \mathcal{O}(d_{\mathbb{R}^{d}}^{\varphi}(\mathbf{y},\mathbf{x})^{2}).$ (36) $$

Moreover, we define

$$ $C:=\int_{\mathbb{R}^{d}}e^{-\frac{1}{2}\varphi(\mathbf{x})^{\top}\mathbf{A}^{- 1}\varphi(\mathbf{x})}\mathrm{d}\mathbf{x}.$ (37) $$

Subsequently, notice that

$$ $\mathbb{E}_{\mathbf{X}\sim p}[\|D_{\varphi(\mathbf{X})}\varphi^{-1}[\varphi(D_ {\varepsilon}(E_{\varepsilon}(\mathbf{X})))-\varphi(\mathbf{X})]\|_{2}^{2}]\\ =\frac{1}{C}\int_{\mathbb{R}^{d}}\|D_{\varphi(\mathbf{x})}\varphi^{-1}[\varphi (D_{\varepsilon}(E_{\varepsilon}(\mathbf{x})))-\varphi(\mathbf{x})]\|_{2}^{2}e ^{-\frac{1}{2}\varphi(\mathbf{x})^{\top}\mathbf{A}^{-1}\varphi(\mathbf{x})} \mathrm{d}\mathbf{x}\\ \leq\frac{1}{C}\int_{\mathbb{R}^{d}}\|D_{\varphi(\mathbf{x})}\varphi^{-1}\|_{2 }^{2}\|\varphi(D_{\varepsilon}(E_{\varepsilon}(\mathbf{x})))-\varphi(\mathbf{x })\|_{2}^{2}e^{-\frac{1}{2}\varphi(\mathbf{x})^{\top}\mathbf{A}^{-1}\varphi( \mathbf{x})}\mathrm{d}\mathbf{x}\\ \leq\frac{\sup_{\mathbf{x}\in\mathbb{R}^{d}}\{\|D_{\varphi(\mathbf{x})}\varphi ^{-1}\|_{2}^{2}e^{-\frac{\beta}{2}\varphi(\mathbf{x})^{\top}\mathbf{A}^{-1} \varphi(\mathbf{x})}\}}{C}\int_{\mathbb{R}^{d}}\|\varphi(D_{\varepsilon}(E_{ \varepsilon}(\mathbf{x})))-\varphi(\mathbf{x})\|_{2}^{2}e^{-(\frac{1}{2}-\frac {\beta}{2})\varphi(\mathbf{x})^{\top}\mathbf{A}^{-1}\varphi(\mathbf{x})} \mathrm{d}\mathbf{x}\\ \overset{\text{\lx@cref{creftype~refnum}{eq:RAEinvdiffeoRegConstantB}}}{=} \frac{B_{\beta,\varphi^{-1}}}{C}\int_{\mathbb{R}^{d}}\|\varphi(D_{\varepsilon} (E_{\varepsilon}(\mathbf{x})))-\varphi(\mathbf{x})\|_{2}^{2}e^{\frac{\beta}{2} \varphi(\mathbf{x})^{\top}\mathbf{A}^{-1}\varphi(\mathbf{x})}e^{-\frac{1}{2} \varphi(\mathbf{x})^{\top}\mathbf{A}^{-1}\varphi(\mathbf{x})}\mathrm{d}\mathbf {x}\\ =B_{\beta,\varphi^{-1}}\mathbb{E}_{\mathbf{X}\sim p}[d_{\mathbb{R}^{d}}^{ \varphi}(D_{\varepsilon}(E_{\varepsilon}(\mathbf{X})),\mathbf{X})^{2}e^{\frac{ \beta}{2}\varphi(\mathbf{X})^{\top}\mathbf{A}^{-1}\varphi(\mathbf{X})}]\\ \overset{\text{\lx@cref{creftype~refnum}{lem:rae-bound-phi-metric}}}{\leq} \varepsilon\frac{B_{\beta,\varphi^{-1}}C_{\beta,\varphi}C_{\beta,\varphi^{-1}} }{1-2\beta}\Bigl{(}\frac{1+\beta}{1-2\beta}\Bigr{)}^{\frac{d}{2}}\sum_{i=1}^{d }\mathbf{a}_{i}.$ (38) $$

Then,

$$ $\mathbb{E}_{\mathbf{X}\sim p}[\|D_{\varepsilon}(E_{\varepsilon}(\mathbf{X}))- \mathbf{X}\|_{2}^{2}]=\mathbb{E}_{\mathbf{X}\sim p}[\|\varphi^{-1}(\varphi(D_{ \varepsilon}(E_{\varepsilon}(\mathbf{X}))))-\varphi^{-1}(\varphi(\mathbf{X})) \|_{2}^{2}]\\ \overset{\lx@cref{creftype~refnum}{eq:thm-l2-bound-rae-taylor-phiinv-phi}}{=} \mathbb{E}_{\mathbf{X}\sim p}[\|D_{\varphi(\mathbf{X})}\varphi^{-1}[\varphi(D_ {\varepsilon}(E_{\varepsilon}(\mathbf{X})))-\varphi(\mathbf{X})]+\mathcal{O}(d _{\mathbb{R}^{d}}^{\varphi}(D_{\varepsilon}(E_{\varepsilon}(\mathbf{X})), \mathbf{X})^{2})\|_{2}^{2}]\\ =\mathbb{E}_{\mathbf{X}\sim p}[\|D_{\varphi(\mathbf{X})}\varphi^{-1}[\varphi(D _{\varepsilon}(E_{\varepsilon}(\mathbf{X})))-\varphi(\mathbf{X})]\|_{2}^{2}+ \mathcal{O}(d_{\mathbb{R}^{d}}^{\varphi}(D_{\varepsilon}(E_{\varepsilon}( \mathbf{X})),\mathbf{X})^{3})]\\ \overset{\text{\lx@cref{creftype~refnum}{eq:thm-rae-ell2-bound-first-order-ter m}}}{\leq}\varepsilon\frac{B_{\beta,\varphi^{-1}}C_{\beta,\varphi}C_{\beta, \varphi^{-1}}}{1-2\beta}\Bigl{(}\frac{1+\beta}{1-2\beta}\Bigr{)}^{\frac{d}{2}} \sum_{i=1}^{d}\mathbf{a}_{i}+o(\varepsilon),$ (39) $$

which yields the claim as $\beta$ was arbitrary.
∎

###### Proof of theorem 1 .

## Appendix B Dataset Construction Details

In this section, we provide a detailed explanation of the construction of the datasets used in our experiments. We organize the datasets into two categories based on the experimental sections in which they are used.

### B.1 Datasets for Manifold Mapping Experiments

Figure: (a) Single Banana
Refer to caption: extracted/5896858/results/visualisations/datasets/single_banana.png

In our manifold mapping experiments (Section 6.1), we use the following datasets illustrated in Figure 6:

- •
Single Banana Dataset: A two-dimensional dataset shaped like a curved banana.
- •
Squeezed Single Banana Dataset: A variant of the Single Banana with a tighter bend.
- •
River Dataset: A more complex 2D dataset resembling the meandering path of a river.

Each dataset is constructed by defining specific diffeomorphisms $\varphi$ and convex quadratic functions $\psi$, then sampling from the resulting probability density using Langevin Monte Carlo Markov Chain (MCMC) with Metropolis-Hastings correction. The probability density function is defined as:

$$ $p(\mathbf{x})\propto e^{-\psi(\varphi(\mathbf{x}))},$ (40) $$

where the strongly convex function $\psi$ is given by:

$$ $\psi(\mathbf{v})=\frac{1}{2}\mathbf{v}^{\top}A^{-1}\mathbf{v},$ (41) $$

and $A$ is a positive-definite diagonal matrix. The specific choices of $\varphi$ and $A$ for each dataset determine its geometric properties.

#### B.1.1 Diffeomorphisms and Convex Quadratic Functions

The key differences between the datasets arise from the diffeomorphism $\varphi$ and the covariance matrix $\mathbf{A}$ used in the sampling process. Below, we describe the specific settings for each dataset.

##### 1. Single Banana Dataset

- •
Diffeomorphism:
$\varphi(\mathbf{x})=\begin{pmatrix}x_{1}-ax_{2}^{2}-z\\
x_{2}\end{pmatrix}$
where $a=\frac{1}{9}$ and $z=0$.
- •
Covariance matrix:
$\mathbf{A}=\begin{pmatrix}\frac{1}{4}&0\\
0&4\end{pmatrix}$

##### 2. Squeezed Single Banana Dataset

- •
Diffeomorphism: Same as the Single Banana Dataset.
- •
Covariance matrix:
$\mathbf{A}=\begin{pmatrix}\frac{1}{81}&0\\
0&4\end{pmatrix}$

##### 3. River Dataset

- •
Diffeomorphism:
$\varphi(\mathbf{x})=\begin{pmatrix}x_{1}-\sin(ax_{2})-z\\
x_{2}\end{pmatrix}$
where $a=2$ and $z=0$.
- •
Covariance matrix:
$\mathbf{A}=\begin{pmatrix}\frac{1}{25}&0\\
0&3\end{pmatrix}$

#### B.1.2 Dataset Generation Algorithm

Algorithm 1 outlines the dataset generation process for all three datasets. The specific diffeomorphisms and quadratic functions differ for each dataset.

Figure: Algorithm 1 General Dataset Generation Algorithm

### B.2 Datasets for Riemannian Autoencoder Experiments

In the Riemannian autoencoder experiments (Section 6.2), we use the following datasets:

- •
Hemisphere(d’, d) Dataset: Samples drawn from the upper hemisphere of a d’-dimensional unit sphere and embedded into $\mathbb{R}^{d}$ via a random isometric mapping.
- •
Sinusoid(d’, d) Dataset: Generated by applying sinusoidal transformations to d’-dimensional latent variables, resulting in a complex, nonlinear manifold in $\mathbb{R}^{d}$.

### B.3 Hemisphere( d ′ superscript 𝑑 ′ d^{\prime} italic_d start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT , d 𝑑 d italic_d ) Dataset

The Hemisphere($d^{\prime}$, $d$) dataset consists of samples drawn from the upper hemisphere of a $d^{\prime}$-dimensional unit sphere, which are then embedded into a $d$-dimensional ambient space using a random isometric embedding. Below are the steps involved in constructing this dataset.

##### 1. Sampling from the Upper Hemisphere

We begin by sampling points from the upper hemisphere of the $d^{\prime}$-dimensional unit sphere $S^{d^{\prime}}_{+}\subset\mathbb{R}^{d^{\prime}+1}$. The upper hemisphere is defined as:

$$ $S^{d^{\prime}}_{+}=\left\{\mathbf{x}\in\mathbb{R}^{d^{\prime}+1}:\|\mathbf{x} \|=1,\,x_{1}\geq 0\right\}.$ $$

The first angular coordinate $\theta_{1}$ is sampled from a Beta distribution with shape parameters $\alpha=5$ and $\beta=5$, scaled to the interval $\left[0,\frac{\pi}{2}\right]$. This sampling method emphasizes points near the “equator” of the hemisphere. The remaining angular coordinates $\theta_{2},\ldots,\theta_{d^{\prime}}$ are sampled uniformly from the interval $\left[0,\pi\right]$:

$$ $\theta_{1}\sim\text{Beta}(5,5)\cdot\left(\frac{\pi}{2}\right),\quad\theta_{i} \sim\text{Uniform}(0,\pi),\,\text{for }i=2,\ldots,d^{\prime}.$ $$

##### 2. Conversion to Cartesian Coordinates

Next, each sampled point in spherical coordinates is converted into Cartesian coordinates in $\mathbb{R}^{d^{\prime}+1}$ using the following transformation equations:

$$ $x_{1}=\cos(\theta_{1}),\quad x_{2}=\sin(\theta_{1})\cos(\theta_{2}),\quad\dots ,\quad x_{d^{\prime}+1}=\sin(\theta_{1})\sin(\theta_{2})\cdots\sin(\theta_{d^{ \prime}}).$ $$

This conversion ensures that the sampled points lie on the surface of the unit sphere in $(d^{\prime}+1)$-dimensional space.

##### 3. Random Isometric Embedding into ℝ d superscript ℝ 𝑑 \mathbb{R}^{d} blackboard_R start_POSTSUPERSCRIPT italic_d end_POSTSUPERSCRIPT

After sampling points on the hemisphere in $\mathbb{R}^{d^{\prime}+1}$, the points are embedded into a $d$-dimensional ambient space ($d\geq d^{\prime}+1$) using a random isometric embedding. The embedding process is as follows:

- 1.
Generate a random matrix $\mathbf{A}\in\mathbb{R}^{d\times(d^{\prime}+1)}$, where each entry is sampled from a standard normal distribution $\mathcal{N}(0,1)$.
- 2.
Perform a QR decomposition on matrix $\mathbf{A}$ to obtain $\mathbf{Q}\in\mathbb{R}^{d\times(d^{\prime}+1)}$:
$\mathbf{A}=\mathbf{Q}\mathbf{R}.$
The columns of $\mathbf{Q}$ form an orthonormal basis for a $(d^{\prime}+1)$-dimensional subspace of $\mathbb{R}^{d}$, ensuring that $\mathbf{Q}$ defines an isometric embedding from $\mathbb{R}^{d^{\prime}+1}$ into $\mathbb{R}^{d}$. This guarantees that distances and angles are preserved during the mapping, maintaining the geometric structure of the original space within the higher-dimensional ambient space.
- 3.
Use matrix $\mathbf{Q}$ to map each sample $\mathbf{x}\in\mathbb{R}^{d^{\prime}+1}$ into the ambient space:
$\mathbf{y}=\mathbf{Q}\mathbf{x},$
where $\mathbf{y}\in\mathbb{R}^{d}$ are the embedded samples.

Figure: Algorithm 2 Hemisphere($d^{\prime}$, $d$) Dataset Generation

### B.4 Sinusoid( d ′ superscript 𝑑 ′ d^{\prime} italic_d start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT , d 𝑑 d italic_d ) Dataset

The Sinusoid($d^{\prime}$, $d$) dataset represents a $d^{\prime}$-dimensional manifold embedded in $d$-dimensional space through nonlinear sinusoidal transformations. Below are the detailed steps involved in constructing this dataset.

##### 1. Sampling Latent Variables

The latent variables $\mathbf{z}\in\mathbb{R}^{d^{\prime}}$ are sampled from a multivariate Gaussian distribution with zero mean and isotropic variance, as follows:

$$ $\mathbf{z}\sim\mathcal{N}\left(0,\sigma_{m}^{2}I_{d^{\prime}}\right),$ $$

where $\sigma_{m}^{2}$ controls the variance along each intrinsic dimension, and $I_{d^{\prime}}$ is the $d^{\prime}\times d^{\prime}$ identity matrix. The value of $\sigma_{m}^{2}$ is set to 3 for our experiments.

##### 2. Defining Ambient Coordinates with Sinusoidal Transformations

For each of the $d-d^{\prime}$ ambient dimensions, we construct a shear vector $\mathbf{a}_{j}\in\mathbb{R}^{d^{\prime}}$, with its elements drawn uniformly from the interval $[1,2]$:

$$ $\mathbf{a}_{j}\sim\text{Uniform}(1,2)^{d^{\prime}},\quad\text{for }j=1,\dots,d -d^{\prime}.$ $$

The shear vectors $\mathbf{a}_{j}$ apply a fixed linear transformation to the latent space $\mathbf{z}\in\mathbb{R}^{d^{\prime}}$, determining how the latent variables influence each ambient dimension. These vectors, sampled once for each of the $d-d^{\prime}$ ambient dimensions, modulate the scale and periodicity of the sinusoidal transformation.

Each ambient coordinate $x_{j}$ is generated as a sinusoidal function of the inner product between $\mathbf{a}_{j}$ and $\mathbf{z}$, with a small Gaussian noise added for regularization.

$$ $x_{j}=\sin\left(\mathbf{a}_{j}^{\top}\mathbf{z}\right)+\epsilon_{j},$ $$

where $\epsilon_{j}\sim\mathcal{N}(0,\sigma_{a}^{2})$ is Gaussian noise with variance $\sigma_{a}^{2}$. In our experiments, we set $\sigma_{a}^{2}=10^{-3}$.

##### 3. Constructing the Dataset Samples

The final samples $\mathbf{y}\in\mathbb{R}^{d}$ are formed by concatenating the ambient coordinates $x_{1},x_{2},\dots,x_{d-d^{\prime}}$ with the latent variables $z_{1},z_{2},\dots,z_{d^{\prime}}$:

$$ $\mathbf{y}=\left[x_{1},x_{2},\dots,x_{d-d^{\prime}},\,z_{1},z_{2},\dots,z_{d^{ \prime}}\right]^{\top}.$ $$

Figure: Algorithm 3 Sinusoid($d^{\prime}$, $d$) Dataset Generation

## Appendix C Error Metrics for Evaluation of Pullback Geometries

##### Geodesic Error.

The geodesic error measures the difference between geodesics on the learned and ground truth pullback manifolds. Given two points $\mathbf{x}_{0},\mathbf{x}_{1}\in\mathbb{R}^{d}$, let $\gamma_{\mathbf{x}_{0},\mathbf{x}_{1}}^{\varphi_{\theta_{2}}}(t)$ and $\gamma_{\mathbf{x}_{0},\mathbf{x}_{1}}^{\varphi_{\mathrm{GT}}}(t)$ denote the geodesics induced by the learned map $\varphi_{\theta_{2}}$ and the ground truth map $\varphi_{\mathrm{GT}}$, respectively, where $t\in[0,1]$.

The geodesic error is calculated as the mean Euclidean distance between the learned and ground truth geodesics over $N$ pairs of points:

$$ $\text{Geodesic Error}=\frac{1}{N}\sum_{i=1}^{N}\frac{1}{T}\sum_{k=1}^{T}\left \|\gamma_{\mathbf{x}_{0}^{(i)},\mathbf{x}_{1}^{(i)}}^{\varphi_{\theta_{2}}}(t_ {k})-\gamma_{\mathbf{x}_{0}^{(i)},\mathbf{x}_{1}^{(i)}}^{\varphi_{\mathrm{GT}} }(t_{k})\right\|_{2},$ $$

where $T$ is the number of time steps used to discretize the geodesic, and $t_{k}=\frac{k-1}{T-1}$ for $k=1,\dots,T$.

This metric captures the average discrepancy between the learned and ground truth geodesics, reflecting the accuracy of the learned pullback manifold.

##### Variation Error.

The variation error quantifies the sensitivity of the geodesic computation under small perturbations to one of the endpoints. For two points $\mathbf{x}_{0},\mathbf{x}_{1}\in\mathbb{R}^{d}$, let $\mathbf{z}=\mathbf{x}_{1}+\Delta\mathbf{x}$, where $\Delta\mathbf{x}$ is a small perturbation applied to $\mathbf{x}_{1}$. Define $\gamma_{\mathbf{x}_{0},\mathbf{x}_{1}}^{\varphi_{\theta_{2}}}(t)$ and $\gamma_{\mathbf{x}_{0},\mathbf{z}}^{\varphi_{\theta_{2}}}(t)$ as the geodesics from $\mathbf{x}_{0}$ to $\mathbf{x}_{1}$ and $\mathbf{z}$, respectively, induced by the learned map $\varphi_{\theta_{2}}$.

The variation error is calculated as the mean Euclidean distance between the geodesic from $\mathbf{x}_{0}$ to $\mathbf{x}_{1}$ and the perturbed geodesic from $\mathbf{x}_{0}$ to $\mathbf{z}$:

$$ $\text{Variation Error}=\frac{1}{N}\sum_{i=1}^{N}\frac{1}{T}\sum_{k=1}^{T}\left \|\gamma_{\mathbf{x}_{0}^{(i)},\mathbf{x}_{1}^{(i)}}^{\varphi_{\theta_{2}}}(t_ {k})-\gamma_{\mathbf{x}_{0}^{(i)},\mathbf{z}^{(i)}}^{\varphi_{\theta_{2}}}(t_{ k})\right\|_{2},$ $$

where $N$ is the number of sampled point pairs, $T$ is the number of time steps used to discretize the geodesic, and $t_{k}=\frac{k-1}{T-1}$ for $k=1,\dots,T$.

This metric evaluates the robustness of the learned geodesic against small perturbations, providing insight into the stability of the learned manifold.

## Appendix D Training Details

The following section describes the important configuration parameters for reproducing the experiments on manifold mappings. All experiments share some common parameters, which are listed below, while dataset-specific parameters are provided in Table 2.

Common Parameters:

- •
Optimizer: Adam with betas = (0.9, 0.99), eps = $1\times 10^{-8}$, and weight decay of $1\times 10^{-5}$.
- •
Learning Rate Schedule: Warm-up cosine annealing with 1000 warm-up steps.
- •
Gradient Clipping: Gradient norm clipped to 1.0.
- •
Model Architecture: A composition of affine coupling layers is used, where each layer transforms part of the input while keeping the other part unchanged. The transformation function in each layer is modeled by a residual network (ResNet) consisting of 64 hidden features, 2 residual blocks, ReLU activations, and no batch normalization. Dropout is set to 0, and transformations alternate across different dimensions at each layer.

**Table 2: Training configurations for each experiment.**
| Dataset | Flow Steps | Epochs | Batch Size | $\lambda_{\text{iso}}$ | $\lambda_{\text{vol}}$ | Learning Rate |
| --- | --- | --- | --- | --- | --- | --- |
| Sinusoid(1,3) | 8 | 1000 | 64 | 1.0 | 1.0 | $3\times 10^{-4}$ |
| Sinusoid(2,3) | 8 | 1000 | 64 | 1.0 | 1.0 | $3\times 10^{-4}$ |
| Sinusoid(5,20) | 24 | 2000 | 128 | 1.2 | 2.5 | $4\times 10^{-4}$ |
| Hemisphere(2,3) | 8 | 2000 | 64 | 1.0 | 1.0 | $4\times 10^{-4}$ |
| Hemisphere(5,20) | 12 | 2000 | 64 | 0.75 | 1.2 | $4\times 10^{-4}$ |

## Appendix E Data Manifold Approximation

The learned manifold, shown in orange in Figure 1, is the set $D_{\epsilon}(\mathcal{U})$, where $D_{\epsilon}$ is the RAE decoder eq. 21, the set $\mathcal{U}$ in the latent space is the open set given by

$$ $\mathcal{U}=\prod_{i=1}^{d_{\epsilon}}(-3\sqrt{\mathbf{a}_{u_{i}}},3\sqrt{ \mathbf{a}_{u_{i}}})$ $$

and $\mathbf{a}_{u_{1}},\dots,\mathbf{a}_{u_{d_{\epsilon}}}$ are the $d_{\epsilon}$ highest learned variances corresponding to the ones used in the RAE construction.

To visualize this in practice, we construct a mesh grid by linearly sampling each latent dimension from $-3\sqrt{\mathbf{a}_{u_{i}}}$ to $+3\sqrt{\mathbf{a}_{u_{i}}}$, for $i=1,\dots,d_{\epsilon}$, where $d_{\epsilon}$ is the number of significant latent dimensions. Practically, the off-manifold latent dimensions (those corresponding to negligible variances) are set to zero. The decoder $D_{\epsilon}$ then maps this grid from $\mathcal{U}$ back to $\mathbb{R}^{d}$, generating an approximation of the data manifold, as illustrated in Figure 1.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: In Table 1, how are the starting and goal points for geodesics, as well as the perturbations, selected when evaluating the geodesic and variation errors?
soundness: 2
strengths: - The proposed geometry seems novel, with intriguing closed-form solutions derived from the geometry.
- Most of the derivations seem correct (I could not verify all the details).
summary: This paper introduces a score-based pullback Riemannian geometry—a data-driven geometry in which geodesics follow the data supports. Assuming specific unimodal densities, it derives closed-form expressions for geodesics, distances, the exponential map, and the logarithmic map, and formulates a Riemannian autoencoder with an error bound. Additionally, it proposes a generative model by adapting normalizing flow as a learning algorithm. The model is tested on synthetic datasets, demonstrating its capability to compute geodesics along the data manifold, estimate intrinsic dimensions, and provide global coordinate charts.
weaknesses: - The properties of the proposed geometry are insufficiently discussed and lack intuitive explanations. For example, it would be helpful to explain how the geodesics behave and the form of the Riemannian metrics for different distributions. Without such clarification, the current title may overstate the paper’s contributions.
- Practical applications of the proposed geometry, beyond the Riemannian autoencoder, are not clearly identified. For the Riemannian autoencoder, it remains unclear in what contexts it outperforms alternative methods.
- Furthermore, the unimodal density assumption is restrictive; because computational efficiency depends heavily on this assumption, extending the method to multimodal cases while retaining computational advantages seems challenging.
- The experiments are limited to synthetic data, without comparisons to similar Riemannian geometry methods, such as those mentioned in the introduction (e.g., Arvanitidis et al. 2016).

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. As I mentioned in the weakness section, although the quasi-Gaussian assumption of the data is comprehensible, can we generalize it to another class of density functions?
2. Do we have any computational advantage when using this method comparing with the other ones (for ex, NF, Anisotropic NF, Isometric NF)? 
3.  Can you elaborate on the intuition of using the score function in this article?
soundness: 3
strengths: In overall, the approach in this article is both innovative and compelling. Specifically:
1. The Riemannian metric structure proposed in the article is notably original, with a comprehensive examination of its fundamental geometric properties. 
2. The proof supporting auto-encoder and decoder mechanism is well-articulated and appears to be mathematically rigorous.
3. The paper is clearly written.
summary: In the article, the authors introduce the scored-base Riemannian structure that utilizes pullback Riemannian geometry in conjunction with generative models. Specifically:
(1) Based on Gaussian-like distribution in the Euclidean space, they define a Riemannian metric, which is linked to the pullback of the score function.
(2) A comprehensive study about this metric is provided, including the geodesics, logarithmic map, exponential map, distance, and barycenter. 
(3) They propose a novel auto-encoder and decoder methods, based on the metric, offering theoretical guarantees on the consistency of estimation (Theorem 1). 
(4) A mechanism for learning probability densities is introduced, utilizing an adapted normalizing flow loss function, with multiple experiments.
weaknesses: 1. The data used in the study is somewhat artificial. The work would benefit from employing more realistic datasets.
2. Although the methodology is quite straightforward, it would be preferable if the authors included step-by-step derivations to verify the  geometric properties (mentioned in Proposition 1). 
3. Given that the data is generated from quadratic structures, the model appears to perform well only for data with similar characteristics (such as the banana-shaped data in the paper). For more complex and diverse data structures, the model may require further development to capture such complexities.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. What is the computational complexity of the proposed approach? How does it scale with the intrinsic dimension of the data dimension and the ambient dimension? 
2. What is the main difficulty to extend it to multimodal data?
soundness: 3
strengths: 1.	This framework facilitates the learning of interpretable representations by leveraging Riemannian geometry, which can capture complex data structures more effectively than traditional methods.
2.	Under the assumption of the unimodal distribution, the score-based approach facilitates the construction of the pullback geometry with closed-form manifold mappings.
3.	The paper also showed that the resulting geodesics always pass through the support of data probability
summary: This work introduces score-based pullback Riemannian geometry, a scalable framework that combines elements of pullback Riemannian geometry and generative models to enhance interpretable representation learning. Focusing on unimodal distributions, it features a Riemannian autoencoder (RAE) with closed-form geodesics that traverse the data's probability density, allowing for accurate estimation of the data manifold's dimension. Numerical experiments show that the framework effectively produces high-quality geodesics and reliably assesses intrinsic dimensions, even in high-dimensional spaces.
weaknesses: 1.	The proposed approach has a strong assumption on the data distribution (unimodal distributions), which limits its applicability to more complex or multimodal data.
2.	In the experimental results, the proposed method is only compared with some discrete-time normalizing flow methods (NF, Anisotropic NF, Isometric NF). It is too limited.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 1
questions: Can you provide the motivation why tangent space similarity obtained from density function is related to metric?
Can you provide a detailed procedure for learning?
The figures should be explained better.
soundness: 2
strengths: The notion of point similarity using Riemanian metric from generative models is interesting. The algorithm works well with synthetic data upon complex manifold.
summary: The authors propose a metric between data samples based on the score information of probability density. Many notations are used without clear definitions. I can assume p in line 135 represents a point in the data manifold, and the score is a vector in the tangent space according to Eq. (5). The tangent space similarity determines the metric between sample points. Because the probability density in use adopts the form of energy-based model with a convex energy psi, the equations can be formulated using the gradient of psi. Along with the learning objective function from normalizing flow (NF), a data generating algorithm can be developed. When an NF is employed, the NF neural network is the diffeomorphism mentioned in the text. It is unclear how the metric obtained from D_X\phi_{\theta_2} information makes sense.
weaknesses: The authors need to provide clear explanations of how the equations are motivated and how the algorithms can be implemented.
Real-world experiment is missing. The authors made a strong manifold assumption, and it is necessary to know how the algorithm works with data having noisy manifold.
What do the contour lines in Figure 2 represent? There is no information what the horizontal and vertical axes represent nor the meaning of the contour lines in the figure.
The paper lacks motivation for why the score information should be used in the metric. Why not simply use \nabla p?

Meta-review: In the paper, the authors proposed score-based pullback Riemannian geometry, a generative data-driven framework that integrates Riemannian geometry.

While the reviewers agree that the proposed method is quite interesting and intuitive, there are several weaknesses of the current paper: (1) the computational efficiency depends heavily on the unimodal density assumption, which is restrictive. Adapting the current method to handle multimodal cases while maintaining computational advantages seems challenging. (2) The experimental results are limited to only synthetic data. It is unclear about the practical performance of the proposed method in the real-world settings. (3) The properties of the proposed geometry are insufficiently discussed and not intuitive. Furthermore, the writing of the current paper is also not clear.

Given the above concerns raised by the reviewers which have not been fully addressed after the rebuttal, I recommend rejecting the paper at the current stage. I encourage the authors to include the suggestions and feedback of the reviewers into the revision of their manuscript.

Decision: Reject