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

Yanyi Li; Yimu Zhang; Cong Fang

cs.LGFeb 26, 2026

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

Authors: Yanyi Li, Yimu Zhang, Cong Fang

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

Abstract

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.

Figures & tables

Panel of Figure 1 — Figure 1: The proposed PRAC projects activations onto both the principal and random subspaces and yields the minimum variance unbiased estimator, thus achieving up to 36% total memory reduction with negligible performance degradation. (Left) The flowchart of PRAC; (Middle) A conceptual comparison of subspace strategies; (Right) Performance comparison on LLaMA-1B.

Figure 2: Singular value spectrum (Left) and cumulative energy ratio (Right).

Algorithm	Form	Scaling	Unbiasedness	Variance
PAC	$\tilde{X}_{\text{pac}}=(XQ_{1})Q_{1}^{\top}$	-	$\times$	-
RAC	$\tilde{X}_{\text{rac}}=(XQ_{2})Q_{2}^{\top}$	$n/r_{2}$	$\checkmark$	High
PRAC	$\tilde{X}_{\text{prac}}=(XP)P^{\top}$	$(n-r_{1})/r_{2}$	$\checkmark$	Minimum

Table 1: Comparison of Estimation Unbiasedness and Variance across Projection Methods.

	LLaMA-130M			LLaMA-350M			LLaMA-1B			GPT-2-124M			GPT-2-355M
Method	PPL	MEM	$\Delta$	PPL	MEM	$\Delta$	PPL	MEM	$\Delta$	PPL	MEM	$\Delta$	PPL	MEM	$\Delta$
Baseline	24.41	21.26	-	18.77	46.10	-	15.33	OOM		19.83	62.27	-	16.26	55.94	-
GaLore	25.12	21.08	-1%	19.65	45.37	-2%	15.64	OOM		21.01	62.15	-0%	17.88	55.64	-1%
RSO	25.41	19.57	-8%	19.57	40.36	-13%	15.72	79.32	-16%	22.31	54.79	-12%	18.29	46.90	-16%
PRAC	24.66	15.65	-27%	18.92	32.21	-30%	15.41	60.48	-36%	19.95	49.68	-20%	16.35	44.21	-21%
B / S	128/256			128/256			128/256			64/1024			32/1024

Table 2: Comparison with memory-efficient algorithms on pretraining tasks. Validation perplexity (PPL, lower is better) and peak memory consumption (MEM, in GB, lower is better) are reported. For methods marked out-of-memory (OOM), perplexity is measured using half the batch size specified above.

B,S,\Delta

denote the micro-batch size, sequence length and total memory reduction respectively.

	Baseline	PRAC
GeLU Input	$bsn$	$bs(r_{1}+r_{2})$
$W_{\text{down}}\;\text{Input}$	$bsn$	$bs(r_{1}+r_{2})$
Total	$2bsn$	$2bs(r_{1}+r_{2})$

Table 3: Activation memory footprint comparison for the GELU and linear layer.

Figure 3: Loss curves of pre-training LLaMA-130M and LLaMA-350M model.

Method	Batch Size	Memory (GB)	Training Time (h)
Baseline	64	50.35	156
PRAC	64	36.90	179
PRAC	96 ( $\bf{\uparrow 50\%}$ )	48.70	117 ( $\bf{\downarrow 25\%}$ )

Table 4: Training Time Comparison on

4\times

A800 GPU for LLaMA-1B.

Figure 4: Scaling behavior of PRAC across model sizes (35M to 1B parameters)

Method	Memory (GB)	COLA	STSB	MRPC	RTE	SST2	MNLI	QNLI	QQP	AVG
Full Fine-Tuning	0.42	62.24	90.92	91.30	79.42	94.57	87.18	92.33	92.28	86.28
RSO (rank=4)	0.29	62.47	90.62	92.25	78.70	94.84	86.67	92.29	90.94	86.10
PRAC (rank=4)	0.26 ( $\bf{\downarrow 38\%}$ )	63.56	90.83	93.28	78.71	94.72	86.60	92.35	90.95	86.38

Table 5: Comparison with memory-efficient algorithms on fine-tuning RoBERTa models. Average scores across the GLUE benchmark and the peak memory usage of activations in MRPC task are provided.

Figure 5: Loss curve of using PRAC, PAC, RAC in LLaMA series pre-training. RAC reports the result of

r=0.6

due to the divergence of

r=0.3

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Operation	Activation Saved
Operation	Baseline	PRAC $(0\leq r_{1}+r_{2}\leq 0.6)$
$X=\mathrm{LayerNorm}(\tilde{X})$	$\tilde{X}\in\mathbb{R}^{b\times s\times n}$ $(\mu(\tilde{X}),\sigma(\tilde{X})^{2})\in\mathbb{R}^{b\times s}$	${\color[rgb]{1,0,0}\tilde{X}P_{x}\in\mathbb{R}^{b\times s\times(nr_{1})},\tilde{X}Q_{x}\in\mathbb{R}^{b\times s\times(nr_{2})}}$ $(\mu(X),\sigma(X)^{2})\in\mathbb{R}^{b\times s}$
$q=XW_{q}\$ $k=XW_{k}$ $v=XW_{v}$	$X\in\mathbb{R}^{b\times s\times n}$	${\color[rgb]{1,0,0}XP_{q}\in\mathbb{R}^{b\times s\times(nr_{1})}}$ ${\color[rgb]{1,0,0}XQ_{q}\in\mathbb{R}^{b\times s\times(nr_{2})}}$
reshape $q,k,v$ to $(b,h,s,n/h)$	None	None
$A_{h}=\mathrm{flash\_attn}(q,k,v)$	$q,k,v\in\mathbb{R}^{b\times h\times s\times n/h}$ two buffers $\in\mathbb{R}^{b\times s\times h}$	$q,k,v\in\mathbb{R}^{b\times h\times s\times n/h}$ two buffers $\in\mathbb{R}^{b\times s\times h}$
reshape $A_{h}$ to $(b,s,n)$	None	None

