Low-rank & Tensorized Optimizer State

• The paper presents low-rank and tensorized optimizer state techniques that reduce memory from O(d²) to O(dr) while preserving convergence in large-scale models.
• It leverages factorization and projection methods, as seen in LoRA-Pre, GaLore, and PLUMAGE, to efficiently update first- and second-moment statistics.
• Empirical evaluations show significant memory savings and enhanced stability, enabling effective training and fine-tuning of billion-parameter models on affordable hardware.

Low-rank and tensorized optimizer state approaches exploit the empirical observation that the optimizer statistics (first and second moments, or higher-order tensors of gradients) encountered in deep learning admit strong low-rank or structured approximations. These techniques substantially reduce the memory and computational barriers associated with storing and manipulating optimizer state in large-scale models, such as LLMs with billions of parameters, without degrading convergence or generalization. Key modern instances include LoRA-Pre, GaLore, and PLUMAGE, which leverage linear algebraic insights and stochastic estimators to compress and efficiently update the state tracked by optimizers such as Adam or Muon, enabling high-fidelity training and fine-tuning on affordable hardware (Wang et al., 27 Feb 2026, Zhao et al., 2024, Haroush et al., 23 May 2025).

1. Reinterpretation of Optimizer State: From EMA to Online Linear Regression

The foundation for low-rank optimizer states is the reinterpretation of the exponential moving averages (EMA) in adaptive optimizers as the online training of a linear regressor via gradient flow. For example, the first-moment vector update in Adam,

$m_t \leftarrow \beta_1 m_{t-1} + (1-\beta_1) g_t ,$

can be rewritten as a gradient-descent step on the squared loss

$L(m; g_t) = \frac{1}{2} \| m - g_t \|^2 .$

Similarly, the second-moment matrix EMA,

$$M_t \leftarrow \beta_2 M_{t-1} + (1-\beta_2) (g_t g_t^\top) ,$$

corresponds to a least-squares update of a matrix-valued regressor. This equivalence enables viewing momentum accumulation as a problem of online linear regression on the instantaneous gradients or their various algebraic forms (Wang et al., 27 Feb 2026).

2. Low-rank and Tensorized Approximations of Momentum States

Storing the full first- or second-moment matrices is rapidly prohibitive as parameter dimensionality grows (e.g., $d \sim 4$K yields millions of elements per matrix). Approaches such as LoRA-Pre and GaLore thus introduce low-rank factorizations or tensor projections to the optimizer state.

• LoRA-Pre parameterizes the momentum matrix as $M_t \approx U_t V_t^\top$ with $U_t, V_t \in \mathbb{R}^{d \times r}$, $r \ll d$. Closed-form “EMA-style” updates for $U_t$ and $V_t$ yield an implicit online low-rank regressor, leveraging quasi-Newton updates to maintain alignment with the full-rank EMA dynamics (Wang et al., 27 Feb 2026).
• GaLore projects gradients onto a dynamically adapted low-rank subspace. At each SVD update interval, projectors $P_t \in \mathbb{R}^{m \times r}$ and $L(m; g_t) = \frac{1}{2} \| m - g_t \|^2 .$0 (for weights $L(m; g_t) = \frac{1}{2} \| m - g_t \|^2 .$1) capture the top singular vectors of the gradient. All optimizer updates (Adam, etc.) are performed on projected “core” statistics, and reconstructed to the ambient space, dramatically compressing the state (Zhao et al., 2024).
• PLUMAGE generalizes GaLore by introducing a probabilistic stochastic low-rank estimator. It samples singular modes according to a “water-filling” solution that achieves minimal estimator variance under a rank constraint, yielding an unbiased gradient estimator and compatible state update. PLUMAGE also includes a principled realignment mechanism for low-rank state when projectors are refreshed, addressing misalignment instability seen in previous methods (Haroush et al., 23 May 2025).

3. Algorithmic Frameworks and Upgrade Path for Adam-like Optimizers

Three contemporary methodologies crystalize operational differences and algorithmic flows for low-rank/tensorized optimizer state management.

