Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Geoopt: Riemannian Optimization in PyTorch

Max Kochurov, Rasul Karimov, Serge Kozlukov

TL;DR

Geoopt tackles integrating Riemannian optimization into PyTorch by treating manifold-valued parameters as first-class citizens and providing geometry-aware optimizers. It defines a robust abstract interface and a ManifoldTensor, enabling efficient updates via retractions and parallel transport within standard PyTorch workflows. The package supports multiple manifolds (Sphere, Stiefel, Birkhoff, Poincaré ball, Lorentz) and optimizers (RiemannianAdam, Sparse variants, RiemannianSGD), with advanced hyperbolic learning capabilities. Its design emphasizes numerical stability, broadcasting, and extensibility, making Geoopt a versatile tool for geometric deep learning across NLP, CV, time-series, and hyperbolic representations.

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.

Geoopt: Riemannian Optimization in PyTorch

TL;DR

Geoopt tackles integrating Riemannian optimization into PyTorch by treating manifold-valued parameters as first-class citizens and providing geometry-aware optimizers. It defines a robust abstract interface and a ManifoldTensor, enabling efficient updates via retractions and parallel transport within standard PyTorch workflows. The package supports multiple manifolds (Sphere, Stiefel, Birkhoff, Poincaré ball, Lorentz) and optimizers (RiemannianAdam, Sparse variants, RiemannianSGD), with advanced hyperbolic learning capabilities. Its design emphasizes numerical stability, broadcasting, and extensibility, making Geoopt a versatile tool for geometric deep learning across NLP, CV, time-series, and hyperbolic representations.

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.

Paper Structure

This paper contains 13 sections, 1 equation, 2 figures.

Table of Contents

  1. Introduction
  2. Riemannian optimization
  3. Design goals
  4. Implementation details
  5. Features
  6. Advanced Usage
  7. Other Applications
  8. Language models.
  9. Computer vision.
  10. Time series.
  11. Hyperbolic deep learning.
  12. Summary.
  13. Related projects

Figures (2)

  • Figure 1: Creation of a manifold valued parameter.
  • Figure 2: A gradient descent step on the Poincaré disk. Contour lines visualize the objective function; $p_t$ is the current estimate; $-\flat(\operatorname{d}\mathcal{J})$ is the descent direction, visualized as a geodesic curve; $p_{t+1}$ is the final point of that curve and the new estimate; $\partial_{x_1},~\partial_{x_2}$ are basis vectors in the space of directions at $p_t$; stroked line visualizes the (downscaled) "Euclidean" gradient.