Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation

AI summary

Overview

Research area: Memory-efficient optimization for large language model training, specifically compressing optimizer states (momentum buffers) using low-rank factorization.

Technical level: Intermediate. The motivating idea is intuitive, but the closed-form update derivation relies on Newton's method, Kronecker-product Hessians, and matrix calculus.

Scope: The paper proposes LoRA-Pre, a low-rank optimizer that recasts momentum updates as online linear regression, compresses momentum into two small matrices, and validates the approach on Llama pre-training (60M–1B parameters) and fine-tuning (7B/8B parameters).

What This Paper Is About

Training large language models with Adam-style optimizers requires storing not just model weights but also first- and second-order momentum estimates, roughly tripling memory use and limiting how large a model can be trained on fixed hardware. Existing fixes project gradients into a low-rank subspace, but they update that subspace only periodically, which introduces lag and accumulated error. This paper's goal is to compress optimizer state without that lag, by showing that momentum bookkeeping is secretly a small regression problem that can be factored and solved continuously at every step.

Key Contributions

1. A theoretical equivalence between momentum and online linear regression. The authors prove that the standard exponential moving average update m ← β·m + (1−β)·g is mathematically identical to one gradient descent step on the objective L(m; g) = ½‖m − g‖²_F, where the momentum m plays the role of the model weight, (1−β) is the learning rate, and (m − g) is the gradient.

2. LoRA-Pre, a low-rank optimizer built on that equivalence. Instead of storing the full momentum matrix m ∈ ℝ^{p×q}, the method stores m = m_B · m_A with m_B ∈ ℝ^{p×r} and m_A ∈ ℝ^{r×q}, cutting memory from p×q to (p+q)×r.

3. Closed-form, backpropagation-free update rules. Theorem 3.1 derives explicit update equations for the two factors using Newton's method, yielding updates that retain the EMA form while operating in the compressed space. A separate reparameterization v = (v_B v_A)^{∘2} keeps the second-order momentum element-wise positive, which Adam requires for its square root.

4. Optimizer-agnostic variants and broad empirical validation. The framework is instantiated for both AdamW and Muon, and tested on pre-training from scratch (Llama 60M, 130M, 350M, 1B on C4) and fine-tuning (Llama-3.1-8B, Llama-2-7B on MetaMathQA, evaluated on GSM8K and MATH-500).

Main Findings

• Best or near-best pre-training perplexity at every scale. LoRA-Pre Adam and LoRA-Pre Muon place first or second across all four model sizes. LoRA-Pre Adam improves over the previous best efficient baseline by 0.81, 2.45, and 1.6 perplexity points at 130M, 350M, and 1B parameters respectively.

• Strong rank efficiency. On the 60M model, LoRA-Pre Adam at rank 16 matches GaLore at rank 128 — an 8× reduction. On the 130M model, rank 16 matches GaLore at rank 256, a 16× reduction. The authors attribute this to continuous subspace adaptation versus GaLore's periodic updates.

• Small ranks converge to competitive quality. Training curves for 130M models show that low-rank runs start with worse perplexity but close the gap rapidly, indicating the dynamic update captures the evolving momentum structure even under tight rank budgets.

• Fine-tuning gains over LoRA and other PEFT baselines. On Llama-3.1-8B, LoRA-Pre Adam reaches 47.05 average (GSM8K + MATH-500), 2.59 points above the second-best method and 3.14 points above standard LoRA. On Llama-2-7B, it reaches 32.15 average, 4.88 points above second-best and 6.17 above standard LoRA.

• Projection-based methods degrade on Muon. When GaLore and Fira are adapted to the Muon optimizer, they perform worse than Muon with no momentum at all (34.39 and 34.45 vs. 32.15 perplexity at 60M), while LoRA-Pre Muon reaches 30.76. The authors attribute this to periodic subspace updates corrupting the momentum used in Muon's orthogonalization step.

Methodology in Plain English

The starting observation is a rewrite of the momentum update. Adam's momentum step is m ← β·m + (1−β)·g, which can be rearranged into m ← m − (1−β)·(m − g). That has the exact shape of a gradient descent step: the "parameter" is the momentum, the "learning rate" is 1−β, and the "gradient" is the difference between the current momentum and the incoming gradient. The loss being minimized is simply the squared distance between the momentum and the newest gradient. So an optimizer is, in effect, running a tiny online regression that fits the momentum to the stream of gradients.

Once momentum is a regression weight, standard model-compression tricks apply. The authors treat the momentum matrix as a low-rank product of two skinny matrices, just as LoRA does for weight updates in fine-tuning. The trick is solving for those two matrices efficiently. Rather than differentiating through the product each step, they apply Newton's method, which happens to produce a closed-form update that looks exactly like the original EMA — a blend of the old factor and a correction term built from the current gradient. Because the update is closed-form and matrix-shaped rather than a full backprop pass, it runs cheaply and can be applied at every single optimization step, so the subspace never goes stale.

