Leveraging Neighbor Attention Initialization (NAI) for Efficient Training of Pretrained LLMs

Abstract

In the realm of pretrained language models (PLMs), the exponential increase in computational resources and time required for training as model sizes expand presents a significant challenge. This paper proposes an innovative approach named neighbor attention initialization (NAI) to expedite the training process of larger PLMs by leveraging smaller PLMs through parameter initialization. Our methodology hinges on the hypothesis that smaller PLMs, having already learned fundamental language structures and patterns, can provide a robust foundational knowledge base for larger models, which is called function preserving. Specifically, we present a comprehensive framework detailing the process of transferring learned features on transformer-based language models mainly using the neighbor attention head and neighbor layer. We conducted experiments in GPT-2 and demonstrated that our method yields considerable savings in training costs compared to standard approaches, including learning from scratch and bert2BERT, indicating a notable improvement in training efficiency for large PLMs.

1. Introduction

Recent advancements in Natural Language Processing (NLP) have been significantly driven by PLMs such as BERT [1], GPT series [2,3,4], XLNet [5], T5 [6], RoBERTa [7] and LLaMa series [8,9]. These models have demonstrated exceptional performance across a range of NLP tasks. However, the pre-training phase of these large-scale PLMs entails substantial computational resources. The field has observed a growing inclination towards constructing and training exceptionally large models. This trend aims to explore the upper bounds of PLMs’ capabilities. Illustrative examples include GPT-2 (117M, 345M, 762M, 1542M parameter); GPT-3 (175B parameters); PanGu-$\alpha$ [10] with 200B parameters; Switch Transformers [11] comprising an unprecedented 1571B parameters; LLaMa series with 7B, 13B, 33B, 70B parameters [8,9]; and Qwen (1.8B, 7B, 14B, 72B parameters) [12]. These models have shown promising results in language comprehension and generation. Notwithstanding their performance, a notable limitation is that these colossal models are generally trained from scratch, not leveraging the already acquired knowledge from their smaller counterparts. However, our empirical studies suggest the presence of overlapping knowledge domains across PLMs of varying sizes, indicating potential training efficiencies to be gained from small PLMs. For example, in Figure 1, the attention patterns of the two GPT-2s with different sizes are quite similar in the same layer and same attention head. And, similar phenomena have been found in previous studies [13]. Furthermore, in Figure 2, we compare the attention patterns of identical heads across different layers in GPT-2 (117M) and GPT-2 (345M). Notably, the 12th layer of GPT-2 (117M), which is its final layer, closely resembles the 13th layer of GPT-2 (345M). This observation underpins the potential for depth expansion in smaller PLMs.

Drawing from these observations, our research presents neighbor attention initialization (NAI), a knowledge transfer approach from smaller to larger PLMs, aimed at boosting PLMs training efficiency and cost-effectiveness. NAI mainly extends function preserving training [13,14] to Transformer-based PLMs by reuse the parameters of the existing smaller PLMs, and in the process, NAI strictly retains the ability of small models, thereby providing the larger model with an optimized starting point for subsequent training phases. Additionally, NAI employs the closest attention head for width expansion and the nearest layer for depth expansion, for which we observe that there is a large number of similar attention distributions across both neighbor attention heads and neighbor layers.

To demonstrate the superiority of our method, we conducted experiments on GPT-2. We employed the GPT-2 model for the following reasons: (1) The GPT-2 model architecture is publicly available in four different size variants, facilitating our research on knowledge transfer across models of varying scales within the same family. (2) While previous work, such as bert2BERT [13], has primarily focused on encoder-based models like BERT [1] as a research foundation, there has been relatively limited exploration of knowledge transfer techniques for decoder-only models. As a seminal decoder-only model, GPT-2 provides an ideal testbed for investigating efficient pre-training methods tailored to this important class of language models. (3) The model architecture and parameters of the GPT-3 model, a larger successor to GPT-2, are not publicly available. Furthermore, the sheer scale of GPT-3’s parameter exceeds the computational resources at our disposal for pre-training models of such magnitude.

Our results show that our method can save a significant amount of computation in pre-training compared to the traditional way of learning from scratch and outperform the existing initialization methods bert2BERT [13]. One typical example is that, using GPT-2 (117M) as a base for GPT-2 (345M) initialization, NAI reduces computation costs by 31% compared to conventional pre-training from scratch.

