Title: Riemannian SAM: Sharpness-Aware Minimization on Riemannian Manifolds
Venue: NeurIPS 2023
Accepted: yes
Keywords: optimization

riemannian

manifolds

sharpness-aware
Abstract: Contemporary advances in the field of deep learning have embarked upon an exploration of the underlying geometric properties of data, thus encouraging the investigation of techniques that consider general manifolds, for example, hyperbolic or orthogonal neural networks. However, the optimization algorithms for training such geometric deep learning models still remain highly under-explored. In this paper, we introduce Riemannian SAM by generalizing conventional Euclidean SAM to Riemannian manifolds. We successfully formulate the sharpness-aware minimization on Riemannian manifolds, leading to one of a novel instantiation, Lorentz SAM. In addition, SAM variants proposed in previous studies such as Fisher SAM can be derived as special examples under our Riemannian SAM framework. We provide the convergence analysis of Riemannian SAM under a less aggressively decaying ascent learning rate than Euclidean SAM. Our analysis serves as a theoretically sound contribution encompassing a diverse range of manifolds, also providing the guarantees for SAM variants such as Fisher SAM, whose convergence analyses are absent. Lastly, we illustrate the superiority of Riemannian SAM in terms of generalization over previous Riemannian optimization algorithms through experiments on knowledge graph completion and machine translation tasks.

Riemannian SAM: Sharpness-Aware Minimization on Riemannian Manifolds

Abstract

Contemporary advances in the field of deep learning have embarked upon an exploration of the underlying geometric properties of data, thus encouraging the investigation of techniques that consider general manifolds, for example, hyperbolic or orthogonal neural networks. However, the optimization algorithms for training such geometric deep models still remain highly under-explored. In this paper, we introduce Riemannian SAM by generalizing conventional Euclidean SAM to Riemannian manifolds. We successfully formulate the sharpness-aware minimization on Riemannian manifolds, leading to one of a novel instantiation, Lorentz SAM. In addition, SAM variants proposed in previous studies such as Fisher SAM can be derived as special examples under our Riemannian SAM framework. We provide the convergence analysis of Riemannian SAM under a less aggressively decaying ascent learning rate than Euclidean SAM. Our analysis serves as a theoretically sound contribution encompassing a diverse range of manifolds, also providing the guarantees for SAM variants such as Fisher SAM, whose convergence analyses are absent. Lastly, we illustrate the superiority of Riemannian SAM in terms of generalization over previous Riemannian optimization algorithms through experiments on knowledge graph completion and machine translation tasks.

1 Introduction

Deep learning incorporating the underlying geometry of data, referred to as geometric deep learning (GDL), has emerged as a significant research area in recent years due to its remarkable capability to effectively capture intrinsic structural properties. As a significant direction in this line, hyperbolic representation learning has been shown to offer several advantages over conventional Euclidean geometry. For example, the hyperbolic space allows for more efficient representations of highdimensional data by offering a more flexible and natural way to model hierarchical structures, which are commonly encountered in network embeddings [1, 2, 3], computer visions [4, 5], and natural language processing [6, 7]. Adding to the fascinating array of approaches in geometric deep learning, the orthogonal neural networks enforcing the orthogonality in model parameters emerge as a promising research area. In another dimension, the orthogonal neural networks that constrain the model parameter to satisfy the orthogonality (known as Stiefel manifold) are proposed [8, 9, 10] under the motivation that they prevent the vanishing/exploding gradient problem and theoretically enhance the model generalization [11]. Furthermore, more generally, the Riemannian extension of deep learning technique on Euclidean space continues to be proposed in many fields including Riemannian normalizing flows [12, 13], Riemannian diffusion models [14], Poincaré ResNet [15], hyperbolic deep reinforcement learning [16].

Along with the attempts to learn non-Euclidean representations in deep learning, Riemannian optimization has also been greatly studied to train non-Euclidean deep models. As a pioneering example, Riemannian (stochastic) gradient descent (R(S)GD) [17] is an extension of (stochastic) gradient descent to Riemannian manifolds, which updates the gradient computed on (a random subset of) the

training data at each iteration and then projects the gradient onto the tangent space of the manifold before taking a descent step. Starting with Riemannian gradient descent, several popular optimization algorithms in Euclidean space, such as conjugate gradient and trust region, have been generalized to Riemannian manifolds [18, 19, 20, 21, 22, 23, 24]. In addition to these previous works, there have been many studies to incorporate the momentum and variance reduction technique into the Riemannian manifolds. In this line of work, many optimization algorithms are proposed including the stochastic variance reduction scheme (R-SVRG) [25], stochastic recursive momentum (R-SRG) [26], and stochastic path-integrated differential estimator (R-SPIDER) [27], extended from their Euclidean counterparts. Going beyond the first-order algorithms, Riemannian (quasi)-Newton methods [28, 29] are a family of second-order methods using a (quasi)-Newton approach on Riemannian manifolds. In the context of deep learning, Riemannian extensions of adaptive gradient methods (ex. AdaGrad/Adam/AMSGrad) are proposed [30, 31, 32], which combines the benefits of adaptive learning rate methods with the efficiency of Riemannian optimization techniques.

However, the exploration of optimization algorithms for training deep learning models on nonEuclidean geometry has been considerably limited, highlighting the need to generalize successful optimizers in Euclidean space to Riemannian manifolds. Under this motivation, we introduce a new class of optimization schemes on Riemannian manifolds. Toward this, as our motivating optimization algorithm, we consider the sharpness-aware minimization (SAM) [33] in Euclidean space, which efficiently improves model generalization by considering the underlying geometry of loss landscapes. With the great success of SAM, several SAM variants including adaptive SAM [34], Fisher SAM [35], Efficient SAM [36], and GSAM [37], are proposed in recent years. In this paper, we propose Riemannian SAM, a sharpness-aware minimization on Riemannian manifolds that can be applied to various manifolds. Our Riemannian SAM considers the sharpness of the loss defined on manifolds, thereby effectively improving model generalization. We believe that our framework could further bring out the potential of the Riemannian deep models and enable more accurate evaluation.

Our contributions are summarized as follows:

We introduce Riemannian SAM, a sharpness-aware minimization scheme on Riemannian manifolds. Under our framework, we present a novel instance, Lorentz SAM, mainly used in our empirical evaluations. Furthermore, one of the SAM variants, Fisher SAM, which considers the underlying distribution space of neural networks can be derived as a special example under our Riemannian SAM framework.

We provide the convergence analysis of Riemannian SAM. Our convergence analysis achieves the first-order optimal rate of SGD and we highlight the challenges in our analysis. We allow for a less aggressively decaying ascent learning rate than the condition in the convergence of Euclidean SAM. Also, we provide the convergence guarantee for the SAM variants such as Fisher SAM whose convergence proofs are absent.

We validate the Riemannian SAM on knowledge graph completion and machine translation tasks for hyperbolic neural networks. The state-of-the-art hyperbolic architecture equipped with our Riemannian SAM improves the performance of the baselines trained with Riemannian Adam, which is a conventional optimizer in Riemannian deep learning.

2 Preliminaries

Before introducing the sharpness-aware minimization on Riemannian manifolds, we organize the necessary concepts and notations for Riemannian geometry and Riemannian optimization.

2.1 Riemannian Geometry for Optimization

