Riemannian SAM: Sharpness-Aware Minimization on Riemannian Manifolds
Authors: Jihun Yun
Eunho Yang
Venue: NeurIPS 2023 poster (2023)
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
Jihun Yun KAIST
Eunho Yang KAIST, AITRICS eunhoy@kaist. ac.kr
arcprimeCkaist.et
Abstract
Contemporary advances in the field of deep learning have embarked upon exploration the underlying geometric properties of data, thus encouraging the in- vestigation 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 Rie- mannian manifolds. We successfully formulate the sharpness- aware minimization on Riemannian manifolds, leading to one of a novel instantiation, Lorentz SAM. addition, SAM variants proposed in previous studies such Fisher SAM can be derived as special examples under our Riemannian SAM framework. We provide the convergence analysis of Riemannian SAM under less aggressively decaying ascent learning rate than Euclidean SAM. Our analysis serves as a theoretically sound contribution encompassing 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 Riemannian SAM in terms of gener- alization over previous Riemannian optimization algorithms through experiments on knowledge graph completion and machine translation tasks.
Introduction
Deep learning incorporating the underlying geometry of data, referred to 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 significant direction in this line, hyperbolic representation learning has been shown to offer several advantages over conventional Fuclidean geometry. For example, the hyperbolic space allows for dimensional 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 5], and natural language processing [6. 7J. 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 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 Fuclidean 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].
efficient representations of high-
Along with the attempts learn -on-Euclidean representations in deep learning, Riemannian opti- mization has also been greatly studied to train mon-Euclidean deep models. As pioneering example, Riemannian (stochastic) gradient descent (R(S)GD) [17] is extension of (stochastic) gradient descent to Riemannian manifolds, which updates the gradient computed on (a random subset of) the
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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 Eu- clidean counterparts. Going beyond the first-order algorithms, Riemannian (quasi)-Newar methods approach on Riemannian
[28, 29] are a family of second-order methods using a manifolds. In the context of deep learning, Riemannian extensions of adaptive gradient methods (ex. AdaGraad/ Adam/ A MSGrad) 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 non- Euclidean 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. Riemannian SAM, a sharpness-awarr minimization on Riemannian manifolds that can be applied to various manifolds. Our Riemannian SAM considers the sharpness of the loss defined 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.
propose
Our contributions are summarized as follows:
We introduce Riemannian SAM, a sharpness-await minimization scheme on Riemannian manifolds. Under our framework, we present 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 Fuclidean 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 trans- lation 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 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 Tw denote the tangent space at w. An inner product on tangent space Tw is a bilinear, symmetric, positive definite function
gw( == (;:)) TwM x TwM + R. If a metric (-/) smoothly varies with w E M, we call ( a Riemannian metric. An induced norm on T is [|[| == VVC.C) Riemannian manifold is a pair (M,g) of the manifold M and the associated Riemannian metric tensor g.
2

---

Definition 2 (Geodesic). A geodesic is a curve 7(-): [0, + 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 expw T + M maps tangent vector Sw E T T onto .M along a geodesic curve such that 7(0) and 7(1) = with (0) Su- Specifically, 7 = expr.(Sut) represents a geodesic. A retraction (- (compulationally efficientt generalization of an exponential map salisfying the following properties:
(0) and (O) I.TT.. M where DRw represents the derivatives of R and I. denotes an identity map on T.
Definition 4 (Parallel translation/ Vector transport). A parallel translation (-) TTMM + T transport a tangent vector in T T&M to T T in parallel while preserving norm and direction (i.e.. along a geodesic). A vector transport (()) map R. maps a vector S2 E T.M to Sw E T along a retraction curve (t) = Rw (Ewt) ) for some Eu E Tw which is computationally efficient approximation of parallel translation. In
T.MM- + T with respect retraction
this work, we only consider isometric vector transport. i.e., (T." S2. ".a.T." = (Sws quo) w for all Suz Tl E T.
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 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 IP" = (B"", sp) is a Riemannian manifold with
2
B (I E ||||| < 1}, gp(x)
(z)
||
where ga represents an Euclidean metric tensor. The associated distance on Poincaré ball is given by
dp(2; v) arcosh (1 2 (1
|| vl
1
21 1b1
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 I (H", .e)) is cemi-Riemannian manifold, but it still possible to employ Riemannian optimization. The Lorentz model consist of
0
0
H fx R*** (x) 11 0}, ge(x)= =
where (x. y)2 = v0 + Σラ= IJbý is known as the Lorenzzian sealar product. The associated
distance function is given by
de(x,y) arcosh( (z, v)
Note that an n-dimensional Lorentz model requires more redundant dimension in Euclidean space. Importantly, Poincaré ball and the Lorentz model are equivalent under the diffeomorphism
: " - pm (with the corresponding inverse mapping 4p p - H defined
p( T1, ... ,n,) (21, F2,
2nn
(2)
Po
8 (21, ".. [.))= (n)) (1 + x
+ ..... + 222 2x1 ... 2x
(3)
- (x²
2)
3

---

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 equívalent, 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 &p in (2) and (3).
Stiefel manifold. The orthogonality, i.e. W I for model parameter W, plays role to circumvent the vanishing/exploding gradient problem [8, 10, 43] and delivers theoretically enhanced generalization error bounds [1]]. The Stiefel manifold St(n. p) (V" .gw) for 2 p prevalent for orthogonal neural networks, which is also Riemannian manifold defined by
4 = (X E R X l}. 9w (2 Z2) = Tr Z2).
for tangent vectors Z1. Z2 E Sw p). The tangent space of St(n, p) at W is defined by T Tw:St (Z Z + =
2.2 Riemannian Optimization
We are interested in the following optimization problem over the Riemannian manifold (M, g)
min C(w).
WEM
where L : MM- R a smooth function defined on manifold M. Following the work [17]. Riemannian stochastic gradient descent (RSGD) updates the model parameter E M as
Wt++1 = exp (
o.grad((w:))
(4)
where grad[(w;) E M is a Riemannian gradient at w, and 0g 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 Rw, (-). yielding the update rule Wz+1 Rw., (zgrad(((w)) Generally, the Riemannian gradient grad[(w)) in (4) computed with the Riemannian metric tensor g
gradC(w.) = gw:
(5)
where (((w,) 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 Tw, M. In this case, we should project the gradient onto the tangent space since the exponential map not defined.
Sharpness-Aware Minimization on Riemannian Manifolds
In empirical risk minimization (ERM) including deep learning tasks, generally minimize the finite-sum objective (or equivalently expected objective) for training dataset D ((I,, y)) as
1 n Σ xw)
min ((w)= = WEM
where the smooth loss function L(-) M R is defined Riemannian manifold M. We formulate the sharpness-awar minimization in terms of loss function values on the manifold M as
min WEM S
max -
- [(w)
Sharpness in termis of loss function
In contrast to Euclidean space, the Riemannian manifold is not vector space in general. Hence, the familiar concepts defined in Euclidean space may not be well-defined. Therefore, we restrict the perturbation 5 in the tangent space T M. To solve the inner subproblem, we resolve the inner optimization problem in a different manner as
max L (R)) L(w) such that ||8112 < p². TETw.M
(8)

---

for a fixed point w E M.. For ease of computations, we approximate the perturbed loss function C (6)) via Taylor's 's expansion as
(RRw ~ [(wr) + (gradC(w), 6)
(9)
Then, our inner maximization problem comes in hand:
max (gradC(w), 6) such that ||ōl| S p². S''
(10)
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
gradC(w)
8 P (w)
(11)
Under the optimal perturbation 8 in (11), we further approximate the gradient of the sharpness-awan minimization (7) (for outer minimization problem) as
gradd((R..
gradd(w))
(12)
since the left-hand side of (12) requires a higher-order Riemannian gradient, which is not compu- tationally feasible in practice. The remaining that the approximated Riemannian SAM gradient gradf(w) (8+) in (12) is not on the tangent space at w. T. To perform an actual parameter update on w, we should transport the Riemannian SAM gradient to the tangent space T..M via vector transport T (8+) with respect to the retraction R.. We summarize the overall optimization procedure in Algorithm 1. Note that, in order to consider the most practical case, 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, can use different batches and lines 8 ~ 10 respectively or full-batch gradient for both ascent and descent steps.
for lines
Remarks on Algorithm 1. In fact, it might be most natural to choose perturbation region at the cur- rent point as in the conventional Euclidean SAM, (p(w)) == {a E . (w,, x) S p) where d.M represents the distance the manifold. However, adopting the constraint in this manner may pose challenges in utilizing the standard assumptions for analy saalyzing non-convex Riemannian optimization, such geodesic or retraction smoothness (see condition (C-4) in Section 4), which makes it difficult to guarantee convergence. Moreover, the computation of d is often computationally inefficient in practice. Another possible extension is apply the vector transport operation from line of Algorithm i to line 9. The following outlines the modified procedure: (i) ado = A A(gradL(us (w; S) homardde) and (ii) A T T of Riemannian optimization can be adopted (c.g., Riemannian SGD). In the meanwhile, when the vector transport is applied after constructing dfdh via the momentum- based optimizer A, the momentum construction takes place on the tangent space T M M at the perturbed point wadv, while the pa- rameter update occurs on the different tangent space at the point. As consequence, this might introduce another challenges in understanding and analyzing the overall optimization process. 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-Aaae Minimization general manifolds for Riemannian optimization. Then, we have successfully demon- strated both convergence analysis (Section 4) and empirical studies (Section 5) to corroborate our Riemannian SAM.
gā For base optimizer A, any optimization algorithm commonly used in
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 neuromanifold [45, 46] equipped with the KL divergence metric between two points. Hence, let w E M be the point on neuromanifold (or statistical manifold) M which realized by Fuclidean network parameter E Rdd. 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 and 7 could be computed as
gradL(w) F(0) -V
grad[(w) 8 p Igrad((w)]l = - V VgradC(w)) F (()gra(((w)) VVC() F(0))~Z10
((C))
5

---

Algorithm Riemannian SAM: Sharpness-Awaid Minimization on Riemannian Manifolds
1: Input: Descent learning rate dt, ascent learning rate Pt. and a base Riemannian optimizer A (such as Riemannian SGD, Riemannian Adam).
2: Initialize: wo E MM
for = 0, ....., do
4: Draw minibatch sample S = {x1, 11 : sis.
5:
9t
gradL(wwza)) m.).
D Stochastic Riemannian gradient
6:
6,
ligelluee
> Perturbation (11)
7:
wgdv &- Ru, (prdz).
> SAM ascent step
8:
godv + - updee grad (aelCCus 8) (w: wwwwed.
b Gradient at ws and transport T T An update vector via base optimizer . Final descent step
9:
Andu + A(ggddo)
10: W++1 f- Rus Cr A,A0d Addv).
11:
end for
Output: WT
which is entirely identical the perturbation Fisher SAM [35].
Novel example: Lorentz SAM on hyperbolic geometry. We derive the novel instance of Riemannian SAM called Lorentz SAM over the Lorentz model introduced in 2.1. First, derive the Riemannian gradient on the Lorentz model L_" (HI", gr). As Section 2.2, the Riemannian gradient could be computed
Dt V(=).
(13)
Since de in (1) a diagonal matrix, is easy to compute the vector h with Euclidean gradient VL(w). However, the vector h is not on the tangent space w, Tw.L", thus should have to project the vector h onto the tangent space T L. The projection is casily computed in a closed-form as
(roj..() = v + (w, v)...
(14)
Hence, the Riemannian gradient on Lorentz model computed by
(15) As next step, we should normalize the Riemannian gradient as in line 6 in Algorithm 1, and this is easy to compute Ilgrad[(w)]le V(gradC(w) grad[(w)) (graad((()ggrad[((ee))
gradf(w) = proj VV
via Lorentzian scalar product
(1 defined in Section 2.1.
We utilize the Lorentz SAM derived above for our primary empirical studies conducted on hyperbolic space. In similar way, one derive the Riemannian SAM on Poincaré ball Stiefel manifold, which we defer to Appendix.
3.1 Riemannian SAM Illustration: Toy Illustration
Toy Esample
We illustrate the 3-dimensional toy example on the sphere manifold. Let us define the 3D sphere manifold S with the tangent space w. T M as
2 = {w R 3
|lll| 1),
s, = (u 3 NTwTo 0}
9
We consider the regression problem with the neural-net- like objective function on the randomly generated syn- thetic dataset. Toy 3D optimization problem with ob- ReLU(X w) Il where X E 500xx3 and E 300 are drawn from NO 12) and (1,, 1) respectively with the model parameter w (x,, y, =) E I under ///l 1. (a) Comparison of con- verged points for each method. plot the contour plots
jective function f(w) = 2y
2
0
0
Figure 1: Toy 3D illustration
with the spherical coordinates under the relation (z,y, v. ++ (r. 0, (1,0,pp.

---

Hessinn Trace
ManiFold-A Aood---wrrreee Sharpness
3.10-
0.07
RiemsnniinnAA
3.05
Euclidean SAM
0.06
3.00
0.05
0.04
0.033
0.02
2.85
RiemanminnnAA
0.01
2.80
Euclidlea SAM
0.00
20
60
80
100
40
60
80
100
Iteration
Iteration
(a) Trace of Hessian, TN ((w))
(b) Manifold-Aware Sharpness, [[grad((w)
Figure 2: Comparison of two sharpness measures.
The Figure 1 corresponds to Cartesian coordinates (z, (r, 2) to spherical coordinates (r, 0, yp = (1, 0, sp), 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 maximum iteration budget 100, the purple point (Riemannian SAM) attains loss value of 0.3800 while the pink point (conventional Euclidean SAM) converges with 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 [gradd((w)] Notably, the manifold- aware sharpness aligns with information-geoment manifolds, where the Riemannian metric is defined by the Fisher information. For the aforementioned problem, compare two metries 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 deep learning.
sharpness [47] when dealing with statistical
4 Convergence Analysis
In this section, we present the convergence guarantees for the RiemSAM framework. Our goal is to find a first-order e-approximate solution: the output こ such that E (() I = which is generalized convergence criterion of Euclidean e-stationary point. To guarantee the convergence for e-approximate solution, require the following mild assumptions.
(C-1) (Upper-Hessiee bounded) The objective function L is said to be upper-Hlessian bounded in U C. MM with respect to retraction R. if there exists some positive constant C such that & (t)) < C, for all w E l and r E T with || = 1, and for all such that Rw dez for all 7 t0.
(C-2) (Lower-bounded) The objective function C(-) is differentiable and has bounded suboptimality. L(w > -00.
for the optimal point w* M. (C-3) (Unbiasedness and bounded variance) The stochastic Riemannian gradient unbiased and has a bounded variance:
E blw.eelgradd .=] grad((w).
E ( [ [llgrad[( x) - (w)dd((w)le)) 02.
where gradC(w) is true Riemannian gradient evaluated on a full batch of training dataset (C-4) (Retraction smoothness) We assume that there exists a constant Ls > such that
1 Ls slll
[(x)= < L(w) + (grad[(w), n)u, + 2
7

---

for all w, 2 E M and 7() = (ty(t represents a retraction curve on M for E Tw with the starting point 7(0) = w and the terminal point y(1) = 2.
(C-5) (Individual Retraction Lipschitzness) We assume that there exists Ln > 0 such that gradC(w; S LR|
7 [7))
for all w, M. As in condition (C-4), rt) denotes retraction eurve and T(7) is vector transport associated with this retraction curve.
The function class with condition (C-1) corresponds to the continuous function family with Lipschitz continuous gradients in the Euclidean space [26, 32, 48). The assumptions (C-2) (C-4) 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 Lg in (C-4) and LR in (C-5) may be different, According to [26, 50], the condition (C-5) can be derived under the standard assumption on retraction I 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 (C-5), 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 {wi, W2, WT) }. Further, we let L = max Lg, Lu} where the constants Lg and Ln are defined in condition (C-4) and (C-5) respectively. Under the conditions (C-1) (C-5) with descent learning rate 04 <TI and ascent learning rate Pt TaseT
we have the following complexity for
the constant batch size b:
Q VT
Q20²
Q30²
E [[]grad((M)1)
UT5/6 O(I/VT).
(16)
byT
where A = [(wo) L(w") and the constants (Q,) are irrelevant to the total iteration T or the manifold dimension d.
We make some remarks on convergence results and relationship to conventional Euclidean SAM.
Theoretical implications. Our key observation of Theo- rem 1 lies in the alignment between the Riemannian gradi- ent grad (line in Algorithm 1) and the Riemannian SAM gradient (ergg.grrC[[[g (w(d))) (line 9 in Algorithm 1)
On Statistionl Manifold
25-
4M2 (ord" (Ours)
/WFyophh) (rool) [Budalen 300
Degror betwoon Vorrar
for the perturbed point widdo, w. wadee The previous study [51] on Euclidean SAM says that Euclidean SAM gradient should be well-aligned with the true gradient step for con- vergence. Unlike the theoretical claim in [5]], we stress that for convergence guarantee those gradients should be well-alignes within the preconditioned space (by inverse Riemannian metric) regardless of alignment in Euclidean space. To verify this insight, we directly measure the an- gles between two vectors with 2D toy example, illustrat- ing how they align in practice. Toward this, we consider two angles: (i) L(Vf(w f((w)) (Euclidean Align- ment) and (ii) <(T A..add f g (Riemannian Alignment, Ours). In this example, we consider the logistic regression where 200 data samples are generated with 100 of them sampled from (V(-1, 12) and the remaining sampled from (11, 12). The labels are assigned such that if sample was drawn from a Gaussian distribution with mean of 11, the label was set to O, 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 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 Fuclidean 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 closcly with deviations
10
15
Iteration
Figure 3: Toy 2D illustration
8

---

Table 1: Link prediction results ) in the filterated setting for WN18RR and FB15k-237 datasets. For hyperbolic architectures, 3 E {200, 400, 500} and we report the best result. The results baselines taken from [52] except for HYBONET whose results reproduced by ourselves. The best results among hyperbolic architectures with the same dimensions are in boldface. Our Riemannian SAM (denoted by rSAM) shows the superior performance compared to Riemannian Adam (denoted by rAdam).
WN18RR
FBISk-237
Manifold
Model
#Dims
MRR
H@ 10
1103
11@1
#Dims
MRR
11@ 10
H603
H@
MURP
32
46.5
54.4
48.4
42.0
32
32.3
50.1
35.3
23,5
ROTH
32
47.2
55.3
49.0
42.8
32
31.4
49.7
34.6
22.3
HYBBNET
48.8
55,4
50,3403
45,5=03
32
33.5
51,5
36,5493
24.2
w/ ISAM
HYBONET
32
29.3
56.0
50,7
46.2
32
34.3
$2.0
37.3
25.1 2M
Hyperbolic
MURP
48.1
56.6
49.5
44.0
33.5
51.8
36.7
24.3
ROTH
0
49.6
38.6
51.4
44,9
34.4
53.5
38.0
24.6
38.4
HYBBONET
51.2#0.2
$7.1
$2.540.2
48.3
35.2
52.9
38.7
26.3
HYBONET
8 $1.6
58,7
53.3
48.6
36.0 +02
54,3
39.6
26.6
w/ ISAM
only 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 Q3. but still achieve the optimal complexity of SGD, (L1/c*) for e-approximate solution. Note that the presence term involving the constant Q3 in our bound comes from the fact that (i) smoothness condition (C-4) and Lipschitzness condition (C-5) are not equivalent on manifolds and (ii) we should handle the trctor-rannppotted gradients (see line 8 in Algorithm 1) at each iteration, which 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 evaluate 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 linear, attention, residual, and positional encoding layer. We implement our Riemannian SAM upon Geoopt framework [53] written in Py Torch library [54]. Regarding hyperparameters, we basically adhere the same experiment settings [52] and the details are provided in each section and Appendix.
5.1 Knowledge Graph Completion
knowledge graph completion aims to predict missing relationships within a knowledge graph. which represents structured information collection of entities, their attributes, and the relationships between them. More precisely, knowledge in a graph is the form of triplets (h, r, where h, r. and t denote the head entity, the relationship knowledge graph completion task, given a partially populated knowledge graph, the goal to predict the missing entity relationship in triplet: solving (h, r, ?) and (?, r,t).
predicate, and the tail entity respectively. In the
In our experiments, we use two popular benchmark datasets; WNISRR [41] and Fb15k-237 [55]. We employ the same data preprocessing in [3] and two standard metries for evaluations: (i) Mean Recipro- cal Rank (MRR), the average of the inverse of the true entity ranking, and (ii) Precision K (H@ K), the proportion of test instances where the correct answer appears in the top- ranked predictions. For

---

Table 2: The BLEU scores the test of IWSLT '14 and WMT '14 under low-dimensional setting following the hyperbolic study [56] with the word vector dimension d = 64. The results on baselines taken from [52] except for HYBONET whose results reproduced by ourselves.
Manifold
Model
IWSLT '14
WMT '14
CoNvSEo2SEQ [57]
23.6
14.9
Euclidean
TRANSFORMER [58]
23.0
17.0
HYPERNN++ [56]
22.0
17.0
HATT [42]
23.7
18.8
Hyperbolic
HYBONET (with Riemannian Adam)
25.5
19.3
HYBONET (with Riemannian SAM, Ours)
26.0
20.1
marginal hyperparameter tuning, we tune the ascent learning rate Pt E {10-1, $10-1,10-3 111111100-.,........1111111 11102,10011 10,10-13 for Riemannian SAM and the other hyperparameters are the same as HYBONET [52] for fair comparisons.
Table illustrates the results on WN18RR and 5b15k-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 8 E {200, 400, 500) for both datasets. Regarding the large dimension, we report the best results among the dimension candidates. In Table Riemannian SAM on the Lorentz model achieves the best performance with great margins for all comparison metries. In the same way, Riemannian SAM shows the state-of-the-ar 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], 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 Pr {10-5,10* .......,, 10-2 } for marginal hyperparameter tuning.
Table 2 demonstrates the results. HYBONET baseline trained with Riemannian Adam already outper- forms the Euclidean Transformer in both IWSLT 14 and WMT "4 datasets. Upon this baseline, we only substitute our Riemannian SAM for Riemannian Adam with other hyperparameters unchanged. As seen 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 single iteration loop (see Algorithm 1). Thus, we report the wall-clock time compar- ison 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 employ the stochastic weight perturbation (SWP) and saurpness-sensitiie data selection (SDS) suggested in [36], which do not depend on the manifold structure. Another practical consideration is 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 Rieman- nian SAM. Under our framework, we presented novel examples of Riemannian SAM including 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 SAM can 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 net- works. As future work, we plan to study the technique to reduce the computations and analyze the generalization error bounds of Riemannian SAM theoretically.
10

---

Acknowledgement
This work was supported by the National Research Foundation of Korea (NRF) grants (0.2008RIASAAISS RS-2023-00209060), Institute of Information & Communications Technol- ogy Planning Evaluation (IITP) grants (No. N0.02019-0000055, Artificial Intelligence Graduate School Program(KAIST) funded by the Korea government (MSIT).
References
[1] Maximillian Nickel and Douwe Kiela. Poincaré embeddings for learning hicrarchical representations. Advances in neural information processing systems, 30, 2017.
[2] Maximillian Nickel and Douwe Kiela. Learning continuous hierarchies the lorentz model of hyperbolic geometry. In International conference on machine learning, pages 3779-3788. PMLR, 2018.
[3] Ivana Balazevic, Carl Allen, and Timothy Hospedales. Multi-rclational poincaré graph embeddings. Advances in Neural Information Processing Systems, 32, 2019.
[4] Keegan Lensink, Bas Peters, and Eldad Haber. Fully hyperbolic convolutional neural networks. Research in the Mathematical Sciences, (4))600, 2022.
[5] 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.
[6] Zhe Liu and Yibin Xu. Thg: Transformer with hyperbolic geometry. arXiv preprint arXiv:2106.07350. 2021.
[7] Boli Chen, Yao Fu, Guangwei Xu, Pengjun Xie, Chuanqi Tan, Mosha Chen, and Liping Jing. Probing (hert) in hyperbolic spaces. In International Conference on Learning Representations, 2021.
[8] 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 Yc. Orthogonal recurrent neural networks with scaled cayley transform. In International Conference on Machine Learning, pages 1969-1978. PMLR, 2018.
[10] Jiayun Wang. Yubei Chen, Rudrasis Chakraborty, and Stella X Yu. Orthogonal convolutional neural networks. In Proceedings of the IEFF/CVF conference on computer vision and pattern recognition, pages 15505-1155,, 2020.
[11] Shuai Li, Kui Jia, Yuxin Wen, Tongliang Liu, and Dacheng Tao. Orthogonal deep neural networks. JEEE transactions on pattern analysis and machine intelligence, 43(4):1352 1368, 2019.
[12] Mevlana C Gemici, Danilo Rezende, and Shakir Mohamed. Normalizing flows on riemannian manifolds. arXiv preprint arXiv:1611.02.30 2016.
[13] Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows. Advances in Neurul Information Processing Systems, 33:2503-251 2515, 2020.
[14] Chin-We Rein-Wei Huang, Milad Aghajohari, Joey Bose, Prakash Panangaden, and Aaron Courville. Riemannian diffusion models. Advances in Neural Information Processing Systems, 35:2750-2761, 2022.
[15] Max van Spengler, Erwin Berkhout, and Pascal Mettes. Poincar) 'e resnet. arXiv preprint 14024 2023.
[16] Edoardo Cetin, Benjamin Paul Chamberlain, Michacl M. Bronstein, and Jonathan J Hunt. Hyperbolic deep reinforcement learning. In The Eleventh International Conference Learning Representations, 2023.
[17] Silvere Bonnabel. Stochastic gradient descent on ricmannian manifolds. IEEE Transactions on Automatic Control, 58()):217--2299 2013.
[18] 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.
[19] Nicolas Boumal. Riemannian trust regions with mmite-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.
11

---

[20] P-A Absil. Trust-region methods riemannian manifolds. Foundations of Computational Mathematics. 7:303-330, 2007.
[21] Nilesh Tripuraneni, Nicolas Flammarion, Francis Bach, and Michael Jordan. Averaging stochastic gradient descent riemannian manifolds. In Conference On Learning Theory, pages 650-687. PMLR. 2018.
[22] Orizon Ferreira, Mauricio S Louzeiro, and LF401 8420 Prudente. Gradient method for optimization on riemannian manifolds with lower bounded curvature. SIAM Journal on Optimization, 2914)22517-2541 2019.
[23] Xiaojing Zhu and Hiroyuki Sato, Riemannian conjugate gradient methods with inverse retraction. Compu- tational Optimization and Applications, 77:779-810, 2020.
[24] Melanie Weber and Suvrit Sra. Riemannian optimization via frank-wolfe methods. Mathematical Programming, pages 1-32, 2022.
[25] Hongyi Zhang, Sashank Reddi, and Suvrit Sra. Riemannian svrg: riemannian manifolds. Advances in Neural Information Processing Systems, 29, 2016.
Fast stochastic optimization on
[26] Hiroyuki Kasai, Hiroyuki Sato, and Bamdev Mishra. Riemannian stochastic recursive gradient algorithm. In International Conference on Machine Learning, pages 2516-2524. PMLR, 2018.
[27] Jingzhao Zhang, Hongyi Zhang, and Suvrit Sra. R-spider: fast riemannian stochastic optimization algorithm with curvature independent rate. arXiv preprint 181 1.04194, 2018.
[28] 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.
[29] 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.
[30] Soumava Kumar Roy, Zakaria Mhammedi, and Mchrtash Harandi. Geometry aware constrained optimiza- tion techniques for deep learning. In Proceedings of the IEEE Conference on Computer Vision and Puttern Recognition, pages 4460-4469, 2018
[31] Gary Beeigneul and Octavian-Fugent Ganea. Riemannian adaptive optimization methods. In International Conference on Learning Representations, 2019.
[32] Hiroyuki Kasai, Pratik Jawanpuria, and Bamdev Mishra. Riemannian adaptive stochastic gradient algo- rithms matrix manifolds. In International Conference on Machine Learning, pages 3262-3271, PMIRR 2019.
[33] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. International Conference on Learning Representations, 2021.
[34] Jungmin Kwon, Jeongscop Kim, Hyunsco Park, and In Kwon Choi. Asam: Adaptive sharpness-awan minimization for scale-invariant learning of deep neural networks. In International Conference on Machine Learning, pages 3905-59144 PMLR, 2021.
[35] Minyoung Kim, Li, Shell Hu, and Timothy Hospedales. Fisher sam: sharpness aware minimisation. In International Conference on Machine Learning, pages 11148-11161 PMI.R, 2022.
Information geometry and
[36] Jiawci Du, Hanshu Yan, Jiashi Feng, Jocy 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.
[37] Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha C Dvornek, sekhar tatikonda, James Duncan, and Ting Liu. Surrogate gap minimization improves sharpness-awar training. In International Conference on Learning Representations. 2022.
[38] John M Lee and John M Lee. Smooth manifolds. Springer, 2012.
[39] P-A Absil, Robert Mahony, and Rodolphe Sepulchre. Optimization algoritlums matrix manifolds. Princcton University Press, 2008
[40] Nicolas Boumal. An introduction optimization on smooth manifolds. Cambridge University Press, 2023.
12

---

[41] Tim Dettmers, Pasquale Minervini, Pontus Stenetorp, and Sebastian Riedel. Convolutional 2d knowledge graph embeddings. In Proceedings the AAAI conference on artificial intelligence, volume 32, 2018.
[42] 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. International Conference on Learning Representations, 2019.
[43] Asher Trockman and J Zico Kolter. Orthogonalizing convolutional layers with the cayley transform. In International Conference on Learning Representations, 2021.
[44] Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2):251-276. 1998.
[45] Ovidiu Calin. Neuromanifolds, pages 465-504. Springer International Publishing, Cham, 2020.
[46] Shun-ichi Amari, Hyeyoung Park, and Tomoko Ozeki. Geometrical singularities in the neuromanifold of multilayer perceptrons. Advances in neural information processing systems, 14, 2001.
[47] Cheongjae Jang, Sungyoon Lee, Frank Park, and Yung-Kyun Noh. A www.aannttriiatiiiiiairiie sharpness measure hased on information geometry. Advances neural information processing systems, 35:27893- 27905, 2022.
[48] Andi Han and Junbin Gao. Riemannian stochastic recursive momentum method for non-convex optimiza- tion. 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 Organi- zation, 8 2021. Main Track.
[49] 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.
[50] Andi Han and Junbin Gao. Variance reduction for riemannian non-convex optimization with batch size adaptation. arXiv preprint arXiv:2007.014 014444 2020.
[51] Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware minimization.
In International Conference Machine Learning, pages 639-668. PMLR, 2022.
[52] 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 Asxociation for Computational Linguistics (Volume 1: Long Papers). pages 5672-5686, Dublin, Ireland, May 2022. Association for Computational Linguistics.
[53] Max Kochurov, Rasul Karimov, and Serge Kozlukov. Geoopt: Riemannian optimization in pytorch. arXiv preprint arXiv:2005.02811 022899 2020.
[54] 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.
[55] Kristina Toutanova and Danqi Chen. Observed versus latent features for knowledge base and text inference. In Proceedings of the 3rd workshop on continuous vector spuce models and their compositionality, pages 57-66, 2015.
[56] Ryohei Shimizu, YUSUKE Mukuta, and Tatsuya Harada. Hyperbolic neural networks++. In International Conference on Learning Representations, 2021.
[57] Jonas Gehring, Michael Auli, David Grangier, Denis Yarats, and Yann N Dauphin, Convolutional sequence to sequence lcarning. In International conference on machine learning, pages 1243-1252. PMLR, 2017.
[58] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems. 30, 2017.
13

---

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 (Properties of Retraction Smoothness). Let w' Rw (pgrad((w)) and the curve 7(t) (ty) with the endpoints y(0) w and ((1) w". Then, we have
(T w ((>)w-gradC4) gradC(w" gradd(((())) 2 pbx)lerddddd
Proof. By the condition (C-5), we have
(T(")-grad&/ -radd((w) gradC(w))l LR||n|
By the Cauchy-Sehie Schwarz inequality, we have
(() gradC(w") gradL(w), m. < (T(T))-grad(((") -raddww)lllll
I Lx||vll
s p² Ln (wllgr
Therefore, we obtain
grad gradC(w), pw)ad gradd((e) -p² La Lxllgrad((ww)]]
Removing the constant p. the above incquality becomes
(T()) (+ gradC(w") grad[(w), gral 2 - pL (wradd(((w) Lastly, arrive the final result
(T()) ggrdd grad(("))) grad [[grad[(w)]
gradL(w))
N [ga))
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 by gradLr(w) = + Eie,, gradf(w; 24) T. and
adv RRw (pgradL. (w)). Further, let ((t) = (tr) be a retraction curve with ((0) = w and ((1) = ade Then, we have the following inequality
(T()) gradC: "erd(,(("""""rrrd(((
grad((w)) )_]
- pLR 1ap2 LR) grrad ( 2p 121
Proof. Let w² Re (pgrad ((w)) evaluated the loss function. We first add and subtract (T(()) gradC: grad gradd(())) where ((t) R., (tE) is a retraction curve where <(0) w and ((1) = w".
(T(1) grad (T()) grad - T(C) gwwdLo(w"), grad((
((a))-gradC..e grv)
T2
We will bound two terms, Ti and T2. separately. Regarding the term Ti, derive
Ti (T(r) graneeeggrarrrree ade) T()) gradC.(w" (w") g.g))
T(7) w dinds gr T() gradLe(w") I
2
< (T) Wine gradL: gw) || ((())- grad + 1 (w0r
(w)) _
< L
LA (w [krrad(())]..
< p² The (2llgradC((w) gradd(ww) (l) + ( + x2 (lrradd(()
2p Lio
(g 3pp²LL) LR) [grad((w))
b
14

---

From the above inequality, we could finally bound the term 77 as
n 2p Li
+ L
[[arad((w)ll
b
2
Regarding the term T2, we just use the lemma as
T3 (T(4) gradLi grad((w)) 2 (11-pEx)leradd(ee
Hence, we arrive at (T())) made gracC: (rraddl ] 3 pLn - -30² 30² [rrad((w) | - p-R------0
b
According to Algorithm 1, we follow the notation as
gradLe
Σ (w)) IE
adv
Rw, (pgradCe(w:))
We assume the stochastic w--AM where the same batch is used for both inner and outer updates.
Lemma (Descent inequality). Under the assumptions in Theorem I, have
B [[([w++)]] [[[(w:)]
30
E [[gradCC(as))
a² Lso 20 Lo2'sp'o'
2ap2 Lio
b
b
Proof. Using the condition (C-1), we have
f(we+1) = C. Rws (T(()) wi grad (ew
< [(wrs) a gradd((ws), (i)) mide gradL((w8 +
Lg
gradss
For the last term in RHS, we can bound as (ord || - C(wx) + [<) grad -grad((w))
+ +2(T(>) gradd grad[(wx))
Again, we have
[(w(11) s [(w.)) al (gradd((u)), (()d)(())),T)>)>)))) (rad (
+ a² T graddCC
= [(w))) L'Ls [grad[(w)) ll + a²
[(>) Braadder Cr -grad((w))]
2
- -a(1 - -aLs) ((>)
( gradC(w)))
my w
a² aLLs
< <((w))
[grad[(w)) ll + a LS [ (T) gradCi (rad - (rxd
2
+ + La2Lsllgradde LslggradC.( gradL(w))) a(l - aLs) (T)) ma gaagdada (wid")) grad(((w))
<((w)) a²Ls [gn) + t (rr)) + a² [[gradC((w)) -grad((w)) lee
2
- a(1 - -aLs) (T)) grad grad(((v)) in
S ((w)) LLLs Igr + 2a LSP ([rrad 2a L3P 1 (wr)) grad le,
2
+a' Lelgalgga)CC gradL( a(1 a01-(aaall aLs) (T)) grad (wa grad[(wx)) w
=[(w)) a²Is(1 P
IgradC(ws))l., a² LL + 213 IlgradC. ggrdd(ww))l
2
- -a(1 - -als) (T()) and gradL, gradC(ws))
15

---

Taking the expectation on both sides, we have
3 [[[(w[+1]] [[[[w]]] a² L² AT
E [[]erad + a²Tss(1 2T )o?
2
b
a0(1 - alss]
grad((w)))
B[L(wr)] a²Lss11 - -4
L² + )00
E [[ +
2
b
2p" Li
a(1 aals) aLs)
pla 3p (1)err [ll
b
For sufficiently large number of total iteration T, the condition max{LR. Ls} (defined in Theorem 1). Hence, we obtain
S
1 is casily satisfied where d7
30
E (w:)) [[[(w:]] E [[[(w++1) + Ls + +2L3p"
+ all--ad ao(l-as)ppLL
8
b
+ a Loo² 2a Lzp'02
2xx2 LiO0²
S [[(e E. [[(((++1)]
b
By telescoping the above inequality from = 0
1, we arrive at
8A
L
16a Lippo?
16ap Lig
E (6)
E [wnrad
3aT
3b
35
3b
Under the step size condition aa VIT and pa
Ti finally get
QiIA
Q20²
Q20²
El [1])ra
VT
bvT
bT5/6
for appropriate constants (Q.)2...
16

---

B Hyperparameter Details
We use the almost same hyperparameters in the study [32] and implement our experiments in Section 5 upon its official implementation. For completeness, summarize the hyperparameter configurations in Table and Table
Table 3: Hyperparameter configtttion for knowledge graph completion.
WN18RR
FB15k-237
Dimension
32
8
Batch Size
1000
1000
500
500
Negative Samples
50
50
50
50
Margin
8.0
8.0
8.0
8.0
Epochs
1000
1000
500
500
Max Norm
1.5
2.5
1.5
1.5
Max Scaler
3.5
2.5
2.5
2.5
Learning Rate
0.005
0.003
0.003
0.003
Gradient Norm
0.5
0.5
0.5
0.5
Table 4: Hyperparameter configurations for machine translation.
Hyperparameter
IWSLT'14
1MT'14
GPU Numbers
4
Embedding Dimension
64
64
Feed-forward Dimension
256
256
Batch Type
Token
Token
Batch Size
10240
10240
Gradient Accumulation Steps
1
Training Steps
40000
200000
Dropout
0.0
0.1
Attention Dropout
0.1
0.0
Max Gradient Norm
0.5
0.5
Warmup Steps
8000
6000
Decay Method
noam
noam
Label Smoothing
0.1
0.1
Layer Number
6
Head Number
8
Learning Rate
5.0
5.0
Adam B2
0.998
0.998
17


Reviews:
[{"review_id": "eM2iy61H3t", "reviewer": "Reviewer_SuW3", "summary": " This 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  ", "questions": "\n\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 \n \n   \n  \n  ", "limitations": "N/A", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "\n- novel formulation allowing to extend SAM to Riemannian manifolds (Algorithm 1)\n\n- theoretically verified model (Theorem 1)\n\n- various experiments \n\n\n", "weaknesses": "- 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\n\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your constructive feedbacks.\n\n**[On Contributions]**\n\nThank you for your important feedbacks, but we believe that our main example, hyperbolic representation learning, is not a simple “one example” but has a sufficiently general use cases, and that proposing an optimization technique that can advance it can be a significant contribution to the community.\n\nDeep learning in conventional Euclidean space often falls short in effectively capturing hierarchical relationships, efficiently representing tree-like structures, addressing dimensionality challenges (ex. embeddings), and etc. Under such limitations, one of main approach in geometric deep learning, ***hyperbolic representation learning***, has gained significant importance in various fields due to its ability to capture such relations in various data, for example, images, languages, graphs, and so on. The hyperbolic representation learning has been able to address important challenges in Euclidean space more intelligently by taking into account the underlying geometric characteristics of the data, thereby enabling us to devise significantly advanced techniques. We introduce several previous studies.\n\nWhile it holds true that the majority of modern deep learning architectures and techniques continue to be proposed in Euclidean space, the extension of successful models and foundational approaches from Euclidean to (Riemannian) manifolds, irrespective of the domain, is gaining attention. Moreover, (Riemannian) manifold extensions that consider the underlying data structure often outperform their Euclidean counterparts. For instance, in computer vision fields, the methodologies extended to hyperbolic spaces [VLB+19, KMU+20, GMK23, SKW+23, SBM23] have been actively studied for the important problem including image embedding, classification, segmentation, and etc. \n\nAlso, the well-regarded generative models in Euclidean space such as normalizing flow and diffusion models have successfully been extended to Riemannian Manifolds [LFC21, MN20, HAB+22].  Notably, even in language modeling, Transformer which is a main backbone, has been extended with fully hyperbolic modules [CHL+22], and hyperbolic Transformer outperforms the Euclidean counterpart. Moreover, in this year, the pioneering research [CCB+23] that incorporated hyperbolic geometry into deep reinforcement learning has been published in ICLR as a spotlight. More recently, avenues for enhancing numerical stability in hyperbolic representation learning have been published [MWW+23] in ICML 2023. \n\nIn this sense, we do not perceive our contribution to be limited solely to a few examples. Rather, we believe that our optimization approach has a great potential for broader applicability across various domains on manifold-aware deep learning.\n \n**References**\n\n- [VLB+19] Manifold Mixup: Better representations by Interpolating Hidden States, ICML 2019\n       \n- [KMU+20] Hyperbolic Image Embeddings, CVPR 2020\n        \n- [ELM+19] Continuous hierarchical representations with poincare ́ variational auto-encoders, NeurIPS 2019\n        \n- [LFC21] Hyperbolic Generative Adversarial Networks, IEEE Access 2021\n        \n- [MN20] Riemannian Continuous Normalizing Flows, NeurIPS 2020\n        \n- [HAB+22] Riemannian Diffusion Models, NeurIPS 2022\n        \n- [GSA+22] Hyperbolic Image Segmentation, CVPR 2022\n        \n- [CHL+22] Fully Hyperbolic Neural Networks, ACL 2022\n        \n- [GMK23] Hyperbolic Contrastive Learning for Visual Representations beyond Objects, CVPR 2023\n       \n- [SKW+23] Robust Hierarchical Symbolic Explanations in Hyperbolic Space for Image Classification, CVPR 2023\n        \n- [CCB+23] Hyperbolic Deep Reinforcement Learning, ICLR 2023\n        \n- [MWW+23] The Numerical Stability of Hyperbolic Representation Learning, ICML 2023\n        \n- [SBM23] Poincare ResNet, ArXiv 2023\n\n---\n**[On Riemannian SAM formulation]**\n\nDue to space constraint of rebuttal, we answer the reviewer’s concern in **1. Riemannian SAM is a non-trivial extension of Euclidean SAM with novel theoretical insights** in general response.\n\n---\n**[On Algorithm 1]**\n\nIn Algorithm 1, the term “base optimizer $\\mathcal{A}$” refers to any optimizer that can be employed on a manifold, such as Riemannian SGD or Riemannian Adam. Regarding line 9 in Algorithm 1, $\\Delta_t^{adv}$ signifies the final update vector constructed using A with Riemannian SAM gradient. In other words, the Riemannian SAM gradient $g_t^{adv}$ (line 8 in Algorithm 1) can be utilized to construct the first-order/second-order momentum of Riemannian Adam. Or alternatively, it can be directly employed for performing Riemannian SGD. We will clarify it.\n\n---\n**[On Lorentz Model]**\n\nThis is a mistake in our statement, and we apologize for the confusion. As pointed out by the reviewer, while it is true that the Lorentz model involves a (semi)-Riemannian manifold, we can still employ the Riemannian optimization introduced in our paper to optimize the objective function defined on this manifold. In light of this context, we will provide a clearer description in the revision.\n\n---\n**[On Role of Vector Transport in Line 8]**\n\nIn Algorithm 1 and according to our Riemannian SAM formulation, the Riemannian gradient computed at $w=w_t^{adv}$ is not a vector on the tangent space $T_{w_t} \\mathcal{M}$ defined at the point $w=w_t$, making the algebraic operations impossible for directly updating $w_t$. Consequently, in order to update the original parameter $w_t$ using the Riemannian gradient computed at the perturbed point, an operation is required to bring it into the tangent space at $w_t$, which can be achieved via the vector transport $\\mathcal{T}_{w_t^{adv}}^{w_t}$  at line 8 in Algorithm 1."}, {"review_id": "TQOViKGfGq", "reviewer": "Reviewer_zxyx", "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.", "questions": "line 101: should <\\xi,\\eta> be evaluated on the tangent space at z rather than w?\nline 112: the angle is not formally defined.", "limitations": "N/A", "rating": 3, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "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:\n1) The derivations of the new objective relies on standard tools from optimization on manifold and there is limited technical contribution.\n2) 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.\n3) Only the best test accuracies are reported. The authors are encouraged to conduct cross-valiation and/or report average performance based on randomized experiments.\n4) In my opinion, the improvement of the test accuracies is visible but not quite significant.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your constructive feedbacks.\n\n**[On Technical Limitations]**\n\nWe appreciate your valuable feedbacks, but we do not agree with the reviewer’s opinion. Due to space constraint of rebuttals, we answer this concern in **1. Riemannian SAM is a non-trivial extension of Euclidean SAM with novel theoretical insights** in general response.\n\n---\n**[On Intuition of Riemannian SAM]**\n\nThe \"flat minima hypothesis\" in Euclidean space was put forth, suggesting that neural networks achieve better generalization when they are trained to converge to flatter regions within the loss landscape. In this sense, the Sharpness-Aware Minimization (SAM), which encourages flat minima by (informally) minimizing the Euclidean gradient norm, has been demonstrated to effectively enhance model generalization in large-scale deep learning scenarios such as Vision Transformer and MLP-Mixer. \n\nBuilding on this, due to the fact that an $n$-dimensional manifold locally bears a resemblance to an $n$-dimensional Euclidean space, we expect that the flat minima hypothesis proposed in Euclidean space would hold to some extent for objective functions defined on general manifolds. The Riemannian SAM algorithm can be informally understood as reducing the Riemannian gradient norm $\\lVert \\mathrm{grad} f(w)\\rVert_w$, which reflects the local curvature and geometry of the manifold (hence, the underlying structure of the data) at the point $w$.\n\nIn order to validate our intuition, we address the challenge of optimizing a neural-net-like objective function defined on the unit sphere manifold $\\mathbb{S}^{2}$.\n\nThe synthetic dataset is generated by drawing a total of $500$ samples from a standard Gaussian distribution $\\mathcal{N}(0, 1^2)$ for $X$ and and a uniform distribution $\\mathcal{U}(0,1)$ for $y$ respectively, resulting in $X\\in\\mathbb{R}^{500 \\times 3}$ and $y\\in\\mathbb{R}^{500}$. We choose a non-linear regression MSE loss, specifically $f(w)=\\frac{1}{2n}\\Vert y - \\mathrm{ReLU}(Xw) \\Vert_2^2$. In order to craft an objective function on the unit sphere, we impose the constraint $\\mathcal{C} = \\lbrace w\\in\\mathbb{R}^3: \\lVert w \\rVert_2=1 \\rbrace$ on the model parameter $w \\in \\mathbb{R}^3$.\n\nThe Figure 1-(a) in the attached pdf file corresponds to Cartesian coordinates $w=(x,y,z)$ to spherical coordinates $(r,\\theta,\\varphi)=(1,\\theta,\\varphi)$, rendering contour plots. In Figure 1-(a), 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).\n        \nWithin a maximum iteration budget $100$, we search best hyperparameters for each optimization algorithm. In Figure 1-(a), 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$.\n\nFurthermore, in terms of sharpness measures, we considered 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 $\\lVert \\mathrm{grad} f(w)\\rVert_w$. Notably, Manifold-Aware Sharpness aligns with Information-Geometric Sharpness [JLP+22] 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 1-(b,c) depict the results. In both metrics, Riemannian SAM achieves smaller sharpness values than Euclidean SAM, implying convergence toward flatter regions. \n\nIn 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.\n\n**Reference**\n\n- [JLP+22] A Reparametrization-Invariant Sharpness Measure Based on Information Geometry, NeurIPS 2022.\n\n---    \n**[On Experiments]**\n\nContrary to the reviewer’s belief, we believe that our improvement is significant compared to baselines. For detailed response, due to space constraint of rebuttals, we answer the concerns in “**2. On Experiments”** in general response and include the tables for the additional experiments in the attached PDF file."}, {"review_id": "iqb2Q6uWQY", "reviewer": "Reviewer_wvpH", "summary": "The 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.", "questions": "Q1. 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?", "limitations": "The 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.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The extension of SAM to Riemannian manifolds is natural and meaningful.\n- Theoretical guarantees for convergence are provided.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your constructive feedbacks.\n\n**[On Q1 and Q2]**\n\nAs the reviewer pointed out, in most practical cases, the manifolds under consideration are often embedded in Euclidean space or are subsets of $\\mathbb{R}^d$ with an appropriate Riemannian metric. For instance, in the case of one of an elementary manifold, the $n$-dimensional hypersphere $\\mathbb{S}^{n-1}$, it is embedded submanifold of Euclidean space. Thus, the inner product defined on the tangent space $T_x \\mathbb{S}^{n-1}$ for any point on the hypersphere $x\\in \\mathbb{S}^{n-1}$ is simply the Euclidean inner product (i.e., $\\langle u, v\\rangle_x = u^\\mathsf{T} v$ for any $x$ and any tangent vectors $u,v\\in T_x \\mathbb{S}^{n-1}$). As the reviewer mentioned, the computations of gradients on the hypersphere involve calculating the Euclidean gradient first and then projecting it onto the tangent space. Furthermore, in deep learning on manifolds, prevalent manifold structures include the Poincaré ball and Lorentz model on the hyperbolic space, along with the Stiefel manifold that encourages orthogonality among parameters. All three of these manifolds are defined on subsets of $\\mathbb{R}^d$ with suitable Riemannian metrics. In this case, the computation of the Riemannian gradient, as the reviewer described, involves preconditioning the Euclidean gradient with the inverse metric.\n\n---\n**[On Q3]**\n\nIndeed, the Lorentz model is not a Riemannian manifold but rather a (semi)-Riemannian manifold. Nevertheless, Lorentz model remains amenable to Riemannian optimization. In consequence, due to the absence of a guarantee that the Euclidean gradient preconditioned with the inverse metric (i.e., the Riemannian gradient) resides on the tangent space, the projection operation onto the tangent space becomes necessary. The equation (8) in our paper illustrates this projection procedure.\n\n--- \n**[On Q4]**\n\nWhile the Stiefel manifold was not considered in our experiments in this paper, it is mainly employed to encourage ***parameter orthogonality***. The several studies have shown that parameter orthogonality influences model generalization and serves as a remedy for the vanishing/exploding gradient problem. We include some references below.\n\n**Reference**\n\n- [LJW+21] Orthogonal Deep Neural Networks, TPAMI 2021.\n    \n- [TK21] Orthogonalizing Convolutional Layers with the Cayley Transform, ICLR 2021 Spotlight.\n    \n- [FT21] Efficient Riemannian Optimization on the Stiefel Manifold via the Cayley Transform, ICLR 2020.\n    \n- [WCC20] Orthogonal Convoluiontal Neural Networks, CVPR 2020.\n\n---\n**[On Q5]**\n\nYes, you are right. The retraction already involves the model parameter $w$, hence we should apply chain rule with respect to $w$, which requires the (infeasible) higher-order Riemannian gradient in practice.\n\n---\n**[On Minor Points]**\n\nThanks for the correction of typos in the paper. In line 122, $\\varphi$ is known to be isometric. In another context, the isometricity of the vector transport exhibits the characteristics of symmetry/angle preservation and also the preservation of vector lengths. These attributes facilitate the theoretical analysis of Riemannian optimization and the isometricity of vector transports stands as one of the standard assumptions in Riemannian optimization analysis."}, {"review_id": "PHVs9QsGuR", "reviewer": "Reviewer_nQre", "summary": "This 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.", "questions": "N/A", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. 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.", "weaknesses": "1. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your constructive feedbacks.\n\n**[On Motivation on Riemannian SAM formulation]**\n\nWe believe that considering sharpness-aware minimization in Riemannian optimization could help the optimization. The \"flat minima hypothesis\" in Euclidean space was put forth, suggesting that neural networks achieve better generalization when they are trained to converge to flatter regions within the loss landscape. In this sense, the Sharpness-Aware Minimization (SAM), which encourages flat minima by (informally) minimizing the Euclidean gradient norm, has been demonstrated to effectively enhance model generalization in large-scale deep learning scenarios such as Vision Transformer and MLP-Mixer. Building on this, due to the fact that an $n$-dimensional manifold locally bears a resemblance to an $n$-dimensional Euclidean space, we expect that the flat minima hypothesis proposed in Euclidean space would hold to some extent for objective functions defined on general manifolds. The Riemannian SAM algorithm can be informally understood as reducing the Riemannian gradient norm $\\lVert \\mathrm{grad} f(w)\\rVert_w$, which reflects the local curvature and geometry of the manifold (hence, the underlying structure of the data) at the point $w$. \n\nIn addition, for some optimization problems that can be solved via both Euclidean SAM and Riemannian SAM, the Euclidean algorithm might not adequately take into account the underlying geometry of the manifold. \n\nIn order to validate our intuition, we address the challenge of optimizing a neural-net-like objective function defined on the unit sphere manifold $\\mathbb{S}^{2}$.\n\nThe synthetic dataset is generated by drawing a total of $500$ samples from a standard Gaussian distribution $\\mathcal{N}(0, 1^2)$ for $X$ and a uniform distribution $\\mathcal{U}(0,1)$ for $y$, resulting in $X\\in\\mathbb{R}^{500 \\times 3}$ and $y \\in\\mathbb{R}^{500}$. We choose a non-linear regression MSE loss, specifically $f(w)=\\frac{1}{2n}\\lVert y - \\mathrm{ReLU}(Xw)\\rVert_2^2$. In order to craft an objective function on the unit sphere, we impose the constraint $\\mathcal{C} = \\lbrace w\\in\\mathbb{R}^3: \\lVert w\\rVert_2=1\\rbrace$ on the model parameter $w\\in \\mathbb{R}^3$.\n\nThe Figure 1-(a) in the attached PDF file corresponds to Cartesian coordinates $w=(x,y,z)$ to spherical coordinates $(r,\\theta,\\varphi)=(1,\\theta,\\varphi)$, rendering contour plots. In Figure 1-(a) in the attached PDF file, 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).\n\nWithin a maximum iteration budget $100$, we choose hyperparameters for each optimization algorithm. In Figure 1-(a), 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$.\n\nFurthermore, in terms of sharpness measures, we considered 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 $\\lVert\\mathrm{grad} f(w)\\rVert_w$. Notably, Manifold-Aware Sharpness aligns with Information-Geometric Sharpness [JLP+22] when dealing with statistical manifolds, where the Riemannian metric is defined by the Fisher information. For the aforementioned problem, we compare two metrics and the Figure 1-(b,c) in the attached PDF file depict the results. As seen in Figure 1-(b,c), Riemannian SAM achieves smaller sharpness values than Euclidean SAM, implying convergence toward flatter regions. \n\nIn 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.\n\n**Reference**\n\n- [JLP+22] A Reparametrization-Invariant Sharpness Measure Based on Information Geometry, NeurIPS 2022.\n\n---\n**[On Theoretical Insights]**\n\nDue to the space constraint of rebuttals, we answer this concern for theory in general response. Please refer to “**on theoretical side**” in **1. Riemannian SAM is a non-trivial extension of Euclidean SAM with novel theoretical insights** in general response."}]

Author rebuttal:
Due to the space constraints on each rebuttal, we answer some important questions in general response here.

**1. Riemannian SAM is a non-trivial extension of Euclidean SAM with novel theoretical insights**

* **[On technical side]**

We would like to emphasize that our Riemannian SAM does NOT merely mimic the original SAM either theoretically or methodologically. In regarding methodology, there could be several extensions of conventional Euclidean SAM to a Riemannian optimization in different manners.

For example, it is most natural to choose a perturbation region at the current point $w_t$ as in the conventional Euclidean SAM, $\delta \in B_\rho(w_t) = \lbrace x \in \mathcal{M}: d_\mathcal{M}(w_t,x) \leq \rho \rbrace$ where $d_\mathcal{M}$ represents the distance on the manifold. However, adopting the constraint on $\delta$ 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 (C-4) in our paper), which makes it difficult to guarantee convergence. Moreover, the computation of $d_\mathcal{M}$ is often computationally inefficient in practice. Another possible extension is to apply the vector transport operation from Equation 8 of Algorithm 1 to Equation 9. The following outlines the modified procedure. (i) $g_t^{adv} = \mathcal{A}(grad \mathcal{L}(w; \mathcal{S})\lvert_{w=w_t^{adv}})$ and (ii) $\Delta_t^{adv} = \mathcal{T}_{w_t^{adv}}^{w_t} g_t^{adv}$.

For base optimizer $\mathcal{A}$, any optimization algorithm commonly used in Riemannian optimization can be adopted (e.g., Riemannian SGD, Riemannian momentum, etc.). However, when vector transport is applied after constructing $g_t^{adv}$ via the momentum-based optimizer $\mathcal{A}$, the momentum construction takes place on the tangent space $T_{w_t^{adv}} \mathcal{M}$ at the perturbed point $w_t^{adv}$, while the parameter update occurs on the different tangent space $T_{w_t} \mathcal{M}$ at the point $w_t$. As a result, this might introduce another challenges in understanding and analyzing the overall optimization process.
    - In this perspective, various alternative extensions are also possible, but among them, we have carefully designed a ***theoretically valid, computationally practical, and non-trivially extended Sharpness-Aware Minimization to a manifold for Riemannian optimization.*** Then, we have successfully demonstrated both convergence analysis and empirical studies to corroborate our Riemannian SAM.
    
* **[On theoretical side]**

In terms of theory, our key observation of Theorem 1 in the paper lies in ***the alignment between the Riemannian SAM gradient $\mathcal{T}_{w_t^{adv}}^{w_t} \mathrm{grad} f(w_t^{adv})$ and the Riemannian gradient $\mathrm{grad} f(w_t)$*** for the perturbation point, $w_t^{adv} = R_{w_t}(\rho_t \mathrm{grad}f(w_t))$. The previous study [50] on Euclidean SAM says that Euclidean SAM gradient should be well-aligned with the true gradient step for convergence. Unlike the theoretical claim in [50], 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) $\angle(\nabla f(w_t^{adv}), \nabla f(w_t))$ (Euclidean Alignment) and (ii) $\angle (\mathcal{T}_{w_t^{adv}}^{w_t} \mathrm{grad} f(w_t^{adv}), \mathrm{grad} f(w_t))$ (Riemannian Alignment, Ours)

In this example, we consider the logistic regression where $200$ data samples are generated with $100$ of them sampled from $\mathcal{N}(-1, 1^2)$ and the remaining $100$ sampled from $\mathcal{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 2 in the attached PDF file 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-dimensions, we expect that the angles would become significantly larger, corroborating our theoretical insight.

**2. On Experiments**

We did not select the best model based on the test data; instead, we opt for the best model based on the validation loss. Some baselines have reported their best test performance, which might have led reviewers to misunderstand our results. In fact, we report the average performance for knowledge graph completion and choose hyperparameters based on validation. Thus, our evaluation is comparatively more stringent when contrasted with those previous studies, placing us at a potential disadvantage in terms of evaluation. Also, as the reviewer's suggestion, we conduct additional experiments by varying seeds. Please refer to Table 1 in the attached PDF file.

During the rebuttal period, we conduct additional experiments with expanded size of Transformer on machine translation tasks. Given this large-scale experiments, it was hard to run randomized simulations for the baselines as well, due to rebuttal time constraints. As a result, we report test BLEU scores by selecting the model *based on validation loss*. Please refer to Table 2 in the attached PDF file.

According to Table 1 and 2 in the attached PDF file, we emphasize that our results are NOT marginal for all experiments considered in our paper.

Meta-review: This paper proposes 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. This extension is non-trivial because there are several other design choices, but for the extension presented in the current paper, the authors are able to prove convergence under related assumptions, which is the main strength of the paper.  This paper also demonstrated the efficiency of the method via empirical results. 

One main weakness is the inadequate discussion on the intuition of performing SAM on riemannian manifold, as pointed out by the reviewers. I suggest that the authors carefully revise the paper according to the various useful suggestions given by the reviewers.
Decision: Accept (poster)