PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2610.19844v1cs.LGcs.AI

TL;DR

Memory-efficient subspace methods can lack convergence guarantees when noise dominates the selected gradient directions. The paper adds orthogonal residual-gradient correction to momentum SGD and Muon, with analytical strategies for fixed and periodically refreshed bases. Both methods and their multi-stage variants achieve convergence rates of O(T^-1/4), matching the corresponding full-space optimizer rates.

  • Problem

    Gradient noise can dominate leading singular subspaces, and convergence guarantees remain lacking for memory-efficient methods using gradient-dependent dominant subspaces under standard stochastic assumptions.

    Source: §1 · p. 2

  • Method

    SMRC and PMRC keep optimizer states in projected subspaces and add the orthogonal residual gradient to parameter updates.

    Source: §1 · p. 2

  • Results

    SMRC and PMRC achieve O(T^-1/4) convergence rates matching full-space momentum SGD and Muon, respectively, with multi-stage variants retaining the same rate order.

    Source: §1 · p. 2, §7 · p. 10

  • Takeaways & Limitations

    The framework provides convergence guarantees for memory-efficient subspace optimization with refreshed bases, including a boundary correction for multi-stage SMRC.

    Source: §1 · p. 2, §7 · p. 10

    The analysis assumes a lower-bounded loss and, for multi-stage bases, that each basis is determined by optimization history before its stage.

    Source: §3 · p. 4, §4.2 · p. 6

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper addresses convergence gaps in memory-efficient subspace optimization when stochastic-gradient noise dominates the selected directions. It introduces residual-corrected momentum SGD and Muon methods, proves convergence rates matching their full-space counterparts, and reports favorable practical comparisons for PMRC.

1 INTRODUCTION

  • Gradient noise can dominate leading singular directions later in training, leaving subspace methods without convergence guarantees under standard stochastic assumptions.Random projections have guarantees but do not exploit dominant gradient directions and are less effective in practice.Source: §1 · p. 2
  • SMRC and PMRC achieve O(T^-1/4) convergence rates matching full-space momentum SGD and Muon, respectively.Their multi-stage variants retain the same rate order.Source: §1 · p. 2
  • Residual correction maintains optimizer states in a low-dimensional subspace and adds the orthogonal residual gradient to the update.SMRC uses a subspace virtual iterate to address coupling, while PMRC uses weighted momentum tracking for the residual gradient.Source: §1 · p. 2
  • PMRC performs favorably against prior memory-efficient methods across LLaMA pretraining at four scales from 60M to 1.3B parameters and six supervised finetuning settings.Synthetic experiments examine SMRC residual correction and periodic basis refreshes.Source: §1 · p. 2

2 RELATED WORK

The related work surveys parameter-efficient and subspace-based memory reduction, then reviews Muon and memory-efficient variants. It highlights limitations in convergence guarantees for gradient-dependent subspaces and in worst-case memory savings for some Muon variants.

2.1 MEMORY-EFFICIENT OPTIMIZATION FOR LLM TRAINING

  • LoRA reduces finetuning memory with low-rank updates, but its update expressiveness is limited in pretraining.GaLore instead reduces memory by maintaining optimizer states in a low-dimensional gradient subspace.Source: §2.1 · p. 3
  • GaLore stores optimizer states in a low-dimensional gradient subspace, but noisy leading singular directions can miss informative gradients and prevent convergence under standard assumptions.GoLore guarantees convergence with random projections, but these do not favor dominant directions and lead to slow convergence.Source: §2.1 · p. 3
  • Fira adds a norm-scaled discarded-residual correction but lacks a convergence guarantee; other subspace methods have guarantees limited by basis, update, or assumption conditions.The cited analyses do not directly cover periodically refreshed SVD subspaces, or require stronger assumptions and an unweighted residual update.Source: §2.1 · p. 3

2.2 MUON AND ITS VARIANTS

  • Muon applies the matrix sign of momentum for matrix-shaped parameters but retains O(mn) optimizer-state memory.Its matrix-sign direction is steepest descent under the spectral norm and is approximated through Newton–Schulz iterations.Source: §2.2 · p. 3
  • GUM combines low-rank projection with layerwise sampling, but can lose memory reduction over Muon when all layers receive full-rank updates.The related-work section directs readers to Appendix A.3 for additional Muon variants.Source: §2 · p. 2, §2.2 · p. 3

3 NOTATION AND ASSUMPTIONS

This section sets matrix notation, defines projection and norm conventions, and states assumptions for the stochastic nonconvex analysis. The stated conditions include smoothness, a lower-bounded loss, stochastic gradients, and a history-dependent orthonormal basis for stage k.