We refer to the definitions in literature [38, 39, 40] where one can find more details.

Definition 1 (Riemannian manifold) . For each w ∈ M , let T w M denote the tangent space at w . An inner product on tangent space T w M is a bilinear, symmetric, positive definite function g w ( · , · ) := ⟨· , ·⟩ w : T w M× T w M→ R . If a metric ⟨· , ·⟩ w smoothly varies with w ∈ M , we call ⟨· , ·⟩ w a Riemannian metric . An induced norm on T w M is ∥ ζ ∥ w := √ ⟨ ζ, ζ ⟩ w . A Riemannian manifold is a pair ( M , g ) of the manifold M and the associated Riemannian metric tensor g .

Definition 2 (Geodesic) . A geodesic is a curve γ ( · ) : [0 , 1] →M that locally minimizes the distance between two points on a manifold with constant speed, which is the generalization of a straight line in Euclidean space.

Definition 3 (Exponential maps/Retraction) . An exponential map exp w : T w M→M maps a tangent vector ζ w ∈ T w M onto M along a geodesic curve such that γ (0) = w and γ (1) = z with ˙ γ (0) = ζ w . Specifically, γ ( t ) := exp w ( ζ w t ) represents a geodesic. A retraction R w ( · ) is a (computationally efficient) generalization of an exponential map satisfying the following properties:

R w (0) = w and D R w (0) = Id T w M where D R w represents the derivatives of R w and Id T w M denotes an identity map on T w M .

Definition 4 (Parallel translation/Vector transport) . A parallel translation P w z ( · ) : T z M→ T w M transport a tangent vector in T z M to T w M in parallel while preserving norm and direction (i.e., along a geodesic). A vector transport T ( γ ) w z ( · ) : T z M → T w M with respect to retraction map R maps a vector ζ z ∈ T z M to ζ w ∈ T w M along a retraction curve γ ( t ) = R w ( ξ w t ) for some ξ w ∈ T w M , which is computationally efficient approximation of a parallel translation. In this work, we only consider isometric vector transport, i.e., ⟨T w z ζ z , T w z η z ⟩ w = ⟨ ζ w , η w ⟩ w for all ζ w , η w ∈ T w M .

We introduce important examples of Riemannian manifolds in deep learning. The initial two instances represent dominant manifolds within the realm of hyperbolic deep learning. Hyperbolic space is a natural geometry for capturing underlying tree-like, graph-shaped, or hierarchical structures, which are properties existing in many real datasets. Owing to this characteristic, there have been many approaches to hyperbolic deep learning encompassing network embeddings [1, 2, 3, 41], computer vision [4, 5], and natural language processing [42, 6]. For manifolds in hyperbolic space, we mainly follow the definitions in [1, 2].

Poincaré Ball. The Poincaré ball P n = ( B n , g p ) is a Riemannian manifold with

where g e represents an Euclidean metric tensor. The associated distance on Poincaré ball is given by

We remark on some properties of the Poincaré ball. The Poincaré ball is a conformal model, meaning that angles between curves are preserved under conformal transformations. This property enables the Poincaré ball to accurately represent the local geometry of complex spaces. Additionally, the Poincaré ball offers an intuitive visualization of hyperbolic spaces, particularly in two or three dimensions. Despite these advantageous properties, the Poincaré ball also presents challenges in computing mathematical concepts such as geodesics and distances.

Lorentz Model. The Lorentz model L n = ( H n , g ℓ ) is a semi-Riemannian manifold, but it is still possible to employ Riemannian optimization. The Lorentz model consist of

where ⟨ x, y ⟩ L = -x 0 y 0 + ∑ n j =1 x j y j is known as the Lorentzian scalar product . The associated distance function is given by

Note that an n -dimensional Lorentz model requires one more redundant dimension in Euclidean space. Importantly, Poincaré ball and the Lorentz model are equivalent under the diffeomorphism φ : H n → P n (with the corresponding inverse mapping φ -1 : P n → H n ) defined as

We also remark on some characteristics of the Lorentz model. The main advantage is that it allows for more stable Riemannian optimization under relatively simple formulas for mathematical quantities, such as geodesics and distances. However, the visualization is difficult due to less intuitive projection onto a lower-dimensional space. As noted in each manifold, both two manifolds on hyperbolic space are equivalent, but they have different purposes. For this reason, some studies [2] use Lorentz manifolds for their model design and training, then visualize the results on Poincaré ball using the diffeomorphism φ in (2) and (3).

Stiefel manifold. The orthogonality, i.e., W T W = I for model parameter W , plays a role to circumvent the vanishing/exploding gradient problem [8, 10, 43] and delivers theoretically enhanced generalization error bounds [11]. The Stiefel manifold St( n, p ) = ( V n , g W ) for n ≥ p is prevalent for orthogonal neural networks, which is also a Riemannian manifold defined by

for tangent vectors Z 1 , Z 2 ∈ T W St( n, p ) . The tangent space of St( n, p ) at W is defined by T W St( n, p ) = { Z : Z T W + W T Z = 0 } .

2.2 Riemannian Optimization

We are interested in the following optimization problem over the Riemannian manifold ( M , g )

where L : M → R is a smooth function defined on manifold M . Following the work [17], Riemannian stochastic gradient descent (RSGD) updates the model parameter w ∈ M as

where grad L ( w t ) ∈ T w t M is a Riemannian gradient at w t and α t is the learning rate. In some practical cases, the exponential map is computationally inefficient. Hence, it may be replaced with (more computationally efficient) suitable retraction map R w t ( · ) , yielding the update rule w t +1 = R w t ( -α t grad L ( w t ) ) . Generally, the Riemannian gradient grad L ( w t ) in (4) is computed with the Riemannian metric tensor g as

where ∇L ( w t ) denotes the Euclidean gradient. This is also known as natural gradient descent [44]. The quantity on the right-hand side in (5) may not be on T w t M . In this case, we should project the gradient onto the tangent space since the exponential map is not defined.

3 Sharpness-Aware Minimization on Riemannian Manifolds

In empirical risk minimization (ERM) including deep learning tasks, we generally minimize the finite-sum objective (or equivalently expected objective) for training dataset D = { ( x i , y i ) } n i =1 as

where the smooth loss function L ( · ) : M → R is defined on a Riemannian manifold M . We formulate the sharpness-aware minimization in terms of loss function values on the manifold M as

In contrast to Euclidean space, the Riemannian manifold is not a vector space in general. Hence, the familiar concepts defined in Euclidean space may not be well-defined. Therefore, we restrict the perturbation δ in the tangent space T w M . To solve the inner subproblem, we resolve the inner optimization problem in a different manner as

for a fixed point w ∈ M . For ease of computations, we approximate the perturbed loss function L ( R w ( δ ) ) via Taylor's expansion as

Then, our inner maximization problem comes in hand:

This problem could be easily solved since it finds the steepest direction δ on Riemannian manifold M , whose solution is known to be just Riemannian gradient. Therefore, we have

Under the optimal perturbation δ ∗ in (11), we further approximate the gradient of the sharpness-aware minimization in (7) (for outer minimization problem) as

