EMNLP 2025Oral

CoLA: Compute-Efficient Pre-Training of LLMs via Low-Rank Activation

Ziyue Liu^*1, Ruijie Zhang^*1, Zhengyang Wang^*1, Mingsong Yan¹, Zi Yang², Paul Hovland³, Bogdan Nicolae³, Franck Cappello³, Sui Tang¹, Zheng Zhang¹

¹University of California, Santa Barbara²University at Albany, SUNY³Argonne National Laboratory^*Equal contribution

Paper Code Talkfrom 14:11 Slides Poster BibTeX

In short

The activations inside a pre-trained LLM are low-rank, yet every layer computes them at full size. CoLA builds the low rank into the architecture: each linear layer becomes a small auto-encoder. Models come out 2× smaller and training takes half the compute, at full-rank perplexity.

Full-size layer

the activation is as wide as the layer

CoLA layer

a rank-r bottleneck with the nonlinearity inside

Every linear layer, with or without a nonlinearity after it, becomes an auto-encoder B σ(Ax) with A of size r × d and B of size d_ff × r, so the activation is low-rank throughout training.

2×

less training compute than full-rank, at the default rank

2×

smaller models, with full-rank perplexity from 60M to 7B

1.86×

training throughput on LLaMA-1B, and 1.64× at inference

⅔

less pre-training memory for LLaMA-7B with CoLA-M, still faster than full-rank

Abstract

The full-size MLPs and the projection layers in attention introduce tremendous model sizes of large language models (LLMs), consuming extensive computational resources in pre-training. We empirically observe that the activations of pre-trained LLMs exhibit low-rank property. Motivated by such observations, we propose CoLA and its memory-efficient implementation, CoLA-M, to replace these full-size layers with compute-efficient auto-encoders that naturally enforce low-rank activations throughout training. This fundamental architectural change eliminates the activation redundancy and significantly boosts model capacity and training efficiency. Experiments on LLaMA models with 60 million to 7 billion parameters show that CoLA reduces the computing cost by 2× and improves training throughput by 1.86× while maintaining full-rank level performance. CoLA-M further squeezes memory cost without sacrificing throughput, offering a pre-training approach with collectively superior parameter, computing, and memory efficiency. The LLMs produced are also 2× smaller, enabling faster inference with lower memory cost on resource-constrained platforms.

Low-rank activation by construction

In a pre-trained GPT-2, the singular values of the MLP activations fall off quickly: in every block, the rank that keeps 95% of the spectral energy is a small fraction of the full dimension. Full-size layers spend most of their compute producing that redundancy.

Effective rank of the MLP activation in GPT-2 small

singular values needed to keep 95% of the spectral energy, per block

effective rankfull dimension, 3,072

01,0002,0003,072

decoder block

Show the numbers

Block	1	2	3	4	5	6	7	8	9	10	11	12
Effective rank	≈200	≈240	≈250	≈320	≈370	≈400	≈400	≈400	≈390	≈370	≈310	≈220

In every block, about 7% to 13% of the 3,072 activation dimensions hold 95% of the energy. Activations measured on WikiText-2.

CoLA replaces each linear layer with h = B σ(Ax), where r is well below both input and output sizes. The nonlinearity sits between the two factors, so the activation itself stays low-rank. That differs from factorizing W into BA, which compresses parameters but keeps a full-size activation. All MLP layers and all attention projections become CoLA layers; attention itself, residual connections and the gated product stay as they are.

With the default rank r = d/4, a decoder layer needs about half the compute of full-rank training, and CoLA saves compute for any r below about 0.62d. The other efficient pre-training methods each cover only part of the picture:

	CoLA(-M)	SLTrain	GaLore	ReLoRA
Fewer parameters	✓	✓	–	–
Less compute, training	✓	–	–	✓
Less compute, inference	✓	–	–	–
Less memory, training	✓	✓	✓	✓
Less memory, inference	✓	✓	–	–
Higher throughput, training	✓	–	–	–
Higher throughput, inference	✓	–	–	–

Which kinds of efficiency each method provides.

What the theory says

Never worse than plain low rank. A CoLA layer approximates a layer at least as well as a standard low-rank layer of the same rank. Proposition 3.1
Strictly better under data-dependent conditions. The nonlinearity in the middle lets it beat the standard low-rank layer. Theorem 3.3
Benefits from low effective rank. Its error bound depends on the output's effective rank, which is small in practice. Theorem 3.4