3 NOTATION AND ASSUMPTIONS

  • The analysis uses an orthonormal basis P, with PP^⊤ projecting onto span(P) and Q = I − PP^⊤ projecting onto its orthogonal complement.Matrices use the trace inner product and Frobenius, spectral, and nuclear norms.Source: §3 · pp. 3–4
  • The loss is assumed L-smooth and lower bounded, with initial optimality gap Δ = F(W1) − F*.The lower-boundedness condition is stated as F* > −∞.Source: §3 · p. 4
  • For multi-stage analysis, each orthonormal basis Pk is determined by optimization history prior to stage k.The section also introduces standard stochastic nonconvex optimization assumptions.Source: §3 · p. 4

4 SUBSPACE MOMENTUM SGD WITH RESIDUAL CORRECTION (SMRC)

SMRC combines projected momentum with an orthogonal residual-gradient correction, resolving their analytical coupling through a virtual iterate. Its multi-stage extension refreshes bases and corrects stage boundaries while retaining the convergence order of full-space momentum SGD.

4.1 SINGLE-STAGE SMRC

  • The update combines mapped-back projected momentum with γ-weighted orthogonal residual gradient, recovering directions that a noisy subspace may miss.The residual correction requires no additional optimizer state, and total basis-plus-momentum memory is O((m + n)r), versus O(mn) for full-space momentum.Source: §4.1 · pp. 4–5
  • A subspace virtual iterate absorbs projected momentum into an auxiliary trajectory, decoupling it analytically from residual correction across iterations.The virtual sequence is used only in the analysis, not computed by SMRC.Source: §4.1 · p. 5
  • For any β ∈ (0,1), γ ∈ (0,1] and the step size prescribed in Theorem 1, SMRC attains O(T^-1/4) convergence, matching full-space momentum SGD while storing only low-rank momentum.SMRC maintains momentum m_t ∈ R^(r×n) in projected space rather than ambient space R^(m×n). Storing basis P ∈ R^(m×r) and momentum requires O((m+n)r) memory, versus O(mn) for full-space momentum SGD.Source: §4.1 · p. 5

4.2 MULTI-STAGE SMRC

  • Multi-stage SMRC periodically refreshes bases to track the evolving dominant gradient subspace while using a fixed basis within each stage.The theorem allows any sequence of orthonormal stage bases, including bases from dominant stochastic-gradient subspaces.Source: §4.2 · pp. 5–6
  • Boundary correction aligns consecutive virtual iterates after a basis refresh, resolving the inconsistency caused by directly carrying over the terminal parameter iterate.The correction modifies the iterate at each stage transition.Source: §4.2 · pp. 5–6
  • Multi-stage SMRC achieves O(N^-1/4) convergence over N iterations, matching the convergence order of full-space momentum SGD.The guarantee holds for any sequence of orthonormal stage bases under the stated assumptions.Source: §4.2 · p. 6

5 PROJECTED MUON WITH RESIDUAL CORRECTION (PMRC)

PMRC extends residual correction to Muon by applying the matrix sign to projected momentum and adding the orthogonal residual gradient. Single- and multi-stage analyses establish convergence rates of the same order as full-space Muon.

5.1 SINGLE-STAGE PMRC

  • PMRC projects stochastic gradients, maintains low-rank momentum, applies the matrix sign, and adds the orthogonal residual gradient to the update.The residual weight γ controls the correction term.Source: §5.1 · p. 6
  • Weighted momentum tracking accounts for momentum accumulation and iterate variation from residual correction across PMRC's different norm geometries.The matrix-sign step is characterized by the nuclear norm, while the residual gradient is measured in the Frobenius norm.Source: §5.1 · p. 6
  • Under standard stochastic assumptions, single-stage PMRC achieves an O(T^-1/4) rate under a mixed stationarity measure, matching full-space Muon's convergence order.With a full basis, PMRC reduces to Muon, and the mixed measure upper-bounds the full gradient's Frobenius norm.Source: §5.1 · p. 7

5.2 MULTI-STAGE PMRC

  • Multi-stage PMRC refreshes the projection basis between fixed-basis stages, resets momentum in the new projected coordinates, and retains an O(N^-1/4) rate.N denotes the total number of stochastic-gradient evaluations.Source: §5.2 · pp. 7–8

6 EMPIRICAL STUDY

The empirical study evaluates PMRC in LLaMA pretraining and supervised finetuning. PMRC reports favorable results against memory-efficient methods across both settings.

6 EMPIRICAL STUDY

  • The study uses a practical PMRC variant with a residual cap for training stability and an efficient approximate SVD basis refresh.Pretraining uses LLaMA models from 60M to 1.3B parameters; the finetuning experiments cover three models and two tasks.Source: §6 · p. 8, §6.1 · p. 8, §6.2 · p. 9

