H ad S kip: Homotopic and Adaptive Layer Skipping of Pre-trained Language Models for Efficient Inference (2023)

Abstract: Pre-trained language models (LMs) have brought remarkable performance on numerous NLP tasks. However, they require significant resources and entail high computational costs for inference, making them challenging to deploy in real-world and real-time systems. Existing early exiting methods aim to reduce computational complexity by selecting the layer at which to exit, but suffer from the limitation that they have to sequentially traverse through all layers prior to the selected exit layer, which lacks flexibility and degrades their performance. To solve this problem, we propose a h omotopic and ad aptive layer skip ping fine-tuning method named HadSkip. HadSkip adaptively selects the layers to skip based on a predefined budget. Specifically, we introduce a learnable gate before each layer of the LM to determine whether the current layer should be skipped. To tackle various challenges in training such as discrete gates and the budget constraint, we propose a fine-grained initialization strategy and homotopic optimization strategy. We conduct extensive experiments on the GLUE benchmark, and experimental results demonstrate the proposed HadSkip outperforms all state-of-the-art baselines significantly.

## introduction pre-trained language models (lms), such as bert @xcite , roberta @xcite , and deberta @xcite , have significantly improved performance across various natural language processing (nlp) tasks, including paraphrase identification @xcite , natural language inference @xcite , sentiment classification @xcite , and so on. pre-trained lms are typically built by stacking transformer @xcite layers or their variants, where the self-attention mechanism within each layer often exhibits high computational complexity. consequently, the inference complexity of pre-trained lms has been a bottle-neck for deploying them in latency-sensitive or latency-constrained scenarios @xcite . to comply with the demand for inferenceefficient models, several approaches have been proposed to accelerate pre-trained lm inference. these methods include weight quantization @xcite , pruning @xcite , knowledge distillation @xcite , and early exiting @xcite @xcite . among them, early exiting methods become more appealing as they do not rely on specific hardware support, such as custom chips @xcite or gpus, and they do not require training an efficient model from scratch on large pre-training corpora, thus providing a more cost-effective solution. early exiting methods attach a classifier to each transformer layer and aim to determine if the model should exit (stop inference) from the current layer based on a criterion or learned decision. both the backbone model and the classifiers are jointly trained in the finetuning stage. although early exiting methods enable the selection of a layer to exit based on the input sequence, they still require the input to pass through the transformer layers sequentially before exiting. this prevents the model from choosing desired uncontiguous layers. for example, with a 2-layer budget, for some input sequences, choosing the second and fifth layers of the model could achieve the most superior performance. however, early exiting methods can only traverse the layers sequentially and exit at most at the second layer, thereby missing out on the optimal combination of the second and fifth layers and potentially leading to degraded model accuracy as a result of suboptimal layer selection. to address this issue, we propose a homotopic and adaptive layer skipping fine-tuning method named hadskip. the proposed hadskip selects which layers can be skipped based on the difficulty of input sequences and pre-defined budget. specifically, we introduce a binary learnable gate before each transformer layer of the lm to control whether the model can skip the current layer. this allows the model to use different gate values for different input sequences, satisfying the inference budget requirements. however, learning binary gates is a non-trivial task. the binary gates are discrete and it is difficult to optimize by gradient-based optimizers directly. besides, the presence of a budget constraint can lead to early optimization collapse, particularly when the budget value is small. to tackle them, we design an elaborated three-stage training strategy to tackle this challenge. the first stage finetunes the backbone model and uses it to initialize the transformer layers. next, these transformer layers are frozen and the gate parameters are initialized. in the third stage, the transformer layers and gates are trained jointly. to compute the gradient of gates, we employ a second-order gradient approximation method reinmax @xcite . further, we propose a homotopic optimization strategy, which converts the optimization process into a sequence of problems ranging from a large budget to a small budget. on one hand, this ensures a smoother and more stable optimization process. on the other hand, it avoids skipping a large number of layers at the beginning of training and losing valuable information from the pre-trained model. the contributions of the paper are summarized as follows: 1) we propose a homotopic and adaptive layer skipping fine-tuning method to accelerate the inference of pre-trained lms. it is more flexible and economical compared to existing lm inference acceleration methods. 2) we design a three-stage and homotopic optimization strategy. this strategy makes the optimization more stable and easier to solve. 3) we conduct extensive experiments on the glue benchmark. results show that the proposed hadskip outperforms baselines significantly and achieves a good efficiency-accuracy trade-off, e.g., preserving 95.7% of bert's performance with only half the layers on average. ## related work existing approaches for accelerating the inference of pre-trained language models can be broadly categorized into two types: 1) model compressionbased methods and 2) early exiting based methods. the proposed method is more relevant to early exiting based methods, which are discussed in more detail as follows. the discussion about the model compression-based methods can be found in the appendix. early exiting methods enable the production of inference-efficient models for specific inference budgets during the fine-tuning stage. these methods draw inspiration from adaptive computation @xcite techniques in recurrent neural networks and branchynet @xcite in computer vision. the fundamental concept behind early exiting methods is to identify the layer at which to exit early, rather than sequentially passing through all layers. based on the criteria of choosing the exiting layer, they can be organized into three strategies. the first is scorebased early exiting. @xcite @xcite use the entropy of the prediction probability and the maximum of the predicted distribution as the score for exit determination respectively. the second is learning to exit. berxit @xcite used a fully connected layer right after each layer to produce decisions if the model needs to exit from the current layer. the third is patience-based early exiting. for example, pabee @xcite , sentee @xcite and leebert @xcite were designed to use scores of multiple layers to determine if the model could exit. recently, pcee-bert @xcite also combined the score-based method with the patience-based early exiting method to improve accuracy. despite these advancements, early exiting methods still require sequential traversal through multiple layers. however, for many input sequences, the optimal layers for inference could be non-contiguous. this lack of flexibility in early exiting methods can limit their performance. to address this limitation, inspired by blockdrop @xcite , we propose a homotopic and adaptive layer skipping fine-tuning method named hadskip, which can dynamically select which layers can be skipped based on the difficulty of input sequences and the pre-defined budget. concurrently, smart-bert @xcite proposes to combine layer skipping and early exiting based on cross-layer contrastive learning. ## preliminaries in this section, we first introduce the architecture of the pre-trained language model (plm) and analyze its complexity for inference. ## methodology given an l-layer pre-trained language model (plm) and a training set @xmath0, y i )} n @xmath1, where x i is the input sequence, y i is the corresponding label and n is the number of training data, the aim is to train a model f (x i ) → y i which can be efficient for inference while preserving the model performance. ## experiment in this section, we evaluate the proposed hadskip and answer the following questions: rq1) how does hadskip perform compared to state-of-theart early exiting baselines? rq2) what are the roles of initialization, knowledge distillation, and homotopic optimization strategy in model performance improvements respectively? rq3) can the proposed hadskip generalize well with respect to different backbones? rq4) is it possible to combine the proposed hadskip method with other model acceleration techniques? rq5) how does the performance change with varying budget values? rq6) can the proposed hadskip capture the difficulty differences among input sequences? ## conclusion in this paper, we propose a homotopic and adaptive layer skipping fine-tuning method named hadskip for efficient inference. we introduce a learnable gate before each transformer layer of the pretrained language model to determine whether the model can skip the respective layer. for training the binary gate, we introduce a three-stage learning strategy to initialize and update model parameters. additionally, we propose a homotopic optimization strategy to stabilize the model training. we also utilize knowledge distillation to guide model training and achieve an improved trade-off between efficiency and accuracy. we conducted extensive experiments on the glue benchmark, and the results demonstrate that hadskip outperforms state-of-theart baselines. moreover, hadskip complements other acceleration methods, including tinybert.