模块化流形 —— Thinking Machines 解读

模块化流形封面 — 原文封面：把不同层的权重各自约束在不同的流形（曲面）上，再像积木一样组合成整个网络。

一句话抓住重点

我们早就会归一化激活（LayerNorm）和梯度更新（Muon），却很少归一化权重矩阵本身。这篇主张把权重也"管起来"——锁在一个流形（曲面）上并与优化器协同设计，让"学习率"真正等于"这一步走多远"。具体例子是 Manifold Muon（权重锁在奇异值全为 1 的 Stiefel 流形），CIFAR-10 上精度超过 AdamW。再用模块化流形这套积木规则，自动跨层分配学习率预算。

激活✓ 梯度✓
权重✗→✓

这篇补上"约束权重"
这一块缺口

$\sigma\approx1$

训练后权重奇异值
全部收在 1 附近

> AdamW

CIFAR-10 训练 & 测试
精度均超 AdamW

< 1 分钟

每次实验运行
（小 MLP）

01 · 动机

为什么要"管住"权重

训练大网络时，我们不希望张量（权重、激活、梯度）变得太大或太小——否则不只是数值上溢/下溢，权重矩阵在训练中"忽大忽小"还会让"更新相对大小"漂移，直接拖累学习速度。保持张量健康的黄金法则是归一化。但目前的覆盖并不完整：

✅

激活——已归一化

LayerNorm 早已是标配，把激活放在好的尺度上传给下一层

✅

梯度更新——已归一化

Muon 可理解为对更新做"谱归一化"

⬜

权重本身——很少归一化

这正是本文要补的缺口：把权重也约束起来

🎁

好处一箩筐

杜绝范数爆炸、小条件数、行为可预测、Lipschitz 鲁棒保证

02 · 核心直觉

在曲面上做梯度下降

用最简单的例子——把一个向量参数 $w\in\mathbb{R}^d$ 锁在"超球面" $S^n$ 上。

流形就是"凑近看会变平"的曲面，那块局部平面叫切空间。

SVD 奇异值分解 — **奇异值分解（SVD）：**秩为 $k$ 的矩阵 $M\in\mathbb{R}^{m\times n}$ 总能分解为 $M = U\Sigma V^\top$，其中 $U,V$ 列正交，$\Sigma$ 是只含正值的对角阵。$\Sigma$ 的对角元就是**奇异值**——衡量矩阵在各方向上的"拉伸倍数"。

💡 关键洞见：在切空间里走，而不是"走偏再硬拉回来"

与其用普通优化器走一步、再投影回球面，不如直接在切空间里找方向。这样学习率才真正等于"这一步的实际长度"——若总是走偏再投影，这个漂亮性质就没了。

把"走一步"写成带约束的优化问题

在超球面上，约束是 $a^\top w = 0$（更新方向 $a$ 要在切空间里）外加步长 $\|a\|_2=\eta$。要"在损失上下降最多"，就是解：

$$\min_{a\in\mathbb{R}^d} \quad \underbrace{a^\top g}_{\text{损失的线性变化}} \quad \text{s.t.} \quad \underbrace{\|a\|_2 = \eta}_{\text{步长约束}},\ \underbrace{a^\top w = 0}_{\text{切空间约束}} \tag{$\star$}$$

用拉格朗日乘子法求解，引入 $\lambda,\mu$：$\mathcal{L}(a,\lambda,\mu)=a^\top g+\tfrac{\lambda}{2}(a^\top a-\eta)+\mu(a^\top w)$

解出来的最优更新方向非常直观——从梯度里减掉"径向分量"，再归一化：

$$a_\mathrm{opt} = -\,\eta \times \frac{g - ww^\top g}{\|g - ww^\top g\|_2}$$

最后用"回拉映射（retraction）"把权重拉回流形：$w \gets \dfrac{1}{\sqrt{1+\eta^2}}\left[w - \eta\times\dfrac{g-ww^\top g}{\|g-ww^\top g\|_2}\right]$

距离怎么量，决定了往哪走

一个反直觉的点：你用什么距离度量（欧氏 $\ell_2$ 还是曼哈顿 $\ell_1$……）会改变"最佳前进方向"。若用非欧距离，固定步长下不沿梯度方向走反而能在梯度方向走得更远。

乘积流形 — **笛卡尔积粘合流形：**一条线 × 一个圆盘 = 一个圆柱。线上每一点都"长出"一份圆盘的副本——这正是后面"模块化流形"组合多层的几何基础。

三步走，万变不离其宗

STEP 1

找方向

在切空间里找单位长度、最沿梯度方向的向量

→

STEP 2

走一步

乘以学习率 $\eta$，从权重里减掉

→

STEP 3

回拉

用回拉映射把权重拉回流形

只要换"用哪个流形"和"用哪种范数量长度"，就能生成不同优化器：

流形	范数	得到的优化器
欧氏空间 $\mathbb{R}^n$	欧氏范数	普通梯度下降
欧氏空间 $\mathbb{R}^n$	无穷范数	符号梯度下降
超球面 $S^n$	欧氏范数	超球面下降
矩阵空间 $\mathbb{R}^{m\times n}$	谱范数	Muon
Stiefel 流形	谱范数	Manifold Muon ★ 本文新成果

03 · 主菜

Manifold Muon：把权重锁在"奇异值都=1"的流形上

Transformer 里权重矩阵 $W$ 本质是"向量变换器"（$y=Wx$）。我们希望它既不把输入放太大也不缩太小——也就是让所有奇异值都等于 1。这样的矩阵集合正是 Stiefel 流形：

$$\mathsf{Stiefel}(m,n) := \left\{\, W \in \mathbb{R}^{m\times n} \mid W^\top W = I_n \,\right\}$$

