Mira Murati 离开 OpenAI 之后，其团队近日接连发布了两项硬核研究成果。首先是那篇广为关注的论文《Defeating Nondeterminism in LLM Inference》，该文明确指出大语言模型推理中的不确定性根源并非 GPU 并发运算本身，而在于批处理大小变化所导致的数值差异。紧接着，团队昨天又放出一篇聚焦训练稳定性的新文章。
这项新研究有一个颇具哲学意味的名称：“模块化流形”（Modular Manifolds）。简而言之，其核心目标就是让神经网络训练过程更加稳定、更加可预测。
Image
训练大型神经网络从来都不是一件轻松的事。权重可能突然爆炸或消失，梯度忽大忽小，整个过程如同在走钢丝。目前业界普遍依赖各种归一化技术来控制激活值和梯度，但很少对权重矩阵本身加以约束。
说到“流形”这一概念，学术界已经热议了相当长的时间。为什么我们将嵌入表示称为 embedding 而非简单的词向量？正是因为其与流形有着密切的内在关联。当你训练一个拥有数十亿参数的语言模型时，这些参数并非在整个高维空间里随意游荡，而是被限制在远低于原维度的低维子空间中运动。有效的参数配置大多集中在高维空间里某些低维子结构上——这些子结构正是流形的体现。
Thinking Machines 团队的想法非常直接：既然参数天然倾向于沿着某些流形运动，何不主动设计这些流形？他们提出将权重约束在特定的数学“形状”上——好比让权重沿着预定的轨道滑动，而不是在高维空间中漫无目的地游荡。
核心思路是将权重矩阵约束在 Stiefel 流形上——这是由所有正交列矩阵构成的空间。数学上表示为 W^T W = I，即矩阵的所有奇异值均为 1。这样做最直接的收益是：矩阵的条件数始终为 1，从根本上避免了数值不稳定的问题。
这一点对于大模型尤为关键。当模型拥有 100 层 Transformer 时，每一层微小的不稳定性都会被逐级放大。如果每层的权重都稳定地处于 Stiefel 流形上，整个网络的行为就会显著改善：前向传播时信号既不会爆炸也不会消失，反向传播时梯度流动更加平稳，不同层之间的相互影响也变得更加可预测。
他们设计的 Manifold Muon 优化器，本质上是在求解一个带约束的优化问题：
min trace(G^T A)
s.t. "|A||_spectral ≤ η
A^T W + W^T A = 0
第一个约束限制了更新步长，第二个约束确保更新方向位于切空间内。这是一个凸优化问题，团队采用对偶上升法进行求解。具体步骤为：先引入拉格朗日乘子 Λ，将约束最小化问题转化为无约束最大化问题；然后通过梯度上升更新对偶变量：Λ ← Λ + α × H(Λ)；接着计算最优更新：A_opt = -η × msign(G + 2W(Λ + Λ^T))；最后更新权重并投影回流形：W ← msign(W + A_opt)。
这里的 msign 是矩阵符号函数，可以将矩阵的所有奇异值归一化到 1。在计算上，可以通过 Newton-Schulz 迭代或最新的 Polar Express 算法在 GPU 上高效实现。
有意思的是，团队选择了谱范数而非 Frobenius 范数作为距离度量。原因在于：谱范数直接决定了矩阵作为线性算子的最大放大倍数，更契合神经网络中权重矩阵的实际作用。这一选择让优化器能够感知流形的精确形状，从而设计出更优的更新方向。
更精妙之处在于“模块化”设计。他们提出了一套组合规则：
- 前向函数：f₃((w₁,w₂),x) = f₂(w₂, f₁(w₁,x))
- 流形约束：M₃ = M₁ × M₂（笛卡尔积）
- 范数定义："|(w₁,w₂)||₃ = max(s₁||w₁||₁, s₂||w₂||₂)
这里的系数 s₁ 和 s₂ 能够自动调节不同层的学习率分配。该框架可以追踪网络的 Lipschitz 常数，确保每一层的更新都在合理范围之内。
笛卡尔积是组合流形的一种简洁方法。打个比方：一条线和一个圆盘的乘积就是圆柱体——在线上的每个点都有一个圆盘的副本。这样的构造使得研究者可以为网络的不同部分设计各异的约束，然后优雅地将它们整合在一起。
在 CIFAR-10 数据集上，他们使用一个小型 MLP 进行了实验，结果显示 Manifold Muon 相比 AdamW 获得了更高的训练和测试准确率。训练结束后，权重矩阵的奇异值分布高度集中在 1 附近，充分证明了约束的有效性。
这种流形视角还能解释大模型训练中的许多现象。例如，LoRA 之所以有效，本质上是因为它假设参数更新发生在一个低秩流形上。模型坍缩往往是由于参数被推到了流形的退化区域，而训练不稳定则源于优化轨迹偏离了合理的流形结构。
当然，计算开销是一个现实问题。每一步需要执行多次对偶上升迭代，并计算矩阵符号函数。有人怀疑这种方法能否扩展到十亿参数级别的网络，毕竟 SVD 分解和投影操作的复杂度并不低。不过团队也提到，可以通过减少迭代次数或加入动量等策略进行优化。
从理论层面看，仍有许多开放问题值得探索。注意力头应该采用什么流形？Query、Key、Value 矩阵的内在平衡是否需要特定的流形结构？嵌入层和输出层的约束是否应该不同？这些方向都有待未来的研究。
从更大的图景来看，流形约束可能带来更多好处：更好的泛化能力（被约束在流形上的模型不容易过拟合）、更强的可解释性（流形的几何性质可能对应模型的语义特征）、以及更高效的训练（了解参数空间的结构有助于设计更优的优化路径）。
这种方法也与现有技术存在呼应。EDM2 扩散模型已经使用了权重约束，BiT 则采用了权重标准化。但将流形优化、模块化设计和自动学习率分配三者结合，确实提供了一种新的思路。
目前这仍只是一个研究原型，要真正应用于大模型训练，还需要解决不少工程难题。不过这个方向确实值得关注。与其在巨大的参数空间中进行盲目搜索，不如在精心选择的流形上优雅地滑行。如果真能让训练过程变得更容易预测、更稳定，那对整个领域而言都是一件值得期待的事情。
论文原文：https://thinkingmachines.ai/blog/modular-manifolds
