0. TL; DR

本文提出了一个统一框架，用于分析深度学习中的Grokking、U型尺度定律和涌现能力现象。通过研究记忆和泛化之间的竞争，作者揭示了模型大小和训练数据量对模型训练动态和最终性能的影响。该框架不仅解释了Grokking现象，还预测了U型尺度定律现象的发生，并通过多任务学习实验展示了涌现能力的形成机制。

1. 背景介绍

近年来，深度学习领域涌现出许多令人困惑的现象，例如Grokking（模型在过度拟合后突然泛化）、U型尺度定律（模型性能随模型大小或数据量增加先下降后上升）以及大型语言模型中的涌现能力（某些能力仅在模型规模足够大时出现）。这些现象挑战了人类对神经网络行为的直观理解，因此深入研究它们对于揭示深度学习的机制至关重要。

以往的研究大多集中于单独解释这些现象，但本文作者提出了一种全新的统一视角，将这三种现象纳入一个框架进行分析。这一框架基于记忆和泛化之间的竞争关系，通过研究模型大小和训练数据量对模型性能的影响，揭示了四种不同的训练动态：进展（Progression）、未Grokking（Ungrokking）、半Grokking（Semi-Grokking）和Grokking。这些动态不仅解释了Grokking和U型尺度定律现象，还为理解大型语言模型中的涌现能力提供了新的思路。

2. 方法介绍

本文的核心框架基于以下两个假设：

模型的临界数据集大小与模型大小负相关：较大的模型需要较少的数据即可发生Grokking。
模型的记忆容量与模型大小正相关：较大的模型能够记忆更多的训练数据。

基于这两个假设，作者构建了一个图形化的分析框架，将模型大小和训练数据量的关系划分为四个区域，每个区域对应一种独特的训练动态。

作者分析了四种训练动态：

进展（Progression）：当训练数据量超过模型的记忆容量时，模型无法完全记忆所有训练数据，因此会先部分记忆训练数据（此时验证性能为零），随后逐渐泛化，训练准确率和验证准确率同时提高。这种动态的关键在于模型在部分记忆后被迫发展泛化能力，以降低交叉熵损失。
未Grokking（Ungrokking）：当训练数据量小于模型的记忆容量且小于临界数据集大小时，模型会选择纯记忆，验证性能几乎为零。这是因为在这种情况下，记忆的效率高于泛化。
Grokking：当训练数据量介于记忆容量和临界数据集大小之间时，模型会先通过记忆达到完美的训练准确率，随后切换到泛化以提高效率。这种动态的关键在于模型在完全记忆后，由于泛化的效率更高，最终转向泛化。
半Grokking（Semi-Grokking）：当训练数据量接近临界数据集大小时，记忆和泛化的效率相当，模型会同时包含记忆和泛化，最终的验证准确率介于两者之间。这种动态表现为验证准确率的多次波动，表明模型在记忆和泛化之间反复切换。

U型尺度定律现象是指模型性能随模型大小（或数据量）增加先下降后上升。根据上述框架，当训练数据量远低于临界数据集大小时，模型会经历进展、未Grokking、半Grokking和Grokking四个阶段。在这个过程中，模型性能先随着模型大小增加而提高（进展阶段），然后下降（未Grokking阶段），最后再次上升（Grokking阶段）。这种先下降后上升的趋势正是U型尺度定律现象的体现。

3. 实验分析

作者通过一系列实验验证了上述理论框架。实验任务包括模加法任务（$(a + b) \mod P$）和更复杂的任务（$(a + b)^2 \mod P$）。模型采用单层简化版的Transformer架构，通过调整隐藏层大小来改变模型大小。实验中，训练数据量从2600到5000不等，模型隐藏层大小从4到100不等。每个实验重复11次，取平均值作为最终结果。

实验显示，较大的模型在较小的数据集上更容易发生Grokking。例如，隐藏层大小为64的模型在3000个训练样本上即可达到接近完美的验证准确率，而隐藏层大小为32的模型则需要4000个样本。结果表明，临界数据集大小与模型大小负相关。

当训练数据量为2600时，模型表现出明显的U型尺度定律现象：随着模型大小增加，验证准确率先上升（进展阶段），然后下降（未Grokking阶段），最后再次上升（Grokking阶段）。当训练数据量增加到3000时，U型尺度定律现象仍然存在，但不如2600时明显。进一步增加数据量到4000和5000时，U型尺度定律现象消失，模型性能随模型大小增加而单调上升。这验证了关于U型尺度定律现象的预测。

为了进一步验证框架的预测能力，作者通过增加任务的泛化难度（使用更复杂的任务$(a + b)^2 \mod P$），将原本不表现出U型尺度定律现象的数据集（3000个样本）转变为表现出明显U型尺度定律现象的数据集。这表明通过增加泛化难度，可以扩大U型尺度定律现象的出现范围。

作者将模加法任务与纯记忆任务（随机标签的减法任务）结合，发现较小的模型在多任务学习中难以发展出模加法任务的泛化能力，直到模型规模增大到约1570倍于单一任务训练时的规模，才表现出明显的泛化能力。这种现象类似于大型语言模型中的涌现能力，即某些能力仅在模型规模足够大时才出现。这表明多任务学习中记忆任务和泛化任务之间的竞争可能是涌现能力形成的重要原因。