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

MpSub: A Momentum $p$ -Dimensional Subspace Trust-Region Method for Derivative-Free Fine-Tuning of Large Language Models
Sep 7, 2026Yuyang Wang, Haoyu Yao, Pengcheng XieMemory-Efficient Fine-Tuning Zeroth-Order Optimization
Xi’an Jiaotong University, Xi’an, China · Lawrence Berkeley National Laboratory, Berkeley, CA, USA
Denoise First, Orthogonalize Later: Understanding Momentum in Muon via Spectral Filtering
Jun 2, 2026Xianliang Li, Zihan Zhang, Weiyang Liu +1Deep Learning Optimization Momentum Methods
1The Institute of Statistical Mathematics · 2The Graduate Institute for Advanced Studies, SOKENDAI · 3National Institute of Informatics +3
Omni-Masked Gradient Descent: Memory-Efficient Optimization via Mask Traversal with Improved Convergence
Mar 6, 2026Hui Yang, Tao Ren, Jinyang Jiang +2Gradient Descent Memory-Efficient Training
1Guanghua School of Management, Peking University · 2Wangxuan Institute of Computer Technology, Peking University.

Sep 7, 2026J/K move · Enter open · S save

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

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

Open paper AskPDF arXiv

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.

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