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

Q1

I don't understand the role of the tau variable in the result figures. For example, could you explain how it is chosen, and how it relates to the results?

As we mentioned in W1, tau is the Kendall tau correlation coefficient. We chose it following previous works that used it as a metric to assess correlation between generalization and various proposed measures to predict it, [2], [3], [4]. A Kendal tau of 1.0 indicates perfect rank correlation and a kendal tau of -1.0 indicates perfect anti-correlation. We improve on this measure (increase its magnitude) compared to adaptive sharpness in our results

Q2

I don't understand why in Figure 3 (left), the average-case adaptive sharpness is claimed to reveal a correlation with the test loss (generalization gap). To me, it looks not obvious, and the dependency is reversed.

The reviewer makes a good point about the strength of the correlation. It is not very high, but a Kendal tau of 0.4 normally indicates a moderately strong correlation between the ranks of the data points. The plot does not look exactly linear, but for the most part higher values of the average-case adaptive sharpness do produce higher values of the test error. The dependency is indeed reversed, which happens in diagonal networks because average-case adaptive sharpness is the $L^1$ norm of the predictor, $||u\odot v||_1$, which correlates with generalization when the data is sparse in this case, while the maximum-case sharpnesses are (or are close to being) the $L^{\infty}$ norm of the predictor, which anti-correlates.

Q3

c inconsistently appears in Appendix D. Do the experiments of geodesic sharpness also choose a vector by taking to normalize each parameter tensor? If not how to guarantee fair comparison? In Appendix D is preset, but I don't understand whether the geodesic is invariant over scaling?

Good point. c is a hyperparameter for adaptive sharpness that geodesic sharpness does not require. We included c in the appendix merely to highlight the connection with adaptive sharpness. We expanded/re-wrote Appendix D in a way that hope makes clear what the connection is and the role of c.

Q4

I wonder if normalizing QK embeddings will fix the concern in the paper that Transformers have GL group symmetry? It is shown in practice that it stabilizes training (such as [1]). [2] also uses spherical projection for theoretical convenience. Or if not, are there alternative ways to address the scaling symmetry?

Interesting question, and we thank the reviewer for bringing our attention to [2], which we were not aware of and seems quite intriguing. On normalizing QK embeddings: they wouldn’t entirely fix it. The attention operation ($XW_Q@W_K^T$) becomes (X LayerNorm($W_Q$) (LayerNorm($W_K)^T X^T)$, which removes issues with scaling symmetry, but any transformation of $W_Q$ and $W_K$ such that LayerNorm($W_Q$) -> LayerNorm($W_Q$) $G^{-1}$,LayerNorm($W_K$)->LayerNorm($W_K$) $G^T$ would still result in the same attention matrix. In more intuitive terms, transforming a weight matrix by multiplying with an invertible matrix $\in GL_n$ encompasses besides re-scalings ($G = diag(\alpha)$), things (among others) like rotations (for G orthogonal), we expect that the introduction of LayerNorm would fix re-scalings but not necessarily the rotations for example. While we are not entirely familiar with [2] we suspect the same would hold there. We are not aware of any alternative ways of addressing this additional scaling symmetry during training.

[1] Andriushchenko et al. (2023) A Modern Look at the Relationship between Sharpness and Generalization

[2] Jiang et al. (2019) Fantastic Generalization Measures and Where to Find Them

[3] Dziugaite et al. (2020). In search of robust measures of generalization

[4] Kwon et al. (2021). ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks

We thank the reviewer for their feedback and for their suggestions on how to improve the paper. Below we address your questions:

W1

My main concern is the practical effectiveness of the proposed method: does it really improve existing adaptive sharpness? In Figure 4 the improvements are not quite visible, and the parameter is neither explained nor aligned by column.

We have added Table 1 at the start of the appendix with this information. The parameter tau is kendall tau correlation coefficient, which measures rank correlation of the results: if the sharpness of a model a is larger than that of model b, then, if sharpness is indeed correlated with generalization, we expect the generalization gap of model a to be larger than that of model b if the correlation is positive and negative otherwise.

W2

The experiments are not described in detail, which affects reproducibility. For self-containedness, the method used in the experiment from Wortsman et al. (2022) could be explained better in detail. For example, what are the hyperparameters, how are the models trained, what is the precise algorithm, etc.?

Thank you for this suggestion. We have added an appendix (Appendix C) with a description of the algorithm we use for solving for geodesic sharpness, which hews closely to that used to solve for adaptive sharpness. In terms of the models from Wortsman (and the new BERT models) we take the weights of the already trained models from the authors. The hyperparameters we use for solving for geodesic sharpness are identical to the ones used by [1]. We hope this provides sufficient clarity, but we are of course happy to provide any additionally required details.

W3

It is questionable whether the attention's symmetry argued by the author is indeed the practical reason rather than a hypothesis. It would be helpful if the author could give a stronger argument or some illustrations. To give an example, ensure that the scaling symmetry causes irregularity in pre-trained models by measuring the conditional number, and modulating out this symmetry indeed corrects the sharpness---or any other way you prefer.

We have added a small example in Appendix D.1 (Figure 7) extending Figure 1 that we hope will go some way to addressing this. In it, we plot, for our scalar toy model, the level sets of: a) the loss function; b) the euclidean and riemannian gradient; c) the traces of the euclidean and riemannian network hessian. The reason we picked the trace of the network Hessian is that it is a quantity that has been used to quantify flatness. Of note in that figure is that a) the Riemannian version of the gradient and Hessian have the same level set geometry as the loss function; b) both the riemannian gradient norm and the trace of the riemannian hessian have smaller values throughout than their euclidean equivalents; c) the trace of riemannian hessian actually reaches 0 when at the local minimum, whereas the euclidean hessian actually attains its highest value there; d) the Euclidean trace of the hessian is unable to distinguish between a local minimum and a local maximum (upper left vs bottom left in the figure) whereas the Riemannian trace can actually do so, something which is desirable if we want to study generalization.

