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

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 a 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 [57, 76] to general Riemannian manifolds.
Revised: Riemannian geometry is a generalization of the Euclidean geometry [57, 76] to general Riemannian manifolds.

Original: It includes several nonlinear spaces such as the set of positive definite matrices [20, 108],
Revised: It includes several nonlinear spaces such as the set of positive definite matrices [20, 108],

Original: Grassmann manifold of subspaces [5, 17, 38], Stiefel manifold of orthogonal matrices [5, 29, 38],
Revised: Grassmann manifold of subspaces [5, 17, 38], Stiefel manifold of orthogonal matrices [5, 29, 38],

Original: Kendall shape spaces [68, 69, 81], hyperbolic spaces [27, 110, 111], and special Euclidean and orthogonal group [44, 100, 104], 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 [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.
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 [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].
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 a feature is absent in existing Riemannian optimization libraries [18, 24, 72, 80, 82, 103, 109].
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 [25, 42].
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, 22, 43] 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, 22] 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 [21, 114] 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 [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].
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 [34–37].
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, 11, 33, 39, 85].
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].

Original: 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].
Revised: 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 [94, 112] and more generally for empirical risk minimization problems on Riemannian manifolds [52].
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].

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

Original: More recently, [10] has proposed efficient sampling procedures from the tangent Gaussian distribution for large scale and stochastic optimization scenarios.
Revised: 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 [25, 42] is recently introduced machine learning framework which support automatic differentiation capabilities [15] 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 [48] 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 has 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 [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.
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 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.
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().
Revised: • Automatic and efficient vectorization with vmap().

Original: 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: 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.g.,
Revised: This is useful in many cases, e.g.,

Original: [Table]
Revised: [Table]

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

Original: • Per-example gradient clipping.
Revised: • Per-example gradient clipping.

Original: A key process in differentially private optimization is perexample gradient clipping 1 n (cid:80) ni
Revised: 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 [79], Pytorch [89] 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 [47, 75, 96] or come up with algorithmic modifications [26] 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 [107].
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) paradigm.
Revised: • Single Source Multiple Devices (SSMD) paradigm.

Original: JAX follows the SSMD paradigm, and therefore, the code written for CPUs can be run on GPU/TPUs without any additional 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 is available at https
Revised: Rieoptax is available at https

Original: /github.com/
Revised: /github.com/

Original: SaitejaUtpala/Rieoptax
Revised: SaitejaUtpala/Rieoptax

Original: 2. Design and 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.

Original: 2.1. Core
Revised: 2.1. Core

Original: • rieoptax.core.
Revised: • rieoptax.core.

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

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

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

Original: 2.2. Geometries
Revised: 2.2. Geometries

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

Original: 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() .
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, 22, 111].
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 [90]: 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 [12]: 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: the Poincare-ball model of Hyperbolic space with Poincare metric [111], 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: 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 

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

Original: + u 1 v 1 + · · · u d − 1 v d − 1 .
Revised: + u 1 v 1 + · · · u d − 1 v d − 1 .

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

Original: 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 ) .
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: 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 ) .
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.3. Optimizers
Revised: 2.3. Optimizers

Original: Optimizers module contains Riemannian optimization algorithms.
Revised: Optimizers module contains Riemannian optimization algorithms.

Original: Design of optimizers follows
Revised: Design of optimizers follows

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

Original: Where every optimizer
Revised: Where every optimizer

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

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

Original: • riepotax.optimizers.first order.rsrg: Riemannian stochastic recursive gradient descent [66].
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 [67].
Revised: Riemannian adaptive stochastic gradient algorithm [67].

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

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

Original: 2.4. Privacy 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: the Riemannian Laplace mechanism [94] 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: the Log-Euclidean mechanism [112] 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: noise calibration for differentially private Riemannian gradient descent [52] based on moments accountant [2] in the autodp library [113].
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: noise calibration for Differentially private Riemannian stochastic gradient descent [52] based on moments accountant [2] in autodp library [113].
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 [89] based Mctorch [80] and Geoopt [72], Tensorflow [103]
Revised: These include Pytorch [89] based Mctorch [80] and Geoopt [72], Tensorflow [103]

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

