Evolving Subnetwork Training for Large Language Models

TL;DR

This paper introduces Evolving Subnetwork Training (EST), a cost-saving training paradigm for large language models that samples and gradually enlarges subnetworks, reducing FLOPs while improving downstream task performance.

Contribution

EST is a novel training method that dynamically samples and enlarges subnetworks during training, reducing computational costs and enhancing model generalization.

Findings

26.7% FLOPs reduction for GPT2

25.0% FLOPs reduction for TinyLlama

Improved downstream task performance

Abstract

Large language models have ushered in a new era of artificial intelligence research. However, their substantial training costs hinder further development and widespread adoption. In this paper, inspired by the redundancy in the parameters of large language models, we propose a novel training paradigm: Evolving Subnetwork Training (EST). EST samples subnetworks from the layers of the large language model and from commonly used modules within each layer, Multi-Head Attention (MHA) and Multi-Layer Perceptron (MLP). By gradually increasing the size of the subnetworks during the training process, EST can save the cost of training. We apply EST to train GPT2 model and TinyLlama model, resulting in 26.7\% FLOPs saving for GPT2 and 25.0\% for TinyLlama without an increase in loss on the pre-training dataset. Moreover, EST leads to performance improvements in downstream tasks, indicating that it benefits generalization. Additionally, we provide intuitive theoretical studies based on training dynamics and Dropout theory to ensure the feasibility of EST. Our code is available at https://github.com/OpenDFM/EST.