⏶6
NorMuon：使Muon更有效和可扩展
发表
由 
Chen Liang 提交
Chen Liang 提交作者: Zichong Li, Liming Liu, Chen Liang, Weizhu Chen, Tuo Zhao
摘要
AI 生成总结
NorMuon 是一种将正交化与神经元级自适应学习率相结合的新颖优化器，可提高训练效率并平衡大型语言模型中的参数利用率。优化器的选择对大语言模型 (LLMs) 的训练效率和计算成本有显著影响。最近，Muon 优化器通过正交化参数更新，通过更好的条件化来改善优化几何，已显示出有希望的结果。尽管 Muon 已成为 Adam 的潜在继承者，但系统地探索它们优势的联合利用尚未被探索。在这项工作中，我们通过提出 NorMuon（神经元维度归一化 Muon）来弥合这一差距，该优化器将正交化与神经元级别的自适应学习率协同组合。我们的分析表明，虽然 Muon 有效地减小了条件数，但由此产生的更新显示出高度不均匀的神经元范数，导致某些神经元主导优化过程。NorMuon 通过为每个神经元维护二阶动量统计量并在正交化后应用行向归一化来解决这种不平衡，从而确保参数利用的平衡，同时保留 Muon 的条件化优势。为了实现大规模的实际部署，我们在 FSDP2 框架下开发了一种高效的分布式实现，该框架将正交化计算策略性地分布到各个设备上。跨越多种模型规模的实验表明，NorMuon 在 1.1B 预训练设置上，比 Adam 提高了 21.74% 的训练效率，比 Muon 提高了 11.31%，同时保持了与 Muon 相当的内存占用，始终优于 Adam 和 Muon。我们的发现表明，正交化和自适应学习率是互补的，而不是竞争的方法，为大规模深度学习中的优化器设计开辟了新的途径。
优化器的选择对大型语言模型（LLM）的训练效率和计算成本有显著影响。最近，Muon 优化器通过正交化参数更新，并改善优化几何结构，以更好的条件数，展现出了有前景的结果。尽管 Muon 已成为 Adam 的潜在继任者，但系统性地结合两者优势的潜力尚未得到充分探索。在本研究中，我们弥合了这一差距，提出了 NorMuon（Neuron-wise Normalized Muon），一种协同结合了正交化与神经元级别自适应学习率的优化器。我们的分析表明，虽然 Muon 有效地降低了条件数，但由此产生的更新在神经元范数上高度不均匀，导致某些神经元主导了优化过程。NorMuon 通过为每个神经元维护二阶动量统计量，并在正交化后进行行向归一化，来解决这种不平衡，从而在保持 Muon 条件数优势的同时，确保了参数的均衡利用。为了实现大规模的实际部署，我们在 FSDP2 框架下开发了一个高效的分布式实现，该实现策略性地将正交化计算分布到不同设备上。跨越多种模型规模的实验表明，NorMuon 在 1.1B 预训练设置下，持续优于 Adam 和 Muon，在训练效率方面比 Adam 提升了 21.74%，比 Muon 提升了 11.31%，同时保持了与 Muon 相当的内存占用。我们的研究结果表明，正交化和自适应学习率是互补而非竞争的方法，为大规模深度学习中的优化器设计开辟了新的途径。