Class-aware Initialization of Early Exits for Pre-training Large Language Models

Authors: alperen  görmez
erdem  koyuncu
Venue: ICML_cc_2024_Workshop_WANT
Published: 2024-06-02 18:46:13

Abstract: We propose a novel class-aware weight initialization technique for early exit large language models with the purpose of accelerating pre-training. Our design utilizes the neural collapse phenomenon combined with a Gaussian mixture model for the distribution of feature vectors at a given layer. Specifically, we calculate the average of token representations at the early exit point and use the resulting vectors together with class probabilities for initializing the early exit vectors. The next token prediction accuracy of our class-aware initialization technique is up to five times higher than other baselines at epoch zero and matches or surpasses them in later epochs throughout the pre-training process.

Main paper: 
 introduction: 
 State-of-the-art large language models (LLMs) have a large number of parameters, and generally, the higher the number of parameters, the better the performance (Sutton, 2019; Brown et al., 2020; Zhang et al., 2022; Touvron et al., 2023; Jiang et al., 2024; Gemini Team, 2024; OpenAI, 2024; Meta, 2024) . However, their large size and autoregressive design results in high inference latency, which is not desirable for low resource environments and time sensitive settings.The vast majority of LLMs have a "tunnel-like" architecture:The input to the model is processed by all of the layers in a sequential manner, regardless of the input's inherent difficulty (Kaya et al., 2019; Görmez et al., 2022) . On the other hand, not all inputs have the same level of difficulty. Early exit networks exploit this heterogeneous difficulty of inputs. One or more intermediate classifiers are attached to the model, allowing token-level conditional computation (Panda et al., 2016; Teerapittayanon et al., 2016; Kaya et al., Accepted to the Workshop on Advancing Neural Network Training at International Conference on Machine Learning (WANT@ICML 2024).2019; Görmez et al., 2022; Schuster et al., 2022; Del Corro et al., 2023; Bae et al., 2023; Zhu et al., 2024) . Easy tokens can exit early from the LLM in order to save computation.While the addition of early exits can reduce the inference latency, initially they do not possess the optimal weights. The early exit layers have to be trained first, before being effective at inference time. Early exit layers are typically trained together with the backbone model, with two primary approaches: Training only the early exit and the final exit layers while freezing the non-exit layers, or training the backbone and the exits together (Teerapittayanon et al., 2016; Kaya et al., 2019) . Generally, the latter performs better since everything is optimized jointly, but the cross-talk between the exits of the network may lead to suboptimal learning and long training times (Kaya et al., 2019) . Ideally, we would like to initialize the weights of the early exit layers in such a way that the cross-talk is minimized, the joint training is facilitated and training time is reduced.The sizes of both the state-of-the-art LLMs and their training data lead to long training time and high costs. This makes training an early exit LLM even more difficult and costly. In this work, we propose a novel class-aware early exit initialization technique for early exit LLMs to reduce the pre-training costs. We make connections to the optimal detection for the vector additive white Gaussian noise (AWGN) channel from the digital communications domain and utilize the neural collapse phenomenon (Papyan et al., 2020) . Specifically, we calculate the average of token representations at the early exit point and use the resulting vectors for the initialization. While calculating the average of vector representations has been shown to work well as a decision mechanism for early exit networks (Görmez et al., 2022; Görmez & Koyuncu, 2024) , our work is the first to apply it to the initialization of early exit LLMs with the purpose of accelerating pre-training.We demonstrate the effectiveness of our novel weight initialization technique on WikiText-2 (Merity et al., 2016) and BookCorpus (Zhu et al., 2015) datasets using OPT (Zhang et al., 2022) and TinyLlama (Zhang et al., 2024) The rest of the paper is organized as follows: In Section 2, we provide a summary of the literature. Then, we establish the notation used throughout the paper in Section 3 and provide a background on the optimal detection problem from the digital communications domain, which we will make connections to later on. In Section 4, we describe our class-aware early exit initialization technique. Finally, we present the results of our experiments in Section 5. 
