轻量快速的优化器，专为大模型训练设计
LiMuon: Light and Fast Muon Optimizer for Large Models
- 基于动量与随机SVD，降低内存和样本复杂度
- 理论证明其样本复杂度达O(ε⁻³)，优于现有方法
- 适合大规模模型训练，尤其在资源受限场景
大模型在机器学习中广泛应用，高效训练备受关注。近期提出的Muon优化器针对大模型中的矩阵参数结构设计，但现有版本仍存在高样本复杂度或高内存开销问题。为此，我们提出轻量快速的LiMuon优化器，结合基于动量的方差缩减技术与随机奇异值分解（SVD）。LiMuon在保持更低内存的同时，样本复杂度显著低于原有Muon及其变体。理论上，在广义光滑条件下，其寻找非凸随机优化ε-平稳解的样本复杂度为O(ε⁻³)。进一步地，采用Newton-Schulz迭代步骤时，其样本复杂度也优于原Muon。在Mamba-130M、Qwen2.5-0.5B和ViT模型的预训练实验中，结果验证了LiMuon的有效性。
原文摘要 · Abstract (English)
Large models recently are widely applied in machine learning, so efficient training of large models has received widespread attention. More recently, the useful Muon optimizer is specifically designed for matrix-structured parameters of large models. Although some works have begun to study the Muon optimizer, the existing Muon and its variants still suffer from high sample complexity or high memory for large models. To fill this gap, we propose a light and fast Muon (LiMuon) optimizer for training large models, which builds on the momentum-based variance reduced technique and randomized Singular Value Decomposition (SVD). In particular, our LiMuon simultaneously has a lower memory and lower sample complexity than the Muon and its variants. Moreover, we prove that our LiMuon with lower memory has a lower sample complexity of $O(ε^{-3})$ for finding an $ε$-stationary solution of non-convex stochastic optimization under the generalized smoothness condition. To further narrow practice and theory gap, we also prove that our LiMuon with Newton-Schulz steps has a lower sample complexity than the Muon with Newton-Schulz steps. Numerical experimental results on pre-training Mamba-130M, Qwen2.5-0.5B and ViT models demonstrate effectiveness of our LiMuon.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。