AngularMuown:显式角度步长衰减的矩阵感知优化器

Muown Implicitly Performs Angular Step-size Decay

精选理由

想加速Transformer预训练?这篇论文把优化器角度步长显式化,新方法AngularMuown在nanoGPT竞赛中领先,还在0.5B和1.1B模型上验证了效果。

AI 摘要

Muown优化器将权重矩阵分解为行幅度和未归一化的方向变量,分别用Adam和Muon更新。研究表明,Muown的方向更新等价于对归一化方向的黎曼步长,而幅度仅调制角度步长,这解释了其稳定性。基于此提出的AngularMuown直接优化归一化方向,并使用可调度的角度乘子与径向幅度更新解耦。在modded nanoGPT速度竞赛中,初步版AngularMuown领先每优化器类别。在Qwen2-0.5B和1.1B参数MoE模型上的实验证实该算法可扩展至更大规模。

原文 · arXiv cs.LG

Muown Implicitly Performs Angular Step-size Decay

Matrix-aware optimizers such as Muon and Muown have recently shown strong empirical performance for pre-training Transformers. In particular, Muown separates each weight matrix into row magnitudes and an un-normalized direction variable, updating the former with Adam and the latter with Muon. We show that the directional update of Muown is equivalent to a Riemannian step on the normalized directions, while the magnitude of the un-normalized parameterization only modulates the angular step size. This explains the step-size stability of Muown and suggests making the angular step size explicit. The resulting method, AngularMuown, optimizes directly over the normalized directions and uses a schedulable angular multiplier decoupled from the radial magnitude update. AngularMuown improves over Muown and, at the time of writing, a preliminary version is leading the per-optimizer category of the modded nanoGPT speedrunning competition. Further experiments on Qwen2-0.5B, and 1.1B parameter mixture-of-experts models confirm the algorithm scales beyond small models. An implementation of the algorithm is available at https://github.com/fhueb/angular-muown