Learning Normal Coordinates on Riemannian Manifolds with Flow Charts

Edmond Cunningham

cs.LGOct 20, 2026

Learning Normal Coordinates on Riemannian Manifolds with Flow Charts

Authors: Edmond Cunningham

Organizations: University of Massachusetts Amherst

Abstract

Riemannian manifolds are powerful models for problems with non-Euclidean structure. Recent work in geometric machine learning has focused on training deep neural networks to learn Riemannian metrics from high-dimensional data. However, many geometric algorithms require repeated evaluations of the exponential and logarithmic maps, whose high computational cost restricts the utility of learned Riemannian manifolds in practice. We propose to address this cost by training a conditional invertible neural network, called a flow chart, to approximate normal coordinates around points on a given Riemannian manifold. With this approach, the exponential and logarithmic maps can be approximated by fast Euclidean operations in the model's learned coordinate space. Our main contribution is an efficient loss function for training flow charts that is nonnegative and, under suitable domain assumptions, is zero precisely when the chart is normal. Across sixteen synthetic and data-derived geometries, flow charts are competitive with a specialized neural geodesic interpolator for logarithmic map estimation, while also being able to estimate the exponential map. When pretrained, flow charts provide a median speedup of 60-70x over numerical baselines at comparable logarithmic map estimation error, and their learned coordinates improve the accuracy of a numerical geodesic solver.

Figures & tables

Figure 1: Illustration of normal coordinate charts and our method, using our notation in Section 3 and Section 4 . The purple and green symbols correspond to the same points

\mathbf{p}_{0}

and

\mathbf{p}

on the manifold in the top left, and the red curve is the geodesic between them. The tangent-space and coordinate plots depict the same blue region of the manifold in different numerical representations. The top row shows the standard approach to defining normal coordinates through the logarithmic map and the bottom shows our approach. Our method introduces flow charts and trains them to satisfy the identity in purple to learn normal coordinates.

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.

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

coupling layers, the GP uses

K=16

intermediate accelerations, and GEORCE uses

K=32

intermediate knot points. All methods achieve low error, and the individually fitted flow chart achieves error comparable to the GP and GEORCE baselines, with mean error below

0.0032

on every geometry. The pretrained chart has higher error than the individually fitted chart on six of the seven high-dimensional geometries.

Figure 4: Amortized logmap accuracy relative to MFM. Each bar denotes the mean

\mathcal{R}_{\mathrm{logmap}}

of the

L=4

flow chart or a parameter matched or time matched MFM at the same held out base points. Error bars indicate one standard error. In most settings the flow chart outperforms the baselines, but the parameter matched method performs best in the high dimensional synthetic geometries.

Figure 5: Comparison of time needed to achieve similar error between flow charts and numerical methods. For each comparison,

\mathcal{R}_{\star}

is the lowest error level that both methods can reach among the settings we tested. Each dot is the numerical method’s solve time divided by the chart’s solve time, using each method’s fastest setting with error at most

\mathcal{R}_{\star}

. Values above one favor the chart. Flow charts are faster than the numerical baselines, except for exponential-map computation on the low-dimensional data geometries, whose metrics are fast to evaluate.

Table 1: GEORCE in data coordinates and in the coordinates of a pretrained flow chart. Columns vary GEORCE’s knot count

K

. Each cell reports the mean of

100\times\mathcal{R}_{\mathrm{logmap}}

across base points in data coordinates followed by an arrow and the value in chart coordinates. The bold numbers above the arrows show the ratio between the two numbers. Solving in chart coordinates consistently lowers error.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Synthetic data generation. We sample in normal coordinates at the data center $\mathbf{p}_{\mu}$ (orange), express the same points in normal coordinates at the reference point $\mathbf{p}_{\mathrm{ref}}$ (blue), and apply $\Psi$ to obtain observed coordinates (purple). The learned chart is centered at a sampled base point $\mathbf{p}_{0}$ , whose observed coordinates are $\mathbf{x}_{0}=x(\mathbf{p}_{0})$ .

Residual network architecture
Property	Alanine dipeptide
Dimension $n$	$30$
Width $w$	$512$
Residual blocks $L$	$8$
Embedding size $c$	$128$
Total parameters (M)	$5.32$

Appendix

