Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

tensorflow-riemopt: A Library for Optimization on Riemannian Manifolds

Oleg Smirnov

TL;DR

The paper introduces tensorflow-riemopt, a TensorFlow-based library for optimization on Riemannian manifolds that targets both research prototyping and production deployment via TensorFlow Extended. It provides manifold-constrained variables, a low-level API for custom geometric algorithms, and high-level neural network layers across manifolds such as hyperbolic spaces, Grassmannian, Stiefel, and SPD, all with Riemannian SGD/RMSprop/Adam optimizers. The design emphasizes seamless integration with TensorFlow, automatic differentiation, and serialization, while prioritizing closed-form manifold operators with numerical fallbacks. Benchmark results indicate strong performance, especially on high-dimensional manifolds, highlighting its suitability for scalable, production-grade geometric ML applications like hyperbolic embeddings and manifold-based feature learning.

Abstract

This paper presents tensorflow-riemopt, a Python library for geometric machine learning in TensorFlow. The library provides efficient implementations of neural network layers with manifold-constrained parameters, geometric operations on Riemannian manifolds, and stochastic optimization algorithms for non-Euclidean spaces. Designed for integration with TensorFlow Extended, it supports both research prototyping and production deployment of machine learning pipelines. The code and documentation are distributed under the MIT license and available at https://github.com/master/tensorflow-riemopt

tensorflow-riemopt: A Library for Optimization on Riemannian Manifolds

TL;DR

The paper introduces tensorflow-riemopt, a TensorFlow-based library for optimization on Riemannian manifolds that targets both research prototyping and production deployment via TensorFlow Extended. It provides manifold-constrained variables, a low-level API for custom geometric algorithms, and high-level neural network layers across manifolds such as hyperbolic spaces, Grassmannian, Stiefel, and SPD, all with Riemannian SGD/RMSprop/Adam optimizers. The design emphasizes seamless integration with TensorFlow, automatic differentiation, and serialization, while prioritizing closed-form manifold operators with numerical fallbacks. Benchmark results indicate strong performance, especially on high-dimensional manifolds, highlighting its suitability for scalable, production-grade geometric ML applications like hyperbolic embeddings and manifold-based feature learning.

Abstract

This paper presents tensorflow-riemopt, a Python library for geometric machine learning in TensorFlow. The library provides efficient implementations of neural network layers with manifold-constrained parameters, geometric operations on Riemannian manifolds, and stochastic optimization algorithms for non-Euclidean spaces. Designed for integration with TensorFlow Extended, it supports both research prototyping and production deployment of machine learning pipelines. The code and documentation are distributed under the MIT license and available at https://github.com/master/tensorflow-riemopt

Paper Structure

This paper contains 8 sections, 2 figures.

Table of Contents

  1. Introduction
  2. Usage modes
  3. Machine learning applications.
  4. Riemannian geometry operations.
  5. Implementation details
  6. Comparison to related software
  7. Conclusion
  8. Benchmarks

Figures (2)

  • Figure 1: Geometric operations on $\mathbb{S}^2$
  • Figure 2: Benchmarks for geodesic distance $d(x,y)$, exponential map $\mathrm{Exp}_x(v)$, and logarithmic map $\mathrm{Log}_x(y)$ across five manifolds: Euclidean space ($\mathbb{R}^n$), Poincaré ball ($\mathbb{H}^n$), hypersphere ($\mathbb{S}^{n-1}$), special orthogonal group ($\mathrm{SO}(n)$), and symmetric positive definite matrices ($\mathcal{S}^n_{++}$). In high-dimensional settings, tensorflow-riemopt performs comparably to or better than existing alternatives. The observed differences stem from implementation choices: exact geometric operations versus approximations such as retractions.