提出新型归一化优化器MOGA,实现跨模型宽度的稳定学习率迁移。
On the Width Scaling of Neural Optimizers Under Matrix Operator Norms I: Row/Column Normalization and Hyperparameter Transfer
- 基于矩阵算子范数视角设计可层间组合的归一化优化方法
- 新方法使平滑性常数与网络宽度无关,突破原有优化器瓶颈
- 适用于大模型训练,显著提升长序列和低损失场景下的效率
现代深度学习中一个核心问题是:如何设计在网络宽度 $w$ 增加时仍保持行为稳定的优化器。本文将 AdamW、Muon 等常用优化器解释为矩阵算子范数下的最速下降法,揭示了优化器几何与网络前向映射利普希茨结构的联系,实现了对利普希茨和光滑性常数的宽度无关控制。然而,标准 $p \to q$ 算子范数诱导的最速下降规则缺乏层间可组合性,无法在深层架构中提供宽度无关的界。为此,本文引入一族均值归一化算子范数 $\pmean \to \qmean$,具备层间可组合性,获得宽度无关的光滑性界,并导出实用优化器如重缩放 AdamW、行归一化、列归一化。由此得到的学习率宽度感知缩放规则恢复了 $μ$P 缩放作为特例,提供了跨宽度学习率迁移的合理机制。我们进一步证明,Muon 可能导致光滑性常数以 $\mathcal{O}(\sqrt{w})$ 的最坏情况增长,而本文提出的行归一化优化器则实现宽度无关的光滑性保证。基于此,提出 MOGA(Matrix Operator Geometry Aware)——仅依赖行/列归一化的宽度感知优化器,支持跨宽度的稳定学习率迁移。GPT-2 和 LLaMA 大规模预训练结果表明,尤其采用行归一化时,MOGA 在性能上可媲美 Muon,且在大令牌和低损失场景下显著更快。
原文摘要 · Abstract (English)
A central question in modern deep learning is how to design optimizers whose behavior remains stable as the network width $w$ increases. We address this question by interpreting several widely used neural-network optimizers, including \textrm{AdamW} and \textrm{Muon}, as instances of steepest descent under matrix operator norms. This perspective links optimizer geometry with the Lipschitz structure of the network forward map, and enables width-independent control of both Lipschitz and smoothness constants. However, steepest-descent rules induced by standard $p \to q$ operator norms lack layerwise composability and therefore cannot provide width-independent bounds in deep architectures. We overcome this limitation by introducing a family of mean-normalized operator norms, denoted $\pmean \to \qmean$, that admit layerwise composability, yield width-independent smoothness bounds, and give rise to practical optimizers such as \emph{rescaled} \textrm{AdamW}, row normalization, and column normalization. The resulting learning rate width-aware scaling rules recover $μ$P scaling~\cite{yang2021tensor} as a special case and provide a principled mechanism for cross-width learning-rate transfer across a broad class of optimizers. We further show that \textrm{Muon} can suffer an $\mathcal{O}(\sqrt{w})$ worst-case growth in the smoothness constant, whereas a new family of row-normalized optimizers we propose achieves width-independent smoothness guarantees. Based on the observations, we propose MOGA (Matrix Operator Geometry Aware), a width-aware optimizer based only on row/column-wise normalization that enables stable learning-rate transfer across model widths. Large-scale pre-training on GPT-2 and LLaMA shows that MOGA, especially with row normalization, is competitive with Muon while being notably faster in large-token and low-loss regimes.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。