特征排斥与谱锁定:两层网络 Grokking 现象的实证研究

arXiv cs.LG 论文

摘要

这项实证研究验证了关于两层神经网络在 Grokking 现象期间的特征排斥和谱锁定理论发现,展示了激活函数如何影响从记忆到泛化的过渡。

arXiv:2605.08119v1 公告类型:new 摘要:Tian (2025) 证明了 Grokking 互动特征学习阶段矩阵 $ B = (\widetilde{F}^\top \widetilde{F} + \eta I)^{-1} $ 的排斥定理(定理 6):相似特征具有负非对角线元素 $ B_{j\ell} $,产生一种有效的排斥力,将它们推开。然而,该定理并未明确指出此机制何时在经验上可观察,也未说明它是否在参数更新中留下可测量的谱签名。我们在 Tian 的模块化加法设置($ M = 71 $,$ K = 2048 $,MSE 损失)上直接测试了这一点,并观察到明显的结构-机制分离。预测的符号规则在激活函数(empirical sign-match 从 $ \sigma = x^2 $ 的 0.865 上升到 0.985,跨越 5 个种子,并在 $ \sigma = \operatorname{ReLU} $ 时饱和至 1.000)的前 200 个最相似特征对上稳健成立。然而,参数更新中的谱签名强烈依赖于激活函数。当 $ \sigma = x^2 $ 时,简单斜率探测器在 15/15 个 Grokking 种子的第 174 个 epoch 处触发(IQR [173,174]),而在 0/15 个非 Grokking 对照组中未触发,晚期幅度分离为 229$ \times $;谱为秩 2。相比之下,当 $ \sigma = \operatorname{ReLU} $ 时,探测器从未触发,谱保持在有效秩 1。这种分离与 Tian 定理 5 中集中(幂律)和扩散(ReLU)记忆的区别一致:虽然 $ B $ 的符号结构仅取决于 $ \widetilde{F}^\top \widetilde{F} $,但特征排斥如何转化为权重更新关键取决于激活导数 $ \sigma' $。
查看原文
查看缓存全文

缓存时间: 2026/05/12 06:44

# 两层网络Grokking的实证研究代码位于 https://github.com/skydancerosel/grokking-integrability/tree/main/tian_eigengap。
来源:https://arxiv.org/html/2605.08119
## 特征排斥与谱锁定:两层网络Grokking的实证研究††致谢:代码位于 https://github.com/skydancerosel/grokking-integrability/tree/main/tian_eigengap。

###### 摘要

