热力学权重衰减：通过注意力比热探讨顿悟加速

# 探索通过注意力比热加速Grokking
来源：https://arxiv.org/html/2607.20552
## 热力学权重衰减：探索通过注意力比热加速Grokking

Chitraansh Pandey 独立研究员 chitraanshpandey@gmail\.com github\.com/baymaxbyte/cbo\_core

###### 摘要

Grokking——神经网络在完全记忆训练数据后很久才出现的延迟泛化——浪费了数千个训练轮次，且其发生时间 notoriously 不可预测。基于最近的研究成果，即Transformer注意力在形式上与热力学系统同构，我们将注意力对数（logits）的方差视为*比热* \(C_v\)，并表明其峰值可靠地出现在泛化转变之前。我们引入 CvAdamW，这是一种即插即用的 AdamW 变体，它在线监控 \(C_v\) 并在检测到相变时通过动态缩放权重衰减来注入“热能”。通过严格的迭代开发过程，我们识别出三种失败模式——初始化噪声、小批量微波动和“弹弓致盲（slingshot blinding）”——并通过记忆门和指数移动平均减震器加以解决。在模算术任务 \((a+b \bmod 97)\) 上，CvAdamW 在 4000 轮预算内于第 2802 轮实现 grokking，而基线从未 grok。我们进一步提出一个*尺度不变*的 z 分数重构，消除了任务特定的超参数，并在 10 个配对种子上进行了评估。配对分析显示，冷启动变体将平均 grokking 延迟减少了 257 轮（6.0%；中位数 166 轮；Wilcoxon p=0.049，Cohen's d=0.68，bootstrap 95% CI [53, 489]），改进了 10 个种子中的 8 个；在此单一任务上，所有 10 个种子中 \(C_v\) 的峰值均出现在 grokking 之前。我们的结果表明，神经网络可能暴露出可检测的即将发生的泛化转变前兆信号，并且一种基于物理原理的比例干预可以在固定计算预算内促进泛化。代码和数据已公开。

## 1 引言