W4-5

Agreed, we have made some changes that we hope improve the presentation of the paper.

W4

Also, the final takeaway message from Figure 4 (see line 488) is that the "...geodesically sharpest models studied on ImageNet are those that generalise best." This seems counter to the typically held belief that less sharp or more flat models generalize better. It would be nice to have more experiments or theory to explain why we see the opposite behavior here than previously in the literature. Or it would be nice to recreate the experiments in the literature which demonstrate that less sharp models generalize better (e.g. the original SAM paper and followups) and show that the geodesic sharpness behaves in the opposite way. To address this concern, could the authors

1. provide a theoretical explanation or justification as to why 'geodesic sharpness' might behave differently from traditional sharpness measures
2. conduct a direct comparison with previous sharpness measures on the same models and datasets. This, in particular, would be very nice.
3. And discuss potential implications of this result for providing better understanding of generalization in deep learning

This is an excellent question that we have been thinking about (and recently excited about, as we indicated in the general response to all reviewers).

1. Theoretical Explanation While we do not have a fully fleshed out theoretical explanation to offer, we can hope that what follows at least provides some justification. In diagonal networks when we train on sparse data, the $L^1$ norm of the predictor $u \odot v$ is a good predictor of generalization, whereas when we train on dense data, the $L^{\infty}$ norm tends to be more useful. Maximum adaptive sharpness (as compared to average adaptive sharpness) with a specific hyperparameter choice, becomes $||u \odot v||_{\infty}$ and so is positively correlated when the networks are trained on dense data, and negatively so when trained on sparse data. The reverse happens for average adaptive sharpness. We speculate something similar might be happening here, but both the architecture and the nature of the data itself is considerably more complex and so we cannot, yet, fully understand it. Teasing out the exact nature of these interactions is definitely something we intend to explore, and we feel that doing so properly will involve systematic investigation along several axes, so to speak, and this will be beyond the scope of this current paper. We are very excited to explore this in future work.
2. Comparison with previous sharpness measures (e.g. SAM) This is an intriguing idea. As mentioned in W.3 for CNNs and MLPs we don’t expect much of a difference with conventional adaptive sharpness, and so the experiments in the original SAM paper will probably not be particularly enlightening. Does the reviewer have a specific SAM followup paper dealing more with transformers in mind? We would be happy to try to reproduce as much as we can.
3. Potential implications of this result. Indeed, potential implications for better understanding generalization is what we are very much interested in. Our view is that this result, as it stands, an essential first step in this direction. We carefully introduced a new metric, and this metric allowed us to recover a previously elusive correlation between sharpness and generalization. Doing so has already opened exciting new questions around it, the sign of the correlation being one of them. Our perspective is that if we want to understand more deeply how sharpness relates to generalization, we at the very least needed to have a notion of sharpness that is useful for transformers and that actually does correlate with generalization. In this work, we have empirically verified this. Building on this in future, we hope (speculate!) that the general approach we used to produce this notion of sharpness: a) provides a more intrinsic view of the parameter space, allowing for access to objects such as the Riemannian curvature and Hessian on the quotient manifold, which we haven’t touched in this paper and might yet prove useful; and b) will help us to fold in more general symmetries and possibly the role of data into this quotient framework, which in turn might allow us to further refine what sharpness ought to mean.