In general, PLMs within the same series often have varying model sizes, as larger models tend to demonstrate enhanced overall capabilities while incurring greater computational resource requirements. However, for relatively simple tasks such as speech generation, smaller models may suffice. Existing common methods for training PLMs of varying scales often involve training each model from scratch, without leveraging the knowledge encapsulated in previously trained smaller PLMs. Therefore, the primary research question we aim to address is as follows: For PLMs within the same series but with different model sizes, how can we leverage the existing smaller PLMs to train larger PLMs in a computationally efficient manner?

2. Related Work

Efficient pre-training in NLP has been a topic of exploration in previous research. Some approaches aim to enhance efficiency by reusing parameters from existing PLMs and employing progressive learning techniques. These works, such as [15,16,17], are motivated by the observation that different layers of PLMs possess similar knowledge, such as attention patterns. The approach involves initially pre-training a smaller model with fewer Transformer layers and subsequently expanding the model iteratively by adding the already-trained layers on top. Another line of research proposes “knowledge inheritance” or “back distillation”, wherein the knowledge of small models is transferred into larger models [18]. Another aspect of efficient pre-training focuses on data efficiency [19], where rare words are given special attention during the pre-training process to enhance the model’s understanding when encountering them subsequently.

Efficiency in pre-training can also be improved through the use of specific pre-training tasks. One notable example is ELECTRA [20], which introduces a task called replaced token detection. This task involves predicting whether each token in a corrupted input has been replaced or not. Our proposed method is orthogonal to this line of work. Additionally, there are several other techniques that can be employed for efficient pre-training, including mixed precision training [21], large batch optimization [16], model architecture innovation [22], and layer dropping training techniques [23].

Another relevant topic is reusable neural networks, which is closely related to transfer learning [24] and has primarily been explored in the field of computer vision. Net2Net [14] is a seminal work in this area, which introduces the concept of function-preserving transformations to enable the reuse of neural networks. However, Net2Net randomly selects neurons to be split, and subsequent work [25] addresses this issue by employing a functional steepest-descent approach to determine the optimal subset of neurons for splitting. The pruning technique [26] has also been employed for reusable neural networks [27]. In addition to [27], another notable study introduces the concept of hierarchical pre-training. This approach effectively reduces both the time required for pre-training and enhances overall performance by leveraging an already pre-trained vision model as an initialization step in the pre-training process. Recently, ref. [13] extended the previous function-preserving [14] Transformer-based language model and further improved it by proposing advanced knowledge for a large model’s initialization, which uses the parameters of the next transformer layer to expand each layer. In this paper, we focus on the reusable pre-training of language models and present a new method called NAI, which accelerates the pre-training of GPT-2.

3. Method

3.1. Preliminary

Before discussing our methodology, we briefly outline the GPT-2 architecture, which comprises an embedding layer followed by several Transformer decoders [3,28]. The architecture is shown in Figure 3.

3.1.1. Embedding Layer

The embedding layer consists of token embedding and positional embedding, which maps the tokens in a sentence into vectors with two embedding matrices ${\mathit{W}}^{TE}$ and ${\mathit{W}}^{PE}$. Then, we obtain the initial hidden states ${\mathit{H}}_0$.

3.1.2. Transformer Layer

The hidden states are iteratively processed by multiple Transformer layers:

$${\mathit{H}}_l={\mathrm{Transformer}}_l\left({\mathit{H}}_{l-1}\right),l\in [1,L]$$

(1)

where L denotes the number of Transformer layers, each including a multi-head casual attention (MHCA) and a feed-forward network (FFN), and both of them have residual connections [29].

Layer Normalization. Both MHCA and FFN modules utilize a layer normalization [30] to stabilize hidden state dynamics. Formally, it is written as:

$$\mathrm{LayerNorm}\left(\mathit{H}\right)=\left(\frac{\mathit{H}-\mu }{\sqrt{{\sigma }^2+ϵ}}\right)\odot {\mathit{W}}^{LN}+{\mathit{b}}^{LN},$$

(2)

where ⊙ means the element-wise multiplication, the statistics of $\mu$ and $\sigma$ are the calculated mean and variance of hidden states $\mathit{H}$ respectively, $ϵ$ is a small constant added for numerical stability, and ${\mathit{W}}^{LN}$ and ${\mathit{b}}^{LN}$ are the scale and shift parameters.

