Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation

Paper: 2602.24283 Authors: Zhengbo Wang, Jian Liang, Ran He, Zilei Wang, Tieniu Tan Categories: cs.LG, cs.AI, cs.CL

The Gap

Training billion-parameter language models requires optimizers like Adam or Muon that maintain momentum states—exponential moving averages of gradients and their squares. The problem? These momentum matrices are the same size as the model parameters themselves. For a 1B parameter model, Adam stores three copies: the parameters, first-order momentum, and second-order momentum. That’s 12GB just for optimizer state if you’re using float32. This memory wall hits hardest during pre-training, where you can’t afford to waste a single byte of GPU memory. Prior work like GaLore and ReLoRA attempted low-rank training, but they focused on parameter updates rather than optimizer states, leaving the momentum bottleneck untouched.

The Increment

Before: Momentum states consume as much memory as the model itself, forcing practitioners to choose between optimizer quality and scale. After: LoRA-Pre compresses momentum into low-rank subspaces, cutting memory by 8× while matching or beating full-rank performance.

The core insight is elegant: exponential moving averages (EMAs) used in momentum are mathematically equivalent to training a linear regressor via online gradient descent. Think of momentum like a streaming data warehouse. Traditional optimizers (the warehouse) store every raw transaction (gradient) in full detail. LoRA-Pre instead maintains a compressed index—a low-rank projection that captures the essential patterns. The “query engine” (parameter updates) reconstructs what it needs on-the-fly from this compact representation. The projection matrices ($A$ and $B$) are the index schema, the low-rank dimension $r$ is the compression ratio, and the online gradient flow is the incremental update protocol that keeps the index fresh without rebuilding from scratch.

Concretely, instead of storing momentum $m_t \in \mathbb{R}^{d \times d}$, LoRA-Pre decomposes it as $m_t \approx A_t B_t^T$ where $A_t \in \mathbb{R}^{d \times r}$ and $B_t \in \mathbb{R}^{d \times r}$ with $r \ll d$. The authors prove this decomposition emerges naturally when you view EMA as solving $\min_{m} \sum_{i=1}^t \beta^{t-i} \|m - g_i\|^2$—a weighted least squares problem that admits a low-rank solution via online gradient descent on $A$ and $B$.

Key Concepts

Exponential Moving Average as Regression: When Adam computes $m_t = \beta m_{t-1} + (1-\beta) g_t$, it’s not just smoothing—it’s solving an optimization problem. Imagine you’re trying to predict tomorrow’s weather by averaging past observations, but recent days matter more. That’s EMA: a weighted average where weights decay exponentially ($\beta^{t-i}$). The authors show this is equivalent to minimizing $\sum_{i=1}^t \beta^{t-i} \|m - g_i\|^2$—finding the single vector $m$ that best approximates all past gradients $g_i$, weighted by recency. This reframing is crucial because least squares problems have well-studied low-rank structure. If your gradients lie near a low-dimensional subspace (which they often do in neural networks), you don’t need the full $d$-dimensional $m$—a rank-$r$ approximation suffices.

Low-Rank Factorization via Online Gradient Flow: Instead of directly storing $m_t$, LoRA-Pre maintains $A_t, B_t$ such that $m_t \approx A_t B_t^T$. But how do you update $A$ and $B$ as new gradients arrive? The trick is to apply gradient descent on the factorized objective: $\min_{A,B} \|AB^T - g_t\|^2 + \text{regularization from past}$. This yields update rules: $A_{t+1} = A_t - \eta abla_A \mathcal{L}, \quad B_{t+1} = B_t - \eta \nabla_B \mathcal{L}$ The beauty is that these updates are local—you never materialize the full $d \times d$ matrix. It’s like updating a compressed JPEG by tweaking the frequency coefficients, not the raw pixels. The online gradient flow ensures $A_t B_t^T$ tracks the true momentum $m_t$ without ever computing $m_t$ explicitly.

Rank Efficiency: The paper shows you can use $r = d/8$ (one-eighth the parameter dimension) and still match full-rank performance. Why? Neural network gradients exhibit low effective rank due to correlated parameter updates—layers learn similar features, creating redundancy. LoRA-Pre exploits this by projecting momentum onto the dominant $r$ directions. Think of it like principal component analysis (PCA) on the gradient history: most variance concentrates in a few eigenvectors. By keeping only those top-$r$ components, you discard noise while retaining signal. The 8× compression isn’t arbitrary—it’s where the eigenvalue spectrum drops off for typical LLM training runs.

Expert Assessment

Problem significance: This is a real pain point. Memory is the primary constraint in scaling LLM training, and optimizer states are low-hanging fruit—they’re necessary but not fundamental to the model. The affected community is large: anyone pre-training or fine-tuning models beyond a few billion parameters. However, the problem is more acute for pre-training than fine-tuning, where methods like LoRA already reduce parameter counts.

Method maturity: The theoretical foundation (EMA as regression) is solid, but the practical implementation raises questions. The authors don’t discuss how to choose rank $r$ adaptively—they test fixed ratios. What happens if the gradient subspace dimensionality changes during training (e.g., early vs. late stages)? The paper also lacks wall-clock time comparisons. Low-rank updates might save memory but add computational overhead from matrix multiplications. Is the memory-compute tradeoff favorable on real hardware? The experiments stop at 1B parameters; scaling to 70B+ models (where memory matters most) is unvalidated.

Experimental rigor: The baselines are fair—GaLore, ReLoRA, and standard LoRA are appropriate comparisons. The Llama architecture is a good testbed. However, the pre-training experiments use relatively small datasets (the paper doesn’t specify token counts clearly). Pre-training convergence is sensitive to hyperparameters; did they tune learning rates separately for each method, or use the same schedule? The fine-tuning results (3.14 and 6.17 point improvements) are impressive but lack error bars or multiple runs. One red flag: the paper claims “highest performance across all model sizes” but doesn’t show loss curves—only final perplexity. Did LoRA-Pre converge faster, or just to a better final point?

Verdict: Weak accept — The core idea is sound and addresses a real bottleneck, but the experimental validation is incomplete for production use at scale, and key practical details (adaptive rank selection, wall-clock efficiency) are underexplored.

Takeaways

Reframe accumulators as learnable systems: Momentum isn’t just a running average—it’s a model being trained. This perspective unlocks compression techniques from representation learning. Anywhere you maintain exponential moving averages (batch norm statistics, policy gradients in RL, Kalman filters), ask: can I factorize this accumulator?

Online gradient flow for streaming compression: The update rules for $A_t, B_t$ are a template for maintaining low-rank approximations of streaming data. If you’re building systems that aggregate high-dimensional signals over time (sensor networks, recommendation systems), this pattern—update factors incrementally rather than recompute the full matrix—is directly applicable.

Memory-compute tradeoffs are hardware-dependent: Low-rank methods save memory but add FLOPs. On memory-bound workloads (large batch training), this tradeoff favors compression. On compute-bound workloads (small batch inference), it might not. Always profile before assuming “fewer parameters = faster.”

Rank as a hyperparameter deserves more attention: The paper shows $r = d/8$ works, but doesn’t explore dynamic rank adaptation. In practice, you could monitor the reconstruction error $\|AB^T - m_{\text{full}}\|$ (computed occasionally on a subset) and adjust $r$ mid-training. This adaptive approach could generalize beyond optimizers to any low-rank approximation problem.