On the Width Scaling of Neural Optimizers Under Matrix Operator Norms I: Row/Column Normalization and Hyperparameter Transfer

该论文通过引入具有层可组合性的均值归一化算子范数，将 AdamW 和 Muon 等优化器统一为矩阵算子范数下的最速下降法，从而提出了能实现宽度无关平滑度保证及跨宽度超参数迁移的 MOGA 优化器，并在 GPT-2 和 LLaMA 的大规模预训练中展现出比 Muon 更优的效率与稳定性。

要点

这篇论文探讨了一个深度学习中的核心难题：当我们把神经网络的“大脑”（即网络宽度）变大时，如何调整它的“学习步长”（学习率），才能让它既学得快，又不走火入魔？

想象一下，你正在教一群学生（神经网络）学习。

• 窄网络：就像只有 10 个学生的小班。
• 宽网络：就像有 1000 个学生的大班。

目前的困境是：在小班里行之有效的教学节奏（学习率），一旦直接套用到大班，要么学生听得太快跟不上了（发散/崩溃），要么听得慢吞吞效率极低。这篇论文就是为了解决这个“大班教学”的难题。

以下是用通俗语言和比喻对这篇论文核心内容的解读：

1. 核心问题：为什么“大班”难教？

在数学上，这被称为**“宽度缩放”（Width Scaling）**问题。
现有的优化器（如 AdamW 或 Muon）就像不同的“教学策略”。

• 旧策略的缺陷：以前的策略在计算“步长”时，没有考虑到学生人数的变化。当学生人数（网络宽度 $w$）增加时，整个班级的“噪音”和“混乱度”会非线性地增加。这就好比在 10 人的教室里大声说话很合适，但在 1000 人的礼堂里，同样的音量会让所有人听不清，甚至引发骚乱。
• 后果：如果你把小模型调好的学习率直接用到大模型上，大模型往往会训练失败。

2. 新视角：把优化看作“几何导航”

作者提出了一种全新的视角：把优化过程看作是在一个复杂的“地形”上寻找最低点（损失最小值）。

• 地形（Loss Landscape）：想象成一片起伏的山地，我们要找谷底。
• 优化器：就是向导。
• 步长（学习率）：向导决定迈多大一步。

作者发现，很多流行的优化器（如 AdamW, Muon）其实都是在特定的“几何规则”下，寻找下降最快的方向（最陡下降法）。

• Muon：像是一个拿着精密罗盘的向导，试图在二维平面上找最陡的路（基于谱范数/2 范数）。
• AdamW：像是一个更随性的向导，主要看每个方向上的符号（基于无穷范数）。

关键发现：
作者发现，Muon 这种“精密罗盘”策略，当学生人数（宽度）增加时，地形的**“崎岖程度”（平滑度常数）**会急剧上升（大约增加 $\sqrt{w}$ 倍）。这意味着，人越多，路越陡、越乱，向导必须把步子迈得极小才能不摔跟头。这就是为什么 Muon 在大模型上很难调参的原因。

3. 解决方案：引入“平均化”的新地图

为了解决这个问题，作者发明了一种新的**“平均化”几何规则**（Mean-Normalized Operator Norms）。

通俗比喻：
想象你在测量一群人的身高。

• 旧方法（标准范数）：直接看最高那个人有多高。如果人多了，最高的人可能更高，导致测量结果随人数膨胀。
• 新方法（平均化范数）：看平均身高。无论人多少，平均身高保持在一个稳定的范围内。

作者通过这种“平均化”处理，重新定义了地形的测量方式。

• 效果：在这种新地图下，无论学生人数（网络宽度）怎么变，地形的**“崎岖程度”（平滑度）**都保持不变！
• 结论：既然地形不再随人数变乱，那么**“步长”（学习率）也就无需随人数调整了**。

4. 新优化器：MOGA

基于这个理论，作者提出了一个新的优化器家族，叫 MOGA（Matrix Operator Geometry Aware，矩阵算子几何感知）。

MOGA 的核心操作非常简单，就像给每个学生的“动作”做**“行归一化”**（Row Normalization）：

• 它不关心整个矩阵的复杂结构，而是简单地让每一行（代表一个神经元或特征）的更新幅度保持平衡。
• 行归一化（Row Normalization）：就像给每个小组分配任务时，确保每个小组的总工作量是固定的，不会因为小组人数多了就乱套。

MOGA 的两大优势：

