Early Exit with Disentangled Representation and Equiangular Tight Frame (2023)

Abstract: Dynamic early exit has demonstrated great potential in coping with the sharply increasing number of pre-trained language model parameters, which can achieve a good trade-off between performance and efficiency. The existing early exit paradigm relies on training parametrical internal classifiers at each intermediate layer to complete specific tasks. Based on the predictions of these internal classifiers, different methods are designed to decide when to exit. Under this circumstance, each intermediate layer takes on both generic language representation learning and task-specific feature extraction, which makes each intermediate layer struggle to balance two types of backward loss signals during training. To break this dilemma, we propose an adapter method to decouple the two distinct types of representation and further introduce a non-parametric simplex equiangular tight frame classifier (ETF) for improvement. Extensive experiments on monolingual and multilingual tasks demonstrate that our method gains significant improvements over strong PLM backbones and early exit methods.

## introduction in recent years, fundamental models that rely on the scaling effect have penetrated different nlp scenarios @xcite @xcite @xcite @xcite @xcite . however, with the increasing number of the pre-trained model parameters, the expensive inference cost hinders their usage in practical applications. besides, overthinking problem @xcite also restricts the ability of plms. specifically, since plms are overparameterized, they can give correct answers according to the shallow representations at earlier layers, while the high-level representation * equal contribution. † corresponding author. our method uses an additional adapter module to undertake the learning of task-specific representation. to achieve a good trade-off between inference cost and performance, the early exit mechanism @xcite @xcite @xcite @xcite , a kind of adaptive inference strategy, has been proposed. these methods insert an internal classifier after each layer of the plms to predict the label of a given instance. in the inference stage, if the prediction is confident enough earlier, the sample will end the inference without going through the entire plm. nevertheless, how to train a competitive early exit plm is not a trivial problem. since each classifier tries to be optimized, different optimization procedures from different classifiers may conflict and interfere with each other @xcite . exist-ing methods can be divided into two categories: (1) jointly training all classifiers and using heuristics or learnable methods to weight loss functions @xcite @xcite ; (2) training classifiers in different stages and freezing the transformer layer when training internal classifiers. however, in the early exit mechanism, each transformer layer plays two roles: providing classifiable representation for the corresponding internal classifier and semantic features for subsequent transformer layers. the former class of methods is obsessed with improving the classification ability of the middle layer of plms while ignoring its ability to capture other linguistic features. the latter class focuses on maintaining the ability to extract semantic representations while constraining the performance of internal classifiers. although some work @xcite has attempted to balance the two roles of the transformer layer by alternate training, they have not explicitly decoupled the two distinct types of representation. in this work, we propose a novel early exit method that enables the plm's intermediate representation to consider both the classification ability and the ability to capture semantic information. specifically, our approach hands off the taskspecific representation of the intermediate transformer layer to an adapter module @xcite with a small number of parameters, thus disentangling the task-specific and universal representation (see figure 1 ). in addition, due to the limited expressive power of the adapter, we use a simplex equiangular tight frame classifier (etf) @xcite to enhance the classification ability of the internal classifiers. experimental results demonstrate that our proposed early exit method significantly improves the performance of monolingual and multilingual tasks over existing strong early exit methods. overall, our contributions are shown below: ## preliminary study it is controversial whether the last classifier and the internal classifiers should be trained jointly in early exit. some studies @xcite @xcite divide the training into two stages to train the last classifier and the internal classifiers, respectively, to preserve the best model ability for the final layer. while others @xcite ) train the whole model simultaneously. in this section, we study the question:"can the same representation serve both the classifier and subsequent layers?" following @xcite , we train layerwise probes to check how much linguistic knowledge is preserved in each layer after finetuning. we evaluate the model on two linguistic tasks: pos tagging using the universal dependencies v2.5 english dataset from xtreme @xcite and syntactic chunking using conll 2000 shared task dataset @xcite . we conduct experiments on pabee @xcite , which jointly trains the last and internal classifiers. as shown in figure 2 , the red lines refer to the performance of bert without any finetuning. the orange and blue lines represent the performance of bert finetuned on sst-2 with and without pabee. the model preserves the information of the downstream tasks in the higher layers with a corresponding loss of the linguistic knowledge learned in the pre-training after finetuning. moreover, the gaps between the blue and orange lines are apparent, indicating that the training with pabee further disturbs linguistic knowledge. on the other hand, a representation trained only to serve the subsequent layers often has a poor performance for classification, especially in the lower layers. as a result, to maintain the internal representations' capability and improve the classification performance, it is essential to train disentangled representations in the intermediate layers. ## method in this section, we introduce our method for early exit, which first utilizes adapter modules to disentangle generic and task-specific representation. to further improve the classification representation, we replace learnable classifiers with equiangular tight frame (etf) classifiers. throughout this section, we consider the case of multi-class classification with samples {(x i , y i )} n @xmath0, where x i is a token sequence and y i is its label. ## related work early exit early exit focuses on enabling inputadaptive inference to reduce the computational cost, which has been proven effective on various nlp tasks @xcite @xcite @xcite @xcite . because the model needs to decide whether stop inference at a specific layer when using early exit, so there are two main problems: 'how to decide whether stop inference at each layer?' and 'how to induce a better internal classifier for each layer?'. to the former problem, @xcite ; @xcite ; @xcite use confidence-based criterion; @xcite propose a novel and effective patience mechanism; @xcite propose a voting-based strategy. in addition, @xcite ; @xcite utilize a latent variable to predict the exit layer's index. for the latter one, @xcite al. (2020); geng et al. (2021) adapter with the increase in the number of parameters of plms, full-parameter fine-tuning faces difficulties in computing resources and is prone to the over-fitting problem @xcite @xcite . therefore, research on efficient fine-tuning @xcite @xcite is developing rapidly. adapter is a promising efficient fine-tuning method. adapter-based methods inject small-scale adapters to the transformer layers and only tune these adapters. many studies @xcite @xcite @xcite have shown that adapter-based methods can achieve comparable performance to the full-parameter finetuning. to further reduce the parameter amount of the adapter module, karimi @xcite propose compacter, a more lightweight adapter that utilizes a combination of hypercomplex multiplication and parameter sharing. due to the modularity of the adapter, adapter-based tuning is suitable for multi-task learning @xcite and can provide different task-specific representations for various tasks. adapterfusion @xcite proposes a fusion method named adapterfusion that fuses adapter representations of different tasks and makes full use of cross-task knowledge. ## conclusion in this paper, we investigate the internal representation of early exit models and observe that the internal layers have difficulty providing good generic linguistic representations for subsequent layers and good task-specific representations for internal classifiers. we propose an adapter-based method to disentangle the two conflicting representations and utilize equiangular tight frame classifiers to improve representations for classification. experiments on the glue benchmark and two cross-lingual transfer tasks demonstrate that our proposed method performs better than existing methods. for future work, we would like to explore: (1) strengthening the information interaction between layers to make full use of the previous layers' predictions, thus optimizing the representation for the internal classifiers; (2) optimizing the early exiting mechanism to further improve performance and accelerate inference; (3) applying our method to more advanced pre-trained models such as deberta @xcite , elasticbert @xcite , etc.