SkipBERT: Efficient Inference with Shallow Layer Skipping

Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 7287 - 7301 May 22-27, 2022 c⃝2022 Association for Computational Linguistics SkipBERT: Efficient Inference with Shallow Layer Skipping Jue Wang1, Ke Chen1, Gang Chen1,2, Lidan Shou1,2∗ and Julian McAuley3 1College of Computer Science and Technology, Zhejiang University 2State Key Laboratory of CAD&CG, Zhejiang University 3University of California, San Diego {zjuwangjue,chenk,cg,should}@zju.edu.cn, jmcauley@ucsd.edu Abstract In this paper, we propose SkipBERT to accel- erate BERT inference by skipping the compu- tation of shallow layers. To achieve this, our approach encodes small text chunks into inde- pendent representations, which are then mate- rialized to approximate the shallow representa- tion of BERT. Since the use of such approx- imation is inexpensive compared with trans- former calculations, we leverage it to replace the shallow layers of BERT to skip their run- time overhead. With off-the-shelf early exit mechanisms, we also skip redundant computa- tion from the highest few layers to further im- prove 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.1 1 Introduction Pre-trained language models, such as ELMo (Pe- ters et al., 2018), GPT (Radford et al., 2018), BERT (Devlin et al., 2019), XLNet (Yang et al., 2019), and RoBERTa (Liu et al., 2019), have yielded sig- nificant improvements to NLP tasks. Despite the gain in accuracy, these models have significant de- mands in computation and inference time, limit- ing their use in resource-constrained or latency- sensitive applications. Therefore, it is desirable to reduce the computational overhead of these models while retaining acceptable accuracy. Knowledge distillation (KD, Hinton et al. 2015) facilitates the transfer of knowledge embedded in pre-trained language models into small student models (Sanh et al., 2019; Sun et al., 2019; Jiao et al., 2020), which usually reduces the redundant parameters of BERT in a uniform manner. Early exit mechanisms (Xin et al., 2020; Zhou et al., ∗ Corresponding author 1Source code is available at https://github.com/ LorrinWWW/SkipBERT. (a) BERT (b) SkipBERT Output Output Transformer Layer Transformer Layer Transformer Layer Transformer Layer Transformer Layer Transformer Layer Transformer Layer Transformer Layer Transformer Layer Transformer Layer Transformer Layer Transformer Layer Skipped by searching precomputed lookup tables (PLOT) Figure 1: Comparison of inference between BERT and SkipBERT. The computation of shallow (lower) layers of SkipBERT are skipped by searching PLOT. 2020; Liu et al., 2020) then use an adaptive num- ber 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. …

Content truncated to 500 words