Source-linked AI summary
geomstats: a Python Package for Riemannian Geometry in Machine Learning
Nina Miolane, Johan Mathe, Claire Donnat, Mikael Jorda, Xavier Pennec
TL;DR
Machine-learning practitioners lack a modular framework for flexible, reusable Riemannian computations across manifolds and learning settings. Geomstats provides a tested, backend-flexible package with geometric operations, losses, gradients, and learning integrations, and demonstrates usability across manifold-valued applications. Its current scope is limited to manifolds with closed-form Exponential and Logarithm maps.
Problem
Machine-learning adoption of Riemannian geometry is limited by a lack of modular reusable implementations supporting flexible metrics and diverse manifold settings.
Method
Geomstats implements Riemannian manifolds, metrics, geodesic operations, losses, gradients, multiple computing backends, and machine-learning integrations.
Results
Geomstats demonstrates efficient and user-friendly manifold learning across inputs, outputs, and parameters, including a 61.2% out-of-sample accuracy for log-Euclidean connectome similarities.
Takeaways & Limitations
Geomstats offers an open-source, educational toolkit for applying Riemannian geometry across manifolds of increasing complexity and machine-learning use cases.
Takeaways & Limitations
The package currently implements manifolds whose Riemannian Exponential and Logarithm maps have closed forms; broader coverage is future work.
Abstract
from arXiv · showhide
We introduce geomstats, a python package that performs computations on manifolds such as hyperspheres, hyperbolic spaces, spaces of symmetric positive definite matrices and Lie groups of transformations. We provide efficient and extensively unit-tested implementations of these manifolds, together with useful Riemannian metrics and associated Exponential and Logarithm maps. The corresponding geodesic distances provide a range of intuitive choices of Machine Learning loss functions. We also give the corresponding Riemannian gradients. The operations implemented in geomstats are available with different computing backends such as numpy, tensorflow and keras. We have enabled GPU implementation and integrated geomstats manifold computations into keras deep learning framework. This paper also presents a review of manifolds in machine learning and an overview of the geomstats package with examples demonstrating its use for efficient and user-friendly Riemannian geometry.
1 Introduction
Riemannian geometry is increasingly relevant to machine learning because inputs, outputs, or parameters may naturally lie on manifolds. Geomstats addresses adoption barriers with a modular, tested, backend-flexible package for manifold computations.
- Motivation: Machine-learning inputs, outputs, or parameters can naturally belong to Riemannian manifolds.Examples include surfaces representing shapes, camera poses in SE(3), and neural-network weights constrained on Stiefel manifolds.
- Motivation: Modeling data on lower-dimensional manifolds can reduce degrees of freedom and memory allocation while providing more intuitive coordinates.The paper illustrates this with geographic locations represented by longitude and latitude rather than 3D Cartesian coordinates.
- Problem: Adoption has been inhibited by the lack of a modular, reusable framework, as existing packages target narrower settings or lack flexible metrics.The introduction contrasts custom code and packages focused on differential tensors or specific Lie groups with missing non-canonical metrics.
- Contribution: Geomstats provides extensively unit-tested manifold computations, NumPy and TensorFlow backends, GPU-capable vectorization, Keras Riemannian gradient descent, and educational material.The package reports more than 90% code coverage and targets machine-learning users.
- Paper organization: The paper combines a package overview with progressively complex manifold use cases, literature reviews, visualizations, and machine-learning implementations.The applications span manifolds embedded in flat spaces, SPD matrices, and Lie groups.
2 The Geomstats Package
Geomstats organizes Riemannian manifolds and metrics into reusable object-oriented classes and exposes core geometric operations for machine learning. It also supplies statistical tools, manifold-aware losses and gradients, constrained training, tested backends, and examples.
- Architecture: Geomstats separates manifold classes from Riemannian-metric classes in an object-oriented design.Manifold subclasses include embedded manifolds and Lie groups, while metrics are attached as manifold attributes.
- Core operations: RiemannianMetric provides inner products, norms, distances, Exponential and Logarithm maps, and geodesics; InvariantMetric adds left- and right-invariant Lie-group metrics.These operations support the package’s common geometric interface.
- Implementation: The implementations have more than 90% code coverage, NumPy and TensorFlow backends, vectorized batch computation, and TensorFlow GPU execution.Arrays are used to facilitate batch computations.
- Riemannian statistics: Geomstats includes weighted Fréchet means, manifold variance, and tangent principal component analysis for Riemannian statistics.The weighted Fréchet mean is computed through a Gauss-Newton gradient-descent iteration.
- Machine learning: For manifold-valued prediction, the package provides squared geodesic-distance losses and closed-form Riemannian gradients for back-propagation.These losses are designed to remain consistent with the manifold’s geometric structure.
- Training and examples: Modified Keras and TensorFlow versions constrain neural-network weights to manifolds during training, and the paper demonstrates the implemented manifolds through literature reviews and concrete use cases.The examples are organized by manifold and include implementation descriptions.
3 Embedded Manifolds - Hypersphere and Hyperbolic Space
Geomstats implements hypersphere and hyperbolic-space computations through embedded-manifold representations, metrics, optimization, visualization, and machine-learning examples. These manifolds support applications ranging from directional data and shape analysis to hierarchical representations and constrained neural-network training.
- Embedded-manifold foundations: The hypersphere and hyperbolic space are implemented as embedded manifolds with analogous geometric logic.The hypersphere is embedded in Euclidean space, whereas hyperbolic space is embedded in Minkowski space.
- Embedded-manifold foundations: Hypersphere and HyperbolicSpace provide coordinate conversion, projection to the manifold, and projection to tangent spaces.Their metrics are derived from the Euclidean and Minkowski embedding metrics, respectively.
- Hypersphere applications: Hyperspheres appear in circular, directional, orientation, protein-structure, crystallography, shape-statistics, and parameter-constraint applications.Examples include Kendall pre-shape spaces, s-rep representations, and neural-network parameters constrained on hyperspheres.
- Hypersphere use cases: Geomstats demonstrates hypersphere optimization by minimizing a positive quadratic form with Riemannian gradient descent and by training neural-network weights on the hypersphere.The implementation modifies Keras stochastic-gradient updates to use the Exponential map and supports parallel manifold optimization.
- Hyperbolic-space applications: Hyperbolic spaces model Fisher-metric Gaussian distributions, hierarchical data, and social-network graphs through their geometry.The paper describes hyperbolic spaces as continuous versions of trees and discusses hyperbolic geometric graphs.
- Hyperbolic-space visualization: Geomstats visualizes hyperbolic geometry in the Poincare disk using deformed geodesic grids and geodesic squares.The visualization toolbox is intended to let users test geometric intuition interactively.
4 Manifold of Symmetric Positive Definite (SPD) Matrices
SPD matrices support applications ranging from medical imaging and brain-connectome analysis to covariance-based descriptors and deep-learning features. Geomstats provides manifold computations and demonstrates connectome classification using Riemannian similarities.
- SPD manifold: Geomstats equips the SPD manifold with an affine-invariant Riemannian metric through its SPDMetric implementation.SPDMatricesSpace is embedded in the General Linear group and inherits from EmbeddedManifold.
- Use cases in machine learning: SPD matrices represent diffusion tensors, connectivity patterns, anatomical shape changes, covariance descriptors, and learned image features across several machine-learning applications.Examples include 3D diffusion tensors in DTI, regularized graph Laplacians in fMRI, Jacobian-based shape changes, covariance clustering, and neural-network features.
- Connectivity graph classification: The package demonstrates supervised brain-connectome learning using regularized Laplacian SPD matrices derived from resting-state fMRI data.The dataset contains 86 patients divided into balanced control and schizophrenia categories, and classification uses an SVM on pairwise brain similarities.
- Connectivity graph classification: Geomstats computes connectome similarities using both affine-invariant Riemannian and log-Euclidean distances, with the latter serving as a computationally lighter proxy.The distances use matrix logarithms and Frobenius norms and are available through the SPDMetric API.
- Connectivity graph classification: 61.2% out-of-sample accuracy makes the log-Euclidean distance the best-performing similarity against the standard Frobenius benchmark.The comparison evaluates similarities for graph classification on the connectome task.
- Connectivity graph classification: The affine-invariant Riemannian distance reveals greater connectome variability and identifies clusters that Frobenius and log-Euclidean distances largely miss.The associated clustermap shows red blocks on the diagonal, while the other distances place connectomes nearly uniformly far apart.
5 Lie Groups SO(n) and SE(n) - Rotations and Rigid Transformations
The paper presents SO(n) and SE(n) as useful representations for rotations, poses, articulated objects, and robotic trajectories, and demonstrates geomstats implementations for geodesic interpolation and manifold-valued prediction. In deep learning, geomstats supports SE(3)-valued CNN outputs through squared Riemannian-distance losses and Riemannian gradients.
- SO(n) represents rotations, while SE(n) represents rotations and translations through homogeneous coordinates.
- Use cases in machine learning: SO(3) and SE(3) model articulated objects, camera orientations, and camera poses in robotics and computer vision.
- Geomstats use case: geodesics on SO(3): Geomstats interpolates robot-arm orientations by computing geodesics on SO(3) between initial and desired final orientations.
- Geomstats use case: deep learning predictions on SE(3): Geomstats trains CNNs to predict poses in SE(3) using a left-invariant metric and squared Riemannian geodesic distance as the loss.
- Geomstats use case: deep learning predictions on SE(3): Experiments show significant accuracy improvements for image-based 2D-to-3D registration with the Riemannian loss.
6 Conclusion and Outlook
The conclusion presents geomstats as an open-source, educational package supporting manifold-based machine learning across diverse geometries and applications. Its current scope is limited to manifolds with closed-form Riemannian Exponential and Logarithm maps, with broader manifold support and a PyTorch backend planned for future work.
- Geomstats provides an open-source, hands-on tool for learning and applying Riemannian geometry across manifolds of increasing complexity.
- The paper demonstrates manifold-valued inputs, outputs, and parameters through concrete machine-learning use cases across varied application areas.
- Current implementations require closed-form Exponential and Logarithm maps for the Riemannian metrics.
- Future work will implement manifolds without necessarily available closed forms and add a PyTorch backend.