Is your batch size the problem? Revisiting the Adam-SGD gap in language modeling

众所周知，Adam 在语言模型中的表现显著优于随机梯度下降（SGD），对此现象已有多种解释被提出。在本研究中，我们通过一系列针对 Transformer 语言建模精心调整的基线训练运行，重新审视了这一“优化器差距”。我们全面分析了动量、梯度裁剪和批量大小如何影响 SGD 与 Adam 之间的性能差异。我们的实证结果表明，在小批量场景下，如果正确调参，带动量的 SGD 实际上可以表现出与 Adam 类似的性能。我们重新评估了现有的关于 Adam 优势的解释，包括重尾类别不平衡、方向性锐度以及 Hessian 异质性等，但这些理论难以直接解释该现象。为了弥合我们对这一问题理解上的差距，通过对 Transformer 训练过程及受文献启发的简单二次函数场景进行分析，我们基于随机微分方程模型，提供了关于批量大小对训练动力学影响的新见解。

论文试图解决Adam优化器在语言模型训练中表现优于SGD的现象（即所谓的'optimizer gap'），并探讨这一现象是否可以通过调整SGD的超参数或重新分析优化动态来解释。这是一个长期存在的问题，但本研究通过系统性实验和理论分析提供了新的视角。

关键思路是通过全面调优SGD（包括动量、梯度裁剪和批量大小等因素）来缩小与Adam的性能差距，并结合随机微分方程模型分析不同批量大小对训练动态的影响。相比以往研究，该论文不仅验证了SGD在小批量设置下的潜力，还提出了批量大小如何影响优化器行为的新见解。

论文设计了详尽的实验，覆盖多种超参数配置和Transformer架构，使用公开的语言建模数据集（如WikiText-103）。此外，作者引入了简单的二次函数设置以简化分析，并结合Hessian矩阵特性进行讨论。代码已开源，为未来研究提供了复现基础。值得深入研究的方向包括进一步探索批量大小对泛化性能的影响以及更复杂的非凸优化场景下的动态分析。

相关研究包括：1) 'On the Convergence of Adam and Beyond' 提出了Adam收敛性的理论分析；2) 'Sharpness-Aware Minimization for Efficiently Improving Generalization' 探讨了方向锐度对泛化的影响；3) 'Why Does收集 More Data Improve Generalization?' 研究了数据规模与泛化的关系；4) 'The Heavy-Tail Phenomenon in SGD' 分析了SGD在重尾分布下的特性。这些研究共同构成了理解优化器行为的基础。