Grokking Transformer中的权重衰减机制：廉价在线诊断

# Grokking Transformer中的权重衰减阶段：廉价在线诊断

来源：https://arxiv.org/html/2605.20441

###### 摘要

在模算术上训练的Grokking Transformer在记忆、泛化和崩溃阶段之间表现出急剧的转变。我们表明，权重衰减作为这些阶段的标量经验控制参数，并引入两种廉价的在线诊断方法——平均成对注意力头余弦相似度和熵标准差——它们仅从注意力激活中跟踪这些仅解码器模算术设置中的底层训练动态，并以更低的计算成本补充损失景观诊断。在11个实验条件和三种模型规模（0.82M到85M参数）下，权重衰减轴将记忆区（λ<λc，接近零的grokking）、发展性grokking（λ≥λc，在λ∼0.1时达到∼100%的grok率，时间到grok从1090减少到83个epoch，λ∈[0.1,2.0]）和崩溃区（λ=10，注意力模式相同）分开。一个接近转变的Logistic拟合将记忆-发展边界定位在λc=0.0158（95% CI [0.0109,0.0200]，N=210）；时间到grok的幂律拟合给出经验指数ν=0.757（CI [0.725,0.799]）。测试的参考指数ν=1/2和3D Ising ν≈0.63在我们的四箱网格下位于经验CI之外；我们将ν报告为经验值，并将普适类识别推迟到更密集的有限尺寸缩放数据塌陷（Bi et al., 2026 (https://arxiv.org/html/2605.20441#bib.bib4)）。一个水平匹配的多任务复制实验（n=280，四个模运算，两种规模）保留了加法之外的WD控制模式；一个配对注意力头重新初始化实验在规范的后转变点λ=0.05处改变了Phase-2振幅（Cohen’s d=−1.190，n=10配对，pt=4.5×10^{-3}），而匹配的权重范数剪枝则没有，从而将效果隔离到头模式结构而非权重幅度上。每头维度d/H以饱和单调的方式调节分化振幅。三个跨架构范围探测（4层MLP h=512，4层LSTM h=512，和4层Mamba d=128；每个规范n=70）复制了非注意力架构中WD控制的转变，每个架构的λc大约跨越一个数量级（MLP λc=0.0511 [0.0495,0.0591]；LSTM λc=0.0365 [0.0299,0.0473]；Mamba λc=0.0144 [0.0106,0.0159]，其CI与Transformer规范CI重叠而非高于它）。声明范围限定于小型Transformer注意力模型中的模算术；跨架构、语言模型和普适类的声明不在范围内。

## 1 引言

当在模算术上训练Transformer时，它们表现出*grokking*：在训练准确率饱和很久之后从记忆到泛化的突然转变（Power et al., 2022 (https://arxiv.org/html/2605.20441#bib.bib35)）。机制分析表明，这种转变对应于内部电路（基于傅里叶的表示）的形成，这些电路逐渐出现但突然显现（Nanda et al., 2023 (https://arxiv.org/html/2605.20441#bib.bib29)）。类似地，Transformer训练期间归纳头的形成构成了从两层模型到70B+参数系统的“相变”（Olsson et al., 2022 (https://arxiv.org/html/2605.20441#bib.bib32)）。跨越个体训练现象（grokking（Power et al., 2022 (https://arxiv.org/html/2605.20441#bib.bib35)；Nanda et al., 2023 (https://arxiv.org/html/2605.20441#bib.bib29)；Kumar et al., 2024 (https://arxiv.org/html/2605.20441#bib.bib18)）、涌现能力（Wei et al., 2022 (https://arxiv.org/html/2605.20441#bib.bib56)）、彩票假说（Frankle & Carbin, 2019 (https://arxiv.org/html/2605.20441#bib.bib10)）、神经崩溃（Papyan et al., 2020 (https://arxiv.org/html/2605.20441#bib.bib33)）），最近的学习力学视角将分阶段学习、经验定律、极限模型和廉价可测量诊断作为未来深度学习理论的核心对象（Simon et al., 2026 (https://arxiv.org/html/2605.20441#bib.bib42)）。统计力学界已将训练与相变联系起来（Bahri et al., 2020 (https://arxiv.org/html/2605.20441#bib.bib3)；Saxe et al., 2014 (https://arxiv.org/html/2605.20441#bib.bib39)；Ziyin & Ueda, 2023 (https://arxiv.org/html/2605.20441#bib.bib67)；Žunkovič & Ilievski, 2024 (https://arxiv.org/html/2605.20441#bib.bib68)），最近的工作将振荡器式同步模型应用于神经激活（Miyato et al., 2025 (https://arxiv.org/html/2605.20441#bib.bib26)）。然而，现有的grokking研究很少提供廉价的在线激活空间诊断和权重衰减控制的注意力动态的阶段图。我们做出四项贡献：

1. **具有bootstrap CI和形式化良构性检查的量化权重衰减临界阈值**。一次密集的权重衰减扫描加上一次稀疏的三尺寸规模探测得到了一种水平匹配的Logistic转变估计，规范mod+密集WD队列的λc=0.0158（95% CI [0.0109,0.0200]，N=210），以及λc以上时间到grok的幂律指数ν=0.757（CI [0.725,0.799]）。双轴（λ,N）图展示了三个性质不同的阶段（记忆、发展、崩溃），同时具有“太少”和“太多”两种失败模式，补充了同样将权重衰减视为相变相关变量但没有确定数值阈值的有限尺寸缩放和权重几何工作。四个诊断良构性恒等式（A1, B1, C1, E1）已在Lean 4中经过机器检查；它们证明了诊断的界和代数恒等式，而非经验阶段声明。

2. **注意力头协调的廉价在线序参量**。我们定义了两个可在每个训练步骤计算的标量量：平均成对余弦相似度s̄(t)和熵标准差σH(t)。它们通过训练跟踪注意力头协调，补充了精炼的局部学习系数（Wang et al., 2024 (https://arxiv.org/html/2605.20441#bib.bib51)），同时仅需要前向传递的注意力权重，并且在相同的检查点上进行头对头基准测试。

3. **两个顺序阶段加上种子依赖的后期保持失败**。沿着规范训练轨迹，我们记录了：阶段1（注意力头协调，接近grokking）其中头收敛且s̄从0.93上升到0.995；阶段2（分化，grokking后）其中头发散，s̄下降到0.88，而准确率保持不变。在20000个训练epoch时，规范种子42轨迹扩展为一个五阶段模式，表现出类似于反grokking的后期准确率崩溃；20000 epoch长视野保持队列（E8，n=20）显示这是种子依赖的脆弱性，而非普适循环（在λ∈{0.1,0.5,1.0,2.0}时保持率分别为5/5, 4/5, 3/5, 4/5）。崩溃轨迹在性质上类似于Prakash & Martin (2026a (https://arxiv.org/html/2605.20441#bib.bib36))报告的后期阶段循环，尽管我们的谱证据（§4.5 (https://arxiv.org/html/2605.20441#S4.SS5)）在时间上与它们不同。

4. **每头维度作为振幅调制器**。在固定模型维度下，变化头的数量将每头维度d/H隔离为分化振幅的经验主导变量：峰值σH随着d/H以饱和单调方式增加（在d/H∈{2,4,8,16}上单调，在d/H≥16时平台化，其中d/H=16和d/H=32的箱均值有重叠的95%箱CI），在d/H≈2时达到与随机标签空控制相同的数量级，同时在统计上仍然可区分（置换检验p=0.009，Cohen’s d=1.11）。因此，我们在此设置中将其标记为经验架构阈值，并将因果机制声明推迟到后续工作。

## 2 相关工作

#### Grokking和训练相变。
Power et al. (2022 (https://arxiv.org/html/2605.20441#bib.bib35))发现小型Transformer在算法任务上在记忆之后很久泛化，这一现象被Nanda et al. (2023 (https://arxiv.org/html/2605.20441#bib.bib29))从机制上解释为电路形成后跟清理，并被Varma et al. (2023 (https://arxiv.org/html/2605.20441#bib.bib49))理论上解释为记忆电路和泛化电路之间受权重衰减驱动的竞争。Kumar et al. (2024 (https://arxiv.org/html/2605.20441#bib.bib18))将grokking框架化为从懒惰到丰富阶段的转变，而Liu et al. (2023 (https://arxiv.org/html/2605.20441#bib.bib19))将grokking扩展到算法数据之外。2025年和2026年的工作已汇聚到将grokking框架化为量化相变：Bi et al. (2026 (https://arxiv.org/html/2605.20441#bib.bib4))应用有限尺寸缩放与Binder累积量交叉以及谱头尾对比作为序参量；Wang (2026a (https://arxiv.org/html/2605.20441#bib.bib52))和Wang (2026b (https://arxiv.org/html/2605.20441#bib.bib53))通过级联维度指数分析了有效维度和自组织临界性；Acharya & Dhakal (2026 (https://arxiv.org/html/2605.20441#bib.bib1))将grokking与方差限制的谱门控联系起来；Truong Xuan Khanh et al. (2026b (https://arxiv.org/html/2605.20441#bib.bib48))提出了表示协方差的归一化谱熵作为标量阈值（在泛化前穿过∼0.61）；Hennick & Corlouer (2026 (https://arxiv.org/html/2605.20441#bib.bib16))研究了约化密度矩阵谱作为早期预警；Golwala (2026 (https://arxiv.org/html/2605.20441#bib.bib12))使用保留的表示质心几何进行早期检测；Tian (2025 (https://arxiv.org/html/2605.20441#bib.bib46))在两层网络中提供了可证明的三阶段缩放定律。Xu et al. (2026 (https://arxiv.org/html/2605.20441#bib.bib57))证明了权重衰减下岭回归中的延迟泛化，而Zhang et al. (2026 (https://arxiv.org/html/2605.20441#bib.bib65))给出了grokking的SLT/算法复杂度抽象视图；Song & Ye (2026 (https://arxiv.org/html/2605.20441#bib.bib44))将模算术上的grokking与作为参数计数的函数的竞争记忆和泛化时间尺度联系起来，补充了我们在固定协议下的经验WD×N图。这些结果强化了廉价序参量和权重衰减控制阶段的价值，但没有提供本文研究的WD×N注意力头阶段图。Lyu et al. (2024 (https://arxiv.org/html/2605.20441#bib.bib21))建立了规范的早期/晚期隐式偏差二分法；Musat (2025 (https://arxiv.org/html/2605.20441#bib.bib27))将grokking刻画为由权重衰减驱动的零损失流形上的范数最小化；Manir & Rupa (2026 (https://arxiv.org/html/2605.20441#bib.bib22))经验发现grokking主要由正则化和优化决定，而非架构。Prakash & Martin (2026a (https://arxiv.org/html/2605.20441#bib.bib36))报告了“先前未报告的第三阶段”（后期泛化崩溃），通过谱密度重尾性诊断，并在Prakash & Martin (2026b (https://arxiv.org/html/2605.20441#bib.bib37))中通过RMT/相关陷阱框架将反grokking处理为一个长视野过拟合阶段。我们使用注意力相似性诊断观察到一个性质上相似的后期崩溃轨迹，并刻画了其权重衰减依赖性，尽管我们的谱证据（§4.5 (https://arxiv.org/html/2605.20441#S4.SS5)）与它们不同，因为重尾结构在grokking发作期间形成，而非后期循环期间。

#### Grokking几何的并行工作。
Xu (2026d (https://arxiv.org/html/2605.20441#bib.bib61))研究了模算术上的多任务grokking，并将权重衰减识别为一个*阶段参数*，控制grokking时间尺度和曲率深度，报告了两个定性阶段（λ≥0.5快 vs. λ≤0.3慢），并在Xu (2026c (https://arxiv.org/html/2605.20441#bib.bib60); a (https://arxiv.org/html/2605.20441#bib.bib58); e (https://arxiv.org/html/2605.20441#bib.bib62); f (https://arxiv.org/html/2605.20441#bib.bib63); b (https://arxiv.org/html/2605.20441#bib.bib59))中有补充结果。我们的工作在三个方面不同。首先，Xu (2026d (https://arxiv.org/html/2605.20441#bib.bib61))在整个过程中固定模型规模N；我们联合扫描λ和N，揭示了一个水平匹配的转变估计λc，在测试的小/中规模对上具有重叠的95% CI（从我们密集WD扫描队列的N=210次运行（阶段A后）拟合得到λc=0.0158，95% CI [0.0109,0.0200]），并且还有一个第三*崩溃*阶段（λ>5）在Xu的图中不存在。其次，Xu的诊断在权重和更新空间操作（PCA轨迹方差、对易子范数||[WQ,WK]||F、谱边缘梯度/衰减分解、功能模式谱）并且需要完整的检查点访问；我们的序参量（注意力激活的平均成对余弦相似度和跨头熵标准差）是激活空间诊断，可在每评估步O(H²)内在线计算。第三，Xu没有涉及同步或置换对称性约简动态；阶段1的同步和阶段2的头专门化框架、五阶段反grokking轨迹以及d/H振幅调制在他们的工作中没有出现。Yıldırım (2026 (https://arxiv.org/html/2605.20441#bib.bib64))提出了一个反框架，其中grokking可通过统一注意力架构消融（无权重衰减）被绕过；我们并不声称权重衰减是泛化的*必要条件*，仅声称在标准注意力架构内，权重衰减表现为一个具有良好定义的阶段图的经验控制参数。Tang et al. (2026 (https://arxiv.org/html/2605.20441#bib.bib45))报告了模算术上grokking发作时H1持久同调特征的急剧上升，这是一种离线的事后几何诊断，补充了我们的在线注意力协调序参量；这两种诊断类针对相同的阶段转变，但信号位置不同（表示拓扑 vs. 头模式协调）且评估成本不同。Wang et al. (2026 (https://arxiv.org/html/2605.20441#bib.bib55))同样通过分布谱坐标（Wasserstein/分位数、Hankel DMD残差、有效秩）在模加法Transformer轨迹上定位grokking转变；这是另一种在权重/激活谱上的转变定位诊断，而非注意力头激活序参量。Ali (2026 (https://arxiv.org/ht...