Investigating Low-Rank Training in Transformer Language Models: Efficiency and Scaling Analysis

TL;DR

This paper investigates applying low-rank parametrization to feedforward networks in Transformer-based language models, demonstrating efficiency gains and improved scaling behavior up to 1.3 billion parameters.

Contribution

It introduces large-scale low-rank FFN parametrization for Transformers trained from scratch, showing efficiency and better scaling compared to original models.

Findings

2.6× FFN speed-up with 32% parameters

Structured FFNs exhibit steeper scaling curves

Structured networks outperform medium and large models in perplexity and throughput

Abstract

State-of-the-art LLMs often rely on scale with high computational costs, which has sparked a research agenda to reduce parameter counts and costs without significantly impacting performance. Our study focuses on Transformer-based LLMs, specifically applying low-rank parametrization to the computationally intensive feedforward networks (FFNs), which are less studied than attention blocks. In contrast to previous works, (i) we explore low-rank parametrization at scale, up to 1.3B parameters; (ii) within Transformer language models rather than convolutional architectures; and (iii) starting from training from scratch. Experiments on the large RefinedWeb dataset show that low-rank parametrization is both efficient (e.g., 2.6$\times$ FFN speed-up with 32\% parameters) and effective during training. Interestingly, these structured FFNs exhibit steeper scaling curves than the original models. Motivated by this finding, we develop the wide and structured networks surpassing the current medium-sized and large-sized Transformer in perplexity and throughput performance. Our code is available at https://github.com/CLAIRE-Labo/StructuredFFN/tree/main.