Method	Projection	Momentum State	Subspace Refresh
LoRA-Pre	None (per-weight matrix, factored U,V)	Low-rank EMA factors $L(m; g_t) = \frac{1}{2} \| m - g_t \|^2 .$2	Each step (no SVD)
GaLore	$L(m; g_t) = \frac{1}{2} \| m - g_t \|^2 .$3 via SVD	Projected statistics $L(m; g_t) = \frac{1}{2} \| m - g_t \|^2 .$4	Every $L(m; g_t) = \frac{1}{2} \| m - g_t \|^2 .$5 steps (SVD)
PLUMAGE	$L(m; g_t) = \frac{1}{2} \| m - g_t \|^2 .$6 via stochastic SVD	Projected, realigned $L(m; g_t) = \frac{1}{2} \| m - g_t \|^2 .$7	Every $L(m; g_t) = \frac{1}{2} \| m - g_t \|^2 .$8 steps (SVD)

All techniques yield $L(m; g_t) = \frac{1}{2} \| m - g_t \|^2 .$9 optimizer state memory (with $$M_t \leftarrow \beta_2 M_{t-1} + (1-\beta_2) (g_t g_t^\top) ,$$0 the parameter dimension and $$M_t \leftarrow \beta_2 M_{t-1} + (1-\beta_2) (g_t g_t^\top) ,$$1 the rank or projection width), versus $$M_t \leftarrow \beta_2 M_{t-1} + (1-\beta_2) (g_t g_t^\top) ,$$2 for the exact, unfactored matrix statistics.

• In LoRA-Pre, first- and second-moment state is stored as compact $$M_t \leftarrow \beta_2 M_{t-1} + (1-\beta_2) (g_t g_t^\top) ,$$3 factors (one per moment), updated with Newton or quasi-Newton steps after observing each gradient.
• GaLore and PLUMAGE maintain projectors and low-rank optimizer moments, with periodic SVD-based update of the subspace to adapt to the changing gradient structure; PLUMAGE’s estimator is probabilistically unbiased, and uses moment realignment when subspaces shift (Wang et al., 27 Feb 2026, Zhao et al., 2024, Haroush et al., 23 May 2025).

4. Memory, Computational Trade-offs, and Implementation

Major efficiency gains derive from dramatically lowering optimizer state size—from $$M_t \leftarrow \beta_2 M_{t-1} + (1-\beta_2) (g_t g_t^\top) ,$$4 to $$M_t \leftarrow \beta_2 M_{t-1} + (1-\beta_2) (g_t g_t^\top) ,$$5 per parameter matrix.

• Memory Cost: For LoRA-Pre, only $$M_t \leftarrow \beta_2 M_{t-1} + (1-\beta_2) (g_t g_t^\top) ,$$6 floats are needed per matrix (two factors). GaLore and PLUMAGE require the storage of projectors ($$M_t \leftarrow \beta_2 M_{t-1} + (1-\beta_2) (g_t g_t^\top) ,$$7) and low-rank moment statistics (e.g., $$M_t \leftarrow \beta_2 M_{t-1} + (1-\beta_2) (g_t g_t^\top) ,$$8 for $$M_t \leftarrow \beta_2 M_{t-1} + (1-\beta_2) (g_t g_t^\top) ,$$9 for GaLore).
• Computational Cost: For LoRA-Pre, per-step extra work involves $d \sim 4$0 for Gram matrix computations and $d \sim 4$1 for $d \sim 4$2 matrix inversions, negligible for $d \sim 4$3. For GaLore and PLUMAGE, per-step cost is $d \sim 4$4 for projection and reconstruction, with SVD refresh amortized over $d \sim 4$5 or $d \sim 4$6 steps and contributing less than $d \sim 4$7 overhead versus standard backpropagation in practice.
• Tensorization: Each layer’s weight matrix is handled independently, with its own low-rank state/projection; further tiling or structured sparsity is conceptually compatible.
• Subspace Refresh: For GaLore and PLUMAGE, SVD or stochastic SVD is performed at fixed intervals ($d \sim 4$8), trading between subspace freshness and computational load; LoRA-Pre does not require such periodic SVDs as its factor updates are closed-form at each step (Wang et al., 27 Feb 2026, Zhao et al., 2024, Haroush et al., 23 May 2025).

