0. TL; DR

本文研究了语言模型中的临界数据规模（Critical Data Size, CDS），即模型从快速记忆转向缓慢泛化的阈值。作者通过调整初始化和权重衰减，成功在简单的语言模型任务上复现了Grokking现象，并提出了数据效率假设（Data Efficiency Hypothesis），揭示了数据不足、数据充足和数据过剩三种训练动态。实验表明，随着模型规模的增加，临界数据规模也增大，表明更大规模的模型需要更多数据才能实现泛化。这一发现为理解语言模型的训练机制提供了新的视角。

1. 背景介绍

在深度学习中，尤其是语言模型领域，数据规模对模型泛化能力的影响一直是一个关键问题。近年来，随着大型语言模型（LLMs）的发展，研究人员发现模型的泛化能力与数据规模密切相关。然而，这种关系背后的机制仍然不甚清晰。本文从Grokking现象出发，研究了语言模型在不同数据规模下的训练动态，揭示了从记忆到泛化的转变过程。

Grokking现象指的是神经网络在过拟合训练数据后，经过长时间的优化，突然开始泛化到测试数据的现象。这一现象表明，模型在训练过程中可能存在一个临界数据规模，超过这个规模后，模型的泛化能力会显著提升。本文通过实验验证了这一假设，并提出了数据效率假设，为理解语言模型的训练机制提供了新的思路。

2. 方法介绍

本文提出了数据效率假设，定义了临界数据规模（CDS），即模型在给定数据分布和训练程序下，达到约100%训练准确率和收敛测试性能所需的最小样本数量。

在数据分布 $D$ 和模型 $M$ 下，临界数据规模 $\text{CDS}(D,M)$ 定义为：

\[\text{CDS}(D,M) := \min \{ n \mid \mathbb{E}_{S \sim D^n} [\text{Acc}_S(M)] \geq \epsilon \}\]

其中，$\text{Acc}_S(M)$ 表示模型 $M$ 在测试样本 $S_{\text{test}}$ 上的平均准确率，$\epsilon$ 是收敛的测试性能阈值。

为了在语言模型中复现Grokking现象，作者提出了一个Grokking配置，通过调整模型参数的初始化和权重衰减来诱导Grokking。具体方法如下：

参数初始化：将模型的每个权重 $w$ 重新缩放为 $w = \alpha \cdot w_0$，其中 $\alpha = \frac{||w||_2}{||w_0||_2}$，$w_0$ 是标准PyTorch初始化的模型参数。这种调整可以放大Grokking现象。
优化目标：优化目标包括交叉熵损失 $L_{\text{CE}}(w)$ 和权重衰减 $L_{\text{WD}}(w)$：

\[\begin{aligned} &L(w)= L_{\text{CE}}(w) + L_{\text{WD}}(w)\\ &L_{\text{CE}}(w)= -\frac{1}{n} \sum_{(x,y) \in D} \log \left( \frac{\exp(h(x,y))}{\sum_{y' \in Y} \exp(h(x,y'))} \right)\\ &L_{\text{WD}}(w)= \frac{\lambda}{2} \|w\|_2 \end{aligned}\]

权重衰减：权重衰减在训练过程中起到关键作用，推动模型从记忆转向泛化。权重衰减的公式为：

\[w(t) \approx \exp(-\gamma t) w_0,t \approx \frac{\ln(w_0 / w^*)}{\gamma} \propto \gamma^{-1}\]

其中，$\gamma$ 是衰减常数，决定了衰减速度。$w^*$ 是泛化时的最优权重的L2范数。

为了验证数据效率假设，作者通过数据剪枝分析训练动态。具体方法如下：

均匀数据剪枝：从大小为 $N$ 的数据集中随机选择 $n$ 个样本，每个样本被选中的概率为 $\frac{1}{N}$。这种方法可以保持原始数据集的比例特性。
训练动态分析：通过在不同数据比例上训练模型，观察训练动态，包括训练准确率和测试准确率的变化。

3. 实验分析

1. 样本级Grokking（Sample-wise Grokking）

1.1 模块加法任务（Modular Addition）

在模块加法任务中，作者观察到Grokking现象与数据规模密切相关。当数据规模接近临界数据规模时，模型开始缓慢泛化。实验结果表明：

临界数据规模：在约3500步时，模型达到100%的测试准确率，表明泛化发生。
数据充足时的快速泛化：当数据规模超过临界数据规模时，泛化步骤显著减少，从3500步减少到500步。
数据不足时的无泛化：当数据规模低于临界数据规模时，模型只能记忆训练样本，无法泛化到测试样本。

模块加法任务验证了数据效率假设，表明临界数据规模是模型从记忆到泛化的关键转折点。

1.2 IMDB情感分类任务

在IMDB数据集上，作者通过调整初始化和权重衰减，成功诱导了Grokking现象。实验结果表明：

临界数据规模：在约350,000步时，模型突然从记忆训练数据转向泛化到测试数据，测试准确率从约50%提升到约87%。
数据充足时的平滑过渡：与模块加法任务相比，IMDB任务的相变更为平滑，表明语言任务的复杂性导致了更平滑的过渡。

IMDB任务验证了数据效率假设，并表明语言任务的复杂性会影响Grokking现象的表现。

1.3 Yelp情感分类任务

在Yelp数据集上，作者通过数据剪枝和Grokking配置，成功诱导了Grokking现象。实验结果表明：

临界数据规模：在约10%的数据子集上，模型表现出Grokking现象。随着数据规模的增加，泛化步骤逐渐减少。
数据充足时的快速收敛：与IMDB任务类似，Yelp任务的相变也较为平滑，表明更大的数据规模可以加速模型的收敛。

Yelp任务验证了数据效率假设，并表明更大的数据规模可以加速模型的泛化过程。

2. 模型级Grokking（Model-wise Grokking）

在IMDB数据集上，作者通过调整模型的隐藏层大小，研究了临界数据规模与模型规模的关系。隐藏层大小从16到256不等，数据集比例从10%到100%不等。实验结果表明：

临界数据规模随模型规模增加：随着隐藏层大小的增加，临界数据规模也增加。例如，隐藏层大小为16时，临界数据规模约为30%；而隐藏层大小为256时，临界数据规模约为70%。
平均准确率随模型规模减少：在固定的数据规模下，较大的模型不一定表现更好。实验结果表明，平均准确率随隐藏层大小的增加而减少。

模型级Grokking实验表明，更大的模型需要更多的数据才能实现泛化。这一发现强调了在学习单一复杂任务时，模型规模和数据规模之间的平衡。

3. Grokking机制分析

作者通过可视化分类层的权重，分析了模型从记忆到泛化的转变过程。实验结果表明，模型的权重在训练过程中经历了六个阶段：

随机初始化：初始权重具有高度随机性。
记忆阶段：权重开始收敛，模型开始记忆训练数据。
完全记忆：权重的L2范数增加，模型完全记忆训练数据。
L2范数减少和参数重新分配：权重的振幅减少，权重衰减开始发挥作用。
Grokking：模型突然泛化，测试准确率达到85%。
结束训练：权重在权重衰减的作用下进一步减少，但仍然保持在固定范围内。

Grokking机制分析表明，模型从记忆到泛化的转变是一个动态过程，权重衰减在其中起到了关键作用。