Original: Rieoptax: Riemannian Optimization in JAX
Revised: Rieoptax : Riemannian Optimization in JAX

Original: author names withheld

Original: Under Review for OPT

Original: Abstract
Revised: Abstract

Original: We present Rieoptax, an open source Python library for Riemannian optimization in JAX.
Revised: We present Rieoptax, an open source Python library for Riemannian optimization in JAX.

Original: 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.
Revised: 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.

Original: We support various range of basic and advanced stochastic optimization solvers like Riemannian stochastic gradient, stochastic variance reduction, and adaptive gradient methods.
Revised: We support a range of basic and advanced stochastic optimization solvers like Riemannian stochastic gradient, stochastic variance reduction, and adaptive gradient methods.

Original: A distinguishing feature of the proposed toolbox is that we also support differentially private optimization on Riemannian manifolds.
Revised: A distinguishing feature of the proposed toolbox is that we also support differentially private optimization on Riemannian manifolds

Original: 1. Introduction
Revised: 1. Introduction

Original: Riemannian geometry is a generalization of the Euclidean geometry [49, 66] and includes several nonlinear spaces such as  positive definite matrices [18, 92], Grassmann manifold  [5, 15, 33], Stiefel manifold  [5, 25, 33], kendall shape spaces [59, 60, 71], hyperbolic spaces [94, 95], and special
Revised: Riemannian geometry is a generalization of the Euclidean geometry [57, 76] to general Riemannian manifolds. It includes several nonlinear spaces such as the set of positive definite matrices [20, 108], Grassmann manifold of subspaces [5, 17, 38], Stiefel manifold of orthogonal matrices [5, 29, 38],

Original:  Euclidean and orthogonal group [38, 87], to name a few.
Revised: Kendall shape spaces [68, 69, 81], hyperbolic spaces [27, 110, 111], and special Euclidean and orthogonal group [44, 100, 104], to name a few.

Original: Optimization with manifold based constraints has become increasingly popular and has been employed in various applications such as low rank matrix completion [21], learning taxonomy embeddings [76, 77], neural networks [53– 55, 75, 81], optimal transport [9, 26, 46, 73, 88], shape analysis [52, 90], and topological dimension reduction [56], among others.
Revised: Optimization with manifold based constraints has become increasingly popular and has been employed in various applications such as low rank matrix completion [23], learning taxonomy embeddings [87, 88], neural networks [45, 62–64, 86, 92], density estimation [51, 59], optimal transport [9, 31, 54, 84, 101], shape analysis [61, 106], and topological dimension reduction [65], among others.

Original: In addition, privacy preserving machine learning [29, 31] has become crucial in real applications, which has been generalized to manifold-valued problems very recently [44, 82, 96].
Revised: In addition, privacy preserving machine learning [2, 30, 34–36, 83, 105] has become crucial in real applications, which has been generalized to manifold-constrained problems very recently [52, 94, 112].

Original: Nevertheless, such  feature is absent in existing Riemannian optimization libraries.
Revised: Nevertheless, such a feature is absent in existing Riemannian optimization libraries [18, 24, 72, 80, 82, 103, 109].

Original: In this work, we introduce Rieoptax ( Rie mannian Opt imization in J ax ), an open source Python library for Riemannian optimization in JAX [23, 37].
Revised: In this work, we introduce Rieoptax ( Rie mannian Opt imization in J ax ), an open source Python library for Riemannian optimization in JAX [25, 42].

Original: The proposed library is mainly driven by the needs of efficient implementation of manifold-valued operations and optimization solvers, readily compatible with GPU and even TPU processors as well as the needs of privacy-supported Riemannian optimization.
Revised: The proposed library is mainly driven by the needs of efficient implementation of manifold-valued operations and optimization solvers, readily compatible with GPU and even TPU processors as well as the needs of privacy-supported Riemannian optimization.

Original: To the best of our knowledge, Rieoptax is the first library to provide privacy guarantees within the Riemannian optimization framework.
Revised: To the best of our knowledge, Rieoptax is the first library to provide privacy guarantees within the Riemannian optimization framework.

Original: 1.1. Background on Riemannian optimization, privacy, and JAX Riemannian Optimization.
Revised: 1.1. Background on Riemannian optimization, privacy, and JAX Riemannian optimization.

Original: Riemannian Optimization [5, 20] considers the following problem
Revised: Riemannian optimization [5, 22, 43] considers the following problem

Original: [Equation]
Revised: [Equation]

