Introduction
Many real-world problems can be modeled effectively with Riemannian manifolds whose metric encodes feasible state transitions [11 , 49 ] (Wurgaft et al. 2026). These transitions are represented by shortest paths, called geodesics, so that problems like trajectory inference [25 ] and probabilistic inference [19 , 52 ] can be expressed through geodesic computations. When the geometry is known a priori through domain constraints, geodesics can be computed in closed form or simulated numerically [32 , 36 ] . However, recent work in geometric machine learning has focused on learning metrics directly from data [49 , 23 , 4 ] using neural networks. While these methods offer exciting new avenues for geometric deep learning [6 ] , these neural metrics present computational challenges that limit their downstream utility.
The two computations that are fundamental to Riemannian algorithms are the exponential and logarithmic maps. These maps play the role of addition and subtraction of points on a manifold, so Riemannian extensions of classical algorithms such as PCA and regression [18 , 17 ] evaluate them repeatedly. On general manifolds, neither map has a closed form and both must be solved for numerically. The exponential map is often computed as the solution to a differential equation, like the geodesic equation or Hamilton’s equations [20 ] , both of which can be simulated efficiently with automatic differentiation [52 ] . However, the simulation updates depend on the metric’s derivative, which can be unreliable for learned metrics whose value may change rapidly [2 ] . The logarithmic map is the inverse of the exponential map, and requires solving for the geodesic interpolation between two points at a time [22 , 47 ] . Although this cost can be amortized with neural networks [25 , 28 ] , these methods take both endpoints as input and so they are not able to evaluate the exponential map. We explore a different approach to computing both maps based on learning a coordinate system that simplifies their computations, called Riemann normal coordinates.
Normal coordinates are an important tool for calculations in Riemannian geometry [32 page 78] . Within a suitable neighborhood of a point, geodesics through that point become straight lines in coordinate space, and so the exponential and logarithmic maps can be computed in closed form with addition and subtraction. However, normal coordinates are most often defined in terms of the exponential and logarithmic maps, and methods for computing them rely on the numerical approximations of these maps [7 , 34 ] . We build on an alternative characterization of normal coordinates that does not depend on either map.
We propose a novel approach for parameterizing and learning normal coordinates on a Riemannian manifold using flow charts , which are conditional invertible neural networks that represent coordinate charts. Our main contribution is an efficient loss function for training flow charts whose value is nonnegative and equals zero over a geodesic ball if and only if the chart is normal. In our experiments, we first show that flow charts trained with the normal loss can recover the logarithmic map accurately at fixed base points on both synthetic manifolds and geometries derived from data. We then show that charts trained over many base points generalize to unseen ones and achieve median logmap query speedups of 60 × 60\mathord{\times} 60 × [2 ] and 72 × 72\mathord{\times} 72 × [47 ] across sixteen geometries, and exponential-map query speedups of 1.4 1.4 1.4 8.2 × 8.2\mathord{\times} 8.2 ×
Related work
Our work is related to methods that learn parametric curves to approximate geodesic interpolations between two points in the setting where the derivatives of the metric are assumed to be unreliable [2 ] . These curve parameterizations include splines [3 , 26 ] , the posterior mean of a Gaussian process [2 ] , and neural networks [25 ] . While these methods do not learn normal coordinates directly, they offer normal coordinate estimates as their initial velocities. This approach is also how classical data-driven methods for learning normal coordinates operate, where solutions are derived from shortest paths on a graph over data points [7 , 34 ] . These approaches are different from our method in that we learn normal coordinates directly whereas the other methods recover them as a byproduct.
More broadly, flow charts are related to methods that use neural networks to amortize the cost of expensive geometric computations. The most relevant method is metric flow matching (MFM), which parameterizes the geodesic between two endpoints with a neural network [25 ] . This method learns a geodesic interpolant by minimizing the energy of a parametric curve between the two endpoints under the given metric. Other methods parameterize the geodesic distance function instead. Kelshaw and Magri Kelshaw and Magri [28 ] and Li et al. Li et al. [33 ] train a network to satisfy the eikonal equation of a given metric. These methods are able to recover the logarithmic map using the gradient of the distance function. Flow charts are also neural network based, but they parameterize normal coordinate charts, which are a more general structure that naturally model the exponential map, logarithmic map and Riemannian distance at once.
Like flow charts, there are other methods that represent a geometry for data using invertible neural network architectures. The pullback metric through a normalizing flow has been used as a geometric model of data [10 ] (Diepeveen 2024) and is useful for counterfactual explanations [13 ] and interpolations [16 ] . However, each of these methods represents its geometry with a single chart. Multi-chart methods instead cover the data with a finite mixture of flow-based charts [24 , 48 , 40 , 54 ] , but are limited in the geometries they can represent because they use a fixed number of charts. Our method differs from these approaches in that it is defined on a given Riemannian manifold, instead of defining its own geometry, and parameterizes a continuous family of charts that can be expanded at any point on the manifold, instead of a discrete set of points.
Background
In this section we review basic notation and terminology that will be used in the paper. We intentionally emphasize the coordinate dependence of geometric objects because our method will parameterize a coordinate system. See Lee [31 , 32 ] for a reference on differential geometry.
Let M \mathcal{M} M n n n coordinate chart is a bijective map from a neighborhood of an abstract point p ∈ M \mathbf{p} \in\mathcal{M} p ∈ M R n \mathbb{R}^n R n M \mathcal{M} M x : M → M x = R n x:\mathcal{M} \to\mathcal{M}_x = \mathbb{R}^n x : M → M x = R n x = x ( p ) \mathbf{x}=x(\mathbf{p}) x = x ( p ) p ∈ M \mathbf{p} \in\mathcal{M} p ∈ M x x x x x x M \mathcal{M} M V x ∈ R n V_x \in\mathbb{R}^n V x ∈ R n V ∈ T p M V \in T_{\mathbf{p}}\mathcal{M} V ∈ T p M ω x ∈ R n \omega_x \in\mathbb{R}^n ω x ∈ R n ω ∈ T p ∗ M \omega\in T_{\mathbf{p}}^*\mathcal{M} ω ∈ T p ∗ M D f ( x ) = ∂ f ( x ) ∂ x Df(\mathbf{x}) = \frac{\partial f(\mathbf{x})}{\partial\mathbf{x}} D f ( x ) = ∂ x ∂ f ( x ) f f f x \mathbf{x} x
A Riemannian metric is a point-dependent inner product between two tangent vectors represented by a symmetric positive-definite matrix g x ( x ) ∈ SPD ( n ) g_x(\mathbf{x}) \in\operatorname{SPD}(n) g x ( x ) ∈ SPD ( n ) x ∈ M x \mathbf{x} \in\mathcal{M}_x x ∈ M x cometric defines an inner product between cotangent vectors and is represented by g x ( x ) − 1 g_x(\mathbf{x})^{-1} g x ( x ) − 1 ( M , g ) (\mathcal{M},g) ( M , g ) geodesics , are uniquely defined by g g g
A curve x ( ⋅ ) : R → M x \mathbf{x}_{(\cdot)}:\mathbb{R} \to\mathcal{M}_x x ( ⋅ ) : R → M x x 0 ∈ M x \mathbf{x}_0 \in\mathcal{M}_x x 0 ∈ M x V x ∈ R n V_x \in\mathbb{R}^n V x ∈ R n exponential map . Given the initial conditions ( x 0 , V x ) (\mathbf{x}_0,V_x) ( x 0 , V x ) x 1 = exp x 0 ( V x ) \mathbf{x}_1 = \exp_{\mathbf{x}_0}(V_x) x 1 = exp x 0 ( V x ) t ∈ R t \in\mathbb{R} t ∈ R x t : = exp x 0 ( t V x ) \mathbf{x}_t := \exp_{\mathbf{x}_0}(tV_x) x t := exp x 0 ( t V x )
Let B = { V x ∈ R n : V x T g x ( x 0 ) V x < ε 2 } ⊂ T x 0 M x B = \{V_x \in\mathbb{R}^n : V_x^{\mathsf{T}}g_x(\mathbf{x}_0)V_x < \varepsilon^2\} \subset T_{\mathbf{x}_0}\mathcal{M}_x B = { V x ∈ R n : V x T g x ( x 0 ) V x < ε 2 } ⊂ T x 0 M x g g g ε \varepsilon ε exp x 0 \exp_{\mathbf{x}_0} exp x 0 B B B λ x 0 : = exp x 0 ( B ) ⊂ M x \lambda_{\mathbf{x}_0} := \exp_{\mathbf{x}_0}(B) \subset\mathcal{M}_x λ x 0 := exp x 0 ( B ) ⊂ M x
centered at x 0 \mathbf{x}_0 x 0 logarithmic map , log x 0 : X x 0 → B \log_{\mathbf{x}_0}:\mathcal{X}_{\mathbf{x}_0}\to B log x 0 : X x 0 → B exp x 0 \exp_{\mathbf{x}_0} exp x 0 X x 0 \mathcal{X}_{\mathbf{x}_0} X x 0 x \mathbf{x} x x 0 \mathbf{x}_0 x 0 x \mathbf{x} x
Riemann normal coordinates (or “normal coordinates”) are a coordinate system in which geodesics through the base point are radial lines from the origin. Let X x 0 ⊂ M x \mathcal{X}_{\mathbf{x}_0}\subset\mathcal{M}_{\mathbf{x}} X x 0 ⊂ M x x 0 ∈ M x \mathbf{x}_0\in\mathcal{M}_{\mathbf{x}} x 0 ∈ M x J x ∈ R n × n J_{\mathbf{x}}\in\mathbb{R}^{n\times n} J x ∈ R n × n g g g x 0 \mathbf{x}_0 x 0 J x ⊤ g x ( x 0 ) J x = I n J_{\mathbf{x}}^{\top}g_{\mathbf{x}}(\mathbf{x}_0)J_{\mathbf{x}}=I_n J x ⊤ g x ( x 0 ) J x = I n normal coordinate chart on X x 0 \mathcal{X}_{\mathbf{x}_0} X x 0 x \mathbf{x} x s = J x − 1 V x ∈ S ⊂ R n \mathbf{s}=J_{\mathbf{x}}^{-1}V_{\mathbf{x}}\in S\subset\mathbb{R}^n s = J x − 1 V x ∈ S ⊂ R n V x = log x 0 ( x ) V_{\mathbf{x}}=\log_{\mathbf{x}_0}(\mathbf{x}) V x = log x 0 ( x ) x 0 \mathbf{x}_0 x 0 0 0 0 Figure 1 . Normal coordinate charts are invertible on their domains, with inverse s ↦ exp x 0 ( J x s ) \mathbf{s}\mapsto\exp_{\mathbf{x}_0}(J_{\mathbf{x}}\mathbf{s}) s ↦ exp x 0 ( J x s ) J x J_{\mathbf{x}} J x D exp x 0 ( 0 ) D\exp_{\mathbf{x}_0}(0) D exp x 0 ( 0 ) J x J_{\mathbf{x}} J x
Figure 1. Illustration of normal coordinate charts and our method, using our notation in § 3 and § 4 . The purple and green symbols correspond to the same points p 0 p_0 p 0 p p p
Method
We present a novel approach for learning normal coordinate charts on Riemannian manifolds. Our method represents coordinate charts using a conditional invertible neural network and is trained with a loss that renders the flow chart approximately normal.
Flow chart
We parameterize coordinate charts using a conditional invertible neural network, called a flow chart .
Definition 1 (Flow chart). Let M x ⊂ R n \mathcal{M}_{\mathbf{x}}\subset\mathbb{R}^n M x ⊂ R n M \mathcal{M} M x \mathbf{x} x h ( ⋅ ) θ : M x × M x → R n h^{\theta}_{(\cdot)}:\mathcal{M}_{\mathbf{x}}\times\mathcal{M}_{\mathbf{x}}\to\mathbb{R}^n h ( ⋅ ) θ : M x × M x → R n θ \theta θ
For every base point x 0 ∈ M x \mathbf{x}_0\in\mathcal{M}_{\mathbf{x}} x 0 ∈ M x h x 0 θ : M x → R n h^{\theta}_{\mathbf{x}_0}:\mathcal{M}_{\mathbf{x}}\to\mathbb{R}^n h x 0 θ : M x → R n f x 0 θ : = ( h x 0 θ ) − 1 f^{\theta}_{\mathbf{x}_0}:=(h^{\theta}_{\mathbf{x}_0})^{-1} f x 0 θ := ( h x 0 θ ) − 1
Every base point maps to the origin of its own chart, h x θ ( x ) = 0 h^{\theta}_{\mathbf{x}}(\mathbf{x})=0 h x θ ( x ) = 0 x ∈ M x \mathbf{x}\in\mathcal{M}_{\mathbf{x}} x ∈ M x
We denote z = h x 0 θ ( x ) \mathbf{z}=h^{\theta}_{\mathbf{x}_0}(\mathbf{x}) z = h x 0 θ ( x ) x \mathbf{x} x h x 0 θ h^{\theta}_{\mathbf{x}_0} h x 0 θ
In practice, we use normalizing flow architectures with fast inverses to parameterize flow charts [46 , 41 , 14 ] .
Flow charts and normalizing flows are closely related in that they both learn coordinate spaces where expensive computations on M x \mathcal{M}_{\mathbf{x}} M x g x g_{\mathbf{x}} g x g z θ g_{\mathbf{z}}^{\theta} g z θ x 0 ∈ M x \mathbf{x}_0\in\mathcal{M}_{\mathbf{x}} x 0 ∈ M x x ∈ M x \mathbf{x}\in\mathcal{M}_{\mathbf{x}} x ∈ M x [32 ] :
g z θ ( z ) = D f x 0 θ ( z ) ⊤ g x ( x ) D f x 0 θ ( z ) (1) g_{\mathbf{z}}^{\theta}(\mathbf{z})=Df_{\mathbf{x}_0}^{\theta}(\mathbf{z})^{\top}g_{\mathbf{x}}(\mathbf{x})Df_{\mathbf{x}_0}^{\theta}(\mathbf{z}) \tag*{(1)} g z θ ( z ) = D f x 0 θ ( z ) ⊤ g x ( x ) D f x 0 θ ( z ) (1) Note that the particular coordinate representation depends on the base point x 0 \mathbf{x}_0 x 0 g z θ ( z ) g_{\mathbf{z}}^{\theta}(\mathbf{z}) g z θ ( z ) x 0 \mathbf{x}_0 x 0
Normal flow chart
Our main observation is that normal coordinate charts can be identified as flow charts satisfying the identity in the following proposition.
Proposition 2 (Normal chart identity). Let ( M , g ) (\mathcal{M},g) ( M , g ) x : M → M x \mathbf{x}:\mathcal{M}\to\mathcal{M}_{\mathbf{x}} x : M → M x x 0 ∈ M x \mathbf{x}_0\in\mathcal{M}_{\mathbf{x}} x 0 ∈ M x X x 0 ⊂ M x \mathcal{X}_{\mathbf{x}_0}\subset\mathcal{M}_{\mathbf{x}} X x 0 ⊂ M x x 0 \mathbf{x}_0 x 0 U x 0 ⊂ X x 0 \mathcal{U}_{\mathbf{x}_0}\subset\mathcal{X}_{\mathbf{x}_0} U x 0 ⊂ X x 0 x 0 \mathbf{x}_0 x 0 h x 0 : U x 0 → R n h_{\mathbf{x}_0}:\mathcal{U}_{\mathbf{x}_0}\to\mathbb{R}^n h x 0 : U x 0 → R n z = h x 0 ( x ) \mathbf{z}=h_{\mathbf{x}_0}(\mathbf{x}) z = h x 0 ( x ) x ∈ M x \mathbf{x}\in\mathcal{M}_{\mathbf{x}} x ∈ M x h x 0 = 0 h_{\mathbf{x}_0}=0 h x 0 = 0 Z x 0 : = h x 0 ( U x 0 ) \mathcal{Z}_{\mathbf{x}_0}:=h_{\mathbf{x}_0}(\mathcal{U}_{\mathbf{x}_0}) Z x 0 := h x 0 ( U x 0 ) g z ( z ) = D f x 0 ( z ) ⊤ g x ( x ) D f x 0 ( z ) g_{\mathbf{z}}(\mathbf{z})=Df_{\mathbf{x}_0}(\mathbf{z})^{\top}g_{\mathbf{x}}(\mathbf{x})Df_{\mathbf{x}_0}(\mathbf{z}) g z ( z ) = D f x 0 ( z ) ⊤ g x ( x ) D f x 0 ( z ) z \mathbf{z} z h x 0 h_{\mathbf{x}_0} h x 0 U x 0 \mathcal{U}_{\mathbf{x}_0} U x 0 U x 0 \mathcal{U}_{\mathbf{x}_0} U x 0 X x 0 \mathcal{X}_{\mathbf{x}_0} X x 0
h x 0 ( ⋅ ) is normal on U x 0 if and only if g z ( z ) z = z , ∀ z ∈ Z x 0 (2) h_{\mathbf{x}_0}(\cdot)\text{ is normal on }\mathcal{U}_{\mathbf{x}_0}\text{ if and only if }g_{\mathbf{z}}(\mathbf{z})\mathbf{z}=\mathbf{z},\quad\forall\mathbf{z}\in\mathcal{Z}_{\mathbf{x}_0} \tag*{(2)} h x 0 ( ⋅ ) is normal on U x 0 if and only if g z ( z ) z = z , ∀ z ∈ Z x 0 (2) This identity is Gauss’s lemma written in normal coordinates and serves as a definition of normal coordinates in physics [38 ] , but, to our knowledge, has not been used in a machine learning context. See Section A for a proof.
The advantage of this characterization of normality is that the identity can be checked with any flow chart and only requires matrix-vector products with the metric and Jacobian-vector products with the flow chart. Additionally, as illustrated by Figure 1 , when a flow chart satisfies the normal chart identity, the exponential and logarithmic maps can be recovered in closed form on the domain where the identity holds:
exp x 0 ( V x ) = f x 0 θ ( D h x 0 θ ( x 0 ) V x ) , log x 0 ( x ) = D f x 0 θ ( 0 ) h x 0 θ ( x ) (3) \operatorname{exp}_{x_0}(V_x) = f_{x_0}^{\theta}\left(Dh_{x_0}^{\theta}(x_0)V_x\right), \qquad\operatorname{log}_{x_0}(x) = Df_{x_0}^{\theta}(0)h_{x_0}^{\theta}(x) \tag*{(3)} exp x 0 ( V x ) = f x 0 θ ( D h x 0 θ ( x 0 ) V x ) , log x 0 ( x ) = D f x 0 θ ( 0 ) h x 0 θ ( x ) (3) We train a flow chart to satisfy Equation (2) at a given x 0 x_0 x 0
Proposition 3 (Normal loss). Let ( M , g ) (\mathcal{M},g) ( M , g ) h x 0 θ h_{x_0}^{\theta} h x 0 θ θ \theta θ U x ⊂ X x 0 \mathcal{U}_x \subset\mathcal{X}_{x_0} U x ⊂ X x 0 x 0 x_0 x 0 ρ x ( ⋅ ∣ x 0 ) \rho_x(\cdot\mid x_0) ρ x ( ⋅ ∣ x 0 ) U x \mathcal{U}_x U x ρ x ( x ∣ x 0 ) > 0 \rho_x(x \mid x_0) > 0 ρ x ( x ∣ x 0 ) > 0 U z = h x 0 θ ( U x ) \mathcal{U}_z = h_{x_0}^{\theta}(\mathcal{U}_x) U z = h x 0 θ ( U x ) ρ z ( ⋅ ∣ x 0 ) \rho_z(\cdot\mid x_0) ρ z ( ⋅ ∣ x 0 ) h x 0 θ h_{x_0}^{\theta} h x 0 θ U z \mathcal{U}_z U z U x \mathcal{U}_x U x h x 0 θ h_{x_0}^{\theta} h x 0 θ
L n o r m a l ( θ ) = E z ∼ ρ z [ e T g z θ ( z ) e + e T g z θ ( z ) − 1 e − 2 ] (4) \mathcal{L}_{\mathrm{normal}}(\theta) = \mathbb{E}_{z \sim\rho_z}\left[\mathbf{e}^{\mathsf{T}}g_z^{\theta}(z)\mathbf{e} + \mathbf{e}^{\mathsf{T}}g_z^{\theta}(z)^{-1}\mathbf{e} - 2\right] \tag*{(4)} L normal ( θ ) = E z ∼ ρ z [ e T g z θ ( z ) e + e T g z θ ( z ) − 1 e − 2 ] (4) where e = z / ∥ z ∥ 2 \mathbf{e} = z/\lVert z\rVert_2 e = z / ∥ z ∥ 2 z = 0 z = 0 z = 0 L n o r m a l ≥ 0 \mathcal{L}_{\mathrm{normal}} \ge0 L normal ≥ 0 L n o r m a l = 0 \mathcal{L}_{\mathrm{normal}} = 0 L normal = 0 h x 0 θ h_{x_0}^{\theta} h x 0 θ
See Section A for a proof. The integrand is the squared g g g z z z § 4.3 . Note that one can train with multiple x 0 x_0 x 0 U x \mathcal{U}_x U x ρ x ( ⋅ ∣ x 0 ) \rho_x(\cdot\mid x_0) ρ x ( ⋅ ∣ x 0 ) U x \mathcal{U}_x U x x ∼ ρ x ( ⋅ ∣ x 0 ) x \sim\rho_x(\cdot\mid x_0) x ∼ ρ x ( ⋅ ∣ x 0 )
Practical considerations
Suppose the user has access to a dataset of points D = { x ( i ) } i = 1 : N \mathcal{D} = \{x^{(i)}\}_{i=1:N} D = { x ( i ) } i = 1 : N D 0 ⊆ D \mathcal{D}_0 \subseteq\mathcal{D} D 0 ⊆ D x 0 ∈ D 0 x_0 \in\mathcal{D}_0 x 0 ∈ D 0 z z z x 0 x_0 x 0 x x x x 0 x_0 x 0 x 0 x_0 x 0 ( x 0 , x ) (x_0,x) ( x 0 , x ) 0 = h x 0 θ ( x 0 ) \mathbf{0} = h_{x_0}^{\theta}(x_0) 0 = h x 0 θ ( x 0 ) z = h x 0 θ ( x ) z = h_{x_0}^{\theta}(x) z = h x 0 θ ( x ) ρ z \rho_z ρ z N ( ⋅ ∣ t h x 0 θ ( x ) , γ 2 I n ) \mathcal{N}(\cdot\mid t h_{x_0}^{\theta}(x), \gamma^2 I_n) N ( ⋅ ∣ t h x 0 θ ( x ) , γ 2 I n ) t ∼ U ( 0 , 1 ) t \sim\mathcal{U}(0,1) t ∼ U ( 0 , 1 ) γ \gamma γ
In theory, every x x x x 0 x_0 x 0 x x x D \mathcal{D} D x 0 x_0 x 0 k k k Figure 2 illustrates this.
Figure 2. The scale of coordinates seen by a flow chart during amortized training depends jointly on how endpoint pairs are selected and the underlying geometry. (a) shows a diffusion geometry and (c) the pullback metric through a normalizing flow with a uniform prior, overlaid with the coordinate curves of learned flow charts contained in fixed radius coordinate balls. (b) and (d) show the distributions of coordinate norms for the nearest neighbors of each base point. The coordinates of endpoints selected with nearest neighbors have a scale that varies substantially across base points in (b), but is more comparable in (d) because in (a), coordinate balls tend to expand in the data space where observations are concentrated.
The figure shows two geometries built from data with coordinate curves of learned flow charts, contained within a fixed radius coordinate ball that approximate geodesic balls. Panel (a) shows a diffusion geometry Jones and Lanners [23 ] , whose metric is designed so that paths through regions supported by data are shorter than paths away from data [21 , 1 , 25 , 49 ] . Panel (c) shows the pullback metric through a normalizing flow with a uniform prior that is fit to the data [13 ] (Diepeveen 2024), whose Riemannian volume density approximates the probability density of data. The colored background in (c) shows the intensity of the volume density in green, and the inverse volume density in blue for (a). We can see that the size of the coordinate balls and the Riemannian volume density on each geometry relate to the density of data in opposite ways. This behavior has practical implications for the scale of z z z k k k x x x x 0 x_0 x 0 z z z
Experiments
Our experiments test the accuracy of charts trained at fixed base points, their performance after pretraining over base points, and their use as coordinates for a numerical solver. The experimental details are given in Section D and the full results for all tested settings are given in Section E.
Geometries We evaluate on six families of manifolds with closed-form exponential and logarithmic maps: spheres, flat tori, hyperbolic spaces, manifolds of symmetric positive-definite matrices, Grassmann manifolds, and special orthogonal groups. Each manifold is used in a low dimensional version, with n ∈ { 2 , 3 , 4 } n \in\{2, 3, 4\} n ∈ { 2 , 3 , 4 } n ∈ { 28 , 30 } n \in\{28, 30\} n ∈ { 28 , 30 } [25 ] , and a diffusion geometry built from an airline dipdeptide data (n = 30 n = 30 n = 30 [4 ] . For these four data geometries, the reference logmap is computed with a GP solver [2 ] in double precision. Each geometry has at most 128 test base points with 1,024 queries per base point on the closed form manifolds and 512 on the data geometries, before discarding queries without a valid reference solution. For each base point x 0 x_0 x 0 J ( x 0 ) \mathcal{J}(x_0) J ( x 0 ) I ( x 0 ) ⊆ J ( x 0 ) \mathcal{I}(x_0) \subseteq\mathcal{J}(x_0) I ( x 0 ) ⊆ J ( x 0 )
Flow chart architecture The flow chart we use in these experiments is a conditional neural spline flow ([14 ] ) with L L L [ − 10 , 10 ] [-10, 10] [ − 10 , 10 ] [43 ] ) residual network with three residual blocks. The conditioning is computed with a separate residual network that encodes the base point into a 64-dimensional feature vector for all coupling layers. The network widths and other architectural choices are given in Section D.1.1.
We optimize the normal loss with the Adam optimizer ([30 ] ) for 2,000,000 steps with a learning rate of 3 × 10 − 4 3 \times10^{-4} 3 × 1 0 − 4
Evaluation
Logarithmic map The baselines we consider are all based on geodesic interpolation. The first is metric flow matching (MFM) ([25 ] ), which uses a neural network to parameterize a curve that is guaranteed to start at x 0 x_0 x 0 x x x [2 ] and GEORCE [47 ] , which solve for the geodesic between the endpoints by refining a curve parameterized by K ∈ { 4 , 8 , 16 , 32 } K \in\{4, 8, 16, 32\} K ∈ { 4 , 8 , 16 , 32 }
We train flow charts using the normal loss, and MFM by minimizing the energy of its curves. In the non-amortized setting, we train a separate chart with L = 32 L = 32 L = 32 L ∈ { 4 , 8 , 16 , 32 } L \in\{4, 8, 16, 32\} L ∈ { 4 , 8 , 16 , 32 }
Let U x ( i ) U_x^{(i)} U x ( i ) V x ( i ) V_x^{(i)} V x ( i ) i i i x 0 x_0 x 0 E l o g m a p ( x 0 ) \mathcal{E}_{\mathrm{logmap}}(x_0) E logmap ( x 0 ) R l o g m a p ( x 0 ) \mathcal{R}_{\mathrm{logmap}}(x_0) R logmap ( x 0 ) x 0 x_0 x 0
E l o g m a p = 1 ∣ I ∣ ∑ i ∈ I ∥ U x ( i ) − V x ( i ) ∥ g 2 , r r e f = 1 ∣ I ∣ ∑ i ∈ I ∥ V x ( i ) ∥ g 2 , R l o g m a p = E l o g m a p r r e f \mathcal{E}_{\mathrm{logmap}} = \sqrt{\frac{1}{|\mathcal{I}|}\sum_{i \in\mathcal{I}}\lVert U_x^{(i)} - V_x^{(i)} \rVert_g^2},\quad r_{\mathrm{ref}} = \sqrt{\frac{1}{|\mathcal{I}|}\sum_{i \in\mathcal{I}}\lVert V_x^{(i)} \rVert_g^2},\quad\mathcal{R}_{\mathrm{logmap}} = \frac{\mathcal{E}_{\mathrm{logmap}}}{r_{\mathrm{ref}}} E logmap = ∣ I ∣ 1 i ∈ I ∑ ∥ U x ( i ) − V x ( i ) ∥ g 2 , r ref = ∣ I ∣ 1 i ∈ I ∑ ∥ V x ( i ) ∥ g 2 , R logmap = r ref E logmap For each geometry, we report the mean of R l o g m a p ( x 0 ) \mathcal{R}_{\mathrm{logmap}}(x_0) R logmap ( x 0 ) τ \tau τ
Exponential map We use the same geometries to evaluate the exponential map. Let ( x 0 , V x , x 1 ) (x_0, V_x, x_1) ( x 0 , V x , x 1 ) V x = log x 0 ( x 1 ) V_x = \log_{x_0}(x_1) V x = log x 0 ( x 1 ) x 0 x_0 x 0 V x V_x V x x 1 x_1 x 1 x 0 x_0 x 0 V x V_x V x [50 ] and Dopri8 [44 ] with ranges of steps selected to cap the total number of cometric evaluations at 129. We use the implementation of these solvers provided by the Diffrax library [29 ] .
Let γ 1 ( i ) \gamma_1^{(i)} γ 1 ( i ) i i i x 0 \mathbf{x}_0 x 0 E e x p m a p ( x 0 ) \mathcal{E}_{\mathrm{expmap}}(\mathbf{x}_0) E expmap ( x 0 ) R e x p m a p ( x 0 ) \mathcal{R}_{\mathrm{expmap}}(\mathbf{x}_0) R expmap ( x 0 )
E e x p m a p = 1 ∣ I ∣ ∑ i ∈ I ∥ γ 1 ( i ) − x 1 ( i ) ∥ 2 2 , r c h o r d = 1 ∣ J ∣ ∑ i ∈ J ∥ x 1 ( i ) − x 0 ∥ 2 2 , R e x p m a p = E e x p m a p r c h o r d \mathcal{E}_{\mathrm{expmap}} = \sqrt{\frac{1}{|\mathcal{I}|}\sum_{i\in\mathcal{I}}\left\lVert\gamma_1^{(i)}-\mathbf{x}_1^{(i)}\right\rVert_2^2},\quad r_{\mathrm{chord}} = \sqrt{\frac{1}{|\mathcal{J}|}\sum_{i\in\mathcal{J}}\left\lVert\mathbf{x}_1^{(i)}-\mathbf{x}_0\right\rVert_2^2},\quad\mathcal{R}_{\mathrm{expmap}} = \frac{\mathcal{E}_{\mathrm{expmap}}}{r_{\mathrm{chord}}} E expmap = ∣ I ∣ 1 i ∈ I ∑ γ 1 ( i ) − x 1 ( i ) 2 2 , r chord = ∣ J ∣ 1 i ∈ J ∑ x 1 ( i ) − x 0 2 2 , R expmap = r chord E expmap For this experiment, a query is successful if the method returns a finite endpoint. The reference scale uses all queries in J ( x 0 ) \mathcal{J}(\mathbf{x}_0) J ( x 0 ) R e x p m a p ( x 0 ) \mathcal{R}_{\mathrm{expmap}}(\mathbf{x}_0) R expmap ( x 0 ) τ \tau τ
Runtime comparisons We also compare the runtimes of the methods needed to achieve a common accuracy. For the logarithmic map, we compare pretrained flow charts with GP and GEORCE, and for the exponential map we compare them with the ODE solvers. All methods are timed on the same NVIDIA A16 gpu in the same way, and we sweep over various batch sizes to solve multiple GP and GEORCE queries at once, and report the time for the fastest batch size. For each geometry and comparison, we find the lowest error threshold that both methods can reach among their tested settings, then compare their fastest settings that meet this threshold. See Section D.4 for details.
Results
Figure 3 shows that a flow chart trained at a fixed base point can recover the logarithmic map to a degree comparable to the baselines. Over 4 base points and 1024 and 512 test queries for closed form and data-driven geometries, respectively, the flow chart achieved an average error below 0.0032, and the pretrained flow chart outperformed GEORCE on 11 of the 16 geometries. The chart is more accurate than GEORCE on all sixteen geometries, but slightly less accurate than GP on the high dimensional manifolds. These results empirically validate the correctness of training flow charts for the normal loss.
Figure 3. Mean logmap normalized RMSE at four fixed base points with error bars indicating one standard error. All methods were evaluated on the same four randomly selected test base points, and use their most favorable settings. The charts use L = 32 L = 32 L = 32 K = 16 K = 16 K = 16 K = 32 K = 32 K = 32
Figure 4 compares the performance of a flow chart with L = 4 L = 4 L = 4 n ≤ 4 n \le4 n ≤ 4 n ≥ 28 n \ge28 n ≥ 28
Figure 4. Amortized logmap accuracy relative to MFM. Each bar denotes the mean R l o g m a p \mathcal{R}_{\mathrm{logmap}} R logmap L = 4 L = 4 L = 4
For the logarithmic map, Figure 5 shows that the amortized flow chart is up to 364 × 364\times 364 × 766 × 766\times 766 × 60 × 60\times 60 × 72 × 72\times 72 × 50 50 50 150 × 150\times 150 × 41 41 41 283 × 283\times 283 × 100 × R l o g m a p 100 \times\mathcal{R}_{\mathrm{logmap}} 100 × R logmap 364 × 364\times 364 × 766 × 766\times 766 ×
Figure 5. Comparison of time needed to achieve similar error between flow charts and numerical methods. For each comparison, R ⋆ \mathcal{R}_{\star} R ⋆ R ⋆ \mathcal{R}_{\star} R ⋆
For the exponential-map comparison in Figure 5 , we vary the chart size and the number of steps used by each ODE solver. For each geometry, we select the fastest chart and the fastest setting across all four ODE solvers whose mean R e x p m a p \mathcal{R}_{\mathrm{expmap}} R expmap R ⋆ \mathcal{R}_{\star} R ⋆ n ≥ 28 n \ge28 n ≥ 28 Figure 5 shows that flow charts are 1.4–8.2× faster than the ODE solvers to achieve comparable error, and on the two dimensional point clouds, the ODE solvers are faster than the charts. Here, the fastest solver takes 0.14–0.49 times the runtime of the fastest chart. Section D.4.1 reports the exponential-map errors and the comparison at a fixed runtime budget.
Flow chart preconditioner
Finally, we show that pretrained flow charts can be used to refine the logmap solution of GEORCE. Flow charts learn a coordinate space where geodesics are approximated by radial lines through the origin. We hypothesize that this behavior can make the task of solving for geodesics simpler to boost the performance of numerical methods. We test this by comparing the performance of GEORCE in the data space and in the coordinate space. Let x 0 x_0 x 0 x x x 0 0 0 h x 0 θ ( x ) h_{x_0}^{\theta}(x) h x 0 θ ( x ) g z θ g_z^{\theta} g z θ L = 4 L = 4 L = 4 K ∈ { 4 , 8 , 16 , 32 } K \in\{4, 8, 16, 32\} K ∈ { 4 , 8 , 16 , 32 } Table 1 shows that solving GEORCE in the learned coordinate space consistently improves its performance.
Table 1. GEORCE in data coordinates and in the coordinates of a pretrained flow chart. Columns vary GEORCE’s knot count K K K 100 × R l o g m a p 100 \times\mathcal{R}_{\mathrm{logmap}} 100 × R logmap
Conclusion
We introduced flow charts, which are conditional invertible neural networks that learn normal coordinates on a given Riemannian manifold. Our training objective is non-negative and vanishes precisely at normal coordinate charts, under suitable domain assumptions, enabling direct approximations of both the exponential and logarithmic maps. Across synthetic and data-derived geometries, flow charts approximate the maps on par with numerical methods, while being substantially faster when pretrained on data, and they are comparable in performance to MMF for computing the logarithmic map. The main limitation of our method is that it requires heuristics to select training points and that our method can only answer local geodesic queries. Future work could address these limitations by investigating different training strategies, using flow charts to define their own geometry, and by re-centering charts along longer geodesics to answer queries beyond a single chart’s domain.
AI use statement
The research ideas and the design of the method and loss functions originated with the authors. We used ChatGPT and Claude to assist with designing and carrying out experiments, writing and debugging code, running jobs, analyzing results, and producing figures. These tools also helped optimize computational performance, including the efficient computation of geodesic accelerations described in the appendix. We used them to suggest draft proofs, check derivations, and proofread the manuscript. We personally checked the experimental setup, results, and proofs, and take responsibility for the final content of this work, including all AI-assisted text, claims, code, and figures.
Proofs
Lemma 4 (Contracted Christoffel symbols). Let ( M , g ) (\mathcal{M},g) ( M , g ) n n n x : M → M x x:\mathcal{M}\to\mathcal{M}_x x : M → M x x ∈ M x x\in\mathcal{M}_x x ∈ M x V x = V x i ∂ ∂ x i ∈ T x M x V_x=V^i_x\frac{\partial}{\partial x^i}\in T_x\mathcal{M}_x V x = V x i ∂ x i ∂ ∈ T x M x Γ x ( x ) i j k \Gamma_x(x)^k_{ij} Γ x ( x ) ij k x ∈ M x x\in\mathcal{M}_x x ∈ M x Γ x ( x ) ( V x , V x ) ∈ R n \Gamma_x(x)(V_x,V_x)\in\mathbb{R}^n Γ x ( x ) ( V x , V x ) ∈ R n Γ x ( x ) ( V x , V x ) k = Γ x ( x ) i j k V x i V x j \Gamma_x(x)(V_x,V_x)^k=\Gamma_x(x)^k_{ij}V_x^iV_x^j Γ x ( x ) ( V x , V x ) k = Γ x ( x ) ij k V x i V x j
Γ x ( x ) ( V x , V x ) = g x ( x ) − 1 [ D V ( g x ( x ) V x ) − 1 2 ∇ x ( V x T g x ( x ) V x ) ] (5) \Gamma_x(x)(V_x,V_x)=g_x(x)^{-1}\left[D_V(g_x(x)V_x)-\frac{1}{2}\nabla_x\left(V_x^{\mathsf{T}}g_x(x)V_x\right)\right] \tag*{(5)} Γ x ( x ) ( V x , V x ) = g x ( x ) − 1 [ D V ( g x ( x ) V x ) − 2 1 ∇ x ( V x T g x ( x ) V x ) ] (5) where V V V D V D_V D V ∇ x \nabla_x ∇ x
Proof. The Christoffel symbols of the Levi-Civita connection are given by
Γ x ( x ) i j k = 1 2 g x ( x ) k m ( ∂ g x , j m ( x ) ∂ x i + ∂ g x , i m ( x ) ∂ x j − ∂ g x , i j ( x ) ∂ x m ) (6) \Gamma_x(x)^k_{ij}=\frac{1}{2}g_x(x)^{km}\left(\frac{\partial g_{x,jm}(x)}{\partial x^i}+\frac{\partial g_{x,im}(x)}{\partial x^j}-\frac{\partial g_{x,ij}(x)}{\partial x^m}\right) \tag*{(6)} Γ x ( x ) ij k = 2 1 g x ( x ) k m ( ∂ x i ∂ g x , j m ( x ) + ∂ x j ∂ g x , im ( x ) − ∂ x m ∂ g x , ij ( x ) ) (6) Contracting with V x i V x j V_x^iV_x^j V x i V x j
Γ x ( x ) ( V x , V x ) k = g x ( x ) k m ( ∂ g x , i m ( x ) ∂ x j V x i V x j − 1 2 ∂ g x , i j ( x ) ∂ x m V x i V x j ) (7) \Gamma_x(x)(V_x,V_x)^k=g_x(x)^{km}\left(\frac{\partial g_{x,im}(x)}{\partial x^j}V_x^iV_x^j-\frac{1}{2}\frac{\partial g_{x,ij}(x)}{\partial x^m}V_x^iV_x^j\right) \tag*{(7)} Γ x ( x ) ( V x , V x ) k = g x ( x ) k m ( ∂ x j ∂ g x , im ( x ) V x i V x j − 2 1 ∂ x m ∂ g x , ij ( x ) V x i V x j ) (7) Converting this expression to matrix notation, and passing the derivatives through V x V_x V x
Proposition 2 (Normal chart identity). Let ( M , g ) (\mathcal{M},g) ( M , g ) x : M → M x x:\mathcal{M}\to\mathcal{M}_x x : M → M x x 0 ∈ M x x_0\in\mathcal{M}_x x 0 ∈ M x X x 0 ⊂ M x X_{x_0}\subset\mathcal{M}_x X x 0 ⊂ M x x 0 x_0 x 0 U x ⊆ X x 0 U_x\subseteq X_{x_0} U x ⊆ X x 0 x 0 x_0 x 0 h x 0 : U x → R n h_{x_0}:U_x\to\mathbb{R}^n h x 0 : U x → R n z = h x 0 ( x ) z=h_{x_0}(x) z = h x 0 ( x ) x ∈ M x x\in\mathcal{M}_x x ∈ M x h x 0 ( x 0 ) = 0 h_{x_0}(x_0)=0 h x 0 ( x 0 ) = 0 Z x 0 : = h x 0 ( U x ) Z_{x_0}:=h_{x_0}(U_x) Z x 0 := h x 0 ( U x ) g z ( z ) = D f x 0 ( z ) T g x ( x ) D f x 0 ( z ) g_z(z)=Df_{x_0}(z)^{\mathsf{T}}g_x(x)Df_{x_0}(z) g z ( z ) = D f x 0 ( z ) T g x ( x ) D f x 0 ( z ) z z z h x 0 h_{x_0} h x 0 U x U_x U x U x U_x U x X x 0 X_{x_0} X x 0
h x 0 ( ⋅ ) is normal on U x if and only if g z ( z ) z = z , ∀ z ∈ Z x 0 (8) h_{x_0}(\cdot)\text{ is normal on }U_x\text{ if and only if }g_z(z)z=z,\quad\forall z\in Z_{x_0} \tag*{(8)} h x 0 ( ⋅ ) is normal on U x if and only if g z ( z ) z = z , ∀ z ∈ Z x 0 (8) Proof. First we will show that when h x 0 h_{x_0} h x 0 g z ( z ) z = z g_z(z)z=z g z ( z ) z = z h x 0 h_{x_0} h x 0 f x 0 ( z ) = exp x 0 ( J x z ) f_{x_0}(z)=\exp_{x_0}(J_xz) f x 0 ( z ) = exp x 0 ( J x z ) J x = D f x 0 ( 0 ) J_x=Df_{x_0}(0) J x = D f x 0 ( 0 ) D f x 0 ( z ) = D exp x 0 ( J x z ) J x Df_{x_0}(z)=D\exp_{x_0}(J_xz)J_x D f x 0 ( z ) = D exp x 0 ( J x z ) J x V x = J x z V_x=J_xz V x = J x z W x = J x w W_x=J_xw W x = J x w w ∈ R n w\in\mathbb{R}^n w ∈ R n x 0 ∈ M x x_0\in\mathcal{M}_x x 0 ∈ M x V x ∈ R n V_x\in\mathbb{R}^n V x ∈ R n exp x 0 \exp_{x_0} exp x 0 W x ∈ R n W_x\in\mathbb{R}^n W x ∈ R n
( D exp x 0 ( V x ) V x ) T g x ( x 1 ) ( D exp x 0 ( V x ) W x ) = V x T g x ( x 0 ) W x , x 1 = exp x 0 ( V x ) (9) \left(D\exp_{x_0}(V_x)V_x\right)^{\mathsf{T}}g_x(x_1)\left(D\exp_{x_0}(V_x)W_x\right)=V_x^{\mathsf{T}}g_x(x_0)W_x,\qquad x_1=\exp_{x_0}(V_x) \tag*{(9)} ( D exp x 0 ( V x ) V x ) T g x ( x 1 ) ( D exp x 0 ( V x ) W x ) = V x T g x ( x 0 ) W x , x 1 = exp x 0 ( V x ) (9) where D exp x 0 ( V x ) D\exp_{x_0}(V_x) D exp x 0 ( V x ) V x V_x V x x = f x 0 ( z ) x=f_{x_0}(z) x = f x 0 ( z )
V x T D exp x 0 ( V x ) T g x ( x ) D exp x 0 ( V x ) W x = z T J x T D exp x 0 ( V x ) T g x ( x ) D exp x 0 ( V x ) J x w (10) V_x^{\mathsf{T}}D\exp_{x_0}(V_x)^{\mathsf{T}}g_x(x)D\exp_{x_0}(V_x)W_x=z^{\mathsf{T}}J_x^{\mathsf{T}}D\exp_{x_0}(V_x)^{\mathsf{T}}g_x(x)D\exp_{x_0}(V_x)J_xw \tag*{(10)} V x T D exp x 0 ( V x ) T g x ( x ) D exp x 0 ( V x ) W x = z T J x T D exp x 0 ( V x ) T g x ( x ) D exp x 0 ( V x ) J x w (10) = z T D f x 0 ( z ) T g x ( x ) D f x 0 ( z ) w (11) =z^{\mathsf{T}}Df_{x_0}(z)^{\mathsf{T}}g_x(x)Df_{x_0}(z)w \tag*{(11)} = z T D f x 0 ( z ) T g x ( x ) D f x 0 ( z ) w (11) = z T g z ( z ) w (12) =z^{\mathsf{T}}g_z(z)w \tag*{(12)} = z T g z ( z ) w (12) And the RHS can be written as
V x T g x ( x 0 ) W x = z T J x T g x ( x 0 ) J x w = z T w (13) V_x^{\mathsf{T}}g_x(x_0)W_x=z^{\mathsf{T}}J_x^{\mathsf{T}}g_x(x_0)J_xw=z^{\mathsf{T}}w \tag*{(13)} V x T g x ( x 0 ) W x = z T J x T g x ( x 0 ) J x w = z T w (13) Since w w w g z ( z ) z = z g_z(z)z=z g z ( z ) z = z
Next, we will prove that the identity g z ( z ) z = z g_z(z)z=z g z ( z ) z = z h x 0 h_{x_0} h x 0
① The identity near 0 0 0 g z ( 0 ) = I n g_z(0)=I_n g z ( 0 ) = I n J x : = D f x 0 ( 0 ) J_x:=Df_{x_0}(0) J x := D f x 0 ( 0 ) g g g
② The identity on Z x 0 Z_{x_0} Z x 0 Γ z ( z ) ( z , z ) = 0 \Gamma_z(z)(z,z)=0 Γ z ( z ) ( z , z ) = 0 f x 0 ( z t ) f_{x_0}(z_t) f x 0 ( z t )
③ ① + ② ⟹ h x 0 ( x ) = J x − 1 log x 0 ( x ) \Longrightarrow h_{x_0}(x)=J_x^{-1}\log_{x_0}(x) ⟹ h x 0 ( x ) = J x − 1 log x 0 ( x )
::: ① First we show that the identity implies that J x J_x J x g g g e ∈ R n e \in\mathbb{R}^n e ∈ R n ∥ e ∥ 2 = 1 \lVert e\rVert_2 = 1 ∥ e ∥ 2 = 1 s > 0 s > 0 s > 0 z = s e z = se z = se Z x 0 \mathcal{Z}_{x_0} Z x 0 g z ( s e ) s e = s e g_z(se)se = se g z ( se ) se = se s s s s → 0 s \to0 s → 0 g z ( 0 ) e = e g_z(0)e = e g z ( 0 ) e = e e e e g z ( 0 ) = I n g_z(0) = I_n g z ( 0 ) = I n D f x 0 ( 0 ) ⊤ g x ( x 0 ) D f x 0 ( 0 ) = I n D f_{x_0}(0)^\top g_x(x_0)D f_{x_0}(0) = I_n D f x 0 ( 0 ) ⊤ g x ( x 0 ) D f x 0 ( 0 ) = I n J x : = D f x 0 ( 0 ) J_x := D f_{x_0}(0) J x := D f x 0 ( 0 ) g g g
② Next we will show that the identity implies that f x 0 ( z t ) f_{x_0}(z_t) f x 0 ( z t ) V z = z V_z = z V z = z
Γ z ( z ) ( z , z ) = g z ( z ) − 1 [ D z ( g z ( z ) z ′ ) − 1 2 ∇ z ( z ′ ⊤ g z ( z ) z ′ ) ] , where z ′ = stop_gradient ( z ) (14) \Gamma_z(z)(z,z) = g_z(z)^{-1}\left[D_z\left(g_z(z)z'\right) - \frac{1}{2}\nabla_z\left(z'^\top g_z(z)z'\right)\right], \qquad\text{where } z' = \operatorname{stop\_gradient}(z) \tag*{(14)} Γ z ( z ) ( z , z ) = g z ( z ) − 1 [ D z ( g z ( z ) z ′ ) − 2 1 ∇ z ( z ′⊤ g z ( z ) z ′ ) ] , where z ′ = stop_gradient ( z ) (14) Under the identity, each of the two terms in the bracket equals zero.
Let δ ( z ) = g z ( z ) z − z \delta(z) = g_z(z)z - z δ ( z ) = g z ( z ) z − z D z δ ( z ) D_z\delta(z) D z δ ( z )
D z ( g z ( z ) z ′ ) = D δ ( z ) z − δ ( z ) (15) D_z\left(g_z(z)z'\right) = D\delta(z)z - \delta(z) \tag*{(15)} D z ( g z ( z ) z ′ ) = D δ ( z ) z − δ ( z ) (15) where D δ ( z ) D\delta(z) D δ ( z ) δ \delta δ z z z
For the second term, consider ∇ z ( z ⊤ δ ( z ) ) \nabla_z(z^\top\delta(z)) ∇ z ( z ⊤ δ ( z ))
∇ z ( z ′ ⊤ g z ( z ) z ′ ) = D δ ( z ) ⊤ z − δ ( z ) (16) \nabla_z\left(z'^\top g_z(z)z'\right) = D\delta(z)^\top z - \delta(z) \tag*{(16)} ∇ z ( z ′⊤ g z ( z ) z ′ ) = D δ ( z ) ⊤ z − δ ( z ) (16) Plugging both terms in gives the expression for the Christoffel contraction in terms of δ \delta δ
Γ z ( z ) ( z , z ) = g z ( z ) − 1 [ D δ ( z ) z − 1 2 D δ ( z ) ⊤ z − 1 2 δ ( z ) ] (17) \Gamma_z(z)(z,z) = g_z(z)^{-1}\left[D\delta(z)z - \frac{1}{2}D\delta(z)^\top z - \frac{1}{2}\delta(z)\right] \tag*{(17)} Γ z ( z ) ( z , z ) = g z ( z ) − 1 [ D δ ( z ) z − 2 1 D δ ( z ) ⊤ z − 2 1 δ ( z ) ] (17) Since Z x 0 \mathcal{Z}_{x_0} Z x 0 δ ( z ) = 0 \delta(z) = 0 δ ( z ) = 0 z ∈ Z x 0 z \in\mathcal{Z}_{x_0} z ∈ Z x 0 D δ ( z ) = 0 D\delta(z) = 0 D δ ( z ) = 0 Γ z ( z ) ( z , z ) = 0 \Gamma_z(z)(z,z) = 0 Γ z ( z ) ( z , z ) = 0
Now suppose that z ∈ Z x 0 z \in\mathcal{Z}_{x_0} z ∈ Z x 0 z t = t z z_t = tz z t = t z t ∈ ( 0 , 1 ] t \in(0,1] t ∈ ( 0 , 1 ] z t ∈ Z x 0 z_t \in\mathcal{Z}_{x_0} z t ∈ Z x 0 z t z_t z t ( M z , g z ) (\mathcal{M}_z,g_z) ( M z , g z ) V t = z ˙ t i ∂ ∂ z i V_t = \dot{z}_t^i \frac{\partial}{\partial z^i} V t = z ˙ t i ∂ z i ∂ z ˙ t = z \dot{z}_t = z z ˙ t = z z ¨ t = 0 \ddot{z}_t = 0 z ¨ t = 0 z t z_t z t ∇ V t V t \nabla_{V_t}V_t ∇ V t V t
∇ V t V t = ( z ˙ t i ∂ z ˙ t k ∂ z i + Γ z , i j k z ˙ t i z ˙ t j ) ∂ ∂ z k (18) \nabla_{V_t}V_t = \left(\dot{z}_t^i\frac{\partial\dot{z}_t^k}{\partial z^i} + \Gamma^k_{z,ij}\dot{z}_t^i\dot{z}_t^j\right)\frac{\partial}{\partial z^k} \tag*{(18)} ∇ V t V t = ( z ˙ t i ∂ z i ∂ z ˙ t k + Γ z , ij k z ˙ t i z ˙ t j ) ∂ z k ∂ (18) = ( d z ˙ t k d t ∣ z ˙ t = 0 + Γ z ( z t ) ( z ˙ t , z ˙ t ) k ) ∂ ∂ z k (19) = \left(\left.\frac{d\dot{z}_t^k}{dt}\right|_{\dot{z}_t=0} + \Gamma_z(z_t)(\dot{z}_t,\dot{z}_t)^k\right)\frac{\partial}{\partial z^k} \tag*{(19)} = ( d t d z ˙ t k z ˙ t = 0 + Γ z ( z t ) ( z ˙ t , z ˙ t ) k ) ∂ z k ∂ (19) = ( 1 t 2 Γ z ( z t ) ( z t , z t ) k ) ∂ ∂ z k (20) = \left(\frac{1}{t^2}\Gamma_z(z_t)(z_t,z_t)^k\right)\frac{\partial}{\partial z^k} \tag*{(20)} = ( t 2 1 Γ z ( z t ) ( z t , z t ) k ) ∂ z k ∂ (20) = 0 (21) = 0 \tag*{(21)} = 0 (21) The second to last step follows from bilinearity of Γ z ( z ) ( ⋅ , ⋅ ) \Gamma_z(z)(\cdot,\cdot) Γ z ( z ) ( ⋅ , ⋅ ) z t z_t z t z t z_t z t x t = f x 0 ( z t ) x_t = f_{x_0}(z_t) x t = f x 0 ( z t )
③ Recall the initial value form of geodesics, x t = exp x 0 ( t V x ) x_t = \exp_{x_0}(tV_x) x t = exp x 0 ( t V x ) V x V_x V x x 0 x_0 x 0 d d t ∣ t = 0 f x 0 ( z t ) = D f x 0 ( 0 ) z = J x z \left.\frac{d}{dt}\right|_{t=0}f_{x_0}(z_t) = Df_{x_0}(0)z = J_xz d t d t = 0 f x 0 ( z t ) = D f x 0 ( 0 ) z = J x z
f x 0 ( z t ) = exp x 0 ( t J x z ) , t ∈ ( 0 , 1 ] (22) f_{x_0}(z_t) = \exp_{x_0}(tJ_xz), \qquad t \in(0,1] \tag*{(22)} f x 0 ( z t ) = exp x 0 ( t J x z ) , t ∈ ( 0 , 1 ] (22) If we can invert the expression at t = 1 t = 1 t = 1 J x z J_xz J x z exp x 0 \exp_{x_0} exp x 0 X x 0 \mathcal{X}_{x_0} X x 0 X x 0 = exp x 0 ( B ) \mathcal{X}_{x_0} = \exp_{x_0}(B) X x 0 = exp x 0 ( B ) B = { V x ∈ R n : V x T g x ( x 0 ) V x < ϵ 2 } ⊂ T x 0 M B = \{V_x \in\mathbb{R}^n : V_x^{\mathsf{T}}g_x(x_0)V_x < \epsilon^2\} \subset T_{x_0}\mathcal{M} B = { V x ∈ R n : V x T g x ( x 0 ) V x < ϵ 2 } ⊂ T x 0 M J x z ∈ B J_xz \in B J x z ∈ B (22) can be inverted. By ①, J x J_x J x g g g ( J x z ) T g x ( x 0 ) ( J x z ) = z T z (J_xz)^{\mathsf{T}}g_x(x_0)(J_xz) = z^{\mathsf{T}}z ( J x z ) T g x ( x 0 ) ( J x z ) = z T z J x z ∈ B J_xz \in B J x z ∈ B ∥ z ∥ 2 2 < ϵ \lVert z\rVert_2^2 < \epsilon ∥ z ∥ 2 2 < ϵ
Consider the curve X t = log x 0 ( f x 0 ( z t ) ) X_t = \log_{x_0}(f_{x_0}(zt)) X t = log x 0 ( f x 0 ( z t )) [ 0 , 1 ] [0,1] [ 0 , 1 ] B B B t ∥ z ∥ 2 2 < ϵ t\lVert z\rVert_2^2 < \epsilon t ∥ z ∥ 2 2 < ϵ X t X_t X t t J x z tJ_xz t J x z B B B exp x 0 ( X t ) = f x 0 ( z t ) = exp x 0 ( t J x z ) \exp_{x_0}(X_t) = f_{x_0}(zt) = \exp_{x_0}(tJ_xz) exp x 0 ( X t ) = f x 0 ( z t ) = exp x 0 ( t J x z ) exp x 0 \exp_{x_0} exp x 0 B B B X t = t J x z X_t = tJ_xz X t = t J x z
If ∥ z ∥ 2 2 ≥ ϵ \lVert z\rVert_2^2 \ge\epsilon ∥ z ∥ 2 2 ≥ ϵ t < t ∗ = ϵ / ∥ z ∥ 2 2 ≤ 1 t < t^* = \epsilon/\lVert z\rVert_2^2 \le1 t < t ∗ = ϵ / ∥ z ∥ 2 2 ≤ 1 X t ∗ = t ∗ J x z X_{t^*} = t^*J_xz X t ∗ = t ∗ J x z X t ∗ T g x ( x 0 ) X t ∗ = ϵ 2 X_{t^*}^{\mathsf{T}}g_x(x_0)X_{t^*} = \epsilon^2 X t ∗ T g x ( x 0 ) X t ∗ = ϵ 2 B B B ∥ z ∥ 2 2 < ϵ \lVert z\rVert_2^2 < \epsilon ∥ z ∥ 2 2 < ϵ J x z ∈ B J_xz \in B J x z ∈ B f x 0 ( z ) = exp x 0 ( J x z ) = x f_{x_0}(z) = \exp_{x_0}(J_xz) = x f x 0 ( z ) = exp x 0 ( J x z ) = x
h x 0 ( x ) = J x − 1 log x 0 ( x ) , x ∈ U x (23) h_{x_0}(x) = J_x^{-1}\log_{x_0}(x), \qquad x \in\mathcal{U}_x \tag*{(23)} h x 0 ( x ) = J x − 1 log x 0 ( x ) , x ∈ U x (23) Proposition 3 (Normal loss). Let ( M , g ) (\mathcal{M},g) ( M , g ) h x 0 θ h_{x_0}^{\theta} h x 0 θ θ \theta θ U x ⊂ X x 0 \mathcal{U}_x \subset\mathcal{X}_{x_0} U x ⊂ X x 0 x 0 x_0 x 0 ρ x ( ⋅ ∣ x 0 ) \rho_x(\cdot\mid x_0) ρ x ( ⋅ ∣ x 0 ) U x \mathcal{U}_x U x ρ x ( x ∣ x 0 ) > 0 \rho_x(x \mid x_0) > 0 ρ x ( x ∣ x 0 ) > 0 U z = h x 0 ( U x ) \mathcal{U}_z = h_{x_0}(\mathcal{U}_x) U z = h x 0 ( U x ) ρ z ( ⋅ ∣ x 0 ) \rho_z(\cdot\mid x_0) ρ z ( ⋅ ∣ x 0 ) h x 0 h_{x_0} h x 0 U z \mathcal{U}_z U z U x \mathcal{U}_x U x h x 0 θ h_{x_0}^{\theta} h x 0 θ
L n o r m a l ( θ ) = E z ∼ ρ z [ e T g z θ ( z ) e + e T g z θ ( z ) − 1 e − 2 ] (24) \mathcal{L}_{\mathrm{normal}}(\theta) = \mathbb{E}_{z\sim\rho_z}\left[e^{\mathsf{T}}g_z^{\theta}(z)e + e^{\mathsf{T}}g_z^{\theta}(z)^{-1}e - 2\right] \tag*{(24)} L normal ( θ ) = E z ∼ ρ z [ e T g z θ ( z ) e + e T g z θ ( z ) − 1 e − 2 ] (24) where e = z / ∥ z ∥ 2 e = z/\lVert z\rVert_2 e = z / ∥ z ∥ 2 z = 0 z = 0 z = 0 L n o r m a l ≥ 0 \mathcal{L}_{\mathrm{normal}} \ge0 L normal ≥ 0 L n o r m a l = 0 \mathcal{L}_{\mathrm{normal}} = 0 L normal = 0 h x 0 θ h_{x_0}^{\theta} h x 0 θ U x \mathcal{U}_x U x
Proof. We drop the dependence on θ \theta θ δ ( z ) = g z ( z ) z − z \delta(z) = g_z(z)z - z δ ( z ) = g z ( z ) z − z z ∈ U z z \in\mathcal{U}_z z ∈ U z
δ ( z ) \delta(z) δ ( z ) z z z g z ( z ) z g_z(z)z g z ( z ) z V z b = g z ( z ) V z V_z^b = g_z(z)V_z V z b = g z ( z ) V z V z = z V_z = z V z = z z = ∇ z 1 2 ∥ z ∥ 2 2 z = \nabla_z \frac{1}{2}\lVert z\rVert_2^2 z = ∇ z 2 1 ∥ z ∥ 2 2 z z z d 1 2 ∥ z ∥ 2 2 \mathrm{d}\frac{1}{2}\lVert z\rVert_2^2 d 2 1 ∥ z ∥ 2 2 δ ( z ) \delta(z) δ ( z ) g g g δ ( z ) T g z ( z ) − 1 δ ( z ) \delta(z)^{\mathsf{T}}g_z(z)^{-1}\delta(z) δ ( z ) T g z ( z ) − 1 δ ( z )
First we will show that the integrand of Equation (4) is this squared g g g ∥ z ∥ 2 2 \lVert z\rVert_2^2 ∥ z ∥ 2 2
δ ( z ) T g z ( z ) − 1 δ ( z ) = ( g z ( z ) z − z ) T g z ( z ) − 1 ( g z ( z ) z − z ) = z T g z ( z ) z − 2 z T z + z T g z ( z ) − 1 z (25) \begin{aligned}
\delta(z)^{\mathsf{T}}g_z(z)^{-1}\delta(z) = (g_z(z)z-z)^{\mathsf{T}}g_z(z)^{-1}(g_z(z)z-z) \\
&= z^{\mathsf{T}}g_z(z)z - 2z^{\mathsf{T}}z + z^{\mathsf{T}}g_z(z)^{-1}z \tag*{(25)}
\end{aligned} δ ( z ) T g z ( z ) − 1 δ ( z ) = ( g z ( z ) z − z ) T g z ( z ) − 1 ( g z ( z ) z − z ) = z T g z ( z ) z − 2 z T z + z T g z ( z ) − 1 z (25) Dividing by ∥ z ∥ 2 2 \lVert z\rVert_2^2 ∥ z ∥ 2 2 e = z / ∥ z ∥ 2 e = z/\lVert z\rVert_2 e = z / ∥ z ∥ 2 e T g z ( z ) e + e T g z ( z ) − 1 e − 2 e^{\mathsf{T}}g_z(z)e + e^{\mathsf{T}}g_z(z)^{-1}e - 2 e T g z ( z ) e + e T g z ( z ) − 1 e − 2 z ≠ 0 z \ne0 z = 0
Since g z ( z ) − 1 g_z(z)^{-1} g z ( z ) − 1 z ≠ 0 z \ne0 z = 0 z = 0 z = 0 z = 0 ρ z \rho_z ρ z L n o r m a l ≥ 0 \mathcal{L}_{\mathrm{normal}} \ge0 L normal ≥ 0 z ≠ 0 z \ne0 z = 0 δ ( z ) = 0 \delta(z) = 0 δ ( z ) = 0
Next we will show that L n o r m a l = 0 \mathcal{L}_{\mathrm{normal}} = 0 L normal = 0 h x 0 h_{x_0} h x 0 U x \mathcal{U}_x U x U x \mathcal{U}_x U x x 0 x_0 x 0 X x 0 \mathcal{X}_{x_0} X x 0 U z \mathcal{U}_z U z h x 0 h_{x_0} h x 0 U x \mathcal{U}_x U x δ \delta δ U z \mathcal{U}_z U z δ \delta δ U z \mathcal{U}_z U z ρ z \rho_z ρ z L n o r m a l = 0 \mathcal{L}_{\mathrm{normal}} = 0 L normal = 0 L n o r m a l = 0 \mathcal{L}_{\mathrm{normal}} = 0 L normal = 0 U z ∖ { 0 } \mathcal{U}_z \setminus\{0\} U z ∖ { 0 } ρ z \rho_z ρ z ρ z \rho_z ρ z U z \mathcal{U}_z U z ρ x ( ⋅ ∣ x 0 ) \rho_x(\cdot\mid x_0) ρ x ( ⋅ ∣ x 0 ) U x \mathcal{U}_x U x f x 0 f_{x_0} f x 0 δ ( z ) = 0 \delta(z) = 0 δ ( z ) = 0 z ∈ U z ∖ { 0 } z \in\mathcal{U}_z \setminus\{0\} z ∈ U z ∖ { 0 } δ ( 0 ) = 0 \delta(0) = 0 δ ( 0 ) = 0
Datasets
All of the datasets are split into a training, validation and test set using a ( 70 , 10 , 20 ) (70, 10, 20) ( 70 , 10 , 20 )
Alanine dipeptide The dataset consists of 750,000 snapshots of a molecule over three 250 ns trajectories [39 , 51 ] . Each snapshot contains the three-dimensional positions of 10 heavy atoms (the full molecule has 12 extra hydrogen atoms that are not measured) so that each data point is a vector in R 30 \mathbb{R}^{30} R 30
This dataset is highly structured. As mentioned, the data is preprocessed to shift and orient each trajectory to the same view, so the data manifold is restricted to a 24 dimensional submanifold. Additionally, the mean distance between pairs of atoms across snapshots can be used to determine whether or not each pair is bonded. This yields a graph structure over the atoms in 3d space that characterizes the shape of the molecule. Most of the shape is captured by two main backbone dihedral angles, ϕ \phi ϕ ψ \psi ψ [45 ] , which is a 2d torus whose axes are each of the dihedral angles. The projection of data onto this plane organizes into clusters that, at a high level, represent the different shapes that the molecule transitions in and out of during a trajectory. For this reason, the Ramachandran plane is often used during evaluation to interpret trajectories.
Synthetic datasets The three synthetic datasets are point clouds of 500,000 points in R 2 \mathbb{R}^{2} R 2 make_moons, with a noise standard deviation of 0.06, and make_circles, with an inner radius of 0.5 and a noise standard deviation of 0.04, from Scikit-learn [42 ] . We generate the disk data from a density proportional to exp ( − 1.8 ( x 1 + 1 ) ) \exp(-1.8(x_1 + 1)) exp ( − 1.8 ( x 1 + 1 ))
Licenses and data permissions The alanine trajectories are publicly provided by the Computational Molecular Biology Group at Freie Universität Berlin through mdshare under CC BY 4.0, which permits reuse with attribution.1
Geometries
We test our method on synthetic metrics that admit closed form solutions, RBF metrics on three point clouds, and a metric learned from alanine dipeptide with Riemannian metric matching [4 ] .
Synthetic metrics
We generate data on the six synthetic manifolds in Table 2 , whose exponential maps, logarithmic maps and geodesic distances are known in closed form. Normal coordinates would already give a valid solution at their center, so we apply a nonlinear change of coordinates before presenting the data to our method. For each manifold, we construct one fixed observed coordinate system x x x x 0 \mathbf{x}_0 x 0 h x 0 θ h^{\theta}_{\mathbf{x}_0} h x 0 θ Figure 6 for a visual overview.
Figure 6. Synthetic data generation. We sample in normal coordinates at the data center p μ p_\mu p μ p r e f p_{\mathrm{ref}} p ref Ψ \Psi Ψ p 0 p_0 p 0 x 0 = x ( p 0 ) x_0 = x(p_0) x 0 = x ( p 0 )
Table 2. Overview of synthetic manifolds. V , W V,W V , W
Let n = dim M n = \dim\mathcal{M} n = dim M ι \iota ι p ∈ M \mathbf{p} \in\mathcal{M} p ∈ M J p : R n → T p M \mathbf{J}_{\mathbf{p}}: \mathbb{R}^{n} \to T_{\mathbf{p}}\mathcal{M} J p : R n → T p M g g g
We center the normal coordinates at a fixed reference point p r e f p_{\mathrm{ref}} p ref p μ p_\mu p μ p μ p_\mu p μ ℓ \ell ℓ p r e f p_{\mathrm{ref}} p ref e ∈ S n − 1 e \in S^{n-1} e ∈ S n − 1
p μ = exp p r e f ( ℓ J p r e f e ) , ℓ = { 0.3 ι , ι < ∞ , 1 , ι = ∞ (26) p_\mu= \exp_{p_{\mathrm{ref}}}(\ell J_{p_{\mathrm{ref}}}e), \qquad
\ell=
\begin{cases}
0.3\iota, & \iota< \infty,\\
1, & \iota= \infty
\end{cases} \tag*{(26)} p μ = exp p ref ( ℓ J p ref e ) , ℓ = { 0.3 ι , 1 , ι < ∞ , ι = ∞ (26) For a point p p p ι \iota ι p r e f p_{\mathrm{ref}} p ref s s s p r e f p_{\mathrm{ref}} p ref x x x
s ( p ) = J p r e f − 1 log p r e f ( p ) , x = x ( p ) = Ψ ( s ( p ) ) (27) s(p) = J_{p_{\mathrm{ref}}}^{-1}\log_{p_{\mathrm{ref}}}(p), \qquad\mathbf{x} = x(p) = \Psi(s(p)) \tag*{(27)} s ( p ) = J p ref − 1 log p ref ( p ) , x = x ( p ) = Ψ ( s ( p )) (27) Ψ : R n → R n \Psi:\mathbb{R}^n \to\mathbb{R}^n Ψ : R n → R n ι = ∞ \iota= \infty ι = ∞ M \mathcal{M} M g x = ( x − 1 ) ∗ g g_x = (x^{-1})^*g g x = ( x − 1 ) ∗ g
To construct Ψ \Psi Ψ Q 0 , … , Q 3 Q_0,\ldots,Q_3 Q 0 , … , Q 3 f 1 , f 2 , f 3 f_1,f_2,f_3 f 1 , f 2 , f 3
Ψ = Q 3 ∘ f 3 ∘ Q 2 ∘ f 2 ∘ Q 1 ∘ f 1 ∘ Q 0 (28) \Psi= Q_3 \circ f_3 \circ Q_2 \circ f_2 \circ Q_1 \circ f_1 \circ Q_0 \tag*{(28)} Ψ = Q 3 ∘ f 3 ∘ Q 2 ∘ f 2 ∘ Q 1 ∘ f 1 ∘ Q 0 (28) Each f k f_k f k c c c
c ↦ m + s sinh ( arcsinh ( c − m s ) + b ) , s > 0 (29) c \mapsto m + s \sinh\left(\operatorname{arcsinh}\left(\frac{c-m}{s}\right)+b\right), \qquad s>0 \tag*{(29)} c ↦ m + s sinh ( arcsinh ( s c − m ) + b ) , s > 0 (29) We draw b ∼ Uniform ( − 0.8 , 0.8 ) b \sim\operatorname{Uniform}(-0.8,0.8) b ∼ Uniform ( − 0.8 , 0.8 ) f k f_k f k m m m s s s Ψ \Psi Ψ
To sample base points, we draw tangent coordinates at the data center p μ p_\mu p μ
u 0 ∼ N ( 0 , σ 2 I n ) , σ = 1 n { 0.3 ι , ι < ∞ , 1 , ι = ∞ (30) u_0 \sim N(0,\sigma^2 I_n), \qquad\sigma= \frac{1}{\sqrt{n}}
\begin{cases}
0.3\iota, & \iota< \infty,\\
1, & \iota= \infty
\end{cases} \tag*{(30)} u 0 ∼ N ( 0 , σ 2 I n ) , σ = n 1 { 0.3 ι , 1 , ι < ∞ , ι = ∞ (30) Around each base point we train on a geodesic ball of radius r max = 0.25 ι r_{\max}=0.25\iota r m a x = 0.25 ι ι < ∞ \iota<\infty ι < ∞ r max = 1 r_{\max}=1 r m a x = 1 ι \iota ι u 0 \mathbf{u}_0 u 0 x \mathbf{x} x
d g ( p r e f , exp μ ( J p μ u 0 ) ) + r max < ι (31) d_g\left(\mathbf{p}_{\mathrm{ref}},\exp_{\mu}\left(\mathbf{J}_{p\mu}\mathbf{u}_0\right)\right)+r_{\max}<\iota \tag*{(31)} d g ( p ref , exp μ ( J p μ u 0 ) ) + r m a x < ι (31) We then generate n b p n_{\mathrm{bp}} n bp
x 0 = x ( exp μ ( J p μ u 0 ) ) (32) \mathbf{x}_0=\mathbf{x}\left(\exp_{\mu}\left(\mathbf{J}_{p\mu}\mathbf{u}_0\right)\right) \tag*{(32)} x 0 = x ( exp μ ( J p μ u 0 ) ) (32) At each x 0 \mathbf{x}_0 x 0 h x 0 ϑ h_{\mathbf{x}_0}^{\vartheta} h x 0 ϑ
U x ( x 0 ) = { x ∈ M x : d g ( x 0 , x ) < r max } (33) U_{\mathbf{x}}(\mathbf{x}_0)=\left\{\mathbf{x}\in\mathcal{M}_{\mathbf{x}}:d_g(\mathbf{x}_0,\mathbf{x})<r_{\max}\right\} \tag*{(33)} U x ( x 0 ) = { x ∈ M x : d g ( x 0 , x ) < r m a x } (33) To sample a point in this ball, we independently draw v ∼ N ( 0 , I n ) \mathbf{v}\sim\mathcal{N}(\mathbf{0},I_n) v ∼ N ( 0 , I n ) t ∼ Uniform ( 0 , 1 ) t\sim\operatorname{Uniform}(0,1) t ∼ Uniform ( 0 , 1 )
u = r max t 1 / n v ∥ v ∥ 2 , x = exp x 0 ( g x ( x 0 ) − 1 / 2 u ) . (34) \mathbf{u}=r_{\max}t^{1/n}\frac{\mathbf{v}}{\lVert\mathbf{v}\rVert_2},\qquad\mathbf{x}=\exp_{\mathbf{x}_0}\left(g_{\mathbf{x}}(\mathbf{x}_0)^{-1/2}\mathbf{u}\right). \tag*{(34)} u = r m a x t 1/ n ∥ v ∥ 2 v , x = exp x 0 ( g x ( x 0 ) − 1/2 u ) . (34) This generative process defines ρ x ( ⋅ ∣ x 0 ) \rho_{\mathbf{x}}(\cdot\mid\mathbf{x}_0) ρ x ( ⋅ ∣ x 0 ) U x ( x 0 ) U_{\mathbf{x}}(\mathbf{x}_0) U x ( x 0 ) n t r a i n n_{\mathrm{train}} n train n v a l n_{\mathrm{val}} n val n t e s t n_{\mathrm{test}} n test
RBF Metric
The synthetic datasets of the exponential map experiment in Section D.2 use the RBF metric of metric flow matching (MFM) ([25 ] ). This metric is constructed so that geodesics avoid regions without data.
g x ( x ) = 1 h ( x ) + ρ I n , h ( x ) = ∑ k = 1 K w k exp ( − λ k 2 ∥ x − c k ∥ 2 2 ) (35) g_{\mathbf{x}}(\mathbf{x})=\frac{1}{h(\mathbf{x})+\rho}I_n,\qquad h(\mathbf{x})=\sum_{k=1}^{K}w_k\exp\left(-\frac{\lambda_k}{2}\lVert\mathbf{x}-\mathbf{c}_k\rVert_2^2\right) \tag*{(35)} g x ( x ) = h ( x ) + ρ 1 I n , h ( x ) = k = 1 ∑ K w k exp ( − 2 λ k ∥ x − c k ∥ 2 2 ) (35) where h h h K = 50 K=50 K = 50 c k \mathbf{c}_k c k λ k \lambda_k λ k w k w_k w k ρ > 0 \rho>0 ρ > 0 [25 ] state a more general diagonal metric, their released code fits a single network and applies it to every coordinate.
The centers are selected as the K K K σ k \sigma_k σ k k k k λ k = 1 / ( 2 ( κ σ k ) 2 ) \lambda_k=1/(2(\kappa\sigma_k)^2) λ k = 1/ ( 2 ( κ σ k ) 2 ) κ = 1.5 \kappa=1.5 κ = 1.5
Let x ( 1 ) , … , x ( M ) \mathbf{x}^{(1)},\ldots,\mathbf{x}^{(M)} x ( 1 ) , … , x ( M ) M = 50,000 M=50{,}000 M = 50 , 000
min w 1 , … , w K ≥ 10 − 4 J , J = 1 M ∑ i = 1 M ( 1 − h ( x ( i ) ) ) 2 . (36) \min_{w_1,\ldots,w_K\geq10^{-4}}J,\qquad J=\frac{1}{M}\sum_{i=1}^{M}\left(1-h\left(\mathbf{x}^{(i)}\right)\right)^2. \tag*{(36)} w 1 , … , w K ≥ 1 0 − 4 min J , J = M 1 i = 1 ∑ M ( 1 − h ( x ( i ) ) ) 2 . (36) The fit makes h ≈ 1 h\approx1 h ≈ 1 g x g_{\mathbf{x}} g x ( 1 + ρ ) − 1 I n (1+\rho)^{-1}I_n ( 1 + ρ ) − 1 I n ρ − 1 I n \rho^{-1}I_n ρ − 1 I n ρ = ( 1 − J ) / 2.75 \rho=(1-J)/2.75 ρ = ( 1 − J ) /2.75 J J J M M M
Diffusion geometry
Figure 2 (a) was built with non-parametric diffusion geometries built from 2d point clouds ([23 Section 3.1] ). The pointclouds we use come from m = 50,000 m=50{,}000 m = 50 , 000 exp ( − 1.8 x 1 ) \exp(-1.8x_1) exp ( − 1.8 x 1 ) ∥ x ∥ 2 < 1 \lVert\mathbf{x}\rVert_2<1 ∥ x ∥ 2 < 1
Let { x ( j ) } j = 1 m ⊂ R n \{\mathbf{x}^{(j)}\}_{j=1}^{m}\subset\mathbb{R}^n { x ( j ) } j = 1 m ⊂ R n
k ( x , y ) = exp ( − ∥ x − y ∥ 2 2 ϵ ρ ( x ) ρ ( y ) ) , K ~ j l = k ( x ( j ) , x ( l ) ) ( q ( x ( j ) ) q ( x ( l ) ) ) ( 2 − d ) / 4 , P j l = K ~ j l ∑ l ′ = 1 m K ~ j l ′ (37) k(\mathbf{x},\mathbf{y})=\exp\left(-\frac{\lVert\mathbf{x}-\mathbf{y}\rVert_2^2}{\epsilon\rho(\mathbf{x})\rho(\mathbf{y})}\right),\qquad\widetilde{K}_{jl}=\frac{k\left(\mathbf{x}^{(j)},\mathbf{x}^{(l)}\right)}{\left(q\left(\mathbf{x}^{(j)}\right)q\left(\mathbf{x}^{(l)}\right)\right)^{(2-d)/4}},\qquad P_{jl}=\frac{\widetilde{K}_{jl}}{\sum_{l'=1}^{m}\widetilde{K}_{jl'}} \tag*{(37)} k ( x , y ) = exp ( − ϵ ρ ( x ) ρ ( y ) ∥ x − y ∥ 2 2 ) , K j l = ( q ( x ( j ) ) q ( x ( l ) ) ) ( 2 − d ) /4 k ( x ( j ) , x ( l ) ) , P j l = ∑ l ′ = 1 m K j l ′ K j l (37) where ρ ( x ) = q ( x ) − 1 / 2 \rho(\mathbf{x}) = q(\mathbf{x})^{-1/2} ρ ( x ) = q ( x ) − 1/2 q q q q q q ϵ \epsilon ϵ d ^ \hat{d} d ^ [23 ] . For each ϵ ∈ { 2 − 10 , 2 − 9.75 , … , 2 9.75 } \epsilon\in\{2^{-10}, 2^{-9.75}, \ldots, 2^{9.75}\} ϵ ∈ { 2 − 10 , 2 − 9.75 , … , 2 9.75 } ϵ \epsilon ϵ d ^ \hat{d} d ^ ϵ \epsilon ϵ d ^ \hat{d} d ^ K ~ \widetilde{\mathbf{K}} K P \mathbf{P} P μ \mu μ
We compute the leading, n 0 = 240 n_0 = 240 n 0 = 240 n 0 = 120 n_0 = 120 n 0 = 120 P \mathbf{P} P m × n 0 m \times n_0 m × n 0 U \mathbf{U} U U T diag ( μ ) U = I n 0 \mathbf{U}^{T}\operatorname{diag}(\mu)\mathbf{U} = \mathbf{I}_{n_0} U T diag ( μ ) U = I n 0 x c ∈ R m \mathbf{x}_c \in\mathbb{R}^{m} x c ∈ R m U U T diag ( μ ) x c \mathbf{U}\mathbf{U}^{T}\operatorname{diag}(\mu)\mathbf{x}_c U U T diag ( μ ) x c [23 Computational Note 7] ).
The geometry corresponding to the point cloud is given through the carré du champ of the smoothed coordinates. It is a covariance under one step of the Markov chain ([23 Equation 6] ) given by,
Γ j ( x c , x c ′ ) = 2 ϵ ρ ( x ( j ) ) 2 Cov [ x c ( X ) , x c ′ ( X ) : X ∼ P j ] . (38) \Gamma_j(x_c,x_{c^{\prime}}) = \frac{2}{\epsilon\rho(x^{(j)})^2}\operatorname{Cov}\left[x_c(X),x_{c^{\prime}}(X):X\sim P_j\right]. \tag*{(38)} Γ j ( x c , x c ′ ) = ϵ ρ ( x ( j ) ) 2 2 Cov [ x c ( X ) , x c ′ ( X ) : X ∼ P j ] . (38) We average the per-point tensors once under P \mathbf{P} P
Finally, the cometric is defined as a soft, Gaussian weighted average of the CDC over the point cloud
g x ( x ) − 1 = δ I n + ( 1 − δ ) a ( x ) ∑ j = 1 m w ( j ) ( x ) Γ j ( x , x ) , g_x(\mathbf{x})^{-1} = \delta I_n + (1-\delta)a(\mathbf{x})\sum_{j=1}^{m}w^{(j)}(\mathbf{x})\Gamma_j(\mathbf{x},\mathbf{x}), g x ( x ) − 1 = δ I n + ( 1 − δ ) a ( x ) j = 1 ∑ m w ( j ) ( x ) Γ j ( x , x ) , w ( j ) ( x ) = exp ( − ∥ x − x ( j ) ∥ 2 2 / 2 σ 2 ) ∑ l = 1 m exp ( − ∥ x − x ( l ) ∥ 2 2 / 2 σ 2 ) . (39) w^{(j)}(\mathbf{x}) = \frac{\exp\left(-\lVert\mathbf{x}-\mathbf{x}^{(j)}\rVert_2^2/2\sigma^2\right)}{\sum_{l=1}^{m}\exp\left(-\lVert\mathbf{x}-\mathbf{x}^{(l)}\rVert_2^2/2\sigma^2\right)}. \tag*{(39)} w ( j ) ( x ) = ∑ l = 1 m exp ( − ∥ x − x ( l ) ∥ 2 2 /2 σ 2 ) exp ( − ∥ x − x ( j ) ∥ 2 2 /2 σ 2 ) . (39) where δ = 2 × 10 − 3 \delta= 2 \times10^{-3} δ = 2 × 1 0 − 3 σ \sigma σ a ( x ) a(\mathbf{x}) a ( x ) a ( x ) = 1 a(\mathbf{x}) = 1 a ( x ) = 1
a ( x ) = S ( x ) S ( x ) + τ , S ( x ) = ∑ j = 1 m exp ( − ∥ x − x ( j ) ∥ 2 2 2 σ s 2 ) , (40) a(\mathbf{x}) = \frac{S(\mathbf{x})}{S(\mathbf{x})+\tau}, \qquad S(\mathbf{x}) = \sum_{j=1}^{m}\exp\left(-\frac{\lVert\mathbf{x}-\mathbf{x}^{(j)}\rVert_2^2}{2\sigma_s^2}\right), \tag*{(40)} a ( x ) = S ( x ) + τ S ( x ) , S ( x ) = j = 1 ∑ m exp ( − 2 σ s 2 ∥ x − x ( j ) ∥ 2 2 ) , (40) with ( σ s , τ ) = ( 0.02 , 1 ) (\sigma_s,\tau) = (0.02,1) ( σ s , τ ) = ( 0.02 , 1 )
Chart fitting for the diffusion illustration For Figure 2 , we fit one chart per displayed base point to its 4,096 nearest neighbors within the same connected component, excluding the base point. Each chart has four affine coupling layers with dense linear mixing, two gated condition blocks of width 32, and a one-block base point encoder of width 16 with 8 output features. We use metric preconditioning and center inputs at the base point, scaling by the median neighbor distance measured with g x ( x 0 ) g_x(x_0) g x ( x 0 ) 3 × 10 − 4 3 \times10^{-4} 3 × 1 0 − 4 r 0 r_0 r 0 γ = 0.1 r 0 / 2 \gamma= 0.1r_0/\sqrt{2} γ = 0.1 r 0 / 2 γ = 0.5 r 0 \gamma= 0.5r_0 γ = 0.5 r 0 γ \gamma γ
Normalizing-flow pullback illustration For the pullback panel of Figure 2 , we fit an unconditional flow h h h H = Φ ∘ h H = \Phi\circ h H = Φ ∘ h g x = D H T D H g_x = DH^{T}DH g x = D H T D H
For both geometries, we decode coordinate grids restricted to disks of a common radius within each panel, orient the grids using rotation-only ICA on encoded neighbors, and overlay the same 600 observations. The backgrounds use separate logarithmic color scales: blue shows inverse volume density for the diffusion metric, and green shows volume density for the pullback metric. Panels (b) and (d) show normalized histograms of the learned coordinate norms of each base point’s 4,096 nearest neighbors, using 40 equally spaced bins shared across base points within each geometry. Colors identify corresponding base points across the geometry panels and histograms.
Riemannian metric matching
For higher dimensional datasets, we use Riemannian metric matching (RMM) (Bamberger et al. [4 ] ) to learn its geometry. Let p ε ∣ 0 ( x ε ∣ x ) = N ( x ε ∣ x , ε I n ) p_{\varepsilon|0}(x_{\varepsilon} \mid x) = \mathcal{N}(x_{\varepsilon} \mid x, \varepsilon I_n) p ε ∣0 ( x ε ∣ x ) = N ( x ε ∣ x , ε I n ) ε \varepsilon ε x x x p 0 , ε ( x , x ε ) = p d a t a ( x ) p ε ∣ 0 ( x ε ∣ x ) p_{0,\varepsilon}(x, x_{\varepsilon}) = p_{\mathrm{data}}(x) p_{\varepsilon|0}(x_{\varepsilon} \mid x) p 0 , ε ( x , x ε ) = p data ( x ) p ε ∣0 ( x ε ∣ x ) p ε ( x ε ) p_{\varepsilon}(x_{\varepsilon}) p ε ( x ε ) p 0 ∣ ε ( ⋅ ∣ x ε ) p_{0|\varepsilon}(\cdot\mid x_{\varepsilon}) p 0∣ ε ( ⋅ ∣ x ε )
x ˉ 0 ( x ε ) = E x 0 ∼ p 0 ∣ ε ( ⋅ ∣ x ε ) [ x 0 ] (41) \bar{x}_0(x_{\varepsilon}) = \mathbb{E}_{x_0 \sim p_{0|\varepsilon}(\cdot\mid x_{\varepsilon})}[x_0] \tag*{(41)} x ˉ 0 ( x ε ) = E x 0 ∼ p 0∣ ε ( ⋅ ∣ x ε ) [ x 0 ] (41) and the geometry is the family of cometrics, indexed by ε \varepsilon ε
g ε , x ( x ε ) − 1 = 1 2 ε Cov x 0 ∼ p 0 ∣ ε ( ⋅ ∣ x ε ) [ x 0 ] = 1 2 D x ˉ 0 ( x ε ) (42) g_{\varepsilon,x}(x_{\varepsilon})^{-1} = \frac{1}{2\varepsilon}\operatorname{Cov}_{x_0 \sim p_{0|\varepsilon}(\cdot\mid x_{\varepsilon})}[x_0] = \frac{1}{2}D\bar{x}_0(x_{\varepsilon}) \tag*{(42)} g ε , x ( x ε ) − 1 = 2 ε 1 Cov x 0 ∼ p 0∣ ε ( ⋅ ∣ x ε ) [ x 0 ] = 2 1 D x ˉ 0 ( x ε ) (42) where the second equality is Tweedie’s formula for the posterior covariance (Efron [15 ] ). This is the mean centered carré du champ of RMM. At a noised point x ε x_{\varepsilon} x ε x ε x_{\varepsilon} x ε g ε , x θ g_{\varepsilon,x}^{\theta} g ε , x θ
L ( θ ) = E ε ∼ p ( ε ) E x ε ∼ p ε [ ∥ g ε , x θ ( x ε ) − 1 − g ε , x ( x ε ) − 1 ∥ F 2 ] = E ε ∼ p ( ε ) E ( x , x ε ) ∼ p 0 , ε [ ∥ g ε , x θ ( x ε ) − 1 − 1 2 ε ( x − x ˉ 0 ( x ε ) ) ( x − x ˉ 0 ( x ε ) ) ⊤ ∥ F 2 ] + const. (43) \begin{aligned}
\mathcal{L}(\theta) = \mathbb{E}_{\varepsilon\sim p(\varepsilon)} \mathbb{E}_{x_{\varepsilon} \sim p_{\varepsilon}} \left[\left\|g_{\varepsilon,x}^{\theta}(x_{\varepsilon})^{-1} - g_{\varepsilon,x}(x_{\varepsilon})^{-1}\right\|_F^2\right] \\
&= \mathbb{E}_{\varepsilon\sim p(\varepsilon)} \mathbb{E}_{(x,x_{\varepsilon}) \sim p_{0,\varepsilon}} \left[\left\|g_{\varepsilon,x}^{\theta}(x_{\varepsilon})^{-1} - \frac{1}{2\varepsilon}\left(x-\bar{x}_0(x_{\varepsilon})\right)\left(x-\bar{x}_0(x_{\varepsilon})\right)^{\top}\right\|_F^2\right] + \text{const.} \tag*{(43)}
\end{aligned} L ( θ ) = E ε ∼ p ( ε ) E x ε ∼ p ε [ g ε , x θ ( x ε ) − 1 − g ε , x ( x ε ) − 1 F 2 ] = E ε ∼ p ( ε ) E ( x , x ε ) ∼ p 0 , ε [ g ε , x θ ( x ε ) − 1 − 2 ε 1 ( x − x ˉ 0 ( x ε ) ) ( x − x ˉ 0 ( x ε ) ) ⊤ F 2 ] + const. (43) The second expression is tractable to optimize provided that we can compute the denoiser. For this, we either pretrain a denoiser, or download a pretrained denoiser from the internet, which we denote by D ε ( x ε ) ≈ x ˉ 0 ( x ε ) D_{\varepsilon}(x_{\varepsilon}) \approx\bar{x}_0(x_{\varepsilon}) D ε ( x ε ) ≈ x ˉ 0 ( x ε )
Low rank parameterization
We parameterize the cometric using the approximately low rank parameterization of RMM:
g ε , x θ ( x ε ) − 1 = δ I n + ( 1 − δ ) M ε θ ( x ε ) M ε θ ( x ε ) ⊤ (44) g_{\varepsilon,x}^{\theta}(x_{\varepsilon})^{-1} = \delta I_n + (1-\delta)M_{\varepsilon}^{\theta}(x_{\varepsilon})M_{\varepsilon}^{\theta}(x_{\varepsilon})^{\top} \tag*{(44)} g ε , x θ ( x ε ) − 1 = δ I n + ( 1 − δ ) M ε θ ( x ε ) M ε θ ( x ε ) ⊤ (44) where M ε θ ( x ε ) ∈ R d × n M_{\varepsilon}^{\theta}(x_{\varepsilon}) \in\mathbb{R}^{d \times n} M ε θ ( x ε ) ∈ R d × n d ≪ n d \ll n d ≪ n
We train with δ \delta δ
L ( θ ) = E ε ∼ p ( ε ) E ( x , x ε ) ∼ p 0 , ε [ ( 1 − δ ) 2 ∥ M ε θ ( x ε ) M ε θ ( x ε ) ⊤ ∥ F 2 + 2 δ ( 1 − δ ) ∥ M ε θ ( x ε ) ∥ F 2 − 1 − δ ε ∥ M ε θ ( x ε ) ( x − D ε ( x ε ) ) ∥ 2 2 ] + const. (45) \mathcal{L}(\theta) = \mathbb{E}_{\varepsilon\sim p(\varepsilon)} \mathbb{E}_{(x,x_{\varepsilon}) \sim p_{0,\varepsilon}} \left[(1-\delta)^2\left\|M_{\varepsilon}^{\theta}(x_{\varepsilon})M_{\varepsilon}^{\theta}(x_{\varepsilon})^{\top}\right\|_F^2 + 2\delta(1-\delta)\left\|M_{\varepsilon}^{\theta}(x_{\varepsilon})\right\|_F^2 - \frac{1-\delta}{\varepsilon}\left\|M_{\varepsilon}^{\theta}(x_{\varepsilon})(x-D_{\varepsilon}(x_{\varepsilon}))\right\|_2^2\right] + \text{const.} \tag*{(45)} L ( θ ) = E ε ∼ p ( ε ) E ( x , x ε ) ∼ p 0 , ε [ ( 1 − δ ) 2 M ε θ ( x ε ) M ε θ ( x ε ) ⊤ F 2 + 2 δ ( 1 − δ ) M ε θ ( x ε ) F 2 − ε 1 − δ M ε θ ( x ε ) ( x − D ε ( x ε )) 2 2 ] + const. (45) Denoiser architecture
We train a denoiser for alanine dipeptide. The denoiser is parameterized by a residual network, and predicts the clean point under additive Gaussian corruption, so its noise argument is the variance ε \varepsilon ε [27 ] , where D ϕ ( x ε , ε ) = c s k i p ( ε ) x ε + c o u t ( ε ) F ϕ ( c i n ( ε ) x ε , ε ) D_{\phi}(x_{\varepsilon}, \varepsilon) = c_{\mathrm{skip}}(\varepsilon)x_{\varepsilon} + c_{\mathrm{out}}(\varepsilon)F_{\phi}(c_{\mathrm{in}}(\varepsilon)x_{\varepsilon}, \varepsilon) D ϕ ( x ε , ε ) = c skip ( ε ) x ε + c out ( ε ) F ϕ ( c in ( ε ) x ε , ε ) F ϕ F_{\phi} F ϕ c i n \mathrm{c}_{\mathrm{in}} c in c o u t \mathrm{c}_{\mathrm{out}} c out c s k i p \mathrm{c}_{\mathrm{skip}} c skip σ d a t a = 1 \sigma_{\mathrm{data}} = 1 σ data = 1
c i n ( ε ) = 1 ε + 1 , c o u t ( ε ) = ε ε + 1 , c s k i p ( ε ) = 1 ε + 1 (46) \mathrm{c}_{\mathrm{in}}(\varepsilon) = \frac{1}{\sqrt{\varepsilon+ 1}}, \qquad\mathrm{c}_{\mathrm{out}}(\varepsilon) = \sqrt{\frac{\varepsilon}{\varepsilon+ 1}}, \qquad\mathrm{c}_{\mathrm{skip}}(\varepsilon) = \frac{1}{\varepsilon+ 1} \tag*{(46)} c in ( ε ) = ε + 1 1 , c out ( ε ) = ε + 1 ε , c skip ( ε ) = ε + 1 1 (46) The network F ϕ F_{\phi} F ϕ R n \mathbb{R}^{n} R n w w w L L L w w w R n \mathbb{R}^{n} R n ε \varepsilon ε e ( ε ) ∈ R c e(\varepsilon) \in\mathbb{R}^{c} e ( ε ) ∈ R c log ε \log\varepsilon log ε 1000 − k / 32 1000^{-k/32} 100 0 − k /32 k = 0 , … , 31 k = 0, \ldots, 31 k = 0 , … , 31 [43 ] ),
h ↦ h + W 2 SiLU ( ( 1 + a ) ⊙ SiLU ( W 1 h ) + b ) , ( a , b ) = W c SiLU ( e ( ε ) ) (47) \mathbf{h} \mapsto\mathbf{h} + W_{2}\operatorname{SiLU}\left((1+\mathbf{a}) \odot\operatorname{SiLU}(W_{1}\mathbf{h})+\mathbf{b}\right), \qquad(\mathbf{a},\mathbf{b}) = W_{c}\operatorname{SiLU}(e(\varepsilon)) \tag*{(47)} h ↦ h + W 2 SiLU ( ( 1 + a ) ⊙ SiLU ( W 1 h ) + b ) , ( a , b ) = W c SiLU ( e ( ε )) (47) where W 2 W_{2} W 2 w = 512 w = 512 w = 512 L = 8 L = 8 L = 8 c = 128 c = 128 c = 128 Table 3 for the denoiser and cometric architectures and training details.
Table 3. Denoiser and cometric architectures and training on alanine dipeptide.
Denoiser training We train the denoiser on the training split with the denoising loss of Karras et al. [27 ] , which weights the squared error by c o u t ( ε ) − 2 \mathrm{c}_{\mathrm{out}}(\varepsilon)^{-2} c out ( ε ) − 2
L ( ϕ ) = E ( x , x ε ) ∼ p 0 , ε [ c o u t ( ε ) − 2 ∥ D ϕ ( x ε , ε ) − x ∥ 2 2 ] (48) \mathcal{L}(\phi) = \mathbb{E}_{(x,x_{\varepsilon}) \sim p_{0,\varepsilon}}\left[\mathrm{c}_{\mathrm{out}}(\varepsilon)^{-2}\left\|D_{\phi}(x_{\varepsilon},\varepsilon)-x\right\|_{2}^{2}\right] \tag*{(48)} L ( ϕ ) = E ( x , x ε ) ∼ p 0 , ε [ c out ( ε ) − 2 ∥ D ϕ ( x ε , ε ) − x ∥ 2 2 ] (48) We draw ε \varepsilon ε [ 0.01 , 5 ] [0.01, 5] [ 0.01 , 5 ] { 5 × 10 − 5 , 10 − 4 , 2 × 10 − 4 , 4 × 10 − 4 , 8 × 10 − 4 } \{5 \times10^{-5}, 10^{-4}, 2 \times10^{-4}, 4 \times10^{-4}, 8 \times10^{-4}\} { 5 × 1 0 − 5 , 1 0 − 4 , 2 × 1 0 − 4 , 4 × 1 0 − 4 , 8 × 1 0 − 4 }
Cometric architecture The cometric shares the exact same architecture as the denoiser, except that the last layer maps to a larger vector of size R d × n \mathbb{R}^{d \times n} R d × n M θ ε M_{\theta}^{\varepsilon} M θ ε c i n \mathrm{c}_{\mathrm{in}} c in c o u t \mathrm{c}_{\mathrm{out}} c out c s k i p \mathrm{c}_{\mathrm{skip}} c skip M θ ε ( x ε ) M_{\theta}^{\varepsilon}(x_{\varepsilon}) M θ ε ( x ε ) x ε x_{\varepsilon} x ε d = n d = n d = n δ = 1 × 10 − 4 \delta= 1 \times10^{-4} δ = 1 × 1 0 − 4
Cometric training We train the cometric on the training data with the objective from Section C.4.1, centered on the selected denoiser. The optimization procedure is identical to how we trained the denoiser.
Experiment details
Here we give the details of our experiments.
Timing details We measure wall clock time for all of the experiments in this paper in the same way. This method takes a compiled JAX program as input, with its inputs already in device memory, so that compilation time and other preprocessing steps are not included in the reported time. The program is run once as a warmup, and then 5 times after, from which we report the median time. All timings are reported on one NVIDIA A16 in float32 and every program is implemented in JAX (Bradbury et al. [5 ] ) and compiled by XLA with its default settings.
Logarithmic map experiment details
We generate data on each manifold following the procedure in Section C.1. For each manifold, we generate 2048 basepoints for training the amortized models, 64 for validation and 128 for testing.
For each base point, we sample 1024 associated endpoints to train the MFM and flow chart with. These methods are then evaluated on endpoints of the 128 test base points. All methods use the same test endpoints. The closed form manifolds provide exact references. The data geometries use the reference solutions described in Section D.2. After reference filtering, the two moons has queries at 124 test base points and the other geometries have queries at all 128. We report the mean of R logmap ( x 0 ) \mathcal{R}_{\operatorname{logmap}}(x_0) R logmap ( x 0 )
Methods
Our baselines employ geodesic interpolation to retrieve normal coordinate and logmap estimates. These methods parameterize a curve γ ( ⋅ ) ϕ : [ 0 , 1 ] → M x \gamma^{\phi}_{(\cdot)} : [0,1] \to\mathcal{M}_x γ ( ⋅ ) ϕ : [ 0 , 1 ] → M x x 0 x_0 x 0 x \mathbf{x} x ϵ s t o p = 1 × 10 − 5 \epsilon_{\mathrm{stop}} = 1 \times10^{-5} ϵ stop = 1 × 1 0 − 5
GP We implement the fixed-point solver of Arvanitidis et al. [2 ] , which represents a curve by a Gaussian-process posterior mean conditioned on its endpoints and K K K λ = 1 / ( 2 ( K − 1 ) ) \lambda= \sqrt{1/(2(K-1))} λ = 1/ ( 2 ( K − 1 )) 10 − 2 10^{-2} 1 0 − 2
GEORCE We also use the GEORCE Rygaard and Hauberg, [47 ] algorithm to solve for a geodesic interpolation. This approach parameterizes the geodesic with K K K K K K [37 ] for evaluation. The initial velocity of GEORCE is the forward difference ( K + 1 ) ( y 1 − x 0 ) (K+1)(\mathbf{y}_1-\mathbf{x}_0) ( K + 1 ) ( y 1 − x 0 ) y 0 = x 0 , y 1 , … , y K , y K + 1 = x \mathbf{y}_0=\mathbf{x}_0,\mathbf{y}_1,\ldots,\mathbf{y}_K,\mathbf{y}_{K+1}=\mathbf{x} y 0 = x 0 , y 1 , … , y K , y K + 1 = x t k = k / ( K + 1 ) t_k=k/(K+1) t k = k / ( K + 1 )
Metric flow matching (MFM) Our neural network baseline is MFM [25 ] . This method parameterizes each curve as follows:
x c u r v e , t θ = ( 1 − t ) x 0 + t x 1 + t ( 1 − t ) N N θ ( t , x 0 , x 1 ) (49) \mathbf{x}^{\theta}_{\mathrm{curve},t}=(1-t)\mathbf{x}_0+t\mathbf{x}_1+t(1-t)\mathrm{NN}_{\theta}(t,\mathbf{x}_0,\mathbf{x}_1) \tag*{(49)} x curve , t θ = ( 1 − t ) x 0 + t x 1 + t ( 1 − t ) NN θ ( t , x 0 , x 1 ) (49) where N N θ \mathrm{NN}_{\theta} NN θ N N θ \mathrm{NN}_{\theta} NN θ E E E
We choose the number of blocks and hidden size separately for two MFM baselines. The first has a parameter count modestly close to that of the four-layer flow chart, and the second has a comparable time to compute the fast normal coordinate estimates z ( i ) \mathbf{z}^{(i)} z ( i ) h x 0 θ ( x ( i ) ) h_{\mathbf{x}_0}^{\theta}(\mathbf{x}^{(i)}) h x 0 θ ( x ( i ) ) J x 0 − 1 x ˙ 0 ( i ) J_{\mathbf{x}_0}^{-1}\dot{\mathbf{x}}_0^{(i)} J x 0 − 1 x ˙ 0 ( i ) Table 4 lists the model sizes and computation times. Times are microseconds per query on an NVIDIA A16 in FP32.
Table 4. Sizes of the amortized models and the wall clock cost of computing 1024 normal coordinate estimates at a fixed base point.
Flow chart The flow chart used for the amortized MFM comparison and the preconditioner is a conditional neural spline flow [14 ] with 4 rational quadratic spline coupling layers, each with 8 bins on [ − 10 , 10 ] [-10,10] [ − 10 , 10 ] [43 ] residual network with three residual blocks of width 32. The conditioning is computed with a separate residual network that encodes the base point into a 64-dimensional feature vector for all coupling layers with 8 gated residual blocks of width 64. The fixed-base-point and runtime experiments use the wider coupling networks described in Section D.3. We train our method with the normal loss in Proposition 3 under the sampling distribution of § 4.2 , with γ = 0.1 r max / n \gamma=0.1\,r_{\max}/\sqrt{n} γ = 0.1 r m a x / n
We minimize this loss with Adam in single precision at a learning rate of 3 × 10 − 4 3\times10^{-4} 3 × 1 0 − 4 t ∼ U ( 0 , 1 ) t\sim U(0,1) t ∼ U ( 0 , 1 ) z t \mathbf{z}_t z t § 4.2 . Gradient updates are clipped to have unit norm. We keep an exponential moving average of the parameters with decay 0.999, and this average is the model we evaluate and report. The validation loss is evaluated on held out pairs every 500 steps. The learning rate is constant until step 1,000,000 and then decays linearly to zero at step 2,000,000, where training ends, and we keep the average with the lowest validation loss.
Evaluation metrics
Recall that each method we test takes as input a base point x 0 x_0 x 0 { x ( 1 ) , … , x ( N ) } \{x^{(1)}, \ldots, x^{(N)}\} { x ( 1 ) , … , x ( N ) } { z ( 1 ) , … , z ( N ) } \{z^{(1)}, \ldots, z^{(N)}\} { z ( 1 ) , … , z ( N ) } { U x ( 1 ) , … , U x ( N ) } \{U_x^{(1)}, \ldots, U_x^{(N)}\} { U x ( 1 ) , … , U x ( N ) } x ( i ) x^{(i)} x ( i ) x ( i ) x^{(i)} x ( i ) t = 1 t = 1 t = 1
Method solutions The fast logmap estimate is simply the initial velocity of the curve, and the corresponding normal coordinate is the logmap scaled by a g g g J x 0 = g x 0 ( x 0 ) − 1 2 J_{x_0} = g_{x_0}(x_0)^{-\frac{1}{2}} J x 0 = g x 0 ( x 0 ) − 2 1 x 0 x_0 x 0 x ( i ) x^{(i)} x ( i )
U x ( i ) = x ˙ 0 ( i ) , z ( i ) = J x 0 − 1 x ˙ 0 ( i ) (50) U_x^{(i)} = \dot{x}_0^{(i)}, \qquad z^{(i)} = J_{x_0}^{-1}\dot{x}_0^{(i)} \tag*{(50)} U x ( i ) = x ˙ 0 ( i ) , z ( i ) = J x 0 − 1 x ˙ 0 ( i ) (50) For the flow chart, whose curve is x t ( i ) = f x 0 θ ( t h x 0 θ ( x ( i ) ) ) x_t^{(i)} = f_{x_0}^{\theta}(t h_{x_0}^{\theta}(x^{(i)})) x t ( i ) = f x 0 θ ( t h x 0 θ ( x ( i ) )) z ( i ) = h x 0 θ ( x ( i ) ) z^{(i)} = h_{x_0}^{\theta}(x^{(i)}) z ( i ) = h x 0 θ ( x ( i ) ) U x ( i ) = D f x 0 θ ( 0 ) z ( i ) U_x^{(i)} = Df_{x_0}^{\theta}(0)z^{(i)} U x ( i ) = D f x 0 θ ( 0 ) z ( i ) § 3 .
The GP and GEORCE methods fail if the solution is not finite, or the solver stops before meeting its stopping criterion within the maximum number of iterations. We write I ( x 0 ) ⊂ { 1 , … , N } \mathcal{I}(x_0) \subset\{1, \ldots, N\} I ( x 0 ) ⊂ { 1 , … , N } x 0 x_0 x 0 J ( x 0 ) \mathcal{J}(x_0) J ( x 0 )
Evaluation metrics We evaluate the solutions against the ground truth coordinates and logarithmic maps, { s ( 1 ) , … , s ( N ) } \{s^{(1)}, \ldots, s^{(N)}\} { s ( 1 ) , … , s ( N ) } { V x ( 1 ) , … , V x ( N ) } \{V_x^{(1)}, \ldots, V_x^{(N)}\} { V x ( 1 ) , … , V x ( N ) }
Wallclock Since we are concerned with speed, we time how long it takes (in wall clock time) to compute the solutions, as measured by Section D. We let τ ( x 0 ) \tau(x_0) τ ( x 0 ) z ( i ) z^{(i)} z ( i ) U x ( i ) U_x^{(i)} U x ( i ) i ∈ I ( x 0 ) i \in\mathcal{I}(x_0) i ∈ I ( x 0 ) ∣ I ( x 0 ) ∣ |\mathcal{I}(x_0)| ∣ I ( x 0 ) ∣ τ ( x 0 ) \tau(x_0) τ ( x 0 ) I ( x 0 ) \mathcal{I}(x_0) I ( x 0 ) N N N x 0 x_0 x 0 I ( x 0 ) \mathcal{I}(x_0) I ( x 0 )
For the GEORCE and GP baselines, τ ( x 0 ) \tau(x_0) τ ( x 0 ) τ ( x 0 ) \tau(x_0) τ ( x 0 )
Because a batch of queries iterates until its slowest curve stops, we sweep over batch sizes in { 64 , 128 , 256 , 512 , 1024 } \{64, 128, 256, 512, 1024\} { 64 , 128 , 256 , 512 , 1024 } τ ( x 0 ) \tau(x_0) τ ( x 0 ) Table 5 lists the batch sizes selected for the GP and GEORCE.
Table 5. Batch sizes selected for the GP and GEORCE at each knot count K K K τ ( x 0 ) \tau(x_0) τ ( x 0 )
We report the mean of τ ( x 0 ) \tau(x_0) τ ( x 0 )
Logmap To evaluate the logmap solutions, we compute the normalized RMSE R l o g m a p ( x 0 ) \mathcal{R}_{\mathrm{logmap}}(x_0) R logmap ( x 0 ) § 5.1 . We compute the root mean squared g g g
E l o g m a p ( x 0 ) = ( 1 ∣ I ( x 0 ) ∣ ∑ i ∈ I ( x 0 ) ∥ U x ( i ) − V x ( i ) ∥ g 2 ) 1 2 (51) \mathcal{E}_{\mathrm{logmap}}(x_0) = \left(\frac{1}{|\mathcal{I}(x_0)|}\sum_{i \in\mathcal{I}(x_0)} \left\|U_x^{(i)} - V_x^{(i)}\right\|_g^2\right)^{\frac{1}{2}} \tag*{(51)} E logmap ( x 0 ) = ∣ I ( x 0 ) ∣ 1 i ∈ I ( x 0 ) ∑ U x ( i ) − V x ( i ) g 2 2 1 (51) We divide E l o g m a p ( x 0 ) \mathcal{E}_{\mathrm{logmap}}(x_0) E logmap ( x 0 ) r r e f ( x 0 ) r_{\mathrm{ref}}(x_0) r ref ( x 0 ) J ( x 0 ) \mathcal{J}(x_0) J ( x 0 ) R l o g m a p ( x 0 ) \mathcal{R}_{\mathrm{logmap}}(x_0) R logmap ( x 0 )
Accuracy comparisons
Amortized accuracy In this experiment, we pretrain a flow chart and the two MFMs described in Table 4 on the training set of each geometry, and then evaluate them on the test set. Each is trained once, from one seed, so the error bars measure variation across base points. For the baselines, we evaluate them directly on the test set. We sweep the knot count of the GP and GEORCE, over K ∈ { 4 , 8 , 16 , 32 } K \in\{4, 8, 16, 32\} K ∈ { 4 , 8 , 16 , 32 }
To estimate the uncertainty in the ratios of Figure 4 , we resample the held out base points with replacement 10,000 times, using the same base points for the flow chart and MFM in each sample. We compute the ratio of mean errors for each sample and report the 2.5th and 97.5th percentiles.
Fixed base point accuracy We evaluate the numerical baselines in Figure 3 at the same four base points as the non-amortized flow charts, with indices 26, 44, 67 and 100. For each baseline, we select the knot count with the lowest mean of R l o g m a p R_{\mathrm{logmap}} R logmap K = 16 K=16 K = 16 K = 32 K=32 K = 32
Convergence of the baselines Table 6 shows the number of test queries that fail, for the GP and GEORCE at each knot count over all test queries of all test base points.
Table 6. Number of test queries that fail for the GP and GEORCE at each knot count K K K
Exponential map experiment details
Data and geometry
We use the twelve closed form manifolds of Section C.1, three geometries constructed from synthetic datasets using the MFM RBF metric, and the diffusion geometry of alanine dipeptide learned with RMM as described in Section B. After pretraining each RMM geometry, we fix ε = 0.1 \varepsilon= 0.1 ε = 0.1 g 0.1 , x θ g^{\theta}_{0.1,x} g 0.1 , x θ
Synthetic datasets The synthetic datasets have 500,000 points total, and are split into a training, validation and test set with a ( 70 , 10 , 20 ) (70,10,20) ( 70 , 10 , 20 )
Data selection To select endpoints, we sample from each base point based on proximity. For this, we first run clustering on each dataset (except for alanine dipeptide due to its size) and give each point a label corresponding to its connected component. This clustering is done with the HDBSCAN (Campello et al., [8 ] ; McInnes et al., [35 ] ) algorithm from Scikit-learn (Pedregosa et al., [42 ] ). We run it separately on each split, with a minimum cluster size of 1% of the split and 16 samples per core point, and allow a single cluster. Points that are labeled as noise are given the label of their nearest neighbor with a label. We give all points in alanine dipeptide the same cluster label.
Then, given a base point x 0 x_0 x 0 x x x k = ⌊ ν ⋅ N ⌋ k = \lfloor\nu\cdot N \rfloor k = ⌊ ν ⋅ N ⌋ x 0 x_0 x 0 N N N ν = 0.1 \nu= 0.1 ν = 0.1
Flow chart architecture and training details
For the MFM comparison and the flow chart preconditioner, we use the four-layer architecture of Section D.1.1. We train these flow charts for at most 2,000,000 steps, with early stopping on the data geometries. On the data geometries, the learning rate remains constant after warmup, and training stops after 100 validation checks (50,000 steps) without improvement. For the capacity sweep, we use the training procedure of Section D.3, without early stopping. We train with a batch containing 16 base points, with 32 endpoints for each. During evaluation, we measure the average loss over the base points of the validation set, each with a single endpoint drawn from the same procedure with a fixed random seed.
We set γ = 0.1 r ˉ / n \gamma= 0.1 \bar{r}/\sqrt{n} γ = 0.1 r ˉ / n r ˉ \bar{r} r ˉ k k k
Evaluation set generation
This experiment evaluates different methods on triplets ( x 0 , V x , x 1 ) (x_0, V_x, x_1) ( x 0 , V x , x 1 ) V x ≈ log x 0 ( x 1 ) V_x \approx\log_{x_0}(x_1) V x ≈ log x 0 ( x 1 ) B = 128 B = 128 B = 128 M = 512 M = 512 M = 512 M ε θ M^\theta_{\varepsilon} M ε θ K = 50 K = 50 K = 50 1 / ( 2 ( K + 1 ) ) \sqrt{1/(2(K+1))} 1/ ( 2 ( K + 1 )) 10 − 4 10^{-4} 1 0 − 4
Since the norm of the logmap solution is supposed to equal the geodesic distance between the base point and the query, the logmap solution can be refined by estimating the length and rescaling. Let L ( i ) = ∫ 0 1 ∥ x ˙ t ( i ) ∥ g d t L^{(i)} = \int_0^1 \lVert\dot{x}^{(i)}_t \rVert_g \,\mathrm{d}t L ( i ) = ∫ 0 1 ∥ x ˙ t ( i ) ∥ g d t
V x ( i ) = L T x ˙ 0 ( i ) ∥ x ˙ 0 ( i ) ∥ g (52) V_x^{(i)} = L^\mathsf{T} \frac{\dot{x}_0^{(i)}}{\lVert\dot{x}_0^{(i)} \rVert_g} \tag*{(52)} V x ( i ) = L T ∥ x ˙ 0 ( i ) ∥ g x ˙ 0 ( i ) (52) Once the curve is fit, we save this refined logmap solution as V x V_x V x h h h h ≥ 0.5 h \ge0.5 h ≥ 0.5 10 − 4 10^{-4} 1 0 − 4 Table 7 lists the number of retained test triplets. On these geometries, the errors measure agreement with the numerical reference, so we compare the GP baseline with the same solver using more acceleration nodes.
Table 7. Sizes of the training and test sets. On the synthetic and bio datasets, every point of the training set is a base point and its query points are its ⌊ 0.1 B t r a i n ⌋ \lfloor0.1 B_{\mathrm{train}}\rfloor ⌊ 0.1 B train ⌋ B t e s t N t e s t B_{\mathrm{test}}N_{\mathrm{test}} B test N test
Efficient acceleration computation for GP on RMM cometrics For metric parameterized via Section C.4.1, solving for geodesic interpolations for the GP is expensive because every acceleration evaluation differentiates through a cometric solve. However, by taking advantage of the structure of the metric, we can avoid derivative through the inverse. At each of its nodes, the GP evaluates the geodesic acceleration x ¨ = − Γ x ( x ) ( x ˙ , x ˙ ) \ddot{x} = -\Gamma_x(x)(\dot{x}, \dot{x}) x ¨ = − Γ x ( x ) ( x ˙ , x ˙ ) x x x x ˙ \dot{x} x ˙ (5) , this differentiates the metric g ε , x θ ( x ) g^{\theta}_{\varepsilon,x}(x) g ε , x θ ( x )
Instead, we compute the acceleration from Hamilton’s equations. Let p = g ε , x θ ( x ) x ˙ p = g^{\theta}_{\varepsilon,x}(x)\dot{x} p = g ε , x θ ( x ) x ˙ H ( x , p ) = 1 2 p T g ε , x θ ( x ) − 1 p H(x,p) = \frac{1}{2}p^\mathsf{T}g^{\theta}_{\varepsilon,x}(x)^{-1}p H ( x , p ) = 2 1 p T g ε , x θ ( x ) − 1 p
x ˙ = δ p + ( 1 − δ ) M ε θ ( x ) T M ε θ ( x ) p , p ˙ = − 1 − δ 2 ∇ x ∥ M ε θ ( x ) p ∥ 2 2 (53) \dot{x} = \delta p + (1-\delta)M^\theta_{\varepsilon}(x)^\mathsf{T}M^\theta_{\varepsilon}(x)p,\qquad\dot{p} = -\frac{1-\delta}{2}\nabla_x\lVert M^\theta_{\varepsilon}(x)p\rVert_2^2 \tag*{(53)} x ˙ = δ p + ( 1 − δ ) M ε θ ( x ) T M ε θ ( x ) p , p ˙ = − 2 1 − δ ∇ x ∥ M ε θ ( x ) p ∥ 2 2 (53) where no gradient flows through p p p x ˙ \dot{x} x ˙
x ¨ = ( 1 − δ ) ( D x ˙ M ε θ ( x ) T M ε θ ( x ) + M ε θ ( x ) T D x ˙ M ε θ ( x ) ) p + g ε , x θ ( x ) − 1 p ˙ (54) \ddot{x} = (1-\delta)\left(D_{\dot{x}}M^\theta_{\varepsilon}(x)^\mathsf{T}M^\theta_{\varepsilon}(x) + M^\theta_{\varepsilon}(x)^\mathsf{T}D_{\dot{x}}M^\theta_{\varepsilon}(x)\right)p + g^\theta_{\varepsilon,x}(x)^{-1}\dot{p} \tag*{(54)} x ¨ = ( 1 − δ ) ( D x ˙ M ε θ ( x ) T M ε θ ( x ) + M ε θ ( x ) T D x ˙ M ε θ ( x ) ) p + g ε , x θ ( x ) − 1 p ˙ (54) where D x ˙ M ε θ ( x ) D_{\dot{x}}M^\theta_{\varepsilon}(x) D x ˙ M ε θ ( x ) M ε θ ( x ) M^\theta_{\varepsilon}(x) M ε θ ( x ) x x x x ˙ \dot{x} x ˙ p p p
Methods
The baselines simulate the geodesic with initial position x 0 x_0 x 0 V x V_x V x H ( x , p ) = 1 2 p T g x ( x ) − 1 p H(x,p) = \frac{1}{2}p^\mathsf{T}g_x(x)^{-1}p H ( x , p ) = 2 1 p T g x ( x ) − 1 p
x ˙ t = g x t ( x t ) − 1 p t , p ˙ t = − 1 2 ∇ x ( p t T g x ( x t ) − 1 p t ) , p 0 = g x ( x 0 ) V x (55) \dot{x}_t = g_{x_t}(x_t)^{-1}p_t,\qquad\dot{p}_t = -\frac{1}{2}\nabla_x\left(p_t^\mathsf{T}g_{x}(x_t)^{-1}p_t\right),\qquad p_0 = g_{x}(x_0)V_x \tag*{(55)} x ˙ t = g x t ( x t ) − 1 p t , p ˙ t = − 2 1 ∇ x ( p t T g x ( x t ) − 1 p t ) , p 0 = g x ( x 0 ) V x (55) from t = 0 t=0 t = 0 t = 1 t=1 t = 1 x 1 = exp x 0 ( V x ) x_1 = \exp_{x_0}(V_x) x 1 = exp x 0 ( V x ) p 0 p_0 p 0 [29 ] of increasing order, Euler, Heun, Tsit5 [50 ] , and Dopri8 [44 ] , which cost one, two, six, and thirteen evaluations of Equation (55) per step, plus one more at the first step for Tsit5 and Dopri8 and one evaluation of the cometric at x 0 x_0 x 0 p 0 p_0 p 0 S S S [ 0 , 1 ] [0,1] [ 0 , 1 ] c c c S S S ⌈ 128 / c ⌉ \lceil128/c\rceil ⌈ 128/ c ⌉ S ∈ { 1 , 3 , 11 , 38 , 128 } S \in\{1, 3, 11, 38, 128\} S ∈ { 1 , 3 , 11 , 38 , 128 } { 1 , 3 , 8 , 23 , 64 } \{1, 3, 8, 23, 64\} { 1 , 3 , 8 , 23 , 64 } { 1 , 2 , 5 , 10 , 21 } \{1, 2, 5, 10, 21\} { 1 , 2 , 5 , 10 , 21 } { 1 , 2 , 3 , 5 , 9 } \{1, 2, 3, 5, 9\} { 1 , 2 , 3 , 5 , 9 }
The flow chart computes exp x 0 ( V x ) = f x 0 θ ( D h x 0 θ ( x 0 ) V x ) \exp_{x_0}(V_x) = f^\theta_{x_0}(D h^\theta_{x_0}(x_0)V_x) exp x 0 ( V x ) = f x 0 θ ( D h x 0 θ ( x 0 ) V x ) § 3 .
Evaluation metrics
For each method, we compute E expmap ( x 0 ) \mathcal{E}_{\operatorname{expmap}}(x_0) E expmap ( x 0 ) R expmap ( x 0 ) \mathcal{R}_{\operatorname{expmap}}(x_0) R expmap ( x 0 ) § 5.1 . The numerator uses the queries in I ( x 0 ) \mathcal{I}(x_0) I ( x 0 ) r chord ( x 0 ) r_{\operatorname{chord}}(x_0) r chord ( x 0 ) J ( x 0 ) \mathcal{J}(x_0) J ( x 0 ) R expmap ( x 0 ) \mathcal{R}_{\operatorname{expmap}}(x_0) R expmap ( x 0 )
We time every method with the procedure of Section D. Let τ ( x 0 ) \tau(x_0) τ ( x 0 ) x 0 x_0 x 0 τ ( x 0 ) \tau(x_0) τ ( x 0 ) τ ( x 0 ) \tau(x_0) τ ( x 0 ) x 0 x_0 x 0 τ ( x 0 ) \tau(x_0) τ ( x 0 ) τ \tau τ τ ( x 0 ) \tau(x_0) τ ( x 0 )
Capacity and fixed base point experiments
We vary the capacity of the flow chart and compare amortized and non-amortized training on the closed form manifolds, the three synthetic datasets, and alanine dipeptide. We use both the low and high dimensional closed form manifolds, where the logarithmic and exponential maps are known in closed form. On the data geometries, we use the numerical references of Section D.2. See Section C.1, Section C.2 and Section B for a description of each dataset, the geometry they define, and how the training, validation and test sets are generated. Each dataset produces a set of B B B N N N ( x 0 , V x , x ) (x_0, V_x, x) ( x 0 , V x , x ) V x = log x 0 ( x ) V_x = \log_{x_0}(x) V x = log x 0 ( x ) x = exp x 0 ( V x ) x = \exp_{x_0}(V_x) x = exp x 0 ( V x )
We train flow charts on this data in two different settings. The first is an amortized setting over base points, where we train on B t r a i n B_{\mathrm{train}} B train N t r a i n N_{\mathrm{train}} N train R logmap \mathcal{R}_{\operatorname{logmap}} R logmap R expmap \mathcal{R}_{\operatorname{expmap}} R expmap B t e s t B_{\mathrm{test}} B test N t e s t N_{\mathrm{test}} N test
The second is a non-amortized setting where we train a flow chart at a single test base point. These B s i n g l e = 4 B_{\mathrm{single}} = 4 B single = 4 B t e s t B_{\mathrm{test}} B test N s i n g l e N_{\mathrm{single}} N single N t e s t N_{\mathrm{test}} N test N s i n g l e N_{\mathrm{single}} N single N s i n g l e N_{\mathrm{single}} N single Table 7 lists N s i n g l e N_{\mathrm{single}} N single
For Figure 3 , we compute R logmap \mathcal{R}_{\operatorname{logmap}} R logmap
Flow-chart architecture
In all of these experiments, we use a modified version of the conditional neural spline flow of Section D.1.1. We keep the conditioned elementwise scale and shift and the invertible linear map that are placed before each coupling layer, and change the coupling layer and the width of its conditioner. Each coupling layer transforms an input v v v s ⋅ spline ( v ) + t s \cdot\operatorname{spline}(v) + t s ⋅ spline ( v ) + t [ − 10 , 10 ] [-10, 10] [ − 10 , 10 ] s s s t t t s s s
In the amortized setting, the base point encoder is a residual network with 8 gated residual blocks of width 64 that encodes the base point into a 64-dimensional feature vector for all coupling layers. In the non-amortized setting, the base point is fixed, so the base point encoder has a single gated residual block. We vary the capacity of the amortized flow chart through the number of coupling layers, L ∈ { 4 , 8 , 16 , 32 } L \in\{4, 8, 16, 32\} L ∈ { 4 , 8 , 16 , 32 } L = 32 L = 32 L = 32
Training
We train every flow chart with the procedure of Section D.1.1, with a budget of 2,000,000 steps, learning rate decay and no early stopping. All non-amortized flow charts reported here and all amortized flow charts with L = 32 L = 32 L = 32 L = 4 L = 4 L = 4 L = 8 L = 8 L = 8 ε = 0.1 \varepsilon= 0.1 ε = 0.1 γ = 0.1 r max / n \gamma= 0.1\,r_{\max}/\sqrt{n} γ = 0.1 r m a x / n γ = 0.1 r ˉ / n \gamma= 0.1\,\bar{r}/\sqrt{n} γ = 0.1 r ˉ / n N s i n g l e N_{\mathrm{single}} N single ⌊ 0.1 N v a l ⌋ \lfloor0.1\,N_{\mathrm{val}}\rfloor ⌊ 0.1 N val ⌋ N v a l N_{\mathrm{val}} N val
Evaluation
We evaluate every flow chart in the same way as the flow charts of the experiments above. On every geometry, we compute the mean of R l o g m a p \mathcal{R}_{\mathrm{logmap}} R logmap R e x p m a p \mathcal{R}_{\mathrm{expmap}} R expmap s = g x ( x 0 ) 1 / 2 v ⃗ x \mathbf{s}=g_x(x_0)^{1/2}\vec{v}_x s = g x ( x 0 ) 1/2 v x
In the amortized setting, we evaluate the flow chart at all B t e s t B_{\mathrm{test}} B test B s i n g l e B_{\mathrm{single}} B single B s i n g l e B_{\mathrm{single}} B single
Baselines
Closed form manifolds ODE Solvers We run the ODE solvers of Section D.2 on the closed form manifolds in both low and high dimensions. All N t e s t N_{\mathrm{test}} N test
Closed form manifolds MFM and GP, GEORCE See Section D.1.
Data manifolds ODE Solvers See Section D.2.
Data manifolds MFM and GP, GEORCE. We run the methods of Section D.1.1 on the synthetic and bio datasets, and evaluate them as in Section D.1.2 on the test triplets of Section D.2 that were not discarded. The batch sizes of the GP and GEORCE are selected from { 64 , 128 , 256 , 512 } \{64, 128, 256, 512\} { 64 , 128 , 256 , 512 }
The two MFMs are matched to the flow chart of Section D.2 on each dataset, with the time measured on the 512 queries of a test base point, and Table 8 lists their sizes. They are trained on the same training pairs as that flow chart, with the same training procedure as Section D.1. We train these MFMs for 2,000,000 steps with learning rate decay and no early stopping.
Table 8. Sizes of the MFM baselines on the synthetic and bio datasets.
Experiments
This experiment measures how the errors of a flow chart depend on its capacity, and how much of the error is due to amortization. For every dataset of Table 7 , we use an amortized flow chart at each L ∈ { 4 , 8 , 16 , 32 } L \in\{4, 8, 16, 32\} L ∈ { 4 , 8 , 16 , 32 } B s i n g l e = 4 B_{\mathrm{single}} = 4 B single = 4 L = 32 L = 32 L = 32 L L L L L L
We compare against the baselines of the experiments above, evaluated on the same test triplets. For the logarithmic map, these are the GP, GEORCE and the two MFMs of Section D.1.1, and for the exponential map, the ODE solvers of Section D.2. The two MFMs are selected to have the same parameter count and time as the flow chart used in the original experiment, and not the size of the ablation models.
Table 9 contains the parameter count of the amortized flow chart for each dimension and L L L
Table 9. Parameter counts of the amortized flow charts used in the capacity sweep for each number of coupling layers L L L
Scaling with dimension
Figure 7 shows that the cost of a pretrained flow chart also scales favorably with dimension compared with the numerical baselines. Across the closed form manifolds that have low and high dimensional variants, the smallest and largest chart sizes only take 1.30 to 1.63 times longer to run in high dimensions while the GP grows by 3.0 to 52× and GEORCE by 3.8 to 27× (at K = 16 K = 16 K = 16
Figure 7. Increase in time per logmap query from the low dimensional to the high dimensional version of each manifold. Each method uses the same architecture or knot count in both dimensions.
Matching and solver resolution
To compare runtimes at matched accuracy, we set the error threshold R ⋆ R_\star R ⋆ R l o g m a p R_{\mathrm{logmap}} R logmap R e x p m a p R_{\mathrm{expmap}} R expmap R ⋆ R_\star R ⋆ L = 32 L = 32 L = 32 Table 10 lists the settings selected for the comparisons in Figure 5 .
Table 10: Settings selected for the runtime comparisons in Figure 5 . For each numerical method, L L L K K K S S S R ⋆ \mathcal{R}_{\star} R ⋆
Table 10. Settings selected for the runtime comparisons in Figure 5.
On the two dimensional sphere, several ODE settings have very large finite errors, and the errors do not consistently decrease as the number of steps increases. The matched comparison in Figure 5 uses Heun with 64 steps. These results show instability on this geometry, but we have not identified its cause. The conclusions about the exponential map in high dimensions do not depend on this geometry.
Exponential map errors and runtime budget
For the exponential map, the flow chart with L = 32 L = 32 L = 32 100 × R e x p m a p 100 \times\mathcal{R}_{\mathrm{expmap}} 100 × R expmap n ≥ 28 n \ge28 n ≥ 28 1.99 ± 0.29 1.99 \pm0.29 1.99 ± 0.29 L = 32 L = 32 L = 32 R e x p m a p \mathcal{R}_{\mathrm{expmap}} R expmap Figure 9 gives the errors and runtimes at every chart size and ODE step count.
Figure 9. Exponential map accuracy and runtime across all sixteen geometries. Each point gives the mean of 100 × R e x p m a p 100 \times\mathcal{R}_{\mathrm{expmap}} 100 × R expmap
Flow chart preconditioner
For Table 1 , we use the pretrained flow charts with four coupling layers described in Section D.1.1. For each base point, we express the metric in the chart coordinates using Equation (1) , encode each query, and run GEORCE between the origin and the encoded query. We map the initial velocity back to data coordinates through D f x 0 θ ( 0 ) D f^{\theta}_{x_0}(0) D f x 0 θ ( 0 ) R l o g m a p \mathcal{R}_{\mathrm{logmap}} R logmap
At K = 32 K = 32 K = 32 R l o g m a p \mathcal{R}_{\mathrm{logmap}} R logmap
Complete experimental results
We summarize the results over all geometries and tested settings in three figures. Figure 8 and Figure 9 show the accuracy and time per query for the logarithmic and exponential maps respectively. Figure 10 shows the effect of flow chart capacity and amortization.
Figure 8. Logarithmic map accuracy and runtime across all sixteen geometries. Each point gives the mean of 100 × \times × R l o g m a p R_{\mathrm{logmap}} R logmap
Figure 10. Effect of flow chart capacity and amortization. The panels show the mean of 100 × R logmap 100 \times\mathcal{R}_{\operatorname{logmap}} 100 × R logmap
