从曲率视角出发,通过泰勒展开和曲率分解,发现 Muon 因更低的归一化方向锐度(NDS)而比 Adam 实现更大的一步损失下降,数据不平衡和层内曲率是其主要优势来源。
摘要:
Muon 在大语言模型训练中相比 Adam 将训练效率提升约两倍,但这一优势的局部几何来源尚不清楚。我们的工作首次从曲率视角尝试揭开 Muon 优于 Adam 的原因。首先,我们对训练损失曲面应用二阶泰勒近似,表明在匹配验证损失下,Muon 比 Adam 实现更大的一步损失下降。两种优化器的一阶增益相当,但 Muon 始终承受更小的二阶曲率惩罚。其次,我们将该曲率惩罚分解为更新范数的平方和归一化方向锐度(NDS)。我们发现 Muon 和 Adam 的更新范数相当,因此 Muon 更小的曲率惩罚源于更低的 NDS,而非更新尺度。第三,我们研究训练数据和模型结构如何塑造 Muon 的 NDS 优势。使用具有受控不平衡的 Zipf-概率上下文无关文法(PCFG)数据,我们表明数据不平衡放大了 Muon 相对于 Adam 的 NDS 优势。进一步的层内/跨层分解表明,在训练的中后期,Muon 更低的 NDS 主要由更小的层内曲率维持。除了经验证据,我们还分析了具有异质曲率和梯度对齐于高曲率模式的风格化二次问题。我们证明 Muon 通过平衡曲率组间的更新能量,实现了比 GD 更低的平均 NDS;当曲率异质性足够强时,在相同步数后这也产生更低的局部二次损失。
英文摘要:
Muon improves training efficiency over Adam in large language-model training by about two times, but the local geometric source of this advantage remains unclear. Our work takes a first step toward demystifying Muon's superiority over Adam from a curvature perspective. First, we apply a second-order Taylor approximation to the training landscape and show that Muon achieves a larger one-step loss decrease than Adam at matched validation loss. The two optimizers have comparable first-order gains, but Muon consistently incurs a smaller second-order curvature penalty. Second, we decompose this curvature penalty into the squared update norm and Normalized Directional Sharpness (NDS). We find that Muon and Adam have comparable update norms, so Muon's smaller curvature penalty is driven by lower NDS, not update scale. Third, we study how training data and model structure shape Muon's NDS advantage. Using Zipf-Probabilistic Context-Free Grammar (PCFG) data with controlled imbalance, we show that data imbalance amplifies Muon's NDS advantage over Adam. A within-/cross-layer decomposition further shows that, in the middle and late stages of training, Muon's lower NDS is mainly sustained by smaller within-layer curvature. Beyond empirical evidence, we analyze stylized quadratic problems with heterogeneous curvature and gradient alignment toward high-curvature modes. We prove that Muon attains a smaller average NDS than GD by balancing update energy across curvature groups; when curvature heterogeneity is sufficiently strong, this also yields lower local quadratic loss after the same number of steps.
Muon Learns More Robust and Transferable Features than Adam
Yale University(耶鲁大学);National University of Singapore(新加坡国立大学);University of Chinese Academy of Sciences(中国科学院大学);Academy of Mathematics and Systems Science, CAS(中国科学院数学与系统科学研究院)
本文通过鲁棒性和可迁移性视角,证明 Muon 优化器相比 Adam 和 SGD 能学习到更鲁棒、更可迁移的特征,并通过理论分析支持了经验发现。
Conda: Column-Normalized Adam for Training Large Language Models Faster
State Key Lab of General AI, School of Intelligence Science and Technology, Peking University(通用人工智能国家重点实验室,智能科学与技术学院,北京大学);Singapore Management University(新加坡国立大学);Nankai University(南开大学);Beijing Normal University(北京师范大学);ByteDance Seed(字节跳动种子);Beijing University of Posts and Telecommunications(北京邮电大学);Institute for Artificial Intelligence, Peking University(人工智能研究院,北京大学);Pazhou Laboratory (Huangpu), Guangzhou, Guangdong, China(黄埔实验室(广州),广东,中国)
GWT: Scalable Optimizer State Compression for Large Language Model Training
National Key Laboratory of Parallel and Distributed Computing(并行与分布式计算国家重点实验室);National University of Defense Technology(国防科技大学);Center for Machine Learning Research(机器学习研究中心)
Preconditioning Benefits of Spectral Orthogonalization in Muon
Penn Department of Statistics and Data Science, Wharton School, University of Pennsylvania(宾夕法尼亚大学统计与数据科学系,沃顿商学院,宾夕法尼亚大学);Yale Department of Statistics and Data Science, Yale University(耶鲁大学统计与数据科学系,耶鲁大学);Department of Electrical and Systems Engineering, University of Pennsylvania(宾夕法尼亚大学电气与系统工程系,宾夕法尼亚大学)
Why Adam Can Beat SGD: Second-Moment Normalization Yields Sharper Tails
Department of Electrical and Computer Engineering, The Ohio State University(俄亥俄州立大学电气与计算机工程系);School of Electrical, Computer and Energy Engineering, Arizona State University(亚利桑那州立大学电气、计算机与能源工程学院)
MoE Key Lab of Artificial Intelligence, AI Institute, School of Computer Science, Shanghai Jiao Tong University(人工智能微电子实验室、人工智能学院、计算机科学学院、上海交通大学);Xiaohongshu Inc.(小红书公司)
Clean: Second-order LLM Training at Linear Memory Cost via Nystr{ö}m Sketching
Mila & DIRO, Université de Montréal(蒙特利尔大学米拉与DIRO);Critical ML, University of Waterloo(滑铁卢大学Critical ML);Zhejiang University(浙江大学);Canada CIFAR AI Chair(加拿大CIFAR AI主席)
How the Optimizer Shapes Learned Solutions in Equivariant Neural Networks
University of Stuttgart, Germany(斯图加特大学);International Max Planck Research School for Intelligent Systems, Germany(国际马克斯·普朗克智能系统研究学校);Tudor Vianu High School of Computer Science, Romania(托尔德·维安乌计算机科学高中)
ETH Zurich(苏黎世联邦理工学院);Technical University of Munich(慕尼黑技术大学);ELLIS Institute Tübingen(图宾根ELLIS研究所);MPI-IS Tübingen(图宾根MPI-IS);Tübingen AI Center(图宾根人工智能中心)