Original: where f : M → R , and M denotes a Riemannian manifold.
Revised: where f : M → R , and M denotes a Riemannian manifold.

Original: Instead of considering (1) as a constrained problem, Riemannian optimization [5, 20] views it as an unconstrained problem on the manifold space.
Revised: Instead of considering (1) as a constrained problem, Riemannian optimization [5, 22] views it as an unconstrained problem on the manifold space.

Original: Riemannian (stochastic) gradient descent [19, 98] generalizes the Euclidean© . gradient descent with intrinsic updates on manifold, i.e., w t +

Revised: Riemannian (stochastic) gradient descent [21, 114] generalizes the Euclidean gradient descent with intrinsic updates on manifold, i.e., w t +

Original: =
Revised: =

Original: Exp w t ( − η t grad f ( w t )) , where grad f ( w t ) is the Riemannian (stochastic) gradient, Exp w ( · ) is the Riemannian exponential map at w and η t is the step size.
Revised: Exp w t ( − η t grad f ( w t )) , where grad f ( w t ) is the Riemannian (stochastic) gradient, Exp w ( · ) is the Riemannian exponential map at w and η t is the step size.

Original: Recent years have witnessed significant advancements for Riemannian optimization where more advanced solvers are generalized from the Euclidean space to Riemannian manifolds.
Revised: Recent years have witnessed significant advancements for Riemannian optimization where more advanced solvers are generalized from the Euclidean space to Riemannian manifolds.

Original: These include variance reduction methods [42, 43, 57, 85, 100, 101], adaptive gradient methods [14, 58], accelerated gradient methods [7, 8, 45, 68, 99], quasi-Newton methods [51, 80], zeroth-order methods [67] and second order methods, such as trust region methods [4] and cubic regularized Newton’s methods [6].
Revised: These include variance reduction methods [49, 50, 66, 98, 116, 117], adaptive gradient methods [16, 67], accelerated gradient methods [7, 8, 53, 78, 115], quasi-Newton methods [60, 91], zeroth-order methods [77] and second order methods, such as trust region methods [4] and cubic regularized Newton’s methods [6].

Original: Differential privacy on Riemannian Manifolds.
Revised: Differential privacy on Riemannian manifolds.

Original: Differential privacy (DP) provides a rigorous treatment for data privacy by precisely quantifying the deviation in the model’s output distribution under modification of a small number of data points [29–32].
Revised: Differential privacy (DP) provides a rigorous treatment for data privacy by precisely quantifying the deviation in the model’s output distribution under modification of a small number of data points [34–37].

Original: Provable guarantees of DP coupled with properties like immunity to arbitrary post-processing and graceful composability have made it a de-facto standard of privacy with steadfast adoption in the real applications [3, 10, 28, 34, 74].
Revised: Provable guarantees of DP coupled with properties like immunity to arbitrary post-processing and graceful composability have made it a de-facto standard of privacy with steadfast adoption in the real applications [3, 11, 33, 39, 85]. Further, it has been shown empirically that DP models resist various kinds of leakage attacks that can cause privacy violations [14, 28, 93, 97, 102, 118].

Original: Recently, there is a surge of interest on differential privacy over Riemannian manifolds, which has been explored in the context of Fr´echet mean computation [82, 96] and, more generally, empirical risk minimization problemswhere the parameters are constrained to lie ona Riemannian manifold [44].
Revised: Recently, there is a surge of interest on differential privacy over Riemannian manifolds, which has been explored in the context of Fr´echet mean computation [94, 112] and more generally for empirical risk minimization problems on Riemannian manifolds [52]. [52] proposed differentially private Riemannian (stochastic) gradient descent methods by perturbing the Riemannian gradient with noise from the tangent Gaussian distribution before taking step: ζ = grad f ( w ) + ϵ, ϵ ∼ N w (0 , σ 2 ) . More recently, [10] has proposed efficient sampling procedures from the tangent Gaussian distribution for large scale and stochastic optimization scenarios.

Original: JAX and its ecosystem.
Revised: JAX and its ecosystem.

Original: JAX [23, 37] is recently introduced machine learning framework which support automatic differentiation capabilities [13] via grad() .
Revised: JAX [25, 42] is recently introduced machine learning framework which support automatic differentiation capabilities [15] via grad() .

Original: Further some of the distinguishing features of JAX are just-in-time (JIT) compilation using the accelerated linear algebra (XLA) compiler [41] via jit() , automatic vectorization (batch-level parallelism) support with vmap() , and strong support for parallel computation via pmap() .
Revised: Further some of the distinguishing features of JAX are just-in-time (JIT) compilation using the accelerated linear algebra (XLA) compiler [48] via jit() , automatic vectorization (batch-level parallelism) support with vmap() , and strong support for parallel computation via pmap() .