Table 3: Denoiser and cometric architectures and training on alanine dipeptide.

	Flow chart		MFM, parameter matched			MFM, time matched
$n$	Parameters	$\mu$ s	Width $\times$ blocks	Parameters	$\mu$ s	Width $\times$ blocks	Parameters	$\mu$ s
$2$	$187{,}504$	$0.98$	$64\times 15$	$188{,}674$	$1.51$	$64\times 9$	$113{,}410$	$1.01$
$3$	$191{,}296$	$1.00$	$64\times 15$	$188{,}867$	$1.71$	$64\times 6$	$75{,}971$	$0.97$
$4$	$195{,}104$	$1.05$	$96\times 7$	$197{,}572$	$1.39$	$64\times 6$	$76{,}164$	$1.01$
$28$	$291{,}296$	$1.33$	$96\times 10$	$288{,}604$	$2.05$	$64\times 6$	$80{,}796$	$1.32$
$30$	$299{,}728$	$1.35$	$128\times 6$	$309{,}790$	$1.58$	$64\times 9$	$118{,}814$	$1.30$

Appendix

Table 4: Sizes of the amortized models and the wall clock cost of computing 1024 normal coordinate estimates at a fixed base point.

		GP				GEORCE
Manifold	$n$	$K=4$	$8$	$16$	$32$	$K=4$	$8$	$16$	$32$
Sphere	$2$	$1024$	$1024$	$1024$	$1024$	$1024$	$1024$	$1024$	$512$
Flat torus	$2$	$1024$	$1024$	$1024$	$1024$	$1024$	$1024$	$1024$	$1024$
Hyperbolic	$2$	$1024$	$1024$	$1024$	$512$	$1024$	$1024$	$1024$	$1024$
SPD	$3$	$1024$	$1024$	$1024$	$512$	$1024$	$512$	$256$	$128$
Grassmann	$4$	$512$	$128$	$64$	$64$	$64$	$64$	$64$	$64$

Appendix

Table 5: Batch sizes selected for the GP and GEORCE at each knot count

K

, the size with the smallest mean of

\tau(\mathbf{x}_{0})

over the first

16

test base points.

		GP				GEORCE
Manifold	$n$	$K=4$	$8$	$16$	$32$	$K=4$	$8$	$16$	$32$
Sphere	$2$	$246$	$32$	$3$	$0$	$0$	$0$	$0$	$0$
Flat torus	$2$	$0$	$0$	$0$	$0$	$0$	$0$	$0$	$0$
Hyperbolic	$2$	$0$	$0$	$0$	$0$	$0$	$0$	$0$	$0$
SPD	$3$	$0$	$0$	$0$	$0$	$0$	$0$	$0$	$0$
Grassmann	$4$	$1$	$0$	$0$	$0$	$0$	$0$	$0$	$0$

Appendix

Table 6: Number of test queries that fail for the GP and GEORCE at each knot count

K

, out of

131{,}072

queries per geometry,

1{,}024

at each of the

128

test base points.

Dataset	$n$	$B_{\text{train}}$	$N_{\text{train}}$	$B_{\text{test}}$	$N_{\text{test}}$	$N_{\text{single}}$	Test triplets
Sphere	$2\,/\,30$	$2{,}048$	$1{,}024$	$128$	$1{,}024$	$262{,}144$	$131{,}072$
Flat torus	$2\,/\,30$	$2{,}048$	$1{,}024$	$128$	$1{,}024$	$262{,}144$	$131{,}072$
Hyperbolic	$2\,/\,30$	$2{,}048$	$1{,}024$	$128$	$1{,}024$	$262{,}144$	$131{,}072$
SPD	$3\,/\,28$	$2{,}048$	$1{,}024$	$128$	$1{,}024$	$262{,}144$	$131{,}072$
Grassmann	$4\,/\,30$	$2{,}048$	$1{,}024$	$128$	$1{,}024$	$262{,}144$	$131{,}072$
Rotations	$3\,/\,28$	$2{,}048$	$1{,}024$	$128$	$1{,}024$	$262{,}144$	$131{,}072$

Appendix

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

\lfloor 0.1\,B_{\text{train}}\rfloor

nearest neighbors with the same cluster label. The last column is the number of test triplets that remain out of

B_{\text{test}}N_{\text{test}}

after the discards of Section D.2 .

		Flow chart		MFM, parameter matched			MFM, time matched
