mlsanityData to 7 Oct 2026

Rieoptax: Riemannian Optimization in JAX

arXiv 2210.04840 ↗First posted 10 Oct 2022math.OC cs.LG cs.MScited by 1

Not on the map. Why?

Abstract (arXiv)

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.

Key facts. Cited by 1 paper in the map (data through 7 Oct 2026). Builds on 56 of its 116 references that are in the map. First posted on arXiv 10 Oct 2022.

Builds on

56 of 116 references.

References · 56

  1. 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,160
  2. 2016Deep Learning with Differential Privacycited by 2,682
  3. 2016Membership Inference Attacks against Machine Learning Modelscited by 2,221
  4. 2019Deep Leakage from Gradientscited by 1,004
  5. 2018The Secret Sharer: Evaluating and Testing Unintended Memorization in Neural Networkscited by 847
  6. 2015Automatic differentiation in machine learning: a surveycited by 846
  7. 2017Poincar'e Embeddings for Learning Hierarchical Representationscited by 674
  8. 2016TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systemscited by 660
  9. 2017Renyi Differential Privacycited by 643
  10. 2018Hyperbolic Neural Networkscited by 393
  11. 2021Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulationcited by 375
  12. 2019White-box vs Black-box: Bayes Optimal Strategies for Membership Inferencecited by 265
38 more
  1. 2022On Neural Differential Equationscited by 264
  2. 2014RAPPOR: Randomized Aggregatable Privacy-Preserving Ordinal Responsecited by 262
  3. 2018Learning Continuous Hierarchies in the Lorentz Model of Hyperbolic Geometrycited by 261
  4. 2018Subsampled R'enyi Differential Privacy and Analytical Moments Accountantcited by 234
  5. 2018Riemannian Adaptive Optimization Methodscited by 187
  6. 2021A Practical Method for Constructing Equivariant Multilayer Perceptrons for Arbitrary Matrix Groupscited by 185
  7. 2016A Riemannian Network for SPD Matrix Learningcited by 162
  8. 2017Collecting Telemetry Data Privatelycited by 147
  9. 2021Equinox: neural networks in JAX via callable PyTrees and filtered transformationscited by 143
  10. 2022Reconstructing Training Data with Informed Adversariescited by 143
  11. 2018A micro Lie theory for state estimation in roboticscited by 117
  12. 2013Manopt, a Matlab toolbox for optimization on manifoldscited by 112
  13. 2022Optimal Transport Tools (OTT): A JAX Toolbox for all things Wassersteincited by 96
  14. 2016Pymanopt: A Python Toolbox for Optimization on Manifolds using Automatic Differentiationcited by 93
  15. 2016First-order Methods for Geodesically Convex Optimizationcited by 92
  16. 2018geomstats: a Python Package for Riemannian Geometry in Machine Learningcited by 71
  17. 2016Riemannian SVRG: Fast Stochastic Optimization on Riemannian Manifoldscited by 64
  18. 2016Building Deep Networks on Grassmann Manifoldscited by 58
  19. 2020Enabling Fast Differentially Private SGD via Just-in-Time Compilation and Vectorizationcited by 58
  20. 2019A neural network based on SPD manifold learning for skeleton-based hand gesture recognitioncited by 46
  21. 2020Scaling up Differentially Private Deep Learning with Fast Per-Example Gradient Clippingcited by 46
  22. 2015Efficient Per-Example Gradient Computationscited by 42
  23. 2016Deep Learning on Lie Groups for Skeleton-based Action Recognitioncited by 42
  24. 2021On Riemannian Optimization over Positive Definite Matrices with the Bures-Wasserstein Geometrycited by 32
  25. 2021Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descentcited by 32
  26. 2021FedJAX: Federated learning simulation with JAXcited by 32
  27. 2021Fast and Memory Efficient Differentially Private-SGD via JL Projectionscited by 29
  28. 2019A Continuous-time Perspective for Modeling Acceleration in Riemannian Optimizationcited by 25
  29. 2020From Nesterov's Estimate Sequence to Riemannian Accelerationcited by 24
  30. 2021Transductive Few-Shot Classification on the Oblique Manifoldcited by 23
  31. 2017An Alternative to EM for Gaussian Mixture Models: Batch and Stochastic Riemannian Optimizationcited by 20
  32. 2018Adaptive regularization with cubics on manifoldscited by 17
  33. 2019Efficient Per-Example Gradient Computations in Convolutional Neural Networkscited by 17
  34. 2018McTorch, a manifold optimization library for deep learningcited by 16
  35. 2021Differential Privacy Over Riemannian Manifoldscited by 16
  36. 2018Faster First-Order Methods for Stochastic Non-Convex Optimization on Riemannian Manifoldscited by 12
  37. 2022Differentially private Riemannian optimizationcited by 12
  38. 2021Manifold optimization for non-linear optimal transport problemscited by 8

50 most cited of 56 All 56 (JSON)

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2026: 1.

1 paper.

  1. 2026hyperbolix: Hyperbolic Deep Learning in JAXcited by 0