Momentum Factorized SGD (MoFaSGD)

• MoFaSGD is a memory-efficient optimizer that replaces full momentum buffers with an online low-rank SVD approximation of the first moment.
• It projects new gradients onto the tangent space of the current low-rank momentum, achieving spectrally normalized updates without full-matrix SVDs.
• Empirical results show significant GPU memory savings and competitive performance compared to AdamW and other low-rank techniques in fine-tuning large models.

Momentum Factorized SGD (MoFaSGD) is a memory-efficient first-order optimizer for large-scale fine-tuning that replaces full momentum buffers with a dynamically maintained low-rank singular value decomposition (SVD) of the first moment and uses the normalized singular directions of that low-rank momentum for parameter updates. It was introduced in "Low-rank Momentum Factorization for Memory Efficient Training" (Mahdavinia et al., 10 Jul 2025) for the regime in which optimizer-state memory, especially the first and second moments of AdamW, dominates GPU memory during fine-tuning. The method targets the trade-off between full-parameter optimization and low memory usage by maintaining an adaptive low-rank approximation to the exponential moving average of gradients, rather than restricting training to fixed low-rank adapters or periodically recomputing gradient subspaces by expensive full-matrix SVDs.

1. Definition and problem setting

MoFaSGD is formulated for matrix parameters $W \in \mathbb{R}^{m\times n}$. In the full-momentum baseline, the first moment is

$M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$

where $G_t = \nabla_W \mathcal{L}(W_t)$, and vanilla momentum SGD would update $W_{t+1}=W_t-\eta M_t$. MoFaSGD instead assumes that the exponential moving average of gradients is approximately low rank and maintains a rank-$r$ SVD approximation

$$\hat M_t \triangleq U_{t+1}\Sigma_{t+1}V_{t+1}^\top \approx \sum_{i=1}^t \beta^{t-i} G_i,$$

with $U_{t+1}^\top U_{t+1}=I_r$, $V_{t+1}^\top V_{t+1}=I_r$, and diagonal $\Sigma_{t+1}$. The rank $r$ is a user hyperparameter analogous to LoRA rank and controls optimizer-state size and approximation fidelity (Mahdavinia et al., 10 Jul 2025).

The motivation is specific to fine-tuning large foundation models. AdamW and similar adaptive optimizers maintain full-rank first and second moment buffers $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$0, so optimizer state is $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$1 per parameter matrix and often several times the inference memory. Existing alternatives in the source material fall into three classes: parameter-efficient fine-tuning such as LoRA, optimizer-state compression such as AdaFactor and SM3, and low-rank subspace methods such as GaLore, Flora, ReLoRA, LDAdam, and APOLLO. MoFaSGD is positioned against the limitations of fixed parameter subspaces, full-gradient SVD resampling, and subspace moment accumulation by treating the first moment itself as the low-rank object and updating its factorization online at every iteration.

2. Low-rank momentum representation and online update

The central computational problem is to update $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$2 without forming a full SVD of $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$3. MoFaSGD addresses this by projecting the new gradient onto the tangent space of the rank-$M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$4 manifold at the previous low-rank momentum approximation. With current factors $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$5, the tangent-space projection is

$M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$6

Conceptually, the next low-rank momentum is obtained from

$M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$7

Because both $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$8 and $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$9 have rank at most $G_t = \nabla_W \mathcal{L}(W_t)$0, their sum has rank at most $G_t = \nabla_W \mathcal{L}(W_t)$1. The algorithm therefore avoids a full decomposition by computing $G_t = \nabla_W \mathcal{L}(W_t)$2, $G_t = \nabla_W \mathcal{L}(W_t)$3, and $G_t = \nabla_W \mathcal{L}(W_t)$4, then forming QR decompositions of the augmented subspaces

$G_t = \nabla_W \mathcal{L}(W_t)$5

and a small core matrix