Original: All the above transformations can be composed arbitrarily because JAX follows the functional programming paradigm and implements these as pure functions.
Revised: All the above transformations can be composed arbitrarily because JAX follows the functional programming paradigm and implements these as pure functions.

Original: Given that JAX hasso many interesting features, its ecosystem has been constantly expanding in the last couple of years.
Revised: Given that JAX has many interesting features, its ecosystem has been constantly expanding in the last couple of years.

Original: Examples include neural network modules (Flax [48], Haiku [50], Equinox [62], Jraph [39], Equivariant-MLP [35]), reinforcement learning agents (Rlax [12]), Euclidean optimization algorithms (Optax [12]), federated learning (Fedjax [83]), optimal transport toolboxes (Ott [27]), sampling algorithms (Blackjax [64]), differential equation solvers (Diffrax [61]), rigid body simulators (Brax [36]), differentiable physics (Jax-md [86]).
Revised: Examples include neural network modules (Flax [56], Haiku [58], Equinox [71], Jraph [46], Equivariant-MLP [40]), reinforcement learning agents (Rlax [13]), Euclidean optimization algorithms (Optax [13]), federated learning (Fedjax [95]), optimal transport toolboxes (Ott [32]), sampling algorithms (Blackjax [73]), differential equation solvers (Diffrax [70]), rigid body simulators (Brax [41]), and differentiable physics (Jax-md [99]), among others.

Original: 1.2. Rieoptax
Revised: 1.2. Rieoptax

Original: We believe  the proposed framework for Riemannian optimization in JAX is a timely contribution, bringing several benefits of JAX and new features (such as privacy support) to the manifold optimization community, which are discussed below.
Revised: We believe that the proposed framework for Riemannian optimization in JAX is a timely contribution that brings several benefits of JAX and new features (such as privacy support) to the manifold optimization community discussed below.

Original: • Automatic and efficient vectorization with vmap() : functions that are written for inputs of size 1 can be converted to functions that take batch of inputs by wrapping it with vmap() .
Revised: • Automatic and efficient vectorization with vmap(). Functions that are written for inputs of size 1 can be converted to functions that take batch of inputs by wrapping it with vmap() .

Original: For example, the function def dist(point a, point b) for computing distance between a single point a and a single point b can be converted to function that computes distance between a batch of point a and/or a batch point b by wrapping dist with vmap() without modifying the dist() function.
Revised: For example, the function def dist(point a, point b) for computing distance between a single point a and a single point b can be converted to function that computes distance between a batch of point a and/or a batch point b by wrapping dist with vmap() without modifying the dist() function.

Original: This is useful in many cases, e
Revised: This is useful in many cases, e.g.,

Original: g., Fr´echet mean computation min w ∈M (cid:8) 1 n (cid:80) ni

Original: [-=
 f i ( w ) : = 1 n (cid:80) ni =1 dist 2 ( w, z i ) (cid:9) .-]
Revised: [-=
 f i ( w ) : = 1 n (cid:80) ni =1 dist 2 ( w, z i ) (cid:9) .-][Table]

Original: Furthermore, vectorization with vmap() is usually faster or on par with manual vectorization [23].
Revised: Furthermore, vectorization with vmap() is usually faster or on par with manual vectorization [25].

Original: • Per-example gradient clipping : A key process in differentially private optimization is perexample gradient clipping 1 n (cid:80) ni
Revised: • Per-example gradient clipping. A key process in differentially private optimization is perexample gradient clipping 1 n (cid:80) ni

Original: =1 clip τ ( grad f i ( w )) , where clip τ ensures norm is atmost τ .
Revised: =1 clip τ ( grad f i ( w )) , where clip τ ensures norm is atmost τ .

Original: Here, the order of operations is important: the gradients are first clipped and then averaged.
Revised: Here, the order of operations is important: the gradients are first clipped and then averaged.

Original: Popular libraries including Autograd [69], Pytorch [78] and Tensorflow [1] are heavily optimized to directly compute the mean gradient 1 n (cid:80) ni
Revised: Popular libraries including Autograd [79], Pytorch [89] and Tensorflow [1] are heavily optimized to directly compute the mean gradient 1 n (cid:80) ni

Original: =1 grad f i ( w ) and hence do not expose per-example gradients i e., grad f i ( w ) .
Revised: =1 grad f i ( w ) and hence do not expose per-example gradients i.e., grad f i ( w ) .

