bert2BERT: Towards Reusable Pretrained Language Models

TL;DR

The paper tackles the high cost of pre training large Transformer based PLMs by reusing knowledge from smaller pre trained models. It introduces bert2BERT, a method that expands width and depth of a small model to initialize a larger one and then trains it with a two stage pre training regime. Key contributions include function preserving initialization, advanced knowledge initialization, and depth wise expansion, all demonstrated to substantially reduce pre training compute for BERT and GPT without sacrificing performance. The approach shows strong potential as an energy efficient pathway for scaling up PLMs and is shown to generalize across model families. The work also provides a procedural framework and ablations to guide future efficient pre training efforts.

Abstract

In recent years, researchers tend to pre-train ever-larger language models to explore the upper limit of deep models. However, large language model pre-training costs intensive computational resources and most of the models are trained from scratch without reusing the existing pre-trained models, which is wasteful. In this paper, we propose bert2BERT, which can effectively transfer the knowledge of an existing smaller pre-trained model (e.g., BERT_BASE) to a large model (e.g., BERT_LARGE) through parameter initialization and significantly improve the pre-training efficiency of the large model. Specifically, we extend the previous function-preserving on Transformer-based language model, and further improve it by proposing advanced knowledge for large model's initialization. In addition, a two-stage pre-training method is proposed to further accelerate the training process. We did extensive experiments on representative PLMs (e.g., BERT and GPT) and demonstrate that (1) our method can save a significant amount of training cost compared with baselines including learning from scratch, StackBERT and MSLT; (2) our method is generic and applicable to different types of pre-trained models. In particular, bert2BERT saves about 45% and 47% computational cost of pre-training BERT_BASE and GPT_BASE by reusing the models of almost their half sizes. The source code will be publicly available upon publication.

Figures (13)

• Figure 1: Loss curves of bert2BERT and baselines. Both stackBERT stack and MSLT stack2 are based on the progressive training setting. More details are shown in Table \ref{['tab:glue_results']}.
• Figure 2: The comparisons of attention patterns between small and large PLMs. The upper ones are the attention patterns of BERT$_{\rm BASE}$ model of a architecture of {$L$=$12$, $D$=$768$}, and the lower ones are the attention patterns of one small BERT model whose architecture is {$L$=$12$, $D$=$512$}. We find that there are a large number of similar attention patterns in the same layer of the two models, indicating the possibility of reusing parameters of trained small PLMs to speed up the pre-training of large PLMs.
• Figure 3: Overview of the matrix expansion. It first enlarges the size of in-dimension and then enlarges the size of out-dimension. From the view of neurons, it adds new neurons (red) in both input and output.
• Figure 4: Overview of the function preserving initialization (FPI). Given the same input {$x_1$, $x_2$}, FPI ensures the initialized target model has the same output of {$y_1$, $y_2$} with the source model. The first and the second steps are expanding the in-dimension and out-dimension of the parameter matrix according to both $g_{in}$ and $g_{out}$ mapping functions respectively. From the view of neurons, FPI copies the input and output neurons to expand the network. After the matrix expansion, we can further use the in-dimension expansion to ensure the output same as the original one.
• Figure 5: Overview of the advanced knowledge initialization (AKI). It first performs the in-dimension expansion on both the matrixes of current and upper layers. Then it uses the widened matrix of the current layer as the top part of the new matrix and samples the row of the widened matrix of the upper layer as the bottom part of the new matrix.