技术摘要：加速顿悟的平等梯度下降

问题陈述

本文探讨了“顿悟”（grokking）现象，即深度学习模型在训练过程中性能迅速达到接近零误差，而其测试（泛化）性能却在很长一段时间内停滞在随机水平，随后突然跃升至接近完美的准确率。尽管该现象已在多种架构和任务（如模运算、稀疏奇偶校验、形式语言）中被观察到，但导致这种延迟的根本机制仍知之甚少。现有文献表明，这与病态损失景观、奇异学习理论以及从“懒惰”核机制到特征学习机制的过渡有关，但缺乏一种在不损害最终泛化能力的前提下控制或加速这一过渡的原则性方法。

方法论：平等梯度下降（EGD）

作者提出了平等梯度下降（Egalitarian Gradient Descent, EGD），这是对标准（随机）梯度下降（SGD）的一种改进，旨在均衡梯度所有主方向上的优化速度。

核心机制

该方法通过对每一层的梯度矩阵 $G$ 进行变换来运作。设 $G$ 为维度为 $m \times p$ 的层的梯度矩阵。作者定义变换后的梯度 $\tilde{G}$ 为：
$$\tilde{G} = F^{-1/2}G = (GG^\top)^{-1/2}G$$
其中 $F = GG^\top$ 是经验 Fisher 信息矩阵（FIM）。

谱解释：
利用奇异值分解（SVD）$G = USV^\top$，该变换产生：
$$\tilde{G} = US^{-1}U^\top USV^\top = US^{-1}SV^\top = UV^\top$$
该操作保留了梯度的左（$U$）和右（$V$）奇异向量（即主方向），但强制所有奇异值等于 1。因此，参数的动力学在所有主方向上以完全相同的速度演化，消除了病态景观中特征性的“快”模式与“慢”模式之间的差异。

实际实现

• 计算： 主要成本是对梯度矩阵执行 SVD。论文指出，精确 SVD 计算昂贵，但建议使用随机化 SVD（RSVD），甚至简化的列归一化（按 $L_2$ 范数重新缩放列）可作为有效的近似。
• 切换策略： 为了减轻计算开销，作者提出了一种“预热”策略，即仅应用 EGD 直到达到验证准确率阈值，之后优化器恢复为普通 SGD。
• 与其他方法的关系：
  • 自然梯度下降（NGD）： EGD 是 NGD 的“白化”版本（$\tilde{G} = F^{1/2} \cdot \text{NGD\_update}$）。与使用 $F^{-1}$ 的 NGD 不同，EGD 使用 $F^{-1/2}$。
  • Grokfast： EGD 与 Grokfast（通过滤波放大低频梯度分量）共享归纳偏置，即降低高频方向的权重。然而，EGD 是通过谱归一化而非时间滤波来实现这一点的，因此不需要过去梯度的记忆缓冲区。
  • Muon： EGD 在几何上与 Muon 优化器（对更新进行正交化）相似，但它是基于顿悟的谱理论推导而来，而非基于经验设计选择。

理论洞察

本文利用简化的玩具模型（各向异性高斯数据上的线性回归）提供了理论分析以解释顿悟：

1. 病态作为原因： 在普通 SGD 中，经验协方差矩阵的条件数（由特征方差比率驱动）决定了收敛速度。如果条件数很高（例如 $1/\epsilon$），优化器会快速收敛到一个记忆训练数据的解（快方向），但需要 $O(1/\epsilon)$ 的时间才能与真实决策边界对齐（慢方向），从而导致平台期。
2. 通过归一化加速： 玩具模型中的修改更新规则 (9) 有效地消除了对条件数 $\epsilon$ 的依赖。通过均衡更新矩阵的谱半径，该方法确保在独立于病态条件的步数内收敛到最优解，从理论上保证了消除平台期。

主要贡献

1. 平等梯度下降（EGD）： 一种简单、无超参数的方法，通过对梯度进行归一化以确保沿主方向的各向同性演化，从而加速顿悟。
2. 理论框架： 形式化地证明了在简化设置中，顿悟平台期是由病态梯度引起的，并且 EGD 可证明消除了与这些条件相关的延迟。
3. 实证验证： 大量实验展示了 EGD 在以下方面的有效性：
   • 玩具问题： 各向异性数据的二分类。
   • 算法任务： 模加法、模乘法和稀疏奇偶校验问题。
   • 现实数据集： 使用 MLP、CNN 和 Transformer 的 MNIST 和 CIFAR-10。
4. 效率与鲁棒性： 证明了 EGD（尤其是结合 RSVD 时）相比基线减少了挂钟时间，未引入额外的内存开销（不同于 Grokfast），并提高了在非平稳环境（如标签洗牌）中的适应性。

实验结果

• 加速： 在模运算和稀疏奇偶校验任务中，EGD 仅在几个 epoch 后便达到高测试准确率，而普通 SGD 则表现出持续数千个 epoch 的平台期。
• 与基线比较：
  • 对比普通 SGD： EGD 一致地消除了平台期。
  • 对比 Grokfast： EGD 实现了相当或更优的加速。关键在于，使用 RSVD 的 EGD 变体通常在挂钟时间上优于 Grokfast，且不需要 Grokfast 所必需的过去梯度记忆缓冲区。
  • 对比列归一化： EGD 的简化版本（列归一化）也能显著加速顿悟，尽管带有 RSVD 的完整 EGD 通常在速度和准确率之间提供最佳平衡。
• 非平稳性： 在涉及 CIFAR-10 标签洗牌的实验中，将 EGD 添加到优化器（Adam、RAdam、RMSprop）中显著提高了它们从分布偏移中恢复的能力，增加了恢复曲线下的面积。

意义与主张

本文声称，EGD 提供了一种原则性的、谱感知的方法来解决顿悟延迟问题。通过将顿悟框架化为病态梯度动力学的后果，作者提供了一种方法：

• 加速学习： 大幅减少模型从记忆过渡到泛化所需的时间。
• 保持泛化： 不损害最终的泛化水平；在某些情况下，它完全消除了平台期。
• 实用性强： 无需额外的超参数，与优化器无关，并且可以使用随机近似以低计算开销实现。
• 统一视角： 将顿悟现象与经典优化概念（条件数、自然梯度）和最近的启发式方法（Grokfast、Muon）联系起来，表明谱归一化是加速深度网络学习的基本原语。

作者总结道，虽然目前的理论分析仅限于简化模型，但在不同架构和任务上的实证成功表明，梯度谱归一化是一种稳健且有效的策略，可用于缓解泛化延迟。未来的工作提议将理论分析扩展到 Transformer 等复杂模型，并探索与自适应优化器的即插即用组合。