RSAM: Learning on Manifolds with Riemannian Sharpness-Aware Minimization

Hi, reviewer gTdN. Thank you for your comments. We would like response as follows:

• Firstly, it was a typo in page 4 that $\mathcal{L}_\mathcal{S}$ was omitted, we have fixed the issue.

• The second equation has no RHS, it essentially means that motivated by the theorems above, we should instead solve the optimization problem $min\_{\theta \in \mathcal{M}} \mathcal{L}^{RSAM}\_\mathcal{S}(\theta)$ instead of $min\_{\theta \in \mathcal{M}} \mathcal{L}\_{\mathcal{D}}(\theta)$ as mentioned at the beginning of Section (3.2).

• The operation $\mathcal{R}\_\theta$ is the well-known retraction operation on a manifold. Specifically, it will be a smooth mapping from the tangent space to the manifold: $\mathcal{R}\_{\theta}: \mathcal{T}\_{\theta} \mathcal{M} \to \mathcal{M}$. This retraction map satisfies $\mathcal{R}\_\theta(0\_\theta) = \theta$, where $0\_\theta$ denotes the zero element of the vector space $\mathcal{T}\_\theta \mathcal{M}$, and $\frac{d}{dt}\_{t = 0} \mathcal{R}\_{\theta}(t \xi) = \mathcal{T}\_{0\_\theta} \mathcal{R}\_{\theta}(\xi) = \xi$ for all $\xi \in \mathcal{T}\_{\theta}\mathcal{M}$. In practice, a manifold can have multiple different retraction operations. For Stiefel manifold, we specifically used the operation mentioned in Section (4). In the last line of Alg 1, the old $R\_{\theta_t}$ denotes the same retraction operation at $\theta_t$, which is $\mathcal{R}\_{\theta_t}$. However, it was a typo that $R$ is not capitalized, we have also fixed that in the revision.

• We believe that $D$ having more degree of freedom can also be a good idea to explore because it potentially can capture more information within the loop. In our experiments, we chose $D$ to be a simple matrix based on a single $\theta$ for the sake of computational tradefoff, and we demonstrate that it can improve the final result with only a little tradeoff. Last but not least, we have also included the proper link to our code. Thank you for your feedbacks.

Hi, reviewer efpX. Thank you for your comments. There are a few things we want to clarify and we would like to respond to your comments as follows:

• Firstly, about the phrase “we imposed orthogonality on a single convolutional layer in the middle of the architecture in all settings”, what we mean in this statement is that we impose orthogonality on a convolutional layer, which is different from imposing orthogonality on a single kernel. Specifically, with PyTorch, on ResNet34 we can extract this layer as model.layer1.1.conv1.weight. We found that only imposing orthogonality on this single layer improves the performance notably, while we also have the freedom to put the constraint on multiple layers.

• The matrix $S$ can also be optimized with SAM and SGD. In our experimental results, the "SAM" row means that the whole architecture, including the matrix $S$, is optimized with SAM. We need to clarify that in the decomposition in section (4.1), the matrix $S$ is not constrained to reside on the Stiefel manifolds, so it has the freedom to have any values.

• We will take into account your third comment and provide additional results on those baselines. Among those baselines, we have experimented with Riemannian SGD. Specifically, we optimized the matrix $U$ with RSGD and other parameters with SGD. What we found was that the performance was about the same as SAM in average across the settings on average.

• Also, thank you for informing us about the work "Riemannian SAM: Sharpness-Aware Minimization on Riemannian Manifolds". It was a concurrent work that we were not informed of while working on this direction. Nevertheless, our approach has slightly different mathematical development, theoretical results, and application directions. Indeed, as you mentioned, it is also desirable to include comparisons with FisherSAM, Riemannian Adam, and AdamW. We will take that into account.

Hi reviewer dxfe, thank you for your comments. We would like to respond to your comments as follows:

• In Proposition 1, $**v**_\vtheta^{\top} = \gradtheta\gL(\vtheta)^{\top}**D**_\vtheta$ as we mentioned at the beginning. However, we only used that in the proof for simplicity and is not needed for the proposition statement, so we can safely removed that. We have updated that in the revision. .

• In the proof of Theorem 1, $\tilde{\theta}$ is the image of $\theta$ under the logarithm map. The logarithmic map is a fundamental concept in Riemannian manifolds. Specifically, Let $v \in \mathcal{T}\_\mathcal{\theta} \mathcal{M}$ be a tangent vector to the manifold at $\theta$. Then there is a unique geodesic $\gamma:[0,1] → \mathcal{M}$ satisfying $\gamma(0) = \theta$ with initial tangent vector $\gamma'(0) = v$. The corresponding exponential map is defined by $\exp\_{\theta}(v) = \gamma(1)$. Hence, the logarithmic map is the inverse of this exponential map.

• Regarding the effect of $\rho$ on the accuracy, we will take into account your comment for further improvement. Currently, from what we have tried so far, $\rho$ for RSAM is chosen similarly to that of SAM in the original SAM work, and the accuracies in those sensible values are not differed significantly.

• The loss function $\mathcal{L}$ should be assumed to be Lipschitz, we have included that in the revision and also the typos that you mentioned in the minor comments.