Hide & Seek: Transformer Symmetries Obscure Sharpness & Riemannian Geometry Finds It

We wish to thank the reviewer for their comprehensive review and for their helpful suggestions.

Other Comments Or Suggestions:

The Table 1 should be in the paper

Thank you for pointing this out. We will move it to the main body of the paper.

---

Questions

Q1: What is the complexity and the computation cost of estimating geodesic sharpness? How does it compare to classical and adaptive sharpness?

This is a good question! We will add the following to the paper:

Geodesic sharpness does not significantly differ in time complexity from adaptive sharpness (or classical sharpness with a nearly identical time complexity). In our language model experiments, adaptive sharpness takes 1.5s per step, while our $S_{\text{inv}}$ takes 3s, and our $S_{\text{mix}}$ around 2s per step. For $S_\text{inv}$, the main additional overhead is inverting a $d_{\text{head}} \times d_{\text{head}}$ matrix and performing a Sylvester solve in SciPy on CPU, as there is no such solve in PyTorch.

Q2: How sensitive is the approach to different architectures? Only one architecture is considered for each experimental task, which could introduce bias in the results.

This is indeed a limitation of the experimental setup we inherit from [1]. To facilitate the comparison with [1], we decided to focus on the same architectures that were present there, but this could have a hidden bias. The reason for which this is done in [1] is to always compare models within the same loss surface (albeit at different points). Since the existence of correlation with generalization exists for all tasks we study, we suspect that this is not architecture-dependent, but are conducting further experiments with a broader set of ViT models to determine whether this indeed is the case.

---

We hope we have addressed all of your questions and look forward to any further questions and insights that might come up during the discussion period.

[1] Andriushchenko, M., Croce, F., Müller, M., Hein, M., and Flammarion, N. A modern look at the relationship between sharpness and generalization. 2023.

We wish to thank the reviewer for their thoughtful review and their really interesting suggestions. We had not fully considered the possible connections with relative flatness, but find these to potentially be a really fruitful avenue of research.

Claims And Evidence:

Although the reparameterization-invariance follows from the derivation of geodesic sharpness, it would have been nice to empirically verify this, as well.

This is a good suggestion, and we will add this to any final versions in an appendix (similar to Figure 1 in Andriushchenko et al.).

Experimental Designs Or Analyses:

The result that sharpness is negatively correlated with generalization for vision transformers is curious. [...]

Thank you for this great observation! This could quite possibly help explain this curious phenomenon. One possible way to test for this (at least on synthetic datasets) would be to use a similar experimental setup to that used in [1], modified for classification with diagonal networks, where we can artificially control local label constancy through class separation. More broadly, and with an eye to possible further experiments, if the reviewer happens to be aware of any possibly useful approximations to the local constancy of labels, we would welcome any such suggestions. In principle, we need access to the data-generating process details to estimate the labels' local constancy, something we don't have for ImageNet.

We also suspect that looking into the data distribution is the most promising future direction for understanding the sign of the correlation flip. For instance, in our synthetic regression task, we observe differing behaviours for the correlation between sharpness and generalization in the under or overparametrized regimes. The data itself can introduce additional symmetries in the overparametrized regime since in this regime $n<d$, where $n$ is the number of data points and $d$ is the dimension, and so the data matrix $X$, which is $n \times d$, always has a non-trivial null space by the rank-nullity theorem. This implies that two predictors $\beta = u \odot v, \beta' = u' \odot v'$ should be equivalent if there is $z \in Null(X)$ s.t. $\beta=\beta'+z$. In the underparametrized regime, where $n>d$, this is no longer necessarily the case, and we expect the null-space to be trivial, thus making these symmetries disappear. We are unsure at this moment of what concrete impact these additional symmetries have but we intend to investigate them further. More broadly, data-dependent symmetries are already known in the literature (e.g. [2]), but remain, in our opinion, underexplored.

Relation To Broader Scientific Literature:

That is, regularizing with relative flatness improves generalization also for transformers [2] and its behaviour wrt. adversarial examples is similar-ish for CNNs and transformers [3]. Since computing it for the penultimate layer and the CE loss is very efficient [3], it would be interesting to discuss geodesic sharpness wrt. relative sharpness.

Thank you for the suggestion. We agree it would be interesting to add this to the paper, and we'll endeavour to do so in any final version.

---

Questions

Q1: Have you considered measuring average geodesic sharpness, as well?

This is a good question! We have considered it, especially since [3] reports that, at least for diagonal networks, average sharpness can have different correlation signs with generalization (positively correlated vs anti-correlated and vice-versa). We did not present results on it in our paper due to its mathematical complexity: we would need to find a computationally feasible algorithm to properly integrate over the high-dimensional geodesic ball. It was unclear to us whether this was entirely possible, although a Monte Carlo approach might yield results that are "good enough".

double-check

Q2: Are the symmetries in [4] only instances of scaling and re-scaling?

Exactly!

---

We hope we have addressed all of your questions and look forward to any further questions and insights that might arise during the discussion.

[1] Petzka, Henning, et al. "Relative flatness and generalization." Advances in neural information processing systems 34 (2021): 18420-18432.

[2] Zhao, Bo, et al. "Symmetries, flat minima, and the conserved quantities of gradient flow." International Conference on Learning Representations 2023.

[3] Andriushchenko, M., Croce, F., Müller, M., Hein, M., and Flammarion, N. A modern look at the relationship between sharpness and generalization. 2023.

