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

Xinyu Luo; Jincheng Zhou; Mengbo Wang; Zilin Shen; Nadia Atallah Lanman; Petros Drineas; Bruno Ribeiro; Brian Bullins

cs.LGOct 16, 2026

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

Authors: Xinyu Luo, Jincheng Zhou, Mengbo Wang, Zilin Shen, Nadia Atallah Lanman, Petros Drineas, Bruno Ribeiro, Brian Bullins

Organizations: Department of Computer Science, Purdue University. · Purdue Institute for Cancer Research, Purdue University. · Department of Comparative Pathobiology, Purdue University.

Abstract

A prominent family of low-rank LLM optimizers, pioneered by GaLore, performs gradient-subspace projection: at each step, they project the gradient into a low-dimensional subspace and map the update back to the full model. Most existing work either recomputes the gradient subspace at certain intervals, or performs post-hoc adaptation to catch up to the moving subspace. In this work, we explicitly identify in practice the slow, stable manner in which LLM gradient subspaces drift across time. To this end, we propose Ori, a novel low-rank optimizer for training LLMs that captures the shifting principal gradient directions over time to achieve state-of-the-art accuracy under stringent low-rank constraints. Ori employs a streaming $k$ -PCA algorithm inspired by Oja++, whose low per-step cost enables high-frequency, on-the-fly subspace tracking. In practice, Ori reduces the performance gap of the state-of-the-art low-rank optimizers to full-rank optimizers by up to 69.3% on a 1B LLaMA pretraining task, and achieves the best performance on 10 out of 16 task scenarios in Roberta-Base fine-tuning. Together, these contributions allow Ori to further push the performance and efficiency frontiers of gradient-projecting optimizers, narrowing the gap to full-rank training for LLMs. Code can be found at https://anonymous.4open.science/r/Ori-Submission-2132.

Figures & tables

