Geoopt: Riemannian Optimization in PyTorch

TL;DR

Geoopt is an open-source PyTorch package that facilitates Riemannian optimization by providing a modular manifold interface, supporting various algorithms and enabling geometry-aware neural network layers.

Contribution

It introduces a flexible, modular framework for Riemannian optimization in PyTorch, including algorithms and tools for building geometry-aware neural networks.

Findings

Supports basic Riemannian SGD and adaptive algorithms

Provides manifold arithmetic methods for neural network layers

Enables integration of geometry-aware layers with existing models

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.