Thinking Machines Lab 提出“模块化流形”方法优化权重矩阵

Thinking Machines 研究员 Jeremy Bernstein 在博客中提出了一种新的神经网络优化思路——模块化流形（Modular Manifolds）。该方法通过在权重矩阵上施加几何约束，并结合特定优化策略，来提升模型训练的稳定性和泛化能力。

这一研究为神经网络训练提供了一种新颖的“几何视角”，或许能成为稳定大模型优化的新路径。

传统训练通常依赖激活标准化或梯度归一化，但很少直接对权重矩阵本身施加约束。Bernstein 指出，如果强制权重矩阵的奇异值保持在合理范围（如约束到 Stiefel 流形），就能避免过度拉伸或压缩输入，从而使训练过程更可控。

他提出的 Manifold Muon 优化器 在 CIFAR-10 上的实验显示，相比 AdamW，它能更好地保持权重的奇异值分布稳定，并带来一定性能提升，但计算开销更高。

进一步的“模块化流形”框架将这一思路推广到网络层级。每个模块不仅定义前向计算，还包含参数的流形约束和范数度量。多个模块组合时，其约束和敏感性界可以系统地合成，从而让整个网络在理论上具备更可控的 Lipschitz 性能。

未来方向包括：在不同层使用不同的流形约束、提升对偶优化的效率、探索非 Riemannian 度量，以及在大规模 GPU 训练中落地高效的流形计算。

推荐阅读：Thinking Machines Lab 发文，揭示 LLM 推理过程不确定性的真相