Q1

line 112. in contribution (d) you claim that Figure 1 shows that there is a strong correlation between generalization and sharpness. Is this a typo? Can you better explain the connection between Figure 1 and generalization?

This was definitely a typo, thank you for pointing it out.

Q2

In general, the paper appears to only address symmetries that arise as group actions of GL. Is it obvious that this captures all possible symmetries in the parameter space? What is the effect of your construction of the geodesic sharpness when not all symmetries are removed?

The reviewer is completely right that we are not accounting for all possible symmetries. For instance permutation symmetry, exploited to great effect in [1], is not fully encompassed in our approach since we use a more restricted re-scaling symmetry for the non-attention parameters. It could conceivably be incorporated since permutation matrices are invertible, but we have not done so yet. But besides these there might still be other, more complex symmetries that we do not account for. Our view is that it is better to at least account for some (as shown in the improved results) than none, but we hope to extend this in the future.

Q3

in contribution (f) and in the paper you make a comparison with the CLIP experiments from Andriushchenko et al. However, those models were not trained to convergence and one could argue that the role of sharpness in generalization is only evident for well-tune and converged models. Does that also apply to your measure of geodesic sharpness?

We suspect it does, but in the same way that sharpness is useful during training (for SAM, for example), we think our measure could have the same use in training.

Q4

line 279, why denote the points in total space by I'm confused by this notation. Thank you for pointing this out. It is a typo, as the prime should not be there. We have fixed this in the paper.

Q5

line 280, is this a typo? Do you mean x,y for points in the quotient space?

Thank you for pointing this out. We have fixed this in the paper.

Q6

line 297. Can you clarify the details in the section "Linear embedding space". When you say "On the outer most layer, we are given a linear Euclidean space E…" what does this mean? Do you mean the parameters relating to the outermost layer of the network? Also, why define the loss just on these outer layer parameters? How are and related to each other?

This section was not well written in the paper and the mention of the outer layer is an oversight. We changed it, but in summary: our parameters ultimately live on computers and in principle can be any real matrix, this we call our embedding space; we are, by assumption 5.1 working with full-rank matrices, so our total space is that of full-rank matrices and from this total space we construct the quotient manifold.

Q7.

parameters relating to the outermost layer of the network? Also, why define the loss just on these outer layer parameters? How are and related to each other?

We would answer this roughly in the same way as in Q6. The loss is defined on the network parameters. We have re-worked this section in the paper.

Q8

line 318, i mentioned this above, but here you reference endowing the total space with a smooth inner product. Yes. it is possible to endow any smooth manifold with a Riemannian metric. However, this R. metric is not unique. You reference defining it in Appendix B.4 but that just points to a definition of what a Riemannian metric is. This is confusing and crucial aspect of this work.

We hope the discussion in W1 served to address this.

Q9