Figure 1: Trade-off between evaluation perplexity (Y-axis, lower is better) and peak optimizer state memory (X-axis, lower is better) for various low-rank optimizers in different ranks on the 1B LLaMA pre-training task. Bottom-left corner is the ideal region. Ori (

\star

) forms a better Pareto frontier, achieving a superior performance-to-memory trade-off compared to state-of-the-art baselines like GaLore (

\circ

) and OSD (

\square

Figure 2: Average gradient subspace overlap during pre-training. We plot the overlap ( Equation 1 ) between subspaces at step

t

and

t+\Delta t

(with

\Delta t=10

), averaged over a 60M LLaMA model. The high, stable overlap, significantly above a random baseline, confirms that the subspace evolves gradually.

Full-Rank Methods
Method	60M ( $d_{\text{model}}=256$ )			130M ( $d_{\text{model}}=768$ )			350M ( $d_{\text{model}}=1024$ )			1B ( $d_{\text{model}}=2048$ )
Method	$r=128$	$r=32$	$r=8$	$r=256$	$r=64$	$r=16$	$r=256$	$r=64$	$r=16$	$r=512$	$r=128$	$r=32$
Adam		39.03			27.60			20.68			19.92
Low-Rank Methods
ReLoRA [ 14 ]	144.48	103.96	103.99	85.53	85.54	85.54	72.59	72.72	72.69	65.82	65.83	65.79
GaLore [ 39 ]	39.38	53.27	75.65	27.70	36.82	60.29	21.26	26.31	41.25	16.87	26.49	33.82

Table 1: Comparison with full- and low-rank optimizers on pre-training various sizes of LLaMA models on C4 dataset. Validation perplexity is reported (lower

\downarrow

is better).

k

is the rank for low-rank optimizers, and

d_{\text{model}}

is the model embedding dimension. We show the best scores in bold and second best ones with underlines. Ori is on-par with the best baseline on moderately high rank settings, and achieves superior performance when the rank

k

is severely constrained.

Method	60M ( $d_{\text{model}}=256$ )			130M ( $d_{\text{model}}=768$ )			350M ( $d_{\text{model}}=1024$ )			1B ( $d_{\text{model}}=2048$ )
Method	$r=128$	$r=32$	$r=8$	$r=256$	$r=64$	$r=16$	$r=256$	$r=64$	$r=16$	$r=512$	$r=128$	$r=32$
Adam		39.03			27.60			20.68			19.92
LDAdam ( $T=1$ ) [ 28 ]	34.57	36.19	39.90	29.32	30.91	30.02	19.97	23.52	23.67	17.69	16.29	20.27
LDAdam ( $T=200$ )	37.75	47.90	69.74	27.64	30.95	41.16	22.38	25.26	32.03	161.62	162.92	45.66
O-LDAdam (Ours) ( $T=200$ )	41.60	38.52	39.14	31.20	28.58	27.55	22.84	23.87	22.39	22.47	16.53	15.66
Training Tokens	0.44B			0.88B			2.5B			5.2B

Table 2: Comparison of O-LDAdam versus LDAdam [ 28 ] with different subspace projection update frequencies (

T\in\{1,200\}

) when pretraining LLaMA on C4 dataset. O-LDAdam improves LDAdam’s single-step power iteration update with Ori-inspired multi-step, stochastic power iteration update. In the very low-rank regimes (e.g.

r=32

when

d_{\text{model}}=2048

), O-LDAdam with update frequency

T=200

achieves on-par or even superior performance to LDAdam with per-step updates (

T=1

Method	Rank	CoLA	STS-B	MRPC	RTE	SST2	MNLI	QNLI	QQP	Avg.
Adam	Full	61.61 $\pm$ 0.46	90.88 $\pm$ 0.05	92.52 $\pm$ 0.12	78.45 $\pm$ 1.50	93.88 $\pm$ 0.67	87.41 $\pm$ 0.05	92.63 $\pm$ 0.37	89.21 $\pm$ 0.16	85.82
LoRA [ 10 ]	8	41.85 $\pm$ 1.66	87.07 $\pm$ 0.25	87.95 $\pm$ 2.40	67.03 $\pm$ 2.40	93.12 $\pm$ 0.20	85.14 $\pm$ 0.10	90.62 $\pm$ 0.24	84.66 $\pm$ 0.14	79.68
GaLore [ 39 ]	8	58.24 $\pm$ 0.63	90.58 $\pm$ 0.07	91.06 $\pm$ 0.56	76.30 $\pm$ 1.46	93.88 $\pm$ 0.24	86.14 $\pm$ 0.14	92.41 $\pm$ 0.25	86.10 $\pm$ 0.30	84.34
OSD [ 15 ]	8	60.39 $\pm$ 0.31	90.93 $\pm$ 0.19	91.85 $\pm$ 0.50	77.36 $\pm$ 1.46	94.61 $\pm$ 0.30	86.62 $\pm$ 0.10	92.29 $\pm$ 0.16	87.66 $\pm$ 0.78	85.16
LD Adam ( $T=1$ )	8	60.33 $\pm$ 1.07	90.84 $\pm$ 0.05	91.49 $\pm$ 0.16	76.65 $\pm$ 0.74	94.64 $\pm$ 0.27	86.98 $\pm$ 0.27	92.70 $\pm$ 0.31	89.90 $\pm$ 1.33	85.96
LD Adam ( $T=200$ )	8	59.07 $\pm$ 0.87	90.54 $\pm$ 0.07	91.05 $\pm$ 0.47	76.29 $\pm$ 1.33	94.61 $\pm$ 0.34	87.13 $\pm$ 0.07	92.46 $\pm$ 0.14	88.54 $\pm$ 0.16	84.96

Table 3: Comparison with full- and low-rank optimizers on fine-tuning Roberta-Base on GLUE dataset. Average accuracy score and standard deviation is reported (lower

\downarrow

is better) for each task, optimizer, and rank combination over three random seeds. We also report average accuracy over all tasks for each method. We show the best scores in bold and second best ones with underlines. Ori consistently achieves superior performance on 10 out of 16 scenarios.

Panel of Figure 3 — Figure 3: Ablation Studies on the 60M LLaMA model pre-trained on C4 dataset for 0.44B tokens. ( 3(a) ) Varying the drift coefficient $\rho$ shows peak performance around $\rho=0.99$ , validating the importance of transferring optimizer statistics. ( 3(b) ) Performance degrades as the update frequency $T_{\text{period}}$ increases , demonstrating the cost of a stale projector. ( 3(c) ) Oja++ consistently outperforms the standard Oja algorithm across all ranks.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Task	CoLA	STS-B	MRPC	RTE	SST-2	MNLI	QNLI	QQP
Metric	Matthews Correlation	Pearson	F1	Accuracy	Accuracy	Accuracy	Accuracy	F1

Appendix

Table 4: Accuracy metrics we used for the Roberta-Base fine-tuning experiments of Table 3 .

Ori Hyperparameters	60M			130M			350M			1B
Ori Hyperparameters	$r=128$	$r=32$	$r=8$	$r=256$	$r=64$	$r=16$	$r=256$	$r=64$	$r=16$	$r=512$	$r=128$	$r=32$
$\gamma$	0.5	0.5	1.0	0.25	0.25	0.25	0.75	0.75	0.5	0.5	0.5	0.5
$\rho$	0.99	0.99	0.99	0.99	0.99	0.99	0.99	0.99	0.99	0.99	0.99	0.99
$T_{\text{period}}$	20	20	20	20	20	20	20	20	20	20	20	20
$T_{\text{warmup}}$	20	20	20	20	20	20	20	20	20	20	20	20
$\eta$	1.0	1.0	1.0	1.0	1.0	1.0	1.0	1.0	1.0	1.0	1.0	1.0

Appendix

Table 5: Ori hyperparameter used for the LLaMA pre-training results on the C4 dataset ( Table 1 ).

Ori Hyperparameters	GLUE all tasks
$\gamma$	4.0
$\rho$	0.99
$T_{\text{period}}$	20
$T_{\text{warmup}}$	20
$\eta$	1.0
$N_{\text{update}}$	2

Appendix

Table 6: Ori hyperparameter used for the Roberta-Base fine-tuning results on the GLUE dataset ( Table 3 )

Ori Hyperparameters	60M			130M			350M			1B
Ori Hyperparameters	$r=128$	$r=32$	$r=8$	$r=256$	$r=64$	$r=16$	$r=256$	$r=64$	$r=16$	$r=512$	$r=128$	$r=32$
$T_{\text{period}}$	200	200	200	200	200	200	200	200	200	200	200	200
$T_{\text{warmup}}$	0	0	0	0	0	0	0	0	0	0	0	0
$\eta$	1.0	1.0	1.0	1.0	1.0	1.0	1.0	1.0	1.0	1.0	1.0	1.0
$N_{\text{steps}}$	8	8	8	8	8	8	16	16	16	16	16	16
$N_{\text{vecs}}$	1	1	1	1	1	1	1	1	1	1	1	1

Appendix

Table 7: O-LDAdam hyperparameter related to the Ori-inspired stochastic power iteration update used for the LLaMA pre-training results on the C4 dataset ( Table 2 ).

Ori Hyperparameters	GLUE task MNLI, MRPC, CoLA, QQP, RTE, STSB	GLUE task SST2, QNLI
$T_{\text{period}}$	200	200
$T_{\text{warmup}}$	0	0
$\eta$	1.0	1.0
$N_{\text{steps}}$	4	2
$N_{\text{vecs}}$	1	1

Appendix

Table 8: O-LDAdam hyperparameter related to the Ori-inspired stochastic power iteration update used for the Roberta-Base fine-tuning results on the GLUE dataset ( Table 3 )

Optimizer State Peak Memory (GB)
Optimizer rank	$r=512$	$r=128$	$r=32$
Adam	2.4942	2.4942	2.4942
ReLoRA [ 14 ]	1.1349	0.4670	0.3000
GaLore [ 39 ]	0.3385	0.3385	0.3385
OSD [ 15 ]	0.2832	0.2832	0.2832
Ori (Ours)	0.2583	0.2506	0.2487

Appendix

Table 9: Memory and Wallclock Time Comparison between Ori and non-error-feedback baselines (Adam, ReLoRA, GaLore, OSD). Analysis is conducted by pre-training 1B LLaMA on C4 dataset.

Optimizer State Peak Memory (GB)
Optimizer rank	$r=512$	$r=128$	$r=32$
LDAdam ( $T=1$ )	0.1221	0.1221	0.1221
LDAdam ( $T=200$ )	0.1221	0.1221	0.1221
O-LDAdam (Ours)	0.1221	0.1221	0.1221
Average Update Time (s / Step)
LDAdam ( $T=1$ )	13.9294	13.3475	13.1410

Appendix

Table 10: Memory and Wallclock Time Comparison between O-LDAdam and an error-feedback baseline, LDAdam. Analysis is conducted by pre-training 1B LLaMA on C4 dataset.

Optimizer Peak Memory (GB)
Method	60M			130M			350M			1B
Method	$r=128$	$r=32$	$r=8$	$r=256$	$r=64$	$r=16$	$r=256$	$r=64$	$r=16$	$r=512$	$r=128$	$r=32$
Adam	0.1082	–	–	0.2498	–	–	0.6854	–	–	2.4942	–	–
ReLoRA [ 14 ]	0.0797	0.0657	0.0622	0.1751	0.1125	0.0968	0.3450	0.1779	0.1361	1.1349	0.4670	0.3000
GaLore [ 39 ]	0.0670	0.0670	0.0670	0.1049	0.1049	0.1049	0.1458	0.1458	0.1458	0.3385	0.3385	0.3385
OSD [ 15 ]	0.0635	0.0635	0.0635	0.0970	0.0970	0.0970	0.1318	0.1318	0.1318	0.2832	0.2832	0.2832

Appendix

Table 11: Memory and Wallclock Time Comparison of pre-training various sized LLaMa model on the C4 dataset.

LLM Training

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

See all 342 papers

Optimal low-rank stochastic gradient estimation for LLM training
Mar 21, 2026Zehao Li, Tao Ren, Zishi Zhang +2LLM Training Memory-Efficient Training
Guanghua School of Management, Peking University, Beijing, 100871, China · Stern School of Business, New York University, New York, NY 10012, USA
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
Zeroth-Order Fine-Tuning of LLMs in Learnable Low-Rank Subspaces
Oct 10, 2026Chuang Tong, Ling ChenLLM Fine-Tuning Zeroth-Order Optimization
State Key Laboratory of Blockchain and Data Security, Zhejiang University · College of Computer Science and Technology, Zhejiang University

Mar 21, 2026J/K move · Enter open · S save

Optimal low-rank stochastic gradient estimation for LLM training

Zehao Li, Tao Ren, Zishi Zhang, Xi Chen+1

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

LLM Training Memory-Efficient Training Stochastic Gradient Descent Low-Rank Compression

Open paper AskPDF arXiv

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.

Abstract

LLM Training

Explore similar work

Optimal low-rank stochastic gradient estimation for LLM training

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

Zeroth-Order Fine-Tuning of LLMs in Learnable Low-Rank Subspaces