arXiv:2603.28254cs.LGstat.ML2026-03被引 14

提出轻量级预正交化均衡方法,提升矩阵参数优化效果

MuonEq: Balancing Before Orthogonalization with Lightweight Equilibration

  • 在正交化前用行列归一化重平衡动量矩阵,改善输入几何结构
  • 在LLaMA2预训练中,130M~1B模型均实现更快收敛和更低困惑度
  • 适合需要高效优化矩阵参数的大型语言模型训练场景

正交化更新优化器如Muon能提升矩阵参数的训练效果,但现有方法通常在正交化后缩放更新或使用更复杂的白化预条件器。本文提出{ exttt{MuonEq}},一种轻量级的预正交化均衡方案,包含双侧行列归一化(RC)、行归一化(R)和列归一化(C)三种形式。通过在有限步Newton-Schulz正交化前重新平衡动量矩阵,{ exttt{MuonEq}}改善了正交化所见的几何结构。我们证明有限步正交化受输入谱影响,特别是稳定秩和条件数,而行列归一化可作为白化的零阶近似。对于隐藏层矩阵权重,行归一化(R)为默认变体。理论上,{ exttt{MuonEq}}(R)保留标准$ ilde{ m O}(T^{-1/4})$非凸平稳性保证,支持解耦权重衰减和无边界递减学习率,并将该结果扩展至最多含显式不精确常数的有限步NS5。在C4数据集上对LLaMA2进行预训练时,{ exttt{MuonEq}}(R)在130M、350M和1B模型上均持续优于Muon,具备更快收敛速度和更低验证困惑度。代码已开源。

原文摘要 · Abstract (English)

Orthogonalized-update optimizers such as Muon improve training of matrix-valued parameters, but existing extensions typically either rescale updates after orthogonalization or use heavier whitening-based preconditioners before it. We introduce {\method}, a lightweight family of pre-orthogonalization equilibration schemes for Muon with three forms: two-sided row/column normalization (RC), row normalization (R), and column normalization (C). By rebalancing the momentum matrix before finite-step Newton--Schulz orthogonalization, {\method} improves the geometry seen by orthogonalization. We show that finite-step orthogonalization is governed by the input spectrum, especially stable rank and condition number, and that row/column normalization acts as a zeroth-order surrogate for whitening. For hidden matrix weights, R is the default variant. Theoretically, {\method} (R) retains the standard $\widetilde{\mathcal O}(T^{-1/4})$ Muon-type nonconvex stationarity guarantee with decoupled weight decay and a horizon-free diminishing learning-rate schedule, and extends it to finite-step NS5 up to an explicit inexactness constant. In LLaMA2 pretraining on C4, {\method} (R) consistently outperforms Muon on 130M, 350M, and 1B models, with faster convergence and lower validation perplexity. The code is available at the \href{https://github.com/MaeChd/muon-eq}{MuonEq codebase}.

优化器矩阵参数大模型训练

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。