泛化的奥秘——即神经网络为何能在特定图像或文本上训练后,在未见过的数据上也能表现良好——是深度学习中的“暗物质”问题。

多年来,一个主要的假设是 锐度 (Sharpness) 的概念 (或其反面: 平坦性) 。直觉很简单: 如果神经网络在损失景观的“平坦”山谷中找到解,那么这个解就是鲁棒的。如果训练数据发生轻微偏移 (模拟训练集和测试集之间的差异) ,损失不会剧增。相反,“尖锐”的极小值意味着即使是微小的偏移也会导致高误差。

这一理论对多层感知机 (MLP) 和卷积神经网络 (CNN) 相当有效。但最近出现了一个问题。当研究人员将这些锐度度量应用于 Transformer——驱动 ChatGPT、Claude 和 Gemini 的架构——时,这种相关性破裂了。平坦性似乎不再能可靠地预测泛化能力。

在一篇名为 “Hide & Seek: Transformer Symmetries Obscure Sharpness & Riemannian Geometry Finds It” 的精彩论文中,研究人员 da Silva、Dangel 和 Oore 指出,问题不在于锐度这个概念本身,而在于我们要如何测量它。Transformer 拥有丰富且复杂的 对称性 , 这些对称性扭曲了欧几里得空间。为了真正看清这个景观,我们必须放弃标准几何,转而采用 黎曼几何

在这篇文章中,我们将梳理他们的推导过程,从损失景观的基本直觉出发,深入到商流形的复杂几何,最终看到这种新的“测地线锐度”如何成功预测大规模 Transformer 的泛化能力。

对称性问题

要理解为什么标准锐度会失效,我们首先需要理解神经网络中的 对称性 。 参数空间中的对称性意味着你可以改变网络的权重而不改变网络计算的函数 (因此也不改变损失) 。

欧几里得空间的错觉

考虑一个具有缩放对称性的简单网络层。如果你将一层的权重加倍并将下一层的权重减半,输出在数学上可能保持不变。在标准的欧几里得空间中,这两组权重看起来像是两个相距甚远的点。但就它们所代表的函数而言,它们是同一个点。

这给测量锐度带来了巨大的问题。

Figure 1: Quantities from the Riemannian quotient manifold respect the loss landscape’s symmetry; Euclidean quantities do not.

请看上方的 图 1 。 这个玩具示例展示了一个具有缩放对称性 (具体为 GL(1) 对称性) 的损失景观。

  • (a) 损失 (The Loss): 注意那些双曲线形状的“山谷”。损失值沿着这些曲线是恒定的。
  • (b) 欧几里得梯度 (Euclidean Gradient): 如果你使用标准欧几里得数学计算梯度范数,你会得到同心圆。这产生了一个矛盾: 当你沿着损失函数恒定的曲线移动时,梯度值却在变化!欧几里得度量被参数缩放“迷惑”了。
  • (c) 黎曼梯度 (Riemannian Gradient): 通过调整我们的几何以尊重对称性,梯度范数在恒定损失曲线上变得恒定。

作者认为,标准的锐度度量就像图 1(b) 那样——它们对本不该重要的参数缩放非常敏感。

Transformer 的对称性: GL(h)

虽然 CNN 具有简单的缩放对称性,但 Transformer 拥有更复杂的东西: 一般线性群对称性 , 即 \(\mathrm{GL}(h)\)。

Definition of GL(h)

在 Transformer 的注意力机制中,我们有查询 (\(Q\)) 、键 (\(K\)) 和值 (\(V\)) 的矩阵。注意力的数学运算涉及 \(QK^T\) 这样的乘积。由于这种矩阵乘法,我们可以在权重之间插入一个可逆矩阵 \(A\) 及其逆矩阵 \(A^{-1}\),而不会改变结果。

Action of the symmetry group

如上所示,如果我们使用矩阵 \(A\) 变换权重 \(G\) 和 \(H\) (代表注意力头的一部分) ,函数保持完全一致。这不仅仅是放大或缩小一个数字;这是一个高维变换,它扭曲了整个参数空间。