since the left-hand side of (12) requires a higher-order Riemannian gradient, which is not computationally feasible in practice. The remaining is that the approximated Riemannian SAM gradient grad L ( w ) | w =R w ( δ ∗ ) in (12) is not on the tangent space at w , T w M . To perform an actual parameter update on w , we should transport the Riemannian SAM gradient to the tangent space T w M via vector transport T w R w ( δ ∗ ) with respect to the retraction R w . We summarize the overall optimization procedure in Algorithm 1. Note that, in order to consider the most practical case, we assume that the same minibatch is used for computing SAM perturbation with the ascent step and the actual parameter updates with the descent step (see lines 5, 6, and 8). In fact, one can use different batches for lines 5 ∼ 7 and lines 8 ∼ 10 respectively or full-batch gradient for both ascent and descent steps.

Remarks on Algorithm 1. In fact, it might be most natural to choose a perturbation region at the current point as in the conventional Euclidean SAM, δ ∈ B ρ ( w t ) := { x ∈ M : d M ( w t , x ) ≤ ρ } where d M represents the distance on the manifold. However, adopting the constraint in this manner may pose challenges in utilizing the standard assumptions for analyzing non-convex Riemannian optimization, such as geodesic or retraction smoothness (see condition (C4 ) in Section 4), which makes it difficult to guarantee convergence. Moreover, the computation of d M is often computationally inefficient in practice. Another possible extension is to apply the vector transport operation from line 8 of Algorithm 1 to line 9. The following outlines the modified procedure: (i) g adv t = A ( grad L ( w ; S ) | w = w adv t ) and (ii) ∆ t = T w t w adv t g adv t . For base optimizer A , any optimization algorithm commonly used in Riemannian optimization can be adopted (e.g., Riemannian SGD). In the meanwhile, when the vector transport is applied after constructing g adv t via the momentum-based optimizer A , the momentum construction takes place on the tangent space T w adv t M at the perturbed point w adv t , while the parameter update occurs on the different tangent space at the point. As a consequence, this might introduce another challenges in understanding and analyzing the overall optimization process. In this perspective, various alternative extensions can also be possible, but among them, we have carefully designed a theoretically valid, computationally practical, and non-trivially extended Sharpness-Aware Minimization on general manifolds for Riemannian optimization . Then, we have successfully demonstrated both convergence analysis (Section 4) and empirical studies (Section 5) to corroborate our Riemannian SAM.

Existing example of Riemannian SAM framework: Fisher SAM. Following our Riemannian SAM update rule in Algorithm 1, we can show that Fisher SAM is a special instance of Riemannian SAM. We can view the set of neural networks as a neuromanifold [45, 46] equipped with the KL divergence metric between two points. Hence, let w ∈ M be the point on a neuromanifold (or statistical manifold) M which is realized by Euclidean network parameter θ ∈ R d . On distribution space, the corresponding metric tensor g is known to be Fisher information matrix [44, 45, 46]. According to Algorithm 1, the perturbation at line 6 and 7 could be computed as

Algorithm 1 Riemannian SAM: Sharpness-Aware Minimization on Riemannian Manifolds

1: Input: Descent learning rate α t , ascent learning rate ρ t , and a base Riemannian optimizer A (such as Riemannian SGD, Riemannian Adam).

2: Initialize: w 0 ∈ M

4: Draw a minibatch sample S = { x 1 , · · · , x |S| } .

3: for t = 0 , 1 , . . . , T -1 do

5: g t ← 1 |S| |S| ∑ i =1 grad L ( w t ; x i ) .

7:

6: δ t ← g t ∥ g t ∥ wt .

t

w

t

R

(

ρ

adv

t

8:

9:

w

g

t

adv

∆

t

adv

10:

t

w

+1

w

t

w

)

δ

t

grad adv

t

←A

(

R

11: end for

12: Output: w T

which is entirely identical to the perturbation of Fisher SAM [35].

Novel example: Lorentz SAM on hyperbolic geometry. Wederive the novel instance of Riemannian SAM called Lorentz SAM over the Lorentz model introduced in 2.1. First, we derive the Riemannian gradient on the Lorentz model L n = ( H n , g ℓ ) . As in Section 2.2, the Riemannian gradient could be computed as

Since g ℓ in (1) is a diagonal matrix, it is easy to compute the vector h with Euclidean gradient ∇L ( w ) . However, the vector h is not on the tangent space at w , T w L n , thus we should have to project the vector h onto the tangent space T w L n . The projection is easily computed in a closed-form as

Hence, the Riemannian gradient on Lorentz model is computed by

As a next step, we should normalize the Riemannian gradient as in line 6 in Algorithm 1, and this is easy to compute as ∥ grad L ( w ) ∥ w = √ ⟨ grad L ( w ) , grad L ( w ) ⟩ L via Lorentzian scalar product ⟨· , ·⟩ L defined in Section 2.1.

We utilize the Lorentz SAM derived above for our primary empirical studies conducted on hyperbolic space. In a similar way, one can derive the Riemannian SAM on Poincaré ball or Stiefel manifold, which we defer to Appendix. 0 . 27 0 . 27 0 . 27 30 0 . 33 0 . 33 0 . 36 0 . 39 0 . 39

0

.

48

.

30

3.1 Riemannian SAM Illustration: Toy 3D Illustration 0 . 0 . 51 0 . 57

.

0

69

33 0 . 45 0 . 54 0 . 60 0 . 63 0 . 66 0 . 72 We illustrate the 3-dimensional toy example on the sphere manifold. Let us define the 3D sphere manifold S 2 with the tangent space at w , T w M as

36

0

0

30

.

27

0

0

.

51

0

0

.

54

57

.

48

27

36

42

0

.

48

-

π

-

π

2

0

π

2

-

π

2

0

π

2

ϕ

Toy Example

0

.

27

0

.

27

0

.

30

0

.

30

0

.

33

0

.

33

0

.

36

0

.

36

0

.

39

0

.

39

0

.

42

0

.

42

0

.

45

0

.

45

0

.

45

0

.

45

.

45

0

.

51

0

.

51

0

.

51

0

.

.

54

0

.

57

0

.

57

0

.

57

0

.

0

.

60

0

.

60

0

.

60

0

.

60

0

.

63

0

.

63

0

.

63

0

.

63

0

.

63

0

.

66

0

.

66

0

.

66

0

.

66

0

.

66

0

.

69

0

.

69

0

.

69

0

.

69

0

.

72

0

.

72

0

.

72

0

.

75

Initial Point

Euclidean SAM

.

0

0

.

48

0

.

54

θ

0

.

36

0

.

0

.

60

57

.

39

33

0

.

.

51

0

.

33

48

0

0

27

0

.

.

0

.

.

0

0 . 36 ↔ ( r, θ, φ ) = (1 , θ, φ ) .

0

0 0 . 30 0 . 33 0 . 36 0 . 42 0 . 45 0 . 48 0 . 51 0 . 54 0 . 57 0 . 60 0 . 63 0 . 66 0 . 69 0 . 72 0 . 75 We consider the regression problem with the neural-netlike objective function on the randomly generated synthetic dataset. Toy 3D optimization problem with objective function f ( w ) = 1 2 n ∥ y -ReLU( Xw ) ∥ 2 2 where X ∈ R 500 × 3 and y ∈ R 500 are drawn from N (0 , 1 2 ) and U (0 , 1) respectively with the model parameter w = ( x, y, z ) ∈ R 3 under ∥ w ∥ 2 = 1 . (a) Comparison of converged points for each method. We plot the contour plots with the spherical coordinates under the relation ( x, y, z )

.

30

0

g

w

.

)