6.1 LLM PRETRAINING

  • PMRC achieves the lowest training and validation losses among memory-efficient methods and lower losses than AdamW at every pretraining scale, while using less optimizer-state memory.It also achieves the highest average accuracy among memory-efficient methods on seven downstream benchmarks.Source: §6.1 · p. 9

6.2 LLM SUPERVISED FINETUNING

  • On IFEval, PMRC performs best among projection-based methods across three models and beats GUM by more than 2 percentage points; on GSM8K, it leads memory-efficient methods and remains competitive with full-space optimizers.The reported metrics are instruction-level strict accuracy on IFEval and zero-shot accuracy on GSM8K.Source: §6.2 · pp. 9–10, Table 2 · p. 9

7 CONCLUSION

The paper combines residual correction with subspace optimization for momentum SGD and Muon. Both methods reduce optimizer-state memory while retaining convergence rates of the same order as their full-space counterparts.

7 CONCLUSION

  • SMRC combines projected momentum with orthogonal residual correction, while PMRC applies the framework to Muon; both attain convergence rates of the same order as their full-space counterparts.Multi-stage variants retain that rate order, with boundary correction aligning consecutive SMRC virtual iterates.Source: §7 · p. 10
  • Both SMRC and PMRC reduce optimizer-state memory from O(mn) to O((m + n)r).The paper reports experiments on LLM pretraining and supervised finetuning demonstrating practical effectiveness and efficiency.Source: §7 · p. 10

A.1 ADDITIONAL MEMORY-EFFICIENT OPTIMIZATION METHODS

The appendix surveys memory-efficient optimization approaches and presents proof details for the paper's single- and multi-stage analyses. It also contrasts prior work and records how the proofs handle projected momentum and stage transitions.

A.1 ADDITIONAL MEMORY-EFFICIENT OPTIMIZATION METHODS

  • Prior approaches include sparse-plus-low-rank updates, layer resampling, random or structured projections, online subspace estimation, and compressed optimizer states.The comparison notes that Nguyen & Nguyen's guarantee uses a fixed projection and unit residual weight, and does not cover subspace switching or Muon matrix-sign updates.Source: §A.1 · p. 15

A.2 LOW-RANK ADAPTATION (LORA)

  • LoRA freezes pretrained weights and trains low-rank adapter matrices; later variants add adaptive rank allocation, periodic merging, or weight quantization.Source: §A.2 · p. 15

A.3 MATRIX-AWARE AND MUON-BASED OPTIMIZERS

  • Muon-related memory methods include subspace momentum orthogonalization, randomized-SVD gradient compression, and 8-bit or 4-bit optimizer-state quantization.SUMO's convergence guarantee is stated only for reversible layers.Source: §A.3 · p. 16

B.1 PROOF OF THEOREM 1

  • The single-stage proof defines projected momentum recursively and uses complementary orthogonal projectors whose combined operator has eigenvalues 1 and γ.Source: §B.1 · p. 16

B.2 PROOF OF THEOREM 2

  • The multi-stage proof defines stage-specific virtual iterates and chooses the next stage's initialization to align with the preceding virtual iterate.The basis for each stage is determined by optimization history before that stage.Source: §B.2 · p. 18

B.3 PROOF OF THEOREM 3

The proof uses corrected-momentum tracking quantities for the fixed-basis analysis and extends the bounds to stages with separately selected bases. Its lemmas and parameter choices support the theorem bounds under stated conditions.

B.3 PROOF OF THEOREM 3

  • The fixed-basis proof defines projected and residual gradients, bias-corrected momentum, and a tracking error against projected true gradients.These quantities are introduced for the analysis of Theorem 3.Source: §B.3 · p. 19
  • Under Assumptions 1, 2, and 3, Lemma 5 applies when Lηγ ≤ 1.Source: §B.3 · p. 20
  • For Theorem 3, the proof chooses β, η, and γ to satisfy Lηγ ≤ 1 before applying the preceding bound.Source: §B.3 · p. 20

B.4 PROOF OF THEOREM 4

  • The multi-stage proof uses a predictable orthonormal basis and complementary projection at each stage, then sums bounds across stages.It defines N=KT for K stages of length T.Source: §B.4 · p. 21

C.2 PROOF OF LEMMA 4

  • The projected and residual gradient components are orthogonal because P⊤Q=0.Source: §C.2 · p. 23

C.4 PROOF OF LEMMA 6

  • Under Assumption 1, the cited passage gives the bound ∥A_i−A_t∥_* ≤ √r L∥W_i−W_t∥_F for 1≤i≤t, relating differences in A to differences in W measured in Frobenius norm.Source: §C.4 · p. 24

C.5 PROOF OF LEMMA 7