5. Empirical Performance and Model-Scale Evaluation

Empirically, low-rank and tensorized state methods achieve substantial memory reductions and, in many cases, improve accuracy or stability compared to traditional low-rank adaptation schemes.

• LoRA-Pre matches or exceeds baseline validation perplexity across Llama-scale model pre-training (60M to 1B parameters), surpassing GaLore and other low-rank baselines even at $d \sim 4$9–$M_t \approx U_t V_t^\top$0 the rank budget. It also yields 2–6 point gains in fine-tuning benchmarks over standard LoRA and variants. Notably, continual subspace updates via low-rank Newton steps enable rapid adaptation and avoid the drift characteristic of periodic-projection techniques (Wang et al., 27 Feb 2026).
• GaLore achieves up to 65.5% reduction in optimizer state memory (over BF16 Adam) and, via the 8-bit extension, up to 82.5%; total training memory is reduced by 63.3%. GaLore delivers full-parameter learning fidelity with projected Adam/Muon while making multi-billion parameter training feasible on 24GB consumer GPUs. Performance matches or slightly exceeds full-rank Adam on pre-training and fine-tuning; e.g., 7B LLaMA can be trained on a single consumer GPU (Zhao et al., 2024).
• PLUMAGE closes the evaluation loss gap to full-rank optimization by 33% on pre-training tasks and improves average training loss on GLUE benchmarks by 28% over prior low-rank approaches. Stability gains stem from unbiased estimation, minimum-variance projections, and rigorous optimizer state realignment during projector refresh (Haroush et al., 23 May 2025).

6. State Realignment, Estimator Bias, and Stability Considerations

A central practical challenge is the “state misalignment” that arises when the low-rank projector is updated: the stored optimizer moments may not correspond to the new subspace. PLUMAGE addresses this via re-projection formulas for both first and second moments:

• First-moment coefficients are mapped by $M_t \approx U_t V_t^\top$1, i.e., $M_t \approx U_t V_t^\top$2;
• Second-moment coefficients are elementwise updated via $M_t \approx U_t V_t^\top$3;

A standard diagonal-Fisher approximation is employed for efficiency and stability in the latter. This principled realignment prevents spurious weight updates and instabilities, which can otherwise arise in both deterministic (GaLore) and random (FLORA) projection schedules (Haroush et al., 23 May 2025).

Estimation bias and variance are also controlled: GaLore is deterministic and thus biased (with top-$M_t \approx U_t V_t^\top$4 singular vectors retained); PLUMAGE achieves unbiasedness and minimal variance subject to the rank constraint, with variance-optimal sampling probabilities for singular modes under a given memory budget.

7. Comparative Analysis and Applicability

The following table summarizes salient features:

Approach	Optimizer State Reduction	Estimator Type	Realignment
LoRA-Pre	Factorized low-rank	Deterministic	N/A
GaLore	Projected low-rank	Deterministic	None
PLUMAGE	Probabilistic projected	Unbiased/min-variance	Yes

LoRA-Pre is notable for achieving state efficiency without SVD refresh or requiring projectors, relying solely on online updates to low-rank factors. GaLore and PLUMAGE generalize to matrix- and tensor-shaped parameters, making them particularly suitable for layers with high intrinsic latent rank. PLUMAGE’s unbiased estimation and stability advantages make it advantageous for domains where convergence to the full-rank solution is critical (Wang et al., 27 Feb 2026, Zhao et al., 2024, Haroush et al., 23 May 2025).

• Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation (Wang et al., 27 Feb 2026)
• GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection (Zhao et al., 2024)
• PLUMAGE: Probabilistic Low rank Unbiased Min Variance Gradient Estimator for Efficient Large Model Training (Haroush et al., 23 May 2025)