[Paper Review] Geomstats: A Python Package for Riemannian Geometry in Machine Learning

Nina Miolane, Brigant, Alice Le+17 more|arXiv (Cornell University)|Apr 7, 2020

Morphological variations and asymmetry41 references97 citations

TL;DR

Geomstats is an open-source Python toolbox for computations and statistics on nonlinear manifolds, enabling geometry-aware ML with vectorized operations and GPU backends.

ABSTRACT

We introduce Geomstats, an open-source Python toolbox for computations and statistics on nonlinear manifolds, such as hyperbolic spaces, spaces of symmetric positive definite matrices, Lie groups of transformations, and many more. We provide object-oriented and extensively unit-tested implementations. Among others, manifolds come equipped with families of Riemannian metrics, with associated exponential and logarithmic maps, geodesics and parallel transport. Statistics and learning algorithms provide methods for estimation, clustering and dimension reduction on manifolds. All associated operations are vectorized for batch computation and provide support for different execution backends, namely NumPy, PyTorch and TensorFlow, enabling GPU acceleration. This paper presents the package, compares it with related libraries and provides relevant code examples. We show that Geomstats provides reliable building blocks to foster research in differential geometry and statistics, and to democratize the use of Riemannian geometry in machine learning applications. The source code is freely available under the MIT license at \url{geomstats.ai}.

Motivation & Objective

• Foster research in differential geometry and geometric statistics by providing reusable low-level implementations.
• Democratize geometric statistics by offering user-friendly algorithms via a Scikit-Learn-like API.
• Provide educational tooling and visualizations to facilitate hands-on learning of differential geometry.
• Offer a modular, tested foundation to integrate manifold computations into broader ML workflows.

Proposed method

• Two main modules: geometry for Riemannian geometry with an object-oriented design and learning for manifold-based statistics and learning algorithms.
• Manifolds inherit from a base class Manifold; RiemannianMetric provides distance, exponential/logarithmic maps, etc.; several metrics (e.g., HyperbolicMetric, StiefelCanonicalMetric) extend this base.
• Automatic differentiation (autograd) enables computations when closed-form expressions are unavailable via a Connection class.
• Learning module implements Fréchet mean, KMeans, TangentPCA, and other manifold-aware algorithms with Scikit-Learn-style API.
• Code is vectorized for batch operations, supports backends NumPy, PyTorch, TensorFlow for GPU acceleration, and includes visualization tools.

Experimental results

Research questions

• RQ1 What manifolds and metrics should be supported to cover common ML and geometry needs?
• RQ2 How can manifold computations be integrated into standard ML pipelines via an intuitive API?
• RQ3 What is the performance and usability trade-off when using multiple backends (NumPy, PyTorch, TensorFlow) for manifold operations?
• RQ4 How do manifold-based learning algorithms (e.g., Fréchet mean, KMeans, TangentPCA) perform and behave across different geometries?

Key findings

• Geomstats provides an open-source, object-oriented foundation for Riemannian geometry and geometric statistics.
• The package offers multiple manifolds and families of Riemannian metrics with associated exponential/log maps, geodesics, and parallel transport.
• Learning algorithms on manifolds (e.g., Fréchet mean estimators, KMeans, TangentPCA) are implemented with a Scikit-Learn-compatible API.
• Operations are vectorized for batch computation and backends include NumPy, PyTorch, and TensorFlow for GPU acceleration.
• Code quality is ensured by extensive unit testing and continuous integration (CI).
• Geomstats is designed to be educational and accessible, with visualization tools and documentation, and is available under the MIT license at geomstats.ai.