From SGD to Muon: Adaptive Optimization via Schatten-p Norms¶

会议: NeurIPS 2026
arXiv: 2605.19781
领域: 优化/理论
关键词: Schatten-p范数, Muon优化器, 线性极小化神谕(LMO), 自适应几何, 随机特征回归

一句话总结¶

论文基于线性极小化神谕（LMO）理论将 SGDM 与 Muon 统一进 Schatten-\((p+1)\) 范数约束下降族，推导出利用一阶激活与动量统计闭式求解单层最优几何指数 \(p^*\) 的代理准则，并结合 \(p\) 依赖的二阶预条件与硬件高效 Taylor-Horner 极因子逼近，实现了端到端自适应谱优化。

研究背景与动机¶

深度神经网络优化长期由 Adam 及其变体等坐标级自适应算法主导。这类方法通过逐元素统计梯度的二阶矩来缩放更新，在处理非均一梯度尺度与随机噪声方面极具鲁棒性，但它们将权重矩阵扁平化为一维向量，完全忽视了现代 Transformer 等架构中占据主导地位的二维权重矩阵结构。近期涌现的 Muon 优化器直接针对二维矩阵几何，利用 Frobenius 范数下距离梯度最近的正交矩阵（即极因子）更新权重，通过拉平奇异值谱消除各向异性缩放，在 LLM 预训练中展现出惊人的收敛加速。

然而，Muon 这种将所有奇异值一律等同对待的纯几何正交化操作存在本质缺陷：它剥夺了方差追踪能力，在病态曲率或随机噪声主导的方向上容易产生剧烈扰动甚至放大噪声梯度。更关键的是，目前以 SGDM（\(p=1\)）和 Muon（\(p\to\infty\)）为代表的优化器均依赖人为先验指定的静态更新几何，而深度网络不同层（如注意力投影层与 MLP 层）在训练不同阶段的最佳下降几何形态截然不同，静态几何无法自适应匹配动态变化的损失曲面。

为了填补这一理论与实践鸿沟，本文将 SGDM 和 Muon 形式化为 Schatten 范数约束下的线性极小化神谕（LMO），并在单步随机特征回归框架下推导出仅依赖一阶矩与激活特征的单层最优 \(p^*\) 闭式解。本文的核心 idea 是：将网络各层的矩阵更新几何建模为连续的 Schatten-\((p+1)\) 范数 LMO 空间，通过一阶统计量动态闭式推导各层自适应最优 \(p^*\)，并配合自适应二阶方差预条件和平滑 Taylor 展开，实现跨层、跨训练周期的最优降维几何流转。

方法详解¶

整体框架¶

SMuon 算法将网络参数更新解耦为几何测度与步长执行两个核心维度。在给定的时间步，算法收集前一层的激活投影与当前层的一阶动量矩阵，利用标量线搜索在有界区间内闭式求解最大化局部收益的谱指数 \(p^*\)；为了防止小批量噪声引起几何振荡，谱统计特征被送入 EMA 缓冲区进行时序平滑；随后，基于求得的 \(p^*\) 构造与几何相匹配的 \(p\) 依赖二阶矩预条件子；最后，利用 Newton-Schulz 极因子结合 Taylor-Horner 级数多项式逼近硬件友好地输出分数极因子更新量。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入：当前层梯度 G 与输入激活 A"] --> B["动量累积与奇异系统分解<br/>计算对齐系数 C 与曲率系数 B"]
    B --> C["时序平滑最优几何指数求解<br/>EMA 滤波特征并线搜索求解 p*"]
    C --> D["p 依赖二阶矩缩放预条件<br/>D^-1/(2(p*+1)) 坐标重加权"]
    D --> E["Taylor-Horner 极因子逼近<br/>结合 Newton-Schulz 避免显式 SVD"]
    E --> F["输出：权重矩阵自适应几何更新 δW"]

关键设计¶

1. Schatten-\((p+1)\) 范数统一框架与闭式 \(p^*\) 代理准则：消除静态矩阵几何的次优性