Section 4. When defining the geodesic sharpness, I'm surprised to not see mention of the exponential map to define geodesics in the manifold using vectors in the tangent space. Would it not be possible to reformulate your construction looking only at vectors in the tangent space? And restrict yourself to vectors of a certain size in order to keep them within a fixed ball when projected to the manifold.

Indeed it would. In equation 7, $\gamma_{\bar{\xi}}$(1) is just the exponential map at a horizontal vector $\bar{\xi}$. The reason we omitted this from the paper is that we tried to make the paper as accessible as possible to readers without more than cursory knowledge of riemannian geometry, and so strived to introduce the smallest possible set of concepts from it in the paper. We are open to changing this if the reviewer feels the presentation would gain from doing so.

Q10

In Figure 3, you plot sharpness vs test loss for diagonal models. But only for 50 of the 200 models from the experiment setup? Why only 50 and why those?

This is an oversight for which we apologize: we trained 50 networks and then showed all of them.

W3

It would be nice to see more experiments or at least some more computational examples. The paper only addresses, in a careful way, very simple diagonal networks and attention layers in transformers. Nothing in between. The experiments are limited to correlational studies. It would be good to have more architectures and datasets to help convince the reader of this new technique. To address this, the authors could

1. test this measure of geometric sharpness on intermediate architectures between diagonal networks and transformers, such as MLPs or CNNs; or ideally, architectures that are more closely comparable with other measures of sharpness in the literature.
2. apply the technique to different datasets beyond ImageNet
3. conduct ablation studies to isolate the impact of different components of the 'geodesic sharpness' measure

First, we strongly agree with the reviewer’s overall point of providing additional empirical evidence (which we did, with results further supporting our approach), and we appreciate the reviewer’s depth of engagement to suggest three different approaches to doing so, and we address each one individually):

1. Intermediate architectures. This is an interesting idea, but the tricky part is determining what would be an appropriate “intermediate” architecture. We considered MLPs and CNNs, but, while it is not immediately obvious, we suspect that results would be quite similar to those that would be obtained by adaptive sharpness in those settings. The reason is that in this case, the $GL_N$ re-scaling symmetry reduces to a $GL_1+$ symmetry, and in that setting geodesic sharpness almost fully reduces to the usual adaptive sharpness (we still have the symmetry curvature term). We have added a new appendix, Appendix F.1 in which we show this explicitly in a new appendix (and which at the same time we hope also provides additional evidence for the choice of metric in the paper). As (admittedly merely suggestive) evidence of this being the case, we do note that in the case of diagonal networks where this reduction happens, geodesic sharpness mirrors adaptive sharpness: it does have a slightly higher tau correlation factor but for the most part it essentially behaves as adaptive sharpness.

2. Datasets beyond Imagenet. This was indeed one of our top priorities, and in the revision we have included the results of geodesic sharpness for a pre-trained BERT on MNLI in the paper. Our approach again clearly improves the result obtained over using adaptive sharpness.

3. Ablating for different components of the geodesic measure. Thank you for this good suggestion; we will add the full results of this to our paper in an appendix, but in our preliminary experiments, the new norm impacts results the most, but the adding the second order weight corrections do not matter as much. Both components are necessary for the best performance. Details can be found in Appendix G

We thank the reviewer for their thorough reading of our paper. We are glad they appreciate the creativity of our ideas and the mathematical soundness of our approach, and we truly appreciate their constructive criticism, which has been helpful in improving the paper. We feel that the thoughtful, detailed level of their suggestions has made it possible for us to fully address almost all of their points, and thus significantly strengthen the paper through doing so.

W1

1. Explain the rationale behind choosing the specific metric in equation (15)
2. Provide some discuss around how sensitive the results are to different choices of Riemannian metrics on the total space
3. Clarify if there's a principled way to select an "optimal" metric for a given architecture. These points would help evaluate whether the method is truly "one-size-fits-all" as claimed.

