The Geometry of Attention: Ricci Curvature and Transformers Training and Robustness

We thank the reviewer for their careful review, their encouraging words, their helpful questions, and their constructive comments on the shortcomings. The following addresses their specific points.

---

---

Weaknesses:

1. Since curvature is derived from attention values, we cannot directly manipulate it without modifying the attention scores. Given that we consider the attention scores as edge weights in a graph, manipulating curvature values by changing attention values in an isolated fashion would be a significant departure from our approach. Instead, we explore strategies to analyze the role of curvature on the properties of this graph (and consequently the transformer), and adjust these properties by indirectly modifying the curvature values, e.g., by penalizing attention variance. Analyzing attention maps as graphs does enable us to use a variety of tools and measures that characterize the structure and dynamics of graphs. While we do not suggest that curvature is the only useful one, it has certain properties of interest. We used some of these properties in this paper highlighting the utility of this measure, and others could be used in future research.

• Curvature characterizes the local geometry of the graph. While related to bottleneckedness (hence to Cheeger constant, graph cut, etc.), curvature is more local. The alternatives mentioned do not capture the graph structure at the local level that curvature does.

• Curvature is linked to three (related) notions of robustness on graphs: structural robustness, dynamic/system robustness, and distributional robustness.

• Curvature is linked to fundamental concepts such as entropy, as we discuss; although it is not the only one.

2. We appreciate pointing this out. The assumption you refer to is in fact that $A$ is fixed across layers, not only that the weights are shared. This assumption is borrowed from Dovonon et al. That being said, we agree with your suggestion, we will reconsider this strong assumption, and will work on revising our proof strategy to relax it.

3. Thank you for this valuable suggestion. Variance regularization is a tool we used to modify the curvature distribution and investigate its impact on the behavior of transformers. We will further clarify this upon revision and resubmission. There are other tools for this, and we do not claim that variance is the most versatile tool, but it is sufficiently simple and effective for the purpose of our investigation. The softmax temperature and attention score variance are indeed related, and while variance regularization could also be layer-specific, the tunability and learnability of temperature are helpful advantages. Exploring how to modify the curvature distribution by adjusting the softmax temperature and other measures is an interesting direction that could strengthen and supplement our empirical observations. We will consider this direction in future work.

4. We will report training loss and will report the hyperparameters. For measuring generalizability, it is our understanding that for comparing two cases that do not reach the same level of training accuracy, test accuracy alone does not fully reflect generalizability.

---

---

Questions:

1. Please see our response to weakness 1 above.

2. Please see our response to weakness 2 above.

3. Please see our response to weakness 3 above.

---

---

We will make all the above points clearer in our future submission.

We thank the reviewer for their careful consideration of our work and for their encouraging remarks. The following addresses their specific weaknesses and questions.

---

---

Weaknesses and Questions:

This is a question that could be explored, and we find it intriguing to explore relevant metrics (e.g. entropy). The regularization method we discuss showcases an example of possible ways to modify the curvature distribution and how this could impact the behavior of transformers. We will clarify this and point to other possibilities upon revising our paper.

We are grateful to the reviewer for their careful consideration of our work. We would like to remark that a main contribution of this work is exploring the robustness of transformers (in addition to their performance) using the graph geometry of attention maps. The following addresses the specific points the reviewer has raised.

---

---

Weaknesses:

”It is unclear whether this paper ...”

We believe our paper falls in the category of works in deep learning research that provide theory-informed heuristics for improving understanding, intuition, inspiration, and methodological developments. We rely on results from more theoretical papers to provide understanding and point to promising venues for improvements in more applied directions. To this end, we bring together theoretical results from different areas to present our theoretical arguments for the utility of a computational tool, in particular graph curvature, which enables a better understanding and further improvements of transformers.

*”Provide a more detailed description of Figure 3...”

We agree that the description of Figure 3 seems insufficient. We will elaborate and clarify.

---

---

Questions:

1. We appreciate this question and will clarify this in our resubmission. Note that Eq. (3) applies locally, hence a more positive-leaning ORC distribution is expected to correspond to local neighborhoods with larger lower bounds. Furthermore, our experiments that modify curvature distribution also modify the minimum ORC.

2. In this paper we consider the attention map as a weighted graph. This point of view opens the door to a rich toolset of graph analysis, including graph geometry. A particular advantage of graph curvature is that it characterizes the local geometry of the attention map, while many other graph analysis tools are less sensitive to local geometry and directly capture more global topological features. Moreover, as we explain in the background, curvature is linked to entropy and robustness, which are fundamental concepts.

---

---

We will make all the above points clearer in our future submission.

We sincerely thank the reviewer for their careful and constructive review. We value this feedback and will improve our paper based on the comments provided. We address the main concerns below and questions in the following comment under this thread, all of which will be incorporated and clarified in our future submission.

---

---

Weaknesses:

1. We understand that our comments on relaxing the assumptions may have been unclear and miscommunicated in the writing. We did not intend to suggest that our analysis holds as is (or with minor adjustment) if the assumptions are relaxed, but rather that relaxing the assumption to a less restrictive version and redoing the analysis can lead to similar results, which is admittedly non-trivial. We agree with the reviewer that relaxing the assumption is non-trivial, and our statement could lead the reader to underestimate what it takes to revise the analysis under relaxed assumptions. We will revise this statement and, following your suggestion, we will revise our proof strategy in order to relax this assumption.

