Beyond Scaling Laws: Understanding Transformer Performance with Associative Memory

增加Transformer模型的大小并不总是会提高性能。这种现象无法用经验缩放定律来解释。此外，当模型记忆训练样本时，改进的泛化能力也会出现。我们提出了一个理论框架，阐明了基于Transformer的语言模型的记忆过程和性能动态。我们使用Hopfield网络对Transformer的行为进行建模，使得每个Transformer块有效地进行近似最近邻搜索。基于此，我们设计了一种能量函数，类似于现代连续Hopfield网络中的能量函数，为注意力机制提供了深入的解释。使用majorization-minimization技术，我们构建了一个全局能量函数，捕捉了Transformer的分层结构。在特定条件下，我们证明了最小交叉熵损失的下限约为1。我们通过在各种数据大小上使用GPT-2进行实验以及在包含2M标记的数据集上训练vanilla Transformer来证实我们的理论结果。

论文试图解决Transformer模型的性能提升并不总是随着模型大小增加而增加的问题，并探究了该现象的原因。同时，论文也试图解释Transformer模型的记忆过程和性能动态。

论文使用Hopfield网络模拟了Transformer模型的行为，将每个Transformer block作为一个近似最近邻搜索的实现。在此基础上，构建了一个类似于现代Hopfield网络的能量函数，并使用majorization-minimization技术构建了一个全局能量函数。论文还证明了在特定条件下，最小交叉熵损失的下界约为1。

论文使用GPT-2模型在不同数据集上进行了实验，并在2M个token的数据集上训练了普通Transformer模型。实验结果验证了论文理论结果的正确性。此外，论文还提供了开源代码。

近期的相关研究包括《Attention is not Explanation》、《The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks》等。