[ICLR 2025] LoRA-Pre: Taming Momentum via Low-Rank Online Regression

TL;DR

Training Large Language Models (LLMs) is a memory nightmare, largely due to the "hidden" cost of optimizer states (momenta). LoRA-Pre rethinks these states as parameters of an online linear regressor. By factorizing the momentum matrix into two low-rank matrices and updating them using a derived Newton-step closed-form rule, LoRA-Pre achieves better-than-full-rank performance with a fraction of the memory. It works across both pre-training and fine-tuning, outperforming baselines like GaLore and standard LoRA.

The "Momentum Tax" and the Failure of Periodic Updates

When you train a model with Adam, you aren't just storing the weights ($W$); you are storing the first moment ($m$) and second moment ($v$). For a 7B parameter model, this "Momentum Tax" consumes roughly 56GB of VRAM just for the optimizer.

Prior attempts to cut this tax, such as GaLore, project gradients into a low-rank subspace. However, they usually update this subspace periodically (e.g., every 200 steps) because SVD is expensive. This leads to Subspace Lag, where the optimizer is effectively pushing the weights in an outdated, suboptimal direction between updates, causing error accumulation.

The Insight: Momentum is a Secret Regressor

The authors provide a beautiful mathematical bridge. They show that the standard EMA update: $m_t\leftarrow \beta m_{t-1}+(1-\beta )g_t$ is actually the solution to optimizing a linear model to fit the gradient history: ${\min }_{m}L(m;g)=\frac{1}{2}\Vert m-g{\Vert }_F^2$

If momentum is just a linear model, why not compress the model itself? By parameterizing $m$ as $m_B\cdot m_A$ (where $m_B$ and $m_A$ are low-rank), we turn the optimizer state into a Low-Rank Adaptation (LoRA) problem that evolves online.

Methodology: Continuous Subspace Adaptation

Unlike GaLore, which "freezes" the subspace between SVDs, LoRA-Pre uses Newton's Method to derive closed-form update rules for $m_B$ and $m_A$. These updates occur at every single training step, ensuring the subspace is always perfectly aligned with the current gradient flow.

Figure 1: Illustration of LoRA-Pre. Momentum updates are reframed as training a low-rank regressor without backprop.

For the second-order momentum ($v$), which must be positive, they use a clever Hadamard product re-parameterization: $v=(v_B{v}_A{)}^{\circ 2}$. This ensures the optimizer never tries to take the square root of a negative number while maintaining the low-rank benefits.

Experimental Results: Doing More with Less (Rank)

The most striking result is Rank Efficiency. In 60M and 130M model tests, LoRA-Pre reached the same perplexity at Rank 16 that GaLore required Rank 128 or 256 to achieve.

Table 1: Pre-training results showing LoRA-Pre Adam/Muon consistently beating full-rank and low-rank baselines.

Key Highlights:

• Pre-training: On a 1B Llama model, LoRA-Pre Adam achieved a perplexity of 13.53, significantly better than GaLore (15.64) and even standard Adam (15.56).
• Fine-tuning: In math tasks (GSM8K), LoRA-Pre Adam boosted Llama-3.1-8B performance by 3.14 points over standard LoRA.
• Optimizer Versatility: It integrates seamlessly with Muon, a modern preconditioned optimizer, pushing the SOTA even further.

Figure 2: Performance vs. Rank. LoRA-Pre (red/purple) stays stable and efficient even at extremely low ranks where others fail.

Critical Analysis & Conclusion

LoRA-Pre succeeds because it treats the optimizer's "memory" as a dynamic entity rather than a static projection. By updating the low-rank factors via Newton's method, it effectively "learns" the most important directions of the gradient manifold in real-time.

Limitations:

• The coupling of hyperparameters ($\beta$ and $\gamma$) is elegant but means that if your $\beta$ is very close to 1, the low-rank adaptation becomes very slow, potentially leading to instability.
• While it reduces memory, the matrix inversions in the update rule ($m_A{m}_A^T{)}^{-1}$ add a slight computational overhead, though this is negligible for the large matrix dimensions typically found in LLMs.

Takeaway: LoRA-Pre proves that we don't need "full" momentum to train "full-rank" models. As we move toward 100B+ parameter models, online low-rank state compression will likely become the standard requirement for efficient training.