标准的“自适应锐度” (ASAM) 试图通过归一化权重来修复简单的缩放问题,但在数学上不足以处理这种复杂的 \(\mathrm{GL}(h)\) 对称性。这就解释了为什么现有的度量无法预测 Transformer 的泛化能力。

引入黎曼几何

为了解决这个问题,作者提出了一个根本性的转变: 停止在“全空间” (原始参数 \(\theta\)) 中测量锐度,开始在 商流形 上测量它。

商流形 (The Quotient Manifold)

想象参数空间 \(\overline{\mathcal{M}}\) 是一张纸。由于对称性的存在,这张纸上的许多点代表完全相同的神经网络函数。我们把所有等价点的集合称为 轨道 (orbit)

如果我们将每个轨道坍缩成一个点,我们就得到这就得到了一个新的、弯曲的形状,称为 商流形 (\(\mathcal{M}\))。在这个流形上,每个点代表一个唯一的函数。这里没有歧义。

挑战在于,我们不能直接在这个抽象流形上运行梯度下降。我们必须在原始参数 (全空间) 上运行它。我们需要一种数学方法将弯曲商流形上的概念转换回我们具体的参数空间。

垂直空间与水平空间

为了连接这两个世界,作者引入了切空间 (即权重可能更新的空间) 的分解。

在参数空间的任意一点,我们可以将权重的变化 (\(\delta\)) 分解为两个分量:

  1. 垂直空间 (\(\mathcal{V}\)): 沿此方向移动是在对称轨道上 移动。权重改变了,但网络的功能 (和损失) 保持完全不变。
  2. 水平空间 (\(\mathcal{H}\)): 沿此方向移动是 垂直于 轨道的。这是实际修改网络功能的改变,剔除了对称性噪声。

Figure 2: Illustrative sketch relating total and quotient space and their tangent spaces.

