Pyramid-BERT: Reducing Complexity via Successive Core-set based Token Selection

Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 8798 - 8817 May 22-27, 2022 c⃝2022 Association for Computational Linguistics Pyramid-BERT: Reducing Complexity via Successive Core-set based Token Selection Xin Huang Amazon AWS, USA xinxh@amazon.com Ashish Khetan Amazon AWS, USA khetan@amazon.com Rene Bidart University of Waterloo, Canada rbbidart@uwaterloo.ca Zohar Karnin Amazon AWS, Israel zkarnin@amazon.com Abstract Transformer-based language models such as BERT (Devlin et al., 2018) have achieved the state-of-the-art performance on various NLP tasks, but are computationally prohibitive. A recent line of works use various heuristics to successively shorten sequence length while transforming tokens through encoders, in tasks such as classification and ranking that re- quire a single token embedding for predic- tion. We present a novel solution to this problem, called Pyramid-BERT where we re- place previously used heuristics with a core- set based token selection method justified by theoretical results. The core-set based token selection technique allows us to avoid expen- sive pre-training, gives a space-efficient fine tuning, and thus makes it suitable to handle longer sequence lengths. We provide extensive experiments establishing advantages of pyra- mid BERT over several baselines and exist- ing works on the GLUE benchmarks and Long Range Arena (Tay et al., 2020) datasets. 1 Introduction Transformers (Vaswani et al., 2017) have gradually become a key component for many state-of-the- art natural language representation models. A re- cent Transformer based model BERT (Devlin et al., 2018), and its variations, achieved the state-of-the- art results on various natural language processing tasks, including machine translation (Wang et al., 2019a; Liu et al., 2020), question-answering (De- vlin et al., 2018; Yang et al., 2019), text classifica- tion (Goyal et al., 2020; Xu et al., 2019), semantic role labeling (Strubell et al., 2018), and so on. How- ever, it takes substantial computational resources to pre-train, fine-tune, or infer such models. The com- plexity of Transformers is mainly due to a pipeline of encoders, each of which contains a multi-head self-attention layer. The self-attention operation scales quadratically with the input sequence length, which is a bottleneck especially for long-sequence data. Given this challenge, intensive efforts have been focused on compressing and accelerating Trans- formers to reduce the cost of pre-training and fine- tuning. This work is particularly inspired by the sequence-level NLP tasks such as text classifica- tion and ranking. The state-of-the-art Transformer models for these tasks utilize a single embedding from the top encoder layer, such as the CLS to- ken, for prediction. Under such regime, retaining full-length sequence till the last encoder creates unnecessary complexity. We follow a line of works detailed in Section 2 that aim to gradually reduce the sequence length in the pipeline of encoders. On a high level, the existing works have two compo- nents: Select : a mechanism in charge of reducing the sequence length, either by pruning or pooling, Train-Select : a training or fine-tuning procedure dedicated to this mechanism. Our main contribution is a novel solution for Se- lect , motivated from …

Content truncated to 500 words