Score-based pullback Riemannian geometry

• Strengths:

– We thank the reviewer for highlighting the interpretability our method offers!

• Weaknesses:

– Regarding the first item (see the related question by reviewer WnXv and QfsL):

∗ We agree that the way in which the work is presented, the setup seems limited. Having that said, in training we can actually already have multiple modes because we train it as an adapted NF (and in practice we have visibly good geodesics for more complicated data sets such as MNIST digit 1). For the downstream geometry we could just use the learned diffeomorphism and strongly convex function. One of the main points of the paper is that from theory we would expect that without modifications we might be unable to guarantee stability of manifold mapping and find the right dimension with the RAE. We can confirm that we indeed get too high a dimension as we would hope for in the case of MNIST although we get visibly good geodesics. We added this experiment to the appendix.

∗ Without disclosing too much on ongoing research, there are actually two ways we are working on to work around current practical challenges: a more complicated base distribution (e.g., multimodal Gaussian along the lines of [1]) or post-processing under a less restricted normalizing flow (so that det is not equal to 1). Either way, additional theory is needed to get this to work and there are several extra factors that need to be regularized for, both of which are beyond the scope of this work. So we decided not to include it, but we are happy to disclose that we are actively working on this (and to mention that this can be made computationally feasible!).

∗ Having that said, we would like to emphasize that the subsequent projects really build on top of the insights from this paper. So we feel strongly that having this base case well-understood will give a good intuition for downstream work to build upon it as we can very clearly state what it can and cannot do.

∗ Finally, to showcase to the reader that multimodality is in principle included in the way the NF training is set up, we added a footnote in section 5.

– Regarding the second item:

∗ We agree with the reviewer that it would have been nice to compare to more methods. Having that said, the point of the first experiment is to showcase that the adaptations to normalizing flow are needed to get decent geometry.

∗ For the remainder of the experiments (on RAE performance) it would have made sense to compare to other methods. However, there is to the best of our knowledge no work on Riemannian geometries in which all manifold mappings is (or even can be) computed in an efficient and stable way. In particular, if there is no closed-form solution, you would need very high-accuracy numerical solvers to find approximate manifold mapping to such that the exp and the log are approximate inverses. In our experience, the only consistent way of doing this is by [2], which is not used in the cited ML works. In these works only geodesics are considered and the way these are computed, there is typically little regard for solving the geodesic equation consistently and in a stable fashion (let along accurately). Hence, constructing RAEs in the sense of [3] would be very unstable (and very slow) for any other method, which is why we have not compared to any other mehod for this second part.

[1] Izmailov, Pavel and Kirichenko, Polina and Finzi, Marc and Wilson, Andrew Gordon, Semi-supervised learning with normalizing flows, International conference on machine learning, 4615–4630, 2020

[2] Rumpf, Martin and Wirth, Benedikt, Variational time discretization of geodesic calculus, IMA Journal of Numerical Analysis, 2015

[3] Diepeveen, Willem, Pulling back symmetric Riemannian geometry for data analysis, arXiv preprint arXiv:2403.06612, 2024

• Questions

– Regarding the first question: Please allow us to break this question down into two parts

∗ Computational Complexity of the Proposed Approach In our work, we employed the exact orthogonal regularization, defined as: $\frac{1}{b} \sum_{i=1}^b\\|J_i^{\top} J_i-I\\|_F^2$ where:

· Ji is the Jacobian of the transformation ϕ for the i-th sample,

· I is the identity matrix,

· b is the batch size.

Computational Complexity The complexity of the exact method is:

$O(b × d^3 + b × d × f),$

where:

· d is the ambient dimension,

· f is the cost of a forward and backward pass through ϕ.

This scales cubically with d and is independent of the intrinsic dimension. We leveraged PyTorch’s vmap to efficiently compute this for dimensions up to d = 100 in our experiments.

∗ Approximate Method for Higher Dimensions In preliminary experiments for higher-dimensional data, we used an approximate regularization method. Instead of computing the full Jacobian, we approximate the orthogonality condition using v random orthonormal vectors $(v_j, v_j)_2$ =1. The regularization term is:

$\frac{1}{b} \sum_i^b \sum_{j=1}^v\\|J_i^{\top} J_i v_j-v_j\\|^2 \text {, }$

where $J_i v_j$ is the Jacobian-vector product for sample i and vector $v_j$ .

Complexity of Approximate Method

$O(b × d × v^2 + b × v × f + b × v^3).$