Original: (Tf-Riemopt) [103], Numpy [55]
Revised: (Tf-Riemopt) [103], Numpy [55]

Original: based Pymanopt [109], and Tensorflow based Geomstats [82].
Revised: based Pymanopt [109], and Tensorflow based Geomstats [82].

Original: While Geomstats supports Numpy, Pytorch, and Tensorflow as backend, currently only the Tensorflow backend provides support for GPUs.
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 [24] in Matlab and Manopt.jl [18] in Julia [19].
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 the AMD EPYC 7B1 processor 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: on 16GB Tesla P100.
Revised: on 16GB Tesla P100.

Original: Hypersphere: hypersphere S ( d ) is supported in Geoopt, Tf-Riemopt, Geomstats, McTorch, and Pyamanopt.
Revised: Hypersphere: hypersphere S ( d ) is supported in Geoopt, Tf-Riemopt, Geomstats, McTorch, and Pyamanopt.

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

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

Original: We benchmark all expect McTorch and Geomstats for dimensions d ∈ { 50 , 100 , 500 , 1 000 , 5 000 , 10 000 , 25 000 , 50 000 } .
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: the Loretnz hyperboloid model H ( d ) is supported in Geoopt, Tf-Riemopt, Geomstats, and Mctorch.
Revised: • Lorentz hyperboloid model: the Lorentz hyperboloid model H ( d ) is supported in Geoopt, Tf-Riemopt, Geomstats, and Mctorch.

Original: While the Exp map is available in Mctorch, it does not implement the Log 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 , 1 000 , 5 000 , 10 000 , 25 000 , 50 000 } .
Revised: We benchmark for dimensions d ∈ { 50 , 100 , 500 , 1 000 , 5 000 , 10 000 , 25 000 , 50 000 } .

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

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

Original: McTorch
Revised: McTorch

Original: does not support the Exp and Log maps.
Revised: does not support the Exp and Log maps.

Original: Geomstats represents
Revised: Geomstats represents

Original: Grassmann elements in projector matrices form XX ⊤ ∈ 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, therefore, exclude these three libraries from benchmarking.
Revised: We, therefore, exclude these three libraries from benchmarking.

Original: We benchmark for matrix sizes ( m, r ) ∈ { (100 , 10) , (500 , 10) , (750 , 10) , (1 000 , 10) , (2 000 , 10) , (5 000 , 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 the affine-invariant metric is supported in Geoopt, Tf-Riemopt, Geomstats, and McTorch.
Revised: • SPD with affine-invariant metric: SPD manifold SPD( m ) with the affine-invariant metric is supported in Geoopt, Tf-Riemopt, Geomstats, and McTorch.

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

Original: We benchmark all except McTorch for matrix sizes 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 improvements, especially on GPUs.
Revised: Overall, we observe that Rieoptax offers significant time improvements, especially on GPUs.

Original: For the 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 would 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.

Original: [Figure]
Revised: [Figure]

Original: 4. An example on PCA
Revised: 4. An example on PCA

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

Original: [Equation]
Revised: [Equation]

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

Original: The Rieoptax implementation for solving the problem (2) is shown in Listing 3.
Revised: The Rieoptax implementation for solving the problem (2) is shown in Listing 3.

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

Original: We take 5 000 images and vectorize each image to produce a sample matrix of size ( n, d ) =
Revised: We take 5 000 images and vectorize each image to produce a sample matrix of size ( n, d ) =

Original: We compute the top r = 5 principal components, which leads to an optimization problem on G (12 288 , 5) .
Revised: We compute the top r = 5 principal components, which leads to an optimization problem on G (12 288 , 5) .

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

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

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

Original: 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: 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=

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

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

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

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

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

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

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

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

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

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

Original: # non private PCA lr, epochs =
Revised: # non private PCA lr, epochs =

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

Original: [Table]
Revised: [Table]

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

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

Original: 5. Conclusion and future roadmap
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 architectures.
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, especially with differential privacy
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