OASIS: Online Activation Subspace Learning for Memory-Efficient Training

Sakshi Choudhary; Utkarsh Saxena; Kaushik Roy

cs.LGApr 10, 2026

OASIS: Online Activation Subspace Learning for Memory-Efficient Training

Authors: Sakshi Choudhary, Utkarsh Saxena, Kaushik Roy

Organizations: Purdue University · Advanced Micro Devices, Inc. (AMD)

Abstract

Training large language models (LLMs) is constrained by memory requirements, with activations accounting for a substantial fraction of the total footprint. Existing approaches reduce memory using low-rank weight parameterizations or low-rank gradient subspaces for optimizer states, while activation memory is addressed through architectural modifications or compression schemes based on periodically updated projections. We propose OASIS, an online activation subspace learning algorithm for memory-efficient training that tracks and continuously updates a low-dimensional activation subspace during training. Intermediate activations are projected onto this evolving subspace, reducing memory without modifying forward-pass computations. The evolving activation subspace induces low-rank gradient representations, enabling both gradients and optimizer states to be maintained directly in this subspace, while a projection-aware optimizer consistently transports optimizer states across subspace updates for stable training. Across various finetuning and pretraining tasks, OASIS achieves up to $2\times$ lower peak memory than full fine-tuning while matching its performance and outperforming prior low-rank methods.

Figures & tables

Panel of Figure 1 — Figure 1: Memory breakdown and subspace dynamics during training. (a) As batch size and sequence length scale, memory becomes increasingly dominated by activations. (b) These activations lie in subspaces that evolve continuously during training.

Model	Rank	Method	Accuracy (%)	Peak Memory (GB)
Llama-2 7B	–	Adam	39.37 $\pm$ 0.39	95.18
	32	LoRA	37.13 $\pm$ 0.64	64.15
		GaLore	35.96 $\pm$ 1.18	71.30
		LDAdam	38.74 $\pm$ 0.11	71.48
		OASIS	39.14 $\pm$ 0.72	48.21
	128	LoRA	37.27 $\pm$ 0.36	66.00

Table 1: Finetuning results on GSM8K for Llama-2 7B and Llama-3 1B. For Llama-2 7B, OASIS matches the performance of Adam while reducing peak memory by

\sim 2\times

Figure 2: Memory breakdown across training components for different methods on LLaMA-2 7B GSM8K finetuning at rank 32. OASIS reduces memory across activations, gradients, and optimizer states, leading to the lowest overall memory footprint.

Rank	Method	Accuracy (%)	Peak Memory (GB)
–	Adam	34.01 $\pm$ 1.29	93.26
4	LoRA	28.34 $\pm$ 2.79	61.70
	GaLore	27.42 $\pm$ 1.22	69.00
	LDAdam	27.93 $\pm$ 1.41	68.99
	OASIS	29.28 $\pm$ 0.61	46.16
32	LoRA	31.84 $\pm$ 1.74	62.08

Table 2: Finetuning results on HumanEval for Llama-2 7B (trained on CodeAlpaca). OASIS achieves the best performance among low-rank methods at aggressive compression, and remains competitive with the strongest baseline at higher rank, while consistently using the lowest memory.

Method	Llama-130M		Llama-350M
	Val Loss	Mem	Val Loss	Mem
Adam	3.21	41.00	3.03	89.31
GaLore	3.55	40.73	3.39	88.22
LDAdam	3.23	40.73	2.86	88.22
OASIS	3.28	37.10	3.02	79.46

Table 3: Pretraining results on C4 across Llama model scales. OASIS achieves comparable validation loss to Adam while consistently reducing memory usage.

Panel of Figure 3 — Figure 3: Subspace drift during pretraining on C4 for Llama-130M and Llama-350M. Both models exhibit a transient phase followed by stabilization at a non-zero drift, indicating that the activation subspace evolves throughout training.