MHCA. The Transformer decoder’s key component, MHCA, processes the layer’s hidden states ${\mathit{H}}_l$ through multiple heads, concatenating their outputs to form the final output as follows:

$$\begin{array}{cc}\hfill & {\mathit{Q}}_i,{\mathit{K}}_i,{\mathit{V}}_i={\mathit{H}}_l{\mathit{W}}_{l,i}^Q,{\mathit{H}}_l{\mathit{W}}_{l,i}^K,{\mathit{H}}_l{\mathit{W}}_{l,i}^V\hfill \\ \hfill & {\mathit{h}}_{l,i}=\mathrm{softmax}(\frac{{\mathit{Q}}_i{\mathit{K}}_i^T}{\sqrt{d_k}}\odot {\mathit{W}}^{MASK}){\mathit{V}}_i\hfill \\ \hfill & {\mathit{H}}_l=\mathrm{Concat}({\mathit{h}}_{l,1},{\mathit{h}}_{l,2},\dots ,{\mathit{h}}_{l,a})\hfill \\ \hfill & {\mathit{H}}_l^{MHCA}={\mathit{H}}_l{\mathit{W}}_l^O+{\mathit{H}}_{l-1}.\hfill \end{array}$$

(3)

After the Layer Normalization, the output ${\mathit{H}}_l$ is linearly projected to queries (${\mathit{Q}}_i$), keys (${\mathit{K}}_i$) and values (${\mathit{V}}_i$) using different weights ${\mathit{W}}_{l,i}^Q,{\mathit{W}}_{l,i}^K,{\mathit{W}}_{l,i}^V$, respectively. ${\mathit{h}}_l^i$ indicates the context-aware vector, which is obtained by the scaled dot-product of queries and keys in the i-th casual attention head. a represents the number of self-attention heads. $d_k$ is the head dimension acting as the scaling factor. ${\mathit{W}}^{MASK}$ is the attention mask to ensure that the model, during self-attention calculations, only considers its previous words and not the subsequent words. And, ${\mathit{H}}_l^{\mathrm{MHCA}}$ is the input of the next sub-module FFN.

FFN. FFN consists of two linear layers and one NewGeLU activation function [3], which is formatted as:

$$\begin{array}{cc}\hfill & {\mathit{H}}_l=\mathrm{LayerNorm}\left({\mathit{H}}_l^{MHCA}\right)\hfill \\ \hfill & {\mathit{H}}_l^{FNN}=\mathrm{NewGeLU}({\mathit{H}}_l{\mathit{W}}_l^1+{\mathit{b}}_l^1){\mathit{W}}_l^2\hfill \\ \hfill & {\mathit{H}}_l={\mathit{H}}_l^{FNN}+{\mathit{H}}_l^{MHCA}\hfill \end{array}$$

(4)

3.2. Overview

We aim to accelerate the pre-training of target model T($L^t,D^t$) by transferring the knowledge of an existing pre-trained model S($L^s,D^s$), where $L^s,L^t$ means the number of Transformer layers, and $D^s,D^t$ means the model width (i.e., hidden size), satisfying $L^s\le L^t$ and $D^s\le D^t$. Similar to [13], NAI initializes the target model T using the parameters from the existing model S by the width-wise expansion ($D^s\to D^t$) and depth-wise expansion ($L^s\to L^t$). This expansion enables direct knowledge transfer from the source model’s parameters to the target model. Essentially, the width-wise expansion can be decomposed into two kinds of operations: in-dimension and out-dimension expansion. As illustrated in Figure 4, the matrix expansion enlarges a parameter matrix $\mathit{W}\in {\mathbb{R}}^{d_{in}^w\ast d_{out}^w}$ of source model S to $\mathit{U}\in {\mathbb{R}}^{d_{in}^u\ast d_{out}^u}$ of target model T. In the following sections, we introduce our neighbor attention initialization (NAI), which employs the closest attention head for width expansion and the nearest layer for depth expansion, simultaneously retaining the capabilities of the original model.

3.3. NAI