We thank the reviewer for their thorough review and their suggestions for improving the paper. We'll endeavour to include as much as possible in any final version.

Theoretical Claims

While the derivations are largely convincing, the reliance on approximations and assumptions about the quotient manifold’s structure in complex networks are points that could benefit from further clarification.

Thanks for pointing this out! We'll include a more detailed discussion on this in future versions.

Other Comments Or Suggestions

A more detailed ablation study on the effect of metric choice (invariant vs. mixed) could help clarify when one might be preferred over the other.

This is indeed a philosophically interesting question. From a theoretical perspective, we do not have any reason to prefer one metric over another as long as both of them correctly reflect symmetries.

In practice, the metrics perform very similarly, with the mixed metric tending to perform slightly better and being faster to run. The mixed metric's numerics are also advantageous as we do not need to solve a Sylvester equation to project into the horizontal space.

Expanding the discussion on the conditions under which geodesic sharpness may flip its correlation sign would benefit practitioners.

Thank you for this important question! This is something we're keen on investigating in future work, as it is one of the main factors limiting our approach's utility during training. We believe this will require significant research and taking into account more aspects, e.g. the data distribution (e.g. hypothesized by Reviewer hcv3), which are currently not considered in our framework that purely focuses on parameter space symmetries.

---

We hope we have addressed all of your concerns and look forward to discussing any outstanding concerns in the discussion period.

Thank you for the insightful review and helpful suggestions for improvement. We will make sure to mention the reference [1] that you brought to our attention, and we will contrast it to our paper which goes beyond the re-scaling symmetry. Please find our answers to your remaining concerns below.

Q1

Can you reconcile your empirical and theoretical findings [...]

Thank you for pointing this out! There is indeed some context missing for it to fully make sense, but these findings are not contradictory.

TL;DR: The theoretical derivation (Equation (13)) assumes the underparameterized regime (less parameters than data), while the empirical results (Figure (3)) are for the practically more relevant overparameterized regime (more parameters than data), where deriving a closed-form is intractable.

Here are the details:

• Theoretical assumptions break down. The theoretical expression (Equation (13)) assumes that $X^\top X = I_{d\times d}$ with data matrix $X \in \mathbb{R}^{n \times d}$ (number of data $n$ and parameters $d$). For this to be feasible, $n\geq d$, i.e. we are in the underparameterized regime. This assumption allows us to derive closed-form expressions for the unique optimal predictor and its sharpness (see, e.g. [2]). Unfortunately, analyzing the practically more relevant overparameterized regime exactly is intractable.

• We can reconcile empirical results with the theoretical prediction. If we re-run the experiment from Figure (3) (overparameterized regime) in the underparameterized regime, we indeed obtain the positive correlation between sharpness and generalization, as predicted by the theory in Equation (13). We believe it is practically less interesting because large models typically operate in the overparameterized regime.

• Our presented empirical findings are consistent with other works. The findings we report for overparameterized diagonal networks in the original submission agree with findings from other works. The anti-correlation of worst-case sharpness with generalization was also found in [3].

---

Q2

What is the relationship between sharpness and generalization you can report?

Thank you for this important question! What we can report is that contrary to [3], once we account for symmetry, there is a relationship between sharpness and generalization, and the Kendall-tau significantly differs from zero. This is what we show in our experiments and we will make sure to make this more explicit in the text. We believe that answering how the different correlation signs come into being is beyond the scope of our paper. The full story is more complicated and rightfully deserves future investigation:

Based on very recent insights from [4] which reports that sharpness minimization differs significantly between vision and language tasks, we hypothesize that the data distribution could play a significant role. Also, Reviewer hcv3 pointed out one interesting avenue to understanding this change in correlation sign via stability of the labels that we hope to follow up on in the very near future.

Other strengths and weaknesses

While the paper is easy to follow, it is unclear what the implications of its findings are [...]?

Our findings provide a foundation for future practical applications. Understanding the sign of the correlation between sharpness and generalization is the limiting factor, but once that is done, we expect geodesic sharpness to be useful for regularizing training.

What is the time complexity [...]?

This is a good question!

• Computational cost: Geodesic sharpness does not significantly differ in time complexity from adaptive sharpness. In our language model experiments, adaptive sharpness takes 1.5s per step, while our $S_{\text{inv}}$ takes 3s, and our $S_{\text{mix}}$ takes around 2s.

• Why we do not consider the Riemannian Hessian: Mainly for reasons of computational efficiency. We deal with models with upwards of $100 M$ parameters, and even accessing quantities such as the Hessian trace through multiple Hessian-vector products is computationally expensive (this was already the case in [3]).

---

References:

[1] Rangamani, A., Nguyen, N. H., Kumar, A., Phan, D., Chin, S. P., & Tran, T. D. (2021, June). A scale invariant measure of flatness for deep network minima. In ICASSP 2021-2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) (pp. 1680-1684). IEEE.

[2] Roger Grosse (2022). "A Toy Model: Linear Regression". University of Toronto, Topics in Machine Learning: Neural Net Training Dynamics.

[3] Andriushchenko, M., Croce, F., Müller, M., Hein, M., and Flammarion, N. A modern look at the relationship between sharpness and generalization. 2023.

[4] Sidak Pal Singh, Hossein Mobahi, Atish Agarwala, and Yann Dauphin. Avoiding spurious sharpness minimization broadens applicability of SAM. arXiv preprint arXiv:2502.02407, 2025.