Original: Hence, one has to resort to ad-hoc techniques [40, 65, 84] or come up with algorithmic modifications [24] which inherently have speed versus performance trade-off.
Revised: Hence, one has to resort to ad-hoc techniques [47, 75, 96] or come up with algorithmic modifications [26] which inherently have speed versus performance trade-off.

Original: JAX, however, offers native support for handling such scenarios and JAX-based differentially private Euclidean optimization methods have been shown to be much faster than their non-JAX counterparts [91].
Revised: JAX, however, offers native support for handling such scenarios and JAX-based differentially private Euclidean optimization methods have been shown to be much faster than their non-JAX counterparts [107].

Original: We observe that JAX offer similar benefits for differentially private Riemannian optimization as well.
Revised: We observe that JAX offer similar benefits for differentially private Riemannian optimization as well.

Original: • Single Source Multiple Devices (SSMD) 
Revised: • Single Source Multiple Devices (SSMD) paradigm.

Original: Paradigm : JAX follows  SSMD paradigm, and hence code written for CPU can be run on GPU/TPU without any  modification.
Revised: JAX follows the SSMD paradigm, and therefore, the code written for CPUs can be run on GPU/TPUs without any additional modification.

Original: Rieoptax ismade available at https://anonymous.4open.science/r/Rieoptax/ for review and will be made public.
Revised: Rieoptax is available at https /github.com/ SaitejaUtpala/Rieoptax

Original: 2.  Implementation Overview
Revised: 2. Design and Implementation overview

Original: The package currently implements several commonly used geometries, optimization algorithms and differentially private mechanisms on manifolds.
Revised: The package currently implements several commonly used geometries, optimization algorithms and differentially private mechanisms on manifolds.

Original: More geometries and advanced solvers will be added in the future. 
Revised: More geometries and advanced solvers will be added in the future. 2.1. Core

Revised: • rieoptax.core.

Revised: ManifoldArray: lightweight wrapper of the jax device array with manifold attribute and used to model array constrained to manifold.

Revised: It is registered as Pytree to ensure compatibility jax primitives like grad() and vmap() .

Revised: • rieoptax.core.rgrad: Riemannian gradient operator.

Original: 2.1. Geometries
Revised: 2.2. Geometries

Original: Geometry module contains manifolds equipped with  Riemannian metrics.
Revised: Geometry module contains manifolds equipped with different Riemannian metrics.

Original: Each Geometry contains Riemannian inner product inp() ,and induced norm norm() , Riemannian exponential exp() and logarithm maps log() , induced Riemannian distance dist() , parallel transport pt() , and transformation from  Euclidean gradient to Riemannian gradient egrad to rgrad() .
Revised: Each Geometry contains Riemannian inner product inp() , induced norm norm() , Riemannian exponential exp() , logarithm maps log() , induced Riemannian distance dist() , parallel transport pt() , and transformation from the Euclidean gradient to Riemannian gradient egrad to rgrad() .

Original: Manifolds include symmetric positive definite (SPD) matrices SPD( m ) :
Revised: Manifolds include symmetric positive definite (SPD) matrices SPD( m ) :

Original: = { X ∈ R m × m : X =
Revised: = { X ∈ R m × m : X =

Original: X ⊤ , X ≻ 0 } , hyperbolic space, Grassmann manifold G ( m, r ) :
Revised: X ⊤ , X ≻ 0 } , hyperbolic space, Grassmann manifold G ( m, r ) :

Original: = { [ X ] : X ∈ R m × r , X ⊤ X = I } where [ X ] := { XO : O ∈ O ( r ) } , O ( r ) denotes the orthogonal group and hypersphere S ( d ) :
Revised: = { [ X ] : X ∈ R m × r , X ⊤ X = I } where [ X ] := { XO : O ∈ O ( r ) } , O ( r ) denotes the orthogonal group and hypersphere S ( d ) :

Original: = { x ∈ R d : x ⊤ x = 1 } .
Revised: = { x ∈ R d : x ⊤ x = 1 } .

Original: We use T x M to represent the tangent space at x and ⟨ u, v ⟩ x to represent the
Revised: We use T x M to represent the tangent space at x and ⟨ u, v ⟩ x to represent the

Original: Riemannian inner product.
Revised: Riemannian inner product.

Original: For more detailed treatment on these geometries, we refer to [5, 20, 95].
Revised: For more detailed treatment on these geometries, we refer to [5, 22, 111].