Panel of Figure 4 — Figure 4: Comparison of periodic PCA and OASIS during finetuning on GSM8K with Llama-2 7B. (a) Periodic PCA consistently underperforms OASIS across all tested update intervals. (b) OASIS also converges faster, achieving lower training loss compared to periodic PCA.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	GSM8K (Table 1 )	HumanEval (Table 2 )
Batch size	32	16
Sequence length	512	1024
Learning rate scheduler	Cosine	Cosine
Learning rate (LR)	$\{2\text{e-}4,5\text{e-}4,9\text{e-}4,1\text{e-}3,1\text{e-}2\}$	$\{2\text{e-}4,5\text{e-}4,9\text{e-}4,1\text{e-}3,1\text{e-}2\}$
Warmup steps	5%	5%
Epochs	3	3

Appendix

Table 4: Hyperparameters for finetuning experiments on GSM8K and HumanEval.

Hyperparameter	Llama-130M	Llama-350M
Batch size	512	512
Sequence length	256	256
Training iterations	20k	60k
Learning rate (LR)	$\{5\text{e-}4,9\text{e-}4,1\text{e-}3,5\text{e-}3,1\text{e-}2\}$	$\{5\text{e-}4,9\text{e-}4,1\text{e-}3,5\text{e-}3,1\text{e-}2\}$
Learning rate scheduler	Cosine	Cosine
Warmup	10%	10%

Appendix

Table 5: Hyperparameters for pretraining experiments on C4 presented in Table 3 .

Model	Hidden	Intermediate	Heads	Layers
Llama-130M	768	2048	12	12
Llama-350M	1024	2736	16	24

Appendix

Table 6: Model configurations for pretraining experiments.

Memory-Efficient Training

Memory-efficient training reduces the memory needed to optimize neural networks by compressing, recomputing, or relocating training state.

See all 87 papers

PRAC: Principal-Random Subspace for LLM Activation Compression and Memory-Efficient Training
Feb 26, 2026Yanyi Li, Yimu Zhang, Cong FangLLM Compression Memory-Efficient Training
School of Intelligence Science and Technology, Peking University, China · Institute for Artificial Intelligence, Peking University, China
AIM-ZO: Activation-Informed Subspace Maintenance for Zeroth-Order LLM Fine-Tuning
Sep 28, 2026Yue Xie, Zhi Zheng, Yunpeng Ba +5Subspace Optimization Memory-Efficient Fine-Tuning
Southern University of Science and Technology · National University of Singapore · Huawei Technologies Ltd. +1
ZO-Act: Efficient Zeroth-Order Fine-Tuning via One-Shot Activation-Informed Low-Rank Subspaces
Jul 1, 2026Xun Dong, Yibo Xu, Naigang Wang +3Zeroth-Order Optimization LLM Fine-Tuning
University at Albany, SUNY · IBM T. J. Watson Research Center

Feb 26, 2026J/K move · Enter open · S save

PRAC: Principal-Random Subspace for LLM Activation Compression and Memory-Efficient Training

Yanyi Li, Yimu Zhang, Cong Fang

School of Intelligence Science and Technology, Peking University, China · Institute for Artificial Intelligence, Peking University, China

LLM Compression Memory-Efficient Training Low-Rank Compression Low-Rank Approximation

Open paper AskPDF arXiv

Activations have become the primary memory bottleneck in large-batch LLM training. However, existing compression methods fail to exploit the spectral structure of activations, resulting in slow convergence or limited compression. To address this, we bridge the relationship between the algorithm's fast convergence and the requirements for subspace projection, and show that an effective compression should yield an unbiased estimate of the original activation with low variance. We propose Principal-Random Subspace for LLM Activation Compression (PRAC), which novelly decomposes activations into two components: a principal subspace captured via SVD to retain dominant information, and a random subspace sampled from the orthogonal complement to approximate the tail. By introducing a precise scaling factor, we prove that PRAC yields an unbiased gradient estimator with minimum variance under certain conditions. Extensive experiments on pre-training and fine-tuning tasks demonstrate that PRAC achieves up to 36% total memory reduction with negligible performance degradation and minimal computational cost.

Abstract

Memory-Efficient Training

Explore similar work

PRAC: Principal-Random Subspace for LLM Activation Compression and Memory-Efficient Training

AIM-ZO: Activation-Informed Subspace Maintenance for Zeroth-Order LLM Fine-Tuning

ZO-Act: Efficient Zeroth-Order Fine-Tuning via One-Shot Activation-Informed Low-Rank Subspaces