Evolving Subnetwork Training for Large Language Models

TL;DR

This work tackles the high training cost of large language models by introducing Evolving Subnetwork Training (EST), a paradigm that randomly samples subnetworks across attention heads, MLP dimensions, and layers, and gradually enlarges them through a three-stage scheduler. Empirical results on GPT-2 and TinyLlama show substantial FLOPs savings ($26.7\%$ and $25.0\%$, respectively) with no loss in pre-training performance and improved downstream generalization. The authors provide theoretical intuition grounded in dropout theory and training dynamics to explain cost savings and improved generalization, indicating that structural dropout-like effects enable faster convergence and flatter loss landscapes. Practically, EST reduces environmental impact while enabling faster iteration and better transferability, and the authors suggest extending the approach to even larger models and other architectures.

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.

Figures (7)

• Figure 1: Overview of our EST method with practical sampling scheduler. The router takes $\mathbb{I}_L$ as input and determines whether to activate the current layer. In stage 1, we obtain a subnetwork to train by sampling from $N_H, N_M$ and $N_A$ dimensions. In such subnetworks, only some layers are activated and in each activated layer, and only some attention heads and MLP neurons are used. In stage 2, all layers are activated while in each layer still only a subset of the layer is used. In stage 3, the complete model is activated.
• Figure 2: Loss curves of EST compared with the original training method.
• Figure 3: Training and evaluation loss of EST training with GPT2-base model.
• Figure 4: Training and evaluation loss of EST training with GPT2-base model.
• Figure 5: Computation in each Transformer layer during subnetwork training.

Theorems & Definitions (1)

• Definition 3.1