Original: • rieoptax.geometry.spd.
Revised: • rieoptax.geometry.spd.

Original: SPDAffineInvariant : SPD matrices with the affineinvariant metric [79]: SPD ( m ) with ⟨ U , V ⟩ X = tr( X − 1 UX − 1 V ) for U , V ∈ T X SPD( m ) .
Revised: SPDAffineInvariant: SPD matrices with the affineinvariant metric [90]: SPD ( m ) with ⟨ U , V ⟩ X = tr( X − 1 UX − 1 V ) for U , V ∈ T X SPD( m ) .

Original: • rieoptax.geometry.spd.
Revised: • rieoptax.geometry.spd.

Original: SPDLogEuclidean : SPD matrices with the log-Euclidean metric [11]: SPD ( m ) with ⟨ U , V ⟩ X = tr (cid:0) D
Revised: SPDLogEuclidean: SPD matrices with the Log-Euclidean metric [12]: SPD ( m ) with ⟨ U , V ⟩ X = tr (cid:0) D

Original: U logm( X )D V logm( X ) (cid:1) where D U logm( X ) is the directional derivative of matrix logarithm at X along U .
Revised: U logm( X )D V logm( X ) (cid:1) where D U logm( X ) is the directional derivative of matrix logarithm at X along U .

Original: • rieoptax.geometry.hyperbolic.
Revised: • rieoptax.geometry.hyperbolic.

Original: PoincareBall : Poincare-ball model of Hyperbolic space with Poincare metric [95], i
Revised: PoincareBall: the Poincare-ball model of Hyperbolic space with Poincare metric [111], i

Original: e., D ( d ) := { x ∈ R d : x ⊤ x < 1 } with ⟨ u , v ⟩ x = 4 u ⊤ v / (1 − x ⊤ x ) 2 for u , v ∈ T x D ( d ) .
Revised: e., D ( d ) := { x ∈ R d : x ⊤ x < 1 } with ⟨ u , v ⟩ x = 4 u ⊤ v / (1 − x ⊤ x ) 2 for u , v ∈ T x D ( d ) .

Original: • rieoptax.geometry.hyperbolic.
Revised: • rieoptax.geometry.hyperbolic.

Original: LorentzHyperboloid Lorentz Hyperboloid model of Hyperbolic space [95], i.e., H ( d ) = { x ∈ R d : ⟨ x , x ⟩ L = − 1 } with ⟨ u , v ⟩ x = ⟨ u , v ⟩ L for u , v ∈ T x H ( d ) , where ⟨ u , v ⟩ L := − u 0 v 0{+
+} + u 1 v 1 + · · · u d − 1 v d − 1 .
Revised: LorentzHyperboloid: the Lorentz Hyperboloid model of Hyperbolic space [111], i.e., H ( d ) = { x ∈ R d : ⟨ x , x ⟩ L = − 1 } with ⟨ u , v ⟩ x = ⟨ u , v ⟩ L for u , v ∈ T x H ( d ) , where ⟨ u , v ⟩ L := − u 0 v {+
+} + u 1 v 1 + · · · u d − 1 v d − 1 .

Original: • rieoptax.geometry.grassmann.
Revised: • rieoptax.geometry.grassmann.

Original: GrassmannCanonicalMetric : Grassmann manifold with the canonical metric [33], i.e., G ( m, r ) with ⟨ U , V ⟩ X = tr (cid:0) U T V (cid:1) for U , V ∈ T X G ( m, r ) .
Revised: GrassmannCanonicalMetric: the Grassmann manifold with the canonical metric [38], i.e., G ( m, r ) with ⟨ U , V ⟩ X = tr (cid:0) U T V (cid:1) for U , V ∈ T X G ( m, r ) .

Original: • rieoptax.geometry.hypersphere.
Revised: • rieoptax.geometry.hypersphere.

Original: HypersphereCanonicalMetric : Hypersphere manifold which canonical metric [5, 20], i.e., S ( d ) with ⟨ u , v ⟩ x = u ⊤ v for u , v ∈ T x S ( d ) .
Revised: HypersphereCanonicalMetric: the hypersphere manifold which canonical metric [5, 22], i.e., S ( d ) with ⟨ u , v ⟩ x = u ⊤ v for u , v ∈ T x S ( d ) .

Original: 2.2. Optimizers
Revised: 2.3. Optimizers

Original: Optimizers module contains Riemannian optimization algorithms. 
Revised: Optimizers module contains Riemannian optimization algorithms. Design of optimizers follows

Revised: Optax [13], which implements every optimizer by chaining of few common transformations.