adv

(

-

α

.

L

t

w

(

∆

;

S

adv

t

)

t

w

)

.

|

t

=

w

adv

t

.

.

▷

▷ Stochastic Riemannian gradient

▷ Perturbation (11)

▷

SAM ascent step and transport to

t

t

w

Gradient at

w

T

▷ An update vector via a base optimizer A ▷ Final descent step

M

0

.

adv

33

.

0

0

.

54

0

0

.

54

.

0

36

0

.

66

0

72

.

π

60

39

0

.

30

.

0

0

36

.

0

48

.

0

57

.

45

0

.

30

0

.

0

33

.

0

.

51

63

.

0

.

0

54

0

.

36

.

0

0

.

27

33

0

.

0

.

33

30

0

.

48

0

.

0

54

.

0

0

.

51

57

0

.

.

0

60

.

69

72

0

.

0

0

.

0

66

.

63

45

.

27

0

.

36

0

.

48

0

.

0

51

0

.

54

57

0

.

63

0

.

.

66

.

0

75

0

0

72

.

0

.

69

0

0

.

27

.

60

.

45

.

0

0

20

40

60

80

100

Iteration

2

.

80

2

.

85

2

.

90

2

.

95

3

.

00

3

.

05

3

.

10

Tr(

H

(

θ

))

Euclidean SAM

(a) Trace of Hessian,

Tr(

∇

2

f

(

w

))

0

20

40

60

80

100

Iteration

0

.

00

0

.

01

0

.

02

0

.

03

0

.

04

0

.

05

0

.

06

0

.

07

‖

grad

f

(

θ

)

‖

θ

Manifold-Aware Sharpness

Euclidean SAM

(b) Manifold-Aware Sharpness,

∥

grad

f

(

w

)

∥

w

The Figure 1 corresponds to Cartesian coordinates ( x, y, z ) to spherical coordinates ( r, θ, φ ) = (1 , θ, φ ) , rendering contour plots. In Figure 1, we showcases the converged points on the objective function under the optimization using Riemannian SAM (in purple color) and the conventional Euclidean SAM (in pink color). Within a maximum iteration budget 100 , the purple point (Riemannian SAM) attains a loss value of 0 . 3800 while the pink point (conventional Euclidean SAM) converges with a slightly higher loss value of 0 . 3808 .

Furthermore, in terms of sharpness measures, we consider the following basic two quantities: (i) the trace of the Hessian (sharpness in the context of Euclidean space), and (ii) manifold-aware sharpness, characterized by the Riemannian gradient norm ∥ grad L ( w ) ∥ w . Notably, the manifoldaware sharpness aligns with information-geometric sharpness [47] when dealing with statistical manifolds, where the Riemannian metric is defined by the Fisher information. For the aforementioned problem, we compare two metrics and Figure 2 depicts the results. In both metrics, Riemannian SAM achieves smaller sharpness values than Euclidean SAM, implying convergence toward flatter regions. In other words, since the Euclidean SAM might fail to properly consider the underlying structure of the manifold even for toy examples, this phenomenon is expected to be exacerbated in extremely high-dimensional problems such as deep learning.

4 Convergence Analysis

In this section, we present the convergence guarantees for the RiemSAM framework. Our goal is to find a first-order ϵ -approximate solution: the output ˜ w such that E [ ∥ grad L ( ˜ w ) ∥ 2 ˜ w ] ≤ ϵ 2 , which is a generalized convergence criterion of Euclidean ϵ -stationary point. To guarantee the convergence for ϵ -approximate solution, we require the following mild assumptions.

(C1 ) (Upper-Hessian bounded) The objective function L is said to be upper-Hessian bounded in U ⊂ M with respect to retraction R if there exists some positive constant C such that d 2 L (R w ( tη )) dt 2 ≤ C , for all w ∈ U and η ∈ T w M with ∥ η ∥ w = 1 , and for all t such that R w ( τη ) for all τ ∈ [0 , t ] .

(C2 ) (Lower-bounded) The objective function L ( · ) is differentiable and has bounded suboptimality.

for the optimal point w ∗ ∈ M .

(C3 ) (Unbiasedness and bounded variance) The stochastic Riemannian gradient is unbiased and has a bounded variance:

where grad L ( w ) is a true Riemannian gradient evaluated on a full batch of training dataset D . (C4 ) (Retraction smoothness) We assume that there exists a constant L S > 0 such that

for all w,z ∈ M and γ ( t ) := R w ( tη ) represents a retraction curve on M for η ∈ T w M with the starting point γ (0) = w and the terminal point γ (1) = z .

(C5 ) (Individual Retraction Lipschitzness) We assume that there exists L R > 0 such that

for all w,z ∈ M . As in condition (C4 ), γ ( t ) denotes a retraction curve and T ( γ ) w z is a vector transport associated with this retraction curve.

The function class with condition (C1 ) corresponds to the continuous function family with Lipschitz continuous gradients in the Euclidean space [26, 32, 48]. The assumptions (C2 ) ∼ (C4 ) are standard in convergence analysis of Riemannian optimization algorithms, under which Riemannian SGD is known to be first-order optimal [49]. Note that, unlike in Euclidean space, the constant L S in (C4 ) and L R in (C5 ) may be different. According to [26, 50], the condition (C5 ) can be derived under the standard assumption on retraction Lipschitzness with parallel translation and one additional assumption on the bound between the parallel translation and the vector transport, but we assume the retraction Lipschitzness with the vector transport for simplicity. Lastly, in condition (C5 ), the retraction Lipschitzness is assumed individually with respect to each sample in order to control the alignment of SAM gradient and the original gradient step as in [51].

Now, we are ready to present our main theorem.

Theorem 1 (Convergence of Riemannian SAM) . Let ˜ w denote an iterate uniformly chosen at random from { w 1 , w 2 , · · · , w T } . Further, we let ˜ L = max { L S , L R } where the constants L S and L R are defined in condition (C4 ) and (C5 ) respectively. Under the conditions (C1 ) ∼ (C5 ) with descent learning rate α t = 1 √ T ˜ L and ascent learning rate ρ t = 1 T 1 / 6 ˜ L , we have the following complexity for the constant batch size b :

where ∆ = L ( w 0 ) -L ( w ∗ ) and the constants { Q i } 3 i =1 are irrelevant to the total iteration T or the manifold dimension d .

We make some remarks on our convergence results and relationship to conventional Euclidean SAM.

