EarlyBERT: Efficient BERT Training via Early-bird Lottery Tickets

TL;DR

EarlyBERT introduces a structured, ticket-based approach to accelerate training for large NLP models by identifying Early-Bird Tickets that emerge early in training. Through a three-stage process—Searching, Ticket-drawing, and Efficient-training—it prunes attention heads and FFN neurons to form a compact subnetwork, enabling 35–45% training-time savings during fine-tuning and 30–35% during pre-training with limited performance loss on GLUE and SQuAD. The method leverages network slimming with L1 regularization to jointly learn sparsity, and demonstrates that the resulting sub-networks are non-trivial and emerge early, allowing significant computational savings. This work is the first to apply structured lottery-ticket concepts to both pre-training and fine-tuning of BERT, offering practical improvements over unstructured pruning and prior NLP-LTH approaches.

Abstract

Heavily overparameterized language models such as BERT, XLNet and T5 have achieved impressive success in many NLP tasks. However, their high model complexity requires enormous computation resources and extremely long training time for both pre-training and fine-tuning. Many works have studied model compression on large NLP models, but only focusing on reducing inference time while still requiring an expensive training process. Other works use extremely large batch sizes to shorten the pre-training time, at the expense of higher computational resource demands. In this paper, inspired by the Early-Bird Lottery Tickets recently studied for computer vision tasks, we propose EarlyBERT, a general computationally-efficient training algorithm applicable to both pre-training and fine-tuning of large-scale language models. By slimming the self-attention and fully-connected sub-layers inside a transformer, we are the first to identify structured winning tickets in the early stage of BERT training. We apply those tickets towards efficient BERT training, and conduct comprehensive pre-training and fine-tuning experiments on GLUE and SQuAD downstream tasks. Our results show that EarlyBERT achieves comparable performance to standard BERT, with 35~45% less training time. Code is available at https://github.com/VITA-Group/EarlyBERT.

Figures (4)

• Figure 1: Illustration of mask difference in Hamming distance. Top: mask distance observed in pre-training. Bottom: mask distance observed in fine-tuning. The color represents the normalized mask distance between different training steps. The darker the color, the smaller the mask distance. In both cases, the mask converges quickly, which indicates the early emergence of the tickets.
• Figure 2: Effect of reducing training epochs and up-scaling learning rate for EarlyBERT in fine-tuning. The combination of 2-epoch fine-tuning and $4\times10^{-5}$ turns out to be the optimal choice.
• Figure 3: How various time points of early stopping for searching influences EarlyBERT performance.
• Figure 4: Effect of reducing training steps in pre-training on EarlyBERTBASE.