This reduces the computational cost, scaling linearly with d, and is also independent of the intrinsic dimension. It offers a scalable alternative for high-dimensional datasets. For our main experiments, we used the exact method due to its strong regularization in moderate dimensions (d ≤ 100). However, the approximate method was tested in preliminary high-dimensional experiments and effectively enforced orthogonality, promoting near-isometric mappings as required by our theoretical framework. The exact method ensures robust regularization in lower to moderate dimensions, while the approximate method provides a scalable alternative for higher-dimensional cases. By leveraging a small number of slicing vectors, it reduces the computational burden while preserving key geometric properties, making it effective across varying dimensional regimes.

– Regarding the second question:

∗ Please see the response of first item in weaknesses. In a nutshell: theory and additional regularization, both of which go beyond the scope of this work.

• Strengths:

– We thank the reviewer for highlighting the interesting connection between Riemannian geometry and generative modeling!

• Weaknesses

– Regarding the first item:

∗ Would the reviewer be more specific which equations should be motivated more clearly?

∗ For section 3, we would like to note that the main motivation for all equations come from the fact that these lead to closed-form manifold mappings in proposition 1. We also state in section 3 that this result is the main motivation for considering the metric tensor field of interest in the first place and give intuition (and now a more general proof in proposition 2 in the appendix) for why geodesics and other manifold mappings will move through the data support.

∗ For section 4, all equations are motivated by the definition of a Riemannian autoencoder given by Diepeveen (2024). The specific way of constructing it is motivated by Theorem 1, which is also motivated from a high level in this section.

∗ For section 5, the loss follows directly once we decide that we want to train using normalizing flows, but want to take the theory into account (which gives the regularizers).

– Regarding the second item:

∗ We agree with the reviewer that this was currently not in the paper. We provide results for MNIST in the revised version (appendix G). We need a more expressive architecture for this, which makes it harder to regularize for isometry. In the experiment the geodesics are visually good, but the dimension is still off, since we do not satisfy isometry.

∗ As we also wrote to other reviewers, this is an obstacle that can be overcome and is ongoing work that is beyond the scope of this paper.

– Regarding the third item:

∗ We thank the reviewer for bringing this up. The method is actually by design incorporating noise. The whole idea is that we find the whole distribution (which is noisy) and interpolate through it or select only the dimensions in the RAE that have largest variance (so the non-noise dimensions).

∗ So our method does not directly fit a manifold and move over it, but we can interpolate between any points in space (even if starting at some noisy data point) and for the RAE select an error margin ϵ which controls the dimension of the manifold we will use to approximate the data distribution. We would like to highlight that we start the paper with this, i.e., “Data often reside near low-dimensional non-linear manifolds as illustrated in Figure 1.” In this figure (and for any of the other experiments or toy examples) the data is noisy.

– Regarding the fourth item:

∗ These are the level sets of a distribution of the form we are interested in. This is just a toy example. The axes don’t carry any meaning, which is why there are no labels

∗ We added this in the figure caption

– Regarding the fifth item:

∗ We agree that from first glance it is hard to see why this metric tensor field should be considered. As also stated in the reply for the first item, the motivation for considering this metric tensor field is just that we get geodesics of the right form that move through the data distribution. There is no further reason to attach too much meaning to the metric tensor field apart from the geometry it generates.

∗ If we would use ∇p, we would get a singular metric tensor field at the suprema of the distribution. So this would not be a good choice to start with. On top of that, even if regularization were to be used, there are no closed-form mappings for geodesics etc., which is the main motivation for using the proposed Riemannian structure (as stated in section 3).

∗ We do understand this might be unsatisfying for the reviewer. So we added an extra result (proposition 2) to the appendix, which should make the picture a bit more complete. It states (informally) that this type of geometry gives under conditions on psi geodesics (and therefore other manifold mappings) that move through the data distribution.

• Questions

– Regarding the first question:

∗ Assuming that the reviewer means $(Ξ, Φ)^{∇ψ◦φ}_x$ , this is the metric tensor field that generates the Riemannian structure. The Riemannian geometry on one’s space of interest (in this work just Rd) is generated from just this object. Once the metric tensor field is chosen, all manifold mappings (including the metric/ distance) are defined. This is all written out in section 2.

– Regarding the second question:

∗ For learning, you minimize the loss in section 5. All details for every experiment (hyper-parameters etc.) are provided in appendices. The source code is already available, but the link is not in the paper for anonymity reasons. Could the reviewer be more specific what is still missing?

– Regarding the third question:

∗ We fixed Figure 2. Are there any other figures that need more explanation? Overall we feel that we are fairly verbose in the figure captions. If the reviewer disagrees, it would be great if they could specify which figures would need improvement.

• Strengths

– We thank the reviewer for their kind words, in particular for emphasizing the originality of this work!

• Weaknesses

– Regarding the first item: We conducted additional experiments on more realistic datasets, including the digit ’1’ subset of MNIST and the Gaussian Blobs dataset. While our method produced high-quality geodesics, the Riemannian Autoencoder (RAE) tended to overestimate the intrinsic manifold dimension. These results, along with a detailed discussion, are now included in the newly added Section G of the Appendix in the revised manuscript.

– Regarding the second item: We agree that the reader could use some help with the details to verify proposition 1. We have added the proof proof to the appendix. It should be easier for the reader to go through the steps while having the right references directly there.

– Regarding the third item (see the related question by reviewer WnXv and 5u2b):

∗ Indeed, the current setup focuses on quadratic structures. However, there are actually two ways to make it more realistic: a more complicated base distribution (e.g., multimodal Gaussian along the lines of [1]) or post-processing under a less restricted normalizing flow (so that det is not equal to 1), which also gives multi-modality. Either way, additional theory is needed to get this to work and there are several extra factors that need to be regularized for, both of which are beyond the scope of this work. So we decided not to include it, but we are happy to disclose that we are actively working on this (and to mention that this can be made computationally feasible!).

∗ Having that said, we would like to emphasize that the subsequent projects really build on top of the insights from this paper. So we feel strongly that having this base case well-understood will give a good intuition for downstream work to build upon it as we can very clearly state what it can and cannot do.

• Questions

– Regarding the first question:

∗ We agree that the way in which the work is presented, the setup seems limited. Having that said, in training we can actually already have multiple modes because we train it as an adapted NF (and in practice we have visibly good geodesics for more complicated data sets such as MNIST as mentioned above). For the downstream geometry we could just use the learned diffeomorphism and strongly convex function. One of the main points of the paper is that from theory we would expect that without modifications we might be unable to guarantee stability of manifold mapping and find the right dimension with the RAE. We can confirm that we indeed get too high a dimension as we would hope for in the case of MNIST digit 1 although we get visibly good geodesics. We added this experiment to the appendix.

∗ As already mentioned above, there are several ways of generalizing the theory to accommodate for multimodality without losing out on comutational feasibility, but these are beyond the scope of this work. So rather soon, we believe that the seemingly shortcomings of this paper can be alleviated.

∗ Finally, to showcase to the reader that multimodality is in principle included in the way the NF training is set up, we added a footnote in section 5.

– Regarding the second question:

∗ We would first like to highlight that our method aims to learn Riemannian geometry while standard NF methods aim to generate data! Having that said, we don’t have computational advantage when using them our method compared to the other ones. It’s more computationally expensive than (NF and Anisotropic NF), as we need to compute the isometry regularisation term. However, through our experiments we proved that it’s necessary to combine isometry regularisation with base distribution anisotropy in order to obtain more accurate and stable manifolds maps as evidenced by the significantly improved geodesic and variation error. The adapted framework additionally allows for the construction of an interpretable latent space, which is not possible with the other variations.

– Regarding the third question (see the related question by reviewer WnXv):

∗ We would like to highlight that from just looking at the metric tensor field it will be hard to see why using information related to the score is a good choice. However, as we already explain in section 3, the motivation of the chosen structure comes solely from proposition 1. We also provide interpretations of what geodesics behave like in the paragraph that start with “This special case highlights why”.

[1] Izmailov, Pavel and Kirichenko, Polina and Finzi, Marc and Wilson, Andrew Gordon, Semi-supervised learning with normalizing flows, International conference on machine learning, 4615–4630, 2020

• Strenghts

– We thank the reviewer for highlighting the novelty of our work!

• Weaknesses

– Regarding the first item (see the related question by reviewer QfsL):

∗ We agree that just from looking at the metric tensor field it will be hard to see why this is a good choice. However, as we already explain in section 3, the motivation of the chosen structure comes solely from proposition 1. We also provide interpretations of what geodesics behave like in the paragraph that start with “This special case highlights why”. We invite the reviewer to specify what other properties of the geometry they want discussed.

∗ One can also derive directly what the resulting metric tensor is - which is precisely the squared hessian of the log likelihood. Hessian metrics have been widely studied in the literature, however here by simply squaring the matrix we are able to recover exact manifold mappings - fundamentally not altering the metric itself. In particular, this is always positive (semi)-definite and will always attract geodesics towards stationary points - which in the case of unimodal distributions is unique.