For the second-order momentum, squaring would break Adam's square root unless every entry stays positive. The authors solve this by factoring the momentum as (v_B v_A) squared element-wise, which forces positivity by construction while preserving the low-rank structure, then plug it into the same theorem.

The empirical protocol follows GaLore's setup for fairness: identical rank-to-dimension ratios, learning-rate sweeps, and scaling factor, with only the optimizer state representation changing.

Why This Matters

Impact on research. The paper supplies a conceptual bridge between optimizer design and regression, which reframes momentum compression as a well-studied approximation problem rather than an ad hoc engineering trick. It also documents that projection-based methods do not transfer cleanly to newer optimizers like Muon, offering a concrete reason why (periodic subspace staleness corrupting orthogonalization) and a working alternative.

Real-world applications:

• Pre-training under fixed GPU memory. Smaller optimizer state means larger models or larger batch sizes fit on the same hardware, directly affecting the cost of training runs.
• Fine-tuning on consumer or single-node hardware. The 7B/8B fine-tuning results apply to practitioners adapting open models on limited GPUs for domain tasks.
• Distributed training communication. Optimizer state is often sharded or synchronized across devices; compressing it reduces both memory and inter-device traffic.
• Mathematical and reasoning task adaptation. GSM8K and MATH-500 gains suggest the method is useful for fine-tuning models where reasoning accuracy is the bottleneck.

Industry relevance. Memory, not compute, is frequently the binding constraint in LLM training. A drop-in optimizer state compression that works across Adam and Muon, needs no periodic synchronization, and tolerates aggressive rank reduction is directly deployable in existing training stacks. The public code release lowers adoption friction.

Future Directions

• Theoretical guarantees. The main text asserts empirical benefits; the appendix claims analysis of approximation error and convergence, but the truncated content leaves open how tight those bounds are and whether they hold for the full training trajectory.

• Rank selection policy. All experiments fix rank per model size. Whether rank can be adapted automatically during training — growing where gradients are high-rank, shrinking elsewhere — is unexplored.

• Broader optimizer coverage. Only Adam and Muon variants are demonstrated. The framework's reliance on an EMA-style momentum suggests applicability to Lion, Sophia, or second-order methods, but this is untested.

• Scaling beyond 1B parameters. Pre-training evidence stops at 1B parameters and 13.1B tokens. Whether the rank efficiency advantage persists at 7B+ pre-training scale, where gradient structure may differ, remains an open question.

• Interaction with quantization and distributed sharding. Combining low-rank momentum with 8-bit optimizers or ZeRO-style partitioning could compound memory savings, but is not evaluated.

Target Audience

• Optimization and training-systems researchers interested in the theory and practice of memory-efficient LLM training.
• Machine learning engineers who need to train or fine-tune large models under tight GPU memory budgets and want a drop-in alternative to Adam or GaLore.
• Graduate students studying low-rank methods, momentum-based optimizers, or the mathematical structure of gradient descent.
• Practitioners comparing PEFT methods for fine-tuning open-weight models, since the paper benchmarks directly against LoRA, rsLoRA, and DoRA with matched hyperparameters.

Readers should be comfortable with matrix factorization, exponential moving averages, and basic optimization theory; the Newton's method derivation in the appendix is the most demanding part.

Authors’ abstract

Modern optimizers like Adam and Muon are central to training large language models, but their reliance on first- and second-order momenta introduces significant memory overhead, which constrains scalability and computational efficiency. In this work, we reframe the exponential moving average (EMA) used in these momenta as the training of a linear regressor via online gradient flow. Building on this equivalence, we introduce LoRA-Pre, a novel low-rank optimizer designed for efficient pre-training. Specifically, LoRA-Pre reduces the optimizer's memory footprint by decomposing the full momentum matrix into a compact low-rank subspace within the online linear learner, thereby maintaining optimization performance while improving memory efficiency. We empirically validate LoRA-Pre's efficacy by pre-training models from the Llama architecture family, scaling from 60M to 1B parameters. LoRA-Pre achieves the highest performance across all model sizes. Notably, LoRA-Pre demonstrates remarkable rank efficiency, achieving comparable or superior results using only 1/8 the rank of baseline methods. Beyond pre-training, we evaluate LoRA-Pre's effectiveness in fine-tuning scenarios. With the same rank, LoRA-Pre consistently outperforms all efficient fine-tuning baselines. Specifically, compared to standard LoRA, LoRA-Pre achieves substantial improvements of 3.14 points on Llama-3.1-8B and 6.17 points on Llama-2-7B, validating our approach's effectiveness across both pre-training and fine-tuning paradigms. Our code is publicly available at https://github.com/mrflogs/LoRA-Pre.