Thinking Machines再发文：模块化流形让训练更稳定

Mira Murati 离开 OpenAI 后，最近分享了她团队在神经网络训练推理上的新研究。他们之前在文章《Defeating Nondeterminism in LLM Inference》里指出了大语言模型推理中不确定性的真相：问题不在于GPU并发运算，而是批处理大小变化导致的数值差异。昨天，又发布了一篇关于训练稳定性的文章。

这次的研究叫"模块化流形"（Modular Manifolds），听起来很玄乎，其实就是想让神经网络训练更稳定、更可预测。

Image

训练大型神经网络一直是个头疼的事。权重可能突然爆炸或者消失，梯度忽大忽小，整个过程像在走钢丝。现在大家都在用各种归一化技术来控制激活值和梯度，但权重矩阵本身却很少有人管。

其实，流形是一个学术圈很热点的研究，就比如embedding之所以叫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

第一个约束限制更新步长，第二个约束确保更新方向在切空间内。这是个凸优化问题，他们用对偶上升法来求解。具体步骤是：

1. 引入拉格朗日乘子 Λ，把约束最小化问题转换成无约束最大化问题
2. 通过梯度上升更新对偶变量：Λ ← Λ + α × H(Λ)
3. 计算最优更新：A_opt = -η × msign(G + 2W(Λ + Λ^T))
4. 更新权重并投影回流形：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