Tian(2025)(https://arxiv.org/html/2605.08119#bib.bib5) 证明了关于两层网络在grokking互动特征学习阶段中 $B=(\tilde{F}^\top\tilde{F}+\eta I)^{-1}$ 的非对角结构的排斥定理(定理6),但未具体说明该机制在训练何时在经验上变得可观察。我们在Tian的精确模加设置($M=71, K=2048, n=2016$, MSE)上直接测试了该定理和一个候选谱可观测量。

定理6在激活函数间成立。符号规则 $\operatorname{sgn}(B_{j\ell})=-\operatorname{sgn}(\tilde{f}_j^\top P_{\eta,-j\ell}\tilde{f}_\ell)$ 在五个确定性回放检查点上的前200对最相似特征对上进行了验证,跨越 $n=5$ 个种子。经验一致性从 epoch 50 的 0.865 [IQR 0.865, 0.875] 上升到 epoch 300 的紧密饱和值 0.985 [IQR 0.980, 0.990](当 $\sigma(x)=x^2$ 时)。在 $\sigma=\text{ReLU}$ 上,相同的符号规则**也**成立,饱和速度更快(在 epoch 500 达到 1.000)。该机制具有激活函数通用性。

参数更新的谱签名具有激活函数特异性。参数更新 Gram 矩阵 $\Delta W$ 的滚动窗口特征间隙(eigengap)$\sigma_2/\sigma_3$ 仅在定理6饱和**且**特征坍缩到尖锐峰值(Tian定理5的*专注记忆*机制,幂激活函数的特征)时触发。对于 $\sigma=x^2$,基于斜率的检测器在 15/15 个grok种子中的 epoch 174(IQR=[173,174])触发,而在 0/15 个控制种子中未触发,后期幅度分离达到 $229\times$。对于 $\sigma=\text{ReLU}$——Tian定理5将其置于*扩散记忆*机制中——相同的检测器在 0/15 个grok种子中未触发,后期幅度分离坍缩至 $1.4\times$,且谱以秩1为主导而非秩2。

这两项发现共同勾勒出结构与机制的区别。Tian的定理6通过 $\tilde{F}^\top\tilde{F}$ 的性质支配 $B$ 的非对角符号结构,这仅通过激活函数对 $\tilde{F}$ 的影响依赖于激活函数。然而,参数更新谱中的签名取决于排斥**如何**转化为权重更新——这取决于 $\sigma'$。幂激活函数($\sigma(x)=x^2$)产生专注的特征,这些特征巩固到两个持续的秩2方向上;ReLU 产生扩散的特征,这些特征保持秩1主导。我们将此与定理5的专注与扩散区别联系起来。

我们还报告了支持性发现:锁定检测器对窗口大小敏感($W\leq 10$ 在 $\eta=0$ 控制中产生假阳性;$W\in\{20,30\}$ 给出完美特异性);在小窗口下 $\sigma_3,\sigma_4,\sigma_5$ 一起坍缩,证实了在 finest 时间分辨率下的秩2;水平指标检测器 $\rho_{\text{tian}}$ 在 $\eta=10^{-5}$ 时的提前量符合 Tian 的 $1/\eta$ 预测(提前 567 个 epoch,grokking 在 epoch 1527)。

## 1 引言

Grokking——在记忆之后很久突然出现的泛化现象(Power et al., 2022)(https://arxiv.org/html/2605.08119#bib.bib4)——通过机制可解释性(Nanda et al., 2023)(https://arxiv.org/html/2605.08119#bib.bib3)、权重衰减作为隐式正则化(Liu et al., 2022)(https://arxiv.org/html/2605.08119#bib.bib2) 以及懒惰到丰富的过渡(Kumar et al., 2024)(https://arxiv.org/html/2605.08119#bib.bib1) 积累了各种解释。Tian(2025)(https://arxiv.org/html/2605.08119#bib.bib5) 提供了最原则性的框架:Li2 将两层网络中的 grokking 动力学分解为三个阶段——*懒惰*学习、*独立*特征学习和*互动*特征学习——其特征是反向传播梯度 $G_F$ 和激活 Gram 矩阵 $\tilde{F}^\top\tilde{F}$ 的结构逐渐丰富。

在第三阶段中,Tian的定理6(*相似特征的排斥*)断言 $B:=(\tilde{F}^\top\tilde{F}+\eta I)^{-1}$ 的非对角元素满足

$$
\operatorname{sgn}(B_{j\ell}) = -\operatorname{sgn}\left(\tilde{f}_j^\top P_{\eta,-j\ell}\tilde{f}_\ell\right), \quad P_{\eta,-j\ell}:=I-\tilde{F}_{-j\ell}(\tilde{F}_{-j\ell}^\top\tilde{F}_{-j\ell}+\eta I)^{-1}\tilde{F}_{-j\ell}^\top, \quad (1)
$$

其中 $\tilde{F}_{-j\ell}$ 排除第 $j$ 列和第 $\ell$ 列。机制如下:当两个隐藏节点获得相似的激活($\tilde{f}_j^\top\tilde{f}_\ell$ 为正且较大)时,$B_{j\ell}$ 变为负值,产生一种有效力将它们推开。

该框架在理论上是清晰的。它没有从经验上回答的两个问题是:(i) 在训练的哪个阶段这种排斥变得可观察,以及 (ii) 它是否表现为从业者可以在线计算而无需昂贵离线诊断的量中的可测量签名?本文解决了这两个问题。

#### 两项互补测试。

在Tian的精确设置下($M=71, K=2048, \sigma(x)=x^2$, MSE, 训练比例 $p\approx 0.40, \eta=2\times 10^{-4}$ vs $\eta=0$ 作为无grokking对照),我们运行两项测试。

第一项通过在多个训练检查点进行确定性回放重建,利用 Woodbury 恒等式计算 $B$,并检查前200对最相似特征对上的符号一致性,直接验证方程 (1) 的符号规则。

第二项测试一个候选在线可观测量:参数更新 Gram 矩阵的滚动窗口特征间隙 $\sigma_2/\sigma_3$。如果第三阶段的排斥巩固了冗余特征维度,滚动 $\Delta W$ 谱应成为低秩——两个持续的更新方向对应于幸存的特征巩固,次要方向坍缩为噪声。$\sigma_2/\sigma_3$ 比率是检测这种坍缩的自然指标。

#### 贡献。

1. 1. 在 $\sigma=x^2$ 上对定理6的多种子验证。经验符号匹配(前200对相似对)在 $n=5$ 个种子中从 epoch 50 的 0.865 [IQR 0.865, 0.875] 上升到 epoch 300 的 0.985 [IQR 0.980, 0.990]。$\geq 0.95$ 的饱和在每个种子中均出现在 epoch 175。
2. 2. 定理6推广到幂激活函数之外。在 $\sigma=\text{ReLU}$ 上,相同的符号规则检查在 epoch 100 时为 0.91,在 epoch 300 时为 0.995,在 epoch 500 时为 1.000。排斥机制具有激活函数通用性。
3. 3. 参数更新谱签名是 $\sigma=x^2$ 特有的。在 $\sigma=x^2$ 上,基于斜率的检测器在 epoch 174(IQR=[173,174])在 15/15 个grok种子中触发,与 $\eta=0$ 对照的后期幅度分离达到 $229\times$。在 $\sigma=\text{ReLU}$ 上,它在 0/15 个grok种子中未触发,幅度分离坍缩至 $1.4\times$。ReLU 谱以秩1为主导而非秩2。这种分离与 Tian 定理5中专注(幂激活)与扩散(ReLU/sigmoid)记忆的区别一致。
4. 4. 方法控制。窗口大小敏感性扫描显示 $W\leq 10$ 在 $\eta=0$ 控制中产生假阳性;特异性在 $W\in\{20,30\}$ 时成立。记录 $\sigma_4,\sigma_5$ 后,秩2声称在 finest 窗口大小($W=5$)处是精确的,其中 $\sigma_3,\sigma_4,\sigma_5$ 坍缩到噪声底;在较大窗口下出现几何级联。
5. 5. 跨 $\eta$ 扩展。在 $\eta=10^{-5}$ 处的扩展单种子运行确认了 Tian 的 $1/\eta$ 缩放:grokking 在 epoch 1527,水平指标 $\rho_{\text{tian}}$(下文方程 (7))比测试准确率提前 567 个 epoch。在此慢速 $\eta$ 下,峰值时的锁定幅度 $\sigma_2/\sigma_3$ 从 $\sim 300$ 降至 $\sim 25$,这与测量结束时秩2结构未完全发展一致。

#### 论文大纲。

第2节(https://arxiv.org/html/2605.08119#S2)描述设置。第3节(https://arxiv.org/html/2605.08119#S3)展示跨激活函数和种子的定理6验证——这是主要结果。第4节(https://arxiv.org/html/2605.08119#S4)展示 $\sigma=x^2$ 上的参数更新谱签名,并记录其在 $\sigma=\text{ReLU}$ 上的失败。第5节(https://arxiv.org/html/2605.08119#S5)报告包括扩展运行在内的 $\eta$ 扫描。第6节(https://arxiv.org/html/2605.08119#S6)简要报告一种在 $\sigma=x^2$ 上有效但不推广到 $(M,p,\sigma)$ 的水平指标检测器,限定其适用范围。第7节(https://arxiv.org/html/2605.08119#S7)讨论激活通用机制与激活特定签名的区别对grokking诊断的谱方法意味着什么。

## 2 设置与仪器

### 2.1 架构与训练

我们完全复制 Tian(2025)(https://arxiv.org/html/2605.08119#bib.bib5) 图3。模型为

$$
\hat{Y}=\sigma(XW)V, \quad (2)
$$

其中 $X\in\mathbb{R}^{n\times 2M}$ 为冻结的身份嵌入(两个输入token的连接one-hot编码),无偏线性 $W\in\mathbb{R}^{2M\times K}$, $V\in\mathbb{R}^{K\times M}$,以及可配置激活函数 $\sigma$。损失函数为Tian代码中使用的零均值MSE:

$$
J(W,V)=\tfrac{1}{2}\left\|P_1^\perp(Y-\sigma(XW)V)\right\|_F^2, \quad P_1^\perp:=I_n-\tfrac{1}{n}\mathbf{1}\mathbf{1}^\top. \quad (3)
$$

训练使用学习率为 $10^{-3}$ 的 Adam 以及权重衰减 $\eta$。Tian表示法中的超参数 $\eta$ 是*权重衰减*(而非学习率);我们全程遵循此约定。

默认设置为 $M=71, K=2048, p=n_{\text{train}}/M^2\approx 0.40, \eta\in\{2\times 10^{-4},0\}$,400个epoch(ReLU为800个,因为其grokking较慢),15个种子。匹配种子的 $\eta=0$ 对照隔离了权重衰减的影响。

### 2.2 记录量

在每个epoch,我们记录:训练/测试准确率;$\tilde{F}^\top\tilde{F}$ 的非对角比率;水平指标 $\rho_{\text{tian}}$(方程 (7),第6节);$\|G_F\|$;$\Delta W$(和 $\Delta V$)滚动窗口 Gram 矩阵的前5个特征值,窗口大小 $W=20$;以及 $G_F$ 列解耦的500对独立性代理。

滚动 Gram 矩阵维护为展平参数delta的deque;在每一步,我们形成 $\Delta=[\Delta W_{t-W+1},\dots,\Delta W_t]\in\mathbb{R}^{P\times W}$ 并对 $\Delta^\top\Delta\in\mathbb{R}^{W\times W}$ 调用 `torch.linalg.eigvalsh`,这是一个每epoch可忽略的 $O(W^3)$ 操作。滚动 Gram 的前 $k$ 个特征值为 $\sigma_k(t)$;我们将 $\sigma_2/\sigma_3$ 报告为主要检测器。

### 2.3 复制 Tian 的图3

图1(https://arxiv.org/html/2605.08119#S2.F1) 在15个种子上复制了 Tian(2025)(https://arxiv.org/html/2605.08119#bib.bib5) 图3:训练准确率在 epoch 25 达到 1;在 $\eta=2\times 10^{-4}$ 时,测试准确率在中位数 epoch 102 处跨过 0.5,而在 $\eta=0$ 时从未跨过;$\|G_F\|$ 在 epoch 50 左右达到峰值;$\tilde{F}^\top\tilde{F}$ 的非对角比率始终低于 0.04(在 Tian 的 8% 界限内)。

参见图注
**图1:** 主要15种子扫描的种子间中位数(准确率和水平指标为 $\pm\text{std}$;特征间隙为 IQR)。上:测试准确率复制。中:水平指标 $\rho_{\text{tian}}$ 仅在grok条件的第二阶段上升。下:滚动 $\Delta W$ Gram 上的 $\sigma_2/\sigma_3$(对数尺度)仅在grok条件中在grokking后饱和。每个条件 $N=15$ 个种子。

## 3 跨激活函数和种子的定理6验证

### 3.1 验证协议

我们通过 Woodbury 恒等式精确计算 $B=(\tilde{F}^\top\tilde{F}+\eta I)^{-1}$,

$$
B=\tfrac{1}{\eta}I-\tfrac{1}{\eta^2}\tilde{F}^\top(\tilde{F}\tilde{F}^\top+\eta I)^{-1}\tilde{F}, \quad (4)
$$

这将 $K\times K=2048\times 2048$ 的逆矩阵简化为 $n\times n=2016\times 2016$ 的逆矩阵,在CPU上使用float64计算(PyTorch 2.5 MPS 不支持双精度线性代数)。

对于每个检查点,我们通过余弦矩阵 $S_{j\ell}=\tilde{f}_j^\top\tilde{f}_\ell/(\|\tilde{f}_j\|\|\tilde{f}_\ell\|)$ 识别前200对最相似无序特征对 $(j,\ell)$,并在这些对上评估定理6的符号规则(方程 (1))。

计算 $P_{\eta,-j\ell}$ 需要从 $\tilde{F}$ 中排除列 $j$ 和 $\ell$ 并为每对重新计算投影子。我们使用近似值 $P_{\eta,-j\ell}\approx P_\eta:=I-\tilde{F}(\tilde{F}^\top\tilde{F}+\eta I)^{-1}\tilde{F}^\top$,它使用全投影子。对 epoch 175(种子0)的10对进行近似直接验证显示,即使幅度差异很大(全投影子 $P_\eta$ 几乎消除 $\tilde{f}_\ell$,因为 $\tilde{f}_\ell$ 在 $\tilde{F}$ 的列空间中,而 $P_{\eta,-j\ell}$ 没有),残差相似度 $\tilde{f}_j^\top P_{\eta,-j\ell}\tilde{f}_\ell$ 的*符号*在 10/10 对中被近似值保留。对于定理6验证,我们只需要符号,因此近似值是合适的。

### 3.2 $\sigma=x^2$ 上的多种子结果

表1(https://arxiv.org/html/2605.08119#S3.T1) 报告了 $n=5$ 个种子在五个检查点上的经验符号匹配。可重复性紧密:每个检查点的 IQR $\leq 0.015$,且饱和 epoch(符号匹配 $\geq 0.95$)在每个种子中与 $\sigma_2/\sigma_3$ 斜率触发 epoch 一致。

**表1:** 在 $\sigma=x^2, \eta=2\times 10^{-4}$ 下,$n=5$ 个种子间前200对最相似特征对上的定理6符号匹配 $\Pr[\operatorname{sgn}(B_{j\ell})=-\operatorname{sgn}(\tilde{f}_j^\top P_\eta \tilde{f}_\ell)]$。

相似文章

神经网络可证明学习群组合的谱表示

arXiv cs.LG

本文从理论上证明,在群组合任务上训练的两层神经网络可以学习谱表示,其中神经元收敛到不可约表示并实现旋转秩一对齐,为特征学习提供了表示论的解释。

神经网络可证明地学习群组合的谱表示

Hugging Face Daily Papers

本文提供了神经网络在群组合任务中学习结构化表示的理论分析,证明了训练动态驱动神经元以指数收敛速度收敛到不可约群表示。该工作建立了特征学习的表示理论解释,并刻画了矩阵值群表示的低秩压缩现象。

语言模型中Grokking的预训练类比:追踪延迟的语法泛化

arXiv cs.LG

本文提出了一种基于暴露的框架,用于研究LLM预训练过程中类似Grokking的延迟泛化现象,使用了BLiMP最小对立对和关键短语。作者观察到五种语法现象均出现延迟泛化,并分析了内部变化,如概念向量的可预测性和注意力头的集中。