Phases of Muon: When Muon Eclipses SignSGD¶
会议: NeurIPS 2026
arXiv: 2605.09552
OpenReview: https://openreview.net/forum?id=iXZjCLrCHW
领域: 优化/理论
关键词: Muon优化器, SignSVD, SignSGD, 随机矩阵理论, 学习率缩放律
一句话总结¶
本文通过高维矩阵值线性回归模型与随机矩阵理论推导了谱优化器(SignSVD/Muon)和对角优化器(SignSGD/Adam)的精确确定性风险动力学,揭示了大批次下的平方根谱预条件加速机制、临界批次效应以及幂律各向异性数据下的三阶段相图。
研究背景与动机¶
现代深度学习模型(尤其是大语言模型)的预训练几乎被以 Adam 为代表的对角预条件一阶自适应优化器所垄断。然而,神经网络的权重天然具备矩阵或张量结构,梯度中蕴含着层间反向传播带来的低秩外积相关性,对角预条件无法捕捉跨坐标的二阶几何关联。近年来,基于 Newton-Schulz 迭代近似矩阵梯度谱白化(spectral whitening)的 Muon 优化器在 LLM 预训练实践中展现出显著超越 Adam 的收敛效率与扩展潜力。然而,现有工作对 Muon 在高维参数空间下的学习动力学、临界批次规模(critical batch size)以及数据协方差结构的依赖机制仍缺乏第一性原理的理论刻画。
这种认知缺失的核心矛盾在于:既有理论要么退化至最坏情况凸/非凸收敛率界(如 \(\mathcal{O}(1/\sqrt{T})\)),无法体现不同自适应方法在真实训练轨迹中的结构性差异;要么仅聚焦于单步更新分析或全批次连续梯度流,无法捕捉随机小批次采样噪声、有限矩阵维度与数据各向异性谱之间的动态耦合。对于实践中 Muon 必须采用远大于 Adam 的学习率、以及 Muon 并非在所有任务上都稳定优于 Adam 的经验现象,学术界尚无统一的解析理论。
本文的切入角度是:构建一个保留反向传播外积梯度结构的矩阵值线性回归高维统计模型,利用随机矩阵理论(RMT)与自由概率论中的预解式自洽定点系统,严格推导随机 SignSVD(Muon 的理想白化极限)与随机 SignSGD(Adam 的无动量代理)的闭式确定性风险递归轨迹。核心 idea:通过高维随机矩阵预解式的均质化分析,将谱白化优化器的随机动力学严格解耦为漂移核与波动核,证明了大批次下 SignSVD 具备对数据协方差谱的平方根预条件加速,而在幂律各向异性谱下存在划分 SignSVD 与 SignSGD 相对优势边界的三阶段相图。
方法详解¶
整体框架¶
本文考虑矩阵值参数 \(W \in \mathbb{R}^{N_{\mathrm{out}} \times N_{\mathrm{in}}}\) 的线性回归任务,真实教师为 \(W^\star\),单样本输入特征为 \(x_{\mathrm{in}} \sim \mathcal{N}(0, \Sigma_{\mathrm{in}})\),上游反向传播误差对应特征为 \(x_{\mathrm{out}} \sim \mathcal{N}(0, \Sigma_{\mathrm{out}})\)。小批次规模为 \(B\) 时,参数偏差 \(\Delta_t = W_t - W^\star\) 产生的批梯度具有天然的外积结构 \(G_{t+1}(\Delta_t) = \frac{1}{B} \sum_{i=1}^B (x_{\mathrm{out}}^i \otimes x_{\mathrm{in}}^i) \langle x_{\mathrm{out}}^i \otimes x_{\mathrm{in}}^i, \Delta_t \rangle\)。谱优化器族通过对奇异值施加变换函数 \(\varphi(z)\) 实现更新,其中 SignSVD 对应理想白化映射 \(\varphi(z) = z^{-1/2}\)(即 \(\widetilde{G}_{t+1} = U_t V_t^\top\)),Muon 则是其 Newton-Schulz 多项式截断近似;而 SignSGD 则执行逐元素符号变换 \((\widehat{G}_{t+1})_{ij} = \mathrm{sign}((G_{t+1})_{ij})\)。
在高维比例极限(\(N_{\mathrm{in}}, N_{\mathrm{out}}, B \to \infty\) 且比值收敛)下,定义在两端特征基底上的投影风险 \(\mathscr{Q}_{ij}(t) = (u_i^\top \Delta_t v_j)^2\) 呈现自平均特性,其确定性均质化动力学 end-to-end 遵循由漂移核 \(\mathfrak{d}_{ij}\) 与波动核 \(\mathfrak{v}_{ij}\) 驱动的普适递归。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["外积结构批梯度<br/>G = 1/B ∑ (x_out ⊗ x_in)⟨x_out ⊗ x_in, Δ⟩"] --> B["预解式定点系统<br/>四方程自洽系统解 resolvent trace"]
B --> C["谱分解与核函数计算<br/>漂移核 d_ij 与波动核 v_ij"]
C --> D["投影风险普适递归<br/>Q_ij(t+1) = Q_ij(t) - 2η d_ij Q_ij/√R + η² v_ij"]
D -->|各向同性条件| E["等效收敛与学习率缩放<br/>收敛率一致,最优学习率相差 √N 倍"]
D -->|幂律各向异性| F["谱解析相变与三阶段相图<br/>临界批次 i# 与 (α, β) 平面竞争相"]
关键设计¶
1. 预解式定点系统:高维随机矩阵漂移与波动的闭式表征
针对高维矩阵梯度非线性奇异值截断难以解析的痛点,本文利用留一法与 Stein 引理(Gaussian integration by parts),推导了样本协方差矩阵与梯度外积预解式 \(R(z) = (G_{t+1} G_{t+1}^\top / (2\mathscr{R}) - z I_{N_{\mathrm{out}}})^{-1}\) 的确定性等价物。漂移核 \(\mathfrak{d}_{ij}\) 与波动核 \(\mathfrak{v}_{ij}\) 可完全由输入输出协方差特征值及四个自洽标量积分方程决定: $\(\sigma = \frac{1}{B} \sum_{i} \frac{\mu_i}{\tilde{s}_1 \mu_i - z}, \quad \tilde{\sigma} = \frac{1}{B} \sum_{j} \frac{\lambda_j}{s_1 \lambda_j - z}, \quad \tilde{s}_1 \sigma = s_1 \tilde{\sigma} = f\left(\frac{1}{\sqrt{-2z \sigma \tilde{\sigma}}}\right)\)$ 其中 \(f(\xi) = \mathbb{E}_{V \sim \mathcal{N}(0, 1)} [V^2 / (V^2 + 2\xi^2)]\)。对于 SignSVD,谱白化性质消除了积分对总体风险 \(\mathscr{R}(t)\) 的非线性依赖,使漂移核化简为与尺度无关的稳态形式,为全训练轨迹追踪提供了闭式解。
2. 临界批次效应:批规模充当谱白化分辨率刻度
在小批次非充分采样条件下,批大小 \(B\) 限制了梯度矩阵的秩(\(\mathrm{rank} \le \min\{N_{\mathrm{in}}, N_{\mathrm{out}}, B\}\)),导致谱优化器无法均匀白化所有特征子空间。在半各向异性幂律谱模型 \(\mu_i = i^{-\alpha}\)(其中 \(\alpha > 0\) 为谱衰减指数)下,本文严格导出了谱分辨率阈值 \(i_{\#}\): $\(i_{\#} \sim \begin{cases} B, & \alpha > 1 \\ B^{1/\alpha} / N^{1/\alpha - 1}, & \alpha < 1 \end{cases}\)$ 对于主导模态 \(i \le i_{\#}\),SignSVD 实现有效谱白化,漂移核呈现 \(\mathfrak{d}_i \propto \sqrt{\mu_i}\) 的平方根预条件加速行为;而对于尾部欠分辨模态 \(i > i_{\#}\),梯度噪声淹没子空间方向,动力学退化为标准一阶 SGD/SignSGD 的慢速线性漂移 \(\mathfrak{d}_i \propto \mu_i\)。当且仅当 \(B \ge N\) 时,全谱被完全解析,SignSVD 全面激活全局平方根预条件。
3. 三阶段竞争相图:各向异性下 Muon 与 SignSGD 的相变分离
针对业界关于“Muon 是否在所有情况下均优于 Adam”的争论,本文分析了达到预设目标误差 \(4\epsilon\) 所需的时间 \(t_{4\epsilon}\)(在最优常数学习率下,使稳态噪声基底恰为 \(\epsilon\))。目标初始权重对齐指数设为 \(\mathscr{Q}_i(0) = i^{-\beta}\),在 \((\alpha, \beta)\) 空间中揭示了三大本质不同的优化动力学相: - Phase A (\(\beta < 1\)):SignSVD 全面胜出。此时信号显著分布于尾部小特征值模态(高秩、偏差主导难题),SignSVD 的漂移核 \(\sqrt{\mu_i}\) 相比 SignSGD 的 \(\mu_i\) 在尾部提供巨大加速度,随着 \(\epsilon \to 0\),两者的收敛时间比呈现关于 \(1/\epsilon\) 的多项式级巨大发散优势。 - Phase B (\(1 < \beta < \alpha + 1\)):交叉混合相。大特征值主导前期损失下降,初期阶段 SignSGD 凭借大步长在大特征值方向取得暂时领先;但随着向低误差 \(\epsilon \to 0\) 推进,SignSVD 的谱平衡优势再次反超,两者轨迹发生 cross-over。 - Phase C (\(\beta > \alpha + 1\)):SignSGD 全面胜出。此时目标高度集中于头部极大特征值主成分,优化变为简单主成分拟合问题。两者的理论渐近时间收敛率幂次同为 \(\epsilon^{-1/2}\),但 SignSGD 的噪声壳层系数 \(\mathcal{S}\) 更小,以固定的常数因子在整条轨迹上始终占优。
损失函数 / 训练策略¶
在均方误差目标下,SignSVD 与 SignSGD 在各步均采用最优自适应步长策略以最大化每步下降量。在各向同性数据(\(\Sigma = I\))下,最小化单步确定性损失推导出的最优时变学习率为: $\(\eta_\star^{\mathrm{s\text{-}SVD}}(t) = \frac{2 C(N/B) \sqrt{\mathscr{R}(t)}}{\min\{B, N\}}, \quad \eta_\star^{\mathrm{s\text{-}SGD}}(t) = \frac{2 \mathcal{N}_B \sqrt{\mathscr{R}(t)}}{\sqrt{\pi} N^2}\)$ 其中 \(\mathcal{N}_B \sim \sqrt{B}\)。在 \(B \ge N\) 的大批次下,两者的最优学习率比值为 \(\eta_\star^{\mathrm{s\text{-}SVD}} / \eta_\star^{\mathrm{s\text{-}SGD}} \asymp \sqrt{N}\);而在 \(B \le N\) 的小批次下,该比值扩展为 \(N / \sqrt{B}\)。这从理论底层解释并严格校准了近期工业界在大模型训练中发现的 Muon 学习率 Moonshot 扩展规则。
实验关键数据¶
主实验¶
论文通过多组大规模高维数值模拟验证了确定性动力学方程对真实随机优化器训练轨迹的精确预测。在各向同性数据基准下,对比了 SignSVD 与 SignSGD 的最优学习率缩放及几何收敛因子:
| 评估维度 (\(N=B\)) | 理论预测表达 | \(N=128\) 实验验证 | \(N=256\) 实验验证 | \(N=512\) 实验验证 | 动力学吻合度 |
|---|---|---|---|---|---|
| SignSVD 最优学习率 \(\eta_\star / \sqrt{\mathscr{R}}\) | \(\frac{2 C(1)}{N} \approx \frac{1.77}{N}\) | 0.0138 | 0.0069 | 0.0035 | 确定性曲线穿过置信区间中轴 |
| 学习率比值 \(\eta_\star^{\mathrm{s\text{-}SVD}} / \eta_\star^{\mathrm{s\text{-}SGD}}\) | \(\asymp \sqrt{N}\) | \(11.3 \times\) | \(16.0 \times\) | \(22.6 \times\) | 严格服从 \(\sqrt{N}\) 增益律 |
| 单步收敛率 \(1 - \mathrm{rate}_{\mathrm{s\text{-}SVD}}\) | \(2 C(1)^2 / N\) | 0.0123 | 0.0061 | 0.0031 | 理论线与经验点完全重合 |
| 收敛率比值 \((1-\mathrm{rate}_{\mathrm{s\text{-}SVD}})/(1-\mathrm{rate}_{\mathrm{s\text{-}SGD}})\) | \(\pi C(1)^2 \approx 0.98 \approx 1\) | 0.982 | 0.980 | 0.985 | 各向同性下无本质加速 |
消融实验¶
在半各向异性幂律谱模型(\(\mu_i = i^{-\alpha}, \mathscr{Q}_i(0) = i^{-\beta}\))下,验证大批次(\(B \ge N\))时达到 \(4\epsilon\) 目标误差时间 \(t_{4\epsilon}\) 的渐近标度指数与相变界限:
| 配置区间 / 相区 | 典型参数 \((\alpha, \beta)\) | SignSVD 时间标度 \(t_{4\epsilon}^{\mathrm{s\text{-}SVD}}\) | SignSGD 时间标度 \(t_{4\epsilon}^{\mathrm{s\text{-}SGD}}\) | 优胜优化器与分离形态 |
|---|---|---|---|---|
| 各向同性基准 (\(\Sigma = I\)) | \(\alpha=0, \beta=0\) | \(\mathcal{O}(\epsilon^{-1/2})\) | \(\mathcal{O}(\epsilon^{-1/2})\) | 平手(相差常数比例因子) |
| Phase A (高秩偏差主导) | \(\alpha=1.5, \beta=0.7\) | \(\mathcal{O}(\epsilon^{-0.625})\) | \(\mathcal{O}(\epsilon^{-1.250})\) | SignSVD 绝对优势(随 \(\epsilon \to 0\) 差距发散) |
| Phase B (混合中间相) | \(\alpha=1.5, \beta=1.5\) | \(\mathcal{O}(\epsilon^{-0.500})\) | \(\mathcal{O}(\epsilon^{-0.750})\) | 交叉翻转(大 \(\epsilon\) 偏 SignSGD,小 \(\epsilon\) 偏 SignSVD) |
| Phase C (大特征值主导) | \(\alpha=1.5, \beta=3.0\) | \(\mathcal{O}(\epsilon^{-0.500})\) | \(\mathcal{O}(\epsilon^{-0.500})\) | SignSGD 稳定占优(噪声壳系数 \(\mathcal{S}_{\mathrm{s\text{-}SGD}} < \mathcal{S}_{\mathrm{s\text{-}SVD}}\)) |
关键发现¶
- 真实 Muon 对 SignSVD 理论的定量贴合:采用 5 轮 Newton-Schulz 迭代的随机 Muon 实验曲线,与 SignSVD 的理论均质化标度线呈现极高契合度,表明截断多项式对理想极分解具有高度保真性。
- 动量机制保持核心斜率:在引入重球动量 \(\beta_{\mathrm{mom}} \in (0, 1)\) 后,有效噪声常数以 \(\sqrt{(1+\beta_{\mathrm{mom}})/(1-\beta_{\mathrm{mom}})}\) 缩放,但学习轨迹在各相区内的本征幂律衰减阶数依然完好存续。
- 小批次平坦化:在 \(B \le N\) 时,SignSVD 的最优学习率几乎不随批大小 \(B\) 剧烈波动,这使得小批次下 Muon 的超参数调优显著比 Adam(强烈依赖 \(B\))更加平稳友好。
亮点与洞察¶
- 将随机矩阵均质化扩展至矩阵优化算法:突破了过去仅能处理向量参数标量损失 SGD 的局限,建立了首个刻画矩阵值反向传播外积梯度的自洽预解式定点体系。
- 澄清了 Muon 学习率 \(\sqrt{N}\) 放大律的数理本质:严格从理论上证实了工业界经验法则,揭示了谱范数步长归一化与坐标级符号方差之间的物理维数差异。
- 终结了“非对角预条件必定优于对角优化器”的片面认知:证明了谱优化的有效性强依赖于数据各向异性与目标投影几何,在主成分主导(Phase C)问题上盲目套用 Muon 反而可能落后于 Adam。
局限与展望¶
- 未建模非线性特征学习与稳定边缘:该理论基于固定协方差的矩阵线性回归体系,无法显式刻画深度网络多层特征表征演变以及 Edge-of-Stability 等动态相变。
- 各向同性输入假设:论文主体的闭式相图分析聚焦于半各向异性情况(\(\Sigma_{\mathrm{in}} = I\)),双端同时呈现通用非对齐各向异性时的解析显式表征仍有待深化。
- 拓展方向:将分析框架推向随机特征模型(RFM)与两层非线性网络,进而推导考虑计算量约束(FLOPs-optimal)的谱优化神经缩放律。
相关工作与启发¶
- vs SignSGD / Adam (Bernstein et al., 2018; Xiao et al., 2025): SignSGD 在坐标基底上执行逐元素非线性符号运算,仅具有坐标对角重缩放能力;而 SignSVD 在正交奇异基上操作,天然具备旋转不变性与谱空间非对角预条件能力。
- vs Shampoo / K-FAC (Martens & Grosse, 2015; Gupta et al., 2018): Shampoo 需要周期性显式计算并求逆 Kronecker 协方差张量;Muon/SignSVD 则直接通过 Newton-Schulz 迭代对单步梯度矩阵做隐式极分解,规避了存储与高阶张量求逆开销。
- vs 最坏情况非凸收敛率界 (Chen et al., 2026; Shen et al., 2025): 传统最坏情况分析得出 Muon 仅有类似一阶 SGD 的 \(\mathcal{O}(1/\sqrt{T})\) 界,遮蔽了预条件优势;本文的高维平均场 RMT 理论精准捕捉到了条件数层面的谱加速实质。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 建立了针对谱优化器的高维随机矩阵理论均质化框架,具有坚实的数学创新。
- 实验充分度: ⭐⭐⭐⭐ 理论解析与受控高维数值模拟结合严密,但未涵盖真实大规模 LLM 端到端下游评测。
- 写作质量: ⭐⭐⭐⭐⭐ 结构清晰流畅,理论体系推导严密且物理图像透彻。
- 价值: ⭐⭐⭐⭐⭐ 为新一代基础模型优化器(Muon)的调参、批次扩展与适用边界提供了关键理论基石。