Original:  • riepotax.optimizers.first order.rsgd : Riemannian stochastic gradient descent [19].
Revised: Where every optimizer • riepotax.optimizers.first order.rsgd: Riemannian stochastic gradient descent [21].

Original: • riepotax.optimizers.first order.rsvrg : Riemannian stochastic variance reduced gradient descent [100].
Revised: • riepotax.optimizers.first order.rsvrg: Riemannian stochastic variance reduced gradient descent [116].

Original: • riepotax.optimizers.first order.rsrg : Riemannian stochastic recursive gradient descent [57].
Revised: • riepotax.optimizers.first order.rsrg: Riemannian stochastic recursive gradient descent [66].

Original: • riepotax.optimizers.first order.rasa :
Revised: • riepotax.optimizers.first order.rasa:

Original: Riemannian adaptive stochastic gradient algorithm [58].
Revised: Riemannian adaptive stochastic gradient algorithm [67].

Original: • riepotax.optimizers.zeroth
Revised: • riepotax.optimizers.zeroth

Original: order.zo rgd : Zeroth-order Riemannian gradient descent [67].
Revised: order.zo rgd: zeroth-order Riemannian gradient descent [77].

Original: 2.3. Mechanism
Revised: 2.4. Privacy mechanism

Original: Mechanism module contains differential private mechanisms on Riemannian manifolds.
Revised: Mechanism module contains differential private mechanisms on Riemannian manifolds.

Original: • rieoptax.mechanism.output perturbation.
Revised: • rieoptax.mechanism.output perturbation.

Original: RieLaplaceMechanism : Implements Riemannian Laplace mechanism [82] which is used for privatizing Fr´echet mean computation.
Revised: RieLaplaceMechanism: the Riemannian Laplace mechanism [94] which is used for privatizing Fr´echet mean computation.

Original: • rieoptax.mechanism.output perturbation.
Revised: • rieoptax.mechanism.output perturbation.

Original: LogEuclideanMechanism : Implements log-Euclidean mechanism [96] which is used for differentially private Fr´echet mean on SPD matrices with log-Euclidean metric.
Revised: LogEuclideanMechanism: the Log-Euclidean mechanism [112] which is used for differentially private Fr´echet mean on SPD matrices with log-Euclidean metric.

Original: • rieoptax.mechanism.gradient perturbation.
Revised: • rieoptax.mechanism.gradient perturbation.

Original: DPRGDMechanism :
Revised: DPRGDMechanism:

Original: Implements noise calibration for Differentially private Riemannian gradient descent [44] based on moments accountant [2] in  autodp library [97].
Revised: noise calibration for differentially private Riemannian gradient descent [52] based on moments accountant [2] in the autodp library [113].

Original: • rieoptax.mechanism.gradient perturbation.
Revised: • rieoptax.mechanism.gradient perturbation.

Original: DPRSGDMechanism : Implements noise calibration for Differentially private Riemannian stochastic gradient descent [44] based on moments accountant [2] in autodp library [97].
Revised: DPRSGDMechanism: noise calibration for Differentially private Riemannian stochastic gradient descent [52] based on moments accountant [2] in autodp library [113].

Original: 3. Benchmarking Rieoptax
Revised: 3. Benchmarking Rieoptax

Original: In this section, we benchmark the proposed Rieoptax against existing Riemannian optimization libraries in Python.
Revised: In this section we benchmark the proposed Rieoptax against existing Riemannian optimization libraries in Python.

Original: These include Pytorch [78] based Mctorch [70] and Geoopt [63], Tensorflow [89]
Revised: These include Pytorch [89] based Mctorch [80] and Geoopt [72], Tensorflow [103]

Original: based Tensorflow-Riemopt
Revised: based Tensorflow-Riemopt

Original: (Tf-Riemopt) [89], Numpy [47] based Pymanopt [93], and Tensorflow based Geomstats [72].
Revised: (Tf-Riemopt) [103], Numpy [55] based Pymanopt [109], and Tensorflow based Geomstats [82].

Original: While Geomstats supports Numpy, Pytorch, and Tensorflow as backend, currently only  Tensorflow backend provides support for both CPU and GPU.
Revised: While Geomstats supports Numpy, Pytorch, and Tensorflow as backend, currently only the Tensorflow backend provides support for GPUs.

Original: Other non-Python based libraries include Manopt [22] in Matlab and Manopt.jl [16] in Julia [17].
Revised: Other non-Python based libraries include Manopt [24] in Matlab and Manopt.jl [18] in Julia [19].

