DynMuon：Muon 优化的动态谱整形视角 DynMuon: A Dynamic Spectral Shaping View of Muon
动态调度 Muon 的谱指数 p，跨规模省 10.6–26.5% 训练步且验证损失更低。
前置知识
Muon 优化器
Muon 是面向矩阵参数的优化器：对矩阵形动量梯度 $M=U\Sigma V^\top$，用极因子 $UV^\top$ 替换更新，即把奇异值拍平为 1、保留奇异方向，实现上用少量 Newton-Schulz 迭代近似极因子以避免完整 SVD，在 LLM 训练中收敛更快更稳。
本文全部工作都是对 Muon 的推广与改进，必须先理解 Muon 用极因子替换动量更新（p=0 的固定谱整形）这一核心机制，才能理解作者为何要把它推广为可变指数 p 的谱整形族。
谱整形（spectral shaping）
对更新矩阵 $M=U\Sigma V^\top$ 做奇异值分解后，把奇异值矩阵 $\Sigma$ 替换为 $\Sigma^p$：$p=1$ 恢复原始梯度（SGD 式），$p=0$ 得到 Muon（奇异值全 1），$p=-1$ 为逆谱。本文把这族变换统称为谱整形，研究指数 $p$ 如何影响训练。
整篇论文的理论与算法都围绕 $D(p)=U\Sigma^p V^\top$ 这族变换展开，不理解 p 取不同值对应的几何含义，就无法读懂信号—噪声权衡与动态调度的设计动机。
局部曲率与噪声模型
在局部最优点 $W^\star$ 附近对损失做二次近似 $\nabla L(W)\approx \kappa H E$，其中 $H$ 是归一化曲率矩阵、特征值 $h_i\in(0,1]$ 代表各曲率方向，再叠加零均值随机梯度噪声 $\Xi$。该模型把谱整形解释为按曲率模式对残差信号与噪声的重加权。
第 2 节的全部推导都建立在这个噪声感知局部模型上，模式级信号—噪声权衡方程 (6)(7) 是预测阶段依赖规律的理论基础，读懂它才能理解为何早期正 p、晚期轻负 p。
Newton-Schulz 迭代
一种用低阶多项式迭代逼近矩阵极因子（正交化）$A^{-1/2}X$ 而无需特征分解的方法。Muon 用它对归一化梯度做几次运算得到 $UV^\top$。DynMuon 在其上叠加二阶 Taylor 修正项以高效支持任意分数指数 $p$。
DynMuon 高效性的关键在于把分数谱整形拆成 Newton-Schulz 加一个低价多项式修正，理解这一迭代才能看懂为何动态谱整形能保持与 Muon 相同的渐近复杂度 $O(m^2n+m^3)$。
残差信号与模式分解
残差信号 $E_t=W_t-W^\star$ 衡量参数到局部最优的距离，投影到曲率矩阵特征基后得到各模式坐标 $\delta_{i,t}$。训练后期强（高曲率）模式的残差先衰减，平坦（低曲率）模式残差相对占优，这正是选择轻微负指数的依据。
作者用残差能量 $\delta_{i,t}^2$、噪声水平 $c_{i,t}$ 和残差偏移 $\Pi_t$、信噪比优势 $\Omega_t$ 等量来实测验证理论，理解这些模式级量是看懂 Figure 1 验证证据的前提。
研究动机
Muon 已成为训练大语言模型的主流优化器：对每个矩阵形参数，它构造动量平均梯度 $M=U\Sigma V^\top$，再用极因子 $UV^\top$ 替换，即把奇异值统一拍平为 1。这种 $p=0$ 的固定谱整形在多尺度上提升了收敛与稳定性，但近期分析表明其有效性会随训练条件变化。已有工作虽探索了固定正指数的 power-law 谱整形或抑制主导子空间的变换，却仍把谱整形当作静态规则，试图寻找一个最好的固定变换。然而训练动力学在早期与晚期截然不同——残差信号在不同曲率模式间的分布会随步数演化——单一固定的 $p=0$ 很可能并非全程最优。此外先前工作缺乏描述谱整形如何随训练演化的动力学模型，留下一个未解的核心问题：Muon 式的谱整形究竟应不应随训练阶段自适应调整？
本文的目标是本文目标是系统回答谱整形是否、以及如何应随训练阶段自适应。具体而言：首先把 Muon 推广为含谱指数 $p$ 的幂律族 $D(p)=U\Sigma^p V^\top$，统一刻画 $p=-1$（逆谱）、$p=0$（Muon）、$p=1$（SGD）等情形；其次建立一个噪声感知的局部动力学模型，揭示 $p$ 如何沿局部曲率方向对残差信号收缩与随机噪声放大进行模式级权衡；再次据此发现一个被忽视的阶段依赖规律——早期正 $p$ 有利、晚期轻微负 $p$ 有利；最终提出高效算法 DynMuon，用一条从正到轻负的 logistic 调度曲线，在几乎不增加每步开销的前提下持续超越 Muon，并在多种规模、架构与训练设置下验证。
与已有工作不同的是，本文的独特切入在于三点结合。一是把谱整形从寻找单一最优固定变换，重新框架为随训练演化选择合适指数 $p$ 的自适应问题，这是先前 spectral shaping 工作未曾系统化处理的。二是首次为 Muon 类优化提供噪声感知的局部模型，用曲率模式的信号—噪声分解把抽象的谱操作与训练阶段联系起来，既有机制洞察又能产生可验证的定量预测（如残差向平坦模式集中、平坦模式的信噪比优势 $\Omega_t>0$）。三是工程上把 Newton-Schulz 迭代扩展到任意分数指数 $p$，使动态谱整形保持与 Muon 相同的渐近复杂度而无需在线估计残差分布，兼顾理论与实践。
核心方法
整体思路是先建模型再设计算法。直觉上：Muon 把奇异值拍平为 1（$p=0$），但训练早期信号集中在高曲率方向、应加强；晚期信号集中在平坦方向、应转移权重。技术路线上，作者先把谱整形写成统一族 $D(p)=U\Sigma^p V^\top$，再在局部最优点附近用二次近似 $\nabla L(W)\approx \kappa H E$ 把梯度与残差联系起来，并用曲率对齐代理 $(GG^\top)^{(p-1)/2}\approx \alpha^{(p-1)/2}H^{(p-1)/2}$，推导出残差的逐模式演化方程。该方程显示 $p$ 通过 $h_i^{(p+1)/2}$ 控制信号收缩、通过 $h_i^{(p-1)/2}$ 放大噪声。基于早期正 $p$、晚期轻负 $p$ 的预测，DynMuon 用一条单调递减的 logistic 曲线调度 $p_t$，并用扩展的 Newton-Schulz 高效实现，全程不依赖在线估计。
核心创新是把 Muon 的固定极因子重构为随阶段自适应的谱指数，并用噪声感知的局部模型给出可操作的调度规则。与已有方法的本质区别在于：与固定正指数谱整形不同，DynMuon 允许 $p$ 连续从正变负；与只调标量超参（学习率退火、warmup）的调度不同，它直接调度矩阵更新本身的谱形状，使不同谱方向的相对权重随训练演化；理论给出明确的阶段判据——早期残差集中在强模式、应增 $p$ 加速收缩并抑制噪声；晚期残差集中到平坦模式、应轻微降 $p$ 把收缩力转向这些方向，但不可过负以免噪声放大超过收益。这套正到轻负的 logistic 调度正是其他静态或反向调度所不具备的。
方法步骤详情
输入当前步 $t$、总步 $T$、矩阵形动量梯度 $M=U\Sigma V^\top$。(1) 按 logistic 算进度 $a_t=1/(1+\exp(u_t))$，$u_t=(t/T-\tau)/w$，令 $p_t=p_{\min}+a_t(p_{\max}-p_{\min})$，默认 $p_{\max}=1$、$p_{\min}=-0.25$。(2) 按 $p_t$ 分段：若 $p_t\ge 0.25$ 用原始更新；若 $p_t\in[0,0.25)$ 用 Newton-Schulz 正交化（Muon）；若 $p_t\in[p_{\min},0)$ 用高效谱近似。(3) 对归一化输入 $X_n=X/\|X\|_F$，令 $A=X_n X_n^\top$、$E=A-I$、$\delta=p/2$，做二阶 Taylor 得 $C=I+\delta E+\tfrac12\delta(\delta-1)E^2$。(4) 对 $Y_\mu=A^{-1/2}X_n$（Muon 更新）做 NS 迭代。(5) 输出 $\hat X=\|X\|_F^p\,C\,Y_\mu$ 作为更新。
技术新颖性
技术新颖性体现在三方面。理论新颖：首次为 Muon 类优化建立噪声感知的局部模型，把谱整形分解为模式级信号—噪声权衡，并预测出晚期轻微负指数有利这一被忽视的现象，且模型预测得到实测验证（残差向平坦模式集中、平坦模式信噪比优势 $\Omega_t$ 为正）。算法新颖：用 logistic 曲线调度 $p_t$ 并配合稳定锚定策略——正区间用原始更新/NS、轻负区间用 Taylor 修正——避免 $p$ 较大时 Taylor 近似不可靠带来的不稳定。工程新颖：通过等价分解 $U\Sigma^p V^\top=(X_n X_n^\top)^{p/2}(X_n X_n^\top)^{-1/2}X_n$ 把分数谱整形拆成 NS 加低价多项式修正，复杂度与 Muon 同为 $O(m^2n+m^3)$，实测每步时间仅为 Muon 的 1.003–1.025 倍。
实验结果
核心发现分理论与实验。理论：模型预测的三个量均被实测支持——残差偏移 $\Pi_t$ 在约 500 步后转负（残差向平坦模式集中）、平坦模式噪声调整信号 $u_{i,t}$ 保持显著正、平坦模式信噪比优势 $\Omega_t$ 迅速转正并维持。实验一（阶段依赖）：第 500 步切负指数时，$p=-0.1/-0.25$ 优于 Muon($p=0$)，而 $-0.75/-1$ 失稳变差；早期用 $p=1$ 再切 $-0.1$ 的两段式调度取得最低验证损失，反向(负到正)或翻倍学习率均更差。实验二（主结果，Table 1）：在 127M/601M/1.11B、10B/20B token 下 DynMuon 验证损失持续最低，达到 Muon 在 80% 训练处损失目标的步数减少 10.6%–26.5%（127M/20B 省 26.5%），每步时间仅增 0.3%–2.5%。此外在 Qwen 式架构、2.5B–20B token、学习率 0.003–0.04、FineWeb-Edu 及 NorMuon 基线上均稳定领先，近似与精确 SVD 轨迹几乎重合。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| GPT 式预训练验证损失（127M / 10B token, FineWeb） | Best Val. Loss (↓) 与 Steps to Target | DynMuon 3.171，12500 步达目标 | Muon 3.190，16000 步 | 验证损失更低，达目标步数减少 21.9% |
| GPT 式预训练验证损失（601M / 20B token, FineWeb） | Best Val. Loss (↓) 与 Steps to Target | DynMuon 2.797，25000 步达目标 | Muon 2.808，30400 步 | 验证损失更低，达目标步数减少 17.8% |
| GPT 式预训练验证损失（1.11B / 10B token, FineWeb） | Best Val. Loss (↓) 与 Steps to Target | DynMuon 2.776，14300 步达目标 | Muon 2.788，16000 步 | 验证损失更低，达目标步数减少 10.6% |
| 每步计算开销（多规模） | Per-Step Time 相对 Muon | 1.003–1.025× Muon | Muon 1.0× | 几乎无额外运行时开销，近似与精确 SVD 轨迹重合且快约 3× |
局限与改进
作者明确承认实验目的是验证随训练调 $p$ 可获实际改进这一假设，所提 logistic 调度并非终局，最优 $p$ 可能依赖无法预先完全预知的训练时动力学。模型本身做了多处简化近似（局部二次化、K-FAC 式 Hessian、曲率对齐代理、条件零均值噪声），虽经验验证但非端到端收敛证明。模型分析与实验主要在 decoder-only Transformer、有限规模（≤1.11B）与有限 token 预算（≤20B）上完成，未覆盖更大模型或非 Transformer 架构的稳定性。我观察到：负指数区被刻意限制在 $p\in[p_{\min},0)$ 且依赖 Taylor 近似，$p$ 更负时稳定性论证较薄弱；调度参数 $\tau$、$w$ 及 $p_{\min}$ 仍需人工设定，缺乏跨任务自动选取；与 SOAP/Shampoo 等非对角优化器的直接对比也较少。
独立分析的弱点
弱点一：调度依赖固定超参（$p_{\max}=1$、$p_{\min}=-0.25$、$\tau=0.04$、$w=0.04$），不同数据/规模/损失可能最优值不同，缺乏自适应机制。改进方向：用文中提到的残差偏移 $\Pi_t$ 或信噪比 $\Omega_t$ 作为在线信号，偶尔重估 $p$（作者自己也指出近最优调度器或只需偶尔重估 $p$）。弱点二：稳定区被强制锚定为原始更新/NS（$p\ge0.25$ 用原更新、$p\in[0,0.25)$ 用 Muon），丢失了正区间连续插值潜力，且 Taylor 近似在 $p$ 偏大时不可靠。改进方向：用更高阶或保正的多项式近似扩展连续正区间。弱点三：实验规模与 token 预算偏小（≤1.11B、≤20B），未在百亿参数或长 horizon 验证。改进方向：扩大到更大模型与万亿 token。弱点四：仅横向对比 AdamW/Muon/NorMuon，缺少与 SOAP/Shampoo/PolarGrad 等矩阵感知优化器的直接比较。
未来方向
作者提出的方向：把 $p$ 的选择从固定 logistic 调度升级为基于观测优化统计量的在线自适应，允许偶尔（而非每步）使用更昂贵信号（如 Hessian-vector 乘积、残差能量 $\delta_{i,t}^2$、噪声水平 $c_{i,t}$）来重估 $p$，因为有用 $p$ 的范围相对窄，近最优调度器或只需偶尔重估。基于成果可延伸：将谱整形思想迁移到非矩阵参数（如嵌入、LayerNorm）或非 Transformer 架构；把模式级信号—噪声分解用于诊断与可视化训练健康度；探索更一般的谱变换（非幂律、随方向各异的各向异性整形）；与学习率/动量调度联合优化形成统一调度框架；研究谱整形与二阶方法、自然梯度的理论联系，建立更严格的端到端收敛保证。
复现评估
复现性较好。作者开源代码于 github.com/fzwark/DynMuon。实验配置详细：模型为 GPT 式（rotary 位置编码、RMSNorm、ReLU² MLP，见表 2/3）与 171M Qwen 式（GQA、gated SiLU），序列长 1024、全局批 512，默认学习率 Muon/DynMuon=0.01、AdamW=0.002，主实验用 FineWeb/FineWeb-Edu 的 10B–20B token。关键超参（$p_{\max}$、$p_{\min}$、$\tau$、$w$）与消融齐全。挑战：训练需多卡大算力（1.11B、20B token、批 512），单机难复现主结果；局部模型分析需 Hessian-vector 乘积与小批量梯度方差估计，工程较重；某些近似（曲率对齐）需经验验证。小规模 127M/10B token 实验相对易复现，可作为入口。整体属于可复现但门槛较高的前沿工作。
论文图表
三联图验证模型预测。左：残差偏移 $\Pi_t$ 随训练递减并在约 500 步后转负，说明残差信号从高曲率（强）模式转向低曲率（平坦）模式。中：平坦模式的噪声调整信号 $u_{i,t}=\log\hat\delta^2-\log\hat c$ 保持显著正，表明平坦模式残差相对噪声仍大。右：平坦模式信噪比优势 $\Omega_t=\mathrm{med}_F u-\mathrm{med}_S u$ 迅速上升并维持正值，证明晚期平坦模式相对强模式有更有利的信噪权衡。
这张图是理论模型可被实测验证的直接证据，回答了负指数为何在晚期有利，是理解整篇动机与信号—噪声理论的核心。