$G_t = \nabla_W \mathcal{L}(W_t)$6

A rank-$G_t = \nabla_W \mathcal{L}(W_t)$7 SVD of this $G_t = \nabla_W \mathcal{L}(W_t)$8 matrix,

$G_t = \nabla_W \mathcal{L}(W_t)$9

yields the updated factors

$W_{t+1}=W_t-\eta M_t$0

The resulting per-layer complexity is

$W_{t+1}=W_t-\eta M_t$1

coming from two QR decompositions of size $W_{t+1}=W_t-\eta M_t$2 and $W_{t+1}=W_t-\eta M_t$3 and an SVD of size $W_{t+1}=W_t-\eta M_t$4. For $W_{t+1}=W_t-\eta M_t$5, this replaces the expensive full-matrix SVDs required by methods that periodically resample a gradient-defined subspace (Mahdavinia et al., 10 Jul 2025).

3. Optimization geometry and spectrally normalized updates

A defining feature of MoFaSGD is that the parameter update is not the low-rank momentum itself. Once $W_{t+1}=W_t-\eta M_t$6 has been computed, the update is

$W_{t+1}=W_t-\eta M_t$7

Only the singular directions are used; the singular values $W_{t+1}=W_t-\eta M_t$8 do not appear directly in the step. The source material describes this as using the normalized, spectrally whitened directions of the low-rank momentum and as a low-rank analogue of Muon/Shampoo-style gradient whitening (Mahdavinia et al., 10 Jul 2025).

This design gives MoFaSGD two coupled interpretations. First, the factors $W_{t+1}=W_t-\eta M_t$9 and $r$0 define an implicitly evolving optimization subspace. Unlike GaLore, which defines a projection matrix from the current gradient and resamples it every $r$1 steps, MoFaSGD updates its subspace every iteration through the momentum SVD itself. The subspace is therefore momentum-defined, online, and smooth. Second, the update $r$2 is the polar or sign factor of the low-rank momentum restricted to rank $r$3, so the method substitutes directional structure for explicit second-moment accumulation.

This distinction resolves a common confusion. MoFaSGD is neither a fixed-subspace parameter-efficient method nor a straightforward low-rank version of vanilla momentum descent. It does not freeze base weights as LoRA does, and it does not use $r$4. Its state is a low-rank factorization of the first moment, while its step is a spectrally normalized rank-$r$5 direction extracted from that factorization.

4. Memory, computational profile, and empirical performance

For a single matrix $r$6 with $r$7, the source material gives the following memory and subspace-update costs.

Method	Memory	Subspace update
GaLore	$r$8	full SVD of gradient, $r$9
LoRA	$$\hat M_t \triangleq U_{t+1}\Sigma_{t+1}V_{t+1}^\top \approx \sum_{i=1}^t \beta^{t-i} G_i,$$0	No subspace resampling
MoFaSGD	$$\hat M_t \triangleq U_{t+1}\Sigma_{t+1}V_{t+1}^\top \approx \sum_{i=1}^t \beta^{t-i} G_i,$$1	online, $$\hat M_t \triangleq U_{t+1}\Sigma_{t+1}V_{t+1}^\top \approx \sum_{i=1}^t \beta^{t-i} G_i,$$2