CoLA-M: save only the small activations

At large batch sizes, activations can take most of the memory. Plain gradient checkpointing keeps each block's input and recomputes the whole block in the backward pass, almost a full forward pass. CoLA-M keeps the low-rank activations instead, and recomputes only the up-projections and self-attention, about half of a CoLA forward pass.

At a similar memory saving on LLaMA-1B, that is 4.6× less recomputation than plain checkpointing. CoLA-M needs a quarter of full-rank training's memory at 1B and a third at 7B, and it still trains faster.

One CoLA-M decoder block

LLaMA-style, read from top to bottom

low-rank activation, savedrecomputed in the backward passfull-size SiLU

CoLA-M against full-rank training

one H100 94GB; Vanilla GCP is full-rank training with plain gradient checkpointing

MemoryGB, lower is better

Full-rank

69.84

Vanilla GCP

14.89

CoLA

66.46

CoLA-M

17.33

0306090

ComputeFLOPs relative to full-rank, lower is better

Full-rank

1.00×

Vanilla GCP

1.68×

CoLA

0.40×

CoLA-M

0.55×

00.511.52

Throughputtokens per second, higher is better

Full-rank

12,365

Vanilla GCP

8,799

CoLA

22,979

CoLA-M

16,617

010k20k

Results

LLaMA models from 60M to 7B parameters pre-trained on C4 with no data repetition, following the exact setup of GaLore and SLTrain, plus BERT-Large and an over-trained LLaMA-350M.

Smallest model, full-rank perplexity

At every size CoLA has the fewest parameters and the least memory, and its perplexity is on par with full-rank training: 15.52 against 15.56 at 1B.

LLaMA-1B: perplexity, compute and model size

bubble area is the number of parameters

CoLA, 609Mothers, 646M to 1,339M

Show the numbers

Method	Perplexity	Compute	Parameters
CoLA	15.52	0.41×	609M
Full-rank	15.56	1.00×	1,339M
ReLoRA	18.33	1.06×	1,339M
GaLore	15.64	2.31×	1,339M
SLTrain	16.14	2.96×	646M

Compute is the paper's per-step estimate for one decoder layer with a 256-token batch. SLTrain and GaLore spend extra compute on rebuilding weights and projecting gradients, so they sit above full-rank training.

Compute-optimal pre-training on C4

same setup as GaLore and SLTrain

Method	Perplexity	Parameters (M)	Memory (GB)
Full-rank	34.06	58	0.43
ReLoRA	37.04	58	0.37
GaLore	34.88	58	0.36
SLTrain	34.15	44	0.32
CoLA	34.04	43	0.32

Rank / width 128 / 512, 1.1B training tokens. Memory counts weights, gradients and optimizer states in BF16.

Method	Perplexity	Parameters (M)	Memory (GB)
Full-rank	24.36	134	1.00
ReLoRA	29.37	134	0.86
GaLore	25.36	134	0.79
SLTrain	26.04	97	0.72
CoLA	24.48	94	0.70

Rank / width 256 / 768, 2.2B training tokens. Memory counts weights, gradients and optimizer states in BF16.

Method	Perplexity	Parameters (M)	Memory (GB)
Full-rank	18.80	368	2.74
ReLoRA	29.08	368	1.94
GaLore	18.95	368	1.90
SLTrain	19.42	194	1.45
CoLA	19.40	185	1.38

Rank / width 256 / 1024, 6.4B training tokens. Memory counts weights, gradients and optimizer states in BF16.

Method	Perplexity	Parameters (M)	Memory (GB)
Full-rank	15.56	1,339	9.98
ReLoRA	18.33	1,339	6.79
GaLore	15.64	1,339	6.60
SLTrain	16.14	646	4.81
CoLA	15.52	609	4.54

Rank / width 512 / 2048, 13.1B training tokens. Memory counts weights, gradients and optimizer states in BF16.

LLaMA-7B with a third of the memory

At 150k steps, CoLA-M reaches a perplexity of 12.73 with 26.8 GB of memory. 8-bit Adam and 8-bit GaLore stay around 14.6 while using 65 to 73 GB.

LLaMA-7B on C4

memory for each method in the legend

