Optimal low-rank stochastic gradient estimation for LLM training

Zehao Li; Tao Ren; Zishi Zhang; Xi Chen; Yijie Peng

cs.LGMar 21, 2026

Optimal low-rank stochastic gradient estimation for LLM training

Authors: Zehao Li, Tao Ren, Zishi Zhang, Xi Chen, Yijie Peng

Organizations: Guanghua School of Management, Peking University, Beijing, 100871, China · Stern School of Business, New York University, New York, NY 10012, USA

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.

Figures & tables

Figure 1: Illustration of the proposed low-rank gradient estimator and the lazy-update gradient descent framework. The gradient estimator with respect to $\Theta$ is computed via a rank- $r$ reparameterization using an auxiliary variable $B$ and a randomly sampled projection matrix $V$ (yellow trapezoids), and then lifted back to the original parameter space by multiplying $V^{\top}$ (blue block). In the lazy-update scheme, the same projection direction $V_{0}$ is reused for $K$ inner steps before switching to a new projection direction $V_{1}$ (green trapezoids).

[LowRank-IPA]\hskip 9.24994ptB_{(t,k+1)}=B_{(t,k)}-\alpha_{t}\nabla_{B}F(\xi_{(t,k)},\Theta_{t}+BV_{t}^{\top})\big|_{B=B_{(t,k)}}.

(8)

Algorithm 1 (Low-rank gradient descent with lazy update)

\Pr(i\in J)=\pi_{i}^{\star},\hskip 18.49988pti=1,\ldots,n,

Algorithm 4 Instance-dependent sampling of an optimal low-rank projector

Panel of Figure 2 — Figure 2: MSE versus samples plot of independent Likelihood Ratio estimator

Panel of Figure 3 — Figure 3: MSE versus samples plot of independent Infinitesimal Perturbation Analysis estimator

Panel of Figure 4 — Figure 4: MSE versus samples plot of dependent Likelihood Ratio (LR) estimator

Panel of Figure 5 — Figure 5: MSE versus samples plot of dependent Infinitesimal Perturbation Analysis (IPA) estimator

Panel of Figure 6 — Figure 6: Training loss of Stiefel and Gaussian subspace sampling across six datasets

Method	SST-2	SST-5	SNLI	MNLI	RTE	TREC
Zero-shot	79.0	35.5	50.2	48.8	51.4	32.0
Vanilla LR	86.3	40.8	68.5	56.7	58.6	62.4
Gaussian LowRank-LR	88.0	41.1	73.4	61.6	61.2	77.9
Stiefel LowRank-LR	90.2	43.2	74.3	59.4	63.7	80.9
Coordinate LowRank-LR	91.3	42.6	71.0	57.2	62.2	79.4
Vanilla IPA	91.9	47.5	77.5	70.0	66.4	85.0

Table 1: RoBERTa-large fine-tuning results on six classification datasets.

	Vanilla IPA	LowRank-IPA	Vanilla LR	LowRank-LR
Memory (GB)	16.7	14.3	5.49	3.83

Table 2: Memory profile of different methods for fine-tuning the RoBERTa-large.

	Vanilla IPA	LowRank-IPA	Vanilla LR	LowRank-LR
Time (second)	0.784	0.787	0.468	0.493

Table 3: Per-step wall clock time.

Panel of Figure 7 — Figure 7: Pretraining 20M language model using IPA gradient estimator

Panel of Figure 8 — Figure 8: Pretraining 60M language model using IPA gradient estimator

Panel of Figure 9 — Figure 9: Pretraining 100M language model using IPA gradient estimator

LLM Training

Large language model training studies the data, objectives, optimizers, and systems used to build and adapt language models.

See all 342 papers

Geometrically Principled Randomized Optimization for Efficient LLM Training
Oct 2, 2025Sahar Rajabi, Nayeema Nonta, Sirisha RambhatlaLanguage Model Pretraining Deep Learning Optimization
CriticalML Lab, University of Waterloo
MoARa: Module-Aware Rank Allocation and Structure-Preserving Decomposition for Low-Rank LLM Pre-training
Sep 14, 2026Keunyoung Kim, Nojun KwakLanguage Model Pretraining Low-Rank Approximation
Seoul National University
Ori: Drift-Aware Streaming $k$ -PCA for Very Low-Rank LLM Training
Oct 16, 2026Xinyu Luo, Jincheng Zhou, Mengbo Wang +5LLM Training Low-Rank Adaptation
Department of Computer Science, Purdue University. · Purdue Institute for Cancer Research, Purdue University. · Department of Comparative Pathobiology, Purdue University.

Oct 2, 2025J/K move · Enter open · S save

Geometrically Principled Randomized Optimization for Efficient LLM Training

Sahar Rajabi, Nayeema Nonta, Sirisha Rambhatla

CriticalML Lab, University of Waterloo

Language Model Pretraining Deep Learning Optimization Riemannian Optimization Randomized Sketching LLM Training

Open paper AskPDF arXiv

Low-rank gradient optimization for large language models is currently divided into two categories: structured methods that rigorously identify subspaces, and randomized approaches employed primarily for computational efficiency. In this work, we question the intuition behind why random projections are effective. We trace this phenomenon to the geometry of the gradient subspaces, which exhibits subspace optimization landscape has a nearly flat curvature, while a significant portion of gradient information lies outside the core subspace. Leveraging these insights, and drawing on randomized linear algebra, we theoretically establish that random low-rank projections preserve the geometry, and we introduce GrassWalk and GrassJump, algorithms that navigate the Grassmannian manifold via random walks and jumps. By coupling this randomized exploration with subspace-aware optimizer and recovering the lost gradient signals, we achieve state-of-the-art results on LLaMA-1B, LLaMA-7B, and Qwen-1.5B pretraining. Our findings reframe randomization not merely as a computational shortcut, but as a geometrically principled approach to high-dimensional optimizations.

Abstract

LLM Training

Explore similar work

Geometrically Principled Randomized Optimization for Efficient LLM Training

MoARa: Module-Aware Rank Allocation and Structure-Preserving Decomposition for Low-Rank LLM Pre-training

Ori: Drift-Aware Streaming kkk-PCA for Very Low-Rank LLM Training

Ori: Drift-Aware Streaming $k$ -PCA for Very Low-Rank LLM Training