技术摘要：OrScale——基于层信任比率缩放的正交化优化

1. 问题陈述

Muon 优化器已成为大规模训练中的实用工具，它利用正交化更新，通过牛顿 - 舒尔茨迭代（近似极因子）来固定矩阵值步骤的方向。然而，Muon 存在一个关键局限：虽然它优化了更新的方向，但将幅度的控制权留给了一个单一的、本质上是全局的学习率。

这种全局缩放并非最优，因为不同的层（例如注意力投影、MLP 上/下投影、嵌入层）表现出截然不同的梯度分布和有效曲率。虽然 Moonlight 变体引入了静态的每形状因子以均衡每元素的均方根（RMS），但它无法适应训练过程中层尺度的动态漂移。

此前尝试将层自适应比率缩放（LARS）和层自适应矩缩放（LAMB）移植到 Muon 上的努力均告失败，这两种方法曾成功在 SGD 和 Adam 中将步长与全局学习率解耦。本文指出，简单的移植之所以崩溃，是由于三个具体的结构性原因：

1. 退化的分母：使用正交化方向（$\|Q_\ell\|_F$）的弗罗贝尼乌斯范数作为信任比率的分母会失效，因为它几乎是一个由矩阵形状（秩）决定的常数，而非更新幅度。
2. 单位不匹配与截断饱和：使用原始动量范数（$\|\tilde{M}_\ell\|_F$）会引入单位不匹配（梯度单位与权重单位），导致信任比率几乎普遍触及上限截断值（$r_{max}$），从而使自适应失效。
3. 解耦权重衰减的失控：将校准后的分母与解耦的权重衰减相结合，会形成一个正反馈循环，导致权重范数不受控制地增长（10–50 倍），从而破坏训练的稳定性。

2. 方法论

作者提出了OrScale，这是一个基于单一设计原则构建的框架：层信任比率的分母必须测量优化器将从参数中减去的真实更新方向的弗罗贝尼乌斯范数。

核心原则

设 $D_{\ell,t}$ 为在应用全局学习率之前从权重 $W_{\ell,t}$ 中减去的实际参数空间方向。信任比率定义为：
$$r_{\ell,t} = \frac{\|W_{\ell,t}\|_F}{\|D_{\ell,t}\|_F + \epsilon}$$
关键在于，$D_{\ell,t}$ 必须同时包含权重衰减项和正交化方向。

两种变体

该框架产生了两种具体的实现：

1. OrScale（通用/视觉默认）：

   • 更新方向：$D_{\ell,t} = \lambda W_{\ell,t} + Q_{\ell,t}$（其中 $\lambda$ 是权重衰减，$Q$ 是正交化方向）。
   • 形状因子：$s_\ell = 1$（无静态缩放）。
   • 校准：无。
   • 权重衰减：直接耦合到信任比率缩放的更新步骤中。

2. OrScale-LM（大语言模型默认）：

   • 更新方向：$D_{\ell,t} = \lambda W_{\ell,t} + s_\ell Q_{\ell,t}$。
   • 形状因子：采用 Moonlight 的静态因子 $s_\ell = 0.2\sqrt{\max(m_\ell, n_\ell)}$ 以匹配 AdamW 的每元素 RMS。
   • 校准：引入一个每层惰性校准常数（$c_{\text{denom}, \ell}$），在初始化时计算一次。这将初始信任比率锚定在恰好 1（$r_{\ell,0}=1$）。
   • 校准的意义：这使得 OrScale-LM 能够继承为 Moonlight（进而 AdamW）调优的学习率和权重衰减超参数，而无需进行新的超参数搜索。它确保了初始化时的宽度不变性。

算法细节

• 耦合权重衰减：权重衰减被整合到信任比率计算中（$\lambda W$ 是分母的一部分），而不是作为单独的拉回操作应用。这防止了解耦变体中观察到的失控反馈循环。
• 效率：该方法增加的开销微乎其微（<1% 的挂钟时间），每层矩阵仅需一个额外的标量，每步仅需两次弗罗贝尼乌斯范数计算，而这些计算量远小于现有的牛顿 - 舒尔茨迭代。

