A Global Past-Future Early Exit Method for Accelerating Inference of Pre-trained Language Models (2021)

Abstract: Early exit mechanism aims to accelerate the inference speed of large-scale pre-trained language models. The essential idea is to exit early without passing through all the inference layers at the inference stage. To make accurate predictions for downstream tasks, the hierarchical linguistic information embedded in all layers should be jointly considered. However, much of the research up to now has been limited to use local representations of the exit layer. Such treatment inevitably loses information of the unused past layers as well as the high-level features embedded in future layers, leading to sub-optimal performance. To address this issue, we propose a novel Past-Future method to make comprehensive predictions from a global perspective. We first take into consideration all the linguistic information embedded in the past layers and then take a further step to engage the future information which is originally inaccessible for predictions. Extensive experiments demonstrate that our method outperforms previous early exit methods by a large margin, yielding better and robust performance.

## introduction pre-trained language models (plms), e.g., bert @xcite , roberta @xcite and xlnet @xcite , have obtained remarkable success in a wide range of nlp tasks. despite their impressive performance, plms are usually associated with large memory requirement and high computational cost. such drawbacks slow down the inference and further encumber the application of plms in the scenarios where inference time and computation budget are restricted. to address this issue, a growing number of studies focusing on improving model efficiency have emerged recently. particularly, @xcite point out that the current over-parameterized models conduct excessive computation for simple instances, which is actually undesirable and computationally wasteful. in light of this observation, an increasing amount of work seeks various early exit methods, of which the basic idea is to exit early without passing through the entire model during inference. concretely, for nlp tasks, they couple branch classifiers with each layer of the pre-trained language models and stop forward propagation at an intermediate layer. then the current branch classifier makes a prediction based on the representation of the token that is used as the aggregated sequence representation for classification tasks and is referred to as the state of the layer in this work. however, existing work on early exit has two major drawbacks. first, existing work @xcite uses only local states in the early exit framework. they inevitably lose valuable features that are captured by passed layers but are ignored for prediction, leading to less reliable prediction results. moreover, these methods abandon the potentially useful features captured by the future layers that have not been passed, which may hurt the performance of the instances requiring high-level features embedded in the deep layers. consequently, their performance dramatically declines when the inference exits earlier for a higher speed-up ratio. these two major drawbacks hinder the progress of early exit research and motivate us to develop a new mechanism using the hierarchical linguistic information embedded in all layers @xcite from a global perspective. however, up to now, a global early exit mechanism remains a under-explored challenging problem. we extend the existing methods to their corresponding global versions and find that naive global strategies only result in poor performance. meanwhile, the future states are originally inaccessible in the early exit framework, which also remains a bottleneck for a global prediction considering both past and future states. in this paper, we focus on the aforementioned problems and first put into practice a global past-future early exit mechanism. the term global is two-fold: (1) instead of using one or several local state(s) for prediction in previous work, all the available past states are effectively incorporated in our method; (2) furthermore, to grasp the features embedded in the deep layers, the originally inaccessible future states are approximated by imitation learning and are also engaged for prediction. the comparison of the previous method and our method is illustrated in figure 1 . by combining both past and future states, our model is able to make more accurate predictions for downstream tasks. extensive experiments reveal that the proposal significantly outperforms previous early exit methods. particularly, it surpasses the previous methods by a large margin when the speed-up ratio is relatively high. in addition, extensive experiments with different pre-trained language models as backbones demonstrate consistent improvement over the baseline methods, which verifies the generality of our method. to summarize, our contributions are as follows: ## related work large-scale pre-trained language models @xcite based on the transformer @xcite architecture demonstrate superior performance in various nlp tasks. however, the impressive performance is on the basis of massive parameters, leading to large memory requirement and computational cost during inference. to overcome this bottleneck, increasing studies work on improving the efficiency of overparameterized pre-trained language models. knowledge distillation @xcite @xcite compacts the model architecture to obtain a smaller model that remains static for all instances at the inference stage. @xcite focus on reducing the number of layers since their investigation reveals variations on hidden size dimension have a smaller impact on computation efficiency. @xcite learn from multiple intermediate layers of the teacher model for incremental knowledge extraction instead of only learning from the last hidden representations. further, @xcite design elaborate techniques to drive the student model to mimic the self-attention module of teacher models. @xcite compress model by progressive module replacing, showing a new perspective of model compression. however, these static model compression methods treat the instances requiring different computational cost without distinction. moreover, they have to distill a model from scratch to meet the varying speed-up ratio requirements. to meet different constraints for acceleration, another line of work studies instance-adaptive methods to adjust the number of executed layers for different instances. @xcite select models in different sizes depending on the difficulty of input instance. besides, early exit is a practical method to adaptively accelerate inference and is first proposed for computer vision tasks @xcite . elbayad et al. (2020); xin et al. (2020); schwartz et al. (2020) follow the essential idea and leverage the method in nlp tasks. to prevent the error from one single classifier, @xcite make the model stop inference when a cross-layer consistent prediction is achieved. however, researches on the subject has been mostly restricted to only use the local states around the exit layer. ## method we first introduce the strategies to incorporate multiple states and the imitation learning method for generating approximations of future states. then we introduce the merging gate to adaptively fuse past and future states. at last, we show the training process and the exit condition during inference. ## conclusions we propose a novel past-future early exit method from a global perspective. unlike previous work using only local states for prediction, our model employs all available past states for prediction and propose a novel approach to engage the future states which are originally inaccessible for prediction. experiments illustrate that our method achieves significant improvement over baseline methods with different models as backbones, suggesting the superiority of our early exit method.