early exit llms: 
 Numerous attempts have been made in the past to reduce the inference latency of transformer (Vaswani et al., 2017) based models in the past. Perhaps the most visited idea has been adding early exits to BERT variants (Devlin et al., 2018; Zhou et al., 2020; Xin et al., 2020; 2021; Zhu, 2021) . However, these models are primarily designed for classification tasks such sentiment analysis, rather than language modeling and text generation.Developing early exit LLMs for text generation is more challenging, because token-level early exiting requires careful consideration (Elbayad et al., 2019; Liu et al., 2021) . Copying hidden states of tokens that exited early to the deeper layers for KV-caching, which confidence measure to use and batch inferencing have been tackled in the past (Schuster et al., 2022; Bae et al., 2023; Del Corro et al., 2023) .As the number of parameters in an LLM grows, fine-tuning it on datasets becomes more time-consuming and expensive. To address this challenge, researchers have explored parameter-efficient fine-tuning techniques such as adapter approaches, often coupled with quantization (Hu et al., 2021; Liu et al., 2022; Zhang et al., 2023; Dettmers et al., 2024) . These methods involve training only a subset of the model parameters, effectively reducing the overall training cost. Most recently, parameter-efficient fine-tuning of early exit LLMs has been explored via data, tensor and pipeline parallelism (Chen et al., 2023; Pan et al., 2024) . 
preliminaries and problem formulation: 
 In this section, we establish our notation and lay the foundation for our method by describing the pre-training process of a decoder-only LLM. We then provide background on the problem of optimal detection for the vector AWGN channel from the digital communications domain, which we will make critical connections to later on.As long as you live, keep learning how to live.Embedding Layer Decoder 1 Decoder L ... T 1 T C ... T i ∈ {1,…, V} R C,0 ... R 1,0 R C,L ... R 1,L P C ... P 1 Softmax argmax T 1 T C ... R i,l ∈ ℝ D P i ∈ ℝ VWe focus on the pre-training phase of models belonging to the family of decoder-only LLMs. The model consists of an embedding layer, L decoder blocks and a language modeling (LM) head. Let V , D, C denote the vocabulary size, the embedding dimensionality, and the context length, respectively.During the pre-training process, the tokenizer breaks down a text from the training set into C tokens, denoted as T i ∈ {1, . . . , V }, where i ∈ {1, . . . , C}. Let S v denote the set of all training tokens T i at any position i such that T i = v.The tokens are subsequently passed through the embedding layer in parallel, resulting in corresponding vectors R i,0 ∈ R D . These vectors are then fed into the first decoder block, generating output vectors R i,1 ∈ R D . This iterative process continues sequentially, with the output R i,l of decoder block l being passed to decoder block l + 1 as the input, where l ranges from 1 to L -1.In the final stage of the feed-forward process, the output R i,L of the last decoder block is fed to the LM head, which is a linear layer with weight matrix W ∈ R D×V . The output of the LM head is converted to probability vectors P i ∈ R V via the softmax operation. Suppose the index of the maximum probability in P i is Ti . Since the primary objective of the pre-training phase is next-token prediction, the model is optimized with the cross-entropy loss to ensure Ti = T i+1 . This process is shown in Figure 1 .In order to accelerate inference, one or more early exit LM heads can be integrated to the already pre-trained decoderonly language model. However, the integration of the additional layer(s) necessitates a separate pre-training, which may incur substantial costs as discussed in Section 1. Here, we assume that only one early exit LM head is added. Suppose that this LM head appears after decoder block K where K < L, and its weight matrix is W ∈ R D×V , sharing the same dimensions with the backbone LM head. Our main goal is to find a smart way of initializing W such that the pre-training phase for the early exit LM head will be accelerated, and therefore training costs associated with it will be decreased. Our proposed solution relies on the problem of optimal detection for the vector AWGN channel.The vector AWGN channel can be modeled asEQUATIONwhere r, s m and n are N -dimensional vectors. A message s m is sent to the receiver through the AWGN channel, which adds a noise n to the message. The components of the noise vector are independent and identically distributed Gaussian random variables with zero mean and N0 2 variance. The receiver observes r, and decides which message was sent among {s 1 , . . . , s M }. The goal is to minimize the probability of error. Using the Bayes rule, the optimal detection rule can be written asEQUATIONAs in Equation 4.2-15 from (Proakis & Salehi, 2008 ), the equation above can be simplified further as follows:EQUATIONwhere η m = N0 2 ln P (s m ) -1 2 ∥s m ∥ 2 . The careful reader will notice the striking similarity between the last line of Equation ( 3) and the operational logic of a linear layer serving as a classification head. Given an input x, the linear layer with weights w and biases b classifies the input according to the maximum element of b + x • w. 
