[论文解读] Omnigrok: Grokking Beyond Algorithmic Data

Ziming Liu, Eric J. Michaud+1 more|arXiv (Cornell University)|Oct 3, 2022

Machine Learning and Data Classification被引用 8

一句话总结

本文提出 'LU机制'——即在权重范数上，L形的训练损失与U形的测试损失之间的不匹配——作为神经网络中涌现'领悟'（grokking）现象的根源。通过分析简化的1D损失景观，作者表明，这种损失景观的不匹配导致了缓慢的优化动力学，从而引发'领悟'现象，并在图像、语言和分子等多种任务中展示了该现象的存在，尤其在以表征学习为主的任务（如算法数据集）中效应更显著。

ABSTRACT

Grokking, the unusual phenomenon for algorithmic datasets where generalization happens long after overfitting the training data, has remained elusive. We aim to understand grokking by analyzing the loss landscapes of neural networks, identifying the mismatch between training and test losses as the cause for grokking. We refer to this as the "LU mechanism" because training and test losses (against model weight norm) typically resemble "L" and "U", respectively. This simple mechanism can nicely explain many aspects of grokking: data size dependence, weight decay dependence, the emergence of representations, etc. Guided by the intuitive picture, we are able to induce grokking on tasks involving images, language and molecules. In the reverse direction, we are able to eliminate grokking for algorithmic datasets. We attribute the dramatic nature of grokking for algorithmic datasets to representation learning.

研究动机与目标

• 理解'领悟'现象的起源，即在神经网络过拟合后出现的延迟泛化现象。
• 探究'领悟'是否不仅存在于算法数据集，也出现在视觉、自然语言处理和分子性质预测等任务中。
• 解释为何在算法数据集中'领悟'现象表现得极为显著，而在真实世界数据集（如MNIST）中信号则较微弱。
• 通过损失景观分析，阐明表征学习在增强'领悟'效应中的作用。
• 证明在训练过程中对权重范数施加约束可抑制'领悟'，从而验证该机制的合理性。

提出的方法

• 通过在固定权重范数下对角度方向最小化训练损失，定义一个简化的损失景观，得到关于权重范数 w 的1维函数。
• 可视化并分析由此产生的训练损失（L形）与测试损失（U形）随权重范数的变化，识别出LU机制。
• 使用一个简化的教师-学生模型，模拟并验证由LU机制引发的慢-快动力学行为。
• 将LU机制的直觉应用于真实数据集（MNIST、IMDb、QM9），通过修改训练设置来诱导'领悟'信号。
• 比较算法数据集与真实世界数据集的损失景观，以解释'领悟'严重程度的差异。
• 在训练过程中对权重范数施加约束，实证性地消除'领悟'现象，验证该机制的预测能力。

实验结果

研究问题

• RQ1 是什么导致了在过拟合后出现的延迟泛化现象，特别是在'领悟'现象中？
• RQ2 是否可以在非算法数据集（如图像分类、情感分析和分子性质预测）中诱导出'领悟'？
• RQ3 为何'领悟'在算法数据集中表现得更为剧烈，而在真实世界数据集（如MNIST）中则较弱？
• RQ4 表征学习如何影响'领悟'信号的强度和可观测性？
• RQ5 LU机制能否解释'领悟'对数据集大小、权重衰减的依赖性，以及表征的出现？

主要发现

• LU机制——即在权重范数上，训练损失呈L形而测试损失呈U形——通过损失景观的不匹配，解释了'领悟'现象的起源。
• 在图像分类（MNIST）、情感分析（IMDb）和分子性质预测（QM9）任务中均观察到'领悟'现象，尽管其信号弱于算法任务。
• '领悟'的严重程度与测试性能对所学表征的依赖程度强相关，这种依赖在算法数据集中较高，而在MNIST中较低。
• 在训练过程中对权重范数施加约束，几乎完全消除了'领悟'现象，证实损失景观不匹配是该现象的核心。
• '领悟'的时间尺度满足 t ∝ 1/λ，与实证观察一致，支持LU机制的预测有效性。
• 本研究表明，简化损失景观分析是诊断数据-模型交互作用与表征学习动态的强大工具。