NorMuon: Making Muon more efficient and scalable

优化器的选择对大语言模型（LLM）的训练效率和计算成本具有显著影响。最近，Muon优化器通过参数更新的正交化，在改善优化几何结构（即提升条件数）方面展现出良好效果。尽管Muon被视为Adam的潜在继任者之一，但如何协同发挥两者优势的潜力尚未得到系统探索。在本研究中，我们提出NorMuon（神经元级归一化的Muon），一种将正交化与神经元级别的自适应学习率相结合的优化器，填补了这一空白。我们的分析表明，虽然Muon能有效降低条件数，但其产生的更新向量在不同神经元上的范数差异极大，导致某些神经元在优化过程中占据主导地位。NorMuon通过为每个神经元维护二阶动量统计信息，并在正交化之后进行按行归一化，解决了这一不平衡问题，在保留Muon改善条件数优势的同时，实现了更均衡的参数利用。为了支持大规模实际部署，我们在FSDP2框架下开发了一种高效的分布式实现方案，能够策略性地将正交化计算分布到多个设备上。在多种模型规模上的实验表明，NorMuon在性能上持续优于Adam和Muon：在11亿参数的预训练设置中，其训练效率比Adam高出21.74%，比Muon提升11.31%，同时内存占用与Muon相当。我们的研究结果表明，正交化与自适应学习率并非相互竞争，而是互补的方法，为大规模深度学习中的优化器设计开辟了新的方向。

论文试图解决大型语言模型（LLM）训练中优化器效率与计算成本之间的权衡问题，特别是现有优化器如Adam和新兴的Muon在参数更新过程中存在的优化几何不佳或更新不均衡的问题。尽管Muon通过正交化改善了条件数，但其导致神经元级别的更新范数高度不均，造成某些神经元主导训练过程。这个问题尚未被系统性地研究，尤其在如何结合自适应学习率与正交化优势方面仍属较新领域。

提出NorMuon优化器，将Muon的参数更新正交化与神经元级别的自适应学习率相结合：首先进行正交化以改善优化几何结构，然后对每一行（即每个神经元）进行归一化，并维护二阶动量统计，从而实现更均衡的参数更新。这一思路创新地将正交化与自适应方法视为互补而非竞争机制，突破了传统优化器设计的范式。

在多个模型规模上进行了实验验证，特别是在1.1B参数预训练设置下，NorMuon相比Adam提升21.74%训练效率，相比Muon提升11.31%；同时内存占用与Muon相当。作者开发了基于FSDP2框架的高效分布式实现，支持大规模部署。代码虽未明确提及开源，但其实现细节为后续研究提供了可复现路径。值得深入研究的方向包括将该思想扩展到其他架构（如视觉Transformer）以及探索更细粒度的自适应机制。

近期相关工作包括：'Adam: A Method for Stochastic Optimization'（Kingma & Ba, ICLR 2015）提出的Adam优化器；'MuSG: Momentum-aware Unit-norm SGD Improves Generalization and Training Speed'（Zhang et al., 2023）探讨了单位范数更新的优势；'Orthogonalized SGD: A Simple and Efficient Algorithm for Distributed Neural Network Training'（Wang et al., ICML 2022）研究了正交化在优化中的作用；以及'Layer-wise Adaptive Rates for Deep Learning'（You et al., ICML 2020）提出的LARS优化器，启发了层级别自适应学习率的设计。这些工作共同构成了NorMuon的思想基础。