method: 
 Our aim is to initialize the early exit LM head in such a way that it starts from a reasonably good point and achieve a certain level of next-token prediction accuracy before any pre-training, rather than starting from a random point achieving a low next-token prediction accuracy.We calculate the average of all output vectors after decoder K that correspond to the tokens in S v :EQUATIONNote that the backbone model is already pre-trained at this point, therefore the intermediate representations are not bad representations. The underlying idea behind Equation ( 4) is the neural collapse phenomenon (Papyan et al., 2020) :The intermediate representation of an input belonging to a certain class converges to its corresponding class mean in the final layer of the network. Here, we carry this idea one step further and postulate that, if the input token T i satisfies T i ∈ S v for some class/word v, then the corresponding feature R i,K at layer K is a Gaussian random vector with mean M v (i.e. the class mean in (4)), and covariance N0 2 I, where N 0 is a hyperparameter to be tuned experimentally. Now suppose that the early exit LM head is the receiver we mentioned in Section 3.2. In this context, we can writeEQUATIONwhere R i,K , M v , and ϵ are all the D-dimensional vector. Also, ϵ, is a zero-mean noise vector with covariance N0 2 I. The mean vector M v is sent to the early exit LM head as the message, and noise ϵ has been added during transmission. The early exit LM head observes R i,K , and decides which mean vector was sent among {M 1 , . . . , M V }.Similar to Equation (3), the optimal decision equation for the early exit LM head can be written asEQUATIONwhere N 0 is a hyper-parameter and P (M v ) is the prior probability for each token in the training set, determined using the empirical frequencies in the training set.As a result, the early exit LM head is initialized asEQUATIONwith a separate bias vector η = [η 1 , . . . , η V ]. Our initialization method is shown in Figure 2 .Tokenizer Embedding Layer Decoder 1 Decoder L LM Head Decoder KT 1 T C … … … R 1,K R C,K … Softmax argmax P 1 P C … T 1 T C … M v = 1 | S v | ∑ T i ∈S v R i,K EE LM Head W : M 1 M 2 M V η : η 1 η 2 η V η v = N 0 2 ln P(M v ) - 1 2 M v 2Already pre-trainedTo be pre-trained argmaxT 1 T C … Figure 2.Our proposed method for initializing the early exit LM head using the mean representation vectors for each token in the vocabulary set.  
