模块流形：神经网络训练中的流形优化新突破

近期，由OpenAI前首席技术官Mira Murati创立的Thinking Machines实验室再次公布了其重要研究成果！

这标志着他们继发布《克服LLM推理中的不确定性》论文后，推出了第二篇研究文章——《模块流形》（Modular Manifolds）。

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

训练大规模神经网络的过程犹如「走钢丝」，需要极为谨慎地维护其内部「健康状态」，避免权重、激活值或梯度等关键张量出现数值过大或过小的情况，以防引发数值溢出等连锁问题。

其中一项关键策略，是为大型模型建立一套统一的量级控制体系。

首要任务是巩固基础。

采用Layer Norm技术将每一层的输出调整到适宜范围，对激活向量执行归一化处理，这已成为当前广泛使用的标准方法之一。

对梯度更新进行归一化同样常见，例如Muon优化器通过谱归一化手段控制每一步更新的幅度，确保更新过程稳定。

更进一步，是直接对权重本身实施「管控」。

归一化权重矩阵是一个颇具潜力的探索方向。

文章提出了一种重新审视优化算法的新思路：将权重张量限制在特定的子流形上，从而能够与这些流形约束协同设计优化算法。

这相当于将「事后补救」转变为「事前预防」：

在初始阶段就将参数设置在健康范围内，使得训练过程更加稳定、更易解释，进而助力大模型实现更稳健、高效的训练。

流形优化器的形态

我们知道，流形是一种局部近似平坦的曲面。

当放大到足够尺度时，它看起来就像一个普通的平面。

流形上某一点附近的局部平坦空间被称为「切空间」。

如图1所示，三维球面或更高维的超球面是一个流形，图中红色区域展示了其在某一点上的切平面。

为了使权重能够「停留」在指定的流形内，一种简单的方法是使用常规优化器，并在每次更新后将权重投影回流形。

但问题在于，如果优化步骤偏离流形过远，再强行投影回来，会导致名义学习率不再对应参数在流形上的实际移动距离，从而削弱我们对「步长与效果」关系的直观理解。

要在流形上精心设计训练算法，首先必须明确：在切空间中应如何度量「距离」？

一种解决方案是直接在切空间内进行优化。这样，每一步都沿着流形「表面」前进，学习率能更准确地对应「实际位移」。

欧几里得距离是常见的选择，但也可以采用其他方式测量距离，如图2所示。

值得注意的是，距离度量方式的选择会直接影响最优优化步骤的方向。

图3中，粉色箭头代表原始梯度——即损失函数对权重的偏导数。

这意味着，我们并非必须严格遵循梯度方向移动。

为了用数学语言描述这一过程，我们可以将「在流形约束和特定距离度量下的最优更新方向」视为一个带约束的优化问题，例如以一个配备欧几里得范数的超球面为例。

用g表示梯度，w表示超球面上的当前点，a表示更新方向，η表示学习率，我们需要解决的问题是：

回到图1、2和3的可视化表述，这个公式的含义是：绿色箭头（即a的最优解）必须同时满足两个条件：

一是它必须位于红色的切平面上，二是它必须处于半径为η的黄色圆圈内。

我们可以运用拉格朗日乘数法进行求解。

其中λ和μ是拉格朗日乘子。

对该拉格朗日函数关于a求导并设为零，然后结合两个约束条件求解λ和μ，即可获得最优更新方向。

简而言之，最优更新的做法是：首先从梯度中剔除与w同方向的径向分量，即将梯度投影到切空间上，随后将结果归一化，再乘以学习率。

这样得到的更新方向便位于切空间之中。

图4展示了这个微小的调整过程，被称为「回缩映射」。

完整的流形优化算法如下所示：

总结来说，一阶流形优化器包含三个核心步骤：

寻找一个单位长度的切向量，使其在梯度方向上尽可能延伸；

用学习率乘以这个方向，然后从当前权重中减去；

通过回缩映射将更新后的权重拉回流形上。

在执行这一流程时，我们需要决定选择何种流形作为约束，以及如何定义「长度」的度量方式。

根据这两项选择的不同，我们能够衍生出不同的优化算法，具体见下表。

流形Muon

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

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

理解矩阵如何作用于向量的一个有效方法是使用奇异值分解，如图5所示。

SVD通过分解矩阵来展示矩阵如何沿不同轴向拉伸输入向量。

我们希望矩阵的「拉伸效应」接近1，因此选择了一个所有奇异值均为1的矩阵流形。

这种矩阵流形在数学上被称为Stiefel流形，在假设为高矩阵（m≥n）的情况下，它可以等价地定义为以下集合：

要为Stiefel流形设计优化器，还需选择一个合适的距离函数。

为了限制权重更新对输入向量的最大拉伸作用，谱范数——即矩阵最大奇异值的度量——是一个适宜的选择。

尽管它只约束了最大效应，但由于优化器会使这一上限饱和，因此也能间接防止最小效应过小。

正是这一构想，推动了Muon优化器的诞生。

将这一想法与Stiefel流形约束相结合，便形成了「流形Muon」问题。

文章中的一个关键发现是，这可以转化为一个凸优化问题，能够通过标准方法——对偶上升法来求解。

经过推导，对偶函数的梯度为：

通过一项小型实验，可以验证算法的可行性，实验设置与结果见图6。

模块流形

这里还存在一个重要问题：当我们将多个层组合起来构建完整的神经网络时，会发生什么？

是否需要关注层与层之间的相互作用，并据此调整优化策略？

这需要一种能够将前述推导逻辑推广至整个神经网络的方法——模块流形理论。

该理论的核心思想是：构建一种抽象机制，用于指导如何在各层之间合理分配学习率。

本质上，在不同层之间分配学习率，或对单个层进行缩放，都依赖于我们对网络输出相对于权重的Lipschitz敏感性的理解。

我们在构建网络的过程中会追踪这种敏感性，而流形约束有助于我们更精确地把握它。

参考资料：

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