摘要：近日，由OpenAI前CTO Mira Murati创立的Thinking Machines Lab发布了他们的第二篇研究《模块流形》（Modular Manifolds），试图从根本上改变我们训练大型神经网络的方式。

近日，由OpenAI前CTO Mira Murati创立的Thinking Machines Lab发布了他们的第二篇研究《模块流形》（Modular Manifolds），试图从根本上改变我们训练大型神经网络的方式。

研究者指出，当前训练大规模模型如同“走钢丝”，每一步都需小心翼翼，以防权重、激活值或梯度的数值溢出或消失。

传统做法依赖于Layer Norm或梯度裁剪这类“事后干预”，而新研究则提出一种更具前瞻性的思路，将权重约束在特定的数学“流形”上，从源头控制其行为。

这项研究不仅优化器的设计逻辑，更引入“模块流形”理论，尝试统一管理多层网络之间的学习动态。

这也意味着，未来的大模型训练，将不再依赖启发式调参，而是走向一种具备数学解释性、可严格控制的优化框架？

在人工智能领域里，大型神经网络的训练无疑是核心焦点所在，吸引着无数科研人员为之殚精竭虑。

如今，我们身处一个智能化应用无处不在的时代，语音助手能精准理解我们的话语，图像识别软件能快速分辨各种复杂场景，这些都离不开大型神经网络的强大助力。

然而，在这看似神奇的背后，训练大型神经网络却艰难前行。

训练过程中，权重、激活值、梯度这些关键的张量，稍不留意，它们的数值就可能变得过大或者过小，进而引发诸如数值溢出、梯度消失等一系列棘手问题，随时可能让整个训练成果功亏一篑。

面对这样的困境，传统的解决办法主要是依赖像 Layer Norm、梯度裁剪这类 “事后干预” 手段。

可这就好比亡羊补牢，虽然能起到一定作用，但总归是被动应对，没办法从根源上杜绝问题的产生。

就在大家都在为这些难题愁眉不展之时，OpenAI 前 CTO 团队重磅出击，他们发布的《模块流形》研究打破了这一僵局。

这个研究提出了一种极具前瞻性的思路，不再是等问题出现了再去补救，而是要从源头出发，把权重等关键要素牢牢控制住。

试图打造出一种全新的、具备数学解释性且能够严格控制的优化框架，让大型神经网络的训练告别那小心翼翼、如履薄冰的 “走钢丝” 状态。

而且这背后还有论文作者 Jeremy Bernstein 提出的 “模块化流形” 思路在助力。

要知道，以往传统的归一化方法大多都很少直接对权重矩阵本身进行归一化处理，而 “模块化流形” 思路可不一样。

它站在一个更为宏观和全面的视角，把整个网络看作是一个由多个部分组合而成的结构。

通过统一设计学习率、约束方式以及优化逻辑，有望攻克训练过程中那些因张量数值不稳定带来的难题，给大型神经网络训练开辟出一条崭新的道路。

当我们踏入大型神经网络训练这个神秘领域，想要探寻其中的优化奥秘时，就不得不先了解一个关键的概念 —— 流形。

流形，这个听起来有些抽象的词汇，其实有着独特的形态特点，在局部呈现出看似平坦的模样，可当你试着把观察的视角不断放大，就会发现它仿佛又变成了普通的平面，让人不禁感叹其奇妙之处。

而在流形之上，还有一个重要的 “小天地”，那就是某点附近的 “切空间”。

这 “切空间” 可不容小觑，它在整个流形优化的过程中扮演着至关重要的角色，就像是隐藏在幕后却掌控着全局的 “幕后推手”。

那在实际的神经网络训练中，我们为什么要让权重 “待在” 指定的流形里？

起初科研人员尝试使用普通优化器来解决这个问题，也就是在每一步更新之后，把权重投影回流形当中。

这看起来似乎是个合理的办法，可实际操作下来，却出现了让人头疼的新状况。

一旦优化步骤偏离流形太多，再强行把权重投影回来，就会导致名义学习率和权重在流形上实际产生的位移对不上号了。

好像原本规划好的路线图突然乱了套，让我们对 “步长 — 效果” 之间的关系也变得一头雾水，完全没办法按照预期去把控训练的节奏和效果。

你可别小看这个 “距离” 呀，不同的距离度量方式，它们居然会直接影响到最优优化步骤的方向。

每一种方式都可能引导着权重在流形上走出截然不同的 “轨迹”，这其中的关联既微妙又复杂。

不仅如此，流形优化还有两个关键的 “决策点”，那就是要确定采用哪种流形约束以及如何去测量长度。

这就好比我们在搭建一座复杂的桥梁，选择不同的建筑材料（流形约束）和测量尺度（长度测量方式），最终搭建出来的桥梁（优化算法）结构和性能都会大不相同。

