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

TL;DR

This work identifies a fundamental shortcoming in traditional sharpness-based generalization metrics when applied to transformers, due to rich, continuous symmetries in attention. It introduces geodesic sharpness, a symmetry-corrected sharpness measure defined on the Riemannian quotient manifold obtained by quotienting out these symmetries, and computes perturbations along geodesics rather than Euclidean paths. Through analyses of diagonal networks and empirical studies on Vision Transformers and language models, the authors demonstrate that higher-order geometric corrections restore strong correlations between sharpness and generalization, outperforming adaptive sharpness on transformers. The approach provides a principled framework to account for symmetry-induced ambiguities in parameter space, with implications for diagnosing generalization and guiding symmetry-aware training in large-scale models.

Abstract

The concept of sharpness has been successfully applied to traditional architectures like MLPs and CNNs to predict their generalization. For transformers, however, recent work reported weak correlation between flatness and generalization. We argue that existing sharpness measures fail for transformers, because they have much richer symmetries in their attention mechanism that induce directions in parameter space along which the network or its loss remain identical. We posit that sharpness must account fully for these symmetries, and thus we redefine it on a quotient manifold that results from quotienting out the transformer symmetries, thereby removing their ambiguities. Leveraging tools from Riemannian geometry, we propose a fully general notion of sharpness, in terms of a geodesic ball on the symmetry-corrected quotient manifold. In practice, we need to resort to approximating the geodesics. Doing so up to first order yields existing adaptive sharpness measures, and we demonstrate that including higher-order terms is crucial to recover correlation with generalization. We present results on diagonal networks with synthetic data, and show that our geodesic sharpness reveals strong correlation for real-world transformers on both text and image classification tasks.

Figures (13)

• Figure 1: Quantities from the Riemannian quotient manifold respect the loss landscape's symmetry; Euclidean quantities do not. We illustrate this here for a synthetic least squares regression task with a two-layer NN, where $x \mapsto \theta_2 \theta_1 x$ with scalar parameters ${\bm{\theta}} \in {\mathbb{R}}^2$ and input $x \in {\mathbb{R}}$ (i.e. each layer is a linear function). The NN is re-scale invariant, i.e. has $\mathop{\mathrm{GL}}\nolimits(1)$ symmetry: For any $\alpha \in {\mathbb{R}} \setminus \{ 0 \}$, the parameters $(\theta_1', \theta_2') = (\alpha^{-1}\theta_1, \alpha \theta_2)$ represent the same function. (\ref{['subfig:visualization-scalar-toy-loss']}) The loss function inherits this symmetry and has hyperbolic level sets. (\ref{['subfig:visualization-scalar-toy-euclidean-gradient']}) The Euclidean gradient norm does not share the loss function's geometry and changes throughout an orbit where the NN function remains constant. (\ref{['subfig:visualization-scalar-toy-riemannian-gradient']}) The Riemannian gradient norm follows the loss function's symmetry and remains constant throughout an orbit, i.e., it does not suffer from ambiguities for two points in parameter space that represent the same NN function.
• Figure 2: Illustrative sketch relating total and quotient space and their tangent spaces. A tangent vector at a point in total space, $\bar{\xi}_{\bar{x}}\in T_{\bar{x}} \overline{{\mathcal{M}}}$ can be decomposed into a horizontal component $\bar{\xi}_{\bar{x}}^{{\mathcal{H}}}$ and a vertical component $\bar{\xi}_{\bar{x}}^{{\mathcal{V}}}$. The vertical component points along the direction where the quotient space $x = [\bar{x}]$ remains unaffected. The horizontal component points along the direction that changes the equivalence class. We can use $\bar{\xi}_{\bar{x}}^{{\mathcal{H}}}$ as a representation of the tangent vector $\xi_{x}\in T_x {\mathcal{M}}$ on the quotient space. The component $\bar{\xi}_{\bar{x}}^{{\mathcal{H}}}$ represents the horizontal lift of $\xi_{x}$.
• Figure 3: We show for 50 diagonal models trained on sparse data the generalization gap (here just the test loss) vs. worst-case adaptive sharpness (left, $\tau=-0.68$), $\langle \cdot, \cdot \rangle^{\text{inv}}$-geodesic sharpness (middle, $\tau=-0.83$), and $\langle \cdot, \cdot \rangle^{\text{mix}}$-geodesic sharpness (right, $\tau=-0.86$). The x axis is the value of the sharpness measure being considered, and the y axis is the test error.
• Figure 4: We show for the 72 models from wortsman2022soups the generalization gap on ImageNet vs. worst-case adaptive sharpness (left, $\tau=-0.41$), $\langle \cdot, \cdot \rangle^{\text{inv}}$-geodesic sharpness (middle, $\tau=-0.71$), and $\langle \cdot, \cdot \rangle^{\text{mix}}$-geodesic sharpness (right, $\tau=-0.70$).
• Figure 5: We show for 35 models from mccoy2020bertsfeather the generalization gap on the MNLI dev matched set mnli2018 vs. worst-case adaptive sharpness (left, $\tau=0.06$), $\langle \cdot, \cdot \rangle^{\text{inv}}$-geodesic sharpness (middle, $\tau=0.28$), and $\langle \cdot, \cdot \rangle^{\text{mix}}$-geodesic sharpness (right, $\tau=0.38$).

Theorems & Definitions (11)

• Definition 3.1: Functional $\mathop{\mathrm{GL}}\nolimits$-symmetric building block
• Definition 3.2
• Example 1.1: Self-attention vaswani2017attention
• Example 1.2: Shallow linear net
• Example 1.3: Low-rank adapters (LoRA, hu2022lora)
• Definition 2.1
• Definition 2.2
• Definition 2.3
• Definition 2.4
• Definition 2.5