Original: We benchmark the Riemannian exponential (Exp) and logarithm (Log) maps with the proposed Rieoptax against the aforementioned Python libraries whenever available with 64bitfloat precision.
Revised: We benchmark the Riemannian exponential (Exp) and logarithm (Log) maps with the proposed Rieoptax against the aforementioned Python libraries whenever available with 64bitfloat precision.

Original: For CPU benchmarking we use Processor AMD EPYC 7B1  with 2 cores and 16GB RAM.
Revised: For CPU benchmarking, we use the AMD EPYC 7B1 processor with 2 cores and 16GB RAM.

Original: For GPU benchmarking, we use CUDA version 11 . {+
+}
Revised: For GPU benchmarking, we use CUDA version 11 . {+
+}

Original: 4 on 16GB Tesla P100.
Revised: on 16GB Tesla P100.

Original: • Hypersphere canonical metric: Hypersphere S ( d ) is supported in Geoopt, Tf-Riemopt, Geomstats,  and Pyamanopt. 
Revised: Hypersphere: hypersphere S ( d ) is supported in Geoopt, Tf-Riemopt, Geomstats, McTorch, and Pyamanopt. McTorch does not support the Exp and Log maps.

Original: On GPU however, Geomstats raises an error.
Revised: On GPU, Geomstats raises an error.

Original: we benchmark  for dimensions d ∈ { 50 , 100 , 500 , 1000 , 5000 , 10000 , 25000 , 50000 } .
Revised: We benchmark all expect McTorch and Geomstats for dimensions d ∈ { 50 , 100 , 500 , 1 000 , 5 000 , 10 000 , 25 000 , 50 000 } .

Original: • Loretnz hyperboloid model:  Loretnz hyperboloid model H ( d ) is supported in Geoopt, Tf-Riemopt, Geomstats, and Mctorch.
Revised: • Loretnz hyperboloid model: the Loretnz hyperboloid model H ( d ) is supported in Geoopt, Tf-Riemopt, Geomstats, and Mctorch.

Original: While Riemannian exponential map is available in Mctorch, it does not implement the logarithm map.
Revised: While the Exp map is available in Mctorch, it does not implement the Log map.

Original: We benchmark for dimensions d ∈ { 50 , 100 , 500 , 1000 , 5000 , 10000 , 25000 , 50000 } . 
Revised: We benchmark for dimensions d ∈ { 50 , 100 , 500 , 1 000 , 5 000 , 10 000 , 25 000 , 50 000 } . [Figure] [Figure]

Original: • Grassmann with canonical metric : Grassmann manifold G ( m, r )with canonical metric is supported in Tf-Riemopt, Pymanopt,  Geomstats. 
Revised: • Grassmann: Grassmann manifold G ( m, r ) is supported in Tf-Riemopt, Pymanopt, McTorch, and Geomstats. McTorch

Original: However, we notice that the logarithm map in Tf-Riemopt is incorrectly implemented and  Geomstats representsGrassmann in projector matrix XX ⊤ ∈
Revised: does not support the Exp and Log maps. Geomstats represents

Original:  R m × m instead of X ∈ R m × r , which is prohibitively expensive.
Revised: Grassmann elements in projector matrices form XX ⊤ ∈ R m × m instead of X ∈ R m × r , which is prohibitively expensive.

Original: We thus exclude these two libraries from benchmarking.
Revised: We, therefore, exclude these three libraries from benchmarking.

Original: We benchmark for matrix size ( m, r ) ∈ { (100 , 10) , (500 , 10) , (750 , 10) , (1000 , 10) , (2000 , 10) , (5000 , 10) } .
Revised: We benchmark for matrix sizes ( m, r ) ∈ { (100 , 10) , (500 , 10) , (750 , 10) , (1 000 , 10) , (2 000 , 10) , (5 000 , 10) } .

Original: • SPD with affine invariant metric : SPD manifold SPD( m ) with affine invariant metric is supported in Geoopt, Tf-Riemopt, Geomstats.
Revised: • SPD with affine-invariant metric: SPD manifold SPD( m ) with the affine-invariant metric is supported in Geoopt, Tf-Riemopt, Geomstats, and McTorch.

Revised: McTorch, however, does not support the Exp and Log maps.

Original: We benchmark  for matrix size m ∈ { 10 , 50 , 75 , 100 , 150 , 200 } .
Revised: We benchmark all except McTorch for matrix sizes m ∈ { 10 , 50 , 75 , 100 , 150 , 200 } .

