🤖 machine learning

Muown: Row-Norm Control for Muon Optimization

本文介绍了 Muown，一种即插即用的优化器，它通过显式控制行幅度向量来防止 Muon 中的谱范数漂移，从而提升训练稳定性、降低对权重衰减的敏感性，并在各种模型规模下实现更优的困惑度。

原作者： Kai Lion, Florian Hübler, Bingcong Li, Antonio Orvieto, Niao He

发布于 2026-05-12

📖 1 分钟阅读☕ 轻松阅读

原作者： Kai Lion, Florian Hübler, Bingcong Li, Antonio Orvieto, Niao He

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在训练一个巨大的数字大脑（大型语言模型），让它写故事、解决数学问题或与你聊天。为此，这个大脑必须通过调整其电路内部数百万个微小的旋钮（权重）来学习。

很长一段时间以来，调整这些旋钮的最佳方法是一种名为 AdamW 的技术。最近，一种名为 Muon 的更快新方法出现了。这就像从自行车换成了跑车；它学得更快、更好。但有一个问题：这辆跑车有失控加速的倾向。

问题所在：“漂移”的引擎

该论文指出了 Muon 的一个具体问题。随着模型训练的进行，其内部连接的“大小”（具体而言，是权重矩阵中行的强度）开始不受控制地增长。

想象一下正在被充气的 balloon。

Muon 非常擅长判断往哪个方向吹气球（方向）。
然而，它不断泵入过多的空气，导致气球越变越大，直到可能爆裂（数值错误）或变得不稳定。
标准的解决方法是用一根绳子系住气球，每当它变得太大时就拉紧（权重衰减）。但作者发现，这就像用大锤砸坚果：它拖慢了学习过程，因为它挤压了整个气球，包括那些原本正常的部分。

诊断：问题的两个部分

由 Kai Lion 领导的作者们决定深入气球内部，看看究竟是什么在增长。他们意识到，连接的“大小”由两个截然不同的部分组成：

幅度（Magnitude，即体积）： 这行数字整体的大小。
相干性（Coherence，即形状）： 该行内部数字彼此之间的排列方式。

他们发现，只有幅度（Magnitude）在失控漂移。形状（相干性）实际上表现完美，保持稳定。气球并没有改变形状，只是变大了。

解决方案：Muown（智能泵）

作者们创造了一种名为 Muown 的新优化器。

Muown 不再将整个连接视为一个巨大的整体，而是将工作分为两个独立的任务：

方向团队（Muon）： 这个团队继续做 Muon 最擅长的事：确定更新权重的完美方向。他们保持这部分不变。
体积团队（新）： 这是新的部分。他们提取“幅度”（体积），将其视为一个独立的显式变量。他们为这个体积设定了一套特殊的规则（一种特定的数学几何结构，称为 $\ell_\infty$ ），以防止其增长过快。

类比：
想象你正在驾驶一辆汽车，其中方向盘和油门踏板被卡在一起。如果你转动方向盘，油门踏板也会被推动，导致汽车失控加速。

Muon 就像那辆被卡住的汽车。
Muown 就像将方向盘与油门踏板解耦。你仍然可以完美地转向（使用 Muon 的逻辑），但现在你为油门踏板（幅度）配备了一个独立的、智能的巡航控制系统，它能在不拖慢速度的情况下，将速度精确保持在所需位置。

结果

该论文在从较小（1.24 亿参数）到非常大（27 亿参数）的模型上测试了这种方法。以下是他们的发现：

更好的性能： Muown 始终产生比 Muon、AdamW 和其他竞争对手更好的结果（更低的“困惑度”，即衡量模型困惑程度的指标）。
更具容错性： 使用 Muon 时，你必须非常小心地设置参数（学习率）。如果选错了，模型就会失败。Muown 则稳健得多；它在广泛设置下都能表现良好，就像一辆拥有更宽“最佳驾驶区间”的汽车。
无额外成本： 通常，添加新的控制系统会拖慢汽车。但由于作者将这种新的“音量控制”直接集成到引擎现有的工作流程中，它几乎为训练过程增加了额外的时间。
稳定性： “气球”不再漂移。谱范数（连接的大小）保持稳定，防止了原始 Muon 所遭受的数值错误。

总结

该论文认为，流行的 Muon 优化器中的不稳定性并非其方向定位能力的缺陷，而是缺乏对其更新“音量”的控制。通过将“音量”与“方向”分离，并使用不同的专用工具对它们进行控制，Muown 在无需像权重衰减那样生硬修复的情况下，保持了训练的快速与稳定。它是一个即插即用的替代方案，使训练过程更平滑、更快、更可靠。

技术摘要：Muown - 用于 Muon 优化的行范数控制

问题陈述

Muon 已成为大型语言模型（LLM）预训练中有竞争力的优化器，它通过利用 Transformer 线性层的矩阵结构，往往能超越 AdamW。然而，最近的观察表明，Muon 在大规模应用中的行为对权重衰减（weight decay）敏感，且容易出现数值不稳定性，例如 bfloat16 溢出和注意力 logits 爆炸。

一个关键问题是，当使用 Muon 而未采用解耦权重衰减时，权重矩阵的谱范数（ $\|W_t\|_{S_\infty}$ ）在训练过程中会持续向上漂移。这种漂移违反了稳定特征学习所需的谱条件（ $\|W\|_{S_\infty} = \Theta(\sqrt{m/n})$ ）。虽然解耦权重衰减可以通过均匀缩小权重矩阵来缓解这一问题，但它是一种非针对性的干预措施，无论具体哪些方向驱动了增长，都会对整个谱进行缩放，从而可能减缓训练早期的收敛速度。本文旨在确定导致这种谱范数漂移的具体经验驱动因素，并提出一种更具针对性的控制机制。

