通过低秩正交化提升大模型训练效率,显著优化了基础模型训练性能。
Low-rank Orthogonalization for Large-scale Matrix Optimization with Applications to Foundation Model Training
- 利用梯度的低秩特性实现矩阵正交化,降低计算开销。
- 低秩Muon在GPT-2和LLaMA预训练中超越调优后的原始Muon。
- 理论证明了算法在重尾噪声下的收敛性,适合大规模模型训练者。
神经网络训练本质上是大规模矩阵优化问题,但长期以来参数矩阵结构被忽视。近期优化器Muon因其在基础模型训练中的优异表现受到关注,其成功关键在于矩阵正交化。本文提出低秩正交化方法,利用训练中梯度的低秩特性进行正交化。基于此,我们引入低秩矩阵符号梯度下降(MSGD)及低秩版Muon。数值实验表明,低秩正交化性能优越,低秩Muon在GPT-2和LLaMA预训练中表现优异,优于经精心调优的原始Muon,尤其在大模型场景下。理论上,我们建立了低秩MSGD寻找近似驻点的迭代复杂度,以及低秩Muon在重尾噪声下寻找近似随机驻点的复杂度。代码已开源:https://github.com/dengzhanwang/Low-rank-Muon。
原文摘要 · Abstract (English)
Neural network (NN) training is inherently a large-scale matrix optimization problem, yet the matrix structure of NN parameters has long been overlooked. Recently, the optimizer Muon \citep{jordanmuon}, which explicitly exploits this structure, has gained significant attention for its strong performance in foundation model training. A key component contributing to Muon's success is matrix orthogonalization. In this paper, we propose \textit{low-rank orthogonalization}, which performs orthogonalization by leveraging the low-rank nature of gradients during NN training. Building on this, we introduce low-rank matrix-signed gradient descent (MSGD) and a low-rank variant of Muon. Numerical experiments demonstrate the superior performance of low-rank orthogonalization, with low-rank Muon achieving promising results in GPT-2 and LLaMA pretraining -- surpassing the carefully tuned vanilla Muon on tasks with large model sizes. Theoretically, we establish the iteration complexity of low-rank MSGD for finding an approximate stationary solution, and the iteration complexity of low-rank Muon for finding an approximate stochastic stationary solution under heavy-tailed noise. The code to reproduce our numerical experiments is available at https://github.com/dengzhanwang/Low-rank-Muon.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。