Score-based Pullback Riemannian Geometry: Extracting the Data Manifold Geometry using Anisotropic Flows
Authors: Willem Diepeveen
Georgios Batzolis
Zakhar Shumaylov
Carola-Bibiane Schönlieb
Venue: ICML 2025 poster (2025)
Keywords: data manifold geometry
score-based pullback Riemannian metric
anisotropic normalizing flows
closed-form geodesics
intrinsic dimension estimation
Riemannian auto-encoder
interpretable representation learning
Primary area: deep_learning->generative_models_and_autoencoders
TL;DR: A scalable framework for extracting the geometry of data manifolds

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 diverse datasets, including image data, we demonstrate that the proposed framework produces high-quality geodesics passing through the data support, reliably estimates the intrinsic dimension of the data manifold, and provides a global chart of the manifold. To the best of our knowledge, this is the first scalable framework for extracting the complete geometry of the data manifold.

Score-based Pullback Riemannian Geometry: Extracting the Data Manifold Geometry using Anisotropic Flows
Willem Diepeveen Georgios Batzolis" 2 Zakhar Shumaylov Carola-Bibiane Schönlieb
Abstract
1. Introduction
Data often reside near low-dimensional non-linear mani- folds illustrated in Figure 1. This manifold assumption (Fefferman et al., 2016) has been popular since the early work on non- dimension reduction (Belkin & Niyogi, 2001; Coifiman & Lafon, 2006; Roweis & Saul, 2000; Sam- mon, 1969; Tenenbaum et al., 2000). Learning this non- linear structure, or representation learning, from data has proven to be highly successful (DeMers Cottrell, 1992; Kingma & Welling, 2013) and continues to be recurring theme in modern machine learning approaches and down- stream applications (Chow et al., 2022; Gomari et al., 2022; Ternes et al., 2022; Vahdat & Kautz, 2020; Zhong et al., 2021).
Data-driven Riemannian geometry has emerged as a powerful tool for interpretable representation learning, offering improved efficiency in down- stream tasks. Moving forward, 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 framework for data-driven Riemannian geometry that is scal- able in both geometry and learning: score-based pullback Riemannian geometry. Focusing on uni- modal distributions as first step, we propose score-based Riemannian structure with closed- form geodesies that pass through the data prob- ability density. With this structure, we construct Riemannian autoencoder (RAE) with error bounds for discovering the correct data manifold dimension. This framework can naturally used with anisotropic normalizing flows by adopting isometry regularization during training. Through numerical experiments diverse datasets, includ- ing image data, we demonstrate that the proposed framework produces high-que -ualltyal geodesics pass- ing through the data support, reliably estimates the intrinsic dimension of the data manifold, and provides a global chart of the manifold. To the best of our knowledge, this the first scalable framework for extracting the complete geometry of the data manifold.
Recent advances in data-driven Riemannian geometry have demonstrated its suitability for learning representations. In this context, these representations elements residing in learned geodesic subspace of the data space, governed by a non-trivial Riemannian structure' across the entire ambient space (Arvanitidis et al., 2016; Diepeveen, 2024; Hauberg et al., 2012; Peltonen et al., 2004; Scarvelis & Solomon, 2023; Sorrenson al., 2024; Sun al., 2024). Among these contributions, it is worth highlighting that Sorrenson al. (2024) are the first and only ones far to use information from the full data distribution obtained though generative models (Dinh et al., 2017; Song et al., 2020), even though this seems a natural approach given recent studies such as Horvat & Pfister (2022); Tempczyk et al. (2022); Sakamoto al. (2024); Stanczuk et al. (2024). 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. (2024) also mention themselves that it can be numerically challenging to work with their induced Riemannian geometry.
"Equal contribution 'Department Mathematics, Univer- sity of California, Los Angeles, USA 'Department of Engi- neering, University of Cambridge, UK 'Facuity of Mathemat- ies. University of Cambridge, UK. Correspondence Willem Diepeveen <'ddeepeveen @ waw.t.cca.cccca..mmmm Georgios Batzolis <gb511@ cak..
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 pro- found implications on how to handle data in general. In-
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
'rather than the standard -iinner product

---

Score-based pullback Riemannian geometry
B
(a) Hemisphere (2,3)
(b) Sinusoid (1,3)
(c) Sinusoid (2,3)
Figure 1. Approximate data manifolds learned by the Riemannian autoencoder generated by score-based pullback Riemannian geometry for three datasets. The orange surfaces represent the manifolds learned by the model, while the blue points correspond to the training data. Each manifold provides convincing low-dimensional representation of data, isometric to its respective latent space.
deed, beyond typical data analysis tasks such as comput- ing distances, means, and interollaioon/eettaplllttoo of data points, 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 (Diepeveen et al. (2023); Fletcher et al. (2004) and Chen & Lipman (2023); Huang et al. (2022)) that have proven beneficial by improving upon full black box methods in terms of interpretability (Diepeveen, 2024) or Euclidean counterparts in terms of efficiency (Kapusniak al., 2024; de Kruiff et al., 2024). Here it is worth highlighting that scalability of manifold mappings was completely circum- vented by (Diepeveen, 2024) and (de Kruiff et al., 2024) by using pullback geometry. However, here learning a suitable (and stable) pullback geometry suffers from challenges re- garding scalability of the training algorithm, contrary to the approach by (Sorrenson et al., 2024).
Specifically, we consider a family of unimodal probability densities whose negative log- liglliiessoods are compositions of strongly convex functions and diffeomorphisms. As this work is an attempt bridge between the geometrie data anal- ysis 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.
We demonstrate that this density- based Riemannian struc- ture naturally leads to a Riemannian autoencoder and provide error bounds on the expected reconstruction error, which allows for approximation of the data manifold as illustrated in Figure 1.
We introduce 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.
Motivated by the above, this work aims to address the fol- lowing question: How to strike good balance between scalability of training data-driven Riemannian structure and of evaluating its corresponding manifold mappings?
Practically, this work showcases how two simple adapta- tions to the normalizing flows framework enable data-driven Riemannian geometry. This significantly expands the poten- tial for downstream applications compared to the unadapted framework.
1.1. Contributions
In this paper, we take first steps towards striking such a balance and propose score-based pullback Riemannian metric assuming a relatively simple but generally applicable family of probability densities, which we show result in both scalable manifold mappings and scalable learning algo- rithms. 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.
1.2. Outline
After introducing notation in Section 2, Section 3 considers family of probability distributions, from which we ob- tain suitable geometry, and Section 4 showcases how one can subsequently construct Riemannian Autoencoders with theoretical guarantees. From these observations Section discusses the natural limitations of standard normalizing
in the sense Diepeveen (2024)
2

---

Score-based pullback Riemannian geometry
flows and how to change the parametrisation and training for downstream application in Riemannian geometrie set- ting. Section 6 showcases several use cases of data-driven Riemannian structure on several data sets. Finally, we sum- marize our findings in Section 7.
Throughout the rest of the paper pullback mappings will be denoted similarly Equation (1) with the diffeomorphism & as superscript, i.e., we write dx q). 7p.. expş(Ep) and logp q for p.q E N and Ep E TpN.
3. Riemannian geometry from unimodal probability densities
2. Notation
Here present some basic notations from differential and Riemannian geometry, see Boothby (2003); Carmo (1992); Lee (201)); Sakai (1996) for details.
We remind the reader that the ultimate goal of data-driven Riemannian geometry on Rd is to construct a Riemannian structure such that geodesics always pass through the sup- port of data probability densities. In this section focus constructing Riemannian geometry that does just that from unimodal densities R + IR of the form
Smooth Manifolds and Tangent Spaces. Let M be a smooth manifold. We write C (M) for the space of smooth func- tions over M. The tangent space at p E M, which is defined as the space of all derivations p, is denoted by TpM and for langent vectors we write Ep Tp.M. For the tangent bundle we write TM and smooth vector fields, which are defined as smooth sections of the tangent bundle, are written as 2' (M) C CTM.
will
p(x) x e M((()))
(2)
where 2/ Rd - R is smooth strongly convex function and y IRd - Rd is a diffeomorphism. In particular, we will consider pullback Riemannian structures of the form*
Riemannian Manifolds. A smooth manifold M becomes a Riemannian manifold if equipped with a smoothly varying metric tensor field .): (M) x 2 (M) C (M). This tensor field induces (Riemannian) metric dm: . x M + R. The metric tensor can also be used to construct a unique affine connection, the revi-Civita con- nection, that is denoted by ((-{()) : (M) x 2 (M) + 2 (M). This connection is in turn the cornerstone of a myriad of manifold mappings.
(E,
x x)x Vooy = (Dx o p|E), V o ((()))2
(3)
For proofs of the results below and those of more general statements we refer the reader to Appendix A.
The following result, which is a direct application of (Diepeveen, 2024, Prop. 2.1), gives us closed-form ex- pressions of several important manifold mappings under (., V if we choose
Geodesic. One is the notion of geodesic, which for two points p E M is defined as curve ? [0, 1] + M with minimal length that connects p with q. Another closely
v(x) = 2 A
(4)
where A E Rdxd is symmetric positive definite.
related notion to geodesics is the curve t + (t) for a geodesic starting from p E M with velocity (p.z, (0)
Proposition 3.1. Let up Rdd -> R2d be a smooth diffeo-
morphism and let ub : Rd - IR be a function of the form Equation (4).
Ep T.MM
Exponential Map. This can be used define the exponential map expp: Dp- + M as expp(Ep) := ?p. (I), where Dp CTpMis is the set on which Yp (1) is defined.
Then,
V (x,y) = |A -(((xx) - (2y)))l2,
(5)
3d Veoo (1) = - - ())((x) + (p(()),
Logarithmic Map. Furthermore, the logarithmic map logp: exp(D)) + _ is defined as the inverse of expp so P. it is well-defined on D c 'p where expp is a diffeomor- phism.
xx
(6)
expx Vuop (Ex) (x(x) + Dx
(7)
log Vioo D D(x))p - (p(x) - -xx)].
(8)
Pullback metrics. Finally, if (M, .) is a d-dimensional Riemannian manifold, N is d-dimensional smooth man- ifold and Φ N + M is diffeomorphism, the pullback metric
Remark 3.2. We note that e -tabbiliiy of geodesics inher- ited by (Diepeveen, 2024, Thm. 3.4), if we have (approxi- mate) local l -isometrr of yp on the data support.
A direct result of Proposition 3.1 is that geodesics will pass through the support of p(x) from (2). in the sense that geodesics pass through regions with higher likelihood than the end points. This can be formalized in the following result.
(1)
where DpΦ: : TpN - M denotes the differential of o. defines a Riemannian structure on NV, which we denote by (N, (..)))).. Pullback metrics literally pull back all ge- ometric information from the Riemannian structure on M.
NNote that Vab o y should be read as (Vi) o
3

---

Score-based pullback Riemannian geometry
Theorem 3.3. Let sp: R. + Rd be smooth diffeomor- phism and let db IRd + R. be a function of the form (4).
and consider a threshold e E [0, 1]. We then consider dg [d] defined as the integer that satisfies
Then, mapping
de = dm=iinfa E ld 10111] 1]
Σ
adw
S
au } (12)
1 b (xp (OOP(())))
E [0.
(9)
+1
is strongly convex.
if aug S & Eu=ı ay and de = d otherwise.
The Riemannian structure in Equation (3) is related to the Riemannian structure obtained from the score function V Vlog(p(-)) Rd + R if o is close to smooth local _²_ isometry on the data support, i.e., Dxip is an orthogonal operator:
Finally, we define the encoder (chart) Eg : Rd R
Ex(x) == (logg-1(0) x, Dovp Do 1)==))p-100)
Equation (8)
(f((x), cnt))2 w ....., de, (13)
and define the decoder (inverse chart) De IR - R as
(D&Vlog(b()))))) logg))))))))))))))))))))))))))
= =D2VVW =(D)(()))))
D ( p p)[{))
[E}
D.(p) = exp
Pu Dop
Dopp-)uuu)
Dx ((D()))) T O V&l o p [0);
#==1
~ (D 2 [ D& V ( = x)x Voop. (10)
Equation (7)
(
Σ Pu Pw (14)
For that reason, we call such an approach to data-driven Riemannian geometry: score-based pullback Riemannian geometry.
which generate Riemannian autoencoder and the set De (R=+) C Rd as an approximate data manifold as in the scenario in Figure 1.
4. Riemannian autoencoder from unimodal probability densities
Similarly to classical PCA, 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 B.
Proposition 3.1 begs the question what eh could still be used for. We note that this case comes down to having a data probability density that is deformed Gaussian distribution. In the case of a regular (non-deformed) Gaussian, one can compress the data generated by it through projecting them onto 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 Riemannian au- toencoder (RAE) (Diepeveen, 2024, Sec. 4), whose error we can bound by picking the dimension the autoencoder a clever way, reminiscent of the classical PCA error bound.
Theorem 4.1. Let up IRd Rd be a smooth diffeomor- phism and let tb : Id + IR be a quadratic function of the form Equation (4) with positive definite diagonal matrix A E Rdxdd Furthermore, let p IRd + IK be the corre- sponding probability density of the form Equation (2). Fi- nally, consider & E [0. and the mappings Fig RRd - Rd. and De R.d - Rd in Equations (13) and (14) with de E [d] as in Equation (12).
Then,
EX-pIID, - X
S (T())r(())O)) +o(E), (15)
Concretely, assume that have unimodal density of the form Equation (2) with quadratic strongly convex function v (x) = xx7 A-1x for some diagonal matrix A := diag(a).... a4) with positive entries Next, we de- fine an indexing uw E [d] == [1....,d] d} for w 1....,d d
where
Cl
C2
C8
1 + 8
$
Cp:= inf SE
3.p
3.p
Big
(16)
29
1 28
such that
au3-
Bugz
(11)
for
"Note that the score by itself not always a diffeomorphism. 'Note that this is not restrictive as for a general symmetric positive definite matrix A the eigenvalues can he used as diagonal entries and the orthonormal matrices can he concatenated with the diffeomorphism.
:= sup (I) I1 T 1p(x)), (17) xERd
C3 == sup {|det(Dx())ee f000) 10000) (18)
xERd

