Timezone: »
Rieoptax: Riemannian Optimization in JAX
Saiteja Utpala · Andi Han · Pratik Kumar Jawanpuria · Bamdev Mishra
Event URL: https://openreview.net/forum?id=d5MN87JhdiN »
We present Rieoptax, an open source Python library for Riemannian optimization in JAX. We show that many differential geometric primitives, such as Riemannian exponential and logarithm maps, are usually faster in Rieoptax than existing frameworks in Python, both on CPU and GPU. We support various range of basic and advanced stochastic optimization solvers like Riemannian stochastic gradient, stochastic variance reduction, and adaptive gradient methods. A distinguishing feature of the proposed toolbox is that we also support differentially private optimization on Riemannian manifolds.
Author Information
Saiteja Utpala (IIT Kanpur)
Andi Han (University of Sydney)
Pratik Kumar Jawanpuria (Microsoft)
Bamdev Mishra (Microsoft)
More from the Same Authors
-
2021 Poster: On Riemannian Optimization over Positive Definite Matrices with the Bures-Wasserstein Geometry »
Andi Han · Bamdev Mishra · Pratik Kumar Jawanpuria · Junbin Gao -
2020 : Poster Session 1 (gather.town) »
Laurent Condat · Tiffany Vlaar · Ohad Shamir · Mohammadi Zaki · Zhize Li · Guan-Horng Liu · Samuel Horváth · Mher Safaryan · Yoni Choukroun · Kumar Shridhar · Nabil Kahale · Jikai Jin · Pratik Kumar Jawanpuria · Gaurav Kumar Yadav · Kazuki Koyama · Junyoung Kim · Xiao Li · Saugata Purkayastha · Adil Salim · Dighanchal Banerjee · Peter Richtarik · Lakshman Mahto · Tian Ye · Bamdev Mishra · Huikang Liu · Jiajie Zhu -
2020 Poster: Statistical Optimal Transport posed as Learning Kernel Embedding »
Saketha Nath Jagarlapudi · Pratik Kumar Jawanpuria -
2018 Poster: Inexact trust-region algorithms on Riemannian manifolds »
Hiroyuki Kasai · Bamdev Mishra -
2018 Poster: A Dual Framework for Low-rank Tensor Completion »
Madhav Nimishakavi · Pratik Kumar Jawanpuria · Bamdev Mishra -
2015 Poster: Efficient Output Kernel Learning for Multiple Tasks »
Pratik Kumar Jawanpuria · Maksim Lapin · Matthias Hein · Bernt Schiele