[Paper Review] Geoopt: Riemannian Optimization in PyTorch

Max Kochurov, Rasul Karimov+1 more|arXiv (Cornell University)|May 6, 2020

Advanced Numerical Analysis Techniques35 references49 citations

TL;DR

Geoopt is a PyTorch-based, modular library for Riemannian optimization that integrates seamlessly with PyTorch, supports hyperbolic geometries, and provides manifold-aware optimizers and layers for geometric deep learning.

ABSTRACT

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.

Motivation & Objective

• Provide a PyTorch-integrated, drop-in interface for Riemannian optimization in neural networks.
• Support a range of manifolds and geometry-aware operations for geometric deep learning.
• Offer robust, efficient optimization algorithms including adaptive methods on manifolds.
• Enable composition of geometry-aware layers and compatibility with existing PyTorch models.

Proposed method

• Introduce a standard Manifold interface for generic optimization on manifolds.
• Implement manifold-aware optimizers (e.g., Riemannian Adam, Riemannian SGD) that conform to PyTorch optimizers.
• Represent points and directions via embeddings in ambient vector spaces with egrad2rgrad translation.
• Provide several manifolds (Sphere, Stiefel, Birkhoff Polytope, Stereographic hyperbolic models, Lorentz) and product/scaled variants.
• Enable efficient updates by retractions, exponential maps, and tangent-vector transport on product manifolds.

Experimental results

Research questions

• RQ1 How can Riemannian optimization be integrated into the PyTorch ecosystem as drop-in components?
• RQ2 What manifolds and geometric primitives are most useful for deep learning applications (e.g., hyperbolic spaces, unitary constraints, orthogonality)?
• RQ3 Can adaptive and sparse-gradient optimizers be effectively realized on manifolds to improve training stability and performance?

Key findings

• Geoopt provides a manifold-aware optimization framework that plugs into PyTorch optimizers and training loops.
• It supports hyperbolic deep learning via Poincaré ball, Hyperboloid, and generalized κ-Stereographic models, enabling geometry-aware neural networks.
• Adaptive optimizers (RiemannianAdam, SparseRiemannianAdam) and SGD variants are available for manifold parameters.
• Product and scaled manifolds allow combining multiple geometries and extending standard manifolds in Riemannian optimization.