Original: Figures 1 and 2 present the timing results with CPU- and GPU-based computations, respectively.
Revised: Figures 1 and 2 present the timing results with CPU- and GPU-based computations, respectively.

Original: Overall, we observe that Rieoptax offers significant time improvement, especially on GPU.
Revised: Overall, we observe that Rieoptax offers significant time improvements, especially on GPUs.

Original: For  SPDAffineInvariant Case, Rieoptax is slightly slower than Geoopt because eigh which provides eigen decomposition is slightly slower in JAX compared to Pytorch.
Revised: For the SPDAffineInvariant case, Rieoptax is slightly slower than Geoopt because eigh which provides eigen decomposition is slightly slower in JAX compared to Pytorch.

Original: Given that JAX is a relatively new framework, we believe it will be faster even in this case in the near future. 
Revised: Given that JAX is a relatively new framework, we believe it would be faster even in this case in the near future. [Figure] 4. An example on PCA

Revised: In this section, we consider the problem of principal components analysis (PCA) by viewing it as an optimization problem on the Grassmann manifold [5, 22], i.e.,

Revised: [Equation]

Revised: where z i ∈ R n denote the data points.

Original: 4. Conclusion and future roadmap
Revised: The Rieoptax implementation for solving the problem (2) is shown in Listing 3. We provide timing of Rieoptax on TinyImageNet [74] which has a training set of 10 5 images in dimensions of 3 × 64 × 64 on both CPU and GPU. We take 5 000 images and vectorize each image to produce a sample matrix of size ( n, d ) = We compute the top r = 5 principal components, which leads to an optimization problem on G (12 288 , 5) .

Revised: For non-private PCA, we run the full Riemannian gradient descent method for 400 epochs.

Revised: For private PCA, we run the differentially private Riemannian gradient descent method [52] forepochs with a privacy configuration of ϵ = 0 .

Revised: 1 , δ = 10 − 6 and gradient clipping parameter of 0 . 1 .

Revised: from functools import partial from jax import jit, vmap from jax.numpy.linalg import norm from rieoptax.core import rgrad, ManifoldArray from rieoptax.geometry.grassmann import GrassmannCanonical from rieoptax.optimizer.first_order import rsgd, dp_rsgd from rieoptax.mechanism.gradient_perturbation import DP_RGD_Mechanism from rieoptax.optimizer.update import apply_updates def fit (params, data, optimizer, epochs, private=

Revised: False): @jitdefstep (params, opt_state, data): def cost (params, data): def _cost (params, data): diff = data-params.value @ (params.value.

Revised: T @ data) return norm(diff)**return vmap(_cost, in_axes=(None, 0 ))

Revised: (params, data).mean() rgrad_fn =

Revised: rgrad(cost) if private: data=data

Revised: [:, None] #per-example gradient rgrad_fn =

Revised: vmap(rgrad_fn, in_axes=(None, 0 ))

Revised: #calculates Riemannian gradients rgrads = rgrad_fn(params, data) updates, opt_state = optimizer.update(rgrads, opt_state, params) #update using Riemannian Exp params =

Revised: apply_updates(params, updates) return params, opt_state, loss_value

Revised: opt_state = optimizer.init(params) for i in range(epochs): params, opt_state, loss_value = step(params, opt_state, data) #initializationU_init=

Revised: ManifoldArray(value=init, manifold=GrassmannCanonical())

Revised: # non private PCA lr, epochs =

Revised: ( 3e-3 , 400 ) optimizer = rsgd(lr) non_private_U = fit(U_init, Z, optimizer, epochs)

Revised: [Table]

Revised: For both the private and non-private algorithms, we choose the same initialization and a learning rate of 3 × 10 − 3 .

Revised: Figures 3(a) and 3(b) show excess risk against runtime (in seconds) for non-private and private PCA, respectively.

Revised: 5. Conclusion and future roadmap

Original: In this work, we present a Python library for (privacy-supported)
Revised: In this work we present a Python library for (privacy-supported)

Original: Riemannian optimization, Rieoptax, and illustrate its efficacy on both CPU and GPU evaluation.
Revised: Riemannian optimization, Rieoptax and illustrate its efficacy on both CPU and GPU architectures.

Original: Our roadmap includes adding support for more manifold geometries, optimization algorithms, and a collection of example codes showcasing the usage of Rieoptax in various applications. [Figure] [Figure]
Revised: Our roadmap includes adding support for more manifold geometries, optimization algorithms and a collection of example codes showcasing the usage of Rieoptax in various applications, especially with differential privacy