Dataset	$n$	Parameters	$\mu$ s	Size	Parameters	$\mu$ s	Size	Parameters	$\mu$ s
Disk	$2$	$187{,}504$	$1.85$	$64\times 15$	$188{,}674$	$2.20$	$64\times 12$	$151{,}042$	$1.82$
Two moons	$2$	$187{,}504$	$1.92$	$64\times 15$	$188{,}674$	$2.23$	$64\times 12$	$151{,}042$	$1.87$
Nested circles	$2$	$187{,}504$	$1.89$	$64\times 15$	$188{,}674$	$2.22$	$64\times 12$	$151{,}042$	$1.87$
Alanine dipeptide	$30$	$305{,}008$	$2.34$	$128\times 6$	$309{,}790$	$2.74$	$64\times 7$	$93{,}726$	$2.29$

Appendix

Table 8: Sizes of the MFM baselines on the synthetic and bio datasets.

Datasets	$n$	$L=4$	$8$	$16$	$32$
Sphere, torus, hyperbolic	$2$	$319{,}232$	$533{,}760$	$962{,}816$	$1{,}820{,}928$
SPD, rotations	$3$	$326{,}616$	$548{,}464$	$992{,}160$	$1{,}879{,}552$
Grassmann	$4$	$334{,}016$	$563{,}200$	$1{,}021{,}568$	$1{,}938{,}304$
SPD, rotations	$28$	$516{,}416$	$926{,}464$	$1{,}746{,}560$	$3{,}386{,}752$
Sphere, torus, hyperbolic, Grassmann	$30$	$532{,}032$	$957{,}568$	$1{,}808{,}640$	$3{,}510{,}784$
Disk, two moons, nested circles	$2$	$319{,}232$	$533{,}760$	$962{,}816$	$1{,}820{,}928$

Appendix

Table 9: Parameter counts of the amortized flow charts used in the capacity sweep for each number of coupling layers

L

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.

Geometry	$L$	$K$	$L$	$K$	$L$	Solver	$S$
	Logarithmic map				Exponential map
	GP		GEORCE		ODE
Sphere 2	16	16	4	32	16	Heun	64
Flat torus 2	8	16	4	32	16	Tsit5	2
Hyperbolic 2	8	16	4	32	16	Tsit5	2
SPD 3	16	16	4	32	32	Dopri8	1

Appendix

Table 10: Settings selected for the runtime comparisons in Figure 5 . For each numerical method,

L

is the number of coupling layers in the chart used in that comparison.

K

is the number of knots for GP or GEORCE, and

S

is the number of ODE steps. Each comparison selects the fastest chart and numerical settings whose mean normalized RMSE is at most

\mathcal{R}_{\star}

Figure 8: Logarithmic map accuracy and runtime across all sixteen geometries. Each point gives the mean of

100\times\mathcal{R}_{\mathrm{logmap}}

, with one standard error across base points. The points for the four flow chart sizes and the four knot counts of the GP and GEORCE are joined by lines. The open diamond denotes the smaller flow chart used for the MFM comparison, shown alongside both MFM baselines. Both axes use logarithmic scales, with ranges chosen separately for each geometry. A black cross marks a setting with failed queries, which are excluded when computing its error.

Figure 9: Exponential map accuracy and runtime across all sixteen geometries. Each point gives the mean of

100\times\mathcal{R}_{\mathrm{expmap}}

, with one standard error across base points. The curves show every flow chart size and ODE step count listed in the legend. Both axes use logarithmic scales, with ranges chosen separately for each geometry. A black cross marks a setting with nonfinite endpoints, which are excluded when computing its error. We retain all finite errors. The large errors on the two dimensional sphere show that several ODE settings are unstable on this geometry. All flow chart queries with valid references return finite endpoints.

Figure 10: Effect of flow chart capacity and amortization. The panels show the mean of $100\times\mathcal{R}_{\mathrm{logmap}}$ for amortized and non-amortized flow charts, evaluated on the same four test base points at each chart size. The two panels share a logarithmic color scale.

Riemannian Geometry

Riemannian geometry studies curved spaces equipped with smoothly varying inner products, enabling intrinsic notions of distance, curvature, and optimization.

See all 43 papers

Jul 21, 2026

Riemannian Deep Learning: Modules, Networks, and Geometries