而且这种差异在实际应用中体现得十分明显，比如把流形优化的思路从向量推广到矩阵的情况就是个很好的例子。

在探索大型神经网络训练的优化之路上，流形 Muon 算法有着独特的光彩和价值。

就拿 Transformer 里那个典型的权重矩阵来说吧，它就像是一个肩负重任的 “指挥官”，负责把输入向量转变为输出向量，而如何让这个 “指挥官” 恰到好处地发挥作用，可是个不小的难题。

这时候奇异值分解（SVD）就登场了，能清晰地展现出矩阵沿着不同轴拉伸输入向量的具体情况。

科研人员通过这个 仔细观察、深入思考，最终选择了 Stiefel 流形，这是一种所有奇异值均为 1 的矩阵流形。

可光有流形还不够，还得给它配上合适的 “搭档”，也就是距离函数，在众多的选择中，谱范数脱颖而出，它是什么？

其实它就是对矩阵最大奇异值的一种度量方式。别小看它它既能有效地约束住矩阵对输入向量作用时的最大效应，又能像个巧妙的 “平衡器” 一样，间接防止最小效应过小，让整个权重矩阵的作用发挥得更加稳定、合理。

正是基于这样巧妙的构思，Muon 优化器应运而生，同时也引出了那个充满挑战又趣味十足的 “manifold Muon” 问题。

首先，要在 Stiefel 流形的切空间上进行梯度更新的计算，这个过程就像是在一个特定的 “舞台” 上，按照既定的规则精心编排着每一步动作。

接着，还得动用奇异值分解、sign 函数等这些 “工具”，来确保每一次的更新都是合法合规的，就如同给整个更新过程上了一道道 “保险”。

最后再把更新后的结果稳稳地投影回 Stiefel 流形，让权重矩阵始终能在合适的 “轨道” 上运行。

那这个算法到底效果如何？在 CIFAR - 10 数据集上训练小 MLP 的实验给出了答案。

实验结果显示，流形 Muon 在训练和测试的准确率方面可是略胜一筹，和常用的 AdamW 相比，它有着自己的优势。

而且经过它处理后的权重矩阵，奇异值更加集中，幅度也更加稳定，就像训练有素的士兵，整齐有序。

不过它也有一点点小 “不足”，那就是每一步的运行时间比起 AdamW 来说会稍慢一些。

但这可难不倒科研人员，他们已经在思考后续的改进方向了，比如通过优化 dual ascent 步数、引入动量等方法。

来进一步减少这个算法的额外开销，让它能在大型神经网络训练中发挥出更卓越的性能。

模块流形理论有着一个极为核心的思想，那就是构建一种抽象机制，这个机制就专门用来指导我们如何在各层之间合理地分配学习率。

其实这依赖于我们对网络输出对权重的 Lipschitz 敏感性的深刻理解。这敏感性就好比是隐藏在网络内部的 “脉络”，虽然看不见摸不着，但却实实在在地影响着网络的成长和表现。

它就像是一个精准的 “导航仪”，帮助我们更加精准地去把握这种敏感性。

有了它，我们在分配学习率或者对单个层进行缩放的时候，就仿佛有了可靠的 “指南针”，不再盲目行事了。

具体来说模块流形理论有着独特的构造方式呢。它把神经网络里的每一层或者每一个模块。

都看成是一个单独的 “流形”，然后通过笛卡尔积把这些 “流形” 巧妙地拼接在一起，形成一个庞大且有序的流形空间。

在这个特别的空间里，采用最大范数的方式来统一分配学习率，就像是给这个 “生态系统” 制定了一套公平合理的 “规则”。

设置了全局的更新上限，避免了不同模块 “步伐” 不一致的情况出现，让整个网络的更新过程都能被这个全局机制牢牢地约束住。

这样做带来的好处可真是不容小觑。

它使得大型神经网络的训练变得更加稳定、高效了，就好比给一辆原本容易失控的 “列车” 装上了性能卓越的 “刹车” 和 “调速器”。

以前，训练那些千亿规模的大模型时，总是容易出现各种问题，可现在通过模块流形理论，能把出现这种问题的概率大大降低，甚至能降到千分之一。

这可给相关的公司节省了大量的预算，让他们可以更安心地去探索人工智能更深层次的奥秘，推动整个领域不断向前发展。

在人工智能的浪潮中，《模块流形》研究照亮了大型神经网络训练的前行道路。

从流形优化器的基础探索，到流形 Muon 算法的创新实践，再到模块流形理论的深远影响，一步步为解决训练难题贡献力量。

它不仅是技术的革新，更是开启更智能未来的希望之光，相信会引领人工智能迈向新的高度，绽放出更璀璨的光彩。

来源：快看张同学一点号