We observe that there is a large number of similar attention distributions across both neighbor attention heads and neighbor layers in GPT-2 [2,3,4], which is also mentioned of Bert [1] in previous works [31,32]. For example, in Figure 5, we can find that there are lots of similar distributions across the neighbor attention heads in the same layer.

3.3.1. Width-Wise Expansion

Width-wise expansion includes two kinds of operations: in-dimension and out-dimension expansion. Based on the results of Figure 5, NAI first conducts the in-dimension expansion by copying the weights of the last dimension and then performs the weight normalization to ensure that the output is the same as the original output. Second, NAI conducts the out-dimension expansion, which just does the copy operation. Note that because of the multi-head attention mechanism, all the operations are conducted in the head-wise. Lastly, to break the symmetry more rapidly [14], a small amount of noise is added to the expanded dimension. Figure 6 and Figure 7 illustrate how NAI conducts width-wise expansion. We hereby denote the weight expansion as

$$\mathit{U}={\mathrm{EXPN}}_{\mathrm{out}}\left({\mathrm{EXPN}}_{\mathrm{in}}\left(\mathit{W}\right)\right),$$

(5)

which includes two stages of expansion, in-dimension expansion ${\mathrm{EXPN}}_{\mathrm{in}}$ and out-dimension expansion ${\mathrm{EXPN}}_{\mathrm{out}}$.

Expansion for all modules. We apply width-wise expansion of NAI for all modules of GPT-2 via matrix expansion ${\mathrm{EXPN}}_{\mathrm{in}}$ and ${\mathrm{EXPN}}_{\mathrm{out}}$. Specifically, for the embedding matrix ${\mathit{U}}^{\mathit{TE}}$ and ${\mathit{U}}^{\mathit{PE}}$, we only conduct the out-dimension expansion:

$$\begin{array}{cc}\hfill & {\mathit{U}}^{\mathit{TE}}={\mathrm{EXPN}}_{\mathrm{out}}\left({\mathit{W}}^{\mathit{TE}}\right),\hfill \\ \hfill & {\mathit{U}}^{\mathit{PE}}={\mathrm{EXPN}}_{\mathrm{out}}\left({\mathit{W}}^{\mathit{PE}}\right).\hfill \end{array}$$

(6)

Layer normalization is a frequently employed component that is typically applied on top of MHCA and FFN modules. This normalization technique ensures that the input data undergoes normalization. Its expansion is formulated as follows:

$$\begin{array}{c}\hfill {\mathit{U}}^{\mathit{LN}}={\mathrm{EXPN}}_{\mathrm{out}}\left({\mathit{W}}^{\mathit{LN}}\right)\end{array}$$

(7)

The MHCA module can be decomposed into multiple parallel self-attention heads. It is worth noting that in Equation (3), when we calculate attention score, it will be scaled according to dimension $\sqrt{d_k}$, so after completing the expansion of $Q,K$, we need to multiply the overall attention by the coefficient $\frac{D^t}{D^s}$. At the same time, in order to scale the impact of coefficient $\frac{D^t}{D^s}$ on the original matrix, we allocate it to both Q and K. A more detailed explanation of this operation will be discussed later. The expansion can be formulated as:

$$\begin{array}{cc}\hfill {\mathit{U}}^{Q|K}& =\sqrt{\sqrt{\frac{D^t}{D^s}}}{\mathrm{EXPN}}_{\mathrm{out}}\left({\mathrm{EXPN}}_{\mathrm{in}}\left({\mathit{W}}^{Q|K}\right)\right),\hfill \\ \hfill {\mathit{U}}^V& ={\mathrm{EXPN}}_{\mathrm{out}}\left({\mathrm{EXPN}}_{\mathrm{in}}\left({\mathit{W}}^V\right)\right),\hfill \\ \hfill {\mathit{U}}^O& ={\mathrm{EXPN}}_{\mathrm{out}}\left({\mathrm{EXPN}}_{\mathrm{in}}\left({\mathit{W}}^O\right)\right).\hfill \end{array}$$

(8)

For the FFN module, we perform the expansion on the parameter matrices ${\mathit{W}}^{1|2}$ (Equation (4)) as follows:

$$\begin{array}{c}\hfill {\mathit{U}}^{1|2}={\mathrm{EXPN}}_{\mathrm{out}}\left({\mathrm{EXPN}}_{\mathrm{in}}\left({\mathit{W}}^{1|2}\right)\right).\end{array}$$

