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

TL;DR

This paper introduces CoLA, a low-rank activation-based method for pre-training large language models more efficiently, reducing computational costs and memory usage while maintaining performance.

Contribution

It proposes a novel architecture replacing full-size layers with auto-encoders enforcing low-rank activations, significantly improving training efficiency and model size.

Findings

CoLA halves the computing cost of LLM pre-training.

CoLA improves training throughput by 86%.

Produced models are twice as small, enabling faster inference.

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