The optimizer-state reduction is substantial because MoFaSGD stores $$\hat M_t \triangleq U_{t+1}\Sigma_{t+1}V_{t+1}^\top \approx \sum_{i=1}^t \beta^{t-i} G_i,$$3, $$\hat M_t \triangleq U_{t+1}\Sigma_{t+1}V_{t+1}^\top \approx \sum_{i=1}^t \beta^{t-i} G_i,$$4, and $$\hat M_t \triangleq U_{t+1}\Sigma_{t+1}V_{t+1}^\top \approx \sum_{i=1}^t \beta^{t-i} G_i,$$5, with no second-moment buffer. For LLaMA-3.1-8B fine-tuning in bf16 with no checkpointing, batch size $$\hat M_t \triangleq U_{t+1}\Sigma_{t+1}V_{t+1}^\top \approx \sum_{i=1}^t \beta^{t-i} G_i,$$6, and gradient accumulation $$\hat M_t \triangleq U_{t+1}\Sigma_{t+1}V_{t+1}^\top \approx \sum_{i=1}^t \beta^{t-i} G_i,$$7, the reported total memory is approximately $$\hat M_t \triangleq U_{t+1}\Sigma_{t+1}V_{t+1}^\top \approx \sum_{i=1}^t \beta^{t-i} G_i,$$8 GB for AdamW, $$\hat M_t \triangleq U_{t+1}\Sigma_{t+1}V_{t+1}^\top \approx \sum_{i=1}^t \beta^{t-i} G_i,$$9 GB for MoFaSGD with rank $U_{t+1}^\top U_{t+1}=I_r$0, $U_{t+1}^\top U_{t+1}=I_r$1 GB for LoRA with rank $U_{t+1}^\top U_{t+1}=I_r$2, $U_{t+1}^\top U_{t+1}=I_r$3 GB for fused GaLore with rank $U_{t+1}^\top U_{t+1}=I_r$4, and $U_{t+1}^\top U_{t+1}=I_r$5 GB for a stateless SWAN-like baseline. The MoFaSGD breakdown is $U_{t+1}^\top U_{t+1}=I_r$6 GB parameters, $U_{t+1}^\top U_{t+1}=I_r$7 GB optimizer, $U_{t+1}^\top U_{t+1}=I_r$8 GB gradients, and $U_{t+1}^\top U_{t+1}=I_r$9 GB activations (Mahdavinia et al., 10 Jul 2025).

Empirical evaluation was reported on three settings. In modded NanoGPT pre-training on a FineWeb subset, MoFaSGD consistently outperformed GaLore across ranks $V_{t+1}^\top V_{t+1}=I_r$0, with smoother convergence and better perplexity, while still trailing full-rank Muon and AdamW. In GLUE fine-tuning with RoBERTa-Base, average accuracy across seven tasks was $V_{t+1}^\top V_{t+1}=I_r$1 for full-rank AdamW, $V_{t+1}^\top V_{t+1}=I_r$2 for GaLore $V_{t+1}^\top V_{t+1}=I_r$3, $V_{t+1}^\top V_{t+1}=I_r$4 for LoRA $V_{t+1}^\top V_{t+1}=I_r$5, and $V_{t+1}^\top V_{t+1}=I_r$6 for MoFaSGD $V_{t+1}^\top V_{t+1}=I_r$7; at rank $V_{t+1}^\top V_{t+1}=I_r$8, the corresponding numbers were $V_{t+1}^\top V_{t+1}=I_r$9, $\Sigma_{t+1}$0, and $\Sigma_{t+1}$1 for GaLore, LoRA, and MoFaSGD. In Tulu-3 instruction tuning with LLaMA-3.1 8B and rank $\Sigma_{t+1}$2, the average score over MMLU, TruthfulQA, BigBenchHard, GSM8K, and HumanEval was $\Sigma_{t+1}$3 for AdamW, $\Sigma_{t+1}$4 for LoRA, $\Sigma_{t+1}$5 for GaLore, and $\Sigma_{t+1}$6 for MoFaSGD. Reported throughput on that setup was approximately $\Sigma_{t+1}$7 tokens/sec for MoFaSGD, $\Sigma_{t+1}$8 for GaLore, and $\Sigma_{t+1}$9 for LoRA. A momentum spectral analysis on Tulu-3 with AdamW further indicated that the top-$r$0 singular values of the first moment contain approximately $r$1 of Frobenius energy, and the top-$r$2 contain approximately $r$3, which empirically supports low-rank first-moment factorization.

5. Theoretical guarantees