1. 超参数迁移（Hyperparameter Transfer）：这是最厉害的一点。如果你在小模型（比如 1 亿参数）上调好了学习率，你可以直接把这个学习率用到大模型（比如 100 亿参数）上，不需要重新调参！就像给小班设计的课表，直接拿来给大班用，效果依然完美。
2. 训练效率：在大规模训练（特别是训练后期，损失值很低时），MOGA 比 Muon 和 AdamW 跑得更快、更稳。

5. 实验验证：真的好用吗？

作者在 GPT-2 和 LLaMA 这样的大语言模型上进行了测试：

• 学习率迁移：从 GPT-2 Small（小）到 GPT-XL（大），MOGA 的最佳学习率几乎完全不变。这证明了理论的正确性。
• 训练速度：在训练后期（低损失阶段），MOGA 的表现优于 Muon，收敛得更快。特别是在数据量很大（大 Token 预算）的情况下，MOGA 的优势更明显。

总结：这篇论文告诉我们什么？

1. 旧地图不行：以前用来指导大模型训练的“几何规则”（如 Muon 的谱范数）在人数（宽度）增加时会失效，导致路变陡、步长难调。
2. 新地图很稳：作者发明的“平均化”几何规则，让地形在任何规模下都保持“平坦”和“稳定”。
3. 一招鲜，吃遍天：基于新规则设计的 MOGA 优化器，实现了**“一次调参，通吃大小模型”**。你不需要为不同大小的模型重新调学习率，这大大降低了训练大模型的门槛和成本。
4. 行归一化是王道：在 MOGA 家族中，**行归一化（Row Normalization）**表现最好，它在保持训练稳定的同时，没有过度限制模型的学习能力。

一句话总结：
这篇论文就像给深度学习界提供了一套**“万能课表”**。以前，老师（优化器）需要根据班级人数（模型宽度）不断调整讲课速度（学习率）；现在，有了 MOGA 这套新规则，无论班级是 10 人还是 1000 人，老师都可以用同样的节奏讲课，而且效果一样好，甚至在大班后期跑得更快。

技术摘要

这篇论文《On the Width Scaling of Neural Optimizers Under Matrix Operator Norms I: Row/Column Normalization and Hyperparameter Transfer》深入探讨了神经网络优化器在模型宽度（Width）增加时的行为稳定性问题，并提出了一种基于矩阵算子范数几何的新优化框架。

以下是该论文的详细技术总结：

1. 研究背景与问题 (Problem)

• 核心问题：现代深度学习中的“缩放定律”表明模型性能随规模提升，但优化超参数（特别是学习率）通常高度依赖于模型宽度。例如，为 512 隐藏单元调优的学习率，在扩展到 2048 单元时往往会导致发散或收敛变慢。
• 现有局限：
  • 主流优化器（如 AdamW, Muon）缺乏跨宽度的超参数可迁移性。
  • 现有的理论分析（如 $\mu$P 参数化）主要基于特征学习行为或谱条件，且往往依赖于特定的初始化假设，难以推广到更广泛的优化器类别。
  • 传统的 $p \to q$ 矩阵算子范数在多层网络堆叠时，无法保持 Lipschitz 常数和光滑性常数（Smoothness Constant）的宽度无关性，导致优化几何结构随宽度失真。

2. 方法论 (Methodology)

作者将多种神经网络优化器统一视为**矩阵算子范数下的最速下降（Steepest Descent）**实例，并引入了新的几何视角：

• 统一框架：将 SignSGD, AdamW, Muon 等算法解释为在不同矩阵算子范数（如 $\ell_1 \to \ell_\infty$, $\ell_2 \to \ell_2$）下的最速下降方向。
• 引入均值归一化范数 (Mean-Normalized Norms)：
  • 为了解决经典 $p \to q$ 范数在层间传递时因维度不匹配导致的稳定性破坏，作者定义了均值归一化范数：$\|\mathbf{x}\|_{(p, \text{mean})} = n^{-1/p} \|\mathbf{x}\|_p$。
  • 这种归一化消除了维度 $n$ 对范数大小的影响，使得相邻层之间的范数满足兼容性条件（$\|\mathbf{I}\| \le 1$），从而保证了 Lipschitz 常数在深度网络中的宽度无关性。
• 光滑性分析 (Smoothness Analysis)：
  • 分析了不同几何下的 $L$-光滑性（梯度 Lipschitz 连续性）。
  • 发现经典 Muon 对应的 $(2, \text{mean}) \to (2, \text{mean})$ 几何下，光滑性常数随宽度 $w$ 以 $O(\sqrt{w})$ 增长，导致优化不稳定。
  • 证明了在 $(1, \text{mean}) \to (q, \text{mean})$ ($q \ge 2$) 和 $(p, \text{mean}) \to \infty$ 几何下，光滑性常数与宽度无关。