图 2 清晰地展示了这一点。

  • 马鞍形状代表 全空间 (\(\overline{\mathcal{M}}\))。
  • 连接 \([x]\) 和 \([x']\) 的曲线是一个轨道 (垂直空间) 。
  • 为了测量真实的锐度,我们只关心 水平分量 (\(\xi^H\))。这个向量代表“水平提升 (Horizontal Lift)”——即剔除对称性噪声后的网络功能改变。

定义黎曼度量

为了使之严谨,我们需要一个 黎曼度量——一种尊重对称性的测量距离和角度的方法。作者为注意力层提出了两种特定的度量: 不变量 (Invariant) 度量和 混合 (Mixed) 度量。

Definitions of Invariant and Mixed Metrics

这些度量 (\(\langle \cdot, \cdot \rangle^{\mathrm{inv}}\) 和 \(\langle \cdot, \cdot \rangle^{\mathrm{mix}}\)) 通过权重矩阵 (\(G\) 和 \(H\)) 的逆或转置来加权梯度。请注意这与标准欧几里得点积有多么不同。这些度量确保无论你如何使用对称群缩放或变换权重,测得的距离都保持一致。

测地线锐度 (Geodesic Sharpness)

现在我们来到了核心贡献: 测地线锐度

标准锐度度量 (如 SAM 或 Random Sharpness) 通过添加一个向量 \(\delta\) 来扰动权重:

\[ \theta_{new} = \theta + \delta \]

在几何上,这是沿直线移动。但在一个弯曲空间 (如我们的商流形) 中,“直线”实际上是被称为 测地线 (geodesics) 的曲线。

如果你试图在弯曲的球面上走直线,你会偏离预定路径。同样,添加线性欧几里得扰动到 Transformer 权重中会产生扭曲,因为它忽略了对称群的曲率。

公式化

作者利用测地线重新定义了最坏情况下的锐度:

Equation for Geodesic Sharpness

这里,\(\bar{\gamma}_{\bar{\xi}}(1)\) 代表在以速度 \(\xi\) 开始沿测地线曲线移动一个时间步长后到达的点。

然而,计算精确的测地线在计算上非常昂贵。我们无法为训练的每一步都求解微分方程。作者通过使用二阶泰勒展开来近似测地线路径来解决这个问题。

Approximation of the geodesic path

这个等式是“秘方”。

  • 第一项 \(\bar{\xi}^i t\) 是标准的线性更新。
  • 涉及 \(\Gamma_{kl}^i\)( 克里斯托费尔符号 )的第二项是对曲率的修正。

这揭示了一个惊人的洞见: 标准的自适应锐度只是测地线锐度的一阶近似。它假设空间是平坦的 (克里斯托费尔符号 = 0) 。通过添加二阶项,作者修正了由 Transformer 复杂对称性引起的曲率。

实验: 它有效吗?

理论很美好,但它能预测泛化吗?作者在三种不同的架构上进行了测试。

1. 对角网络 (合成测试)

他们从对角线性网络开始,这是一个数学上易于处理的简单架构。

Figure 3: Generalization gap vs. sharpness for diagonal models.

图 3 中,请看相关系数 (\(\tau\))。

  • 自适应锐度 (左): \(\tau = -0.68\)。还不错的相关性。
  • 测地线锐度 (中/右): \(\tau = -0.83\) 和 \(-0.86\)。

黎曼方法显著增强了相关性。 (注: 此处的负相关意味着在特定的度量设置下,较高的锐度值与较低的测试误差相关,这与通常的“平坦即好”的启发式方法相反,但这是一种有趣的倒置,重要的是预测关系的 强度) 。

2. 视觉 Transformer (ImageNet)

这是真正的测试。他们使用了 72 个在 ImageNet 上微调过的 CLIP 模型 (ViT-B/32)。先前的工作表明,自适应锐度无法预测此处的泛化。

Figure 4: Generalization gap vs. sharpness for ImageNet models.

图 4 展示了结果:

  • 自适应锐度 (左): \(\tau = -0.41\)。数据点是一团散乱的云。信号很弱。
  • 测地线锐度 (中/右): \(\tau = -0.71\) 和 \(-0.70\)。

突然间,一个清晰的趋势出现了。通过考虑对称性曲率,锐度测量中的“噪声”变成了一个清晰的信号。黎曼几何成功“找到”了被对称性掩盖的锐度。

3. 语言模型 (BERT on MNLI)

最后,他们研究了在 MNLI 数据集上微调的 BERT 模型。

Figure 5: Generalization gap vs. sharpness for BERT models.

图 5 对旧方法来说可能是最致命的打击:

  • 自适应锐度 (左): \(\tau = 0.06\)。实际上是零相关。纯粹是随机噪声。
  • 测地线锐度 (中/右): \(\tau = 0.28\) 和 \(0.38\)。

虽然相关性并不完美,但它在欧几里得度量完全失效的地方发现了信号。

“曲率修正”为何重要

为了证明复杂的数学 (克里斯托费尔符号) 是必要的,作者进行了消融研究。他们关掉了方程中的二阶修正项,看看是否仅仅使用黎曼度量范数就足够了。

Ablation study results

图 8 (右图) 中,使用完整的测地线锐度,他们达到了 \(\tau = 0.38\)。在中图中,当他们关掉二阶权重修正 (忽略曲率) 时,相关性下降到 \(\tau = 0.24\)。这证实了 跟随曲线 (测地线) 至关重要,而不仅仅是归一化向量。

结论与启示

论文 “Hide & Seek” 对我们理解深度学习优化提出了深刻的修正。它表明“平坦性”假设没有错,而是我们的欧几里得尺子坏了。

关键要点:

  1. 对称性很重要: Transformer 具有高维对称性 (\(\mathrm{GL}(h)\)),这扭曲了参数空间。
  2. 欧几里得是盲目的: 标准梯度下降和锐度度量将参数视为平坦的欧几里得向量,忽略了这些扭曲。
  3. 黎曼是真实的: 通过将参数空间建模为商流形并沿测地线移动,我们恢复了损失景观几何与模型性能之间的强相关性。

这项工作为新的优化器打开了大门。就像 SAM (锐度感知最小化) 显式地优化平坦性一样, 测地线 SAM 可以优化黎曼平坦性,这可能导致 Transformer 的泛化能力超越当前最先进的模型。数学虽然更重,但正如这篇论文所示,几何学发现了对称性所隐藏的东西。