Memory-Efficient Subspace Optimization with Residual Correction: From Momentum SGD to Muon
Training large language models is memory-intensive, and subspace projection methods like GaLore can diverge when stochastic noise dominates the selected low-rank gradient space. This paper proposes SMRC and PMRC, which keep optimizer states in a low-dimensional subspace while adding the orthogonal residual gradient back into the parameter update. The authors prove $O(T^{-1/4})$ convergence rates for momentum SGD and for Muon under a mixed stationarity measure, and demonstrate empirical gains on LLM pretraining and fine-tuning up to $1.3$B parameters.
The paper makes a solid theoretical contribution by closing a convergence gap for gradient-dependent subspace methods with residual correction, and backs it with practical algorithms that empirically outperform prior memory-efficient optimizers. As the authors note, "there still remains a lack of convergence guarantees for memory-efficient subspace optimization with gradient-dependent dominant subspaces under standard assumptions." Their multi-stage SMRC introduces a genuine boundary correction to align virtual iterates across basis refreshes, and PMRC extends the framework to the matrix-sign geometry of Muon. Empirically, "PMRC consistently achieves the lowest training and validation losses among memory-efficient methods" on LLaMA models from $60$M to $1.3$B. The main limitations are that PMRC’s guarantee relies on an unconventional mixed nuclear/Frobenius stationarity measure, multi-stage PMRC omits the boundary correction used for SMRC, and the experiments do not close the gap to full-space Muon.
The technical innovations are well motivated and cleanly executed. For SMRC, the "subspace virtual iterate strategy that absorbs the momentum into an auxiliary trajectory" elegantly resolves the coupling between projected momentum and residual correction, enabling a standard $O(T^{-1/4})$ rate under fixed bases. The extension to multi-stage SMRC is backed by a stage-to-stage boundary correction that preserves consistency of the virtual sequence. For PMRC, the analysis is more intricate because the matrix sign is nonlinear; the authors overcome this by presenting "a weighted momentum tracking analysis that jointly characterizes momentum accumulation and effects induced by the residual correction." The memory reduction from $O(mn)$ to $O((m+n)r)$ is preserved throughout, and the LLM experiments cover four scales and several downstream tasks.
The PMRC convergence guarantee is measured under a mixed stationarity metric $\|P^{\top}\nabla F(W_t)\|_{*} + \|Q\nabla F(W_t)\|_F$, which is weaker than the usual squared-Frobenius gradient norm and complicates direct comparison with full-space Muon guarantees. The scaling of the residual weight is also troubling: Remark 6 states that "the choice $\gamma = \Theta(T^{1/4})$ corresponds to an effective step size $\eta\gamma = \Theta(T^{-1/2})$ for the residual component," meaning the raw residual coefficient grows without bound with the horizon. A more serious structural gap is that multi-stage PMRC (Algorithm 4) transitions between stages by simply setting "$W_{k+1,1} = W_{k,T+1}$ and update $P_{k+1}$," with no boundary correction, whereas multi-stage SMRC explicitly compensates for the virtual-iterate offset. The paper asserts the same $O(N^{-1/4})$ rate for this variant, but the main text offers no explanation for how the momentum reset and basis switch interact with the weighted tracking argument. Finally, the practical PMRC implementation adds a residual cap and approximate SVD procedures that are deferred to Appendix E.2.1 and are not covered by the theory.
The evidence supports the core practical claims, though with caveats. On LLaMA pretraining, PMRC improves over GaLore, GoLore, Fira, and GUM at all scales, and on supervised fine-tuning it leads memory-efficient methods on IFEval and GSM8K: "On the IFEval benchmark, PMRC achieves the best performance among all projection-based methods across the three models." However, full-space Muon still outperforms PMRC (validation loss $2.495$ vs. $2.570$ at $1.3$B), so the memory-efficient variant does not fully close the gap. Comparisons to GUM are complicated by the fact that GUM is run with different projection ranks and a layer-sampling budget tailored to PMRC’s memory usage, rather than a fixed configuration. The authors also provide a careful theoretical comparison to Nguyen and Nguyen (2025), identifying an omitted term in their Lemma D.1 that undermines their bound, which strengthens the positioning of this work.
Reproduction would be challenging. The practical PMRC variant uses heuristics not in the theory: the main text notes that "we adopt a practical PMRC variant that applies a residual cap to improve training stability and refreshes the basis through an efficient approximation of SVD, as detailed in Appendix E.2.1." No code repository or training logs are mentioned in the main text, and the experiments span multiple GPU types (RTX PRO 6000 and H20) with different per-device batch sizes. Theoretical step sizes for SMRC depend on the unknown smoothness constant $L$ and total iteration count $T$ (e.g., $\eta = \min\{\frac{\gamma}{2L}, \frac{1}{L\sqrt{T}}\}$), and the paper does not clarify whether these formulas are implemented exactly or replaced by a grid search in practice, making independent hyperparameter replication non-trivial.
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.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.