(9)

Additionally, we introduce an extension to the classifier layer, typically comprising a linear layer, as our observations have indicated that this modification can expedite the convergence of the loss function:

$$\begin{array}{c}\hfill {\mathit{U}}^{\mathit{CLS}}={\mathrm{EXPN}}_{\mathrm{in}}({\mathit{W}}^{\mathit{CLS}}).\end{array}$$

(10)

We provide a simplified example to illustrate the knowledge transfer process of NAI. In this example, we consider a GPT-2 structured source model S {12, 128}, which is extended to a target model T {12, 192}. The dimension of the attention head is 64, and the hidden layer state is represented head-wise. That means, for example, we represent the 128-dimensional vector $\mathit{X}$ as [${\mathit{X}}_0,{\mathit{X}}_1$]. We denote the hidden state after passing through each module of the source model S as the left-hand side of the arrow, and the hidden state after passing through T as the right-hand side. Assume that the input data batch is 1; after passing through the embedding layer, we have the initial hidden state:

$$\begin{array}{c}\hfill {\mathit{H}}_0=[{\mathit{H}}_{0,0},{\mathit{H}}_{0,1}]\to {\mathit{H}}_0^{\prime }=[{\mathit{H}}_{0,0},{\mathit{H}}_{0,1},{\mathit{H}}_{0,1}].\end{array}$$

(11)

After the layer normalization layer, if we ignore the influence of the mean and variance of the hidden layer state, we have the following:

$$\begin{array}{c}\hfill {\mathit{H}}_1=[{\mathit{H}}_{1,0},{\mathit{H}}_{1,1}]\to {\mathit{H}}_1^{\prime }=[{\mathit{H}}_{1,0},{\mathit{H}}_{1,1},{\mathit{H}}_{1,1}].\end{array}$$

(12)

Then ${\mathit{H}}_1$ and ${\mathit{H}}_1^{\prime }$ are projected through ${\mathit{W}}^{Q\left|K\right|V}$ and ${\mathit{U}}^{Q\left|K\right|V}$:

$$\begin{array}{cc}\hfill \mathit{Q}=[{\mathit{Q}}_0,{\mathit{Q}}_1]& \to {\mathit{Q}}^{\prime }=\sqrt{\sqrt{\frac{D^t}{D^s}}}[{\mathit{Q}}_0,{\mathit{Q}}_1,{\mathit{Q}}_1]\hfill \\ \hfill \mathit{K}=[{\mathit{K}}_0,{\mathit{K}}_1]& \to {\mathit{K}}^{\prime }=\sqrt{\sqrt{\frac{D^t}{D^s}}}[{\mathit{K}}_0,{\mathit{K}}_1,{\mathit{K}}_1]\hfill \\ \hfill \mathit{V}=[{\mathit{V}}_0,{\mathit{V}}_1]& \to {\mathit{V}}^{\prime }=[{\mathit{V}}_0,{\mathit{V}}_1,{\mathit{V}}_1].\hfill \end{array}$$

(13)

Then the attention score of the i-th head in target model T can be written as:

$$\begin{array}{cc}\hfill \frac{{\mathit{Q}}_i^{\prime }{\mathit{K}}_i^{\prime T}}{\sqrt{D^t}}& =\frac{\sqrt{\sqrt{\frac{D^t}{D^s}}}{\mathit{Q}}_i\sqrt{\sqrt{\frac{D^t}{D^s}}}{\mathit{K}}_i^T}{\sqrt{D^t}}\hfill \\ \hfill & =\frac{{\mathit{Q}}_i{\mathit{K}}_i^T}{\sqrt{D^s}},\hfill \end{array}$$

(14)

which is the attention score of the i-th head in the source model S. The analysis of the subsequent layers follows a similar pattern. Also, we expand a pretrained GPT-2 of {L = 12, D = 768} to a target GPT-2 model of {L = 12, D = 1024}. Figure 8 shows that the newly formed attention heads by NAI basically retains the capabilities of the original neighbor heads.

3.3.2. Depth-Wise Expansion