1. Rationale for choice of metric. Thank you for raising this important question. There do indeed exist multiple choices: this is a general degree of freedom in Riemannian geometry, and it is a priori not clear which metric will yield the best results (e.g. there are natural gradient descent flavors based on the Fisher-Rao or the Wasserstein metric; both are used in practice). In our case we don’t have complete freedom to pick any metric we want: whatever metric we use has to be compatible with the symmetry we are trying to factor out. This requirement immensely restricts the metrics we can pick, but still does not result in a single unique possible choice. The reasons we chose the specific metric that we did are two-fold (in order of importance): a) this metric reduces nicely in the case of re-scaling symmetry - as those of the diagonal networks or adaptive sharpness -, to quantities known to correlate with generalization, providing a signal on the suitability of this particular metric; b) the quotient space geometry was already characterized by Absil et al.

2. Sensitivity to different metrics Ignoring the trivial cases of metrics that are related to metric (15) via re-scalings (as these result in the same geodesic sharpness, see appendix I.3) , we are aware of one other metric that yields a quotient geometry that has previously been characterized. In Appendix F.2 we describe this alternative metric. In our initial experiments the results did not substantially change.

3. Principled approaches to metric selection We are not aware of any principled way of selecting this metric, besides making sure it reduces to a useful measure in the usual settings (such as the scalar re-scalings of adaptive sharpness).

W2

Also the details of how to actually compute the geodesic sharpness for the attention layer example (Sec 5.2) is lacking, see line 468-469. It's not clear to me what that means to "...plug Eq. 16 into Eq. 17 and solve the resulting optimization problem using SGD". Would it be possible to include the details for that computation and optimization solution details/setup? To address this concern, could the authors

1. provide the explicit form of the optimization problem after plugging Eq. 16 into Eq. 17 and clarify the specific SGD algorithm used to solve it and hyperparameters used, and
2. provide any additional constraints or modifications needed to ensure the optimization remains on the manifold

Thank you for pointing this out; we agree that this section was lacking detail. We have added an appendix (appendix C) to the paper explaining in detail how the computation works and have also included some pseudo-code. We hope that this clarifies the paper in this regard.

Q11

How natural or reasonable is Assumption 5.1? There is a low-rank bias in deep learning when training overparameterized models. Can you provide a reference that this assumption is usually satisfied in multi-head attention layers for default choices of dv, dk

In general in non-linear networks there is indeed a tendency towards low-rank representations, which might make Assumption 5.1 seem excessive and counter to realistic situations. However, while the learned $W_Q W_K^T$ tends to be low-rank, $W_Q$ and $W_K$ (on which Assumption 5.1 ought to apply) themselves are usually high/full (column) rank [2]. In practice, this is also the case for the models that we study in this paper.

Q12

line 453: Is it possible to endow the total space with any smooth Riemannian metric? Why this one? Admittedly I'm less familiar with the reference (Absil et al 2008) but it would be nice to see a computation verifying this is indeed a Riemannian metric and why this inner product was chosen.

Added in appendix I.2 a computation showing it is indeed a Riemannian metric. Hope the first part is addressed by our answer to W1.

Q13

line 514: Can you clarify what you mean when you say "..attention weights approach being singular" here? How are you defining singular. Also you mention a relaxation parameter but I'm not clear on the role this 'relaxation parameter' plays. You indicate that "..in practise we found out results were robust to this parameter". Can you explain more what robustness analysis you conducted to justify this claim.

We here mean singular to mean $W_{Q,K,V,O}$ are not full-rank. Since we need to invert matrices of the type of $W_Q^T W_Q$ this can create a numerical issue. Because of rounding errors of floating-point precision, in practice $W_Q^TW_Q$ is always invertible, but sometimes the inverted matrices have large singular values. To combat this we introduce the relaxation parameter to $W_Q^TW_Q$, which dampens the resulting singular values. While we can’t take it to be exactly zero, as long as it’s small enough, numerical stability is improved and the results stay roughly the same (we show some results of the effect of varying this parameter in Appendix H.2 and Figure 11).

typos/nits

We have fixed all of these in the paper. We apologize for the size of the plots, they are the way they are due to page limits. We will endeavour to tighten up the presentation to find space for them. We have removed the “dealt by […] line in question as it does not add anything.

[1] Samuel K. Ainsworth et al. (2022). Git Re-Basin: Merging Models modulo Permutation Symmetries