The convergence analysis is stated for the non-convex stochastic matrix optimization problem

$r$4

The assumptions are: $r$5 is $r$6-smooth with respect to the nuclear norm,

$r$7

the stochastic gradient oracle is unbiased with

$r$8

and the initialization satisfies $r$9 (Mahdavinia et al., 10 Jul 2025).

Two theoretical results are emphasized. The first is an optimality property of the tangent projection. For projection operators of the form

$M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$00

the Frobenius residual is minimized when $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$01 and $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$02, yielding exactly the tangent projection

$M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$03

or, equivalently,

$M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$04

This establishes that the online projection used in MoFaSGD is optimal within the specified family.

The second is the convergence theorem. Under A1-A2, $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$05, and $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$06, the iterates satisfy

$M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$07

Choosing

$M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$08

yields

$M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$09

which matches the optimal $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$10 rate for non-convex stochastic optimization. The proof decomposes the discrepancy between the true full-rank momentum $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$11 and the low-rank approximation $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$12 into temporal error $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$13 and compression error $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$14, and then bounds the latter recursively through the tangent projection residual.

6. Related interpretations, misconceptions, and limitations

The term “momentum factorization” has a broader surrounding literature than the 2025 optimizer itself. In "Compressing gradients by exploiting temporal correlation in momentum-SGD" (Adikari et al., 2021), the paper does not use MoFaSGD as its method name, but it explicitly describes a hypothetical “Momentum Factorized SGD (MoFaSGD)” that would exploit the low-pass-filtering effect of momentum and represent the update as a predictable, low-frequency component plus a small residual. In that setting, the factorization is temporal and communication-oriented: predictive coding factorizes the momentum trajectory into a history-based predictor and a compressed residual. A separate communication-centered perspective appears in "Communication-Efficient Distributed Blockwise Momentum SGD with Error-Feedback" (Zheng et al., 2019), where momentum-based updates are represented blockwise as sign plus scale with error-feedback, yielding nearly $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$15 communication reduction. "Accelerating Single-Pass SGD for Generalized Linear Prediction" (Chen et al., 2 Mar 2026) does not define MoFaSGD, but it presents dual-momentum operator factorization and a decomposition of excess risk into optimization, statistical, and model-misspecification terms as relevant design principles for momentum-factorized methods. "Effects of momentum scaling for SGD" (Pasechnyuk et al., 2022) studies momentum-updated preconditioners and emphasizes the role of the momentum coefficient $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$16, which is directly relevant because practical MoFaSGD runs use momentum decay values around $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$17–$M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$18 even though the current proof assumes $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$19.

Several misconceptions are addressed directly by the source material. MoFaSGD is not a parameter-efficient fine-tuning method of the LoRA type; it makes no architectural change and leaves the base model fully trainable. It is not identical to GaLore-style gradient projection, because its low-rank object is the first moment rather than the current gradient, and its subspace evolves online rather than through periodic offline SVD resampling. It is also not simply low-rank momentum descent, because the update uses the spectrally normalized factor $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$20, not the low-rank momentum $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$21.

The limitations are equally explicit. Performance remains below full-rank AdamW and Muon on the largest instruction-tuning benchmarks. Very low ranks can under-approximate the momentum, slowing convergence or degrading final quality. Although the per-step cost is modest for small $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$22, the $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$23 term grows with rank. The theory uses nuclear-norm smoothness and $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$24, which do not match the practical regime exactly. Current implementations apply MoFaSGD primarily to $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$25D linear layers of transformer blocks, while embeddings and $M_t = \beta M_{t-1} + (1-\beta)G_t,\qquad M_0=0,$26D parameters are optimized with AdamW in bf16. This suggests that MoFaSGD is best understood as a low-rank, memory-efficient variant of spectral momentum normalization with strong empirical trade-offs under tight GPU-memory budgets, rather than as a universal replacement for full-rank optimizers.