Theoretical implications. Our key observation of Theorem 1 lies in the alignment between the Riemannian gradient grad L ( w t ) (line 5 in Algorithm 1) and the Riemannian SAM gradient T w t w adv t grad L ( w adv t ) (line 9 in Algorithm 1) for the perturbed point w adv t , w adv t . The previous study [51] on Euclidean SAM says that Euclidean SAM gradient should be well-aligned with the true gradient step for convergence. Unlike the theoretical claim in [51], we stress that for convergence guarantee those gradients should be well-aligned within the preconditioned space (by inverse Riemannian metric) regardless of alignment in Euclidean space . To verify this insight, we directly measure the angles between two vectors with a 2D toy example, illustrating how they align in practice. Toward this, we consider two angles: (i) ∠ ( ∇ f ( w adv t ) , ∇ f ( w t )) (Euclidean Align-

0

5

10

15

20

Iteration

0

5

10

15

20

25

Degree between Vectors

On Statistical Manifold

(

T

w

t

w

adv

t

grad

f

(

w

adv

t

)

,

grad

f

(

w

t

)) (

Ours

)

(

∇

f

(

w

adv

t

)

,

∇

f

(

w

t

)) (Euclidean SAM)

/negationslash

/negationslash

ment) and (ii) ∠ ( T w t w adv grad f ( w adv t ) , grad f ( w t )) (Riemannian Alignment, Ours). In this example, t we consider the logistic regression where 200 data samples are generated with 100 of them sampled from N ( -1 , 1 2 ) and the remaining sampled from N (1 , 1 2 ) . The labels are assigned such that if a sample was drawn from a Gaussian distribution with a mean of -1 , the label was set to y = 0 , and otherwise, we set y = 1 . We minimize the cross-entropy loss with our Riemannian SAM with the Fisher information matrix as the Riemannian metric. The Figure 3 depicts the comparison of angles. The loss decreases up to 10 -th iteration, after which it remains around the converged point. As evident from the illustration, while the angles between the Euclidean space SAM gradient and the gradient deviate by up to around 25 degrees, the angles between the preconditioned SAM gradient and the preconditioned gradient, influenced by the Fisher information, align more closely with deviations

only up to a maximum of 10 degrees. In high-dimensional loss landscapes, we expect that the angles would become significantly larger, corroborating our theoretical insight.

On upper bound. Distinct from the convergence of Euclidean SAM [51], our upper bound (16) has the additional term involving Q 3 , but we still achieve the optimal complexity of SGD, O (1 /ϵ 4 ) for ϵ -approximate solution. Note that the presence of term involving the constant Q 3 in our bound comes from the fact that (i) smoothness condition (C4 ) and Lipschitzness condition (C5 ) are not equivalent on manifolds and (ii) we should handle the vector-transported gradients (see line 8 in Algorithm 1) at each iteration, which are the main challenges in our proof. Our results can also provide the guarantees for SAM variants such as Fisher SAM [35], whose convergence guarantees are missing.

5 Experiments

We conduct two sets of experiments; (i) knowledge graph completion, and (ii) machine translation. The first experiment aims to evaluate our Riemannian SAM on shallow networks and the second task is for optimizing large-scale deep neural networks. For all our experiments, we consider the Lorentz manifold introduced in Section 2.1 and employ the recent hyperbolic architecture, HYBONET [52]. The HYBONET is a fully hyperbolic neural network, whose each layer is constructed on the Lorentz manifold including a linear, attention, residual, and positional encoding layer. We implement our Riemannian SAM upon Geoopt framework [53] written in PyTorch library [54]. Regarding hyperparameters, we basically adhere to the same experiment settings in [52] and the details are provided in each section and Appendix.

5.1 Knowledge Graph Completion

Aknowledge graph completion aims to predict missing relationships within a knowledge graph, which represents structured information as a collection of entities, their attributes, and the relationships between them. More precisely, knowledge in a graph is of the form of triplets ( h, r, t ) where h , r , and t denote the head entity, the relationship or predicate, and the tail entity respectively. In the knowledge graph completion task, given a partially populated knowledge graph, the goal is to predict the missing entity or relationship in a triplet: solving ( h, r, ?) and (? , r, t ) .

In our experiments, we use two popular benchmark datasets; WN18RR [41] and Fb15k-237 [55]. We employ the same data preprocessing in [3] and two standard metrics for evaluations: (i) Mean Reciprocal Rank (MRR), the average of the inverse of the true entity ranking, and (ii) Precision at K (H@K), the proportion of test instances where the correct answer appears in the topK ranked predictions. For

marginal hyperparameter tuning, we tune the ascent learning rate ρ t ∈ { 10 -4 , 10 -3 , 10 -2 , 10 -1 } for Riemannian SAM and the other hyperparameters are the same as HYBONET [52] for fair comparisons.

Table 1 illustrates the results on WN18RR and Fb15k-237 datasets. As in previous work [52], we test our Riemannian SAM on two different regimes: (i) small embedding dimension 32 and (ii) large dimension β ∈ { 200 , 400 , 500 } for both datasets. Regarding the large dimension, we report the best results among the dimension candidates. In Table 1, Riemannian SAM on the Lorentz model achieves the best performance with great margins for all comparison metrics. In the same way, Riemannian SAM shows the state-of-the-art performance for all metrics considered under both regimes.

5.2 Machine Translation

In this experiment, we evaluate our Riemannian SAM on Lorentz Transformer built with Lorentz components introduced in [52] for IWSLT '14 and WMT '14 benchmark datasets in machine translation. We use the BLEU score as an evaluation metric on the IWSLT '14 test set and the newstest2013 test set of WMT '14 respectively. According to [56], we train hyperbolic models with Riemannian SAM in a low-dimensional setting where the dimension of the word vector is d = 64 . As in the knowledge graph completion task, we choose the ascent learning rate in ρ t ∈ { 10 -5 , 10 -4 , · · · , 10 -2 } for marginal hyperparameter tuning.

Table 2 demonstrates the results. HYBONET baseline trained with Riemannian Adam already outperforms the Euclidean Transformer in both IWSLT '14 and WMT '14 datasets. Upon this baseline, we only substitute our Riemannian SAM for Riemannian Adam with other hyperparameters unchanged. As seen in Table 2, Riemannian SAM significantly outperforms the Riemannian Adam baseline for both datasets. Note that both HYPERNN++ and HATT are partially hyperbolic networks, so we could not evaluate our Riemannian SAM on these models since it is difficult for a fair evaluation.

Wall-clock time. As in Euclidean SAM, Riemannian SAM requires additional forward and backward propagation in a single iteration loop (see Algorithm 1). Thus, we report the wall-clock time comparison for each experiment. For knowledge graph completion (Section 5.1) and machine translation (Section 5.2), Riemannian SAM takes roughly 1 . 6 and 1 . 8 times longer than Riemannian Adam for one epoch, respectively. To alleviate the computational overhead, one can employ the stochastic weight perturbation (SWP) and sharpness-sensitive data selection (SDS) suggested in [36], which do not depend on the manifold structure. Another practical consideration is to use a subset of minibatch in computing perturbation (see line 6 in Algorithm 1) for large-scale models. We leave the study on reducing computational cost as future work.

6 Conclusion

In this study, we proposed a sharpness-aware minimization on Riemannian manifolds, called Riemannian SAM. Under our framework, we presented novel examples of Riemannian SAM including a Lorentz SAM. We analyzed the convergence of the Riemannian SAM for general manifolds with a less aggressively decaying ascent learning rate condition. Moreover, we showed that Riemannian SAMcan provide the convergence guarantee for SAM variants whose convergence proofs are missing such as Fisher SAM. We also illustrated that Riemannian SAM empirically outperforms ERM-based Riemannian optimization algorithms for popular deep learning tasks with hyperbolic neural networks. As future work, we plan to study the technique to reduce the computations and analyze the generalization error bounds of Riemannian SAM theoretically.

References

Maximillian Nickel and Douwe Kiela. Poincaré embeddings for learning hierarchical representations. Advances in neural information processing systems , 30, 2017.

