Memory-efficient methods improve training large language models
Memory-Efficient Subspace Optimization with Residual Correction: From Momentum SGD to Muon
Machine LearningArtificial Intelligence
Summary
Training large language models uses a lot of memory and computing power. This paper describes new methods that reduce memory needs by focusing on important parts of the gradient during training while still making progress efficiently. The authors introduce techniques to correct errors caused by ignoring smaller parts of the gradient and update their approach over time to keep improving. Their methods match the performance of full-memory training but use much less memory, confirmed by tests. This helps train big AI models more efficiently.
What this means in practice
- •For machine learning engineers: Reduce memory requirements when training large language models by applying subspace optimization with residual correction to maintain model accuracy.
- •For ai infrastructure teams: Implement memory-efficient update schemes for large-scale distributed training to lower hardware costs and improve resource utilization.
Authors
Sifan Yang, Hongyi Tao, Dingzhi Yu, Haomin Bai, Wei Jiang, Silin Gao, Pi Bu, Lijun Zhang
Abstract
Memory-efficient subspace optimization has emerged as a practical approach to reducing the memory cost of LLM training while preserving optimization performance. Previous algorithms typically rely on the leading singular subspace of stochastic gradients, yet the gradient noise may dominate the subspace in later training stages and can prevent convergence. To overcome this limitation, we adopt a residual correction framework, which maintains the optimizer states in the subspace and adds the orthogonal residual gradient to the parameter update. We first integrate this scheme into momentum SGD to develop SMRC, which attains an $O(T^{-1/4})$ convergence rate for a fixed orthonormal basis, matching that of full-space momentum SGD. Our analysis introduces a subspace virtual iterate to resolve the coupling between the momentum and the residual correction. Furthermore, we extend SMRC to a multi-stage variant which periodically recomputes bases to track the evolving dominant subspace and retains the same rate. Since each basis refresh breaks the consistency of the virtual iterates, we devise a boundary correction approach that aligns them across stages. Beyond momentum SGD, we incorporate our framework into Muon and propose PMRC, which applies the matrix sign to the projected momentum. Through a weighted momentum tracking analysis that accounts for the residual gradient, we establish an $O(T^{-1/4})$ convergence rate for PMRC, which matches the rate of full-space Muon. We also design a multi-stage PMRC with periodically refreshed bases, which maintains a convergence rate of the same order. Extensive experimental results confirm the effectiveness and efficiency of our algorithms.