• 提出 MOGA 优化器：
  • 基于上述几何分析，提出了 MOGA (Matrix Operator Geometry Aware) 优化器。
  • 核心机制：根据算子范数几何推导出的宽度感知学习率缩放规则（Width-aware scaling rules）。
  • 具体实现：包括行归一化 (Row Normalization) 和 列归一化 (Column Normalization) 的变体。例如，行归一化对应 $(p, \text{mean}) \to \infty$ 几何，通过每行梯度的归一化并乘以特定的宽度缩放因子来更新参数。

3. 关键贡献 (Key Contributions)

1. 理论统一：将 SignSGD, AdamW, Muon, GradPower 等主流优化器统一在矩阵算子范数最速下降的框架下。
2. 几何洞察：
   • 揭示了经典 $p \to q$ 范数在深度网络中无法保持宽度无关稳定性的根本原因（层间几何不匹配）。
   • 证明了均值归一化范数是实现宽度无关 Lipschitz 和光滑性控制的关键。
   • 指出 Muon 在 $(2, \text{mean}) \to (2, \text{mean})$ 几何下存在 $O(\sqrt{w})$ 的光滑性常数增长，这是其在大宽度下可能不稳定的理论根源。
3. 新优化器 MOGA：
   • 提出了一种仅依赖行/列归一化的宽度感知优化器。
   • 推导了通用的学习率缩放规则，使得在 $p \ge 2$ 的行归一化下，学习率在不同宽度模型间可直接迁移。
   • 证明了 MOGA 在特定情况下（如 Adam/SignSGD）能精确恢复 $\mu$P 缩放规则，但其理论基础（优化几何视角）比 $\mu$P 更广泛，不依赖谱条件。
4. 权衡分析：深入探讨了范数选择带来的“优化 - 近似”权衡（Optimization-Approximation Trade-off）。行归一化在保持光滑性（利于优化）的同时，对参数空间的约束比列归一化更宽松，从而保留了更好的模型表达能力。

4. 实验结果 (Results)

作者在 GPT-2 和 LLaMA 架构上进行了大规模预训练实验：

• 学习率迁移 (Learning Rate Transfer)：
  • 实验显示，使用 MOGA（特别是行归一化版本）时，从 GPT-2 Small 到 XL（参数量从 1.2 亿到 15 亿），最优峰值学习率几乎保持不变。
  • 这验证了理论预测：在均值归一化几何下，优化动态是宽度无关的。
• 标准 Token 预算 (Standard Token Budget)：
  • 在 1 倍 Chinchilla 最优 Token 数下，MOGA（行归一化）在 LLaMA-130M 上表现与 Muon 相当，且远快于 AdamW；在 GPT-2 Small 上表现介于两者之间。
• 大 Token 预算 (Large Token Budget)：
  • 在 8 倍 Chinchilla 最优 Token 数下（模拟大规模训练后期），MOGA（行归一化）展现出显著优势。
  • 在低损失（Low-loss）阶段，MOGA 的收敛速度明显快于 Muon 和 AdamW，且损失曲线下降更陡峭，表明其在长训练周期和稳定优化方面更具优势。

5. 意义与影响 (Significance)

• 解决超参数迁移难题：提供了一种原理性的机制，使得在小模型上调优的学习率可以直接用于大模型，大幅降低了大规模模型训练的超参数搜索成本。
• 超越 $\mu$P 的视角：虽然 MOGA 在特定情况下复现了 $\mu$P 的缩放规则，但其基于“优化几何”和“光滑性控制”的理论基础更为通用，适用于不满足 $\mu$P 谱条件的优化器。
• 提升训练效率：实验表明，基于行归一化的 MOGA 优化器在大模型、长训练周期的实际应用场景中，比当前 SOTA 的 Muon 优化器更高效、更稳定。
• 理论指导实践：为设计下一代可扩展的神经网络优化器提供了明确的几何设计原则（即选择能保持宽度无关光滑性的算子范数几何）。

总结：该论文通过引入均值归一化的矩阵算子范数几何，从理论上解决了神经网络优化器随宽度扩展时的不稳定性问题，并据此提出了 MOGA 优化器。实验证明，MOGA 不仅实现了可靠的跨宽度学习率迁移，还在大规模预训练任务中展现了优于现有优化器的性能，特别是在训练后期和低损失区域。