3. 主要贡献

• 算法：一个统一的、参数精简的配方（算法 1），成功将层信任比率缩放整合到 Muon 中。作者识别了设计空间中独特的“可行角落”（真实更新方向分母 + 耦合权重衰减）。
• 失败分析：对 Muon 信任比率设计空间进行了严格的映射（表 1），从分析和实证角度刻画了三种失败模式（退化分母、截断饱和和权重增长失控），阐明了它们失败的原因以及所提出方案成功的原因。
• 理论：
  • 收敛性：证明了核范数下的 $O(1/\sqrt{T})$ 非凸收敛保证。
  • 自适应增益：在可测量的层异质性下，建立了层自适应下降常数 $\kappa_{\text{eff}} > 1$，严格优于 Muon 的同质常数。
  • 分离性：证明了在截断饱和条件下，与 MuTrust/MuScale 的严格分离。
  • 校准：证明了 OrScale-LM 在初始化时保留了 $\mu$P 风格的学习率迁移能力。
• 实证：在视觉和语言模型任务上进行了全面评估。

4. 实验结果

视觉任务（CIFAR-10 / DavidNet）

• 性能：OrScale 实现了 94.05% $\pm$ 0.08 的验证集 top-1 准确率，在五个基线中排名第一。
• 改进：这比标准 Muon（93.70%）提高了 0.35 个百分点，比 Muon+Moonlight（93.75%）提高了 0.30 个百分点。
• 稳定性：与 Muon 相比，OrScale 展示了更宽的稳定学习率窗口，这与层自适应增益的理论预测一致。
• 消融实验：失败模式消融研究证实，对应于已识别失败模式的变体（例如使用 $\|Q_\ell\|_F$ 或原始动量）表现出预测的症状（例如 99.5–100% 的截断饱和或权重范数失控），并且未能达到 OrScale 的性能。

大语言模型（FineWeb-Edu 预训练）

• 扩展性：在四种模型规模（1.25 亿到 11 亿参数）上进行了评估。
• 对比：OrScale-LM 在四种规模中的三种（1.25 亿、5.45 亿、11 亿）上优于 Muon+Moonlight，并在所有规模上击败了 AdamW。
• 损失降低：在 11 亿参数规模下，OrScale-LM 的最终验证交叉熵为 2.6251，而 Muon+Moonlight 为 2.6360，AdamW 为 2.7304。
• 效率：这种优势在计算范围内得以保持，OrScale-LM 在计算对大语言模型最为关键的区间表现出清晰的胜利。

5. 意义与主张

本文声称，OrScale 成功弥合了 Muon 的谱几何特性与 LARS/LAMB 的层自适应特性之间的差距。其主要意义在于：

1. 解决幅度问题：它提供了首个针对 Muon 的、对结构失败模式具有鲁棒性的时变层幅度控制。
2. 实际可迁移性：通过 OrScale-LM 变体和惰性校准，它允许从业者将针对 AdamW 或 Moonlight 调优的学习率直接迁移到 OrScale-LM，而无需重新调优，降低了采用门槛。
3. 理论严谨性：它为以往混合尝试的失败提供了严格的理论依据，并证明了所提出的方法在异质性条件下能产生严格更优的收敛常数。
4. 实证验证：它证明了这些理论改进转化为小规模视觉任务和大语言模型预训练中的切实收益，确立了 OrScale 作为矩阵值层最先进优化器的地位。

作者对大语言模型的结果保持谦逊，指出 11 亿参数规模的优势在小规模下与单次随机种子噪声相当，因此这些结果应被视为该方法有效性的方向性证据，而非在所有条件下具有决定性的统计优越性。然而，跨规模的持续改进以及对失败模式的成功消融，有力地支持了所提出的设计原则。