Thinking Machines Lab新成果：模块流形优化器

近日，由OpenAI前CTO Mira Murati创立的Thinking Machines Lab再度发布了一项研究成果！

继《克服LLM推理中的不确定性》后，他们推出了第二篇研究文章——《模块流形》。

博客地址：https://thinkingmachines.ai/blog/modular-manifolds/

训练大型神经网络如同在钢丝上跳舞，必须小心翼翼，防止权重、激活值或梯度等关键张量过大或过小，以免引发数值溢出等问题。

其中一个重要的策略，是为大模型提供一套统一的量级管理。

首先是稳住基本盘。

使用Layer Norm技术将每层的输出拉回合适范围，对激活向量进行归一化，这也是目前一种普遍的做法。

对梯度更新进行归一化也很常见，例如Muon优化器对更新进行谱归一化处理，使每一步更新的幅度可控。

再进一步，是直接管理权重本身。

归一化权重矩阵是一个值得尝试的方向。

文中提出了一种重新思考优化算法的新视角：将权重张量约束在某个子流形上，以便与这些流形约束协同设计优化算法。

这好比把“救火”变“预防”：

一开始就把参数放在健康区间，让训练更稳、更具解释性，从而使大模型可以更稳定、高效地训练起来。

流形优化器的形态

我们知道，流形只是一个局部看起来很平坦的曲面。

如果放大到足够多，它看起来就像是一个普通平面。

流形上某一点附近的局部平坦空间称为“切空间”。

如图1所示，三维球面或更高维度的超球面是一个流形，图中以红色部分表示其在某点的切平面。

为了让权重能够“待在”指定的流形里，一个简单的方法是使用普通优化器，在每步更新后将权重投影回流形。

但问题是如果优化步骤偏离流形太多，再被强制投影回来，这会导致名义学习率不再对应参数在流形上的实际位移，从而削弱我们对“步长—效果”关系的直觉。

流形Muon

Transformer中的典型权重矩阵W是一个“向量变换器”，即它将输入向量x转换为输出向量y=Wx。

我们希望设计一种流形约束和距离函数，使得该矩阵对输入向量的作用合理：既不应导致输出值过大或过小，也不应在更新权重时引起输出向量剧烈变化或几乎无变化。

模块流形

当我们将多个层组合起来构建完整的神经网络时，是否需要考虑层与层之间的交互，并据此修改优化策略？

这需要一种可以将前文介绍的推导逻辑推广到整个神经网络的方法——模块流形理论。

参考资料：

https://thinkingmachines.ai/blog/modular-manifolds/