现有优化器将 SGDM 与 Muon 视为孤立算法，而本设计证明了它们均属于统一的 Schatten 范数约束线性极小化神谕（LMO）极值点。在单步随机特征回归模型 \(\min_W \frac{1}{2k}\|WA - Y\|_F^2\) 下，参数摄动引起的目标函数变化严格展开为一阶梯度项与二阶激活曲率项。当更新方向设定为动量矩阵 \(M\) 的分数奇异值变换 \(\delta W = -\eta U_M \Sigma_M^{1/p} V_M^T\) 时，\(U_M\) 的正交性使得曲率项交叉项完全抵消。通过在最优学习率 \(\eta^*(p)\) 下解析求解最大损失下降，最优指数 \(p^*\) 简化为一维标量最大化问题：

\[p^* = \operatorname{argmax}_{p \in [1, \infty]} \frac{\left( \sum_{i} \sigma_{M, i}^{1/p} C_{i,i} \right)^2}{\sum_{i} \sigma_{M, i}^{2/p} B_{i,i}}\]

其中 \(C_{i,i} = (U_M^T G V_M)_{i,i}\) 度量未投影梯度与动量主成分之间的方向对齐度，\(B_{i,i} = \|A^T V_M e_i\|^2\) 衡量右奇异方向映射到输入激活上的能量。当对齐系数大且激活谱主方向平坦时，\(p^*\) 自然偏向较大的 Muon 类谱拉平模式；当梯度与动量存在正负抵消或曲率集中时，\(p^*\) 自动回缩到低阶 Euclidean 模式。

2. 谱统计量 EMA 时序平滑机制：平衡几何跃迁与小批量高频噪声

经验观察表明，若直接对每步求得的 \(p^*\) 进行时序平均，由于 \(p^*\) 由非线性分式比例求 \(\operatorname{argmax}\) 产生，均值与极值操作不可交换，会导致平滑后的估计量偏离真正的局部最优；此外，\(p^*\) 在真实训练中常表现出分段平稳、突发相变的特征。为此，设计将平滑点前移至底层谱统计量：

\[\tilde{C}_{i,i} \leftarrow \beta_p \tilde{C}_{i,i} + (1-\beta_p) C_{i,i}, \quad \tilde{B}_{i,i} \leftarrow \beta_p \tilde{B}_{i,i} + (1-\beta_p) B_{i,i}\]

在平滑后的统计量空间上求解 \(p^*\)，既允许底层特征随梯度平稳累积，又能保证在几何相变点到来时，\(p^*\) 能够敏锐且果断地“跃迁”至新的稳态，彻底规避了高频小批量噪声导致的几何在 \(1\) 到 \(50\) 之间杂乱跳动。

3. \(p\) 依赖的二阶矩对称预条件：弥合低 \(p^*\) 区域的坐标级方差脆弱性

当 \(p^*\) 趋近于 \(\infty\) 时，Muon 的矩阵等距正交性天然提供了全局尺度的各向同性约束；但当 \(p^*\) 降至接近 \(1\) 的区域时，模型退化为类似 SGDM 的行为，极易受坐标级梯度尺度悬殊与随机噪声的破坏。为了在全谱段提供均一的数值稳定性，引入基于梯度平方积的二阶矩累加量 \(D \leftarrow \beta_2 D + (1-\beta_2) G \odot G\)，并构造自适应缩放算子：

\[\delta W = -\eta D^{\circ -\frac{1}{2(p^*+1)}} \odot \operatorname{LMO}_{(p^*+1)} \left( D^{\circ -\frac{1}{2(p^*+1)}} \odot M \right)\]

指数项 \(\alpha = \frac{1}{2(p^*+1)}\) 是确保更新方向在损失函数全局缩放变换 \(L \to \lambda L\) 下严格保持尺度不变性的唯一解析解。在 \(p^*=1\) 时它退化为 \(\alpha=1/4\)，配合外层 Hadamard 积精确恢复 Adam 的标准 RMS 归一化；当 \(p^* \to \infty\) 时 \(\alpha \to 0\)，坐标预条件无缝隐退，精确回归纯 Muon 的正交谱更新。

