技术摘要：如何驯服 Grokking：将表示几何学作为控制信号

问题陈述

Grokking 是过度参数化神经网络中的一种现象，即模型最初会记忆训练数据，并在长时间优化后才表现出强大的泛化能力。虽然近期的研究已将 grokking 与权重衰减、特征压缩、电路形成和表示学习动态联系起来，但决定这种延迟泛化发生时机和出现的具体因素仍未被完全理解。现有文献的一个共同主题是，学习到的表示会逐渐压缩到低维结构中。然而，以往的工作主要将表示几何学视为一种涌现属性，而非直接的控制机制。本文研究了显式操纵隐藏表示的几何结构是否可以作为一种控制信号，从而影响 grokking 的发生时机和速度。

方法论

作者引入了 几何维度正则化 (GeomDR)，这是一种表示层级的谱正则化器，旨在直接控制训练过程中隐藏表示的有效维度。

几何特征化

对于隐藏层表示矩阵 $Z \in \mathbb{R}^{N \times d}$，该方法首先对数据进行中心化处理以获得 $\tilde{Z}$。随后计算经验协方差矩阵 $C = \frac{1}{N-1}\tilde{Z}^\top\tilde{Z}$。有效维度 ($D_{eff}$) 使用协方差矩阵特征值 ($\mu_i$) 的参与比（participation ratio）进行计算：
$$D_{eff} = \frac{(\sum_i \mu_i)^2}{\sum_i \mu_i^2}$$
此外，还定义了一个局部邻域距离 ($\rho$)，即归一化表示的 $k$ 个最近邻之间的平均欧几里得距离。

GeomDR 目标函数

GeomDR 对目标维度为 $d^*$ 的子空间之外所包含的方差施加惩罚。给定按降序排列的特征值 ($\mu_1 \ge \mu_2 \ge \dots \ge \mu_d$)，逐层损失函数为：
$$L_{GeomDR}^{(\ell)} = \sum_{j=d^*+1}^{d} \mu_j^{(\ell)}$$
总训练目标结合了任务损失 ($L_{task}$) 和几何正则化项：
$$L = L_{task} + \lambda(t)L_{GeomDR}$$
至关重要的是，该正则化器并非从一开始就激活。它是在一个初始的无约束训练阶段（激活步 $t_s$）之后引入，并使用余弦渐进（cosine ramp）逐渐增加其强度 $\lambda(t)$。这种设计模拟了自然 grokking 的轨迹，即记忆化先于泛化发生。

实验方案

该框架在三个以表现出 grokking 特征而闻名的算法任务上进行了评估：

1. 模加法 (Modular Addition)： 预测 $(a + b) \mod 97$。
2. 模除法 (Modular Division)： 预测 $(a \cdot b^{-1}) \mod 97$。
3. 置换组合 (Permutation Composition)： 预测两个置换在 $S_5$ 中的组合。

实验使用了两种架构：多层感知机 (MLP) 和 Transformer。研究系统地改变了干预方案（强度 $\lambda_{max}$ 和渐进时长 $T_{ramp}$）、目标维度 ($d^*$) 以及干预起始时间 ($t_s$)。

核心贡献

1. 观察到前置坍缩： 作者证明了在基准设置中，维度坍缩一致地发生在从记忆化向泛化转变（g%，即 grokking）之前。这表明几何压缩不仅是一个相关特征，更是这种转变的机械驱动力。
2. 引入 GeomDR： 他们提出了 GeomDR，一种谱正则化器，通过直接抑制目标子空间之外的协方差谱，从而控制隐藏表示的有效维度。
3. 系统性控制动力学： 论文提供了一项全面的实证研究，表明几何干预可以显著改变延迟泛化的动力学过程。根据干预方案和目标维度的不同，GeomDR 可以加速、延迟或完全抑制 grokking。

结果

加速 Grokking

在所有评估的设置中，GeomDR 一致地改变了 grokking 的动力学：

• MLP 性能： 在模加法任务中，GeomDR 加速了 grokking 高达 52 倍（将步骤从约 362K 减少到约 7K）。在模除法（约 39 倍）和置换组合（约 23 倍）中也观察到了类似的提升。在所有 MLP 案例中，GeomDR 实现了 100% 的成功率，而基准模型则表现不稳定或失败。
• Transformer 性能： 虽然 Transformer 的基准 grokking 动力学已经比 MLP 更快，但 GeomDR 仍然提供了持续的加速（例如在模加法上约为 2 倍），并保持了 100% 的成功率。

超参数的影响

• 目标维度 ($d^*$)： 中等到较大的目标维度 ($16 \le d^* \le 64$) 产生了最强的加速效果。过于激进的压缩 ($d^* \le 8$) 往往会减慢或破坏学习过程。
• 干预时机 ($t_s$)： 研究识别出了一个“关键时间窗口”。干预应用过早（在有用表示形成之前）或过晚（在记忆化变得僵化之后）效果均较差。中间阶段的激活时间（例如 MLP 中的 $t_s = 2000$）能产生最快且最稳定的 grokking。
• 方案 (Schedule)： 中等强度的正则化和渐进时长通常表现最好；过强的干预可能会导致训练不稳定。

几何相关性

研究证实，加速 grokking 的干预同时也加速了维度坍缩的发生。在 GeomDR 模型中，维度坍缩与成功泛化之间的时差与基准模型相比显著缩小，这表明几何重组与泛化涌现之间存在更紧密的耦合。

可扩展性与效率

作者证明，即使从表示的随机子集（例如 5%–25% 的数据）中估计协方差谱，GeomDR 仍然有效，这在不实质影响 grokking 加速的情况下，实现了 2.4x–2.7x 的运行速度提升。

意义与主张

本文声称，表示几何学可以作为 grokking 有效的控制信号。通过直接修改隐藏激活的协方差谱，研究人员可以系统地影响网络何时以及如何快速地从记忆化转向泛化。

作者认为，这项工作提供了证据，证明几何干预是研究和影响延迟泛化的实用方法。与其被动观察表示的变化，该方法允许直接操纵几何属性来测试假设。结果表明，表示几何与 grokking 之间的关系在不同任务和架构中都是稳健的，尽管提升幅度有所不同（在 MLP 中比在 Transformer 中更为显著，这可能是由于后者具有更强的归纳偏置）。

论文得出结论，虽然它并未确立表示维度是唯一的因果机制，但它证明了几何干预足以改变 grokking 的时机，从而为优化和泛化研究提供了一种新的实验方法论。