Memory-Efficient Subspace Optimization with Residual Correction: From Momentum SGD to Muon

Sifan Yang; Hongyi Tao; Dingzhi Yu; Haomin Bai; Wei Jiang; Silin Gao; Pi Bu; Lijun Zhang

cs.LGOct 14, 2026

Memory-Efficient Subspace Optimization with Residual Correction: From Momentum SGD to Muon

Authors: Sifan Yang, Hongyi Tao, Dingzhi Yu, Haomin Bai, Wei Jiang, Silin Gao, Pi Bu, Lijun Zhang

Organizations: State Key Laboratory of Novel Software Technology, Nanjing University · School of Artificial Intelligence, Nanjing University · Alibaba Token Hub, Alibaba Group · School of Computer Science and Engineering, Nanjing University of Science and Technology

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.

Figures & tables

Loss Type	Memory-Efficient	Optimizer	$\mathbf{60M}$	$\mathbf{150M}$	$\mathbf{350M}$	$\mathbf{1.3B}$
Training Loss	${\color[rgb]{1,0,0}\times}$	AdamW	$3.445$	$3.156$	$2.843$	$2.551$
	${\color[rgb]{1,0,0}\times}$	Muon	$3.383$	$3.086$	$2.703$	$2.414$
	${\color[rgb]{0,0,1}\checkmark}$	GaLore	$3.484$	$3.180$	$2.972$	$2.842$
	${\color[rgb]{0,0,1}\checkmark}$	GoLore	$3.516$	$3.196$	$3.003$	$3.080$
	${\color[rgb]{0,0,1}\checkmark}$	Fira	$3.601$	$3.180$	$\underline{2.852}$	$2.657$
	${\color[rgb]{0,0,1}\checkmark}$	GUM	$\underline{3.461}$	$\underline{3.141}$	$2.871$	$\underline{2.609}$

Table 1: Training and validation losses during pretraining. Bold and underlined values indicate the best and second-best results among memory-efficient methods, respectively.

	IFEval			GSM8K
Optimizer	Qwen3-8B	LLaMA-3.1-8B	OLMo-2-1124-7B	Qwen3-8B	LLaMA-3.1-8B	OLMo-2-1124-7B
AdamW	$65.83$	$50.12$	$38.97$	$90.14$	$72.10$	$77.63$
Muon	$62.11$	$46.16$	$36.45$	$89.76$	$72.78$	$78.01$
GaLore	$66.67$	$47.24$	$\underline{38.13}$	$89.16$	$\underline{71.27}$	$77.86$
GoLore	$\underline{68.35}$	$\underline{47.96}$	$35.73$	$88.93$	$70.66$	$77.10$
GUM	$62.47$	$43.65$	$36.21$	$\underline{89.61}$	$69.74$	$\underline{78.32}$

Table 2: Accuracy (%) on IFEval and GSM8K benchmarks. Bold and underlined values indicate the best and second-best results among memory-efficient methods, respectively.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 1 — Figure 1: Effect of residual correction.

Figure 2: Results of multi-stage refreshing.

Matrix Shape	Exact SVD (ms)	Ours (ms)
$1024\times 1024$	$58.01$	$3.21$
$2736\times 1024$	$57.75$	$2.98$
$1024\times 2736$	$57.36$	$2.98$

Appendix

Table 3: Per-matrix projector latency on the LLaMA-350M matrix shapes on a single NVIDIA RTX PRO 6000 GPU.

Model Size	Hidden Size	Layers	Heads	FFN Size	Vocabulary Size	Max Sequence Length
$60\mathrm{M}$	$512$	$8$	$8$	$1376$	$32000$	$1024$
$150\mathrm{M}$	$768$	$14$	$12$	$2048$	$32000$	$2048$
$350\mathrm{M}$	$1024$	$24$	$16$	$2736$	$32000$	$2048$
$1.3\mathrm{B}$	$2048$	$24$	$32$	$5461$	$32000$	$1024$

Appendix

Table 4: Architecture configurations of the LLaMA-style models.

Hyperparameter	$60\mathrm{M}$	$150\mathrm{M}$	$350\mathrm{M}$	$1.3\mathrm{B}$
Global Batch Size	$512$	$512$	$640$	$1280$
Device Micro-batch Size	$256$	$256$	$160$	$160$
Sequence Length	$512$	$512$	$1024$	$1024$
Tokens per Step	$262144$	$262144$	$655360$	$1310720$
Training Steps	$4431$	$9900$	$11300$	$20480$
Total Training Tokens	$1.162$ B	$2.595$ B	$7.406$ B	$26.84$ B

Appendix

Table 5: Training configurations of the LLaMA-style models.

Model	AdamW	Muon	GaLore	GoLore	Fira	GUM	PMRC
1.3B	$57.77$	$55.52$	$54.00$	$54.00$	$54.00$	$53.85$	53.44

Appendix

Table 6: Peak allocated GPU memory (GiB) during

1.3

B pretraining.

Model	AdamW	Muon	GaLore	GoLore	Fira	GUM	PMRC
	lr	lr	lr	lr	lr	lr	lr	$\gamma$