CoLA-M, 26.8 GB8-bit GaLore, 65.2 GB8-bit Adam, 72.6 GB

Show the numbers

Method	Memory (GB)	10k	40k	80k	120k	150k
8-bit Adam	72.59	n/a	18.09	15.47	14.83	14.61
8-bit GaLore	65.16	26.87	17.94	15.39	14.95	14.65
SLTrain	60.91	27.59	n/a	n/a	n/a	n/a
CoLA-M	26.82	22.76	16.21	13.82	13.09	12.73

8-bit Adam and GaLore numbers are from the GaLore paper, SLTrain's from its paper. The 7B runs are below the compute-optimal budget.

A larger rank beats full-rank training

Raising the rank to 0.7× of the full-rank compute makes CoLA better than full-rank training at all three sizes. Shrinking a full-rank model to CoLA's default budget instead costs a lot of perplexity.

	60Mperplexity	compute	130Mperplexity	compute	350Mperplexity	compute
Full-rank	34.06	1×	24.36	1×	18.80	1×
Smaller full-rank	37.73	0.4×	27.05	0.5×	20.53	0.4×
CoLA, default rank	34.04	0.4×	24.48	0.5×	19.40	0.4×
CoLA, larger rank	31.52	0.7×	23.97	0.7×	18.32	0.7×

Validation perplexity and compute relative to full-rank training.

Beyond compute-optimal

Trained 8× past the compute-optimal budget (51B tokens), LLaMA-350M with CoLA reaches 13.96 against 14.47 for full rank. A CoLA version of BERT-Large (350M) pre-trained on 85B tokens ends at a lower loss (1.257 against 1.263), and after fine-tuning it averages 83.5 on GLUE against 82.7, ahead or level on 7 of 8 tasks.

	QQP	SST-2	MRPC	CoLA	QNLI	MNLI	RTE	STS-B	Avg
BERT-Large	91.1	92.1	90.7	53.1	91.6	84.3	69.9	88.9	82.7
CoLA	91.2	92.3	90.6	54.1	91.7	84.3	74.2	89.7	83.5

GLUE after pre-training both models with NVIDIA's BERT reproduction; mean of the best 5 of 10 seeds.

Faster to train and to serve

CoLA and CoLA-M are the only methods here that train faster than the full-rank model. At inference, CoLA is up to 1.64× faster and uses up to 1.67× less memory.

Pre-training throughput, LLaMA-1B on one A100 40GB

tokens per second at batch 16

CoLA

full-rank

11,506

CoLA-M

8,569

Full-rank

7,161

ReLoRA

6,942

APOLLO

6,838

SLTrain

6,770

Vanilla GCP

5,278

GaLore

4,556

03,0006,0009,00012,000

Limitations and follow-up work

Distributed training. Scaling CoLA further needs proper model parallelism, and porting it to frameworks such as Megatron-LM and DeepSpeed is not trivial.
Industrial scale. The industrial-setting check stops at BERT-size models. Decoder-only models with billions of parameters trained on trillions of tokens are still to be tested, and we are open to collaborating on that.
Mixture-of-experts. The paper covers dense models only. Extending CoLA to MoE architectures is left for future work.
Follow-up work. LaX (NeurIPS 2025), LOST and CR-Net build on CoLA and improve it further.

Citation

BibTeX

@inproceedings{liu-etal-2025-cola,
    title = "{C}o{LA}: Compute-Efficient Pre-Training of {LLM}s via Low-Rank Activation",
    author = "Liu, Ziyue  and
      Zhang, Ruijie  and
      Wang, Zhengyang  and
      Yan, Mingsong  and
      Yang, Zi  and
      Hovland, Paul D.  and
      Nicolae, Bogdan  and
      Cappello, Franck  and
      Tang, Sui  and
      Zhang, Zheng",
    editor = "Christodoulopoulos, Christos  and
      Chakraborty, Tanmoy  and
      Rose, Carolyn  and
      Peng, Violet",
    booktitle = "Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing",
    month = nov,
    year = "2025",
    address = "Suzhou, China",
    publisher = "Association for Computational Linguistics",
    url = "https://aclanthology.org/2025.emnlp-main.230/",
    doi = "10.18653/v1/2025.emnlp-main.230",
    pages = "4627--4645",
    ISBN = "979-8-89176-332-6"
}