模块化流形
当我们训练大型神经网络时,需要让它们保持“健康”。我们不希望网络中的张量——无论是权重、激活还是梯度——变得过大或过小。过小和过大的张量会带来各种问题,并不只限于数值下溢与上溢。例如,训练过程中权重矩阵的尺度变化会使优化算法的设计更困难——因为相对于权重的更新幅度会显著影响学习速度。
保持张量健康的“金标准”是对其进行归一化。对激活向量进行归一化很常见,例如使用层归一化在将激活传给下一层之前把其尺度调整到合理范围。对梯度更新进行归一化也很常见,我们可以把诸如Muon 优化器等快速训练算法理解为对更新做谱归一化。归一化让我们对张量大小有把握——无需每次都去看 Wandb!——而在包含许多相互作用组件的大型神经网络中,对网络内部情况的把握非常宝贵。
对权重矩阵进行归一化则不那么常见,尽管并非闻所未闻。例如,EDM2 扩散模型代码库采用了权重约束,作者在论文中报告了收益;BiT 使用了权重标准化。还有许多其他技术被提出,但在现代大规模训练中并未成为通用实践。更多示例可见 Salimans 等, 2016、Miyato 等, 2018 以及我们的论文 Liu 等, 2021。对权重矩阵做归一化可能有多方面好处:权重约束使得理解优化更新的相对大小更容易;它们消除了权重范数爆炸的问题;它们让我们能够将超参数调优的精力集中到那些“大小最关键”的张量上;它们可以强制矩阵具有较小的条件数,从而使其行为更可预测;此外,权重约束还能促成关于扰动鲁棒性的Lipschitz 保证。
本文介绍一种约束神经网络权重矩阵的有吸引力方式——让每一层的张量受限于某个子流形。这为重新思考优化打开了大门,因为我们可以与这些流形约束协同设计优化算法。作为示例,我们提出一种 Muon 优化器的流形版本(该算法构建在苏剑林与 Franz Louis Cesista 的工作之上,后文将进一步讨论),其权重受限于 Stiefel 流形:即条件数为 1 的矩阵流形。最后我们提出“模块化流形”(modular manifold)的概念——一种可组合的流形——旨在让大规模网络更易扩展与训练。
我们写这篇文章的目标是为一个我们非常兴奋的研究方向提供入门介绍,并强调许多未来工作方向。我们也很期待社区就文末提到的话题展开更多研究!
流形优化器的形状
本节通过最简单的流形学习示例展开:一个被约束在 超球面上的向量参数。我们在整个空间 上定义损失函数,并训练该向量参数以最小化损失。这种设置在实践中很有用,比如 transformer 中的单个嵌入向量。本节也为下一节讨论矩阵参数的“流形版 Muon”做热身。
我们不对流形做过于形式化的定义:只需理解为“一个在足够近处看起来是平的曲面”。在流形上某一点处的局部线性近似称为该点的“切空间”(tangent space),如图所示:
三维的球面(或更高维的超球面)是一个流形。在流形上一点的局部平坦近似即为该点的切空间,图中以红色平面示意。
我们可以把 维超球面刻画为欧氏范数为 1 的点集 。而超球面上点 处的切空间是与 正交的所有向量 的集合。
为了让权重保持在流形上,我们可以使用“非流形”的优化器,并在每步之后把权重投影回流形。相较之下,我们更感兴趣的是设计在切空间中行走的方法。原因在于,我们希望能够将优化器的学习率与优化步长的真实长度等同起来;但如果优化步显著偏离流形再被投回,这个良好性质将不再成立。类似的动机可见 EDM2 论文 第 2.3 节。
在为该流形设计训练算法之前,一个重要决定是:我们要如何在切空间里度量“距离”(对于一个流形要成为“黎曼流形”,其距离度量必须由内积诱导。欧氏()范数由内积诱导,但曼哈顿()距离则不是)。常见选择是欧氏距离,但我们也可以选择其它度量方式,如图所示。下一节我们会基于模块的功能来讨论如何选择距离度量。
不同距离度量在切空间中内切单位球。(欧氏)单位球是圆,而 (曼哈顿)单位球是菱形。
关键在于,距离度量的选择会改变“最佳优化步”的方向。如果距离度量不是欧氏的,那么在固定步长下,不完全沿着梯度方向,我们反而能在“沿梯度方向的投影距离”上走得更远(此处“梯度”指损失对权重的偏导;在黎曼几何中,数学家将“梯度”一词留给另一种对象)。如图所示。
几何如何影响最佳优化步的方向。粉色箭头表示原始梯度(损失对权重的偏导);黄色菱形表示 单位球;绿色箭头是“在梯度方向上最远的单位向量”。注意绿色箭头并不与粉色箭头平行。(实际中,粉色箭头未必位于切空间,但绿色箭头按构造总在其中。)可尝试拖动粉色箭头看看最优更新方向如何变化。
为了从数学上看清这一点,我们把“给定流形约束与距离度量的最优更新方向”本身表述为一个带约束的优化问题。以下以“装备欧氏范数的超球面”为例:记 为梯度、 为当前处于超球面的点、 为更新方向、 为学习率,我们需要求解:
回到图示化语言,这个公式表示:绿色箭头( 的最优取值)必须位于红色的切平面上(),同时还必须落在半径为 的黄色圆上()。为求解 ,我们可使用拉格朗日乘子法。相应的拉格朗日函数为:
其中 与 为拉格朗日乘子。对 求偏导并令其为零,再结合约束解出 与 ,即可得到最优更新 的表达式:
口头解释:把梯度的“径向分量”去掉,对方向归一化,再乘以学习率。由于该更新位于切空间,为了严格回到流形上,还需要一个非常小的纠正(对于学习率 ,回缩映射的影响仅为 ,因此学习率与实际步长几乎相等)。这个纠正称为“回缩映射”(retraction map),如图所示:
回缩映射示意。绿色箭头是切空间中的更新。由于当步长较大时切空间会偏离流形,需要使用回缩映射把更新后的权重投影回流形——紫色箭头所示。
我们可用勾股定理得到回缩映射:对于单位超球面与步长 ,斜边长度为 ,因此在欧氏范数下超球面的回缩映射就是将更新后的权重除以 。合在一起,完整的流形优化算法为:
读者可做一个小练习:计算更新后权重向量的欧氏范数,验证其确实仍位于超球面上。
小结一下,本节的一阶流形优化器包含三步:
- 在切空间中,找出“在梯度方向上走得最远”的单位向量;
- 乘以学习率并从权重中减去;
- 使用回缩映射把更新后的权重拉回流形。
应用这套流程有两项关键决策:该用什么流形约束、以及如何度量“长度”。做出不同选择,就会导出不同的优化算法,如下表所示。
| 流形 | 范数 | 优化器 | 
|---|---|---|
| 欧氏空间 | 欧氏范数 | 普通梯度下降 | 
| 欧氏空间 | 无穷范数 | 符号梯度下降 | 
| 超球面 | 欧氏范数 | 超球下降 | 
| 矩阵空间 | 谱范数 | Muon | 
| Stiefel 流形 | 谱范数 | 流形版 Muon | 
下一节我们将推导表格中的最后一个算法:流形版 Muon。要为矩阵参数设计流形约束与距离度量,需要认真思考“权重矩阵在网络中的作用”。
流形版 Muon
Transformer 中的典型权重矩阵 是一个“向量乘子”,即将输入向量 变换为输出向量 。我们将设计一个流形约束与距离度量,使该矩阵对输入向量的作用“良好”:矩阵不该产生过小或过大的输出;对矩阵的更新也不应让输出向量变化过多或过少。
理解矩阵如何作用于向量的一个好方式是通过奇异值分解(SVD),如图所示。SVD 以一种能揭示“沿不同轴线拉伸多少”的方式将矩阵分解。
mxnmxkkxkk rank M=kxn=MUΣVT 奇异值分解。秩为 的矩阵 总可以分解为 ,其中 与 的列两两正交归一, 为仅含正条目的对角矩阵。 的对角元素称为 的奇异值;它们刻画了矩阵对与 、 相应列对齐的向量的拉伸效应。
我们希望矩阵的拉伸效应接近 1,因此选择一个“所有奇异值都为 1”的矩阵流形。这个流形正式地称为Stiefel 流形。不失一般性,我们可假设处理的是“高矩阵”(),此时 Stiefel 流形可等价地定义为:
\mathsf{Stiefel}(m,n) := \left\\{ W \in \mathbb{R}^{m \times n} \mid W^T W = I_n \right\\}.
此外,可以证明,矩阵 在 处切于 Stiefel 流形,当且仅当: (注意 Stiefel 约束 直接推广了上一节的超球约束 ;类似地,切空间条件推广为 。)
要在 Stiefel 流形上设计优化器,剩下要做的是选择距离度量。为了限制“权重更新对输入向量的最大拉伸效应”,我们选择“谱范数”(spectral norm),它度量的是矩阵的最大奇异值。尽管这只限制了更新的最大效应,但由于我们推导的优化器会“饱和”这一上界,它也会防止“最小效应”过小(也有例外,比如权重矩阵的输出通道少于输入通道时,矩阵及其更新不可避免地存在零空间,会把某些输入映射为零)。
在谱范数约束下做梯度下降的想法促成了 Muon 优化器;当其与 Stiefel 流形约束结合时,我们得到如下问题,我们称之为“流形版 Muon”:
流形版 Muon 问题 直接推广了上一节的问题 。求解 比 更困难。这里我们给出一个数值解法,灵感来自苏剑林与 Franz Louis Cesista 的工作。(我在去年年底解决了“方阵情形”的流形 Muon,但无法解决完整的“矩形情形”,因此在 Modula 文档中把它作为公开问题提出。今年夏天,苏剑林采用拉格朗日方法,并围绕最优性条件给出不动点迭代,从而解决了该问题。我看到过他早期的版本(当时尚未完全奏效)以及 Franz Louis Cesista 的相关工作,这使我能够推导出本文给出的“对偶上升”算法。)
我们的关键洞见是: 是一个凸优化问题,可以用一种标准方法——对偶上升——来求解。下文只给出主要思路,更详细的推导见此页。
与苏剑林的做法类似,我们首先引入一个拉格朗日乘子矩阵 。随后进行一系列变换,把 从“带约束的最小化问题”转化为“无约束的最大化问题”:
式(1)把问题重写为一个鞍点问题:只要违反切空间条件,关于 的最大化就会把目标送到无穷大。式(2)由迹的性质得到,式(3)由 Sion 极小极大定理得到。式(3)中的内层最小化可通过令 来实现,其中 为矩阵符号函数(它把一个矩阵的奇异值“拉到 1”上;可用 Newton–Schulz 迭代或近期的 Polar Express 算法在 GPU 上高效计算)。把该表达式代回式(3)即可得到式(4)。式(4)称为 的“对偶问题”,我们可以通过梯度上升来求解。进一步可得其梯度为:
其中核范数 度量矩阵奇异值之和。
最后,可以写出流形版 Muon 算法(注意它与苏剑林的解法密切相关:我们使用对偶上升,而他的解法等价于通过不动点迭代求解对偶函数的极大点 ):
- 在对偶变量上做梯度上升:,得到 ;
- 计算更新:;
- 应用更新到权重:;
- 回缩回流形:。
我们做了一个很小的实验来校验该算法,并提供了最小可玩实现,方便学生或研究者上手。每次训练在一分钟内完成。代码见此处,设置与结果见图示。
2025-09-25T11:22:41.551982 image/svg+xml Matplotlib v3.10.5, https://matplotlib.org/ 在 CIFAR-10 数据集上用一个小型 MLP 训练 3 个 epoch。浅蓝色曲线表示 AdamW 的不同权重衰减设置。结果为 3 个随机种子的平均。与 AdamW 相比,流形版 Muon 在训练集与测试集上都取得更高精度。第三张图显示了在最佳学习率下第一层权重矩阵的最终奇异值分布:使用流形版 Muon 训练后,奇异值都接近 1。相较 AdamW,流形版 Muon 的单步墙钟时间更长,但可通过减少对偶上升的步数,或在算法中引入动量并在线执行对偶上升来改善。视系统其他瓶颈而定,这一开销也可能不成问题。
模块化流形
到目前为止,我们讨论了针对单个参数张量的流形约束,并为这些约束协同设计了优化逻辑。尚未回答的问题是:当我们把层组合成网络时会发生什么?我们能否把每一层独立看待——抑或需要注意层与层之间的交互,并据此调整优化逻辑?本节的目标是指出:我们可以把前两节的推理扩展到“整个网络”的情形,我们把这称为“模块化流形理论”(the theory of modular manifolds)。该理论基于我与合作者 Tim Large、博士后导师 Phillip Isola、博士导师岳毅松以及许多优秀同事的研究成果而来;文末提供了一些进一步阅读链接。
“模块化流形”的思想是在网络层面提供一种抽象,用来指导我们如何在层之间“分配学习率预算”。每一层内部的优化逻辑与前文推导一致,只是该层的学习率会因其在网络中的位置而被调整。这个抽象建立在我们关于“模块化范数”论文中的一个关键观察之上:无论是在层之间做学习率预算,还是在扩展单层规模时进行预算,本质上都与“网络输出对权重的 Lipschitz 敏感性”密切相关。这个抽象在搭建网络的过程中追踪这种敏感性,而流形约束则帮助我们更紧地刻画这种敏感性。
这一抽象的起点是:把任意神经网络模块(从单层到整个 transformer)视作具有三项属性的数学对象:
- 前向函数 ,将参数空间 与输入空间 映射到输出空间 ;
- 权重空间的一个子流形 ,权重受其约束;
- 一个范数 ,作为权重空间上的“尺子”。
例如,一个线性模块配备谱范数并受限于 Stiefel 流形(我们已经为其推导了优化器),可写为:
若输入到 模块的 具有单位 范数,则在该模块指定的范数下, 关于其权重是 Lipschitz 的,且常数为 1(该论证可拓展到输入的 RMS 范数以及权重上的 RMS–RMS 算子范数):
这种 Lipschitz 论断有助于指导我们如何缩放该模块的权重更新,因为它给出了“当扰动权重时输出最多能改变量”的上界。但当我们把两个模块进行复合时,能否自动为新模块的联合权重空间编译出一个 Lipschitz 论断?答案是肯定的——前提是我们遵循一套构建新模块的特殊规则:
- 新的前向函数  由已有的  与  复合而成:
- 新的流形约束  是两者流形 、 的笛卡尔积(见图中的一个有趣示例):
- 新的范数函数是两者范数的“加权取最大”。令  表示第一个模块的范数、 表示第二个模块的范数,则新的范数  为:
当我们用该“组合范数”来导出优化器(沿用前两节的配方)时,本质上为每一层导出了各自的优化器;而标量系数 则在层之间为学习率做预算分配。
关于该构造的更多细节(包括扩展到其它模块组合方式)见我们关于模块化范数的论文,尽管该文并未涵盖流形优化。关于在模块化范数下构建优化器,亦可参考我们的模块化对偶论文。Modula 项目正朝着将这一构造程序化实现的方向推进。
=x 笛卡尔积是把两个流形“粘合”在一起的一种简单方式。例如,一条直线与一个圆盘的积是一个圆柱:直线上的每一个点都对应一份圆盘的拷贝。
未来工作方向
我们对任何试图让神经网络训练像前向传播一样“有原则且自动化”的研究都感到兴奋。本文的诸多想法得益于与外部研究者(如苏剑林、Franz Louis Cesista)的交流。我们很期待社区在这些主题上开展更多工作。
一些可能的后续方向:
- 模块化。 注意力头该受限于什么流形?嵌入与反嵌入是否应采用不同约束?我们可以在网络不同部位混搭约束,或让某些张量保持不受约束。
- 数值。 流形约束也会限制单个权重条目的取值范围。这是否影响数值性质,或使低精度训练更容易?
- 凸优化。 流形版 Muon 涉及对偶上升。能否采用更精细的凸优化技术,使对偶问题更快或更可靠地求解?
- 收敛分析。 这些算法收敛得有多快?良好的权重矩阵条件数是否有助于收敛?能否给出更多理论结论?
- 正则化。 流形约束隐式正则化了模型。我们是否可以设计约束或调整其半径以改进泛化?
- 结构–优化器共设。 虽然“硬流形约束”未必是约束权重矩阵的终极方案,但它体现了将优化器与结构组件紧耦合共设的理念。这里是否还有更多机会?
- 非黎曼几何。 大多数流形优化工作处于黎曼世界,距离由内积诱导、范数球是椭球。但神经网络不同:矩阵作为算子起作用,像谱范数这样的算子范数并非由内积自然产生。这意味着范数球可能有尖角,且不存在唯一的梯度流。在这个非黎曼世界里是否有更多可发掘之处?
- 工程实现。 要在规模上应用这些技术,需要在 GPU 上高效实现流形算子。近期的 Polar Express 在快速矩阵符号运算上显示了潜力。我们还需要哪些算法创新?
延伸阅读
- 
流形优化。 Absil、Mahony 与 Sepulchre 的教科书是标准参考。针对 Stiefel 流形,参见 Edelman 等, 1998。这些工作处于黎曼语境下。类似地,多数在 Stiefel 流形上做优化的机器学习论文也采取黎曼视角:例如 Li 等, 2020、Kong 等, 2022、Park 等, 2025。 
- 
机器学习中的“非黎曼”几何。 Thomas Flynn 2017 年的论文关于双对结构的梯度下降,将神经网络的权重空间刻画为“Finsler 流形”(即带有范数的流形),值得一读。另见苏剑林关于 Stiefel Muon 的近期博客以及 Franz Louis Cesista 关于在 Stiefel 流形上给出 Muon 启发式解法的博客。Franz 还写了后续博文,将此处方案推广。Scion 论文通过凸组合以另一种方式施加权重约束;Carlson 等, 2015较早讨论了(无约束的)谱下降。 
- 
Modula 项目。 Modula 项目的目标是构建一个库,能为通用架构自动“编译”最速下降优化器及相应的 Lipschitz 论断。参见项目页 https://modula.systems,以及我们关于模块化范数与模块化对偶的论文。我们的优化论文集也提供了进入这一思想空间的便捷路径。 
- 
Lipschitz 约束的深度学习。 相关工作很多,例如可参阅 Louis Béthune 与 Tsui-Wei Weng 的博士论文。通常这类工作并未把“权重量的 Lipschitz 性”与优化器设计联系起来。另见 Anil 等, 2018 以及我们的论文 Newhouse 等, 2025。 
引用
请按如下格式引用:
Jeremy Bernstein, "Modular Manifolds",
Thinking Machines Lab: Connectionism, Sep 2025.
或使用 BibTeX 引用:
@article{bernstein2025manifolds,
  author = {Jeremy Bernstein},
  title = {Modular Manifolds},
  journal = {Thinking Machines Lab: Connectionism},
  year = {2025},
  note = {https://thinkingmachines.ai/blog/modular-manifolds/},
  doi = {10.64434/tml.20250926}
}