[2] Yu, H., & Wu, J. (2023). Compressing Transformers: Features Are Low-Rank, but Weights Are Not!

W1

Geodesic sharpness in section 5.2 is abbreviated. Providing a clear, explicit formula for geodesic sharpness would improve comprehension.

We have changed this in the paper, and we hope it now reads a bit more clearly. But we can make further adjustments if needed.

W2

The discussion of geodesic sharpness in transformers (Section 5.3) is brief. Including a description of geodesic sharpness for transformers, or a brief explanation of how adaptive sharpness could be applied to each layer, would make the analysis more thorough.

We added this as appendix C.1 in the paper, as we agree it was unclear in the previous version. In summary: we treat the ($W_Q$, $W_K$), and ($W_V$,$W_O$) of each head ($W_O$ requires a bit more care) as $GL_h$ symmetric pairs, and treat all other parameters as having the same symmetry as is accounted for in adaptive sharpness, that of element wise re-scalings. We then follow Algorithm 1 in Appendix C.2 (essentially projected gradient ascent) to obtain the maximum sharpness.

W3

The paper aims to propose a metric that surpasses previous methods. Therefore, I would expect comprehensive experimental results demonstrating both the effectiveness of the proposed approach and its superiority over existing methods. While the experiments provide sufficient evidence of the method’s effectiveness, the comparative analysis with previous methods seems limited.

We have included additional experiments on real-world language data, and again observed an increase in the correlation compared with adaptive sharpness. Encouragingly, we managed to improve a tau of roughly zero (from adaptive sharpness), indicating no correlation at all, to a tau with a clear correlation, and with visible order.

Minor points

We have fixed these issues in the paper.

Q1

In Case A), the second-order term in Equation (12) is omitted because it is considered small compared to the first-order term. However, I wonder if this omission might be overly convenient. Since there is no indication of how small can be, I’m unsure about the relative scale of the first and second orders. Would it be possible to retain the second-order term and, perhaps, combine the results of both cases to reach the conclusion?

That is a good point, yes, it would. We added this analysis to Appendix E. Depending on the residual and rho, the results don’t necessarily change much.

Q2

Could you provide further clarification on why "adaptive sharpness is more appropriate" in Section 5.3?

That section of the paper was slightly re-written to be clearer, and we expand on it in Appendix C.1. In essence however, convolutional or fully connected layers do not have the full GL symmetry of the attention parameters, having only a re-scale symmetry, which is the symmetry which adaptive sharpness accounts for.

Q1-3: These changes are now in the revised paper.

Q4:

Diagonal network experiment: it is not clear to me what Figure 3 is showing. What is the meaning of the x axis and y axis? What is the ideal result? Why only show 50 data points when you have 200 model trained? What is the meaning of the Kendall Tau? Why the points in adaptive average case aligned in different direction in comparison with two other cases? I suggest adding more detailed captions and axis labels to Figure 3, and include explanations for these points in the corresponding text.

We have added these points to the main paper. On the issue of 50 vs 200 models, that was a typo; we apologize and we have fixed this. The x axis are the various sharpness measures, the y axis the test loss. The parameter $\tau$ is kendall tau correlation coeficient, which measures rank correlation of the results: if the sharpness of a model a is larger than that of model b, then, if sharpness is indeed correlated with generalization, we expect the generalization gap of model a to be larger than that of model b if the correlation is positive and negative otherwise. The reason for the difference in sign is already present in previous works, flipping depending on the data regime (dense vs sparse data).

Q5:

Assumption 5.1: Does Assumption 5.1 are too strong? What if the assumption is violated? Can we still calculated the result but with wrong value, or not able to calculate it at all? It would be good if the author can provide the evident and explanation of effect of adding into since there have not any for the assumption. Also, including a discussion of the implications of violating Assumption 5.1 and provide empirical evidence for the effects of adding the term.