Appendix

Table 6: Comparison of activation stored in the GPT series architecture, with decreased activation values highlighted in red for PRAC.

	Params	Hidden	Intermediate	Heads	Layers	Steps
LLaMA	35M	384	1024	8	6	5K
	60M	512	1376	8	8	10K
	130M	768	2048	12	12	20K
	350M	1024	2736	16	24	60K
	1B	2048	5461	24	32	150K
GPT-2	124M	768	3072	12	12	50K

Appendix

Table 7: Hyperparameter settings for pre-traing LLaMA and GPT-2 model.

	COLA	STSB	MRPC	RTE	SST2	MNLI	QNLI	QQP
Batch Size	16	16	16	16	16	16	16	16
Epochs	30	30	30	30	30	30	30	30
Learing Rate	3E-05	3E-05	3E-05	3E-05	1E-05	1E-05	1E-05	1E-05
Rank ( $r_{1}+r_{2}$ )	4
PRAC Interval	200
Max Seq Length	512

Appendix

Table 8: Hyperparameter settings for fine-tuning RoBERTa-Base model on the GLUE benchmark using the PRAC method.

Figure 6: Loss curves of pre-training LLaMA model based on the Muon and Adam-mini optimizers w./w.o. PRAC.

Method	PRAC	LLaMA-130M		LLaMA-350M
Method	PRAC	Ppl	Mem	Ppl	Mem
Muon	$\times$	23.90	21.04	17.98	45.26
Muon	$\checkmark$	24.16	15.37 ( $\bf{\downarrow 27\%}$ )	18.04	30.82 ( $\bf{\downarrow 32\%}$ )
Adam- mini	$\times$	24.22	21.01	18.60	45.42
Adam- mini	$\checkmark$	24.48	15.36 ( $\bf{\downarrow 27\%}$ )	18.92	31.09 ( $\bf{\downarrow 31\%}$ )

Appendix

Table 9: Combining with advanced optimizers on pretraining LLaMA. Report validation perplexity (PPL, lower is better) and the algorithm’s peak memory usage (lower is better).

Figure 7: Loss curve of using different

k

in LLaMA-130M pre-training, where

k_{0}=\frac{n-r_{1}}{r_{2}}

, Setting

k=k_{0}

performs better than larger or smaller settings, whether in linear or nonlinear layers.

LLM Compression

Large language model compression reduces the memory, compute, and context costs of LLMs while aiming to preserve their capabilities.

See all 273 papers

Activation Compression in LLMs: Theoretical Analysis and Efficient Algorithm
May 2, 2026Wen-Da Wei, Han-Bin Fang, Yang-Di Liu +3LLM Compression Gradient Compression
Nanjing University, Nanjing, China · Tsinghua University, Beijing, China · Huazhong University of Science and Technology, Wuhan, China +1
OASIS: Online Activation Subspace Learning for Memory-Efficient Training
Apr 10, 2026Sakshi Choudhary, Utkarsh Saxena, Kaushik RoyMemory-Efficient Training Memory-Efficient Fine-Tuning
Purdue University · Advanced Micro Devices, Inc. (AMD)
Activation-Informed Pareto-Guided Low-Rank Compression for Efficient LLM/VLM
Oct 7, 2025Ryan Solgi, Parsa Madinei, Jiayi Tian +4Model Compression LLM Compression
University of California-Santa Barbara, USA · Amazon, USA

May 2, 2026J/K move · Enter open · S save

Activation Compression in LLMs: Theoretical Analysis and Efficient Algorithm

Wen-Da Wei, Han-Bin Fang, Yang-Di Liu, Jiang-Xin Shi+2

Nanjing University, Nanjing, China · Tsinghua University, Beijing, China · Huazhong University of Science and Technology, Wuhan, China+1

LLM Compression Gradient Compression Low-Rank Compression Activation Compression LLM Training Neural Network Compression+1

Open paper AskPDF arXiv

Training large language models (LLMs) is highly memory-intensive, as training must store not only weights and optimizer states but also intermediate activations for backpropagation. While existing memory-efficient methods largely focus on gradients and optimizer states, activation compression is less well established due to the lack of LLM-tailored theory and guarantees. In this work, we develop a theoretical framework showing that activation compression is safe for linear operators when activation compression is unbiased, but problematic for nonlinear ones. We further derive gradient variance bound and establish convergence guarantees for applying activation compression to all linear operators under the standard

L

-smoothness assumption, showing that it does not change the convergence rate. Guided by the theory, we propose an activation-gradient co-compression method that reuses low-rank activation factors to compress linear-layer gradients without extra computation or additional gradient error. We conduct extensive experiments on Qwen and LLaMA models using a pretraining benchmark and multiple fine-tuning benchmarks to validate our theory and demonstrate competitive performance of our method in both accuracy and compression efficiency. We provide our code in the supplementary material for reproducibility.

Abstract

LLM Compression

Explore similar work

Activation Compression in LLMs: Theoretical Analysis and Efficient Algorithm

OASIS: Online Activation Subspace Learning for Memory-Efficient Training

Activation-Informed Pareto-Guided Low-Rank Compression for Efficient LLM/VLM