Previous findings [31,32] show that adjacent Transformer layers have similar functionality, and we also find this phenomenon in GPT-2 as illustrated in Figure 2. Based on this phenomenon, when we perform depth-wise expansion, we continuously copy the last layer of the original model to reach the number of layers of the target model. From an intuitive perspective, adopting weight copying from the last layer appears to align more consistently with the gradual and progressive learning process observed in each layer of the transformer model. To preserve the capabilities of the original model, we set the output layers of the MHCA and FFN module to zero, leveraging the residual connection architecture. Moreover, it was observed that failing to set the output layers of the MHCA and FFN modules to zero can lead to volatile training behavior, ultimately causing divergence in the learning process. Figure 9 provides a concise and comprehensible depiction of the depth-wise expansion of NAI. By initializing the deeper model with transferred knowledge from the $L^s$-layer trained model, we anticipate accelerated learning compared to training the deeper model from scratch. The whole procedure is summarized in Algorithm 1.

Algorithm 1 NAI algorithm

1: Input: the source model S($L^s,D^s$) and the target model T($L^t,D^t$) 2: ${\mathit{T}}_{L^s}(L^s,D^t)\leftarrow \mathit{S}(L^s,D^s)$ by width-wise expansion of NAI 3: for $t=L^s+1\to L^t$ do 4:       ${\mathit{T}}_t(L^{t+1},D^t)\leftarrow \mathrm{stack}\mathrm{the}\mathrm{last}\mathrm{layer}\mathrm{of}{\mathit{T}}_{L^s}\mathrm{on}\mathrm{top}\mathrm{of}{\mathit{T}}_{t-1}$ 5:       set the output layer of the MHCA and FFN modules to zero of layer t 6: end for 7: output: the target model $\mathit{T}(L^t,D^t$)

4. Experiment

4.1. Experimental Setup

Pre-training Details. All of our experiments are mainly based on our own reimplementation of the GPT-2 model. We use the concatenation of OpenWebText [33] and English Wikipedia as the pre-training data. For dataset OpenWebText, we performed a de-duplication operation and removed a non-English corpus. For both datasets, sentences with a length below 128 were excluded from both datasets. The settings of the pre-training are peak learning rate of $1.5\times {10}^{-4}$, warmup step of ${10}^4$, training iteration of ${10}^6$, and a small batch size of 128 due to limited computing resources. The AdamW optimizer was used with $[{\beta }_1,{\beta }_2]=[0.9,0.999]$, weight decay = 0.01. Unless otherwise noted, all methods, including NAI and baselines, use the same pre-training settings for fair comparisons. In the settings of NAI and bert2BERT, the target models has a GPT-2 architecture of {24, 1024} and {36, 1280}, and an architecture of source models {12, 768} {24, 1024} are evaluated. Both source models are pretrained using the above settings, and they are training from scratch without any optimization. The computations in this research were performed using the CFFF platform of Fudan University with 2 NVIDIA A100 GPUs.

Baselines. We initially introduce a method training from scratch without any optimization, named ${\mathrm{GPT}}_{\mathrm{base}}$. As mentioned above, all the source models used in NAI and bert2BERT are trained by this. Then, we introduce an approach called DirectCopy, which involves the direct replication of the source model into the target model, while employing random initialization for the parameters that have not been populated. bert2BERT is also included as the baseline. During the width-wise expansion in bert2BERT, advanced knowledge from the subsequent layer is utilized for initializing the large model; the detailed process can be found in [13]. In the depth-wise expansion phase of bert2BERT, the target model is constructed using the StackBERT technique [15], and the detailed process of StackBERT is provided in Algorithm 2. This approach involves the continuous replication of the entire source model after the width-wise expansion. Note that here, we only compare various initialization methods, so the two-stage pre-training in bert2BERT is not considered.

Algorithm 2 StackBERT

1: Input: the source model $\mathit{S}(L^s,D)$ and the target model $\mathit{T}(L^t,D)$ 2: $k\leftarrow \lfloor L^t/L^s\rfloor ,{\mathit{T}}_1\leftarrow \mathit{S}(L^s,D)$ 3: for $t=2\to k$ do 4:       ${\mathit{T}}_t(L^s\ast t,D)\leftarrow \mathrm{stack}{\mathit{T}}_1\mathrm{on}\mathrm{top}\mathrm{of} {\mathit{T}}_{t-1}$ 5: end for 6: $\mathit{T}(L^t,D)\leftarrow \mathrm{stack}\mathrm{the}\mathrm{top}{L}^t-L^s\ast k\mathrm{layer}\mathrm{of}{\mathit{T}}_1\mathrm{on}\mathrm{top}\mathrm{of}{\mathit{T}}_k$ 7: output: the target model $\mathit{T}(L^t,D$)

