---
title: Low-rank & Tensorized Optimizer State
url: https://www.emergentmind.com/topics/low-rank-and-tensorized-optimizer-state
type: topic
---

# Low-rank & Tensorized Optimizer State

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 [2602.24283][2403.03507][2505.18313].

## 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 [2602.24283].

## 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 [2602.24283].

- **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 $Q_t \in \mathbb{R}^{n \times r}$ (for weights $W \in \mathbb{R}^{m \times n}$) 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 [2403.03507].

- **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 [2505.18313].

## 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 $(U,V)$ | Each step (no SVD) |
| GaLore      | $P,Q$ via SVD    | Projected statistics $(M^r, V^r)$ | Every $T$ steps (SVD) |
| PLUMAGE     | $P$ via stochastic SVD  | Projected, realigned $(M^\parallel, V^\parallel)$ | Every $\tau$ steps (SVD) |

All techniques yield $O(dr)$ optimizer state memory (with $d$ the parameter dimension and $r$ the rank or projection width), versus $O(d^2)$ for the exact, unfactored matrix statistics.

- In LoRA-Pre, first- and second-moment state is stored as compact $U,V$ 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 [2602.24283][2403.03507][2505.18313].

## 4. Memory, Computational Trade-offs, and Implementation

Major efficiency gains derive from dramatically lowering optimizer state size—from $O(d^2)$ to $O(dr)$ per parameter matrix.

- **Memory Cost:** For LoRA-Pre, only $2rd$ floats are needed per matrix (two factors). GaLore and PLUMAGE require the storage of projectors ($mr+nr$) and low-rank moment statistics (e.g., $2nr$ for $M^r, V^r$ for GaLore).
- **Computational Cost:** For LoRA-Pre, per-step extra work involves $O(dr^2)$ for Gram matrix computations and $O(r^3)$ for $r \times r$ matrix inversions, negligible for $r \ll d$. For GaLore and PLUMAGE, per-step cost is $O(mnr)$ for projection and reconstruction, with SVD refresh amortized over $T$ or $\tau$ steps and contributing less than $10\%$ 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 ($T,\tau$), 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 [2602.24283][2403.03507][2505.18313].

## 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 $1/8$–$1/16$ 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 [2602.24283].
- **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 [2403.03507].
- **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 [2505.18313].

## 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 $B = P_t^\top P_{t-1}$, i.e., $M^\parallel \leftarrow B M^\parallel$;
- Second-moment coefficients are elementwise updated via $(B^{\circ2}) V^\parallel$;
  
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 [2505.18313].

Estimation bias and variance are also controlled: GaLore is deterministic and thus biased (with top-$r$ 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 [2602.24283][2403.03507][2505.18313].

## References

- Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation [2602.24283]
- GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection [2403.03507]
- PLUMAGE: Probabilistic Low rank Unbiased Min Variance Gradient Estimator for Efficient Large Model Training [2505.18313]

Source: https://www.emergentmind.com/topics/low-rank-and-tensorized-optimizer-state