∗ We also agree that the intuition for the proposed geometry in the original version (discussed in the paragraph starting with “This special case highlights why”) mainly focuses on the Gaussian case. A more general case can be stated and proven, and is provided in the appendix (proposition 2). This new result gives conditions for when geodesics (and manifold mappings for that matter) will pass through the data support.

– Regarding the second item:

∗ We agree that this can be emphasized. In the initial version of the paper we have tried to do this by showcasing how widely used VAEs are and what the potential of Riemannian geometry is given all the methods out there that have been generalized to the manifold setting. However, in the paper we did not have a lot of space to elaborate on the shortcomings of VAEs in the mentioned scientific applications nor space to use pullback geometry for a concrete downstream application. Please allow us to discuss this here.

· Regarding shortcomings of VAEs: Very often people try to attach meaning to the VAE latent space (comparing distances, discussing densities or computing interpolants), which is a very problematic practice. For example, in Zhong (2021) cited in the paper, latent space interpolation is used to visualize protein trajectories. While this is nice to get a feeling of what has been learned, it is hard to make any scientific claims on how proteins change from one conformation to the next. Similar arguments can be made for the other papers cited.

· Regarding pullback geometry in downstream applications: Also when using VAEs for sampling, key properties might not be preserved for this reason, but which can be salvaged in a pullback setting. As a recent example consider [1], a follow up work on the work by Diepeveen (2024) (which has been the inspiration for this paper as well), in which this case is actually showcased for the downstream application of generation.

∗ So to emphasize the potential future impact of our methods we added the sentence “After that, we believe that this line of work has wide variety of downstream applications as many of the applications mentioned to motivate this line of work will benefit from more interpretable representation learning.” at the end of the conclusion. This way there is a nice connection back to the introduction, which should give the reader a reminder of the broader impact of this paper.

[1] de Kruiff, Friso and Bekkers, Erik and ¨Oktem, Ozan and Sch¨onlieb, Carola-Bibiane and Diepeveen, Willem, Pullback Flow Matching on Data Man- ifolds, arXiv preprint arXiv:2410.04543, 2024

– Regarding the third item (see the related question by reviewer QfsL and 5u2b):

∗ We agree that the way in which the work is presented, the setup seems limited. Having that said, in training we can actually already have multiple modes because we train it as an adapted NF (and in practice we have visibly good geodesics for more complicated data sets such as MNIST digit 1, which is not expected to be a perfect deformed Gaussian). For the downstream geometry we could just use the learned diffeomorphism and strongly convex function. One of the main points of the paper is that from theory we would expect that without modifications we might be unable to guarantee stability of manifold mapping and find the right dimension with the RAE. We indeed can confirm that the recovered dimension is too high in the case of MNIST, despite visibly good geodesics. We added this experiment to the appendix.

∗ Without disclosing too much on ongoing research, there are actually two ways we are working on to work around current practical challenges: a more complicated base distribution (e.g., multimodal Gaussian along the lines of [2]) or post-processing under a less restricted normalizing flow (so that det is not equal to 1). Either way, additional theory is needed to get this to work and there are several extra factors that need to be regularized for, both of which are beyond the scope of this work. So we decided not to include it, but we are happy to disclose that we are actively working on this (and to mention that this can be made computationally feasible!).

∗ Having that said, we would like to emphasize that the subsequent projects really build on top of the insights from this paper. So we feel strongly that having such a “restrictive case” will give a good intuition for downstream work to build upon it as we can very clearly state what it can and cannot do.

∗ Finally, to showcase to the reader that multimodality is in principle included in the way the NF training is set up, we added a footnote in section 5.

– Regarding the fourth item:

∗ The method actually also works on single digits of MNIST (we didnt have time to check for the whole data set), but we only get good geodesics, not the right dimension. This can also be fixed in post-processing, but that would be beyond the scope of this work (as mentioned above). A detailed discussion and presentation of the geodesics has been added in Appendix G in the revised version of the paper.

• Question:

– The starting and goal points for geodesics are randomly sampled from the test data. For the variation error, a perturbed point z is created by adding a small random Gaussian perturbation to the goal point x1, defined as z = x1+0.1·N(0, I). This detail is now explicitly included in Appendix D of the revised manuscript.

[2] Izmailov, Pavel and Kirichenko, Polina and Finzi, Marc and Wilson, Andrew Gordon, Semi-supervised learning with normalizing flows, International conference on machine learning, 4615–4630, 2020