4.2. Results and Analysis

We present the pre-training loss curves of GPT-2 models with config {24, 1024} in Figure 10. The results reveal the following insights: (1) DirectCopy—This method, which involves directly copying the trained parameters from the source model to the target model, exhibits minimal computational cost savings. This suggests that the simplistic approach of DirectCopy is not effective in achieving cost reductions. (2) bert2BERT—By utilizing advanced knowledge for initializing the large model and employing StackBERT for depth-wise expansion, bert2BERT achieves approximately 21% computational cost savings. However, the departure from the function-preserving initialization principle, caused by the utilization of advanced knowledge and depth-wise expansion, limits the acceleration of loss reduction during the early stages of training. Consequently, the preservation of the source model’s capabilities is not substantially retained. (3) NAI (Our proposed method)—NAI outperforms the baselines in terms of performance. It achieves a 10% improvement in computational cost savings compared to bert2BERT. Notably, NAI adheres to the function-preserving initialization principle during both width-wise and depth-wise operations, resulting in a rapid decrease in loss during the initial stages of training. By preserving the capabilities of the source model, NAI ensures an effective initialization of the target model.

Additionally, we depict the pre-training loss curves of GPT-2 models with config {36, 1280} in Figure 11. In comparison with Figure 10, we observe the following trends: All the methods, including DirectCopy, bert2BERT, and NAI, achieve greater computational cost savings compared to the baseline model ${\mathrm{GPT}}_{\mathrm{base}}$. For example, DirectCopy demonstrates approximately 12% more cost savings compared to ${\mathrm{GPT}}_{\mathrm{base}}$, whereas it had negligible savings in Figure 10. Method bert2BERT achieves a cost reduction of 24%, surpassing the 21% achieved in Figure 10. Similarly, NAI achieves a cost savings of 33%, slightly higher than the 31% in Figure 10. And, it also achieves a 9% improvement compared to bert2BERT. We attribute this phenomenon to the closer scale between the source model and the target model. These detailed analyses provide a comprehensive understanding of the experimental results, highlighting the effectiveness of NAI in achieving computational cost savings while preserving the capabilities of the source model.

We also evaluated the effectiveness of all methods on the LAMBADA, PTB, WikiText2, and WikiText103 test datasets; the following is a brief introduction to them.

LAMBADA: LAMBADA is a dataset designed for testing the ability of language models to perform pronoun resolution in context. It consists of narrative passages followed by a single line continuation. The task is to predict the last word of the continuation, which requires understanding the context provided by the passage.

PTB: PTB is one of the most widely used datasets in natural language processing. It consists of annotated text from various sources, including newswire, transcribed speech, and literature. The dataset is extensively used for tasks such as language modeling, part-of-speech tagging, and parsing.

WikiText2: WikiText2 is a dataset derived from Wikipedia articles. It is designed for language modeling tasks, including next-word prediction. The dataset contains high-quality text with diverse topics and writing styles.

WikiText103: WikiText103 is an extension of WikiText2 and contains a larger corpus of Wikipedia articles. Like WikiText2, it is designed for language modeling tasks and provides a diverse range of textual data.

All these datasets are commonly used in natural language processing research for training and evaluating various language models and algorithms.

The test results are shown in

Table 1. Zero-shot results on different datasets of GPT-2 {24, 1024}. We use the perplexity as the metric.

Model	Flops ($\times {10}^{19}$)	LAMBADA	PTB	WikiText2	WikiText103
${\mathrm{GPT}}_{\mathrm{base}}$	6.2	33.7	68.9	28.7	38.6
DirectCopy	5.8	33.4	68.3	28.4	38.5
bert2BERT	4.9	32.6	67.5	28.1	38.2
NAI	4.3	32.2	67.6	28.3	37.7

and

Table 2. Zero-shot results on different datasets of GPT-2 {36, 1280}. We use the perplexity as the metric.

