Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation

Zhengbo Wang, Jian Liang, Ran He, Zilei Wang, Tieniu Tan

TL;DR

LoRA-Pre is introduced, a novel low-rank optimizer designed for efficient pre-training that 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.

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.

Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation

TL;DR

LoRA-Pre is introduced, a novel low-rank optimizer designed for efficient pre-training that 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.

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.
Paper Structure (29 sections, 5 theorems, 64 equations, 3 figures, 4 tables, 2 algorithms)

This paper contains 29 sections, 5 theorems, 64 equations, 3 figures, 4 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Works
  3. Method
  4. Preliminary
  5. Your Momentum is a Secret Online Linear Regressor
  6. LoRA-Pre: Low-Rank Online Linear Regression
  7. Experimental Results
  8. Memory-Efficient Pre-training
  9. Memory-Efficient Fine-Tuning
  10. Ablation Study
  11. Conclusion
  12. Proof of Theoretical Results
  13. Detailed Algorithms of LoRA-Pre for Adam and Muon Optimizers
  14. Algorithm of LoRA-Pre for Adam
  15. Algorithm of LoRA-Pre for Muon
  16. ...and 14 more sections

Key Result

Theorem 3.1

Assume both matrices $m_B\in\mathbb{R}^{p\times r}$ and $m_A\in\mathbb{R}^{r\times q}$ are full rank. For the objective $\min_{m_B, m_A} L(m_B, m_A; g) = \frac{1}{2}\cdot \|m_B m_A - g\|^2_F$, Newton's method yields the following closed-form update rules: Here, $\gamma_1$ is the learning rate for the factorized optimization problem.

Figures (3)

  • Figure 1: Illustration of our LoRA-Pre method. In this work, we establish a novel connection: the exponential moving average (EMA) update for optimizer momentum is mathematically equivalent to training a linear regressor using online gradient descent. Leveraging this equivalence, we propose compressing the optimizer states (i.e., the momenta) using low-rank matrices to reduce the memory footprint. Finally, the closed-form update rules for these matrices without requiring back-propagation are given by Theorem \ref{['thm:first_order']}.
  • Figure 2: Rank efficiency comparison across efficient optimization methods. Perplexity versus rank for 60M (left) and 130M (right) models, demonstrating LoRA-Pre's superior performance at lower ranks compared to baseline methods.
  • Figure 3: Test perplexity for LoRA-Pre Muon with different ranks during training.

Theorems & Definitions (10)

  • Theorem 3.1
  • proof
  • Theorem
  • proof
  • Lemma C.1
  • proof
  • Lemma C.2
  • proof
  • Theorem C.3
  • proof