S kip BERT : Efficient Inference with Shallow Layer Skipping (2022)

Abstract: In this paper, we propose SkipBERT to accelerate BERT inference by skipping the computation of shallow layers. To achieve this, our approach encodes small text chunks into independent representations, which are then materialized to approximate the shallow representation of BERT. Since the use of such approximation is inexpensive compared with transformer calculations, we leverage it to replace the shallow layers of BERT to skip their runtime overhead. With off-the-shelf early exit mechanisms, we also skip redundant computation from the highest few layers to further improve inference efficiency. Results on GLUE show that our approach can reduce latency by 65% without sacrificing performance. By using only two-layer transformer calculations, we can still maintain 95% accuracy of BERT.

## introduction pre-trained language models, such as elmo @xcite , gpt @xcite , bert @xcite , xlnet @xcite , and roberta @xcite , have yielded significant improvements to nlp tasks. despite the gain in accuracy, these models have significant demands in computation and inference time, limiting their use in resource-constrained or latencysensitive applications. therefore, it is desirable to reduce the computational overhead of these models while retaining acceptable accuracy. knowledge distillation @xcite ) facilitates the transfer of knowledge embedded in pre-trained language models into small student models @xcite @xcite , which usually reduces the redundant parameters of bert in a uniform manner. early exit mechanisms @xcite , 2020; @xcite then use an adaptive number of transformer layers during inference, aiming to reduce redundant calculations from the highest few layers. however, since they build the sequence representation from scratch for each forward pass, they require a certain number of lower layers to capture basic syntactic and semantic information, making it difficult to further reduce inference costs. this naturally raises a question: can we reduce the computation at the lower transformer layers? in this paper, we propose skipbert, a novel scheme that skips the computation at the shallow transformer layers of bert. as revealed by @xcite ; @xcite , the lower layers of bert mainly focus on short-distance context, while the higher layers are able to capture long-range dependencies. therefore, it is reasonable to assume that, at lower layers, even if distant tokens are masked, the representation for each token will not vary dramatically. here, by sweeping over the input text, we get short chunks (n-grams) and use their representations to approximate the hidden states of bert's lower layers. we then precompute and store representations of text chunks in a precomputed lookup table (plot). thus, during inference we only need to access plot to get the representations of short chunks, which is inexpen-sive compared with transformer computation. fig. 1 compares the inference procedure between vanilla bert and our proposed skipbert. in bert, the input text needs to be processed by a large number of transformer layers in turn, leading to high latency in inference. in comparison, skipbert precomputes the hidden states of lower transformer layers, which are accessed via table lookups, rather than computed in inference-time. moreover, skipbert exhibits effective compatibility with early exit mechanisms: since the initial sequence representation in our work is partially contextualized (thanks to plot) rather than individual word embeddings, skipbert allows exiting from a relatively earlier layer than typical bert variants, while maintaining good accuracy. we empirically verify this in section 4.5. therefore, our approach can skip the calculations of lower and higher layers for the same input, thereby further improving the inference speed. our contributions are listed as follows: compared with bert, skipbert is capable of accelerating inference by up to 65% without compromising glue score, or accelerating by 82% while retaining 95% accuracy. ## related work knowledge distillation @xcite provides an effective way to transfer the knowledge embedded in a teacher network to a student network. the student network is usually more lightweight than the teacher network and thus more computationally efficient. the student network can be structurally identical to the teacher but contains fewer layers or hidden units, e.g. bert-pkd @xcite , distilbert @xcite , tiny-bert @xcite , minilm @xcite , and bert-emd @xcite . meanwhile, some work adopts specifically designed networks, e.g. @xcite and mobilebert @xcite , to reduce the computation per layer. transformer layer transformer layer transformer layer chunk aggregation yes no table lookup ... ... input text text chunks ... ... prediction prediction prediction exit? exit? exit? yes ram ssd plot mmap figure 2: the overview of our system during inference. plot stores hidden states of the local transformer layers. the global transformer layers are enhanced with early exit mechanisms. input-adaptive inference allows models to choose different computational paths according to the input during inference. in this way, simpler input samples usually require less calculation to make predictions. recently, deebert @xcite adapts confidence-based branchynet @xcite , which uses entropy as an early-exit criterion. fastbert @xcite uses self-distillation to train the branch classifiers. righttool @xcite leverages the same early-exit criterion as in the shallow-deep network @xcite , i.e., softmax scores of predictions. pabee @xcite stops inference when the intermediate predictions of the internal classifiers remain unchanged consecutively. precomputation has also been studied in information retrieval, where documents are assumed to be stored at local database so their representation can be precomputed @xcite . however, this method may not be suitable for other tasks where the input text is unknown before inference. ## model during training, skipbert consists of two groups of transformer layers, local transformer layers for encoding short-distance context, and global transformer layers for leveraging the full context. once pre-training finishes, our approach will replace local transformer layers with plot, which stores the hidden states of local transformer layers; we also enhance global transformer layers with early exit mechanisms to further accelerate inference speed. fig. 2 presents the overview of our system. ## conclusion in this paper, we proposed skipbert, a straightforward yet effective approach to skip the computation of bert's shallow layers. we used representations of short text chunks to approximate bert's shallow representation, and stored them in plot for fast retrieval during inference. empirical results showed that skipbert could achieve performance comparable to bert while significantly reducing inference time. in the future, we would like to leverage discontinuous text chunks to further improve the accuracy and inference speed. we will also try to reduce storage requirements with appropriate pruning and compression strategies. sst-2, mrpc, sts-b, and rte but not as useful for other tasks. specifically, we train additional 10 epochs with no prediction loss for cola and sts-b, and 1 epoch for sst-2, mrpc, and rte.