experiments and results: 
 In this section we describe our experiments in detail and present the numerical results.In our experiments, we used OPT-125M, OPT-350M, OPT-1.3B models from the OPT model family (Zhang et al., 2022) ; as well as TinyLlama-1.1B (Zhang et al., 2024) .The OPT model family is a series of open-sourced decoderonly language models ranging from 125M to 175B parameters. The largest OPT model performs similarly to GPT-3 (Brown et al., 2020) with approximately 1/7 th of the training cost (Zhang et al., 2022) .The TinyLlama1.1B model is developed with the goal of pretraining such a compact model on 3 trillion tokens (Zhang et al., 2024) . The model shares the same architecture with Llama2 model family (Touvron et al., 2023) .The number of decoder layers (L), embedding dimensionality (D) and vocabulary size (V ) of the models we used in our experiments are shown in Table 1 .In our experiments, we used the WikiText-2 (Merity et al., 2016) and the BookCorpus (Zhu et al., 2015) datasets for pre-training the models.WikiText-2 is a collection of tokens extracted from the verified "Good" and "Featured" articles from Wikipedia (Merity et al., 2016) . For pre-training, we used the "wikitext-2-v1" subset from HuggingFace, which contains 44.8K rows.BookCorpus dataset is a large collection containing more than 11K books and 74M rows (Zhu et al., 2015) . Due to its size, we used 1% of the dataset. We allocated 80% of the 1% for pre-training, and the remaining 20% for evaluation.For a backbone model, we begin by downloading the most recent checkpoint from HuggingFace. This checkpoint is the result of training the model on a large and diverse collection of datasets. However, since we are going to add an early exit layer and pre-train it on only one dataset, we fine-tune the backbone model on our dataset for 3 epochs so the effect of other datasets on the model is minimal. We found out that fine-tuning for more than 3 epochs led to overfitting.After the initial fine-tuning of the backbone model, we add the early exit LM head after decoder K = L/2. Specifically, K is 6, 12, 12 and 11 for OPT-125M, OPT-350M, OPT-1.3B and TinyLlama-1.1B models respectively. The early exit LM heads share the same architecture and number of parameters with the backbone LM head, the only difference is that we allow the early exit LM head to have a bias vector.We train the resulting early exit LLM on two different settings. In the first setting, called "no freezing," all parameters are trainable. In the second setting, called "freezing," we freeze the parameters of the model except the two LM heads. These two settings are how early exit neural networks are trained in the literature (Scardapane et al., 2020; Laskaridis et al., 2021) . The training is done on a single NVIDIA A6000 with a batch size of 32. Due to limited hardware memory, we used a context length of C = 128. We used PyTorch (Paszke et al., 2019) in our experiments.We now present the results of our class-aware early exit initialization method and compare it against two other initialization techniques:1 Random initialization: This is the default weight initialization technique for linear layers in PyTorch (Paszke et al., 2019) . The W ∈ R D×V are initialized asW ∼ U -1 √ D , 1 √ D ,where U is the random uniform distribution.2 Copy-from-backbone: Since the weights W of the backbone LM head is already pre-trained and have the same dimensions as W , copying W into W can serve as a good starting point (Pan et al., 2024) .For our class-aware initialization, we use N 0 = 0.25 and we use the empirical frequencies of tokens in the training set for P (M v ), i.e., the number of occurrences of the token divided by the total number of tokens in the training set.We report the next-token prediction accuracy throughout the pre-training epochs for all initialization techniques. We pre-trained the models for 10 epochs as performance started to drop due to overfitting. The results on the WikiText-2 datasets are shown in Figure 3 .The most important takeaway from Figure 3 is that, our class aware initialization technique achieves 25% next-token prediction accuracy at epoch zero, without any training.On the other hand, random initialization and copying from backbone can achieve at most 5%. This shows that classaware initialization of early exits is a promising technique for resource constrained devices and settings.For the "no freezing" setting, although class-aware initialization starts pretty well, it is surpassed by the copy-frombackbone method easily. There are also some scenarios where random initialization surpasses the class-aware initialization as in Figure 3c and Figure 3e . Here, we can easily match the baselines via a convex combination:EQUATIONwhere W CA is the weights initialized in a class-aware manner, and W B is either random initialized weights or the copied weights from the backbone LM head. This convex combination gets the best of both worlds: It helps preserve the performance of class-aware initialization at epoch zero, and it matches the copy-from-backbone performance at later epochs. In our experiments we evaluated α ∈ {0.2, 0.4, 0.6, 0.8}, and we show the best performing α-curves in Figure 3 .In the "freezing" setting, only the LM heads are trainable, therefore learning is more difficult. As it can be seen from Figure 3b , Figure 3d , Figure 3f , Figure 3h ; the random and copy-from-backbone methods struggle heavily and cannot achieve a good next-token prediction accuracy. On the other hand, our class-aware initialization starts from a pretty good point and keeps performing at the same level throughout the pre-training. Only for the OPT-125M model, there is a sharp drop at the first epoch of the pre-training as seen in Figure 3b . This drop can be somewhat treated by the convex combination equation given in Equation ( 8).The same trends are observed for the BookCorpus dataset as seen in Figure 4 . Specifically, without any training, the classaware initialization starts from a high next-token prediction accuracy and the convex combination allows preserving the high performance throughout the pre-training. Notably, in the "freezing" setting, class-aware initialization performs the best 
conclusion: 
 We developed a novel class-aware weight initialization technique for early exit LLMs based on mean representation of tokens. We made connections to the optimal detection problem for the vector AWGN channel from the digital communications domain. Our method performs better than baselines in both "no freezing" and "freezing" settings. We showed the applicability of our method to various model families and datasets, and its effectiveness on accelerating the pre-training phase.  
 
 Appendix: Department of Electrical and Computer Engineering, University of Illinois Chicago, Chicago, IL, USA. Correspondence to: Alperen Görmez <agorme2@uic.edu>. 


Decision: Overall reviewer sentiment for this submission appears to be positive. Some common strengths and weaknesses pointed out by reviewers:

* (+) Multiple reviewers find that the work addresses an important and relevant problem.
* (+) The proposed method is relatively simple and easy to implement.
* (+) Paper is well-written and easy to follow, although there's room for improvement.
* (+) Reasonable experimental evaluation and zero-shot results; multiple minor issues, as listed below.
* (-) Experimental settings: reviewers point out issues with the size and scale of the datasets, lack of evaluation on downstream tasks, and fairness of comparisons (eg: comparing at epoch 0).
* (-) One reviewer points out that the basis for the work (neural collapse) hasn't been verified on larger networks and datasets.
* (-) Baselines: reviewers note the lack of comparisons to appropriate baselines (eg: Schuster et al.)

I recommend acceptance (poster), and request the authors to incorporate the changes suggested by reviewers.