Maximillian Nickel and Douwe Kiela. Learning continuous hierarchies in the lorentz model of hyperbolic geometry. In International conference on machine learning , pages 3779-3788. PMLR, 2018.

Ivana Balazevic, Carl Allen, and Timothy Hospedales. Multi-relational poincaré graph embeddings. Advances in Neural Information Processing Systems , 32, 2019.

Keegan Lensink, Bas Peters, and Eldad Haber. Fully hyperbolic convolutional neural networks. Research in the Mathematical Sciences , 9(4):60, 2022.

Aleksandr Ermolov, Leyla Mirvakhabova, Valentin Khrulkov, Nicu Sebe, and Ivan Oseledets. Hyperbolic vision transformers: Combining improvements in metric learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 7409-7419, 2022.

Zhe Liu and Yibin Xu. Thg: Transformer with hyperbolic geometry. arXiv preprint arXiv:2106.07350 , 2021.

Boli Chen, Yao Fu, Guangwei Xu, Pengjun Xie, Chuanqi Tan, Mosha Chen, and Liping Jing. Probing {bert} in hyperbolic spaces. In International Conference on Learning Representations , 2021.

Martin Arjovsky, Amar Shah, and Yoshua Bengio. Unitary evolution recurrent neural networks. In International conference on machine learning , pages 1120-1128. PMLR, 2016.

Kyle Helfrich, Devin Willmott, and Qiang Ye. Orthogonal recurrent neural networks with scaled cayley transform. In International Conference on Machine Learning , pages 1969-1978. PMLR, 2018.

Jiayun Wang, Yubei Chen, Rudrasis Chakraborty, and Stella X Yu. Orthogonal convolutional neural networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 11505-11515, 2020.

Shuai Li, Kui Jia, Yuxin Wen, Tongliang Liu, and Dacheng Tao. Orthogonal deep neural networks. IEEE transactions on pattern analysis and machine intelligence , 43(4):1352-1368, 2019.

Mevlana C Gemici, Danilo Rezende, and Shakir Mohamed. Normalizing flows on riemannian manifolds. arXiv preprint arXiv:1611.02304 , 2016.

Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows. Advances in Neural Information Processing Systems , 33:2503-2515, 2020.

Chin-Wei Huang, Milad Aghajohari, Joey Bose, Prakash Panangaden, and Aaron C Courville. Riemannian diffusion models. Advances in Neural Information Processing Systems , 35:2750-2761, 2022.

Max van Spengler, Erwin Berkhout, and Pascal Mettes. Poincar \ 'e resnet. arXiv preprint arXiv:2303.14027 , 2023.

Edoardo Cetin, Benjamin Paul Chamberlain, Michael M. Bronstein, and Jonathan J Hunt. Hyperbolic deep reinforcement learning. In The Eleventh International Conference on Learning Representations , 2023.

Silvere Bonnabel. Stochastic gradient descent on riemannian manifolds. IEEE Transactions on Automatic Control , 58(9):2217-2229, 2013.

Antti Honkela, Tapani Raiko, Mikael Kuusela, Matti Tornio, and Juha Karhunen. Approximate riemannian conjugate gradient learning for fixed-form variational bayes. The Journal of Machine Learning Research , 11:3235-3268, 2010.

Nicolas Boumal. Riemannian trust regions with finite-difference hessian approximations are globally convergent. In Geometric Science of Information: Second International Conference, GSI 2015, Palaiseau, France, October 28-30, 2015, Proceedings 2 , pages 467-475. Springer, 2015.

P-A Absil. Trust-region methods on riemannian manifolds. Foundations of Computational Mathematics , 7:303-330, 2007.

Nilesh Tripuraneni, Nicolas Flammarion, Francis Bach, and Michael I Jordan. Averaging stochastic gradient descent on riemannian manifolds. In Conference On Learning Theory , pages 650-687. PMLR, 2018.

Orizon P Ferreira, Mauricio S Louzeiro, and LF4018420 Prudente. Gradient method for optimization on riemannian manifolds with lower bounded curvature. SIAM Journal on Optimization , 29(4):2517-2541, 2019.

Xiaojing Zhu and Hiroyuki Sato. Riemannian conjugate gradient methods with inverse retraction. Computational Optimization and Applications , 77:779-810, 2020.

Melanie Weber and Suvrit Sra. Riemannian optimization via frank-wolfe methods. Mathematical Programming , pages 1-32, 2022.

Hongyi Zhang, Sashank J Reddi, and Suvrit Sra. Riemannian svrg: Fast stochastic optimization on riemannian manifolds. Advances in Neural Information Processing Systems , 29, 2016.

Hiroyuki Kasai, Hiroyuki Sato, and Bamdev Mishra. Riemannian stochastic recursive gradient algorithm. In International Conference on Machine Learning , pages 2516-2524. PMLR, 2018.

Jingzhao Zhang, Hongyi Zhang, and Suvrit Sra. R-spider: A fast riemannian stochastic optimization algorithm with curvature independent rate. arXiv preprint arXiv:1811.04194 , 2018.

Zhi Zhao, Zheng-Jian Bai, and Xiao-Qing Jin. A riemannian newton algorithm for nonlinear eigenvalue problems. SIAM Journal on Matrix Analysis and Applications , 36(2):752-774, 2015.

Hiroyuki Kasai, Hiroyuki Sato, and Bamdev Mishra. Riemannian stochastic quasi-newton algorithm with variance reduction and its convergence analysis. In International Conference on Artificial Intelligence and Statistics , pages 269-278. PMLR, 2018.

Soumava Kumar Roy, Zakaria Mhammedi, and Mehrtash Harandi. Geometry aware constrained optimization techniques for deep learning. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition , pages 4460-4469, 2018.

Gary Becigneul and Octavian-Eugen Ganea. Riemannian adaptive optimization methods. In International Conference on Learning Representations , 2019.

Hiroyuki Kasai, Pratik Jawanpuria, and Bamdev Mishra. Riemannian adaptive stochastic gradient algorithms on matrix manifolds. In International Conference on Machine Learning , pages 3262-3271. PMLR, 2019.

Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations , 2021.

Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. Asam: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks. In International Conference on Machine Learning , pages 5905-5914. PMLR, 2021.

Minyoung Kim, Da Li, Shell X Hu, and Timothy Hospedales. Fisher sam: Information geometry and sharpness aware minimisation. In International Conference on Machine Learning , pages 11148-11161. PMLR, 2022.

Jiawei Du, Hanshu Yan, Jiashi Feng, Joey Tianyi Zhou, Liangli Zhen, Rick Siow Mong Goh, and Vincent Tan. Efficient sharpness-aware minimization for improved training of neural networks. In International Conference on Learning Representations , 2022.

Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha C Dvornek, sekhar tatikonda, James s Duncan, and Ting Liu. Surrogate gap minimization improves sharpness-aware training. In International Conference on Learning Representations , 2022.

John M Lee and John M Lee. Smooth manifolds . Springer, 2012.

P-A Absil, Robert Mahony, and Rodolphe Sepulchre. Optimization algorithms on matrix manifolds . Princeton University Press, 2008.

Nicolas Boumal. An introduction to optimization on smooth manifolds . Cambridge University Press, 2023.

Tim Dettmers, Pasquale Minervini, Pontus Stenetorp, and Sebastian Riedel. Convolutional 2d knowledge graph embeddings. In Proceedings of the AAAI conference on artificial intelligence , volume 32, 2018.