4. 硬件友好的 Taylor-Horner 极因子逼近：免除 SVD 开销的算子级加速

在大规模分布式训练中，对所有矩阵参数执行精确 SVD 的计算复杂度虽然与矩阵乘法相当，但其顺序计算依赖严重阻碍了 Tensor Core 吞吐。利用 Newton-Schulz 迭代高效估计极因子 \(P = U V^T\) 与谱范数上界 \(\alpha \ge \|G\|_2\) 后，构造半正定矩阵 \(Z = \frac{1}{\alpha} G P^T\)，将分数极因子写为二项式 Taylor 展开：

\[U \Sigma^{1/p} V^T = \alpha^{1/p} \sum_{i=0}^K \binom{1/p}{i} \left( \frac{1}{\alpha} G P^T - I \right)^i P\]

结合 Horner 规则将矩阵多项式转化为纯 Fused Multiply-Add (FMA) 矩阵乘序列。相比直接用 Remez 多项式逼近 \(x^{1/p}\) 导致的低频原点发散（\(p \ge 10\) 时失效），该方案在较小的截断阶数 \(K\) 下即可在对数网格上达到远优于基线的精度，仅带来约 3% 的端到端运行开销。

损失函数 / 训练策略¶

SMuon 在学习率设置上采用双端对齐策略：当 \(p^* = p_{\max} = 50\) 时使用 Muon 的主学习率，当 \(p^* = 1\) 时在对数尺度上插值至辅助 AdamW 优化器学习率，确保经验 RMS 步长稳定。对于低频更新，在语言模型任务中设定每 100 步更新一次 \(p^*\)，\(\beta_p = 0.95\)；在视觉任务中每 1 个 epoch 刷新一次 \(p^*\)，线搜索范围严格截断在 \([1.02, 50.0]\)。

实验关键数据¶

主实验¶

论文在 Modded-NanoGPT 预训练（FineWeb 数据集，6 随机种子）、ImageNette 视觉分类（ViT-Small 与 MLP-Mixer-Small）以及 Qwen-2.5-0.5B-Instruct GSM8K LoRA 低秩微调场景下进行了全方位评估。

任务 / 模型	评测指标	SMuon (Adam) (本文)	SMuon (无预条件)	Muon (强基线)	AdamW 基线
NanoGPT (12L-768D, FineWeb)	验证集损失 (Val Loss)	3.366 ± 0.003	3.374 ± 0.004	3.376 ± 0.004	3.421 ± 0.005
GSM8K LoRA (Qwen-2.5-0.5B)	训练损失 (Train Loss)	0.340 ± 0.001	0.464 ± 0.003	0.387 ± 0.002	0.337 ± 0.002
GSM8K LoRA (Qwen-2.5-0.5B)	验证准确率 (Val Acc %)	29.37 ± 0.7%	28.18 ± 0.2%	28.74 ± 0.7%	29.67 ± 0.1%
ImageNette (MLP-Mixer, 100ep)	训练损失 (Train Loss)	0.0054 ± 0.0018	—	—	(固定p=3: 0.0061)
ImageNette (MLP-Mixer, 100ep)	验证准确率 (Val Acc %)	82.66 ± 0.30%	—	—	(固定p=3: 82.74)

消融实验¶

在 NanoGPT 速度测试中，论文深入消融了自适应几何本身的作用（表5），验证了自适应 \(p^*\) 相比于锁死固定几何 \(p=50\) 的收敛步数缩减。同时在 ImageNette 上消融了谱统计量 EMA 平滑系数与更新周期的鲁棒性（表7）。