---

Score-based pullback Riemannian geometry
and
lows from --isometry, in practice, the flow can only approx- imate local isometry through optimization. Thus, found it beneficial to explicitly include a volume-preservation loss, resulting in an adapted normalizing flow loss that enforces both constraints for improved alignment with the desired Riemannian structure.
C2 == sup (D [x(x))==" A00)). XER
(19)
Remark 4.2. Note that the RAE latent space is interpretable as it is (--isometric to the data manifold if yp is an approx- imate -isomeery 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 VAE (Kingma & Welling. 2013).
[(01, 02) = Ex-puee --logpo ,0 (X)
Aval Xx-le -pea [log([ del D (Dxp0,))
+ XxsoBxxpan [I(()))o,) Dx
Lall'l (21)
where >nal: disa > 0 and the negative log likelihood term reduces to
5. Learning unimodal probability densities
Naturally, we want to learn probability densities of the form Equation (2), which can then directly be inserted into the proposed score-based pullback Riemannian geometry frame- work. In this section we will consider how to adapt normal- izing flow (NF) (Dinh et al., 2017) training setting that is more suitable for our purposes In particular, we will con- sider how training a normalizing flow density p.: Rd - IR given by
1
2 Ex-ran [10g (X) TA 10e2 (X)] + tr(log(Ao,)
2
d
-Exupon [log(( det (Dx )D] + log(2r), (22)
where As, is a diagonal matrix and (Poz is a normalizing flow with affine coupling layers (Dinh et al., 2017). For small ambient dimensions, isometry regularization is feasi- ble, but in high dimensions, computing the full Jacobian is impractical. To address this, we use a scalable sliced isom- etry loss based on Jacobian-vector products, significantly reducing both computational and memory costs while pre- serving effectiveness. See Appendix 1.2 for details.
1 p(x)= = e (-bl)((i)) det(Dxxpp)
(20)
Ci.
where C.. > 0 is a normalisation constant that only de- pends on the strongly convex function vb, yields our target distribution Equation (2).
6. Experiments
From Sections 3 and 4 we have seen that ideally the strongly convex function vi Rdd + I corresponds to a Gaussian with a parameterised diagonal covariance matrix A E Rddd resulting in more parameters than in standard normalizing flows, whereas the diffeomorphism yp : IR - Fd is regu- larized to be isometry. In particular, A ideally allows for learnable anisotropy rather than having 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 (-iisometric, unlike standard normalizing flows which are typically non-volume preserv- ing, enabling stability (Remark 3.2) and practically useful and interpretable RAE (Theorem 4.1 and Remark 4.2).
We conducted two sets of experiments to evaluate the pro- posed scheme from 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 configu- rations are provided in Appendix C and additional results in Appendix D.
6.1. Manifold mappings
In addition, e isismetry (on the data support) implies colume-preservation which means that det(Dx se so that the model density (20) reduces to the desired form of Equation (2)². While volume preservation theoretically fol-
As discussed in (Diepeveen, 2024), the quality of learned manifold mappings is determined by two key metries: the geodesic error and the variation error. The geodesic er- ror measures the average deviation form the ground truth
"We 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.
"We note that the choice for affine coupling layers rather than using more expressive diffcomorphisms such rational quadratic flows (Durkan et al., 2019) is due to our need for high regularity for stable manifold mappings (Remark 3.2) and an interpretable RAE (Remark 4.2), which has empirically shown to be more challenging to achieve for more expressive flows both first-and higher-order derivatives of şo will blow up the error terms in Theorem 4.1.
"We note that without these constraints (accommodating mul- timodality) the learned mappings can in principle be used to con- struct Riemannian geometry and RAF. However, from the theory discussed in this paper we cannot guarantee stability of manifold mappings nor that the RAE has the right dimension.
5

---

Score-based pullback Riemannian geometry
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 E.
low-to-moderate dimensional Eaclidean datasets and higher- dimensional image datasets:
Hemisphere(d', d) and Sinusoid(d', d) Synthetic Eu- clidean datasets with controllable intrinsic dimension d' and ambient dimension d.
Our approach introduces two key modifications to the nor- malizing flow (NF) training framework:
d' Y--uusssian Blobs Image Manifold: synthetic image dataset, also with controllable intrinsic dimension d', em- bedded in 1024-dimensional ambient space.
1. Anisotropic Base Distribution: We parameterize the diagonal elements of the covariance matrix A0 intro- ducing anisotropy into the base distribution.
2. l I--Isomerry Regularization: We regularize the flow ,pg to be approximately -iisometric.
MNIST: dataset of handwritten digits, originally em- bedded a 784-dimensional space, which we further embed into a 1024-dimensional ambient space using bicu- bic rescaling.
To assess the effectiveness of these modifications in learning more accurate and robust manifold mappings, we compare our method against three baselines:
The Hemisphere and Sinusoid datasets are used to evalu- ate RAE's performance on nw-too-moderate dimensional manifolds, while the Gaussian Blobs and MNIST datasets are used to test its scalability to highher-imensional and more complex image manifolds. For further details on each dataset, refer to Appendix F.2.
(1) Normalizing Flow (NF): Standard NF with an isotropic Gaussian base (O,,Ta)) no isometry regularization.
(2) Anisotropic NF: NF with a parameterized diagonal covariance base, isometry regularization.
(3) Isometric NF: NF with an isotropic Gaussian base, regularized be l -iismmetri..
6.2.1. 1. EUCLIDEAN DATASETS
1D and 2D manifolds. In Figures 1 and 5, we present the data manifold approximations by our Riemannian autoen- coder for four low-dimensional manifolds: Hemisphere(2.3), Sinusoid Sinusoid(2.3) and Sinusoid
We conduct experiments on three datasets, illustrated in 7 in F.1: the Single Banana Dataset, the Squeezed Single Banana Dataset, and the River Dataset. Detailed deseriptions of the construction and characteristics of these datasets are provided in F.I.
In Ap- pendix G, we detail the process used to create the data manifold approximations for these experiments. In our ex- periments, we set 0.01, which resulted in de = d' for cases, accurately capturing the intrinsic dimension of each manifold and producing accurate global charts.
Table presents the geodesic and variation errors for each method across the three datasets and Figure visually com- pares 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.
manifolds. To assess the scalability of our method, we conducted experiments on the Hemi- sphere(5,20) (522) and Sinusoid(5,20) datasets. The learned vari- ances effectively indicate the importance each latent di- mension, with high variances corresponding to the intrinsic manifold structure.
Higher-dimensinn
Introducing anisotropy in the base distribution without en- forcing isometry in the flow offers no significant improve- ment over the standard flow. On the other hand, regularizing the flow to be approximately isometric without incorpo- rating anisotropy in the base distribution results in under- fitting, leading to noticeably worse performance than the standard flow. Our results demonstrate that the combina- tion of anisotropy in the base distribution with isometry regularization (our method) yields the most accurate and sta- ble manifold mappings, as evidenced by consistently lower geodesic and variation errors.
On the Hemisphere(5 (5200) dataset, our model correctly iden- tified five non- vonnaniihing latent dimensions, achieving near- zero reconstruction error when selecting them. In contrast, choosing latent dimensions with vanishing variance resulted in no meaningful error reduction, confirming the model's ability separate important from redundant dimensions. A more detailed analysis of this effect is provided in Ap- pendix D.
For the more challenging Sinusoid(5,22) dataset, our method remains highly effective, though slightly less precise than for the Hemisphere dataset. The first six most important latent dimensions account for approximately 97% of the
6.2. Riemannian Autoencoder
To evaluate our method's ability to learn a Riemannian autoencoder (RAF), we conducted experiments on both
6

---

Score-based pullback Riemannian geometry
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. 16)
0.2561 (0. 1338)
0.3859 (0.2568)
Variation Error
0.0763 (0.0486)
0. 0.10 (0.0807)
0.1113 (0.0863)
0.0636 (0.0333)
Table 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.
(a) Gmund Trurh
(b) Our Method
(c) Standard NF
(d) Anisotropic NF
(c) Isometric NF
Figure Comparison of geodesics computed using different methods on the river dataset. The geodesics generated by the proposed method have the fewest artifacts, aligning with expectations from Table 1.
variance, increasing to over 99% with the seventh dimen- sion. The model slightly overestimates the intrinsic dimen- sionality, likely due to the increased optimization challenges of learning a more intricate distribution while regularizing for isometry.
experiments, introducing this term substantially improved the RAE's ability to correctly detect the intrinsic dimen- sion while maintaining smooth geodesic paths and accurate reconstractions.
Computing the Hessian-vectus product impractical in high dimensions. To improve scalability, we instead mini- mize the norm of single randomly chosen column of the Hessian-vector product (R4xd) per iteration. This signifi- cantly reduces computational and memory overhead while preserving effectiveness. See Appendix I.3 for details.
6.2.2. IMAGE DATASETS
In our preliminary experiments on image datasets, we ob- served significant overestimation of the intrinsic dimen- sion, despite the high quality of the geodesic interpolations. This issue stems from the use of non-affine flows," which increase the model's flexibility but also inflate Hessian- vector product term in the expansion the Hessian of the model log-density. This inflation can disrupt the RAE's ability assign high variances to the most important latent dimensions (see Appendix H for details).
d' Ga-Guusiian Blobs. We evaluated our RAE syn- thetic d' d'-Gaussian Blobs manifolds embedded in 1024- dimensional ambient space with intrinsic dimensions d' 20, 50, and 100. The RAE accurately captured the structure, assigning over 99% the total variance to 22. 51, and 101 latent dimensions respectively, closely aligning with the ground-truth intrinsic dimensions.
To address this issue, we repeated the image experiments after including the Hessian vector product term as an ad- ditional regularization term in the loss function. The addi- tional regularization term is given by:
Figure 3 shows the normalized cumultative variance and the I reconstruction error as functions of the number of most important latent dimensions. All curves saturate very close to the ground truth intrinsic dimension demonstrating RAE's ability to effectively capture the intrinsic structure of the data in high dimensional settings.
Absas Ex~epe
|| Dx AT, po, (o)ooo())) [x)
(23)
where Apess is the regularization weight, and D& Pe, denotes the Hessian of the flow. As demonstrated in the following
We evaluated the RAE on the MNIST dataset, a real-world benchmark known for its multimodal distribution (each digit class represents separate "mode"). While our
MNIST.
We used compositions of affine coupling layers and 1 x invertible convolutions, with the latter introducing non-affine formations to flow.

---

Score-based pullback Riemannian geometry
Normalized Cumultative Variance
7
7
7
3
3
3
3
1.0
2
2 L 2
4
4
4
4
/
1
/
7
7
7
7
0.4
d' GGaasssin Blobs
0
O
6
6
2
2
2 2
0.0
100
25
50
100
125
150
175
200
4
4
4
4
7
7
7
7
Reconstruction Error
0.008
Figure 4. RAE reconstructions on MNIST. The leftmost column shows original images, while the thrce show reconstructions for decreasing variance thresholds (0.1, 0,05, 0.01), using 176, 208, and latent dimensions, respectively. Reconstructions are clear 176 dimensions and nearly indistinguishable from the originals 271.
0.005
0.004
0.003
0.002
0.001
0.000
driven Riemannian geometry in high-dimensional settings.
25
50
75
100
125
150
175
200
Number of Dimensions
7. Conclusions
Figure 3. Normalized cumulative variance (top) and L reconstruc- tion error (bottom) for d' -Gaussian Blobs with intrinsic dimensions = 20, 50, and 100. Black dotted vertical lines mark the ground- truth intrinsic dimensions for each dataset. Curves shown for the first 200 dimensions to highlight the critical range, as the be- havior from 200 to 1024 dimensions remains effectively constant.
In this work we have taken a first step towards practical data-driven Riemannian geometry framework, striking bal- ance between scalability of training data-driven Rieman- nian structure and of evaluating its corresponding manifold mappings. We have considered a family of unimodal proba- bility densities whose negative log-likelihoode are compo- sitions of strongly convex functions and diffcomorphisms, and sought learn them. We have shown that once these unimodal densities are learned, the proposed score-based pullback geometry provides closed-form geodesies that pass through the data support and an interpretable Riemannian autoencoder with error bounds that estimates the intrinsic dimension of the data. Finally, to learn the distribution we have proposed an adaptation to normalizing flow training. Through numerical experiments on Euclidean and image datasets, we have shown that these modifications are crucial for extracting geometric information, and that our frame- work not only generates high-quality geodesics across the data support, but also accurately estimates the intrinsic di- mension of the approximate data manifold while construct- ing global chart, even in sigh-dimensional settings.
method theoretically best suited for unimodal distributions of the form in Equation (2), still performs remarkably well on MNIST, with only slight tendency to overestimate the 0.1, the RAE
intrinsie dimension. Specifically, assigns 90% of the variance to 176 latent dimensions, which slightly exceeds the maximum local intrinsic dimension (LID) of 152 estimated by ID-DIFF, a state-of-the-art LID method (Stanczuk et al., 2024).
We further examined two stricter variance thresholds, c 0.05 (208 dimensions) and 0.01 (271 dimensions). Fig- ure 4 shows that reconstructions are visually convincing at 176 dimensions and improve slightly with more dimensions, becoming nearly perfect at 271. We anticipate that incor- porating more expressive transformations, such as rational quadratic spline flows (Durkan et al., 2019), 9), could further refine the dimensionality estimation and improve reconstruc- tion quality by better handling the optimization challenges associated with learning the distribution under the isometry and Hessian constraints.
A Although the framework is theoretically best suited for uni- modal distributions, performs remarkably well on real multimodal distributions, albeit with slight overestima- tion of intrinsic dimensionality. This highlights its practical utility for extracting geometry from real data. However, extending the formulation to better handle multimodal dis- tributions remains an important direction for future work.
Overall, these experiments illustrate that our approach scales to real-worll data, provided we incorporate the Hessian vec- tor product regularization term when using non-affine flows. These results mark significant progress toward robust data-
This work paves the way for scalable learning of data geome-
8

---

Score-based pullback Riemannian geometry
try, unlocking applications that were previously out of reach. It enables efficient computation of manifold maps such geodesics and distances, with significant potential for ad- vancing deep metric learning. Furthernore, it introduces Riemannian auto-encoders with interpretable latent spaces that effectively capture the intrinsic structure of data. Look- ing ahead, it paves the way for Riemannian optimization directly on the data manifold by enabling the computation of intrinsic gradients, with the potential to revolutionize inverse problem-solving and push the boundaries of controllable generative modeling.
general geometries. arXiv preprint arXiv:2302.03660 2023.
Chow, Y. L.,, Singh, S., Carpenter, E., and Way, G. P. Predicting drug polypharmacology from cell morphol- ogy readouts using variational autoencoder latent space arithmetic. PLoS computational biology, 180))ee,09888, 2022.
Coifman, R. R. and Lafon, S. Diffusion maps. Applied and computational harmonic analysis, 21(1):5-30, 2006.
de Kruiff, F., Bekkers, E.,, Öktem, O., Schönlieh, C.-B., and Diepeveen, W. Pullback flow matching on data manifolds. arXiv preprint arXiv:2410 04543, 2024.
Acknowledgements
GB acknowledges support from the EPSRC Rosehips grant. ZS acknowledges support from the Cantab Capital Institute for the Mathematics of Information, Christs College and the Trinity Henry Barlow Schol- arship scheme. CBS acknowledges support from the Philip Leverhulme Prize, the Royal Society Wolf- son Fellowship, the EPSRC advanced career fellow- ship EP/VO29428/1, EPSRC Grants EP/S026045/1 and EP/T003553/1, EP/N014588/1, EP/T017961/1, the Well- come Innovator Awards 215733/7/19/2 and 221633/7/2007, the European Union Horizon 2020 research and innovation programme under the Marie Sklodowska- Curie Grant agree- ment No. 777826 NoMADS, the Cantab Capital Institute for the Mathematics of Information and the Alan Turing Institute.
DeMers, D. and Cottrell, G. Non-linear dimensionality
reduction. Advances in neural information processing systems, 5, 1992.
Diepeveen, W. Pulling back symmetric riemannian geome- try for data analysis. arXiv preprint arXiv:2403.066 2024.
Diepeveen, W., Chew, J., and Needell, D. Curvature cor- rected tangent space-based based approximation of manifold- valued data, arXiv preprint arXiv:2306.Nee 00507, 2023.
Dinh, L., Sohl-Dickstein, J., and Bengio, S. Density estimation using real NVP. In International Confer- ence on Learning Representations, 2017. URL https: /openreview net forum? weeerummiidHkppbnnn H=ppbnn99......
Impact Statement
Durkan, C., Bekasov, A., Murray, I., and Papamakarios, G. Neural spline flows. Advances in neural information processing systems, 32, 2019.
This paper presents work whose goal is to advance the field of Machine Learning and Riemannian Representation Learn- ing. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Fefferman, C., Mitter, S., and Narayanan, H. Testing the manifold hypothesis. Journal of the American Mathemat- ical Society, 29(4):983-1049 2016
Fletcher, P. T., Lu, C., Pizer, S. M., and Joshi, S. Principal geodesic analysis for the study of nonlinear statistics of shape. IEEE transactions on medical imaging, 23(8): 995- 1005, 2004.
References
Arvanitidis, G., Hansen, K., and Hauberg, A locally adaptive normal distribution. Advances in Neural Infor- mation Processing Systems, 29, 2016.
Gomari, D. P., Schweickart, A., Cerchietti, L., Paietta, E., Fernandez, H., Al-Amin, H., Suhre, K., and Krumsiek, J. Variational autoencoders learn transferrable representa- tions of metabolomics data. Communications Biology, (1):645, 2022.
Belkin, M. and Niyogi, Laplacian eigenmaps and spectral techniques for embedding and clustering. Advances in neural information processing systems, 14, 2001.
Boothby, W. M. An infroduction to differentiable manifolds and Riemannian geometry, Revised, volume 120. Gulf Professional Publishing, 2003.
Hauberg, S., Freifeld, O., and Black, M. A geometric take on metric learning. Advances in Neural Information Pro- cessing Systems, 25, 2012.
Carmo, M. P. d. Riemannian geometry. Birkhäuser, 1992. Chen, R. T. and Lipman, Riemannian flow matching on
Horvat, C. and Pfister, J.-P. Intrinsic dimensionality estima- tion using normalizing flows. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.).

---

Score-based pullback Riemannian geometry
Advances in Neural Information Processing Systems, volume 35, pp. 18887-18899, Curran Associates, Inc.,
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13556, 2020.
2022. URL https::/pproceed.ngg.. neurips. ce/paper cttppper_fillee//eeeeeeeeeeessee /ieaeer/files//ppeer//2////
watiBaaaacc3b220g94444bbbbbcllla Caperr-eeeeeeeeeeeeeeeeeee pdf.
Sorrenson, P., Behrend-Uriarte, D., Schnörr, C., and Köthe, U. Learning distances from data with normalizing flows and score matching, 2024. URL https: //arxivaxxivv org / abs 2407 09297.
Huang, C.-W., Aghajohari, M., Bose, J., Panangaden, P., and Courville, A. C. Riemannian diffusion models. Ad- vances in Neural Information Processing Systems, 35:
2750-2761, 2022.
Stanczuk, J. P., Batzolis, G., Deveney, T., and Schönlieb, C.-B. Diffusion models encode the intrinsic dimension of data manifolds. In Salakhutdinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp, F. (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Ma- chine Learning Research, pp. 46412-46440, PMLR, 21- 27 Jul 2024. URL https: /ttp://ccocccceeig/eeeoe mmlr. P press vt..//2255///// / http///stttanc222aaaatatt html.
Kadkhodaie, Z., Guth, F., Simoncelli, E. P., and Mallat, S. Generalization in diffusion models arises from geometry- adaptive harmonic representation. In Proceedings of the 12th International Conference Learning Representa- tions (ICLR), 2024. URL https: http://ppppeeeeeeeeeeeeeeeeee..p net/pdf? Seeriii==AVVVVYYET
Kapusniak, K., Potaptchik, P., Reu, T., Zhang, L., Tong, A., Bronstein, M., Bose, A. J., and Giovanni, F. D. Met- ric flow matching for smooth interpolations on the data abs/
Sun, X., Liao, D., MacDonald, K., Zhang, Y., Huguet, G., Wolf, G., Adelstein, I., Rudner, T. G. J., and Kr- ishnaswamy, S. Geometry-aware autoencoders for met- ric learning and generative modeling data mani- folds. In ICML 2024 Workshop Geometry-grounder Representation Learning and Generative Modeling, 2024. URL https: http:///peeeeeeeeeeeeeeeeeeeeeee net. forum? id=EYQZ jMen41.
manifold, 2024. URL 2405. 14780.
Kingma, D. and Welling, M. Auto-encoding variational bayes. arXiv preprint 1312v:1312.6/11 20133 2013.
Lee, J. M. Smooth manifolds. In Introduction to Smooth Manifolds, pp. 1-31. Springer, 2013.
Peltonen, J., Klami, A., and Kaski, S. Improved learning of riemannian metrics for exploratory analysis. Neural Networks, 17(8-9): 1087-1100, 2004.
Tempezyk, P., Michaluk, R., Garncarek, L.. Spurek, P., Ta- bor, J., and Goliñski, A. LIDL: Local intrinsic dimension estimation using approximate likelihood. In Proceed- ings of the 39th International Conference on Machine Learning (ICML), pp. 21205-21231, PMLR, 2022. URL https: https::/azxiv.orree abs///a////// http://www..g////////eeeeee..
Roweis, S. T. and Saul, L. K. Nonlinear dimensionality duction by locally linear embedding. science, 290(5500): 2323-2326, 2000.
Sakai, T. Riemannian geometry, volume 149. American
Tenenbaum, J. B., Silva, V. d., and Langford, J. C. A global geometric framework for nonlinear dimensionality reduction. science, 290(5500):2319-2320 2000.
Mathematical Soc., 1996.
Sakamoto, K., Tanabe, M., Akagawa, M.. Hayashi, Y., Sakamoto, R., Yaguchi, M., Suzuki, M., and Matsuo, Y. The geometry of diffusion models: Tubular neigh- bourhoods and singularities. In ICMI. 2024 Workshop on Representation Learning and Gener- ative Modeling, 2024. URL https: /ppeneeview. het/Forum?? www.:oiiunnonaaa.... E
Ternes, I.., Dane, M., Gross, S., Labrie, M., Mills, G., Gray, J., Heiser, L., and Chang, Y. H. A multi-encoder vari- ational autoencoder controls multiple transformational features single-cell image analysis. Communications biology, 5(1):255, 2022.
Vahdat, A. and Kautz, J. Nvae: A deep hierarchical vari- ational autoencoder. Advances in neural information processing systems, 33:19667-19679, 2020.
Sammon, J. W. A nonlinear mapping for data structure analysis. IEEE Transactions on computers, 100(5):401- 409, 1969.
Scarvelis, C. and Solomon, J. Riemannian metric learning via optimal transport. In The Eleventh In- ternational Conference on Learning Representations, 2023. URL https::/ setta:://peerreieee htwpiiewwnee//tttee Forum? -
Wenliang, L. K. and Moran, B. Score-based generative
models learn manifold-like structures with constrained mixing. In NeurIPS 2022 Workshop on Score-Based Methods, 2022. URL https: / /arxiv. org/ abs/ 2311.09952.
10

---

Score-based pullback Riemannian geometry
Zhong, E. D., Bepler, T., Berger, B., and Davis, J. H. Cryo- drgn: reconstruction of heterogeneous cryo-em structures using neural networks. Nature methods, 18(2):17-118) 2021.

---

Score-based pullback Riemannian geometry
A. Generalizations and proofs of Proposition 3.1 and Theorem 3.3
Proposition 3.1 is a special case of the result below.
Proposition A.1. Let up R + R be a smooth diffeomorphism and let bb : RR + R be a smooth strongly convex function, whose Fenchel conjugate is denoted by 1b* TRd + R. Next, consider the e E-pullback manifolds (R4, Vooo) and (RR (3.))) defined through metric tensor fields
(E, 2 = (Dx (1 and (2, = (Dxp
(24)
Then,
(i) the distance d Vdop, Fdd Rd + R on (R2 (oou) is given by
de Nooo (x,y) px,)) = |I( o (
(25)
In addition, if is of the form Equation (4)
do (x,y) = l|x(x) - -(y)l = (p(x) - p2y))ll..
(26)
(ii) length-minimising geodesics 7x.y
Vgoop :0,, 1] + RR on (R3 (,.)) voop) are given by
V.Ou Xx.y
(t)
= 1 (1 (x) + V
(27)
In addition, if 1/p is of the form Equation (4)
xx
Voo (t) = xx,y (t) = - -11-1111- ((1-))(xx) + tt(ly))
(28)
(iii) the exponential map expx V you
(-) TR.Rd + Rd on (R. (,, V.lop) is given by
expx
vop
(Ex) = (1 V (VV + Ve
(29)
In addition, if is of the form Equation (4)
expx x
Vijo (Ex) = exp (xx)(Exx) = = - + Dx
(30)
(iv) the logarithmic map logx v Yoop
(-) Rdd TxRa (R4, (.. Voop) is given by
logx
y
(Vib (x)]] (D V (V (Vrh
(31)
In addition, if 1 is of the form Equation (4)
logx
Voop y = logg = Dp(x)pp (1(x)) - - :x(x]]
(32)
Proof of Proposition 3.1. First note that Vib o p is a diffeomorphism with inverse x- o Vib", Then, Equations (25), (27), (29) and (31) follow directly from (Diepeveen, 2024, Prop. 2.1). Next, dp is of the form Equation (4), i.e..
1 A -1x,
(xx)
2
we have that its Fenchel conjugate is given by
3 (y) = ạ Ay.
(33)
So both Vx(()) = A A-'x and (y) = Ay are linear mappings, from which follows that they cancel to identity everywhere and yield Equations (26), (28), (30) and (32).
Similarly, Theorem 3.3 is a special case of the result below.
12

---

Score-based pullback Riemannian geometry
Theorem A.2. Let vp R + Rd be a smooth diffeomorphism and let db Rd + R be smooth strongly convex function, whose Fenchel conjugate is denoted by vi* Rdd R. Next, consider the function f = RRd -Rdxd given by
f(z) : DAVO' + Σ (,,, Vo
(34)
Finally, letx, x, E Rd be vectors and assume that for all vectors
{(1 - (VV +VVV te [0, c CR
the matrix f(z) is positive definite.
Then, mapping
t (Pp (I)), E [0,
(35)
strongly convex, where "xyy
Vaoo is the geodesic between x and y under the Riemannian structure (IR", (..)..))) Vvoo).
In addition, if to is of the form Equation (4) the mapping Equation (35) is strongly convex for any x.y Rd
Proof. By Equation (27) in Proposition 3.1 we have
b(p(1xy Veop (1))
V -
= -VV (1- () (VV
(36)
So the claim holds if on the linear subspace
{(1
Ie [0. c R
(37)
the function v Vel" is convex.
Next, note that the Hessian of Val* satisfies
D&V( V f(z).
(38)
By assumption f(z) is positive definite for all z in the subspace Equation (37). other words, on this subspace e (VVx)())) positive definite, which implies strong convexity and yields the main claim.
The claim for the special case of sb is of the form Equation (4) follows directly, because
f(z) = A,
(39)
which is always positive definite.
B. Proof of Theorem 4.1
Auxiliary lemma
Lemma B.1. Let yp Rd Rd be smooth diffcomorphism and let v/ Rd - P. be quadratic function of the form
Equation (4) with diagonal A E gdxd, Furthermore, let p Rd + R be the corresponding probability density of the form
Equation (2). Finally, consider E E [0, 1] and the mappings Eg Rdd - Rd and D Rdd + Rd in Equations (13) and (14)
with de E [d] as in Equation (12).
Then, for any a [0, 1) and any B E (0,1-a)
(X) x328xx0)) 1p(20 <
C2 Ca
8.p
1 3
Σ
(40)
3 1
i=1
13

---

Score-based pullback Riemannian geometry
where
B.,0 : s dett(((())))))) 50(x) TA A A'(0),
(41)
XER2
and
22 sup (Dt(p) Sr(x)) *p(x)3. xER=
(42)
Proof.
We need to distinet two cases: (i) dg = d and (ii) 1 S de <d
(i) dg = we have that (x(Fx(x)) = x for any x E Rdd In other words
Exmpld
x) e 2x)) A [(0) 0 S 0.0
C2 C3
B.p
1
1+ B
$/
Σ a.
(43)
3
1=1
(ii) Next, we consider the case 1 S d < d. First, notice that we can rewrite
(E)))) xxx]1 Equations (13) (14)
Σ (x(x), )2e - xxxx)]3 =1
Σ (p(x), 2e*
k=
orthogonaity
[1(p(x), )²" 13
Σ (x(x),, e²1 = Σ
x(x)
(44)
k=
k=
Moreover, define
C:= / e-2ppxx]T [ A- dx.
(45)
gdd
Then,
A *0000 Se (lo((D(FF -xx)) ((x (t--)))000))))000 (0xx)) A =x(x)
Jar -b((xx) 1x
Equation (45) 1
II+(D)
((x ( ()=x)) xxx)
-2-
C
Equation (44)
1
/
Σ
p(x) -12 ()=x))(>)) p(x)
1
Σ
p(x) e -1--))(()) - -1-1
0
k=d+11
k=d,
Rd
xxp=(()) 1 d
yic
2
- 21
y det(Dyy Dy DDy0~1)1ds l)dy
k=d++
JR.
Σ
(1 -- f
'y
dot Dy.ip -1)le e)-e- -g -
e
C
k=d++1
supyeRefl det Deip -) Eyy A ,
Σ
2 1----4y A 'ydy
yic C
C
k=d
Equation C3
11 7 2 Σ4 y3
y
(1 e
8 fy
by
Σ
y e
my dy
Bdd
Rdd
kmde+1
R de
C2
d
B.p
Nice e
2 -12
dy 1k
e
-13-4- -)
# dy -iyi-idyii+... dyd
C
k=d++1
Bd-1
2
d
B.p
aip
3 3
dyix
-t - E m -Wy--ddyi+1-. ... .yy
C
11 a - 3)
jB
Rd-1
k=d +1
14

---

Score-based pullback Riemannian geometry
C2 Dip
alc
e -2- 2- 21
A dy
- -3) RR
k=de +1
C2
d
aik
11a-8) )
(1 1+B
h/2 c - A yy
(1- a - -8) 1 11-1- a-
k= +1
y=
C3
Σ
d
(1 - a - 3) 11-1 a1-a-3) )))²
ain
1+B 4 { Rdd
-11+ A *pxx)
dct dct (xxp)) dxx
C2
B.P
Σ
aig
a a- 3) 11 1- 8)
(1
1 1 $
//c
-t A 1 det((x [p(x) A xxxx) dx
(1-a-
Rd
1md,+1
BBpp SuPxER dleel (Dxss)le 5x(x) A
( )p(x)
d
ai. (1 1
B] 1 1- 11-a-3) )² $ J/c. e c- dx dx
C
(1 -
-8)
R= +1
Equation (42)
B.p A.p
d
ais
1+ 3
2
| E - A Pxxx)
C
(1 - a - (3)
1 - 0 - B
Rd
Amdde+1
Equation (45)
23 Co Bie
1+ 3
Σ
ai
1-a -
a
k=
Equation (12)
D,o C3
1+ B
1
Σ
a;
(46)
a B
1
i=
Proof of the theorem
Proof of Theerrmm 4.1. First, consider the Taylor approximation
x (v(y)) - "v((() = Do (><x))=""l(())))>)>>>>>>>>>))>>) "[o(y) - (xx)] + 0(1vp())) - -xx))
= =Do(x010 "(x(y)'(>>>->>>_))>>)-)
- (
+ (x))) (47)
Moreover, we define
C= { (=0(x))==~(0())
(48)
Subsequently, notice that
=0 12
= -1 / ||
-(x))] 10(x)T A xx()
/ |1D - 11211 (x))) ((x))BBe=ttce -(( A --'x dx
< C
Rd
supxER (PPEER<1DDD(M)>> "1Be 2x(x) [x(x)]T-- A-100)) / (E,(((EE - xx())133 (3 -2- (2)(()))T A- =xx()
Equuation (17) S2epp
[ || (FF(()(FF(x)))) ( A -x(x)) Tx(x -+xx0x
C
(D. 2e8x |0(X]
Lemma B.I C C2
C3 8.p
1 3
Σ a (49)
2,3
[ 2,3/
15

---

Score-based pullback Riemannian geometry
Then,
Equation (47)
(D) (0 x)
= (xxx=/11pxx000000001111 - -
Old (D. X)
Equation (49)
C2pp
BiP
Co
8
$
Σ
a; + o (50)
2/3
1
2/3
i=1
which yields the claim 3 was arbitrary.
C. Training Details
This section describes the configuration parameters necessary to reproduce our experiments. The experiments are categorized into two groups: Euclidean datasets (Sinusoid and Hemisphere) and image datasets (Gaussian Blobs and MNIST). All experiments share some common parameters, listed below, while dataset-specifie parameters are provided Table 2.
Common Parameters:
Optimizer: Adam with betas = (0.9, 0.99), eps = 10-* and weight decay x 10-5,
Gradient Clipping: Gradient norm clipped to 1.0.
Model Architecture:
Euclidean Datasets:
normalizing flow composed of affine coupling layers, where each layer transforms part of the input while leaving the remaining dimensions unchanged. The parameters of the affine transformations parameterized by a ResNet with 64 hidden features and residual blocks. Transformations alternate across dimensions at each step.
Image Datasets: A normalizing flow composed of three levels, each beginning with squeeze operation that redis- tributes spatial dimensions into channels. Within each level, apply stack of seven flow steps, leading to total of 21 flow steps across all levels. Each flow step consists of:
ActNorm layer.
x 1 invertible convolution for channel mixing.
* An affine coupling transform, where the parameters of the affine transformation are parameterized by a convolutional ResNet. This ResNet has 96 hidden channels and residual blocks.
Table 2. Training configurations for each experiment.
Dataset
Flow Steps
Epochs
Batch Size
Aiso
Avol
Ahess
Learning Rate
Sinusoid(1.3
1000
64
1.0
1.0
3 10 -0
(in)soid(2,3)
8
1000
64
1.0
1.0
3 10
Sinusoid(5.20)
24
2000
128
1.2
2.5
4 x 10-4
Hemisphere(2.33
2000
64
1.0
1.0
4 x 10-4
Hemisphere(5.20)
12
2000
64
0.75
1.2
4 x 10 -4
20-Blobs
500
64
100
1.0
0.5
2 x 10 -4
50-Blobs
500
64
100
1.0
0.5
10-4
100-Blobs
500
64
100
1.0
0.5
2 x 10-4
MNIST
21
500
64
100
2.0
1.0
2 x 10 10-4
16

---

Score-based pullback Riemannian geometry
D. Additional Experimental Results
(a) Dims (5. 59, 92)
(b) Dims (31, 55, 66)
(c) Dims (64, 72, 90)
Figure 5. Approximate data manifold learned by the Riemannian autoencoder for the Sinusoid(1, 100) dataset. The orange depict the learned manifold, while blue points represent training data.
Hemisphere (5,20)
0.16
0.030
0.14
Average Reconstruction Error
0.025
Diagonal Value
0.020
0.10
0.08
0.06
0.010
0.04
0.005
0.00
0.000
Index
Number of Dimensions Used
(a) Learned variances in decreasing order.
(b) Reconstruction error for three latent orders.
Sinusoid (5,20)
Average Reconstructio Error
1.0
0.8
Diagonal Value
0.6
0.2
0.0
Index
Number of Dimensions Used
(c) Learned variances in decreasing order.
(d) Reconstruction emar for three latent orders.
Figure 6. Learned variances and reconstruction emror for Hemisphere(5,22 and Sinusoid(5, 2,)). Left: variances in decreasing order. Right: average e reconstruction error vs. latent dimensions. Errors shown for three variance- based orders: blue (decreasing variance), green (increasing variance), and red (random).
17

---

Score-based pullback Riemannian geometry
E. Error Metrics for Evaluation of Pullback Geometries
The geodesic error measures the difference between geodesics on the learned and ground truth pullback manifolds. Given two points Xo. x Rd let Xx,,x, (t) and Xx0. (t) denote the geodesics induced by the learned map Po, and the ground truth map PGT, respectively, where E [0, 1].
Geodesic Error.
The geodesic error is calculated as the mean Euclidean distance between the learned and ground truth geodesics over N pairs of points:
N
Geodesic Error
1
990 2 x
(4x)
(4x)
PGT 6
x (1x)
N
k=
where 7' is the number of time steps used to discretize the geodesic, and b = - for k = I.....
This metric captures the average discrepancy between the learned and ground truth geodesics, reflecting the accuracy the learned pullback manifold.
The variation error quantifies the sensitivity of the geodesic computation under small perturbations of the endpoints. For points Xo, X1 E R.d let X1 Ax, where Ax is random variable sampled from the Gaussian distribution:
Variation Error.
Ax ~ N(0,0.131)
with mean 0 and covariance 0., where I is the identity matrix. Define Txo (1) and (() as the geodesics from to X1 and z, respectively, induced by the learned map (e.,
The variation error is calculated as the mean Euclidean distance between the geodesic from Xy to X1 and the perturbed geodesic from Xg to z:
Variation Error
200 10 (1x)
(1x) - 90 2 (Lx)
N
where N is the number of sampled point pairs, T' is the number of time steps used discretize the geodesic, and ty. -=1
for k = 1, I.
This metric evaluates the robustness of the learned geodesic against small perturbations, providing insight into the stability of the learned manifold.
F. 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.
F.I. Datasets for Manifold Mapping Experiments
In our manifold mapping experiments (Section 6.1), we use the following datasets illustrated in Figure 7:
Single Banana Dataset: A two-dimensional dataset shaped like a curved banana.
Squeezed Single Banana Dataset: variant of the Single Banana with tighter bend.
River Dataset: A more complex 2D dataset resembling the meandering path of a river.
Each dataset constructed by defining specific diffeomorphisms & and convex quadratic functions v), then sampling from the resulting probability density using Langevin Monte Carlo Markov Chain (MCMC) with Metropolis Hastins correction. The probability density function is defined as:
18

---

Score-based pullback Riemannian geometry
(a) Single Banana
(b) Squeezed Single Banana
(c) River
Figure 7. Visualization of the datasets used in manifold mapping experiments.
pxx) x e -((((())
(51)
where the strongly convex function & given by:
(vv 1 A 2
(52)
and A is positive-definite diagonal matrix. The specific choices of and A for each dataset determine its geometric properties.
F.1.1. DIFFEOMORPHISMS AND CONVEX QUADRATIC FUNCTIONS
The key differences between the datasets arise from the diffeomorphism &P and the covariance matrix used in the sampling process. Below, describe the specific settings for each dataset.
1. Single Banana Dataset
Diffeomorphism:
21 az3
(0(x)
22
where and 2 = O.
Covariance matrix:
A
Squeezed Single Banana Dataset
Diffeomorphism: Same as the Single Banana Dataset.
Covariance matrix:
0 4
A
0
3. River Dataset
Dificomorphism:
11 - sin ((r2))))) 2
x(x)
F2
where and
19

---

Score-based pullback Riemannian geometry
Covariance matrix:
0
25
3
F.1.2. DATASET GENERATION ALGORITHM
Algorithm outlines the dataset generation process for all three datasets. The specific diffcomorphisms and quadratic
functions differ for each dataset.
Algorithm General Dataset Generation Algorithm
Require: Number of samples N, MCMC steps T, Step size 6, Diffeomorphism yo Covariance matrix A Ensure: Dataset {x,,X2,-.. : XN}
1: Initialize: Set initial state X = 0 E R2.
2: for = to N do
3: x x Xg
4: for k to T do
5:
Compute the score function Vx log Purger (x).
6: Propose x' = x + 2 Vx log Phageet (x) + Ön,, where ~ (O,
7:
Compute the forward kernel:
x
x
6² Vxx
log Phuree (x )
K Kfowwadd
28
8: Compute the reverse kerne::
6² V x log Purget (x)
Kreverse
20
9:
Compute the Netropolis-Hastinne acceptance probability:
A = min 1,, Prarget (x')
exp (- Koowaard + Kreverse)
Pharget (x)
10:
Accept x' with probability A; else set x' =x
11: Update x'.
12: end for
13: Store the final x as sample x,.
14: end for
F2. Datasets for Riemannian Autoencoder Experiments
In the Riemannian autoencoder experiments (Section 6.2), we use the following datasets:
Samples drawn from the upper hemisphere of d'-diimesiional unit sphere and embedded into Ed via random isometric mapping.
Hemisphere(d', d) Dataset:
Sinusoid(d', d) Dataset: Generated by applying sinusoidal transformations to d'-dimensional latent variables, resulting in a complex, nonlinear manifold in d
F.3. Hemisphere(d', d) Dataset
The Hemisphere(d", d) dataset consists of samples drawn from the upper hemisphere of d' dimeensional unit sphere, which then embedded into d-dimensional ambient space using a random isometric embedding. Below are the steps involved in constructing this dataset.
1. Sampling from the Upper Hemisphere unit sphere Sd C +11 The upper hemisphere is defined as:
We begin by sampling points from the upper hemisphere of the d' di-dimensional
5 = fx 2d ||x| = 1, 2 o)
20

---

Score-based pullback Riemannian geometry
The first angular coordinate is sampled from Beta distribution with shape parameters and B 5, scaled the interval [0, 5]- This sampling method emphasizes points the "equator" of the hemisphere. The remaining angular ,....,,0' Orr are sampled uniformly from the interval [0, m]:
coordinates ......,
01 ~ (5)
1T
0, ~ Uniform (0, r), for 2, 2,...,d
2. Conversion to Cartesian Coordinates Next, each sampled point in spherical coordinates is converted into Cartesian coordinates in gd'+1 using the following transformation equations:
31 cos(01), 22 sin(O1) cos(
38'+1 = =si((1)) sin( s.n
This conversion ensures that the sampled points lie on the surface of the unit sphere in (d' + dimeesssio space.
3. Random Isometric Embedding into Rd After sampling points the hemisphere in Rd'+I, the points embedded
into d-dimennsiona ambient space (d 2 d' + 1) using random isometric embedding. The embedding process is follows:
1. Generate a random matrix A E [dd(d'+1), where each entry is sampled from a standard normal distribution (0,,) 1).
2. Perform QR decomposition matrix A obtain Q ddd(''+1);
A = QR.
The columns of form an orthonormal basis for a (d' + 1)ddimensional subspace of IRd ensuring that Q defines an isometric embedding from Rd into Edd This guarantees that distances and angles preserved during the mapping, maintaining the geometric structure of the original space within the higher-dimensional ambient space.
3. Use matrix Q map each sample x dd'+1 into the ambient space:
y = =Qx,
where y Rdd are the embedded samples.
21

---

Score-based pullback Riemannian geometry
Algorithm 2 Hemisphere(d', d) Dataset Generation
1: Input: Intrinsic dimension d', ambient dimension d, number of samples n, Beta distribution parameters a = 5, B =
2: Output: Dataset Y TR
3: Step 1: Generate Random Isometric Embedding
4: Generate a random matrix A E Pd +1) with entries from ((0, 1)
5: Perform QR decomposition on A to obtain Q E Id)((d+1);
A QR
6: Step 2: Construct Dataset
7: for 1 to do
8: Step 2.1: Sample Spherical Coordinates
9: Sample the first angular coordinate 0, from scaled Beta distribution:
0 Beta(o, .3) (2
10: Sample the remaining angular coordinates 02,...,8u Aar from a uniform distribution:
0. Uniform(O, T), for i = .,....
11:
Step 2.2: Convert to Cartesian Coordinates
12:
Convert the spherical coordinates to Cartesian coordinates x; E Rdd using:
c1s
sin cos(O2)..... Id'+1 = sin(0;) sin(O2) ... sini)))
13:
Step 2.3: Embed Sample x; into Ambient Space
14:
Map the sample x; to the ambient space using:
y1 = Qx:
Append y; to the dataset Y
16:
end for
17:
Return: The final dataset Y = [y, Y2. .2.... Vnl
E44 (in d) Dataset
The Sinusoid(d', d) dataset represents a d'ddimensional manifold embedded d-dimensional space through nonlinear sinusoidal transformations. Below are the detailed steps involved in constructing this dataset.
1. Sampling Latent Variables The latent variables z E Rdd sampled from multivariate Gaussian distribution with zero mean and isotropic variance, as follows:
N (0,
where oi controls the variance along each intrinsic dimension, and La is the d' d' identity matrix. The value of a is set to for our experiments.
2. Defining Ambient Coordinates with Sinusoidal Transformations For each of the d - d' ambient dimensions, we construct a shear vector aj E IRd with its elements drawn uniformly from the interval [1, 1,2]:
a ~ Uniform(1,2) 2)d", for j I,.... d'.
The shear vectors aj apply fixed linear transformation to the latent space z E Rd determining how the latent variables influence each ambient dimension. These vectors, sampled once for each of the d - d' ambient dimensions, modulate the scale and periodicity of the sinusoidal transformation.
Each ambient coordinate z, is generated as a sinusoidal function of the inner product between a, and z, with a small Gaussian noise added for regularization.
22

---

Score-based pullback Riemannian geometry
Jj sin (a) z) cjz
where E, ~ NO, o is Gaussian noise with variance o. In our experiments, we set o² = 10-3.
3. Constructing the Dataset Samples The final samples y E R are formed by concatenating the ambient coordinates
Di, ,2......, 2d with the latent variables 21, 22..... Adt
= 12, Ed-d's 21, R2, ,,,,,,..
Algorithm 3 Sinusoid(d', d) Dataset Generation
Input: Intrinsic dimension d', ambient dimension d, number of samples n,, variance o = 3, noise variance o 10-3
2: Output: Dataset Y Rxxd
3: Step 1: Generate Shear Vectors
4: for 1 to d d' do
5: Sample shear vector aj E RRd from Uniform(1, (1, 2)d
6: end for
7: Step 2: Construct Dataset
8: for to n do
9: Step 2.1: Sample Latent Variables
Generate latent variables z; IRd from multivariate Gaussian:
2, NO0, o Ix')
Step 2.2: Compute Ambient Coordinates for Sample for - d' do
13:
Compute ambient coordinate z, for the i-th sample:
Jz sin (a. 2:) + + oj ~~~ ~(0,00))
end for
15:
Step 2.3: Form Final Sample y;
16:
Concatenate the ambient coordinates E IRd:
= [11, 22, ... Wdd-d dr) and the latent variables to form the final sample
y: [21, 22,, Td .2..... Zad
17:
Append y, to the dataset Y
end for Return: The final dataset Y y Y2, yn.
F.5. Gaussian Blobs Image Manifold
The Gaussian Blobs Image Manifold dataset defines image manifold with controllable intrinsic dimension d' (number of Gaussian blobs) embedded in a dizeeinma ambient space (32 x 32 pixel images). The dataset is constructed as follows:
1. d' Gaussian blob centers are randomly selected from a 32 x 32 grid without replacement. These centers are fixed and remain the same for all points in the dataset.
2. For each image, Gaussian mixture is generated by centering Gaussian distributions at the fixed locations, with standard deviations sampled uniformly from specified range [Smin 5max]-
3. The normalized density of the Gaussian mixture defines the intensity values of the image. Fach combination of standard deviations uniquely defines point on the manifold.
23

---

Score-based pullback Riemannian geometry
Algorithm 4 Gaussian Blobs Image Manifold Dataset Generation
Require: Number of samples N, intrinsic dimension d', image size 1, standard deviation range (Smin: Smax random seed S
Ensure: Dataset {x1, X............. xN}
with x; R!x!
1: Initialize: Set random seed S for reproducibility.
2: Generate d' Gaussian centers by sampling without replacement from the / x I pixel grid. These centers are fixed for all samples in the dataset.
3: for to N do
4: Initialize an empty image x; of size I I.
5: for = 1 d' do
6: Select the j-th Gaussian center (zz. My) from the preselected centers.
7: Sample standard deviation 8j uniformly from [8min; 8max].
8:
Compute the Gaussian density over the entire image grid:
(- (x- - xz)² (-(o-
(y - 81)
g v) = V2T8j
exp
283
9: Add the Gaussian density to the image: xi v) += g(r,3).
10: end for
Normalize x; the range [0, 1]:
min(x;)
max(x;) mmn(x;))
12: Store the normalized image x, as a sample.
13:
end for
G. Data Manifold Approximation
The learned manifold, shown in orange in Figure 1, the set (2)),, where D, the RAE decoder Equation (14), the set H in the latent space the open set given by
d_ u = (-33/au 3
and auy aus, are the de 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 3aa,, to +3 a... for = 1,...,d, de, where de is the number of significant latent dimensions. Practically, the off-manifold latent dimensions (those corresponding negligible variances) are set to zero. The decoder D, then maps this grid from U back to Rd. generating an approximation of the data manifold, as illustrated in Figure 1.
H. Explanation for the Need of Higher Regularity in Non-Affine Flows
As described in Section 6.2.2, introducing a small regularization term involving the Hessian vector product of the flow can substantially improve the RAE's ability to assign higher variances to the on-manifold (semantically meaningful) directions and suppress variances in off- manifold directions for non-affine flows. The core issue arises from the fact that, for non-affine normalizing flows, the second-order derivatives of fa, appear in the Hessian of the modeled log-density in way that can adversely impact the performance of the RAE if left unregularized, as we shall see below.
H.1. Hessian of the Log-Density and Intrinsic Dimensionality
Recently, Stanczuk et al. (2024) formally proved-and and Stanczuk et al. (2024); Kadkhodaie et al. (2024); Wenliang & Moran (2022) independently demonstrated empirically-that when data concentrate locally around a lower-dimensional manifold of dimension k. the Jacobian of the score function or equivalently, the Hessian of the log-density has & vanishing singular values. The eigenvectors corresponding to these vanishing singular values span the local tangent space of the data manifold. Consequently, for trained model, large singular values of Volog p0, (x) indicate off-manifold directions, while near-zero
24

---

Score-based pullback Riemannian geometry
singular values reveal on-manifold directions- spanning the locally flat subspace in which the data reside.
H.2. Analyzing the Hessian of the Log-Density of Our Model
With the background knowledge on the connection between the Hessian of the log-density and the local intrinsic structure in mind, let's investigate the Hessian of the logdensity of our model under the assumption that yog is local isometry and hence volume preserving well. The modeled density function is
po,,6 (xx po, (v0,(())) |det( (ecpo)
(53)
where D2PPoz (x) is the Jacobian of Pe,, at x and pel (x)) O, Ao,). Since :oo, is volume preserving, the absolute value of the Jacobian determinant 1, simplifying the density to:
pe, po,,0, ,x)
po. (40,, (x)).
(54)
Taking logarithms and differentiating yields the expression for the Hessian of the logdensity:
log po, = - 12000 A (x) (Dx40o, (x)) T A (Dxto, (x))
(55)
where Ve (x) is the Hessian of pe.. We see that the Hessian of the log-density consists of the sum of two terms:
A Hessian vector product term:
(x)0((x) Ax) (0, (x),
A term that depends on the Jacobian and the learned covariance matrix Ae,
(x)) A (xxi)e((x))
When 4e,, is affine (e.g., composed of affine coupling layers only), the Hessian vector product term in equation 55 vanishes.
Hence, the Hessian the logdensity simplifies to:
V loggpo pal (x) -(xxse, A (x)),
(56)
Noting that the Jacobian of yog,, denoted as Dxip,, (x), is orthogonal matrix due to yo, being local isometry, and considering that Ag, is a diagonal matrix with strictly positive entries, it becomes clear that the right-hand side of Equation equation 56 represents the eigendecomposition of the Hessian of the log-density. The magnitude of the eigenvalues of the Hessian directly influences how variance is allocated in the latent space. Large eigenvalues correspond to off off-manifolds directions and therefore will encoded by latent dimensions small learned variance, while small eigenvalues correspond to on-manifold directions and will be encoded by latent dimensions of high learned variance. This analysis provides a clear and intuitive explanation of why the Riemannian Autoencoder effectively detects and encodes important semantic information in the latent dimensions associated with high learned variances when trained with an affine normalizing flow regularized for local isometry.
However, for non-affine flows-such as those incorporating x 1 invertible convolutions after each affine coupling layer or rational quadratic splines-tie the Hessian vector product term can become significant and may interfere with the Jacobian term. This disruption can distort the learned manifold geometry, leading to an incorrect allocation variances the latent space. Specifically, the model may assign large variances to an increased number of latent dimensions and thus overestimate the intrinsic dimension. To address this, we add Hessian vector product penalty to the loss, minimizing its contribution to the Hessian of the log-density and allowing the Riemannian Autoencoder to accurately capture the data manifold's lower-dimensional structure in Ao,-
Note that Pa, maps Rd + Rd, making its Hessian V2 (x) a rank-3 tensor of shape d x d x d. Multiplying this Hessian by A Gog (x) results in a dl x d matrix, whose computation becomes prohibitively expensive in high dimensions, even with optimized methods. To address this in our MNIST experiments, we randomly selected dimension each iteration,
25

---

Score-based pullback Riemannian geometry
computed the corresponding column of the Hessian vector product, and penalized its Euclidean norm (refer to Appendix I for additional details). This efficient regularization enough keep the Hessian term orders of magnitude smaller in Frobenius norm compared to the Jacobian term, allowing the Riemannian Autoencoder to perform effectively with non-affine flows.
I. Computational Complexity of the Regularization Terms
In this section, we discuss the computational complexity and memory requirements for each of the additional regularization terms introduced in objective. For simplicity, denote the normalizing flow yo, (x) as ((x:) throughout this section:
L0eg(01, 02) = Avol En- (log det((D)))
Volume Regularization
+ Aiso Egropen [ D2O 1all
Isometry Regularization
+ Abess Egcopan D 2p 21))122
Hessian Regularization
The Hessian regularization term is only applied when is a non-affine flow. For affine flows, the volume and isometry regularization terms are sufficient, the second derivatives vanish in this case, making the Hessian penalty unnecessary.
The presented regularization objective is computationally feasible only for low-dimensional data, evaluating the full Jacobian and the -essian-vector product is expensive both computation and memory for high dimensions. To address this, we present efficient approximations designed to significantly reduce computational and memory costs while preserving the effectiveness of the regularization. The details of these approximations, including their implementation and impact on scalability, are provided in this section.
L.1. Volume Regularization
The volume regularization term
(log] (D=x))
relies on the log-determinant of the Jacobian, which typically already computed during the forward pass in normalizing flows. Most flow architectures are designed so that log det(Da)) is tractable and inexpensive obtain (e.g., via coupling transforms autoregressive transforms).
Hence, extra gradient backprop additional memory overhead is needed for this volume penalty-it it essentially comes for free from the standard normalizing flow likelihood computation.
1.2. Isometry Regularization
The isometry regularization term,
[I(()>) T Dæ Lol
Ex--pa
penalizes deviations of the Jacobian DuØ from orthogonality. By enforcing this condition, the regularization ensures that the mapping & approximates local isometry, preserving distances and angles in the vicinity of each point.
Full Objective. The full implementation involves computing the d x d Jacobian matrix for each sample, which is both computationally and memory-intensive in high-dimensional settings. These constraints render the full objective impractical for high-dimensional data, limiting its application to low-dimensional scenarios where such costs remain manageable.
Sliced Objective. To address the scalability limitations of the full objective, we introduce the sliced objective. Instead of computing the full Jacobian, we approximate it using Jacobian-vector products (JVPs) with small number m of randomly
26

---

Score-based pullback Riemannian geometry
sampled orthonormal vectors. This approach significantly reduces computational and memory requirements while retaining the effectiveness of the regularization. The procedure is outlined in Algorithm 5.
Algorithm 5 Sliced Isometry Regularization Objective
1: Input: Mapping Φ, batch of inputs z E Bxd number of orthonormal vectors m, and device.
2: Output: Approximate orthogonality regularization loss Liso-
3: Generate a random matrix R E IR Bxdm with entries sampled from (N 1).
4: Perform QR decomposition on R. for each batch sample to obtain orthonormal vectors Q E IR Bxdxm
5: Compute Jacobian-vectt products (JVPs) for all m orthonormal vectors E Q:
Jo, Vo(x)v for = I....,m.. m.
6: Construct the Gram matrix:
- (Jv) Vi,j E mm..
7: Compute the deviation of G from the identity matrix:
G
where I. is the identity matrix of size x 8: Compute the Frobenius norm of the deviation for each sample:
||G I 11
9: Compute the batch mean of the Frobenius norms:
Ciso
G Im 1
In practice, we employ PyTorch's efficient lacotian-vector vector product implementation to compute Ju, efficiently. To further enhance performance, we leverage torch. vmap to vectorize computations over the batch dimension, enabling parallelized and streamlined evaluation of the regularization term across all samples in the batch. These optimizations substantially improve scalability, making the sliced objective well-suited for larger batch sizes and high nigh-dimensional data.
Notably, our experiments demonstrate that using only m = 2 slicing orthonormal vectors is sufficient for efficient and effective isometry regularization. This approach effectively reduces the computational cost of the regularization term to 2 JVPs per sample, compared to d full backpropagation passes required for the full objective. Consequently, the sliced objective enables the method to scale effectively to dige-dimensional datasets.
I.3. -essian-Vector Product (HVP) Regularization
The Hessian-vecti vector product regularization Egaopan Dio - (0)))(())) limits the influence of second-order derivatives on the Hessian of the log-density of the model distribution, enhancing the Riemannian Auto-encoder's ability to capture the intrinsic geometry of the data manifold with expressive non-affine flows. It is unnecessary for affine flows, where these terms naturally vanish. See Appendix H for details.
Full Objective. For each output dimension j of o, the Hessian D is a d x d matrix, and the full Hessian across all outputs forms a rank-3 tensor of shape Rdxaxd, As result, even with optimized Hessian-vector product implementations, computing this term becomes infeasible for high-dimensional data.
Sliced Objective. To reduce computational overhead, we approximate the Hessian penalty by sampling single dimension j E (1,..., d} at each training iteration. Instead minimizing the Frobenius norm of the full Hessian-vector product matrix, we compute and minimize the norm a single column (of size d x 1) corresponding to the sampled dimension j. This approach lightweight and empirically sufficient for effective regularization.
27

---

Score-based pullback Riemannian geometry
Algorithm 6 Sliced Hessian Regularization Objective
1: Input: Mapping Ø, batch of inputs E RBxx inverse covariance Σgg
2: Output: Approximate Hessian regularization loss Lhess-
3: Randomly sample single output dimension j {I,...,d) d).
4: Compute the Hessian-vector product (HVP) for each batch sample:
HVP; (x) (Ez, ((x))
using torch. autag functional hvp or equivalent.
5: Compute the norm of the resulting d vector for each batch sample:
HVP,
6: Compute the batch mean of the norms:
1
Lhess
HVP, H
B
In our experiments, evaluating a single random column the Hessian-vector product matrix per iteration and penalizing its norm was sufficient to maintain low Frobenius norm of the full d x d matrix, effectively regularizing the flow. The computation of the sliced Hessian-vector product in optimized deep learning libraries is equivalent two backpropagation passes: one for the gradient and one for the vector-Jacobian product. Therefore, the sliced objective enables the method to scale effectively to complex, high-dimensional datasets.
28


Reviews:
[{"review_id": "BpyxLifQFX", "reviewer": "Reviewer_pCEZ", "summary": "This paper proposes to construct a Riemannian structure from unimodal probability densities. Under a specific condition, the constructed pullback Riemannian structure turns out to be related to that obtained from the score function (i.e., the gradient of the log probability density with respect to data). The paper also generalizes the idea of classical PCA to the Riemannian setting, enabling the construction of Riemannian autoencoder (with error bounds on the expected reconstruction error) from unimodal probability densities. Finally, the authors show how train the density in question by adapting normalizing flow to their framework. Experimental results demonstrate the potential of the proposed method in various applications.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive constructive feedback! See below for a discussion on the weaknesses:\n\n- **“The paper only considers a simple setting based on unimodal probability densities. This would limit the capability of the proposed method to construct complex geometries.”**\n\n  - This is definitely a limitation of the current method as we also mention in the paper. Having that said, for subsequent work it is important to understand the base case, which is why we focus in this paper on unimodal distributions as a first step.\n\n- **“The Riemannian geometry induced by unimodal probability densities in Section 3 has close connections with several well-established Riemannian geometries such as Log-Euclidean [A] and Log-Cholesky [B]. It would be helpful to have a discussion on those connections. Currently, such a discussion is missing in the paper.”**\n\n  - We feel that these are somewhat different things. The Log-Euclidean and Log-Cholesky metrics are pullback metrics on the space of positive definite matrices rather than on $\\mathbb{R}^d$ and are not data-driven (which is the case we consider). So apart from these also being pullback metrics (and there exist many other pullback metrics on many other spaces), we don’t see a close connection as you would really do different things with both Riemannian structures. Alternatively, is the reviewer thinking about these metrics as metrics between centered normal distributions (information geometry)? In that case this is still not really related to what we are doing as we are considering metrics between elements in a distribution rather than metrics between distributions. So overall, would the reviewer maybe be more specific what this close connection is apart from the fact that they both use pullback Riemannian geometry?", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "I read the proofs to understand the general ideas but could not follow all the details", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "The paper bases on works in geometric data analysis and generative modeling. It seems to be the first work that constructs the complete geometry of the data manifold", "essential_references_not_discussed": "The main reference is Diepeveen, (2024) which is properly cited in the paper", "other_strengths_and_weaknesses": "Strengs:\n- This is an interesting paper that bridges works in geometric data analysis and generative modeling. \n- The paper is nicely presented. The ideas are expressed in a simple and concise way which makes the paper easy to follow. I enjoy reading it. \n- Experimental evaluation show the potential of the proposed method\n\nWeaknesses:\n- The paper only considers a simple setting based on unimodal probability densities. This would limit the capability of the proposed method to construct complex geometries.\n- The Riemannian geometry induced by unimodal probability densities in Section 3 has close connections with several well-established Riemannian geometries such as Log-Euclidean [A] and Log-Cholesky [B]. It would be helpful to have a discussion on those connections. Currently, such a discussion is missing in the paper.\n\n**References**\n\n[A] V. Arsigny, P. Fillard, X. Pennec, and N. Ayache. Fast and Simple Computations on Tensors with Log-Euclidean Metrics. Technical Report RR-5584, INRIA, 2005.\n\n[B] Lin, Z.: Riemannian Geometry of Symmetric Positive Definite Matrices via Cholesky Decomposition. SIAM Journal on Matrix Analysis and Applications 40(4), 1353–1370 (2019).", "other_comments_or_suggestions": "Please see the weaknesses", "questions_for_authors": "Please see the weaknesses", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TGP3SuJILt", "reviewer": "Reviewer_X13X", "summary": "The paper proposes a novel framework for learning the intrinsic geometry of data manifolds using pullback Riemannian metrics induced by an anisotropic normalizing flow. The key idea is to model the data manifold with a Riemannian autoencoder (RAE), where the encoder function provides a pullback metric through a strongly convex potential function. However, the clarity of presentation and theories should be heavily improved.\n\n## update after rebuttal\n\nI thank the authors for the rebuttal. I have re-read the response carefully and admit there are some misunderstandings in my initial review. I therefore raise my score to 3 accordingly. That said, I still have several concerns and hope the final version can be revised accordingly:\n\n- Presentation: The author should highlight that they implicitly learn the manifold via explicitly learning $R ^d$.\n- Theory: It could be argued that learning via ambient space might be restricted, as a more straightforward way is to learn the manifold directly. This should be discussed in the paper.\n- Experiments: When talking about Riemannian metric learning, it should eventually recover some true Riemannian geometry, such as Frechet mean, Riemannian logarithm, etc. This is the intention of my previous comments and references, and is missing in the current version.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their engagement. Below we address major concerns while clarifying our core contribution: A novel Riemannian geometry framework with closed-form geodesics that provably traverse high-density regions, overcoming limitations of prior data-driven approaches. \n\n- **Clarifications & Definitions**  \n  - **Notation consistency:**  \n    - Added explicit clarifications when citing (Diepeveen, 2024) and introduced footnotes for operator precedence (e.g., $\\nabla \\psi \\circ \\varphi \\equiv (\\nabla \\psi) \\circ \\varphi$)  \n  - **Equation motivations:**  \n    - Eq (2): Exponential family chosen for (1) score function compatibility and (2) theoretical/algorithmic tractability with unimodal distributions  \n    - Eq (4): Inverse SPD matrix preserves covariance interpretation and enables cleaner nonlinear PCA generalization  \n  - **Manifold scope:** Our metric operates on all $\\mathbb{R}^d$ rather than explicit submanifolds $\\mathcal{M} \\subset \\mathbb{R}^d$  \n\n- **Theoretical Contributions**  \n  - **Thm 3.3 significance:** First provable guarantee that geodesics traverse density supports (Eq (9) strong convexity ensures well-posed geometry). An interpretation of the theorem is stated right above the theorem statement.  \n  - Having clarified now that $\\nabla \\psi \\circ \\varphi$ should be read as $(\\nabla \\psi) \\circ \\varphi$, it might be easier to see that we get this term by the chain rule for gradients, which indeed gives that  \n    $$\n    \\nabla (\\psi \\circ \\varphi)(x) = (D_x \\varphi)^\\top \\nabla \\psi(\\varphi(x)).\n    $$\n\n- **Methodological Choices**  \n  - $\\nabla \\psi$ vs. $\\psi$: While $\\psi \\circ \\varphi$ isn't a diffeomorphism, $\\nabla \\psi \\circ \\varphi$ enables pullback geometry while canceling gradient terms in our specific setting  \n\n- **Comparison to [c–d]:**  \n  - [c] requires predefined temporal structure; our approach is agnostic to data ordering. So it cannot be compared in our setting  \n  - [d] is a review paper that acknowledges our method and does not propose any new methods (note that all relevant works cited there are discussed in our paper as well)  \n\n- **Experimental Scope**  \n  - **SPD/scenario recovery:** Possible in principle via diffeomorphism design, but our focus is $\\mathbb{R}^d \\rightarrow \\mathbb{R}^d$ cases central to ML applications  \n  - **Submersion limitations:** Current diffeomorphism requirement precludes dimension changes, but alternative pullback formulations remain possible. However, we feel that one should think about submersions, diffeomorphisms, and immersions as different ways of defining Riemannian geometry. We choose diffeomorphism as we get a lot for free, which is generally not the case for sub- and immersions.  \n\n- **Future Directions**  \n  - We see potential for applications with quotient manifolds (equivariant diffeomorphisms needed) and non-flat geometries (Diepeveen 2024 connections). While curvature/group-action extensions are compelling, they require new flow architectures beyond our present scope and the current RAE result should be generalized to account for curvature — an exciting research trajectory we hope to enable.\n\nWe hope the reviewer finds the clarifications satisfactory and kindly invite them to consider updating their score in light of these revisions. (Note that we already made the mentioned changes, but are unable to upload the updated manuscript. So these changes will become visible for the camera ready version.)", "claims_and_evidence": "see wk", "methods_and_evaluation_criteria": "see wk", "theoretical_claims": "see wk", "experimental_designs_or_analyses": "see wk", "relation_to_broader_scientific_literature": "related to manifold learning", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "**Strengths**:\n\nLearning the Riemannian structure is an interesting topic.\n\n**Weaknesses**:\n\n- Several key motivations and notations lack clarification. The direct reuse of (Diepeveen, 2024) is confusing.\n\n  - Why is Eq. (2) defined in this way, by exponential family?\n\n  - What is $\\nabla$ in Eq. (3), is it the derivative?  What is $\\nabla \\psi \\circ \\varphi(x)$?  There could be multiple understandings:  $(\\nabla_x \\psi) \\circ \\varphi$,  or $\\nabla_x (\\psi \\circ \\varphi)$, or $\\nabla_x \\psi \\circ \\varphi (x)$? The last one is the differential of $\\psi \\circ \\varphi$ on $x$ at $x$.\n\n  - What is $x$ in Eq. (3)? $x \\in \\mathcal{M} \\subset R^d$ or $x \\in R^d$?\n\n  - Why is Eq. (4) defined via an SPD quadratic form? As A is SPD, why is there an inverse?\n\n  - What is the benefit of Thm. 3.3? What are the advantages of Eq. (9) as strongly convex?\n\n- Theoretical results seem to be questionable\n  - Why there is no $\\nabla$ in Eqs. (5-8)? A well-known result is that Riemannian operators under isometry are identical modulo the diffeomorphism, which is $\\nabla \\psi \\circ \\varphi$. I only see $\\psi \\circ \\varphi$ on the left-hand side. Besides, if  $\\nabla$ corresponds to the derivative, there are more issues. This $\\nabla \\psi \\circ \\varphi$ is not a common map, it is $\\nabla _x \\psi \\circ \\varphi (x)$. So, is this a diffeomorphism? only go through the proof, it seems that it is $(\\nabla _x \\psi) \\circ \\varphi (x)$. I still fail to understand why it is defined in this way. Why should we use the $(\\nabla _x \\psi)$? If this is the case, we can simply use $\\psi$ and assume the inner product to be arbitrary (characterized by an SPD matrix). This is common in designing flat metrics induced from Euclidean spaces. Why do the authors choose a more complex presentation?\n\n  - I fail to understand the 2nd equality in Eq. (10). I do not know why there is a $\\left(D_{(\\cdot)} \\varphi\\right)^{\\top}$​.  The 1st equality is already the last one in Eq. (10)\n\n- There is some work for learning metric tensor [c-d]. The differences and advantages of this work from the previous methods should be discussed or compared.\n\nExperiments:\n\nMany metrics are designed via pullback from Euclidean space, such as the Log-Euclidean/Cholesky Metric on the SPD manifold. Can the proposed method recover this synthetic scenario?\n\n\nCan the proposed method handle the Riemannian submersion? This could be a harder question that the current version cannot address.\n-  **Quotient**: Many geometries do not have flat structures. Many of them have a structure [a]. Many cases are orbit spaces, which are submanifolds. A prototype of the construction can be found in [Prop. A.2, b] or typical materials [e].\n- **Submersion**: Many manifest as Riemannian submanifolds, whose tangent space inherits the metric from the total space. These cases are even more common than the quotient or the previous isometry cases in machine learning. Can this framework deal with this?\n\n\n[a] Neural networks on Symmetric Spaces of Noncompact Type\n\n[b] A Grassmann Manifold Handbook: Basic Geometry and Computational Aspects\n\n[c] Riemannian Metric Learning via Optimal Transport \n\n[d] Riemannian Metric Learning: Closer to You than You Imagine\n\n[e] Introduction to Riemannian Manifolds", "other_comments_or_suggestions": "see wk", "questions_for_authors": "see wk", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EalNOTh2Pn", "reviewer": "Reviewer_dUdy", "summary": "This paper introduces a novel score-based pullback Riemannian geometry framework to extract the intrinsic geometry of data manifolds using anisotropic flows. The key contributions include:\n\n1. Score-Based Riemannian Metric: Defines a data-driven Riemannian structure where geodesics pass through high-density regions of data.\n2. Riemannian Autoencoder (RAE): Constructs a Riemannian autoencoder with error bounds to estimate the intrinsic dimension of the data manifold.\n3. Integration with Normalizing Flows: Proposes modifications to normalizing flow training by incorporating isometry regularization, enabling stable manifold learning.\n4. Theoretical Guarantees: Provides closed-form geodesics and error bounds for the manifold approximation.\n5. Empirical Validation: Demonstrates effectiveness on synthetic and image datasets, showing superior manifold preservation and intrinsic dimensionality estimation.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "1.) By “respecting” the data distribution, we mean that the geodesics induced by the metric traverse regions of high data density. For a rigorous mathematical formulation of this concept please refer to our answer to your second question. This property naturally extends to all considered manifold mappings, as they inherit this behavior from the properties of the geodesics.\n\n2.) We thank the reviewer for bringing this up. There are different ways of formalizing this based on assumptions on the data distribution (so there is no standard way of doing this and previous work has been somewhat vague about this!). In our work this is formalized in Theorem 3.3. Right above the theorem we state: \n\n> \"A direct result of Proposition 3.1 is that geodesics will pass through the support of $p(x)$ from (2), in the sense that geodesics pass through regions with higher likelihood than the end points. This can be formalized in the following result.\"\n\nIn other words, this is the statement that tells us that the goal of data-driven Riemannian geometry has been achieved (also from this section: \n\n> \"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.\").\n\nWe hope the reviewer finds the clarifications satisfactory and kindly invite them to consider updating their score in light of our response.", "claims_and_evidence": "I am a bit confused for some of their claims. See questions below.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "I checked Proposition 3.1 and Theorem 3.3, and the derivations for geodesics and Riemannian distances look good to me.", "experimental_designs_or_analyses": "Yes. I think it has clear comparisons to standard normalizing flows (NF), anisotropic NF, and isometric NF.", "relation_to_broader_scientific_literature": "This paper builds on normalizing flows (Dinh et al., 2017) and score-based generative modeling (Song et al., 2021) and extends pullback geometry methods (Diepeveen, 2024) to a score-based setting.", "essential_references_not_discussed": "No as far as I know.", "other_strengths_and_weaknesses": "Strength: Well-founded in Riemannian geometry and a novel combination of generative modeling and pullback metrics.\n\nWeakness: Only tested on synthetic datasets and MNIST.", "other_comments_or_suggestions": "Some of the claims make me a bit confused (see questions below). I would suggest make them more clear.", "questions_for_authors": "1. Regarding the first contribution listed in section 1.1, what does it mean by 'respect' the data distribution? Can you make it formal using a rigorous math statement?\n\n2. At the beginning of Section 3, it claims that 'the ultimate goal is ... such that geodesics always pass through the support of data probability densities'. Could you give a formal mathematical statement about this? It's not clear what it means by 'geodesics pass through the suppoart of a probability distribution'.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5oMWGbL2qz", "reviewer": "Reviewer_M1sQ", "summary": "The manifold assumption states that data often reside in low-dimensional submanifolds of the original ambient space. The primary goal of this work is to learn the structure of this low-dimensional manifold and identify the intrinsic dimensionality. This is solved by considering unimodal densities obtained by deforming Gaussians under smooth diffeomorphisms and inferring the geometry of the sub-manifold given by regions of high likelihoods. By fitting this density on data in a manner akin to normalizing flows, one is able to learn the approximate geometry (e.g. Riemannian metric, geodesics, Levi-Civita connection, etc...) of this underlying manifold by pulling back with respect to the diffeomorphsim. Furthermore, since the learned density is modeled as the pullback of a Gaussian under this diffeomorphism, one can also apply PCA on the deformed space (where the data manifold is Euclidean) to infer the principal components and, hence, the intrinsic dimensionality of the data manifold.\n\n## Update after rebuttal\nThe authors have convincingly addressed my original concerns regarding scalability benefits and the necessity of including anisotropy and isometry regularization in recovering the intrinsic dimensions. Hence, I have increased my evaluation of the paper to a 4 from my original score of 3.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### **Weaknesses**\n\n**1.)** Refer to our response to Q1.\n\n**2.)** While our parametric assumption indeed focuses on unimodal distributions of the form $p(\\mathbf{x}) \\propto e^{-\\psi(\\varphi(x))}$, we observed encouraging robustness in multimodal settings, as demonstrated by the MNIST dataset. Although MNIST was the only multimodal dataset tested, we believe it clearly illustrates the method's effectiveness in multimodal scenarios. A thorough investigation into our model's performance across a wider variety of multimodal datasets is an interesting direction we reserve for future work.\n\n**3.)** Refer to our response to Q2.\n\n---\n\n### **Questions for Authors**\n\n**Q1.)** **Sorrenson et al.** first fit a normalizing flow to the data to construct a density-weighted KNN graph. Dijkstra’s algorithm is applied to estimate initial geodesic paths, which are then refined by numerically solving the geodesic ODE using a separately trained score model. However, this pipeline does not scale beyond low-dimensional toy datasets (~25D), due to key bottlenecks: (i) the graph-based initialization becomes unreliable in higher dimensions (Section 5.2.1), and (ii) solving the geodesic ODE is computationally expensive and highly sensitive to initialization quality (Section 4.5). These limitations are explicitly acknowledged in their paper.\n\n**Diepeveen (2024)** proposes to learn pullback geometry using invertible neural networks, but the method has not been shown to scale beyond toy 2D datasets. Importantly, their training setup also requires computing graph-based geodesic distances—introducing the same scalability bottlenecks as Sorrenson’s method, now shifted to the training phase.\n\nIn contrast, **our method** sidesteps graph-based constructions entirely. We derive closed-form expressions for key manifold maps (see Proposition 3.1), enabling efficient geodesic computation via two forward passes $\\phi(x)$, $\\phi(y)$ and a single inverse pass $\\phi^{-1}((1 - t)\\phi(x) + t\\phi(y))$. This enables scalability to high-dimensional, complex distributions and additionally downstream constructions like RAEs, which are not accommodated by previous methods. We obtained high-quality geodesic paths for datasets like Blobs and MNIST, which were omitted from the original submission due to our decision to focus Section 6.2 on the RAE. We are happy to include both this discussion and the corresponding experimental results in the revised manuscript to further clarify the scalability advantages of our method.\n\n---\n\n**Q2.)** A standard normalizing flow (NF) with a fixed identity covariance in its base distribution lacks any mechanism to distinguish important latent directions from unimportant ones, and thus cannot discover the intrinsic dimension of the data. The same limitation applies to an “isometric” NF, since its base covariance remains fixed at the identity.\n\nConsequently, the only meaningful comparison is an *anisotropic* NF **without** isometric regularization. However, in the absence of a mechanism enforcing approximate isometry, the model has no incentive to align its latent space with the data manifold—making such alignment highly unlikely to occur by chance. By contrast, isometry regularization guides the model to build a latent space approximately isometric to the data manifold, ensuring that higher-variance latent dimensions encode meaningful on-manifold variability, while collapsed-variance dimensions capture off-manifold noise. This mechanism underpins the model’s ability to uncover the intrinsic dimension of the underlying data manifold.\n\nFrom a *theoretical* standpoint (Theorem 4.1), the expected reconstruction error depends on how closely the diffeomorphism $\\phi$ approximates an isometry, underscoring the importance of isometry regularization in the performance of RAE. *Empirically*, we observed that when the isometry weight $\\lambda_{\\mathrm{iso}}$ was too small—effectively an unregularized anisotropic NF—it significantly overestimated the intrinsic dimension. In contrast, setting $\\lambda_{\\mathrm{iso}}$ to a sufficiently large value guided $\\phi$ toward an isometry and enabled the RAE to reliably and accurately uncover the true intrinsic dimensionality, in line with our theoretical analysis and intuition.\n\nWe will add a remark in Section 6.2 clarifying why **both anisotropy and isometry regularization** are crucial for reliably and accurately uncovering the data’s intrinsic structure.\n\n---\n\n**Q3.)** It is the MSE, as stated in the caption.\n\n---\n\n### **Other Comments Or Suggestions**\n\nWe do not fully agree with the reviewer here. This connection exists if the diffeomorphism is a smooth local isometry. It is true that in the original version, we did not mention “smooth local”. We will add this in the revised version. Moreover, as we regularize for local isometry, we do have the connection to score in all of our theory and experiments.", "claims_and_evidence": "From the theoretical side, the main claims made in this paper are:\n- Closed form expressions of the Riemannian structure induced by the distribution $p(\\mathbf{x}) \\propto e^{-\\psi(\\varphi(\\mathbf{x})))}$.\n- Claim that the corresponding geodesics pass through high likelihood regions of $p(\\mathbf{x})$ by demonstrating geodesic convexity of $\\psi$.\n- Error bound for Riemannian PCA.\n\nThe above claims are supported by rigorous proofs.\n\nFrom the numerical side, the authors show that:\n- The proposed model and learning algorithm successfully learns the geometry of the data manifold.\n- The Riemannian PCA is able to learn low-dimensional representations of high-dimensional data.\n\nFor the first point, the authors only demonstrate this on benchmarks given by deformations of a unimodal Gaussian - an idealised setting considered in this model. They consider an ablation with regards to various components of the model, in particular, anisotropy and near-isometry. They demonstrate clearly the benefits of adding both components to achieve successful learning of the underlying geometry.\n\nThe second experiment is performed on various high dimensional datasets, where it is demonstrated that in the unimodal case, the Riemannian PCA successfully learns the intrinsic dimension. In the multimodal case (MNIST), which is not a natural setting for the proposed model, the algorithm is demonstrated to still work well, despite overestimating the intrinsic dimensions.\n\nHowever, comparisons to other manifold learning methods are missing, which makes the claim that the proposed model _successfully balances scalability of manifold mapping and training cost_ (which seems to be the main claim that the paper makes) less evident.", "methods_and_evaluation_criteria": "The data and metrics used to assess the model are chosen appropriately.", "theoretical_claims": "I have checked the correctness of the closed form expressions of the induced Riemannian structure and geodesic convexity of $\\psi$. I have not checked the proof details of the error analysis of Riemannian PCA (Appendix B).", "experimental_designs_or_analyses": "I have checked the soundness of the experimental setup used in the paper, which is given in details in the supplementary. I do not see any obvious issues.", "relation_to_broader_scientific_literature": "The present work is extension of ongoing work in data-driven Riemannian geometry, where the goal is to learn intrinsic low-dimensional manifold representations of data. Previous works have utilised generative models to learn the underlying manifold structure, however have limited applicability, due to difficulties in working with the learned geometry (as in Sorrenson et al. 2024) and/or difficulties in training (Diepeveen, 2004, etc...). The work seeks to find a model/algorithm that is scalable both in training time and deployment. However, in doing so, they limit their models to be coming from probability distributions of the form $p(\\mathbf{x}) \\propto e^{-\\psi(\\varphi(\\mathbf{x})))}$ for quadratic $\\psi$, which is essentially a family of probability distributions obtained by nonlinear deformations of the Gaussian (i.e. normalising flows).\n\nI would however remark that the term \"score-based\" in the title can be a little misleading, since, as the authors also demonstrate, the proposed model is only equivalent to the pullback geometry with respect to the score function if and only if $\\varphi$ is a Euclidean isometry. In the end, the training follows that of normalising flows with minor adaptations to accommodate approximate isometry.", "essential_references_not_discussed": "I believe the authors cover the literature sufficiently well to motivate their work, however, I am not familiar enough with the literature in generative manifold learning to point out specific related works.", "other_strengths_and_weaknesses": "__Strengths:__\n- The paper provides a rigorous error analysis of the Riemannian PCA, which illuminates how the error shrinks with the $\\epsilon$ parameter, which controls the number of principal components to use.\n- Experimental results and figures show clear evidence that the model is capable of learning the Riemannian structure of the underlying data manifold, as well as the intrinsic dimension.\n\n__Weaknesses:__\n- Ultimately, I don't believe the authors adequately support their claim about _striking a good balance between scalability of training and evaluating manifold mappings_. While it makes sense that by considering a simplified model, this should be true, a more rigorous statement about how the proposed method is scalable would be useful. Additionally, it would be useful to see comparisons with the previous \"unscalable\" approaches in their experiments as baselines.\n- It would also be useful to see if there are any limitations posed by restricting to unimodal distributions of the form $p(\\mathbf{x}) \\propto e^{-\\psi(\\varphi(\\mathbf{x})))}$. Are there multimodal examples where the model fails to work? Or is it fairly robust, as shown in the MNIST example?\n- To state simply, the proposed model is a normalizing flow with regularization. Hence, the inclusion of this regularization seems to be the key algorithmic innovation. While the first experiment shows the benefit of this regularization in learning the manifold mapping, it is not demonstrated in the second experiment for learning the intrinsic dimension of the data. It would be useful to understand if the results in the experiments in Section 6.2 are not just due to normalizing flows but if the added regularization is contributing to the observed positive results.", "other_comments_or_suggestions": "I believe the term \"score-based pullback Riemannian geometry\" is used quite loosely here, as the model does not use the score function -  it is only related in the case when $\\varphi$ is an isometry. I believe this can be misleading as readers might come in expecting the use of the score function obtained e.g. using score-based diffusion models, to build a Riemannian geometry that represents the data. However, this is not really what is done in this paper.", "questions_for_authors": "1. Can you demonstrate the benefits in terms of scalability of the proposed method compared to previous works like Sorrenson et al. (2024), Diepeveen (2024), etc? I believe this is important in order to adequately support the claims in this paper.\n2. How do the results in Section 6.2 compare when just using a normalising flow (NF) to obtain the pullback geometry? Having an ablation with NF, Anisotropic NF, Isometric NF would also be useful here to find out if the regularisations are complementing the NF to help discover the intrinsic dimension of the underlying Riemannian geometry.\n3. What is the reconstruction error in Figure 3? Is it the MSE?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper integrates 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. With this structure, the paper constructs an autoencoder to discover the correct data manifold dimension. The proposed framework can be used with anisotropic normalizing flows by adopting isometry regularization. Experimental results are provided demonstrating the performance of the method.

I find the paper to be very well written and interesting, and I agree with all the positive comments of the reviewers. One small drawback of the paper is that it only caters to unimodal distributions.

Given the positive reviews of the reviewer, I am happy to recommend accepting the paper. I urge authors to incorporate all the comments of the reviewers along with the rebuttal discussion in the camera-ready version of the paper.
Decision: Accept (poster)