神经网络优化器背后的原理

可靠的优化是神经网络 (NN) 训练的核心，但现代 LLM 的默认优化器 Adam 却建立在脆弱的基础上。本论文为 Adam 奠定了原则基础并激发了新的设计。首先，我们重新审视 Adam 的发散-收敛争论，并证明存在与问题相关的相变：通过正确选择、与批量大小相关的超参数，Adam 会收敛，而在小 $β_2$ 状态下它可能会发散。其次，我们研究了为什么 Adam 通过 Hessian 结构在 Transformer 上大大优于 SGD。我们发现，Hessian 矩阵随着训练向近块对角形式演化，并伴随着强烈的块异质性。我们证明这种结构使得 Adam 的对角预处理器有效。我们进一步证明这种特殊的 Hessian 结构源于大矩阵变量的连续乘法，并基于随机矩阵理论提供了严格的分析。最后，这些见解激发了 Adam-mini，这是一种新的优化器，可以将 Adam 的内存占用量减少 50%，同时保持其性能。除了 Adam 之外，我们的结果还具有更广泛的影响：它们揭示了基于矩阵的非凸问题中的新局部结构，并且还有助于理解和改进最近的神经网络优化器，例如 Muon。