Caglar Gulcehre, Misha Denil, Mateusz Malinowski, Ali Razavi, Razvan Pascanu, Karl Moritz Hermann, Peter Battaglia, Victor Bapst, David Raposo, Adam Santoro, and Nando de Freitas. Hyperbolic attention networks. In International Conference on Learning Representations , 2019.

Asher Trockman and J Zico Kolter. Orthogonalizing convolutional layers with the cayley transform. In International Conference on Learning Representations , 2021.

Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation , 10(2):251-276, 1998.

Ovidiu Calin. Neuromanifolds , pages 465-504. Springer International Publishing, Cham, 2020.

Shun-ichi Amari, Hyeyoung Park, and Tomoko Ozeki. Geometrical singularities in the neuromanifold of multilayer perceptrons. Advances in neural information processing systems , 14, 2001.

Cheongjae Jang, Sungyoon Lee, Frank Park, and Yung-Kyun Noh. A reparametrization-invariant sharpness measure based on information geometry. Advances in neural information processing systems , 35:2789327905, 2022.

Andi Han and Junbin Gao. Riemannian stochastic recursive momentum method for non-convex optimization. In Zhi-Hua Zhou, editor, Proceedings of the Thirtieth International Joint Conference on Artificial Intelligence, IJCAI-21 , pages 2505-2511. International Joint Conferences on Artificial Intelligence Organization, 8 2021. Main Track.

Yossi Arjevani, Yair Carmon, John C Duchi, Dylan J Foster, Nathan Srebro, and Blake Woodworth. Lower bounds for non-convex stochastic optimization. Mathematical Programming , 199(1-2):165-214, 2023.

Andi Han and Junbin Gao. Variance reduction for riemannian non-convex optimization with batch size adaptation. arXiv preprint arXiv:2007.01494 , 2020.

Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware minimization. In International Conference on Machine Learning , pages 639-668. PMLR, 2022.

Weize Chen, Xu Han, Yankai Lin, Hexu Zhao, Zhiyuan Liu, Peng Li, Maosong Sun, and Jie Zhou. Fully hyperbolic neural networks. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 5672-5686, Dublin, Ireland, May 2022. Association for Computational Linguistics.

Max Kochurov, Rasul Karimov, and Serge Kozlukov. Geoopt: Riemannian optimization in pytorch. arXiv preprint arXiv:2005.02819 , 2020.

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems , 32, 2019.

Kristina Toutanova and Danqi Chen. Observed versus latent features for knowledge base and text inference. In Proceedings of the 3rd workshop on continuous vector space models and their compositionality , pages 57-66, 2015.

Ryohei Shimizu, YUSUKE Mukuta, and Tatsuya Harada. Hyperbolic neural networks++. In International Conference on Learning Representations , 2021.

Jonas Gehring, Michael Auli, David Grangier, Denis Yarats, and Yann N Dauphin. Convolutional sequence to sequence learning. In International conference on machine learning , pages 1243-1252. PMLR, 2017.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Supplementary Materials

A Proofs of Theorem 1

We basically follow the arguments in the convergence of Euclidean SAM [51], but the details are totally different. ∗ ( )

Lemma 1 (Properties of Retraction Smoothness) . Let w = R w ρ grad L ( w ) and the curve γ ( t ) = R w ( tη ) with the endpoints γ (0) = w and γ (1) = w ∗ . Then, we have

Proof. By the condition (C5 ), we have

By the Cauchy-Schwarz inequality, we have

Therefore, we obtain

Removing the constant ρ , the above inequality becomes

Lastly, we arrive at the final result as

In the next lemma, we will show the alignment of the true Riemannian gradient and the true Riemannian SAM gradient.

Lemma 2 (Alignment of the true Riemannian gradient and the true Riemannian SAM gradient) . Let us denote the stochastic Riemannian gradient at time t by grad L t ( w ) = 1 b ∑ i ∈ J t grad L ( w ; x i ) ∈ T w M and w adv = R w ( ρ grad L t ( w ) ) . Further, let γ ( t ) = R w ( tη ) be a retraction curve with γ (0) = w and γ (1) = w adv . Then, we have the following inequality

Proof. Let w ∗ = R w ( ρ grad L ( w ) ) evaluated on the loss function. We first add and subtract ⟨T ( ζ ) w w ∗ grad L t ( w ∗ ) , grad L ( w ) ⟩ w where ζ ( t ) = R w ( tξ ) is a retraction curve where ζ (0) = w and ζ (1) = w ∗ .

We will bound two terms, T 1 and T 2 , separately. Regarding the term T 1 , we derive

From the above inequality, we could finally bound the term T 1 as

Regarding the term T 2 , we just use the lemma as

Hence, we arrive at

According to Algorithm 1, we follow the notation as

We assume the stochastic m -SAM where the same batch is used for both inner and outer updates.

Lemma 3 (Descent inequality) . Under the assumptions in Theorem 1, we have

Proof. Using the condition (C4 ), we have

L ( w t +1 ) = L ( R w t ( -α T ( γ ) w t w adv t grad L t ( w adv t ) ) ) ≤ L ( w t ) -α 〈 grad L ( w t ) , T ( γ ) w t w adv t grad L t ( w adv t ) 〉 w t + α 2 L S 2 ∥ ∥ ∥ T ( γ ) w t w adv t grad L t ( w adv t ) ∥ ∥ ∥ 2 w t For the last term in RHS, we can bound as ∥ ∥ ∥ T ( γ ) w t w adv t grad L t ( w adv t ) ∥ ∥ ∥ 2 w t = -∥ grad L ( w t ) ∥ 2 w t + ∥ ∥ ∥ T ( γ ) w t w adv t grad L t ( w adv t ) -grad L ( w t ) ∥ ∥ ∥ 2 w t +2 〈 T ( γ ) w t w adv t grad L t ( w adv t ) , grad L ( w t ) 〉 w t

Again, we have

Taking the expectation on both sides, we have

For sufficiently large number of total iteration T , the condition ρ ≤ 1 4 ˜ L is easily satisfied where ˜ L = max { L R , L S } (defined in Theorem 1). Hence, we obtain

By telescoping the above inequality from t = 0 ∼ T -1 , we arrive at

Under the step size condition α t = 1 √ T ˜ L and ρ t = 1 T 1 / 6 ˜ L , we finally get

for appropriate constants { Q i } 3 i =1 .

B Hyperparameter Details