方法：Muown

作者引入了Muown，作为 Muon 的即插即用替代品，它通过重参数化策略而非聚合约束来解决谱范数漂移问题。

1. 谱范数分解

核心洞察是将谱范数分解为两个因子（命题 1）：
$\|W_t\|_{S_\infty}^2 = \|g_t\|_\infty^2 \cdot \lambda_{\max}(P_t C_t P_t)$
其中：

$g_t = \|W_t\|_{row}$ 是行幅度向量（行的欧几里得范数）。
$D_t$ 是行归一化矩阵（ $W_t = \text{Diag}(g_t)D_t$ ）。
$C_t = D_t D_t^\top$ 是单位范数行的 Gram 矩阵。
$P_t$ 是归一化行幅度的对角矩阵。
$\lambda_{\max}(P_t C_t P_t)$ 代表行相干性因子。

在 5 亿参数模型上的实证分析显示，行幅度因子（ $\|g_t\|_\infty$ ）是谱范数漂移的主要驱动因素，其变化与 $\|W_t\|_{S_\infty}$ 保持同步。相比之下，行相干性因子在整个训练过程中表现良好且有界。

2. 重参数化与优化

基于这一诊断，Muown 将权重矩阵 $W$ 重参数化为：
$W(g, R) = \text{Diag}\left(\frac{g}{\|R\|_{row}}\right) R$
这里， $g \in \mathbb{R}^m$ 代表显式的行幅度向量， $R \in \mathbb{R}^{m \times n}$ 代表方向分量。

方向分量（ $R$ ）： Muon 以不变的方式应用于 $R$ ，将其置于谱范数几何（最陡谱下降）下处理。
幅度分量（ $g$ ）： 行幅度 $g$ 被提升为显式的优化器变量。它们在 $\ell_\infty$ 几何下更新，这是对角矩阵算子范数的自然对偶范数。
更新规则： $g$ 的更新利用 Adam 优化器（在 $\ell_\infty$ 几何下作为平滑的 SignSGD）来处理随机噪声，而 $R$ 的更新则遵循标准的 Muon 程序，包含 Newton-Schulz 正交化。

该方法在优化器状态中隐式实现，保持模型架构的前向传播不变。

主要贡献

谱漂移的诊断： 本文分解谱范数，确定最大行幅度（ $\|g_t\|_\infty$ ）是 Muon 下谱范数增长的经验驱动因素，将其与稳定的行相干性因子区分开来。
Muown 算法： 一种新的优化器，将行幅度视为显式的可训练变量，在 $\ell_\infty$ 几何下更新，同时保留 Muon 的方向更新。这避免了权重衰减的 indiscriminate（无差别）收缩以及谱重参数化的正交先验。
收敛性保证： 作者证明了 Muown 在确定性和随机设定下均能达到最优的非凸收敛速率。分析是在与更新几何对齐的混合对偶范数（幅度采用 $\ell_1$ ，方向采用 Schatten-1）下进行的。理论界限表明，与 Muon 相比，有效随机噪声系数有所降低。
实证验证： 在 FineWeb-Edu 数据集上，针对从 1.24 亿到 27 亿参数不等的模型规模进行的广泛实验，展示了持续的性能提升。

实验结果

困惑度提升： 在所有测试的模型规模（1.24 亿至 27 亿）和学习率下，Muown 始终比 Muon、SOAP、AdamW 和 Lion 实现更低的困惑度。
对超参数的鲁棒性：
- 学习率： Muown 拓宽了不同模型宽度下接近最优学习率的平台期，降低了对按规模调整学习率的敏感性。
- 权重衰减： Muown 减轻了对权重衰减调优的需求；它在无权重衰减（ $\lambda=0$ ）的情况下表现良好，而 Muon 需要仔细调优（最优 $\lambda=0.1$ ）以防止发散或性能次优。
稳定性： Muown 防止了 Muon 中观察到的系统性谱范数漂移。在高学习率设定下（例如 $4 \times 10^{-3}$ ），Muon 经常发散，而 Muown 保持稳定且接近最优。
梯度属性： 分析显示，与 Muon 相比，Muown 在整个训练过程中保持更高的小批量梯度归一化有效秩，并表现出更低的随机噪声系数。
计算开销： 该方法引入的步长时间开销可忽略不计。当优化器在分布式设置中跨秩分片时，Muown 和 Muon 的步长时间变得无法区分。内存开销极小（与行数成线性关系），每层仅增加几兆字节。

意义与主张

本文声称，Muon 中观察到的谱范数漂移并非优化器不可约简的属性，而是行幅度增长失控的结果。通过隔离这一驱动因素并将其提升为显式的优化器变量，Muown 实现了：

改进的收敛性： 更好的困惑度和稳定性，无需激进的权重衰减。
理论严谨性： 在显式考虑更新混合性质的几何感知框架中，实现了最优收敛速率。
实际效率： 作为即插即用替代品，在分布式设置中匹配 Muon 的运行时间，同时提供对超参数选择的更优越鲁棒性。

作者将 Muown 定位为一种解决方案，弥合了特征学习的理论稳定性要求与大规模 LLM 训练的实际需求之间的差距，提供了一种比解耦权重衰减更具针对性的替代方案。