SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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, we derive the Riemannian gradient on the Lorentz model  $\mathbb{L}^n = (\mathbb{H}^n,g_\ell)$ . As in Section 2.2, the Riemannian gradient could be computed as
$$
h = g _ {\ell} ^ {- 1} \nabla \mathcal {L} (w). \tag {13}
$$
Since  $g_{\ell}$  in (1) is a diagonal matrix, it is easy to compute the vector  $h$  with Euclidean gradient  $\nabla \mathcal{L}(w)$ . However, the vector  $h$  is not on the tangent space at  $w$ ,  $\mathrm{T}_w\mathbb{L}^n$ , thus we should have to project the vector  $h$  onto the tangent space  $\mathrm{T}_w\mathbb{L}^n$ . The projection is easily computed in a closed-form as
$$
\operatorname {p r o j} _ {w} (v) = v + \langle w, v \rangle_ {\mathbb {L}} w. \tag {14}
$$
Hence, the Riemannian gradient on Lorentz model is computed by
$$
\operatorname {g r a d} \mathcal {L} (w) = \operatorname {p r o j} _ {w} \left(g _ {\ell} ^ {- 1} \nabla \mathcal {L} (w)\right). \tag {15}
$$
As a next step, we should normalize the Riemannian gradient as in line 6 in Algorithm 1, and this is easy to compute as  $\| \mathrm{grad}\mathcal{L}(w)\| _w = \sqrt{\langle\mathrm{grad}\mathcal{L}(w),\mathrm{grad}\mathcal{L}(w)\rangle_{\mathbb{L}}}$  via Lorentzian scalar product  $\langle \cdot ,\cdot \rangle_{\mathbb{L}}$  defined in Section 2.1.
We utilize the Lorentz SAM derived above for our primary empirical studies conducted on hyperbolic space. In a similar way, one can derive the Riemannian SAM on Poincaré ball or Stiefel manifold, which we defer to Appendix.
3.1 Riemannian SAM Illustration: Toy 3D Illustration
We illustrate the 3-dimensional toy example on the sphere manifold. Let us define the 3D sphere manifold  $\mathbb{S}^2$  with the tangent space at  $w$ ,  $\mathrm{T}_w\mathcal{M}$  as
$$
\mathbb {S} ^ {2} := \left\{w \in \mathbb {R} ^ {3}: \| w \| _ {2} = 1 \right\},
$$
$$
\mathrm {T} _ {w} \mathbb {S} ^ {2} := \{v \in \mathbb {R} ^ {3}: w ^ {\intercal} v = 0 \}
$$
We consider the regression problem with the neural-net-like objective function on the randomly generated synthetic dataset. Toy 3D optimization problem with objective function  $f(w) = \frac{1}{2n} \| y - \mathrm{ReLU}(Xw)\| _2^2$  where  $X\in \mathbb{R}^{500\times 3}$  and  $y\in \mathbb{R}^{500}$  are drawn from  $\mathcal{N}(0,1^2)$  and  $\mathcal{U}(0,1)$  respectively with the model parameter  $w = (x,y,z)\in \mathbb{R}^3$  under  $\| w\| _2 = 1$ . (a) Comparison of converged points for each method. We plot the contour plots with the spherical coordinates under the relation  $(x,y,z)$
Figure 1: Toy 3D illustration
(a) Trace of Hessian,  $\operatorname {Tr}(\nabla^2 f(w))$
Figure 2: Comparison of two sharpness measures.
(b) Manifold-Aware Sharpness,  $\| \mathrm{grad}f(w)\| _w$
The Figure 1 corresponds to Cartesian coordinates  $(x,y,z)$  to spherical coordinates  $(r,\theta ,\varphi) = (1,\theta ,\varphi)$ , rendering contour plots. In Figure 1, we showcase the converged points on the objective function under the optimization using Riemannian SAM (in purple color) and the conventional Euclidean SAM (in pink color). Within a maximum iteration budget 100, the purple point (Riemannian SAM) attains a loss value of 0.3800 while the pink point (conventional Euclidean SAM) converges with a slightly higher loss value of 0.3808.
Furthermore, in terms of sharpness measures, we consider the following basic two quantities: (i) the trace of the Hessian (sharpness in the context of Euclidean space), and (ii) manifold-aware sharpness, characterized by the Riemannian gradient norm  $\| \mathrm{grad}\mathcal{L}(w)\| _w$ . Notably, the manifold-aware sharpness aligns with information-geometric sharpness [47] when dealing with statistical manifolds, where the Riemannian metric is defined by the Fisher information. For the aforementioned problem, we compare two metrics and Figure 2 depicts the results. In both metrics, Riemannian SAM achieves smaller sharpness values than Euclidean SAM, implying convergence toward flatter regions. In other words, since the Euclidean SAM might fail to properly consider the underlying structure of the manifold even for toy examples, this phenomenon is expected to be exacerbated in extremely high-dimensional problems such as deep learning.
In this section, we present the convergence guarantees for the RiemSAM framework. Our goal is to find a first-order  $\epsilon$ -approximate solution: the output  $\widetilde{w}$  such that  $\mathbb{E}\big[\| \mathrm{grad}\mathcal{L}(\widetilde{w})\|_{\widetilde{w}}^2\big]\leq \epsilon^2$ , which is a generalized convergence criterion of Euclidean  $\epsilon$ -stationary point. To guarantee the convergence for  $\epsilon$ -approximate solution, we require the following mild assumptions.
(C-1) (Upper-Hessian bounded) The objective function  $\mathcal{L}$  is said to be upper-Hessian bounded in  $\mathcal{U} \subset \mathcal{M}$  with respect to retraction  $\mathbf{R}$  if there exists some positive constant  $C$  such that  $\frac{d^2\mathcal{L}(\mathrm{R}_w(t\eta))}{dt^2} \leq C$ , for all  $w \in \mathcal{U}$  and  $\eta \in \mathrm{T}_w\mathcal{M}$  with  $\| \eta \|_w = 1$ , and for all  $t$  such that  $\mathrm{R}_w(\tau \eta)$  for all  $\tau \in [0, t]$ .
(C-2) (Lower-bounded) The objective function  $\mathcal{L}(\cdot)$  is differentiable and has bounded suboptimality.
$$
\mathcal {L} (w ^ {*}) > - \infty .
$$
for the optimal point  $w^{*}\in \mathcal{M}$
(C-3) (Unbiasedness and bounded variance) The stochastic Riemannian gradient is unbiased and has a bounded variance:
$$
\mathbb {E} _ {(x, y) \in \mathcal {D}} [ \operatorname {g r a d} \mathcal {L} (w; x) ] = \operatorname {g r a d} \mathcal {L} (w),
$$
$$
\mathbb {E} _ {(x, y) \in \mathcal {D}} [ \| \operatorname {g r a d} \mathcal {L} (w; x) - \operatorname {g r a d} \mathcal {L} (w) \| _ {w} ^ {2} ] \leq \sigma^ {2}.
$$
where  $\operatorname{grad}\mathcal{L}(w)$  is a true Riemannian gradient evaluated on a full batch of training dataset  $\mathcal{D}$ .
(C-4) (Retraction smoothness) We assume that there exists a constant  $L_{S} > 0$  such that
$$
\mathcal {L} (z) \leq \mathcal {L} (w) + \langle \operatorname {g r a d} \mathcal {L} (w), \eta \rangle_ {w} + \frac {1}{2} L _ {S} \| \eta \| _ {w} ^ {2}.
$$
for all  $w, z \in \mathcal{M}$  and  $\gamma(t) \coloneqq \mathrm{R}_w(t\eta)$  represents a retraction curve on  $\mathcal{M}$  for  $\eta \in \mathrm{T}_w\mathcal{M}$  with the starting point  $\gamma(0) = w$  and the terminal point  $\gamma(1) = z$ .
(C-5) (Individual Retraction Lipschitzness) We assume that there exists  $L_{\mathrm{R}} > 0$  such that
$$
\left\| \mathcal {T} (\gamma) _ {z} ^ {w} \operatorname {g r a d} \mathcal {L} (z; x) - \operatorname {g r a d} \mathcal {L} (w; x) \right\| _ {w} \leq L _ {\mathrm {R}} \| \eta \| _ {w}.
$$
for all  $w, z \in \mathcal{M}$ . As in condition (C-4),  $\gamma(t)$  denotes a retraction curve and  $\mathcal{T}(\gamma)_z^w$  is a 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  $(\mathrm{C - 2})\sim (\mathrm{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  $L_{S}$  in (C-4) and  $L_{\mathbb{R}}$  in (C-5) may be different. According to [26, 50], the condition (C-5) can be derived under the standard assumption on retraction Lipschitzness with parallel translation and one additional assumption on the bound between the parallel translation and the vector transport, but we assume the retraction Lipschitzness with the vector transport for simplicity. Lastly, in condition (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  $\widetilde{w}$  denote an iterate uniformly chosen at random from  $\{w_{1}, w_{2}, \dots, w_{T}\}$ . Further, we let  $\widetilde{L} = \max\{L_{S}, L_{R}\}$  where the constants  $L_{S}$  and  $L_{R}$  are defined in condition (C-4) and (C-5) respectively. Under the conditions (C-1)  $\sim$  (C-5) with descent learning rate  $\alpha_{t} = \frac{1}{\sqrt{T} \widetilde{L}}$  and ascent learning rate  $\rho_{t} = \frac{1}{T^{1/6} \widetilde{L}}$ , we have the following complexity for the constant batch size  $b$ :
$$
\mathbb {E} \left[ \| \operatorname {g r a d} \mathcal {L} (\widetilde {w}) \| _ {\widetilde {w}} ^ {2} \right] \leq \frac {Q _ {1} \Delta}{\sqrt {T}} + \frac {Q _ {2} \sigma^ {2}}{b \sqrt {T}} + \frac {Q _ {3} \sigma^ {2}}{b T ^ {5 / 6}} = \mathcal {O} (1 / \sqrt {T}). \tag {16}
$$
where  $\Delta = \mathcal{L}(w_0) - \mathcal{L}(w^*)$  and the constants  $\{Q_i\}_{i=1}^3$  are irrelevant to the total iteration  $T$  or the manifold dimension  $d$ .
We make some remarks on our convergence results and relationship to conventional Euclidean SAM.
Theoretical implications. Our key observation of Theorem 1 lies in the alignment between the Riemannian gradient  $\mathrm{grad}\mathcal{L}(w_t)$  (line 5 in Algorithm 1) and the Riemannian SAM gradient  $\mathcal{T}_{w_t^{adv}}^{w_t}\mathrm{grad}\mathcal{L}(w_t^{adv})$  (line 9 in Algorithm 1) for the perturbed point  $w_{t}^{adv}$ ,  $w_{t}^{adv}$ . The previous study [51] on Euclidean SAM says that Euclidean SAM gradient should be well-aligned with the true gradient step for convergence. Unlike the theoretical claim in [51], we stress that for convergence guarantee those gradients should be well-aligned within the preconditioned space (by inverse Riemannian metric) regardless of alignment in Euclidean space. To verify this insight, we directly measure the angles between two vectors with a 2D toy example, illustrating how they align in practice. Toward this, we consider two angles: (i)  $\angle (\nabla f(w_t^{adv}),\nabla f(w_t))$  (Euclidean Align-
Figure 3: Toy 2D illustration
ment) 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 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 3 depicts the comparison of angles. The loss decreases up to 10-th iteration, after which it remains around the converged point. As evident from the illustration, while the angles between the Euclidean space SAM gradient and the gradient deviate by up to around 25 degrees, the angles between the preconditioned SAM gradient and the preconditioned gradient, influenced by the Fisher information, align more closely with deviations
Table 1: Link prediction results (\%) in the filtered setting for WN18RR and FB15k-237 datasets. For hyperbolic architectures,  $\beta \in \{200,400,500\}$  and we report the best result. The results of baselines are taken from [52] except for HYBONET whose results are 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).
Manifold	Model	WN18RR	FB15k-237
#Dims	MRR	H@10	H@3	H@1	#Dims	MRR	H@10	H@3	H@1
Hyperbolic	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
ATTH	32	46.6	55.1	48.4	41.9	32	32.4	50.1	35.4	23.6
HYBONET w/ rAdam	32	48.8±0.2	55.4±0.1	50.3±0.2	45.5±0.3	32	33.5±0.2	51.5±0.2	36.5±0.3	24.2±0.3
HYBONET w/ rSAM	32	49.3±0.2	56.0±0.2	50.7±0.1	46.2±0.3	32	34.3±0.2	52.0±0.1	37.3±0.3	25.1±0.4
MURP	β	48.1	56.6	49.5	44.0	β	33.5	51.8	36.7	24.3
ROTH	β	49.6	58.6	51.4	44.9	β	34.4	53.5	38.0	24.6
ATTH	β	48.6	57.3	49.9	44.3	β	34.8	54.0	38.4	25.2
HYBONET w/ rAdam	β	51.2±0.2	57.1±0.2	52.5±0.2	48.3±0.2	β	35.2	52.9	38.7	26.3
HYBONET w/ rSAM	β	51.6±0.2	58.7±0.3	53.3±0.2	48.6±0.1	β	36.0±0.2	54.3±0.4	39.6±0.2	26.6±0.1
only up to a maximum of 10 degrees. In high-dimensional loss landscapes, we expect that the angles would become significantly larger, corroborating our theoretical insight.
On upper bound. Distinct from the convergence of Euclidean SAM [51], our upper bound (16) has the additional term involving  $Q_{3}$ , but we still achieve the optimal complexity of SGD,  $\mathcal{O}(1 / \epsilon^4)$  for  $\epsilon$ -approximate solution. Note that the presence of term involving the constant  $Q_{3}$  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 vector-transported gradients (see line 8 in Algorithm 1) at each iteration, which are the main challenges in our proof. Our results can also provide the guarantees for SAM variants such as Fisher SAM [35], whose convergence guarantees are missing.
We conduct two sets of experiments; (i) knowledge graph completion, and (ii) machine translation. The first experiment aims to evaluate our Riemannian SAM on shallow networks and the second task is for optimizing large-scale deep neural networks. For all our experiments, we consider the Lorentz manifold introduced in Section 2.1 and employ the recent hyperbolic architecture, HYBONET [52]. The HYBONET is a fully hyperbolic neural network, whose each layer is constructed on the Lorentz manifold including a linear, attention, residual, and positional encoding layer. We implement our Riemannian SAM upon Geoopt framework [53] written in PyTorch library [54]. Regarding hyperparameters, we basically adhere to the same experiment settings in [52] and the details are provided in each section and Appendix.
A knowledge graph completion aims to predict missing relationships within a knowledge graph, which represents structured information as a collection of entities, their attributes, and the relationships between them. More precisely, knowledge in a graph is of the form of triplets  $(h,r,t)$  where  $h$ ,  $r$ , and  $t$  denote the head entity, the relationship or predicate, and the tail entity respectively. In the knowledge graph completion task, given a partially populated knowledge graph, the goal is to predict the missing entity or relationship in a triplet: solving  $(h,r,?)$  and  $(?,r,t)$ .
In our experiments, we use two popular benchmark datasets; WN18RR [41] and Fb15k-237 [55]. We employ the same data preprocessing in [3] and two standard metrics for evaluations: (i) Mean Reciprocal Rank (MRR), the average of the inverse of the true entity ranking, and (ii) Precision at  $K$  (H@K), the proportion of test instances where the correct answer appears in the top- $K$  ranked predictions. For
Table 2: The BLEU scores on the test set 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 are taken from [52] except for HYBONET whose results are reproduced by ourselves.
Manifold	Model	IWSLT '14	WMT '14
Euclidean	CONVSEQ2SEQ [57]	23.6	14.9
TRANSFORMER [58]	23.0	17.0
Hyperbolic	HYPERNN++ [56]	22.0	17.0
HATT [42]	23.7	18.8
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  $\rho_t \in \{10^{-4}, 10^{-3}, 10^{-2}, 10^{-1}\}$  for Riemannian SAM and the other hyperparameters are the same as HYBONET [52] for fair comparisons.
Table 1 illustrates the results on WN18RR and Fb15k-237 datasets. As in previous work [52], we test our Riemannian SAM on two different regimes: (i) small embedding dimension 32 and (ii) large dimension  $\beta \in \{200,400,500\}$  for both datasets. Regarding the large dimension, we report the best results among the dimension candidates. In Table 1, Riemannian SAM on the Lorentz model achieves the best performance with great margins for all comparison metrics. In the same way, Riemannian SAM shows the state-of-the-art performance for all metrics considered under both regimes.
5.2 Machine Translation
In this experiment, we evaluate our Riemannian SAM on Lorentz Transformer built with Lorentz components introduced in [52] for IWSLT '14 and WMT '14 benchmark datasets in machine translation. We use the BLEU score as an evaluation metric on the IWSLT '14 test set and the newtest2013 test set of WMT '14 respectively. According to [56], we train hyperbolic models with Riemannian SAM in a low-dimensional setting where the dimension of the word vector is  $d = 64$ . As in the knowledge graph completion task, we choose the ascent learning rate in  $\rho_t \in \{10^{-5}, 10^{-4}, \dots, 10^{-2}\}$  for marginal hyperparameter tuning.
Table 2 demonstrates the results. HYBONET baseline trained with Riemannian Adam already outperforms the Euclidean Transformer in both IWSLT '14 and WMT '14 datasets. Upon this baseline, we only substitute our Riemannian SAM for Riemannian Adam with other hyperparameters unchanged. As seen in Table 2, Riemannian SAM significantly outperforms the Riemannian Adam baseline for both datasets. Note that both HYPERNN++ and HATT are partially hyperbolic networks, so we could not evaluate our Riemannian SAM on these models since it is difficult for a fair evaluation.
Wall-clock time. As in Euclidean SAM, Riemannian SAM requires additional forward and backward propagation in a single iteration loop (see Algorithm 1). Thus, we report the wall-clock time comparison for each experiment. For knowledge graph completion (Section 5.1) and machine translation (Section 5.2), Riemannian SAM takes roughly 1.6 and 1.8 times longer than Riemannian Adam for one epoch, respectively. To alleviate the computational overhead, one can employ the stochastic weight perturbation (SWP) and sharpness-sensitive data selection (SDS) suggested in [36], which do not depend on the manifold structure. Another practical consideration is to use a subset of minibatch in computing perturbation (see line 6 in Algorithm 1) for large-scale models. We leave the study on reducing computational cost as future work.
In this study, we proposed a sharpness-aware minimization on Riemannian manifolds, called Riemannian SAM. Under our framework, we presented novel examples of Riemannian SAM including a Lorentz SAM. We analyzed the convergence of the Riemannian SAM for general manifolds with a less aggressively decaying ascent learning rate condition. Moreover, we showed that Riemannian 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 networks. As future work, we plan to study the technique to reduce the computations and analyze the generalization error bounds of Riemannian SAM theoretically.
This work was supported by the National Research Foundation of Korea (NRF) grants (No.2018R1A5A1059921, RS-2023-00209060), Institute of Information & Communications Technology Planning & Evaluation (IITP) grants (No.2019-0-00075, Artificial Intelligence Graduate School Program(KAIST)) funded by the Korea government (MSIT).
A Proofs of Theorem 1
We basically follow the arguments in the convergence of Euclidean SAM [51], but the details are totally different.
Lemma 1 (Properties of Retraction Smoothness). Let  $w^{*} = \mathrm{R}_{w}\big(\rho \mathrm{grad}\mathcal{L}(w)\big)$  and the curve  $\gamma (t) = \mathrm{R}_w(t\eta)$  with the endpoints  $\gamma (0) = w$  and  $\gamma (1) = w^{*}$ . Then, we have
$$
\langle \mathcal {T} (\gamma) _ {w ^ {*}} ^ {w} \operatorname {g r a d} \mathcal {L} (w ^ {*}), \operatorname {g r a d} \mathcal {L} (w) \rangle \geq (1 - \rho L _ {\mathrm {R}}) \| \operatorname {g r a d} \mathcal {L} (w) \| _ {w} ^ {2}
$$
Proof. By the condition (C-5), we have
$$
\| \mathcal {T} (\gamma) _ {w ^ {*}} ^ {w} \operatorname {g r a d} \mathcal {L} (w ^ {*}) - \operatorname {g r a d} \mathcal {L} (w) \| _ {w} \leq L _ {\mathrm {R}} \| \eta \| _ {w}
$$
$$
\begin{array}{l} \left| \langle \mathcal {T} (\gamma) _ {w ^ {*}} ^ {w} \operatorname {g r a d} \mathcal {L} \left(w ^ {*}\right) - \operatorname {g r a d} \mathcal {L} (w), \eta \rangle_ {w} \right| \leq \| \mathcal {T} (\gamma) _ {w ^ {*}} ^ {w} \operatorname {g r a d} \mathcal {L} \left(w ^ {*}\right) - \operatorname {g r a d} \mathcal {L} (w) \| _ {w} \| \eta \| _ {w} \\ \leq L _ {\mathrm {R}} \| \eta \| _ {w} ^ {2} \\ \leq \rho^ {2} L _ {\mathrm {R}} \| \operatorname {g r a d} \mathcal {L} (w) \| _ {w} ^ {2} \\ \end{array}
$$
$$
\langle \mathcal {T} (\gamma) _ {w ^ {*}} ^ {w} \operatorname {g r a d} \mathcal {L} (w ^ {*}) - \operatorname {g r a d} \mathcal {L} (w), \rho \operatorname {g r a d} \mathcal {L} (w) \rangle_ {w} \geq - \rho^ {2} L _ {\mathrm {R}} \| \operatorname {g r a d} \mathcal {L} (w) \| _ {w} ^ {2}
$$
Removing the constant  $\rho$ , the above inequality becomes
$$
\langle \mathcal {T} (\gamma) _ {w ^ {*}} ^ {w} \operatorname {g r a d} \mathcal {L} (w ^ {*}) - \operatorname {g r a d} \mathcal {L} (w), \operatorname {g r a d} \mathcal {L} (w) \rangle_ {w} \geq - \rho L _ {\mathrm {R}} \| \operatorname {g r a d} \mathcal {L} (w) \| _ {w} ^ {2}
$$
Lastly, we arrive at the final result as
$$
\begin{array}{l} \langle \mathcal {T} (\gamma) _ {w ^ {*}} ^ {w} \operatorname {g r a d} \mathcal {L} (w ^ {*}), \operatorname {g r a d} \mathcal {L} (w) \rangle_ {w} = \langle \mathcal {T} (\gamma) _ {w ^ {*}} ^ {w} \operatorname {g r a d} \mathcal {L} (w ^ {*}) - \operatorname {g r a d} \mathcal {L} (w), \operatorname {g r a d} \mathcal {L} (w) \rangle_ {w} + \| \operatorname {g r a d} \mathcal {L} (w) \| _ {w} ^ {2} \\ \geq (1 - \rho L _ {\mathrm {R}}) \| \operatorname {g r a d} \mathcal {L} (w) \| _ {w} ^ {2} \\ \end{array}
$$
In the next lemma, we will show the alignment of the true Riemannian gradient and the true Riemannian SAM gradient.
Lemma 2 (Alignment of the true Riemannian gradient and the true Riemannian SAM gradient). Let us denote the stochastic Riemannian gradient at time  $t$  by  $\mathrm{grad}\mathcal{L}_t(w) = \frac{1}{b}\sum_{i\in J_t}\mathrm{grad}\mathcal{L}(w;x_i)\in \mathrm{T}_w\mathcal{M}$  and  $w^{adv} = \mathrm{R}_w\big(\rho \mathrm{grad}\mathcal{L}_t(w)\big)$ . Further, let  $\gamma (t) = \mathrm{R}_w(t\eta)$  be a retraction curve with  $\gamma (0) = w$  and  $\gamma (1) = w^{adv}$ . Then, we have the following inequality
$$
\mathbb {E} \left[ \left\langle \mathcal {T} (\gamma) _ {w ^ {a d v}} ^ {w} \operatorname {g r a d} \mathcal {L} _ {t} (w ^ {a d v}), \operatorname {g r a d} \mathcal {L} (w) \right\rangle_ {w} \right] \geq \left(\frac {1}{2} - \rho L _ {\mathrm {R}} - 3 \rho^ {2} L _ {\mathrm {R}} ^ {2}\right) \left\| \operatorname {g r a d} \mathcal {L} (w) \right\| _ {w} ^ {2} - \frac {2 \rho^ {2} L _ {\mathrm {R}} ^ {2} \sigma^ {2}}{b}
$$
Proof. Let  $w^{*} = \mathrm{R}_{w}\big(\rho \operatorname{grad}\mathcal{L}(w)\big)$  evaluated on the loss function. We first add and subtract  $\langle \mathcal{T}(\zeta)_{w^{*}}^{w}\operatorname{grad}\mathcal{L}_{t}(w^{*}),\operatorname{grad}\mathcal{L}(w)\rangle_{w}$  where  $\zeta (t) = \mathrm{R}_w(t\xi)$  is a retraction curve where  $\zeta (0) = w$  and  $\zeta (1) = w^{*}$ .