In general in non-linear networks there is a tendency towards low-rank representations, which might make Assumption 5.1 seem excessive and counter to realistic situations. However, while the learned $W_Q W_K^T$ tends to be low-rank, $W_Q$ and $W_K$ (on which Assumption 5.1 ought to apply) themselves are usually high/full (column) rank [1]. In practice, this is also the case for the models that we study in this paper. Since we need to invert matrices of the type of $W_Q^T W_Q$, not being full-rank can create numerical issues. Because of rounding errors of floating-point precision, in practice $W_Q^TW_Q$ is always invertible, but sometimes the inverted matrices have large singular values. To combat this we introduce the relaxation parameter to $W_Q^TW_Q$, which provably dampens the resulting singular values. While we can’t take it to be exactly zero, as long as it’s small enough, numerical stability is improved and the results stay roughly the same (we show some results of the effect of varying this parameter in Appendix H.2 and Figure 11).

[1] Yu, H., & Wu, J. (2023). Compressing Transformers: Features Are Low-Rank, but Weights Are Not!

Q6:

Can the authors describe in detail how you treat the multi-layer schema of transformers model? It is not clear at this time. It would be good if the author can provide the diagram illustrating the approach for multi-layer transformers.

We have added discussion on this to the paper in Appendix C.1, but to briefly summarize here: we treat the $(W_Q, W_K)$, and $(W_V,W_O)$ of each head ($W_O$ requires a bit more care) as $GL_h$ symmetric pairs, and treat all other parameters as having the same symmetry as is accounted for in adaptive sharpness, that of element wise re-scalings. We then follow Algorithm 1 in Appendix C.2 (essentially projected gradient ascent) to obtain the maximum sharpness.

Q7:

Experiment diversity: it would be good if we can demonstrated in the natural language task too, like language modeling with wikitext103, to ensure that it works across domains.

We agree, and we have added some experiments with BERT to the paper. They are not on wikitext103, but we chose them because a) there was an adaptive sharpness baseline available; b) there were a large number of these fine-trained models available from the same source. Again we observe an increase in correlation

Q8:

In order to find the geodesic approximation, you need to solve the optimization using SGD, what is the quantitative performance, i.e. wall clock time or time/memory complexity of this approach in comparison with adaptive sharpness? How does it scale with large models like LLMs?

Added an appendix (Appendix C.3) discussing this, but in short it’s not much more expensive than adaptive sharpness. We require the same amount of forward and backwards passes, which are in practice the main bottleneck. The main difference is that we require some additional matrix multiplications, but we cache as much as is possible to speed this up. The projection onto the horizontal space requires matrix inversions and can be a bit more expensive, but the matrices that are being inverted are d_head*d_head. We haven’t experimented with models that are larger than around 110M parameters.

We thank the reviewer for taking the time to review our paper, for their overall positive response, their very concrete suggestions on how to strengthen the paper, and their willingness to reconsider their score.

1. Limited empirical support: we have included additional experiments on real-world language data, and again observed an increase in the correlation compared with adaptive sharpness. Encouragingly, we managed to improve a tau of roughly zero (from adaptive sharpness), indicating no correlation at all, to a tau with a clear correlation, and with visible order.
2. Effect of assumptions: we discuss this in more detail below, but in summary: while the multiplied together attention weights tend to be low rank, the individual attention weights (on which we need assumption 5.1 to hold) tend to be high/full-rank. In appendix H.2 we provide evidence that the relaxation parameter does not impact the final results.
3. Computational burden: we have provided a complexity analysis of the proposed method [Appendix C for details]; essentially showing that our approach is the same order of computation time as adaptive sharpness.

In the next comment we address your detailed questions.

We thank the reviewer for their helpful and concrete suggestions. We agree with these suggestions and have carried them out:

• we have included additional experiments on real-world language data, and again observed an increase in the correlation compared with adaptive sharpness. Encouragingly, we managed to improve a tau of roughly zero (from adaptive sharpness), indicating no correlation at all, to a tau with a clear correlation, and with visible structure; and
• we have provided a complexity analysis of the proposed method [Appendix C for details]; essentially showing that our approach is the same order of computation time as adaptive sharpness. Both of these are mentioned in the global response and also discussed in more detail in the paper.

We would be happy to further discuss any other questions.