Optimal low-rank stochastic gradient estimation for LLM training

TL;DR

This paper introduces an optimal low-rank stochastic gradient estimator that reduces memory usage and improves training efficiency for large language models by projecting gradients onto low-dimensional subspaces with minimal variance.

Contribution

It proposes a novel unbiased low-rank gradient estimator with optimal projection distribution, applicable across various stochastic gradient paradigms, enhancing memory efficiency and training performance.

Findings

Achieves significant GPU memory savings in RoBERTa-large fine-tuning.

Outperforms traditional gradient estimation methods in LLM pretraining.

Maintains competitive accuracy while reducing memory consumption.

Abstract

Large language model (LLM) training is often bottlenecked by memory constraints and stochastic gradient noise in extremely high-dimensional parameter spaces. Motivated by empirical evidence that many LLM gradient matrices are effectively low-rank during training, we present an unbiased, memory-efficient, low-rank matrix estimator with the lowest variance that is applicable across common stochastic gradient estimation paradigms. The core idea is to project a high-dimensional stochastic gradient estimator onto a random low-dimensional subspace and lift it back, reducing memory while keeping the estimator unbiased and controlling mean-squared error via an optimally designed projection distribution, including Haar--Stiefel projections. The projection distribution is derived by solving a constrained functional optimization problem, yielding an optimal random projector that guides algorithm design. Empirically, the resulting low-rank gradient estimators deliver both practical memory savings and improved training behavior. In RoBERTa-large fine-tuning, our method attains the lowest peak GPU memory among compared methods (e.g., 3.83GB versus 16.7GB for full BP) while remaining competitive in accuracy; in autoregressive LLM pretraining (LLaMA-20M/60M/100M), our method outperforms the traditional methods, supporting the benefit of the proposed optimal projection strategy.