We use the almost same hyperparameters in the study [52] and implement our experiments in Section 5 upon its official implementation. For completeness, we summarize the hyperparameter configurations in Table 3 and Table 4.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper takes the popular idea of SAM from Euclidean space to the Riemannian space, and proposed a new optimization framework known as Riemannian SAM, which is a generalization of an existing technique. At the same time, the authors provided a convergence analysis of this newly proposed framework.\n\nSTRENGTHS:\n1. The authors proposed a general framework that is derived from first-principles, making this paper easy to follow, and provides users with sufficient background to understand.\n\n2. Algorithm 1 is presented nicely, and makes a lot of sense under the SAM framework.\n\n3. The convergence analysis is much needed, as it ensures readers that the new Riemannian SAM is not that much more worse than the Euclidean SAM. \n\n4. Experiments are convincing, and are performed appropriately.\n\nWEAKNESSES:\n1. We need a bit more motivation on why riemannian optimization needs SAM, as some optimization problems can be solved both in the euclidean space and the Riemanian space (say matrix sensing), and if they have similar guarantees, users would obviously avoid constrained optimization. My understanding is that for these problems, the additional structures that Riemannian optimization brings is already important, and SAM might not help a lot. I understand the authors tried to make this point by using the two examples, but it is better to have some theoretical explanations. Also this paper would benefit a lot from explaining why (7) is the best SAM formulation for Riemannian optimization, as it just seems like a direct transport form the euclidean problem.\n\n2. Not a lot of new insights and or proof techniques are proposed in this paper, as it seems like the proof procedures are very much alike the euclidean problem, and only adapted to the riemannian regime by assuming a chain of conditions (C-1) to (C-5). Therefore no specific new observations are made regarding the reimannian problem, and the authors are merely adapting the existing framework to \"make it work\" in riemannian space.\n\nQUESTIONS:\nN/A\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose an extension of the sharpness-aware minimization (SAM) technique on Riemannian manifolds, for which specific computations are feasible as the retraction map and the vector transport. The proposed optimization method considers the nonlinear geometry that is implied due to the manifold when performing a step. The convergence of the method is theoretically proven under assumptions, while empirical results are included to demonstrate the efficiency of the method.\n\nSTRENGTHS:\n- The extension of SAM to Riemannian manifolds is natural and meaningful.\n- Theoretical guarantees for convergence are provided.\n\nWEAKNESSES:\n- Even if the paper is in general well-written, there are parts where the presentation is quite high-level and thus becomes unclear. (See questions for specific examples)\n- Perhaps some figures could have been used to make the paper more accessible and easy to understand.\n- I think that the experimental section implicitly assumes familiarity with the associated techniques, which makes the paper not-self-contained. Perhaps, some details about the settings could have been included in the appendix e.g. the actual models and the corresponding parameter spaces, the retraction maps, the vector transports, etc.\n- As the performance of the proposed method seems to be close to the Riemannian Adam technique, I believe that at least some error-bars could have been included to better justify the difference in performance.\n\nQUESTIONS:\nQ1. It is not clear in many places if the manifold $\\mathcal{M}$ is considered as being embedded within an ambient (Euclidean) space or as (a subset of) $\\mathbb{R}^D$ together with a Riemannian metric as in Fisher SAM. In the embedded case, typically the Euclidean gradient is computed, which is then projected orthogonally to the tangent space followed by a retraction map. In the other case, the gradient is the Euclidean gradient multiplied by the inverse of the metric, while for the update step, simple addition is used.\n\nQ2. Related to Q1. The steps in lines 146-147 are quite confusing. Is it implied here that the actual manifold on which $\\mathcal{L}$ is defined is actually a surface embedded within an ambient Riemannian manifold where the metric $g(w)$ defines the structure? If this is the case, then the \"orthogonal\" projection on the tangent space should be defined with respect to the ambient Riemannian metric?\n\nQ3. Related to Q1 & Q2. In Eq 15, could you explain based on Q1 and Q2 why we need to use both the inverse of the metric and the projection to compute the gradient?\n\nQ4. Is the Stiefel manifold used somewhere as it is presented in line 130?\n\nQ5. In Eq. 12 it is mentioned that higher-order Riemannian gradient is necessary. Could you elaborate on which computation is infeasible? Due to the chain-rule the derivative of the retraction map with respect to the base point $w$?\n\nMinor:\n- Line 202-203: Something is missing at the end of the sentence.\n- Line 210: Probably $\\eta \\in \\mathcal{T}_w\\mathcal{M}$ .\n- Eq 15: probably the $\\nabla$ is missing.\n- Line 164: Capital S at the begging of the sentence.\n- Line 122: Is the function $\\phi$ an isometry?\n- Line 101: Only isometric vector transport implies that the angle between tangent vectors may change? Any influence on the optimization?\n\nLIMITATIONS:\nThe authors discuss some limitations of their work, such as the computational cost, while the standard limitation that comes with such geometric approaches is the access to the associated manifold operations e.g. retraction, vector transport, etc. As a theoretical work, it does not have any direct negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

comments: SUMMARY:
This paper introduce the new objective function SAM (Shapeness-awared minimization) on optimization problems on Riemannian manifolds. The motivation is that, with the success of SAM on the Euclidean space, the new Riemmanian metric, which is typically different from the flat Euclidean metric, can introduce more domain prior knowledge to SAM. Example applications include social network analysis and knowledge graph completion. The paper derived a new objective function based on retraction operations on manifolds, and provide convergence analysis of the Riemannian SGD algorithm. Empirical results show slight improvement over baseselines that are not based on the SAM objective function.

STRENGTHS:
The overall motivations make sense and the paper has made the contribution of combining SAM and optimization on Riemannian manifolds.

WEAKNESSES:
The main weaknesses are as follow:
1) The derivations of the new objective relies on standard tools from optimization on manifold and there is limited technical contribution.
2) Though SAM was published and proved useful, it is still not clear why minimizing a normalized norm of Riemannian gradient helps improve the test accuracy (line 190). The statement on lines 59-60 is not well supported. It is understandable that the main focus of this submission is not on renovating SAM, but some informal explanation can be helpful.
3) Only the best test accuracies are reported. The authors are encouraged to conduct cross-valiation and/or report average performance based on randomized experiments.
4) In my opinion, the improvement of the test accuracies is visible but not quite significant.

QUESTIONS:
line 101: should <\xi,\eta> be evaluated on the tangent space at z rather than w?
line 112: the angle is not formally defined.

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-10

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper extends the Sharpness-Aware Minimization (SAM) algorithm to Riemannian manifolds. They show that the new model subsumes Fisher SAM as a special case, and it leads to a new algorithm (called Lorentz SAM when specified to the Lorentz manifold). The main proposed algorithm is given in Algorithm 1, which works for general Riemannian manifolds, and a convergence theorem is provided for it (Theorem 1). The paper is concluded with some experiments showing the gain of the new algorithm in various settings.\n\nSTRENGTHS:\n- novel formulation allowing to extend SAM to Riemannian manifolds (Algorithm 1)\n\n- theoretically verified model (Theorem 1)\n\n- various experiments\n\nWEAKNESSES:\n- contributions of this paper are unfortunately limited to a few examples\n\n- the proposed algorithm mimics the original SAM after making necessary changes\n\n- Algorithm 1 is not well explained, though it's the main contribution of the paper\n\nQUESTIONS:\nAn interesting result on optimization on manifolds; but I'm concerned about the sufficiency of the contributions of the paper.  Besides that, the method is not well explained and the paper could've been better written. \n\n\n\n\nQuestions/Comments:\n\n- Line 117: \"Lorentz model is a Riemannian manifold\" This is wrong since the Lorentz manifold is a pseudo-Riemannian manifold, as the metric is not positive-definite. I'm wondering how this can change the algorithm's proof and applicability to the Lorentz manifold. \n\n\n\n\n- Line 145-147: This is only true when there is a nice extension of the function on the Euclidean space, and so might not be always feasible.\n  \n - Algorithm 1: the role of the base optimizer is not well explained. \n \n - Line 8: what is the role of the transportation? Why can't one update the parameters after the ascent step (then one does not need transportation?)? This part is also not well explained\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-17"}