根据数据动态调整优化器几何,让神经网络训练更智能高效。
From SGD to Muon: Adaptive Optimization via Schatten-p Norms

- 用梯度与激活统计量推导出自适应更新几何的闭式准则。
- 在三种场景下性能超越或媲美最优的MuAdam和AdamW。
- 仅增加约3%开销,即可实现从SGD到Muon的灵活切换。
现代优化器如Muon在更新时施加矩阵级几何约束,这些约束可统一为线性最小化预言机(LMO)理论。然而,现有方法对更新规则采用固定LMO几何,由设计或经验决定,未必契合问题本身的几何结构。本文提出一种高效的数据驱动准则,可动态选择各层神经网络的近似最优更新LMO几何。该准则基于单步随机特征回归代理模型,从梯度与激活统计量中闭式推导,可实现从SGD到Muon的连续插值。结合参数级预条件机制,该框架能还原SGD、Muon、Adam和MuAdam作为特例。为保证可扩展性,采用高效的计算策略,仅带来约3%的运行时间开销。实验证明,该自适应优化器在三种不同训练场景中表现优于或媲美最佳的MuAdam与AdamW。结果表明,可从运行时数据中成功且高效地适配LMO几何,为突破静态优化器设计开辟新路径。
原文摘要 · Abstract (English)
Modern optimizers, like Muon, impose matrix-wise geometry constraints on their updates. These matrix-wise constraints can be unified under Linear Minimization Oracle (LMO) theory. However, all current methods impose fixed LMO geometries for the update rules, chosen by-design or empirically, which are not necessarily optimal according to the problem's geometry. We introduce a novel efficient datadriven criterion for dynamically choosing proxy-optimal update LMO geometries on individual Deep Neural Network layers. Derived in closed form from gradient and activation statistics using a single-step random feature regression surrogate model, our criterion navigates a design space interpolating from SGD to Muon updates. Moreover, integrating parameter-wise preconditioning allows our framework to recover SGD, Muon, Adam, and MuAdam as specific extrema. To make this adaptive approach scalable, we pair it with efficient computational strategies, achieving only a $\sim$ 3% runtime overhead on highly optimized baselines. As a proof of concept, we show that this data-driven optimizer beats or remains competitive with the performance of the best performing optimizer between Muon and AdamW across three different training scenarios. Ultimately, this work provides evidence that LMO geometry can be successfully and efficiently adapted from runtime data, opening a new pathway for optimizer design beyond static geometries.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。