LLaMA- $60$ M	$3\times 10^{-3}$	$3\times 10^{-3}$	$3\times 10^{-3}$	$5\times 10^{-3}$	$2\times 10^{-3}$	$3\times 10^{-3}$	$2\times 10^{-3}$	$8192$
LLaMA- $150$ M	$2\times 10^{-3}$	$3\times 10^{-3}$	$5\times 10^{-3}$	$5\times 10^{-3}$	$5\times 10^{-3}$	$5\times 10^{-3}$	$2\times 10^{-3}$	$16384$
LLaMA- $350$ M	$1\times 10^{-3}$	$2\times 10^{-3}$	$5\times 10^{-3}$	$5\times 10^{-3}$	$5\times 10^{-3}$	$3\times 10^{-3}$	$2\times 10^{-3}$	$16384$
LLaMA- $1.3$ B	$1\times 10^{-3}$	$1\times 10^{-3}$	$1\times 10^{-3}$	$1\times 10^{-3}$	$1\times 10^{-3}$	$1\times 10^{-3}$	$5\times 10^{-3}$	$16384$

Appendix

Table 7: Hyperparameter settings for different optimizers and model sizes.

Model	Optimizer	ARC-E	ARC-C	OBQA	HellaSwag	PIQA	SIQA	WinoGrande	Avg.
	AdamW	$33.67$	$17.92$	$12.80$	$27.15$	$59.74$	$35.98$	$51.07$	$34.05$
	Muon	$35.77$	$17.83$	$13.20$	$27.46$	$60.66$	$35.82$	$49.80$	$34.37$
	GaLore	$34.39$	$18.60$	$14.40$	$26.94$	$58.87$	$34.75$	$52.33$	$\underline{34.32}$
LLaMA - $60$ M	GoLore	$32.79$	$17.75$	$12.20$	$27.03$	$57.83$	$37.10$	$52.01$	$33.82$
	Fira	$34.93$	$18.86$	$12.20$	$26.98$	$59.19$	$35.47$	$50.75$	$34.05$
	GUM	$35.02$	$18.34$	$13.40$	$27.04$	$58.27$	$35.98$	$51.62$	$34.23$

Appendix

Table 8: Zero-shot accuracy (%) on seven benchmarks. Bold and underlined values indicate the best and second-best average accuracies among memory-efficient methods, respectively. Averages are computed from unrounded scores.

		Prompt-level Strict Accuracy			Instruction-level Strict Accuracy
Memory-Efficient	Optimizer	Qwen3-8B	LLaMA-3.1-8B	OLMo-2-1124-7B	Qwen3-8B	LLaMA-3.1-8B	OLMo-2-1124-7B
${\color[rgb]{1,0,0}\times}$	AdamW	$56.38$	$39.19$	$26.80$	$65.83$	$50.12$	$38.97$
${\color[rgb]{1,0,0}\times}$	Muon	$52.68$	$34.38$	$24.03$	$62.11$	$46.16$	$36.45$
${\color[rgb]{0,0,1}\checkmark}$	GaLore	$57.49$	$\underline{36.60}$	$\underline{26.62}$	$66.67$	$47.24$	$\underline{38.13}$
${\color[rgb]{0,0,1}\checkmark}$	GoLore	$\underline{58.96}$	$32.90$	$24.03$	$\underline{68.35}$	$\underline{47.96}$	$35.73$
${\color[rgb]{0,0,1}\checkmark}$	GUM ( $r=128$ )	$51.76$	$33.46$	$24.95$	$62.47$	$43.65$	$36.21$

Appendix

Table 9: Accuracy (%) on the IFEval (instruction-following) benchmark. Bold and underlined values indicate the best and second-best results among memory-efficient methods, respectively.

Task	Model	AdamW	Muon	GaLore	GoLore	GUM	PMRC
IFEval	Qwen3-8B	$89.16$	$76.23$	$67.71$	$67.71$	$66.01$	65.88
	LLaMA-3.1-8B	$78.32$	$65.32$	$56.63$	$56.63$	$54.90$	54.81
	OLMo-2-1124-7B	$69.74$	$57.43$	$49.23$	$49.23$	$47.79$	47.75
GSM8K	Qwen3-8B	$79.46$	$66.52$	$58.02$	$58.02$	$57.00$	56.18
	LLaMA-3.1-8B	$74.98$	$61.98$	$53.30$	$53.30$	$52.35$	51.48
	OLMo-2-1124-7B	$66.77$	$54.71$	$46.54$	$46.54$	$45.80$	45.03

Appendix

Table 10: Peak allocated GPU memory (GiB) during full-parameter SFT. We use a micro-batch size of

16

with

8

gradient accumulation steps for IFEval and a micro-batch size of

4

with

32

gradient accumulation steps for GSM8K. All memory measurements are performed on a single H20 GPU.

Panel of Figure 3 — Figure 3: Effect of the subspace rank $r$ on PMRC.

Panel of Figure 4 — Figure 4: Effect of the residual weight $\gamma$ on PMRC.