现代过参数化网络经常表现出*grokking* (Power et al., 2022 (https://arxiv.org/html/2607.20552#bib.bib2))：训练准确率在几十轮内饱和，而验证准确率在数百或数千轮内停留在随机水平，然后突然跃升至近乎完美。中间的平稳期浪费计算资源，更糟糕的是，其持续时间难以提前预测。实践者很难判断模型即将泛化还是永久停滞。

最近的理论为这一现象提供了引人注目的重新解释。Kim (2026) (https://arxiv.org/html/2607.20552#bib.bib1) 证明注意力机制在形式上与经典热力学系综同构：softmax 是最小化亥姆霍兹自由能的玻尔兹曼分布，缩放因子 \(1/\sqrt{d_k}\) 是逆温度，注意力对数 \(QK^\top\) 是能级，并且——关键的是——这些对数的*方差*行为类似于*比热*。在统计力学中，比热在相变处发散。如果 grokking 是一种相变，那么网络应该通过这个可观测量的尖峰来宣告它。

我们直接采纳这一预测，并提出一个控制理论问题：*如果网络发出自身相边界的信号，我们能否在线检测该信号，并精确提供跨越该边界所需的能量？* 物理执行器是权重衰减。由于注意力系统的有效温度缩放为 \(T_{\mathrm{eff}} \propto \sqrt{d_k} / \lVert W \rVert^2\)，增加权重衰减会缩小参数范数，从而“加热”系统。我们的贡献是一个闭合此回路的优化器。

#### 贡献.

- • 我们通过实验证明，在模算术任务上，注意力比热 \(C_v = \mathrm{Var}(QK^\top / \sqrt{d_k})\) 在我们测试的每个种子中均出现在 grokking 之前（第 6.1 节 (https://arxiv.org/html/2607.20552#S6.SS1)）。
- • 我们引入 CvAdamW，一种热力学感知优化器，它根据 \(C_v\) 的平滑动量按比例缩放权重衰减（第 3 节 (https://arxiv.org/html/2607.20552#S3)），并记录三种失败模式及其修复。
- • 我们提出一个*尺度不变*的 z 分数公式，消除了任务特定阈值（第 4 节 (https://arxiv.org/html/2607.20552#S4)）。
- • 我们提供了基于 10 个种子的配对统计评估——每个种子的结果、非参数检验、效应量、bootstrap 置信区间、提前时间统计以及前兆相关分析——而非点估计（第 6.3 节 (https://arxiv.org/html/2607.20552#S6.SS3)）。

## 2 背景

### 2.1 注意力的热力学同构

对于单个查询，注意力通过对 \(n\) 个键计算权重，公式为：

\[
p_i = \frac{\exp(z_i / \sqrt{d_k})}{\sum_{j=1}^{n} \exp(z_j / \sqrt{d_k})}, \qquad z_i = (QK^\top)_i. \tag{1}
\]

这恰好是玻尔兹曼分布 \(p_i = e^{-E_i/T} / Z\)，其中能级 \(E_i = -z_i\)，逆温度 \(\beta = 1/\sqrt{d_k}\)，配分函数 \(Z = \sum_j e^{-E_j/T}\) (Kim, 2026 (https://arxiv.org/html/2607.20552#bib.bib1))。在此映射下，正则系综的比热 \(C_v = \mathrm{Var}(E) / T^2\) 对应于缩放注意力对数的方差。遵循 Kim (2026) (https://arxiv.org/html/2607.20552#bib.bib1)，我们将注意力“信息气体”的比热定义为：

\[
C_v = \mathrm{Var}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right). \tag{2}
\]

从热力学角度看，\(C_v\) 衡量改变系统温度所需的能量，并在相变处发散。我们的工作假设是 grokking 正是这样一种相变，因此当模型将其内部表示从记忆重新组织为泛化时，\(C_v\) 应该出现尖峰。

### 2.2 Grokking 与权重衰减

Grokking 最初由 Power 等人 (2022) (https://arxiv.org/html/2607.20552#bib.bib2) 在算法任务上描述，此后与权重范数和正则化相关联 (Liu 等人，2023 (https://arxiv.org/html/2607.20552#bib.bib6); Nanda 等人，2023 (https://arxiv.org/html/2607.20552#bib.bib5))。一种常见观点是，记忆对应一个尖锐的高曲率最小值，而泛化存在于更平坦的盆中；逃离前者需要对参数范数施加隐式或显式的压力。权重衰减 (Loshchilov and Hutter, 2019 (https://arxiv.org/html/2607.20552#bib.bib3)) 是标准工具。热力学图景使这一点精确化：

\[
T_{\mathrm{eff}} \propto \frac{\sqrt{d_k}}{\lVert W \rVert^2}, \tag{3}
\]

增加权重衰减会减小 \(\lVert W \rVert\)，从而提高 \(T_{\mathrm{eff}}\)。更多的权重衰减，字面上就是更多的热量。

## 3 方法：CvAdamW

CvAdamW 是 AdamW 的即插即用替代品，它每轮根据 \(C_v\) 轨迹计算一个*动态*权重衰减 \(\lambda_t\)。基础优化器保持不变；仅调节正则化系数。

### 3.1 比例热量注入（\(\kappa/\tau\) 公式）

令 \(C_v(t)\) 为第 \(t\) 轮测量的比热。我们使用指数移动平均（EMA）跟踪其速度的平滑动量：

\[
\begin{aligned}
\delta_t &= C_v(t) - C_v(t-1), \tag{4} \\
\mu_t &= \alpha \mu_{t-1} + (1-\alpha) \delta_t, \tag{5} \\
\lambda_t &= \lambda_{\mathrm{base}} + \kappa \cdot \max\!\left(0,\; \mu_t - \tau\right). \tag{6}
\end{aligned}
\]

这里 \(\alpha=0.9\) 给出了一个有效窗口 \(1/(1-\alpha)=10\) 轮，\(\tau\) 是一个噪声基底，低于它的波动被忽略，而 \(\kappa\) 将持续的正动量放大为权重衰减单位。当 \(C_v\) 平坦或下降时，\(\mu_t \leq \tau\)，优化器恢复为标准 AdamW，\(\lambda_{\mathrm{base}}=0.1\)。当 \(C_v\) 向转变攀升时，\(\mu_t\) 增长，\(\lambda_t\) 按比例缩放。

#### 记忆门。

上述所有机制仅在模型完成记忆后才激活，由 \(\text{train\_acc} \geq 0.99\) 强制执行。相变只有在系统达到亚稳态（记忆状态）后才有意义。

### 3.2 三种失败模式

CvAdamW 并非一蹴而就。每个设计元素都解决了一个特定的、观察到的失败（完整日志见随附仓库）。

(1) **初始化噪声**。没有任何门控时，所有检测器在第 16 轮触发，将随机初始化的混乱注意力模式误判为相变。此时过早注入热量的效果*比基线更差*：你不能强迫一个系统离开它尚未进入的最小值。记忆门解决了这个问题。

(2) **小批量微波动**。有门控但无平滑时，一个离散的运动学触发器在第 240 轮因一个单轮向上的波动（由批次采样噪声引起）而触发。一次性干预被浪费了。\(\alpha=0.9\) 的EMA吸收了 1-2 轮的噪声，同时传递了真实转变持续 50-200 轮的上升（图 1 (https://arxiv.org/html/2607.20552#S3.F1)）。

(3) **弹弓致盲**。在重大结构重组期间，竞争的电路会暂时降低训练准确率——即“弹弓”效应 (Thilak 等人，2022 (https://arxiv.org/html/2607.20552#bib.bib8))。在一次运行中，一个大的 \(C_v\) 尖峰恰好与 train_acc 降至 0.99 以下同时发生，在峰值处关闭了门，致盲了离散触发器。一个连续的比例优化器规避了这一点：它在门关闭之前已经积累了动量，因此热能是在“山脉”形成时输送的，而非在其经过之后。

这些失败促使了连续性这一核心设计原则。一个跟踪信号幅度的比例响应，对于一次性触发器可能致命的时序误差具有鲁棒性。

参见说明文字
图 1：\((a+b) \bmod 97\) 上的经典 grokking。训练准确率在几十轮内饱和，而验证准确率在数千轮内停留在随机水平（阴影区域“浪费的计算”）。在约束的 4000 轮预算下，基线从未 grok。

## 4 尺度不变重构

\(\kappa/\tau\) 公式有两个任务特定常数：\(\tau\) 是绝对幅度，\(\kappa\) 是维度转换因子。在 \(C_v\) 具有不同范围的任务上，两者都需要重新调整。我们通过将速度 \(v_t = C_v(t) - C_v(t-1)\) 视为随机变量并检测统计异常来消除它们。使用其运行均值和方差的 EMA 估计：

\[
\begin{aligned}
\mu_t &= \beta_z \mu_{t-1} + (1-\beta_z) v_t, \tag{7} \\
\sigma_t^2 &= \beta_z \sigma_{t-1}^2 + (1-\beta_z) (v_t - \mu_{t-1})(v_t - \mu_t), \tag{8} \\
Z_t &= \frac{v_t - \mu_t}{\sqrt{\sigma_t^2} + \epsilon}, \tag{9} \\
\lambda_t &= \lambda_{\mathrm{base}} + \max\!\left(0,\; Z_t - z_{\mathrm{thresh}}\right). \tag{10}
\end{aligned}
\]

唯一的自由常数 \(z_{\mathrm{thresh}}=2.0\) 是一个通用的 \(2\sigma\) 异常阈值。

#### 冷启动 vs. 连续传感器。

记忆门创建了一个关于何时开始积累统计量的选择。*冷启动*变体仅在门打开后（\(\text{train\_acc} \geq 0.99\)）才启动 \(\mu, \sigma^2\)，从而最大化对门开后第一个信号的灵敏度。*连续传感器*变体从第 1 轮开始跟踪统计量（仅对执行器加门），提供一个温热的基线。正如我们将展示的，冷启动变体反而更强：温热的基线稀释了真实转变的相对异常性。

## 5 实验设置

所有实验使用一个 2 层仅解码器 Transformer（\(d_{\mathrm{model}}=128\)，4 个头，\(d_k=32\)，RoPE 位置编码 (Su et al., 2021 (https://arxiv.org/html/2607.20552#bib.bib7))，GELU）在模加法 \((a+b) \bmod 97\) 上训练，数据集包含 \(97^2=9409\) 个样本，采用 \(50/50\) 的训练/验证分割。基础优化器为 AdamW (Loshchilov and Hutter, 2019 (https://arxiv.org/html/2607.20552#bib.bib3))（\(\text{lr}=3 \times 10^{-4}\)，\(\lambda_{\mathrm{base}}=0.1\)，批次大小 512）。我们报告*grokking 轮次*，定义为验证准确率首次超过 0.95 的轮次。尺度不变研究使用 10 个种子 \(\{42,123,256,512,1024,2048,3141,4096,7777,9999\}\)，运行 7000 轮；\(\kappa/\tau\) 研究使用 5 个种子，运行 10,000 轮。

## 6 结果

### 6.1 观察到的 Cv 前兆动力学

在我们对该任务运行的所有配置中，\(C_v\) 都表现出一个显著的峰值，出现在验证准确率转变之前（在第 6.3 节 (https://arxiv.org/html/2607.20552#S6.SS3) 中量化）。图 2 (https://arxiv.org/html/2607.20552#S6.F2) 对比了在 4000 轮预算下的三种情况。基线 (A) 显示出清晰的 \(C_v\) 峰值，但验证准确率仍然被困在随机水平：信号存在，但系统缺乏跨越的能量。阶跃函数干预 (B) 在检测到的峰值处（第 3049 轮）注入一个固定的权重衰减尖峰（\(0.1 \to 1.0\)），并在 284 轮后 grok。CvAdamW (C) 平滑地将权重衰减缩放至峰值 1.61，并且在第 2802 轮最早 grok。在此预算下，基线从未 grok，因此 CvAdamW 的贡献不仅是加速，更是*使泛化成为可能*。

参见说明文字
图 2：4000 轮预算下的主对比（单种子）。(A) 基线 AdamW：\(C_v\)（红色）达到峰值，但验证准确率（蓝色）停留在随机水平。(B) 阶跃函数：在检测到的峰值处施加二元权重衰减尖峰，迫使 grokking。(C) CvAdamW：连续、比例的权重衰减缩放，最早 grok。参见说明文字
图 3：权重衰减调度及相应的验证准确率。阶跃函数（红色）应用二元尖峰并带有固定冷却期；CvAdamW（绿色）应用平滑的比例响应，跟踪 \(C_v\) 动量。连续调度更早地泛化。

### 6.2 连续 vs. 离散干预

图 3 (https://arxiv.org/html/2607.20552#S6.F3) 叠加了两种权重衰减调度。阶跃函数是脆弱的：其单次尖峰必须精确落在峰值处，并且会被弹弓致盲击败。连续调度在整个转变过程中提供能量，并且对单轮门关闭具有免疫力，这与相变是扩展的而非瞬时的物理直觉一致。

### 6.3 尺度不变配对研究

我们在 10 个配对种子上评估尺度不变的冷启动变体与基线。由于每个种子在两种条件下运行，我们使用*配对*分析。令 \(d_i = \text{Baseline}_i - \text{ColdStart}_i\)；正值有利于我们的方法。表 1 (https://arxiv.org/html/2607.20552#S6.T1) 给出了完整的每个种子结果（包括连续传感器变体），以便读者检查方差、异常值和稳健性。