所有奇异值 $\sigma_1=\sigma_2=\dots=1$ —— 矩阵只旋转、不缩放

配上谱范数（衡量最大奇异值）作距离，把"在 Stiefel 流形上沿梯度下降最多"写成带约束优化：

$$\min_{A\in\mathbb{R}^{m\times n}} \quad \underbrace{\operatorname{trace}(G^\top A)}_{\text{损失的线性变化}} \quad \text{s.t.} \quad \underbrace{\|A\|_{\text{spectral}} \leq \eta}_{\text{步长约束}},\ A^\top W + W^\top A = 0 \tag{$\dagger$}$$

⚡ 这是一道"凸优化"题，用对偶上升解

作者的关键洞见：$(\dagger)$ 是凸优化问题，可用经典的对偶上升（dual ascent）求解。引入拉格朗日乘子矩阵 $\Lambda\in\mathbb{R}^{n\times n}$，把"带约束最小化"转成"无约束最大化"，核心运算是矩阵符号函数 $\operatorname{msign}$（把奇异值都"啪"地拍到 1，可在 GPU 上用 Newton-Schulz 或 Polar Express 高效算）。

给定 $\Lambda$，内层最优解有闭式；外层对 $\Lambda$ 做梯度上升直到收敛：

$$A_\mathrm{opt}(\Lambda) = -\,\eta \times \operatorname{msign}\!\big(G + 2W(\Lambda+\Lambda^\top)\big)$$

迭代 $\Lambda \gets \Lambda + \alpha\times H(\Lambda)$ 收敛后，更新权重 $W \gets W + A_\mathrm{opt}$，并 $W\gets\operatorname{msign}(W)$ 拉回流形

🤝 一个漂亮的开放协作案例

作者去年解决了方阵情形但卡在矩形情形，于是把它当作公开难题贴在 Modula 文档上。苏剑林（Jianlin Su）今夏用拉格朗日 + 不动点迭代解出，作者又结合 Franz Louis Cesista 的相关工作整理出这里的对偶上升算法。

小实验：CIFAR-10 上超过 AdamW

CIFAR-10 实验结果 — **CIFAR-10 上训练小 MLP 3 个 epoch（3 个随机种子平均）：**浅蓝细线是不同权重衰减设置的 AdamW；Manifold Muon 取得**更高的训练与测试精度**。第三张图是最佳学习率下第一层权重的奇异值分布——训练后**全部贴近 1**。唯一代价是每步 wall-clock 稍慢（对偶上升开销，可通过减少迭代步数或加动量缓解）。

04 · 更大的图景

模块化流形：把单层"积木式"拼成网络

单层会了，那把很多层拼成网络呢？要不要担心层间互相干扰？

答案是可以"积木式"组合。把任何模块（从一层到整个 Transformer）看成带三个属性的数学对象：前向函数 $f:\mathcal{W}\times\mathcal{X}\to\mathcal{Y}$、流形约束 $\mathcal{M}\subset\mathcal{W}$、范数 $\|\cdot\|:\mathcal{W}\to\mathbb{R}$。例如一个"Stiefel 线性层"：

$$\mathsf{StiefelLinear} = \begin{cases}(W,x)\mapsto Wx, & \text{(前向函数)}\\ \mathsf{Stiefel}(m,n), & \text{(流形)}\\ \|\cdot\|_\mathrm{spectral}. & \text{(范数)}\end{cases}$$

组合两个模块时，三个属性各有简单拼接规则：

$$f_3((w_1,w_2),x) := f_2(w_2,\,f_1(w_1,x)) \qquad \text{（前向：复合）}$$ $$\mathcal{M}_3 = \mathcal{M}_1 \times \mathcal{M}_2 \qquad \text{（流形：笛卡尔积）}$$ $$\|(w_1,w_2)\|_3 := \max\big(s_1\cdot\|w_1\|_1,\; s_2\cdot\|w_2\|_2\big) \qquad \text{（范数：加权 max）}$$

那两个权重 $s_1, s_2$ 正好是跨层分配学习率预算的旋钮

💡 为什么这套抽象成立

它建立在作者团队"模块化范数（modular norm）"论文的核心观察上：跨层、跨规模分配学习率，本质等价于理解"网络输出对权重的 Lipschitz 敏感度"。流形约束让这个敏感度能估得更紧——例如 Stiefel 线性层天然满足 $\|(W+\Delta W)x - Wx\|_2 \leq \|\Delta W\|_\mathrm{spectral}\times\|x\|_2$。

05 · 对你的框架

💡 对 trainer 的落地启示

这篇偏研究方向（作者明确说是"抛砖引玉、招募社区一起做"），但对做 MoE 优化框架的人有几条信号：

⚙️

Muon 系优化器的延伸

框架若已支持/计划支持 Muon，这条"权重约束+流形"路线是同脉络的自然演进

🚀

关注 msign 的 GPU 实现

Newton-Schulz / Polar Express 是把这套理论落到大规模训练的关键算子

📊

跨层学习率预算

"组合即得 Lipschitz 界"的原则化方法，对层数多、规模大的 MoE 调参有参考价值

🔬

作者点名的开放问题

注意力头该约束在什么流形？embedding 该不该不同约束？对低精度训练数值有何影响？

🔥 联网补充：作者与 Muon 的渊源

Jeremy Bernstein 是 Muon 优化器思想脉络的关键研究者之一（Muon 由 Keller Jordan 等推广，近期在大模型预训练里因效率受关注）。配套的 Modula 项目（modula.systems）正把这套理论做成可编程库——自动为通用架构编译最速下降优化器及其 Lipschitz 保证。