Figure 5: Effect of the residual cap on PMRC.

Subspace Optimization

See all 12 papers

Sep 7, 2026

MpSub: A Momentum $p$ -Dimensional Subspace Trust-Region Method for Derivative-Free Fine-Tuning of Large Language Models

Full-parameter fine-tuning of large language models has substantial memory costs because backpropagation stores activations and gradients. Zeroth-order optimization avoids this by estimating update directions from loss evaluations, but existing methods require tuning a sensitive learning rate for each model and task. We propose the momentum

p

-dimensional subspace trust-region method (MpSub). At each iteration, MpSub searches within a

p

-dimensional subspace: one direction preserves historical momentum from the most recent accepted step, while the remaining directions explore via fresh random sampling. The subspace gradient is estimated by central differences, a trial step is computed from a linear trust-region model, and the trust-region radius adapts according to the agreement between predicted and observed loss reduction, eliminating the learning rate. For LLM fine-tuning, evaluations within an iteration share a minibatch, and directions are regenerated in place from seeds, using forward passes alone. For smooth deterministic objectives under unorthogonalized Gaussian directions, we bound the finite-difference error, quantify gradient energy captured by the subspace, and prove that

\lim_{k\to\infty} \|\nabla f(x_k)\|_2 = 0

almost surely under a safeguarded radius update. Under a matched budget of 8,400 training-objective forward passes, we fine-tune OPT-125M and OPT-350M on CommitmentBank. With the same preset parameters at both model sizes, MpSub attains mean test accuracies of 0.673 and 0.690 over three seeds, matching tuned MeZO (0.685) without any learning-rate search.

Yuyang Wang, Haoyu Yao, Pengcheng Xie

Xi’an Jiaotong University, Xi’an, China · Lawrence Berkeley National Laboratory, Berkeley, CA, USA

Memory-Efficient Fine-Tuning Zeroth-Order Optimization LLM Fine-Tuning Trust-Region Optimization

Ask PDF

Jun 2, 2026

Denoise First, Orthogonalize Later: Understanding Momentum in Muon via Spectral Filtering

Muon has recently demonstrated strong empirical performance in large language model training, but the theoretical role of momentum in Muon remains unclear. Existing analyses of Muon either remove momentum to study spectral updates in isolation, or retain momentum without explaining why it improves empirical performance. Our work bridges this gap by showing momentum in Muon acts as a spectral filter. Under a structured signal-plus-perturbation gradient model, we prove that momentum suppresses perturbations while preserving the dominant signal, thereby enlarging the spectral gap between them. This enlarged gap stabilizes the singular subspaces of the matrix passed to Muon's orthogonalization step, making the resulting update more reliable. We further show that applying momentum before orthogonalization achieves provably stronger alignment with the signal component of the gradient than either reversing this order or simply removing momentum. Experiments across diverse tasks, including LLM pretraining, support our theoretical analysis. More broadly, our theory offers a starting point for understanding the benefits of momentum in other matrix-based optimizers.

Xianliang Li, Zihan Zhang, Weiyang Liu +1

1The Institute of Statistical Mathematics · 2The Graduate Institute for Advanced Studies, SOKENDAI · 3National Institute of Informatics +3

Deep Learning Optimization Momentum Methods Muon Optimizer Spectral Methods Gradient Descent Dynamics

Ask PDF

Mar 6, 2026

Omni-Masked Gradient Descent: Memory-Efficient Optimization via Mask Traversal with Improved Convergence

Memory-efficient optimization methods have recently gained increasing attention for scaling full-parameter training of large language models under the GPU-memory bottleneck. Existing approaches either lack clear convergence guarantees, or only achieve the standard

{\mathcal{O}}(ε^{-4})

iteration complexity in the nonconvex settings. We propose Omni-Masked Gradient Descent (OMGD), an optimization method based on mask traversal for memory efficient training, and provide a nonconvex convergence analysis that establishes a strictly improved iteration complexity of

\tilde{\mathcal{O}}(ε^{-3})

for finding an

ε

-approximate stationary point. Empirically, OMGD is a lightweight, plug-and-play approach that integrates seamlessly into most mainstream optimizers, yielding consistent improvements over competitive baselines in both fine-tuning and pre-training tasks.

Hui Yang, Tao Ren, Jinyang Jiang +2

1Guanghua School of Management, Peking University · 2Wangxuan Institute of Computer Technology, Peking University.

Gradient Descent Memory-Efficient Training Nonconvex Optimization

Ask PDF

Abstract

Subspace Optimization

Explore similar work

MpSub: A Momentum ppp-Dimensional Subspace Trust-Region Method for Derivative-Free Fine-Tuning of Large Language Models

Denoise First, Orthogonalize Later: Understanding Momentum in Muon via Spectral Filtering

Omni-Masked Gradient Descent: Memory-Efficient Optimization via Mask Traversal with Improved Convergence

MpSub: A Momentum $p$ -Dimensional Subspace Trust-Region Method for Derivative-Free Fine-Tuning of Large Language Models