---
title: Momentum Factorized SGD (MoFaSGD)
url: https://www.emergentmind.com/topics/momentum-factorized-sgd-mofasgd
type: topic
---

# Momentum Factorized SGD (MoFaSGD)

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" [2507.08091] 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 [2507.08091].

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,V_t \in \mathbb{R}^{m\times n}\), so optimizer state is \(\Theta(mn)\) 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 \((U_t,\Sigma_t,V_t)\) without forming a full SVD of \(G_t\). MoFaSGD addresses this by projecting the new gradient onto the tangent space of the rank-\(r\) manifold at the previous low-rank momentum approximation. With current factors \((U_t,\Sigma_t,V_t)\), the tangent-space projection is
\[
\hat G_t := \operatorname{Proj}_{T_t}(G_t)
= U_tU_t^\top G_t + G_tV_tV_t^\top - U_tU_t^\top G_tV_tV_t^\top.
\]

Conceptually, the next low-rank momentum is obtained from
\[
\hat M_t \approx \operatorname{SVD}_r\bigl(\hat G_t + \beta \hat M_{t-1}\bigr)
= \operatorname{SVD}_r\bigl(\hat G_t + \beta U_t\Sigma_tV_t^\top\bigr).
\]
Because both \(\hat G_t\) and \(\hat M_{t-1}\) have rank at most \(r\), their sum has rank at most \(2r\). The algorithm therefore avoids a full decomposition by computing \(G_tV_t\), \(U_t^\top G_t\), and \(U_t^\top G_tV_t\), then forming QR decompositions of the augmented subspaces
\[
(U_t',R_{U_t})=\mathrm{QR}\big([U_t \;\; G_tV_t]\big),\qquad
(V_t',R_{V_t})=\mathrm{QR}\big([V_t \;\; G_t^\top U_t]\big),
\]
and a small core matrix
\[
S_t=
R_{U_t}
\begin{bmatrix}
\beta \Sigma_t - U_t^\top G_tV_t & I_r\\
I_r & 0_r
\end{bmatrix}
R_{V_t}^\top.
\]
A rank-\(r\) SVD of this \(2r\times 2r\) matrix,
\[
U_t''\Sigma_t''V_t''^\top=\mathrm{SVD}_r(S_t),
\]
yields the updated factors
\[
U_{t+1}=U_t'U_t'',\qquad
V_{t+1}=V_t'V_t'',\qquad
\Sigma_{t+1}=\Sigma_t''.
\]

The resulting per-layer complexity is
\[
\mathcal{O}((m+n)r^2+r^3),
\]
coming from two QR decompositions of size \((m\times 2r)\) and \((n\times 2r)\) and an SVD of size \(2r\times 2r\). For \(r\ll \min(m,n)\), this replaces the expensive full-matrix SVDs required by methods that periodically resample a gradient-defined subspace [2507.08091].

## 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 \(\hat M_t = U_{t+1}\Sigma_{t+1}V_{t+1}^\top\) has been computed, the update is
\[
\boxed{W_{t+1} = W_t - \eta\, U_{t+1}V_{t+1}^\top.}
\]
Only the singular directions are used; the singular values \(\Sigma_{t+1}\) 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 [2507.08091].

This design gives MoFaSGD two coupled interpretations. First, the factors \(U_t\) and \(V_t\) define an implicitly evolving optimization subspace. Unlike GaLore, which defines a projection matrix from the current gradient and resamples it every \(\tau\) steps, MoFaSGD updates its subspace every iteration through the momentum SVD itself. The subspace is therefore momentum-defined, online, and smooth. Second, the update \(U_{t+1}V_{t+1}^\top\) is the polar or sign factor of the low-rank momentum restricted to rank \(r\), 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 \(W_{t+1}=W_t-\eta \hat M_t\). Its state is a low-rank factorization of the first moment, while its step is a spectrally normalized rank-\(r\) direction extracted from that factorization.

## 4. Memory, computational profile, and empirical performance

For a single matrix \(W\in\mathbb{R}^{m\times n}\) with \(m\le n\), the source material gives the following memory and subspace-update costs.

| Method | Memory | Subspace update |
|---|---|---|
| GaLore | \(mn + mr + 2nr\) | full SVD of gradient, \(\mathcal{O}(m^2n)\) |
| LoRA | \(mn + 3mr + 3nr\) | No subspace resampling |
| MoFaSGD | \(mn + mr + nr + r\) | online, \(\mathcal{O}((m+n)r^2 + r^3)\) |

The optimizer-state reduction is substantial because MoFaSGD stores \(U\in\mathbb{R}^{m\times r}\), \(V\in\mathbb{R}^{n\times r}\), and \(\Sigma\in\mathbb{R}^{r\times r}\), with no second-moment buffer. For LLaMA-3.1-8B fine-tuning in bf16 with no checkpointing, batch size \(1\), and gradient accumulation \(8\), the reported total memory is approximately \(70.8\) GB for AdamW, \(29.4\) GB for MoFaSGD with rank \(8\), \(33.6\) GB for LoRA with rank \(8\), \(30.0\) GB for fused GaLore with rank \(8\), and \(43.9\) GB for a stateless SWAN-like baseline. The MoFaSGD breakdown is \(15.5\) GB parameters, \(4.2\) GB optimizer, \(2.1\) GB gradients, and \(7.6\) GB activations [2507.08091].

Empirical evaluation was reported on three settings. In modded NanoGPT pre-training on a FineWeb subset, MoFaSGD consistently outperformed GaLore across ranks \(r\in\{16,32,128\}\), 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 \(85.57\%\) for full-rank AdamW, \(84.36\%\) for GaLore \(r=4\), \(84.29\%\) for LoRA \(r=4\), and \(84.86\%\) for MoFaSGD \(r=4\); at rank \(8\), the corresponding numbers were \(84.52\%\), \(84.71\%\), and \(85.18\%\) for GaLore, LoRA, and MoFaSGD. In Tulu-3 instruction tuning with LLaMA-3.1 8B and rank \(8\), the average score over MMLU, TruthfulQA, BigBenchHard, GSM8K, and HumanEval was \(65.9\) for AdamW, \(59.4\) for LoRA, \(60.9\) for GaLore, and \(61.7\) for MoFaSGD. Reported throughput on that setup was approximately \(4206\) tokens/sec for MoFaSGD, \(3214\) for GaLore, and \(4536\) for LoRA. A momentum spectral analysis on Tulu-3 with AdamW further indicated that the top-\(32\) singular values of the first moment contain approximately \(80\%\) of Frobenius energy, and the top-\(16\) contain approximately \(75\%\), which empirically supports low-rank first-moment factorization.

## 5. Theoretical guarantees

The convergence analysis is stated for the non-convex stochastic matrix optimization problem
\[
\min_{W\in\mathbb{R}^{m\times n}} \mathcal{L}(W)=\mathbb{E}_\xi[\mathcal{L}(W,\xi)].
\]
The assumptions are: \( \mathcal{L}\) is \(L\)-smooth with respect to the nuclear norm,
\[
\|\nabla \mathcal{L}(W_1)-\nabla \mathcal{L}(W_2)\|_*
\le L\|W_1-W_2\|_2,
\]
the stochastic gradient oracle is unbiased with
\[
\mathbb{E}_\xi[\nabla \mathcal{L}(W,\xi)] = \nabla \mathcal{L}(W),\qquad
\mathbb{E}_\xi[\|\nabla \mathcal{L}(W,\xi)-\nabla \mathcal{L}(W)\|_*]\le \sigma,
\]
and the initialization satisfies \(\operatorname{rank}(G_0)\le r\) [2507.08091].

Two theoretical results are emphasized. The first is an optimality property of the tangent projection. For projection operators of the form
\[
\mathrm{Proj}_{(L,R)}(G)=\alpha_1LL^\top G+\alpha_2GRR^\top+\alpha_3LL^\top GRR^\top,
\]
the Frobenius residual is minimized when \(L^\top L=R^\top R=I_r\) and \((\alpha_1,\alpha_2,\alpha_3)=(1,1,-1)\), yielding exactly the tangent projection
\[
\hat G=(I-(I-LL^\top))\,G\,(I-(I-RR^\top))
\]
or, equivalently,
\[
\|G-\hat G\|_F=\|(I-LL^\top)G(I-RR^\top)\|_F.
\]
This establishes that the online projection used in MoFaSGD is optimal within the specified family.

The second is the convergence theorem. Under A1-A2, \(\beta\le 1/3\), and \(\eta\le 1\), the iterates satisfy
\[
\frac{1}{T}\sum_{t=0}^T \mathbb{E}\big[\|\nabla \mathcal{L}(W_t)\|_*\big]
\le
\mathcal{O}\!\left(
\frac{\mathcal{L}(W_0)-\mathbb{E}[\mathcal{L}(W_{T+1})]}{\eta T}
+\eta L
+\frac{\sigma}{\sqrt{T}}
\right).
\]
Choosing
\[
\eta=\Theta\!\left(\sqrt{\frac{\mathcal{L}(W_0)-\mathbb{E}[\mathcal{L}(W_{T+1})]}{TL}}\right)
\]
yields
\[
\frac{1}{T}\sum_{t=0}^T \mathbb{E}\big[\|\nabla \mathcal{L}(W_t)\|_*\big]
\le
\mathcal{O}\!\left(
\frac{(\mathcal{L}(W_0)-\mathbb{E}[\mathcal{L}(W_{T+1})])^{1/2}\sqrt{L}+\sigma}{\sqrt{T}}
\right),
\]
which matches the optimal \(O(1/\sqrt{T})\) rate for non-convex stochastic optimization. The proof decomposes the discrepancy between the true full-rank momentum \(M_t\) and the low-rank approximation \(\hat M_t\) into temporal error \(\|M_t-\bar G_t\|_*\) and compression error \(\|\hat M_t-M_t\|_*\), 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" [2108.07827], 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" [1905.10936], where momentum-based updates are represented blockwise as sign plus scale with error-feedback, yielding nearly \(32\times\) communication reduction. "Accelerating Single-Pass SGD for Generalized Linear Prediction" [2603.01951] 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" [2210.11869] studies momentum-updated preconditioners and emphasizes the role of the momentum coefficient \(\beta\), which is directly relevant because practical MoFaSGD runs use momentum decay values around \(0.85\)–\(0.95\) even though the current proof assumes \(\beta\le 1/3\).

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 \(U_{t+1}V_{t+1}^\top\), not the low-rank momentum \(U_{t+1}\Sigma_{t+1}V_{t+1}^\top\).

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 \(r\), the \(\mathcal{O}((m+n)r^2)\) term grows with rank. The theory uses nuclear-norm smoothness and \(\beta\le 1/3\), which do not match the practical regime exactly. Current implementations apply MoFaSGD primarily to \(2\)D linear layers of transformer blocks, while embeddings and \(1\)D 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.

Source: https://www.emergentmind.com/topics/momentum-factorized-sgd-mofasgd