在大模型训练领域，如何有效管理权重、避免数值异常成为关键挑战。Thinking Machines Lab提出的“模块流形”理论为这一难题提供了创新性解决方案，将传统的事后数值修正转变为事前的约束优化，为大模型训练开辟了新路径。
训练大型神经网络如同精密操作，需严格控制权重、激活值和梯度等关键参数的数值范围。任何参数的过度膨胀或萎缩都可能引发数值溢出，导致训练失败。当前主流方法包括使用Layer Norm技术归一化层输出，以及通过Muon优化器等工具对梯度更新进行谱归一化处理，这些手段旨在维持参数的数值稳定。
研究团队提出更深入的解决方案：将权重张量约束在特定子流形上，并据此设计优化算法。这种“预防式”方法通过初始参数设置确保训练稳定性，使模型具有更强的可解释性。流形优化的核心在于将参数限制在局部平坦的曲面上，通过切空间优化实现精确控制。
切空间作为流形上某点的局部平坦区域，是实施优化的关键场所。以三维球面为例，其切平面构成优化操作的基准面。传统方法在每步优化后强制投影参数回流形，但可能导致实际位移与学习率脱节。研究团队提出直接在切空间进行优化，使学习率能准确反映参数的实际移动距离。
距离度量方式的选择直接影响优化方向。研究采用欧几里得距离作为基准，同时探索其他度量标准。通过数学建模，将流形约束下的最优更新方向转化为带约束的优化问题。以超球面为例，最优更新方向需同时满足切平面约束和半径约束，这一过程可通过拉格朗日乘数法精确求解。
具体操作中，最优更新包含三个步骤：首先将梯度投影到切空间，去除与当前点同方向的径向分量；其次对投影结果进行归一化处理；最后乘以学习率得到更新方向。这种“回缩映射”机制确保参数始终保持在流形约束范围内。
不同流形选择和距离度量组合可衍生出多种优化算法。研究团队特别关注Transformer权重矩阵的优化，提出基于Stiefel流形的约束方法。该流形要求所有奇异值均为1，通过奇异值分解可直观理解矩阵对输入向量的拉伸效应。配合谱范数作为距离函数，有效限制权重更新的最大和最小效应。
结合Stiefel流形约束和谱范数度量，研究团队开发出Muon优化器的流形版本。通过凸优化问题建模和对偶上升法求解，实现了参数更新的精确控制。实验验证表明，该算法能有效维持参数稳定性，防止数值异常。
当扩展至多层神经网络时，研究提出“模块流形”理论。该理论通过追踪网络输出的Lipschitz敏感性，指导各层学习率的合理分配。流形约束为这种敏感性分析提供了精确框架，使不同层之间的参数更新保持协调，避免因层间交互导致的训练不稳定。