Deep neural networks on manifold-valued representations have attracted growing interest, but many basic components remain tied to specific manifolds, rely on Euclidean approximations, or require costly and numerically fragile geometric operations. This thesis develops a unified framework for Riemannian deep learning from three complementary perspectives: reusable neural modules, manifold-specific network architectures, and the design of underlying geometries. It generalizes batch normalization from Euclidean spaces and individual manifolds to broad classes of Lie groups and gyrogroups, and extends multinomial logistic regression from Euclidean space to SPD manifolds and then to general Riemannian manifolds. It further develops neural networks for several important geometric representations, including an unconstrained model of hyperbolic space, Busemann-based hyperbolic learning, and full-rank correlation matrices. Finally, it introduces adaptive and computationally efficient Riemannian metrics on SPD manifolds, including learnable Log-Euclidean geometries and fast, stable Cholesky-based geometries. The proposed methods are supported by theoretical analysis and validated through numerical experiments and applications in vision, signal processing, graph learning, and genomics.

Ziheng Chen

Doctoral School in Information and Communication Technology, Università di Trento

Riemannian Neural Networks Geometric Representation Learning Riemannian Metric Learning Hyperbolic Representation Learning

Ask PDF

Sep 18, 2026

Riemannian Neural Hamiltonian Flows: Geodesic Symplectic Transport and Interpretability

Hamiltonian normalizing flows are attractive generative models because their phase-space maps are invertible and volume preserving, but most neural constructions are formulated in Euclidean space. We introduce Riemannian Neural Hamiltonian Flows, which combine the fixed kinetic energy of a Riemannian manifold, a learned scalar potential, and an explicit geodesic leapfrog integrator. Our analysis explains how the learned Hamiltonian can be made interpretable. Every normalizable potential defines an implicit profile, and the position marginal initially accelerates along the relative score between that profile and the base. The matched potential is the interpretable specialization for which the implicit profile is the target. In the isotropic Gaussian case, the mechanism corresponds to a phase-space rotation. A local harmonic analysis extends this result around each mode of a general target on a manifold. The gap between the learned and the matched potential is the sum of a residual memory of the base and a bias of the model, and the two potentials agree when the position base has been transferred to the momentum. This can be achieved when the former is broader than the target. Numerical experiments on Euclidean, hyperbolic, and spherical spaces show competitive sample quality and numerical cost against a Riemannian continuous normalizing flow, and confirm the interpretability of the learned potential.

Vincent Souveton

CEA, DAM, DIF, Arpajon, F-91297, France

Normalizing Flows Hamiltonian Dynamics Hamiltonian Neural Networks Neural Network Interpretability

Ask PDF

Jun 12, 2026

Riemannian Metric Matching for Scalable Geometric Modeling of Distributions

High-dimensional datasets often concentrate near low-dimensional structures, but estimating their geometry from samples typically relies on graphs and kernels that scale poorly with dataset size and dimension. We propose Riemannian metric matching: a denoising probabilistic framework for learning the Riemannian geometry of data using neural networks. Specifically, we learn the carré du champ operator, which, using diffusion geometry, gives us access to the Riemannian geometry toolkit for downstream machine learning and statistical tasks. Our key observation is that the carré du champ operator can be formulated as a conditional expectation over random perturbations of the data, which can be exploited for sample-wise training and constant cost, amortized inference without explicit kernel construction. Empirically, metric matching rivals or improves the accuracy of

k

-NN-based diffusion geometry estimators, while enabling amortized inference that is up to

400\times

faster, and supports graph-free geometric analysis on high-dimensional images where nearest neighbors break down.

Jacob Bamberger, Adam Gosztolai, Pierre Vandergheynst +2

Department of Computer Science, University of Oxford, Oxford, UK · Institute of Artificial Intelligence, Medical University of Vienna, Vienna, Austria · EPFL, Lausanne, Switzerland +1

Geometric Representation Learning Metric Learning Riemannian Metric Learning Riemannian Geometry

Ask PDF

Abstract

Riemannian Geometry

Explore similar work

Riemannian Deep Learning: Modules, Networks, and Geometries

Riemannian Neural Hamiltonian Flows: Geodesic Symplectic Transport and Interpretability

Riemannian Metric Matching for Scalable Geometric Modeling of Distributions