SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jihun Yun
arcprime@kaist.ac.kr
Eunho Yang
KAIST, ATRICS
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.
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 high-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 [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 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, 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.
- 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.
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 \in \mathcal{M}$ , let  $\mathrm{T}_w\mathcal{M}$  denote the tangent space at  $w$ . An inner product on tangent space  $\mathrm{T}_w\mathcal{M}$  is a bilinear, symmetric, positive definite function  $g_w(\cdot, \cdot) \coloneqq \langle \cdot, \cdot \rangle_w : \mathrm{T}_w\mathcal{M} \times \mathrm{T}_w\mathcal{M} \to \mathbb{R}$ . If a metric  $\langle \cdot, \cdot \rangle_w$  smoothly varies with  $w \in \mathcal{M}$ , we call  $\langle \cdot, \cdot \rangle_w$  a Riemannian metric. An induced norm on  $\mathrm{T}_w\mathcal{M}$  is  $\| \zeta \|_w \coloneqq \sqrt{\langle \zeta, \zeta \rangle_w}$ . A Riemannian manifold is a pair  $(\mathcal{M}, g)$  of the manifold  $\mathcal{M}$  and the associated Riemannian metric tensor  $g$ .
Definition 2 (Geodesic). A geodesic is a curve  $\gamma(\cdot): [0,1] \to \mathcal{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: \mathrm{T}_w\mathcal{M} \to \mathcal{M}$  maps a tangent vector  $\zeta_w \in \mathrm{T}_w\mathcal{M}$  onto  $\mathcal{M}$  along a geodesic curve such that  $\gamma(0) = w$  and  $\gamma(1) = z$  with  $\dot{\gamma}(0) = \zeta_w$ . Specifically,  $\gamma(t) \coloneqq \exp_w(\zeta_w t)$  represents a geodesic. A retraction  $\mathrm{R}_w(\cdot)$  is a (computationally efficient) generalization of an exponential map satisfying the following properties:
-  $\mathrm{R}_w(0) = w$  and  $DR_w(0) = \mathrm{Id}_{\mathrm{T}_w\mathcal{M}}$  where  $DR_w$  represents the derivatives of  $\mathrm{R}_w$  and  $\mathrm{Id}_{\mathrm{T}_w\mathcal{M}}$  denotes an identity map on  $\mathrm{T}_w\mathcal{M}$ .
Definition 4 (Parallel translation/Vector transport). A parallel translation  $P_{z}^{w}(\cdot):\mathrm{T}_{z}\mathcal{M}\to \mathrm{T}_{w}\mathcal{M}$  transport a tangent vector in  $\mathrm{T}_z\mathcal{M}$  to  $\mathrm{T}_{w}\mathcal{M}$  in parallel while preserving norm and direction (i.e., along a geodesic). A vector transport  $\mathcal{T}(\gamma)_z^w (\cdot):\mathrm{T}_z\mathcal{M}\to \mathrm{T}_w\mathcal{M}$  with respect to retraction map  $\mathrm{R}$  maps a vector  $\zeta_z\in \mathrm{T}_z\mathcal{M}$  to  $\zeta_w\in \mathrm{T}_w\mathcal{M}$  along a retraction curve  $\gamma (t) = \mathrm{R}_w(\xi_w t)$  for some  $\xi_w\in \mathrm{T}_w\mathcal{M}$ , which is computationally efficient approximation of a parallel translation. In this work, we only consider isometric vector transport, i.e.,  $\langle T_z^w\zeta_z,T_z^w\eta_z\rangle_w = \langle \zeta_w,\eta_w\rangle_w$  for all  $\zeta_w,\eta_w\in \mathrm{T}_w\mathcal{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  $\mathbb{P}^n = (\mathbb{B}^n,g_p)$  is a Riemannian manifold with
$$
\mathbb {B} ^ {n} = \left\{x \in \mathbb {R} ^ {n}: \| x \| _ {2} <   1 \right\}, \quad g _ {p} (x) = \left(\frac {2}{1 - \| x \| _ {2} ^ {2}}\right) ^ {2} g _ {e} (x).
$$
where  $g_{e}$  represents an Euclidean metric tensor. The associated distance on Poincaré ball is given by
$$
d _ {p} (x, y) = \operatorname {a r c o s h} \left(1 + 2 \frac {\| x - y \| _ {2} ^ {2}}{(1 - \| x \| _ {2} ^ {2}) (1 - \| y \| _ {2} ^ {2})}\right).
$$
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  $\mathbb{L}^n = (\mathbb{H}^n,g_\ell)$  is a semi-Riemannian manifold, but it is still possible to employ Riemannian optimization. The Lorentz model consists of
$$
\mathbb {H} ^ {n} = \left\{x \in \mathbb {R} ^ {n + 1}: \langle x, x \rangle_ {\mathbb {L}} = - 1, x _ {0} > 0 \right\}, \quad g _ {\ell} (x) = \left[ \begin{array}{c c c c} - 1 & \dots & 0 & 0 \\ 0 & 1 & \dots & 0 \\ \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & \dots & 1 \end{array} \right]. \tag {1}
$$
where  $\langle x,y\rangle_{\mathbb{L}} = -x_0y_0 + \sum_{j = 1}^n x_jy_j$  is known as the Lorentzian scalar product. The associated distance function is given by
$$
d _ {\ell} (x, y) = \operatorname {a r c o s h} \left(- \langle x, y \rangle_ {\mathbb {L}}\right)
$$
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  $\varphi : \mathbb{H}^n \to \mathbb{P}^n$  (with the corresponding inverse mapping  $\varphi^{-1} : \mathbb{P}^n \to \mathbb{H}^n$ ) defined as
$$
\varphi \left(x _ {0}, x _ {1}, \dots , x _ {n}\right) = \frac {\left(x _ {1} , x _ {2} , \cdots , x _ {n}\right)}{p _ {0} + 1}, \tag {2}
$$
$$
\varphi^ {- 1} \left(x _ {1}, x _ {2}, \dots , x _ {n}\right) = \frac {\left(1 + x _ {1} ^ {2} + \cdots + x _ {n} ^ {2} , 2 x _ {1} , \cdots , 2 x _ {n}\right)}{1 - \left(x _ {1} ^ {2} + \cdots + x _ {n} ^ {2}\right)}. \tag {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 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  $\varphi$  in (2) and (3).
Stiefel manifold. The orthogonality, i.e.,  $W^{\top}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  $\mathrm{St}(n,p) = (\mathbb{V}^n,g_W)$  for  $n\geq p$  is prevalent for orthogonal neural networks, which is also a Riemannian manifold defined by
$$
\mathbb {V} ^ {n} = \left\{X \in \mathbb {R} ^ {n \times p}: X ^ {\mathsf {T}} X = I \right\}, \quad g _ {W} (Z _ {1}, Z _ {2}) = \operatorname {T r} (Z _ {1} ^ {\mathsf {T}} Z _ {2}).
$$
for tangent vectors  $Z_{1}, Z_{2} \in \mathrm{T}_{W}\mathrm{St}(n,p)$ . The tangent space of  $\mathrm{St}(n,p)$  at  $W$  is defined by  $\mathrm{T}_{W}\mathrm{St}(n,p) = \{Z: Z^{\mathsf{T}}W + W^{\mathsf{T}}Z = 0\}$ .
2.2 Riemannian Optimization
We are interested in the following optimization problem over the Riemannian manifold  $(\mathcal{M},g)$
$$
\min_{w\in \mathcal{M}}\mathcal{L}(w).
$$
where  $\mathcal{L}:\mathcal{M}\to \mathbb{R}$  is a smooth function defined on manifold  $\mathcal{M}$ . Following the work [17], Riemannian stochastic gradient descent (RSGD) updates the model parameter  $w\in \mathcal{M}$  as
$$
w _ {t + 1} = \exp_ {w _ {t}} \left(- \alpha_ {t} \operatorname {g r a d} \mathcal {L} (w _ {t})\right). \tag {4}
$$
where  $\operatorname{grad}\mathcal{L}(w_t) \in \mathrm{T}_{w_t}\mathcal{M}$  is a Riemannian gradient at  $w_t$  and  $\alpha_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  $\mathrm{R}_{w_t}(\cdot)$ , yielding the update rule  $w_{t+1} = \mathrm{R}_{w_t}(-\alpha_t\operatorname{grad}\mathcal{L}(w_t))$ . Generally, the Riemannian gradient  $\operatorname{grad}\mathcal{L}(w_t)$  in (4) is computed with the Riemannian metric tensor  $g$  as
$$
\operatorname {g r a d} \mathcal {L} \left(w _ {t}\right) = g ^ {- 1} \left(w _ {t}\right) \nabla \mathcal {L} \left(w _ {t}\right). \tag {5}
$$
where  $\nabla \mathcal{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  $\mathrm{T}_{w_t}\mathcal{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  $\mathcal{D} = \{(x_i,y_i)\}_{i = 1}^n$  as
$$
\min  _ {w \in \mathcal {M}} \mathcal {L} (w) := \frac {1}{n} \sum_ {i = 1} ^ {n} \mathcal {L} \left(w; x _ {i}\right). \tag {6}
$$
where the smooth loss function  $\mathcal{L}(\cdot):\mathcal{M}\to \mathbb{R}$  is defined on a Riemannian manifold  $\mathcal{M}$ . We formulate the sharpness-aware minimization in terms of loss function values on the manifold  $\mathcal{M}$  as
$$
\min  _ {w \in \mathcal {M}} \max  _ {\| \delta \| _ {w} ^ {2} \leq \rho^ {2}} \underbrace {\left\{\mathcal {L} \left(\mathrm {R} _ {w} (\delta)\right) - \mathcal {L} (w) \right\}} _ {\text {S h a r p n e s s i n t e r m s o f l o s s f u n c t i o n}}. \tag {7}
$$
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  $\delta$  in the tangent space  $\mathrm{T}_w\mathcal{M}$ . To solve the inner subproblem, we resolve the inner optimization problem in a different manner as
$$
\max  _ {\delta \in \mathrm {T} _ {w} \mathcal {M}} \mathcal {L} \left(\mathrm {R} _ {w} (\delta)\right) - \mathcal {L} (w) \quad \text {s u c h t h a t} \quad \| \delta \| _ {w} ^ {2} \leq \rho^ {2}. \tag {8}
$$
for a fixed point  $w \in \mathcal{M}$ . For ease of computations, we approximate the perturbed loss function  $\mathcal{L}\big(\mathrm{R}_w(\delta)\big)$  via Taylor's expansion as
$$
\mathcal {L} \left(\mathrm {R} _ {w} (\delta)\right) \approx \mathcal {L} (w) + \langle \operatorname {g r a d} \mathcal {L} (w), \delta \rangle_ {w}. \tag {9}
$$
Then, our inner maximization problem comes in hand:
$$
\max  _ {\delta \in \mathrm {T} _ {w} \mathcal {M}} \langle \operatorname {g r a d} \mathcal {L} (w), \delta \rangle_ {w} \quad \text {s u c h t h a t} \quad \| \delta \| _ {w} ^ {2} \leq \rho^ {2}. \tag {10}
$$
This problem could be easily solved since it finds the steepest direction  $\delta$  on Riemannian manifold  $\mathcal{M}$ , whose solution is known to be just Riemannian gradient. Therefore, we have
$$
\delta^ {*} = \rho \frac {\operatorname {g r a d} \mathcal {L} (w)}{\| \operatorname {g r a d} \mathcal {L} (w) \| _ {w}}. \tag {11}
$$
Under the optimal perturbation  $\delta^{*}$  in (11), we further approximate the gradient of the sharpness-aware minimization in (7) (for outer minimization problem) as
$$
\operatorname {g r a d} \mathcal {L} \left(\mathrm {R} _ {w} \left(\delta^ {*}\right)\right) \approx \operatorname {g r a d} \mathcal {L} (w) | _ {w = \mathrm {R} _ {w} \left(\delta^ {*}\right)}. \tag {12}
$$
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  $\mathrm{grad}\mathcal{L}(w)|_{w = \mathrm{R}_w(\delta^*)}$  in (12) is not on the tangent space at  $w$ ,  $\mathrm{T}_{w}\mathcal{M}$ . To perform an actual parameter update on  $w$ , we should transport the Riemannian SAM gradient to the tangent space  $\mathrm{T}_{w}\mathcal{M}$  via vector transport  $\mathcal{T}_{\mathrm{R}_w(\delta^*)}^w$  with respect to the retraction  $\mathrm{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\sim 7$  and lines  $8\sim 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,  $\delta \in B_{\rho}(w_t) \coloneqq \{x \in \mathcal{M} : d_{\mathcal{M}}(w_t, x) \leq \rho\}$  where  $d_{\mathcal{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 (C-4) in Section 4), 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 line 8 of Algorithm 1 to line 9. The following outlines the modified procedure: (i)  $g_t^{adv} = \mathcal{A}\big(\mathrm{grad}\mathcal{L}(w;S)|_{w = w_t^{adv}}\big)$  and (ii)  $\Delta_t = \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). In the meanwhile, when the 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  $\mathrm{T}_{w_t^{adv}}\mathcal{M}$  at the perturbed point  $w_t^{adv}$ , 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 \in \mathcal{M}$  be the point on a neuromanifold (or statistical manifold)  $\mathcal{M}$  which is realized by Euclidean network parameter  $\theta \in \mathbb{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
$$
\begin{array}{l} \operatorname {g r a d} \mathcal {L} (w) = F (\theta) ^ {- 1} \nabla \mathcal {L} (\theta) \\ \delta^ {*} = \rho \frac {\mathrm {g r a d} \mathcal {L} (w)}{\| \mathrm {g r a d} \mathcal {L} (w) \| _ {w}} = \rho \frac {F (\theta) ^ {- 1} \nabla \mathcal {L} (\theta)}{\sqrt {\mathrm {g r a d} \mathcal {L} (w) ^ {\top} F (\theta) \mathrm {g r a d} \mathcal {L} (w)}} = \rho \frac {F (\theta) ^ {- 1} \nabla \mathcal {L} (\theta)}{\sqrt {\nabla \mathcal {L} (\theta) ^ {\top} F (\theta) ^ {- 1} \nabla \mathcal {L} (\theta)}}. \\ \end{array}
$$
Algorithm 1 Riemannian SAM: Sharpness-Aware Minimization on Riemannian Manifolds
1: Input: Descent learning rate  $\alpha_{t}$ , ascent learning rate  $\rho_{t}$ , and a base Riemannian optimizer  $\mathcal{A}$  (such as Riemannian SGD, Riemannian Adam).  
2: Initialize:  $w_{0} \in \mathcal{M}$   
3: for  $t = 0,1,\ldots,T-1$  do  
4: Draw a minibatch sample  $\mathcal{S} = \{x_{1},\dots,x_{|S|}\}$ .  
5:  $g_{t} \leftarrow \frac{1}{|S|} \sum_{i=1}^{|\mathcal{S}|} \operatorname{grad} \mathcal{L}(w_{t}; x_{i})$ .  
6:  $\delta_{t} \leftarrow \frac{g_{t}}{\|g_{t}\|_{w_{t}}}$ .  
7:  $w_{t}^{adv} \leftarrow \mathrm{R}_{w_{t}}(\rho_{t} \delta_{t})$ .  
8:  $g_{t}^{adv} \leftarrow \mathcal{T}_{w_{t}^{adv}}^{w_{t}} \operatorname{grad} \mathcal{L}(w; \mathcal{S})|_{w=w_{t}^{adv}}$ .  
9:  $\Delta_{t}^{adv} \leftarrow \mathcal{A}(g_{t}^{adv})$ .  
10:  $w_{t+1} \leftarrow \mathrm{R}_{w_{t}}(-\alpha_{t} \Delta_{t}^{adv})$ .  
11: end for  
12: Output:  $w_{T}$
which is entirely identical to the perturbation of Fisher SAM [35].