Go to ICLR 2025 Conference homepageThe Geometry of Attention: Ricci Curvature and Transformers Training and Robustness
Transformer models have revolutionized machine learning, and the theoretical underpinnings behind their success are only now starting to be explored. In this work, we analyze the performance and robustness of transformers by considering the attention mechanism as a graph operator, focusing on the geometry of attention maps viewed as weighted graphs. Specifically, we investigate the role of Ricci curvature, a metric closely tied to graph spectral properties and system robustness, in shaping the training dynamics and robustness of transformers. Our theoretical analysis establishes a link between Ricci curvature and the convergence of gradient descent on transformers, and consequently, their training and fine tuning. We also show that a higher frequency of more positive values in the Ricci curvature distribution of attention graphs, therefore more system robustness, leads to more robust transformers, highlighting the impact of curvature on the robustness of transformers. Leveraging these insights, we propose an efficient regularization method to train curvature-adjusted transformers. Supporting our theoretical findings, experiments show that our proposed attention curvature manipulation can improve the learning speed, performance, or generalizability of vision and language transformers. Additionally, our observations point to a trade-off between their performance and robustness. This work demonstrates that the geometry of the attention map provides a theoretically elegant and computationally versatile framework for analyzing and manipulating transformers training, generalization, performance, and robustness, opening new avenues for designing models using geometric concepts.