The proof combines bounds using weighted Young’s inequality and Assumption 1, but the residual SGD contribution prevents the direct decomposition from yielding the desired rate.

C.5 PROOF OF LEMMA 7

  • The proof bounds a term using weighted Young’s inequality and then applies Assumption 1.Source: §C.5 · p. 25
  • The residual SGD correction contributes an undesirable term that causes the direct decomposition argument to break down.Source: §C.5 · p. 25
  • Summing over t=2,...,T+1 and rearranging yields Eq. (26).Source: §C.5 · p. 25

D LIMITATIONS OF A DIRECT DECOMPOSITION ANALYSIS

The analysis explains why a direct split between projected momentum and residual descent does not yield the desired stochastic convergence rate. Experiments show benefits from residual correction and periodic basis refresh, alongside favorable PMRC results in pretraining and finetuning.

D LIMITATIONS OF A DIRECT DECOMPOSITION ANALYSIS

  • The direct decomposition into projected momentum tracking and orthogonal residual SGD descent does not establish the desired convergence rate.The residual correction enters the tracking-error bound through the iterate step, complicating the combination of the two bounds.Source: D LIMITATIONS OF A DIRECT… · pp. 26–27
  • Under the proposed scaling, the residual-noise variance contribution does not vanish, preventing this bound from establishing an O(T^-1/2) rate under stochastic gradients.The true-gradient component can be absorbed into descent, but bounding the noise component by σ^2 leaves the nonvanishing term.Source: D LIMITATIONS OF A DIRECT… · p. 27

E.1.1 EFFECT OF RESIDUAL CORRECTION

  • Residual correction restores SMRC convergence when GaLore fails under a noise-dominated selected subspace, bringing performance close to Momentum SGD.The synthetic experiment attributes GaLore's failure to informative gradient directions remaining outside its selected subspace.Source: §E.1.1 · p. 28

E.1.2 BENEFIT OF MULTI-STAGE SUBSPACE REFRESHING

  • Periodic basis refresh makes multi-stage SMRC converge faster than single-stage SMRC and perform comparably to full-space Momentum SGD as dominant directions evolve.Single-stage SMRC continues to progress through residual correction, but slows as the dominant gradient directions change.Source: §E.1.2 · p. 29

E.2.1 A PRACTICAL IMPLEMENTATION OF PMRC

  • The practical PMRC implementation caps residual updates when their scale exceeds the projected Muon direction and uses warm-started basis approximation to reduce refresh latency.On LLaMA-350M matrix shapes, approximate refresh takes around 3 milliseconds versus nearly 58 milliseconds for GaLore's exact SVD.Source: §E.2.1 · pp. 29–30, Table 3 · p. 30

E.2 LANGUAGE MODEL EXPERIMENTS

  • The LLM experiments evaluate PMRC in pretraining on C4 and in supervised finetuning for instruction-following and mathematical reasoning.Pretraining spans four LLaMA-style model scales; the finetuning setup includes three model families.Source: §E.2.2 · p. 30

E.2.2 LLM PRETRAINING TASKS

  • PMRC achieves the highest average accuracy among memory-efficient methods across all tested pretraining model scales.The experiments cover LLaMA-style models from 60M to 1.3B parameters.Source: §E.2.2 · pp. 30–31

E.2.3 LLM FINETUNING TASKS

  • PMRC outperforms GUM at both tested ranks on both IFEval metrics across three models and requires the lowest peak memory across all models.The additional results include prompt-level and instruction-level strict accuracy.Source: §E.2.3 · p. 31

E.2.4 ABLATION STUDIES

The ablations test how PMRC's subspace rank, residual weight, and residual cap affect training. They use LLaMA-150M pretrained on C4 for 9,900 steps, with basis refreshes every 100 steps and a default residual cap of 4.

Effect of the subspace rank

  • Increasing the subspace rank lowers validation loss but increases the memory required for the projection basis and projected momentum.The memory cost scales as O((m + n)r), exposing a trade-off between validation performance and optimizer-state memory.Source: §E.2.4 · pp. 31–32, Figure 3 · p. 33

Effect of the residual weight

  • Across the tested residual weights, validation loss decreases monotonically, with diminishing improvements at larger values; γ = 0 gives the highest validation loss.The residual cap is active, so γ is an upper bound on the effective weight rather than its exact value.Source: §E.2.4 · p. 32, Figure 4 · p. 33

Effect of the residual cap

  • Without a residual cap, training becomes unstable early, whereas the cap limits excessively large residual updates while retaining γ below the threshold.The comparison uses ρ = 4 versus no cap, with rank 256 and γ = 8192 in both runs.Source: §E.2.4 · pp. 33–34, Figure 5 · p. 34
Loading 2610.19844v1…