AI优化器技术演进：从AdamW到流形Muon

在人工智能的世界里，如果说大型神经网络是那个能够学习和思考的“大脑”，那么优化器（Optimizer），就是指导这个“大脑”如何学习、学习多快、以及最终能学得多好的“隐形之手”。从一个随机初始化的网络，到一个能够理解语言、生成代码的强大智能体，优化器在其中扮演着至关重要的角色。

然而，随着模型规模的爆炸式增长，我们熟知的经典优化器正面临着前所未有的挑战。近期，一家由前OpenAI核心团队创立的新公司Thinking Machines Lab（TML），通过其发布的“模块流形”研究，为我们揭示了AI优化器可能即将迎来的一场深刻的范式革命。

从SGD到AdamW

深度学习的早期，最基础的优化器是随机梯度下降（SGD）。它的原理很简单，就像一个蒙着眼睛的登山者，每走一步都沿着当前脚下最陡峭的方向（梯度的反方向）向下走，以期最快到达谷底（损失函数的最小值）。

但SGD的问题也很明显，它容易陷入局部最优的“小山谷”里出不来，或者在曲折的山路上反复震荡。为了解决这些问题，研究者们引入了“惯性”的概念，催生了Momentum等优化器。

随后，一个里程碑式的优化器Adam诞生了。它不仅考虑了惯性，还能为每个参数自适应地调整学习率，使其在各种任务中都表现出色。而我们今天最广泛使用的AdamW，则是在Adam的基础上，改进了权重衰减的处理方式，进一步提升了其性能和泛化能力。可以说，AdamW是当前深度学习优化器领域的“集大成者”。

新时代的挑战与“救火式”修正

当模型参数从百万级跃升至千亿、万亿级别时，AdamW等经典优化器开始面临新的挑战——数值不稳定性。在巨大的网络中，参数的更新很容易出现“爆炸”或“消失”的极端情况，导致训练过程崩溃。

为了应对这个问题，研究者们发明了各种“打补丁”式的修正方法。例如，Layer Normalization通过对每一层网络的激活值进行归一化，来保持其稳定。谱范数归一化等技术，则通过限制权重矩阵的“拉伸”程度来控制更新的幅度。这些方法虽然有效，但都属于“救火式”的被动修正。

TML的“流形优化”

TML发布的“模块流形”研究，提出了一种全新的、从根本上解决问题的思路：不再是在一个无约束的空间里自由下降，然后再想办法把跑偏的参数拉回来；而是一开始，就将参数约束在一个永远不会“跑偏”的、被称为“流形”的特定数学空间内进行优化。

为了验证这一理念，他们设计了一款名为**“流形Muon”的新型优化器。它将神经网络中的权重矩阵，约束在一个名为Stiefel流形**的特殊空间中。在这个空间里，所有矩阵的奇异值（可以理解为矩阵对数据进行拉伸或压缩的程度）都被严格限制为1。这意味着，无论模型如何更新，其内部的参数始终保持在一个“健康”、稳定的状态。

实验结果令人振奋。在经典的CIFAR-10图像分类任务上，使用“流形Muon”训练的神经网络，其性能优于了被广泛使用的AdamW优化器。这初步证明了“流形优化”这一新范式的有效性和巨大潜力。

优化器的未来展望

“流形Muon”的成功，仅仅是这场优化器革命的开端。TML的论文中，还展望了这一新范式未来的多个研究方向，包括：

• 如何为注意力机制、嵌入层等更复杂的网络结构设计合适的流形约束。
• 流形约束对低精度训练的潜在好处。
• 如何设计和调整约束，以提高模型的泛化能力。
• 如何在GPU上高效地实现相关的流形数学运算。

从SGD的简单梯度下降，到AdamW的自适应学习，再到如今“流形Muon”所揭示的几何优化新大陆，AI优化器的发展，正不断向着更稳定、更高效、也更接近数学本质的方向演进。

Thinking Machines Lab的这项研究，无疑为我们打开了一扇通往这个新世界的大门。它预示着，未来的AI大模型训练，可能将不再是一场充满不确定性的“炼丹”，而是一门更加严谨、更加可控的科学。