Source-linked AI summary
Geoopt: Riemannian Optimization in PyTorch
Max Kochurov, Rasul Karimov, Serge Kozlukov
TL;DR
Geoopt addresses the need for Riemannian optimization integrated with PyTorch for geometric deep learning research. It provides manifold-aware optimization and representations through PyTorch-compatible interfaces, supporting standard and adaptive optimizers across several geometries. The package has been used in applications including computer vision, navigation, optimal transport, time-series analysis, and hyperbolic deep learning.
Problem
Riemannian optimization is important for geometric deep learning, but general-purpose manifold optimization is difficult to design for all use cases.
Method
Geoopt combines a standard manifold interface with PyTorch-compatible tensors, optimizers, geometric operations, and implementations for several standard and hyperbolic manifolds.
Results
Geoopt supports widely used optimizers, hyperbolic and stereographic geometry operations, and research applications across computer vision, navigation, optimal transport, time-series analysis, and hyperbolic deep learning.
Takeaways & Limitations
Geoopt provides a PyTorch-integrated foundation for conducting research with Riemannian optimization and hyperbolic representations.
Takeaways & Limitations
Geoopt assumes PyTorch-style integration and ambient-vector-space representations, while other projects provide broader collections of algorithms or manifolds.
Abstract
from arXiv · showhide
Geoopt is a research-oriented modular open-source package for Riemannian Optimization in PyTorch. The core of Geoopt is a standard Manifold interface that allows for the generic implementation of optimization algorithms. Geoopt supports basic Riemannian SGD as well as adaptive optimization algorithms. Geoopt also provides several algorithms and arithmetic methods for supported manifolds, which allow composing geometry-aware neural network layers that can be integrated with existing models.
1. Introduction
Geoopt is a PyTorch-based package motivated by hyperbolic embeddings and hyperbolic neural networks. It supports multiple models of hyperbolic space, including the Poincaré ball, Hyperboloid, and generalized κ-Stereographic models.
- Geoopt is built on PyTorch, enabling automatic differentiation, GPU acceleration, model exporting, and compatibility with native PyTorch functions.Its optimizers can serve as drop-in replacements for native PyTorch optimizers, while manifold parameters are declared differently.
- The package is motivated primarily by experiments with hyperbolic embeddings and hyperbolic neural networks.
- Geoopt provides models of hyperbolic space including the Poincaré ball, Hyperboloid, and generalized κ-Stereographic models.The generalized κ-Stereographic model unifies Hyperbolic, Euclidean, and Spherical geometries.
2. Riemannian optimization
Riemannian optimization replaces ordinary vector-space updates with operations on tangent spaces and the manifold itself. Geoopt obtains ambient derivatives through PyTorch, converts them into manifold directions, and applies geometric updates.
- A manifold gradient step moves from the current point along a tangent-space descent direction using the exponential map.The learning rate η controls the update, and Figure 2 depicts the geodesic path from p_t to p_t+1 on the Poincaré disk.
- The tangent space T_pM is the vector space of directions represented by length-minimizing paths emanating from point p.The logarithmic map recovers a tangent direction from nearby points, while the exponential map maps a direction to its destination.
- For an objective J on a manifold, the derivative at p_t is a linear functional on the tangent space, and the Riemannian metric converts it into a tangent direction.
- Geoopt uses PyTorch’s backward pass to obtain an ambient derivative and converts it through the manifold’s inner product into the corresponding manifold direction.This connects ordinary automatic differentiation with Riemannian optimization.
3. Design goals
Geoopt focuses on geometric deep learning rather than attempting to cover every manifold-optimization use case. Its design goals prioritize PyTorch integration, broadcasting, numerical robustness, efficiency, and extensibility.
- PyTorch integration: Geoopt provides PyTorch-like interfaces so its optimizers can act as drop-in replacements for torch.optim and work with third-party PyTorch packages.
- Broadcasting: The package supports broadcasting for operations and for product-manifold semantics.
- Robustness and numerical stability: Geoopt aims to limit numerical instability in hyperbolic models, including excess NaNs, and to support float32 precision whenever possible.The documentation describes instabilities of specific functions separately.
- Efficiency and extendibility: After satisfying integration, broadcasting, and stability requirements, Geoopt seeks reasonable efficiency and room for extension.
4. Implementation details
Geoopt represents manifold objects within PyTorch tensors and organizes optimization through manifold methods that implement geometric operations. Its implementation also exploits broadcasting, product structure, and efficient update-step transformations.
- geoopt.ManifoldTensor is a PyTorch-compatible tensor that stores a reference to its containing manifold.It inherits from torch.Tensor and torch.nn.Parameter, providing compatibility with the PyTorch ecosystem.
- Product-manifold operations use broadcasting, and adaptive terms are computed per manifold parameter while exploiting product structure.The implementation also merges retractions followed by parallel transport to optimize update steps.
- The geoopt.Manifold base class defines the method set expected by Geoopt optimizers and is captured by PyTorch state dictionaries.
- Geoopt’s inner operation computes inner products of tangent vectors, while egrad2rgrad converts an ambient covector into a tangent vector.
- A manifold implementation must provide operations such as retraction, exponential mapping, vector transport, and tangent-vector handling.Retraction approximates the exponential map, while vector transport approximates parallel transport.
5. Features
Geoopt provides standard Riemannian manifolds and optimizers for common constrained and geometric deep-learning settings. Its feature set includes sparse and adaptive methods, hyperbolic models, and composable product and scaled manifolds.
- Manifolds: Geoopt implements standard manifolds including Sphere, Stiefel, BirkhoffPolytope, Stereographic, Lorentz, Product, and Scaled manifolds.These cover unit-norm constraints, basis reconstruction, permutation inference, hyperbolic learning, and manifold composition.
- Optimizers: Geoopt supports RiemannianAdam and SparseRiemannianAdam as adaptive alternatives to Riemannian stochastic gradient descent.These implement Adam-style optimization with support for sparse gradients.
- Optimizers: Geoopt provides RiemannianSGD and SparseRiemannianSGD, including momentum support and sparse-gradient handling.The sparse optimizer targets models whose gradients are sparse.
6. Advanced Usage
Geoopt supports geometric deep learning through hyperbolic and other manifold representations, while its Riemannian optimization library has been applied across several research domains. Its advanced functionality includes unified curvature-aware stereographic operations and applications in language, vision, time-series, and representation learning.
- 6. Advanced Usage: Geoopt provides robust Poincare Ball operations, unified Mobius arithmetic across negative, zero, and positive curvature, and derivatives that support curvature optimization.The unified stereographic implementation also includes spherical positive-curvature geometry and zero-curvature cases.
- Other Applications: Geoopt is a general-purpose PyTorch optimization library for manifold optimization applications.The package is positioned for research use across geometric deep-learning settings.
- Other Applications: Geoopt's broader application scope includes language models, computer vision, time-series analysis, and hyperbolic deep learning.Examples include unitary or Stiefel constraints, doubly stochastic matrices, SPD neural networks, and hyperbolic embeddings.
- Hyperbolic deep learning: Hyperbolic representations address implicit hierarchical relationships and are used in NLP, image understanding, and general representation learning.Geoopt supports optimization in several models of real hyperbolic spaces together with basic hyperbolic operations.
- Other Applications: Geoopt has supported research in computer vision, navigation, optimal transport, time-series analysis, and hyperbolic deep learning.The summary identifies published applications across these areas.
7. Related projects
Geoopt distinguishes itself from prior Riemannian optimization projects through interfaces designed specifically for PyTorch-based neural networks and geometric deep learning. Unlike PyManOpt's problem-construction workflow, Geoopt lets users define networks and cost functions in the usual PyTorch way.
- Interface distinction: Geoopt is explicitly designed for neural networks and geometric deep learning, whereas PyManOpt and Geomstats are general-purpose Riemannian optimization tools.The distinction is primarily interface-wise rather than a claim that the alternatives are inadequate.
- Interface distinction: Geoopt users define neural networks and cost functions in the usual PyTorch way without constructing a PyManOpt Problem.PyManOpt instead follows solver.solve(Problem(manifold, cost)) semantics.
- Scope: PyManOpt provides a broader collection of algorithms and manifolds than Geoopt, including trusted-region methods and Nelder-Mead.The passage presents both packages as useful for different purposes.