Model	Flops ($\times {10}^{19}$)	LAMBADA	PTB	WikiText2	WikiText103
${\mathrm{GPT}}_{\mathrm{base}}$	14.4	18.2	45.6	23.2	26.6
DirectCopy	12.7	17.6	45.8	23.4	25.7
bert2BERT	11	17.8	45.5	23.1	25.4
NAI	9.6	17.8	45.2	23.7	26.5

. We employed a zero-shot approach for testing on these datasets, which means no training or fine-tuning was performed on the datasets. Perplexity (PPL) was used as the evaluation metric, and the average of three test results was taken as the final result for each test set. From the tables, we can conclude that both bert2BERT and NAI perform comparably to ${\mathrm{GPT}}_{\mathrm{base}}$ and DirectCopy methods. Furthermore, both bert2BERT and NAI demonstrate significant cost savings in terms of computational resources, with NAI exhibiting even better cost savings in training.

5. Potential Limitations and Challenges of NAI

The NAI method, as presented in our study, offers a novel approach to efficiently train larger PLMs by leveraging the knowledge encapsulated in smaller PLMs. Despite its potential to reduce computational costs and expedite training, there are several limitations and challenges associated with applying NAI to different model architectures and domains.

• Architectural Limitations: NAI is inherently tailored for Transformer-based architectures, which are distinguished mainly by their self-attention mechanisms. These architectural elements are pivotal for the function-preserving initialization process that NAI employs. The self-attention mechanism, in particular, enables the model to weigh the importance of different parts of the input data, whereas layer normalization facilitates the stabilization of the learning process by reducing the internal covariate shift. Given the constraints of computational resources, the empirical validation of NAI has been primarily confined to the GPT-2 model, which is a decoder-only PLMs. This limitation necessitates a cautious extrapolation of NAI’s efficacy to other PLM architectures, such as those based on encoder-only or encoder–decoder configurations. Encoder-only models, such as BERT, utilize a bidirectional training strategy that contrasts with the unidirectional nature of GPT-2, while encoder–decoder models, akin to sequence-to-sequence frameworks, are designed to handle input–output mapping tasks that diverge from the autoregressive generation focus of GPT-2. The transference of NAI to these alternative architectures entails a rigorous examination of potential modifications required to align with their distinct operational mechanisms. For instance, the adaptation of NAI to encoder-only models may necessitate strategies that effectively capture and preserve the bidirectional context during the initialization process. Similarly, encoder–decoder models may demand a reconceptualization of the NAI approach to accommodate the dual-input and dual-output structure characteristic of these architectures.
• Data Availability and Quality: NAI’s efficiency is contingent on the availability and quality of the smaller PLMs used for initialization. In domains where high-quality, well-pretrained models are scarce, the effectiveness of NAI may be compromised. Additionally, the method assumes a certain level of overlap in the knowledge domains between the smaller and larger models, which might not always be the case, especially in niche or rapidly evolving domains.
• Generalization and Adaptability: The success of NAI hinges on the assumption that smaller PLMs have learned foundational language structures that are generalizable to larger models. However, the extent to which this generalization holds true across various tasks and languages remains an open question. Furthermore, adapting NAI to non-English languages or multilingual settings may require additional considerations to account for linguistic diversity.

6. Conclusions and Future Work

This paper introduces a novel and efficient pre-training method called NAI. The key idea behind NAI is to employ the closest attention head for width expansion and the nearest layer for depth expansion to get a larger model. By leveraging the knowledge embedded in the small model, NAI aims to improve the efficiency and effectiveness of pre-training large models.

To evaluate the effectiveness of NAI, we applied this method to the popular GPT model. The experimental results demonstrate the benefits of NAI in terms of computational cost reduction and performance improvement on various test datasets. Notably, our method achieves significant savings in computational resources while maintaining competitive performance compared to alternative techniques.

However, it is important to note that due to the limitations of computational resources, we were unable to conduct extensive experiments to fully validate the algorithm. Despite this limitation, the promising results obtained in our initial experiments serve as a strong indication of the potential of NAI in pre-training large-scale language models.

In the future, we plan to extend the application of NAI to train even larger language models. Moreover, we aim to explore the feasibility and effectiveness of applying NAI to other domains and tasks beyond language modeling, thereby expanding its scope to very large PLMs in various fields. This would enable us to further investigate the benefits and limitations of NAI and contribute to the advancement of efficient and effective pre-training techniques.