提出OrScale优化器,通过分层信任比动态调整步长,显著提升模型训练效率。
OrScale: Orthogonalised Optimization with Layer-Wise Trust-Ratio Scaling
- 基于正交化更新方向,设计分层动态缩放因子替代静态校正。
- 在125M至1.1B参数模型上性能媲美或超越原有方法,16B模型中误差降低3.8%。
- 无需额外调参即可迁移现有优化策略,适合大规模语言模型训练者。
Muon将每个矩阵更新的方向固定在其动量的极分解方向,但各层步长仅通过静态形状修正处理。本文提出一种动态分层标量:将LARS/LAMB的信任比原则适配到正交化设置,其中标准分母候选(原始动量范数或极分解范数)或处于错误单位空间,或不携带更新尺度信息。所提方法OrScale采用实际施加的参数空间方向范数作为基准,对每层进行校准使信任比锚定为1,从而实现无需额外调优即可完整迁移Moonlight配方(原为AdamW设计,通过RMS匹配共享给Muon)。组件消融验证了每个设计选择均具独立贡献。理论上,OrScale对任意截断乘子保持核范数 $O(1/ ext{sqrt}{T})$ 的收敛速率,并在两个可由标准训练诊断估计的条件下实现严格分层自适应下降增益 $κ_{ ext{eff}} > 1$,该增益预测会随架构异质性增长。实验验证此预测:在继承Moonlight所有超参数的前提下,OrScale在密集型125M–1.1B FineWeb-Edu预训练中表现与或优于Muon+Moonlight;在16B-A3B混合专家模型上,差距扩大至0.130 nats(相对提升3.8%),且耗时与基线持平。
原文摘要 · Abstract (English)
Muon fixes the \emph{direction} of every matrix-valued update at the polar factor of its momentum, while each layer's step \emph{magnitude} is addressed only by a static shape correction. We derive a dynamic per-layer scalar by adapting the LARS/LAMB trust-ratio principle to the orthogonalized setting, where the standard denominator candidates---the raw momentum norm or the polar-factor norm---either live in the wrong unit space or carry no update-scale information. The resulting method, \emph{OrScale}, uses the norm of the parameter-space direction actually applied and anchors each layer's ratio at one via a per-layer calibration, so that the Moonlight recipe (tuned for AdamW, shared with Muon via RMS matching) transfers with \emph{no additional sweep}; a component ablation confirms each design choice is individually load-bearing. Theoretically, OrScale retains a nuclear-norm $O(1/\sqrt{T})$ convergence rate for any clipped multiplier and achieves a strict layer-adaptive descent gain $κ_{\mathrm{eff}}>1$ under two conditions estimable from standard training diagnostics---a bound that predicts the gain should \emph{grow with architectural heterogeneity}. Experiments confirm the prediction: with every hyperparameter inherited verbatim from the Moonlight recipe, OrScale matches or beats Muon+Moonlight across dense 125M--1.1B FineWeb-Edu pre-training, and on a 16B-A3B mixture-of-experts model---where the logged trust ratios separate cleanly by layer class---the gap widens by an order of magnitude to $0.130$ nats ($3.8\%$ relative) at parity wall-clock cost.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。