Thinking Machines Lab 提出“模块化流形”方法优化权重矩阵
Thinking Machines 研究员 Jeremy Bernstein 在博客中提出了一种新的神经网络优化思路——模块化流形(Modular Manifolds)。该方法通过在权重矩阵上施加几何约束,并结合特定优化策略,来提升模型训练的稳定性和泛化能力。
这一研究为神经网络训练提供了一种新颖的“几何视角”,或许能成为稳定大模型优化的新路径。
传统训练通常依赖激活标准化或梯度归一化,但很少直接对权重矩阵本身施加约束。Bernstein 指出,如果强制权重矩阵的奇异值保持在合理范围(如约束到 Stiefel 流形),就能避免过度拉伸或压缩输入,从而使训练过程更可控。
他提出的 Manifold Muon 优化器 在 CIFAR-10 上的实验显示,相比 AdamW,它能更好地保持权重的奇异值分布稳定,并带来一定性能提升,但计算开销更高。
进一步的“模块化流形”框架将这一思路推广到网络层级。每个模块不仅定义前向计算,还包含参数的流形约束和范数度量。多个模块组合时,其约束和敏感性界可以系统地合成,从而让整个网络在理论上具备更可控的 Lipschitz 性能。
未来方向包括:在不同层使用不同的流形约束、提升对偶优化的效率、探索非 Riemannian 度量,以及在大规模 GPU 训练中落地高效的流形计算。