实验配置	达到目标 Val Loss 步数	相对加速比	最终 Val Loss	说明
固定几何基准 (Fixed \(p = 50\))	4.00 loss: 367.3 步 / 3.40 loss: 1545.3 步	基准 (0.0%)	3.3754	缺乏早期几何适应性
自适应几何 SMuon (Adam)	4.00 loss: 358.2 步 / 3.40 loss: 1507.8 步	节省 2.43%~3.02%	3.3663	前期注意力层自动落入低 \(p^*\)
简化 Naive Optimal-C (无EMA/动量)	4.00 loss: 390.6 步 / 3.40 loss: 1558.0 步	滞后 8.3%	3.3780	几何剧烈抖动破坏收敛
更新步长消融 (Gap=50, \(\beta_p=0.0\))	最终训练损失: 0.0088	较差	—	高频更新无平滑导致震荡
更新步长消融 (Gap=50, \(\beta_p=0.9\))	最终训练损失: 0.0069	最优	—	强 EMA 滤除小批量谱噪声

关键发现¶

在 Transformer 预训练初期（前 100 步），注意力投影矩阵（\(W_q, W_k, W_v\)）的最优谱指数 \(p^*\) 会从 50 骤降至接近低阶区间，此时模型需要更偏向 Adam 的更新结构来拉齐各坐标梯度尺度；而 MLP 层的 \(p^*\) 则全程稳定维持在 \(p \approx 50\) 附近，展现出显著的跨层几何异质性。
在 LoRA 微调这一人造的“Muon 弱势场景”中，低秩分解（\(r=32\)）导致动量奇异谱极度退化，盲目施加 Muon 强行拉平谱分布会放大噪声方向，导致验证集准确率跌至 28.74%；而 SMuon 自动诊断出各层最优 \(p^* \approx 3.1\)，结合二阶预条件取得了 29.37% 的准确率，紧咬针对该场景最优的 AdamW（29.67%）。

亮点与洞察¶

将离散优化器统一为连续几何流：跳出了“Muon 好还是 Adam 好”的二元对立，通过 Schatten 范数 LMO 将经典 SGDM、Muon、Adam 与 MuAdam 严格组织在同一个二维参数几何坐标系中，使优化算法设计具备了连续可微的探索空间。
解析求解而非元学习搜索：利用单步随机特征回归模型直接给出 \(p^*\) 的闭式解，仅需在常规前向反向传播中提取已有激活与动量内积，避免了基于强化学习或元梯度的昂贵黑盒调参。
平滑统计量而非平滑极值点：深刻洞察到 \(\operatorname{argmax}\) 与期望不可交换的数学本质，通过对对齐系数和曲率特征施加 EMA，实现了既能滤除噪声又能捕获突发几何相变的优雅平衡。

局限与展望¶

单步贪心近似的理论局限：最优准则依赖单步二阶泰勒展开，忽略了参数更新在多步长视界下的累积效应和下游层的雅可比回传耦合；在某些特殊曲面上，局部的贪心最优可能会落入多步次优态。
激活特征缓存与通信开销：在张量并行或超大规模分布式训练中，跨卡聚合激活投影矩阵 \(A^T V_M\) 需要一次 \(O(\min(m,n))\) 的额外 All-Reduce 通信；虽然每 100 步一次的频次使开销仅占约 3%，但在极限通信瓶颈集群下仍需进一步工程优化。
展望：未来可探索多步曲率预测模型，或基于模型宽度与深度的 scaling laws 建立跨尺度的 \(p^*\) 先验分布，从而在超大模型上免除在线计算开销直接应用先验几何。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次从理论上建立 SGDM 到 Muon 的连续 Schatten 几何空间并导出闭式 \(p^*\) 解。
实验充分度: ⭐⭐⭐⭐☆ 覆盖 NanoGPT、Vision Transformer、MLP-Mixer 与 GSM8K LoRA，测试设计严密。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨自洽，附录完备，问题定位与实验剖析极为透彻。
价值: ⭐⭐⭐⭐⭐ 为大模型时代的自适应优化器设计指明了新范式，具备重要的理论与工业指导价值。