Score-based Pullback Riemannian Geometry: Extracting the Data Manifold Geometry using Anisotropic Flows

Weaknesses

1.) Refer to our response to Q1.

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.

3.) Refer to our response to Q2.

---

Questions for Authors

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.

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.

In 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.

---

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.

Consequently, 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.

From 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.

We 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.

---

Q3.) It is the MSE, as stated in the caption.

---

Other Comments Or Suggestions

We 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.

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.

2.) 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:

"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."

In other words, this is the statement that tells us that the goal of data-driven Riemannian geometry has been achieved (also from this section:

"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.").

We hope the reviewer finds the clarifications satisfactory and kindly invite them to consider updating their score in light of our response.

We thank the reviewer for their positive constructive feedback! See below for a discussion on the weaknesses:

• “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.”

  • 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.

• “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.”

  • 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?

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.

• Clarifications & Definitions

  • Notation consistency:
    • Added explicit clarifications when citing (Diepeveen, 2024) and introduced footnotes for operator precedence (e.g., $\nabla \psi \circ \varphi \equiv (\nabla \psi) \circ \varphi$)
  • Equation motivations:
    • Eq (2): Exponential family chosen for (1) score function compatibility and (2) theoretical/algorithmic tractability with unimodal distributions
    • Eq (4): Inverse SPD matrix preserves covariance interpretation and enables cleaner nonlinear PCA generalization
  • Manifold scope: Our metric operates on all $\mathbb{R}^d$ rather than explicit submanifolds $\mathcal{M} \subset \mathbb{R}^d$

• Theoretical Contributions

  • 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.
  • 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 $$\nabla (\psi \circ \varphi)(x) = (D_x \varphi)^\top \nabla \psi(\varphi(x)).$$

• Methodological Choices

  • $\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

• Comparison to [c–d]:

  • [c] requires predefined temporal structure; our approach is agnostic to data ordering. So it cannot be compared in our setting
  • [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)

• Experimental Scope

  • 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
  • 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.

• Future Directions

  • 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.

We 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.)