2. We acknowledge that our analysis is conducted in an idealized setting. In our resubmission, we will include further discussion on the assumptions being restrictive and the purpose our analysis serves despite this. It is our understanding that analysis in idealized settings can inspire clues for better understanding or improved heuristics which can be empirically tested and lead to insightful findings. This is the approach we take, which is not only limited to the work by Dovonon et al, but also other works on transformers. For instance, Von Oswald et al. 2023 [1] derive their results under a specific construction, which leads to highly influential and insightful findings. That being said, we do not intend to downplay the restrictiveness of our assumptions, and we will clarify this in our revision.

3. Thank you for bringing this to our attention. We will certainly include a discussion on this point and will also work on supplementing our empirical observations with experiments using the MSE loss.

4. We sincerely thank the reviewer for identifying these points, and will make sure to address them properly.

• Regarding the first major point, we seemed to have made a mistake (possibly a typo) in the indexing of ${\lambda^L}$, incorrectly using ${\lambda^L}_1$ instead of ${\lambda^L}_n$. This mistake has affected the next step of the proof, where we used a lower bound on ${\lambda^L}_1$ instead of an upper bound on ${\lambda^L}_n$ ---both a result from Bauer et al. 2011. We will carefully revise and fix our analysis following this correction. We will also revise and supplement the experiments regarding this result to validate the implications of our revised analysis.

• Regarding the second major point, we recognize that this needs either filling a gap in the proof or a clarification and revision of the presentation of our results to distinguish rigorous claims from educated guesses. From a computational perspective, the main intended outcome of this analysis was to optimize a bound as a surrogate for adjusting a target variable, where the bound is straightforward to compute but the target variable is not. For stating a rigorous statement, we agree with you; there is a missing step in the proof. This gap could be filled under a viable assumption, and we can supplement our experiments to test the practical applicability of this assumption. However, given the feedback provided on the set of assumptions for our analysis, we will instead consider limiting the theoretical results to the bounds and instead discuss our argument for the hypothesized implications that we empirically investigate. We will also supplement our experiments to strengthen this investigation.

5. We will revise the related work to better discuss the literature and clarify the gap we aim to fill.

---

”Altogether, the study of attention layers' geometry shows potential…”

Once again, we thank the reviewer for their valuable and constructive comments and we will revise and make the improvements mentioned above before resubmission.

---

---

We will clarify all the points above in our future resubmission.

---

---

References:

[1] Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., & Vladymyrov, M. (2023, July). Transformers learn in-context by gradient descent. In International Conference on Machine Learning (pp. 35151-35174). PMLR.

We would like to thank the reviewer for their review of our work. Next, we address the reviewer’s specific comments.

---

---

Weaknesses:

1. As you point out and we acknowledge in the paper, the theoretical analysis is performed in an idealized setting, as it is often the case in the community. We do not intend to underestimate the restrictiveness of the assumptions, and we will further clarify this upon revision. Also as you mentioned, we provide extensive empirical results, which aim to assess the applicability of our theoretically educated arguments in practical scenarios. We will also clarify the following points:

“Assumption A0 implies that the attention matrix A is symmetric”

While this implies that the similarities matrix inside the softmax is symmetric, this assumption alone does not necessarily imply that the attention matrix is symmetric, due to the value matrix, $V$.

”A1 is even more restrictive, as it requires A to be invertible”

We agree that invertibility is a restrictive assumption. Meanwhile, it's worth noting Dovonon et al. [1] assume this invertibility and state that this tends to hold in their experiment, and our theoretical analysis partially builds upon their findings and results.

”the authors assume that A does not depend on the input X”

Respectfully, this seems to be misunderstood. As defined in Section 2, $A$ depends on the input through $Q=X W_Q$ and $K = X W_K$.

In summary, while we agree on the restrictiveness of the assumptions, our approach is to utilize theory-informed heuristics for developing a better understanding and pointing to directions for methodological improvements, aligning with a common practice in deep learning research. Our theory aims to inform and inspire, and our empirical investigation seeks to verify whether the trends our theory suggests within idealized settings are observed in practical settings.

---

2. The link between curvature and robustness is indeed established in prior works and this is not against our novelty but in favor/supportive of our work. We discuss this connection in the Background section, citing the references you have provided, and do not claim it to be our contribution. However, this paper’s main contribution is building upon and exploring this link in the context of transformers.

---

3. We considered 3 image and 3 text datasets on 2 vision and 2 language models to provide ample empirical evidence. We will consider supplementing the image datasets to improve this paper for future resubmission.

---

---

Questions:

1. We indeed leverage tools and results from existing papers we cited in areas such as learning theory and graph geometry. The proof we provide (in Appendix A.2) is itself new though, combining these tools for our specific purpose.

2. This paper does not aim to introduce a new model, but rather to improve our understanding of properties of existing architectures. Hence, we did not find a comparison to other baselines informative for our purpose.

---

---

We will clarify all the points above in our future resubmission.

---

References:

[1] Dovonon, G. J., Bronstein, M. M., & Kusner, M. J. (2024). Setting the record straight on transformer oversmoothing. arXiv preprint arXiv:2401.04301.

Thank you for reviewing our response, and thank you again for your constructive feedback. As we explained in our response, we noticed the error following your feedback, and there is a clear path to fixing the error, which we will follow. Also, as we mentioned in our responses, we plan to do this carefully and resubmit in the future with the correction and improvements discussed. To clarify, by that we meant for another venue; we had already planned to withdraw. We did not withdraw right away because we wanted to have the opportunity to openly thank the reviewers for their superb feedback and inform them that all their comments are addressable but we needed more time. Appreciating the rebuttal process and the chance to have a discussion, we decided to respond and wait to hear more feedback on our response, but that was not with the intention to get this version of the paper published.

In conclusion, as we had intended, we will withdraw, correct the error and incorporate other improvements discussed and will resubmit elsewhere.