CoLA: A more efficient way to train large language models
MCS Menu
CoLA might sound like a soft drink, but in computing it’s a new approach for pre-training large language models (LLMs). Developed by researchers at the U.S. Department of Energy’s Argonne National Laboratory, the University of California, Santa Barbara and the University at Albany, SUNY, Compute-Efficient Pre-Training of LLMs via Low-Rank Activation (CoLA) aims to make model training faster, cheaper and more efficient.
Breaking free from an unsustainable trend
LLMs keep becoming bigger, needing more GPUs to provide the necessary power for pre-training over longer periods of time. Researchers have tried to address this trend by simplifying models using “low-rank” techniques, but many of these approaches increase computing costs, use too much memory or reduce model quality.
CoLA takes a different path.
Instead of relying on very large and complex layers commonly used in LLMs, CoLA uses bottleneck-structured auto-encoders, comprising low-weight matrices and a nonlinear activation in the middle.
Normally, the word bottleneck sounds like a bad thing — a place where everything slows down. In CoLA, however, the bottleneck structure is intentional and helpful.
“The key insight is that many of the model activations in pre-trained LLMs are redundant,” said Ziyue Liu, a PhD candidate at UC Santa Barbara and a visiting student at Argonne’s Mathematics and Computer Science (MCS) division. “By selecting the rank to be significantly smaller than the input/output dimensions, we create a narrow ‘bottleneck’ that forces the data into a smaller space. Working in this space, CoLA avoids carrying out inefficient, redundant operations and focuses only on what matters most.”
When tested with popular LLaMA models, CoLA used fewer parameters and cut computations by about half — without hurting performance during training.
Providing a new variant
Just as cola soft drinks come in different flavors, CoLA provides a variant, CoLA-M, specifically designed to save memory.
“Again, we take advantage of the bottleneck structure,” said Ruijie Zhang, a predoctoral appointee in Argonne’s MCS division and a co-author of the CoLA study. “Our CoLA-M implementation stores only the low-rank activations and recomputes the rest when needed, instead of keeping everything in memory. This approach significantly reduces memory demands without slowing down training.”
Compared with other memory-saving techniques — which often add extra computation or reduce accuracy — CoLA-M performs exceptionally well. It cuts overall memory use by about two-thirds, surpassing even the most effective implementation (i.e., “vanilla” gradient checkpointing, GCP, which saves only one intermediate result from each transformer layer) while being 5x cheaper in precomputation and 3x cheaper in total computation.
Faster training, too
Efficiency isn’t just about computation and memory — it’s also about speed. Many existing low-rank methods slow training because they add extra steps. In contrast, both CoLA and CoLA-M deliver higher training throughput, meaning models train faster while using fewer resources.
Across multiple tests, both CoLA and CoLA-M outperformed current state-of-the-art approaches, including the “vanilla” GCP and the full-rank nonredundant version.
Why it matters
“By combining fewer parameters, faster computing and lower memory usage, CoLA and its variant CoLA-M offer a practical and scalable way to pre-train large language models,” Liu said. “That translates into saving GPU resources — a big win for future LLM use.”
Want to learn more? See the paper by Ziyue Liu, Ruijie Zhang, Zhengyang Wang, Mingsong Yan, Zi Yang, Paul D Hovland, Bogdan Nicolae, Franck Cappello, Sui Tang, and Zheng Zhang, “CoLA: Compute-efficient pre-training of LLMs via low-rank activation,” in the Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing, pp. 4627-4645, 2025. https://arxiv.org/pdf/2502.10940
Argonne National Laboratory seeks solutions to pressing national problems in science and technology by conducting leading-edge basic and applied research in virtually every scientific discipline. Argonne is managed by UChicago Argonne, LLC for the U.S. Department of Energy’s Office of Science.
The U.S. Department of Energy’s Office of Science is the single largest supporter of basic research in the physical sciences in the United States and is working to address some of the most pressing challenges of our time. For more information, visit https://energy.gov/science.