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

2025/02/16 by Ziyue Liu, Liu, Ziyue, Ruijie K. Zhang +15 · 12 citations
Mathematics · Medicine · Neuroscience · Physics and Astronomy · #Advanced Radiotherapy Techniques #Artificial Intelligence (cs.AI) #Brain Tumor Detection and Classification #Cola (plant) #Combinatorics #Computer science #Econometrics #FOS: Computer and information sciences #Geography #Internal medicine #Machine Learning (cs.LG) #Mathematics #Medical Imaging Techniques and Applications #Medicine #Rank (graph theory) #Training (meteorology)

paper · pdf · doi:10.48550/arxiv.2502.10940

published in arXiv (Cornell University) (Cornell University)

openalex publication_date 2025/02/16 · openalex created_date 2025/10/10 · openalex updated_date 2026/08/05

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 \bf 2\pmb× and improves training throughput by \bf 1.86\pmb× 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 \bf 2\pmb× smaller, enabling faster inference with lower memory cost on resource-constrained platforms.

Cited by

• BOOST: BOttleneck-Optimized Scalable Training Framework for Low-Rank Large Language Models
• Memory-Efficient Fine-Tuning via Low-Rank Activation Compression
• CR-Net: Scaling Parameter-Efficient Training with Cross-Layer Low-Rank Structure
• LOST: Low-rank and Sparse Pre-training for Large Language Models
• Greedy Low-Rank Gradient Compression for Distributed Learning with Convergence Guarantees
• Reparameterized LLM Training via Orthogonal Equivalence Transformation
• DiaBlo: Diagonal Blocks Are Sufficient For Finetuning
• Scalable Parameter and Memory Efficient Pretraining for LLM: Recent Algorithmic Advances and Benchmarking
• LaX: Boosting Low-Rank Training of Foundation Models via Latent Crossing
• RefLoRA: Refactored Low-Rank Adaptation for Efficient Fine-Tuning of Large Models
• AuroRA: Breaking Low-Rank Bottleneck of LoRA with Nonlinear Mapping
• TEON: Tensorized Orthonormalization Beyond Layer-Wise Muon for Large Language Model Pre-Training

Related

• Scalable Parameter and Memory Efficient Pretraining for LLM: Recent Algorithmic Advances and Benchmarking
• BlockLLM: Memory-Efficient Adaptation of LLMs by Selecting and Optimizing the Right Coordinate Blocks
• Scalable Efficient Training of Large Language Models with Low-dimensional Projected Attention
• CR-Net: Scaling Parameter-Efficient Training with Cross-Layer Low-Rank Structure
• SLTrain: a sparse plus low-rank approach for parameter and memory efficient pretraining