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

Abstract: Transformer-based language models such as BERT (CITATION) 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 require a single token embedding for prediction. We present a novel solution to this problem, called Pyramid-BERT where we replace 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 expensive 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 pyramid BERT over several baselines and existing works on the GLUE benchmarks and Long Range Arena (CITATION) datasets.

## introduction transformers @xcite have gradually become a key component for many state-of-theart natural language representation models. a recent transformer based model bert @xcite , and its variations, achieved the state-of-theart results on various natural language processing tasks, including machine translation @xcite , question-answering @xcite ), text classification @xcite , semantic role labeling @xcite , and so on. however, it takes substantial computational resources to pre-train, fine-tune, or infer such models. the complexity 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 transformers to reduce the cost of pre-training and finetuning. this work is particularly inspired by the sequence-level nlp tasks such as text classification 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 token, 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 components: 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 select , motivated from the following observation: as we consider the token representations in top layers, they have increasing redundancy among themselves. we provide a quantitative study demonstrating this in figure 1 . a collection of tokens with high redundancy can intuitively be represented by a small core-set, composed of a subset of the tokens. inspired by this, our solution for select is based on the idea of core-sets. we provide a theoretically motivated approach that becomes more effective as the redundancy in the representation increases. this concept separates our work from previous art that provide heuristic techniques for sequence length reduction, or approaches that require expensive training. all of these can in fact be shown to fail in toy examples with high redundancy, e.g. the same representation duplicated multiple times. for train-select there is some variety in previous art. some require a full pre-training procedure, and others require a fine-tuning procedure that works on the full uncompressed model, meaning one that keeps all tokens until the final encoder layer. the high quality of our solution to select allows us to simply avoid this additional training phase altogether. the result of this simplification is quite impactful. we obtain a speedup and memory reduction not only to the inference but to the training process itself. this makes it possible to use standard hardware (and training scripts) in the training procedure even for very long sequences. in section 6 we provide an empirical comparison of our technique with sota alternatives and show that it is superior in eliminating redundancy among the tokens, and thus greatly improves the final classification accuracy. in particular, our experiments on the glue benchmarks show that our method achieves up to 3-3.5x inference speedup while maintaining an accuracy (mean value over all glue datasets) drop of 1.5%, whereas the best baseline suffers a 2.5% drop. we show that our method can be combined with long-text transformers such as big bird @xcite and performers @xcite to further alleviate the quadratic space complexity of the selfattention mechanism. empirical experiments on the long range arena (lra) @xcite show that our model achieves a better trade-off between space complexity reduction and accuracy in comparison to its competitors. in particular, when working with performers and reducing the space complexity by 70%, while baselines suffer a drop in accuracy of 4% or more, our technique actually improves the accuracy as it acts as a regularizer. concluding, our paper provides a novel, theoretically justified technique for sequence length reduction, achieving a speedup and memory reduction for both training and inference of transformers, while suffering significantly less in terms of predictive performance when compared to other existing techniques. our methods are vetted via thorough empirical studies comparing it against sota methods and examining its different components. ## related works there have been a number of interesting attempts, that were aimed at model compression for transformers, which can be broadly categorized into three directions. the first line of work focus on the redundancy of model parameters. structure pruning @xcite @xcite @xcite , which removes coherent groups of weights to preserve the original structure of the network, is one common strategy. in addition, various types of distillation techniques @xcite @xcite have been proposed to remove encoders by training a compact transformer to reproduce the output of a larger one. other strategies include weight quantization @xcite @xcite and weight sharing @xcite . the second line of work focus on reducing the quadratic operation of the self-attention matrices. the quadratic time and space complexity of the attention mechanism with respect to the input sequence serves the main efficiency bottleneck of transformers, and thus is prohibitively expensive for training long-sequence data. one popular approach is to sparsify the self-attention operation by restricting each token to only attend a subset of tokens @xcite @xcite @xcite @xcite . in particular, the most recent big bird @xcite and longformer @xcite introduce sparse models which scale linearly with the input sequence. another popular approach @xcite is to approximate the self-attention to reduce its quadratic complexity to linear, where the most recent performers (choromanski et al., 2020) provides an unbiased linear estimation of the attention matrices. the third line, which our work lies in, focus on the redundancy in maintaining a full-length sequence of token-level representation across all encoders @xcite @xcite . this work is particularly inspired by the sequence-level nlp tasks such as text classification and ranking, where the state-of-the-art approach utilizes a single embedding from the top encoder layer of a transformer, such as the cls token, for prediction. under such regime, retaining the full-length sequence till the last encoder creates unnecessary complexity. @xcite applied a simplest strided mean pooling to each sliding window of the sequence to gradually compress tokens. the paper proposed a specialized pre-training procedure for this process that achieves a good balance between speedup and accuracy drop. in comparison we aim to avoid the costly pre-training step and focus on a method that requires only a lightweight fine-tuning. @xcite define a centroid attention operator that given centroids, computes their embeddings via an attention mechanism. these centroids are chosen either as a random subset of the original tokens, or via strided mean pooling. this results in a similar technique to that of @xcite , in terms of having a naive sequence length reduction method. @xcite provide a variant of length 2 strided mean pooling where instead of taking the unweighted average of each pair, they provide learnable weights via a differentiable linear function. in our experiment we did not compare our methods with @xcite since both worked on a limited (nonstandard) collection of datasets and at the time of writing this paper, did not provide code allowing us to reproduce their result. since our paper is focused on techniques to select a subset of tokens and these papers use either random sampling or pooling, we do not believe a thorough comparison is required. @xcite proposed a reinforcement-learning based technique to rank the tokens, thereby allowing it to remove the least important ones. this rl policy must be trained in an expensive process that requires the full network structure. since we focus on methods to improve both training and inference, we do not include it in our experiments. vision transformers @xcite , which apply various types of pooling techniques to reduce input length, are specially designed for image data and thus non-trivial to compare with our method. early exiting approaches @xcite , which allow samples to exit early based on redundancy, are orthogonal to our technique. @xcite developed an attention based mechanism to progressively eliminate tokens in the intermediate encoders in the fine-tuning, while maintaining the classification accuracy. motivation. the state-of-the-art bert utilizes only the cls token from the top encoder layer for tasks such as classification and ranking. a natural question is: do we need to propagate all the token embeddings through the entire pipeline of encoders when only the top layer cls embedding is used for prediction? in general, yes, since the self-attention transforms all the embeddings together, updating each one by capturing information from all the others. however, if two or more tokens are exact duplicates of each other, ignoring the positional embedding, then one can easily remove the duplicates from the input and modify the self-attention appropriately to get the same cls embedding at the top layer, and hence the same prediction. this would reduce the number of flops carried out in each self-attention layer. in general, input sentences do not contain duplicate tokens. however, a preliminary study of token embeddings show that as the embeddings propagate through the pipeline of encoders, they become more similar to the cls token, figure 1 (top row). a deeper investigation shows that they also become more similar with each other and form clusters among themselves, figure 1 (bottom row). in this work, we exploit these observations, and present pyramid bert, a novel bert architecture, that reduces computational and space complexity of fine-tuning and inference of bert while incurring minimal performance degradation. the pyramid bert works for all the downstream tasks that only use the top layer cls token for prediction, such as classification and ranking. architecture. the pyramid bert successively selects subsets of tokens in each encoder layer to propagate them to the next encoder. an illustration of pyramid bert is shown in figure 2 . it involves two main components over bert: (a) a sequencelength configuration: a monotonically decreasing sequence ℓ = (ℓ 1 , ℓ 2 , • • • , ℓ l ) that specifies the number of tokens to retain in each of the l encoders, for all the input examples. (b) a core-set based token selection methodology which in each j-th encoder, given ℓ j-1 token embeddings from the (j -1)-th encoder selects a subset of it of size ℓ j to propagate them to the next encoder. note that the rest of pyramid bert architecture is same as the bert, and it has exactly the same number of parameters as the bert. in section 4, we provide a theoretical derivation of the core-set based token selection methodology by minimizing an upper bound of successive token selection loss. to computationally perform the core-set based token selection, we provide a greedy k-center algorithm in section 5.1. in section 5.2, we present a simple yet effective approach to select the sequence-length configuration ℓ for a desired time/space complexity reduction. let t denote a token selection algorithm. for an example input x, let the input and output embeddings of the token selection algorithm t at encoder j be s j and s j respectively. the size of the two sets are | s j | = ℓ j-1 , and |s j | = ℓ j . in particular, at each encoder j, the algorithm t selects ℓ j embeddings of the input set s j as the output set s j , and eliminates the remaining ℓ j-1 -ℓ j embeddings in s j . let @xmath0, and @xmath1. given the underlying bert parameters w, the sequence length configuration ℓ, and the classification loss l w , pyramid bert solves the token selection problem by minimizing the population risk as follows: for ease of notation, we write l w as l. the above bound follows immediately from the triangle inequality. for the first two terms in the above bound: the bert training error is a constant for fixed parameters w, and the generalization error of models like bert is known to be small @xcite . therefore, we redefine the token selection problem, equation 1, to minimize the third term, the pyramid bert token selection loss to solve equation 3, we first optimize a slightly different token selection algorithm t * for a model pyramid * bert. in pyramid * the embedding sequence length is not reduced across the encoder layers. let s * j and s * j denote the set of input and output embeddings of t * , where the size of the two sets is equal to the input sequence length n , given an input set s * j in encoder j, the algorithm t * first selects a subset of it of size ℓ j , which is exactly same as s j selected by t in pyramid bert. next, instead of eliminating the remaining n -ℓ j embeddings in s * j as is done by t , the t * replaces them with their nearest embedding in s j to form the output set s * j . the unique embeddings of the output set s * j of pyramid * bert are exactly same as the embeddings of output set s j of pyramid bert, that is unique(s * j ) = s j . we make the following observation. remark 1 a pyramid * bert can be reduced to a pyramid bert with a weighted self attention network. the weighted self attention network weighs attention scores according to the duplicity of the tokens in the embeddings of pyramid * bert. given the above remark, we optimize the token selection problem for pyramid * bert. in the theorem below, we give an upper bound for the token selection loss stated in (3) for pyramid * bert. the proof relies on λ-lipschitz continuity of the encoder and classification layers. theorem 1 if the classification layer is λ c -lipschitz, and for all @xmath2, the encoder e j is λ j -lipschitz continuous for all its parameters, and t * is a token selection algorithm such that the unique elements of the output embedding set unique(s * j ) is a δ-cover of the input embedding set s * j , and the n -ℓ j remaining elements in s * j \ s * j are replaced in s * j by their nearest elements in unique(s * j ), then for all i such that x i is bounded, the following holds: we visualize the concept of δ-cover in figure 3 . the set of red points (i.e., token embeddings in our case) with radius δ covers the entire set of points. theorem 1 suggests that we can bound the token selection loss of algorithm t * for pyramid * bert with the δ-cover core-set token selection. the loss goes to zero as the covering radius δ goes to zero. a proof of the theorem is given in appendix a.2. from remark 1, the optimal choice of token selection for pyramid bert is same as the optimal choice of unique tokens selected in pyramid * bert, up to weighing of self-attention. however, in numerical experiments we found that finetuning pyramid bert with the core-set based token selection performs better than weighing the self-attention. the δ-cover core-set token selection problem is equivalent to the k-center problem (also called min-max facility location problem) @xcite . we explain how we solve the k-center problem using a greedy approximation algorithm in §5.1. ## pyramid bert: algorithm given a pre-trained bert, we create a fine-tuned pyramid bert as follows. for the core-set selection module shown in figure 2 , we implement a k-center-greedy-batch-m algorithm, section 5.1, to approximately select the core-set of embeddings. we fine-tune all the trainable parameters of the pyramid bert for different choices of sequencelength configurations ℓ, according to approach given in section 5.2. we select the optimal configuration ℓ satisfying the required inference speed-up or the space complexity reduction. the selected optimal choice of ℓ is kept fixed during inference. we note that our practical implementation of pyramid bert, proposed here, is not exactly the same as the theoretical token selection algorithm analyzed in the previous section. in appendix a.1, we justify the differences between the two. ## experiments we plug-in our core-set based token selection method and the other competitive methods into encoder layers of a backbone transformer, and after fine-tuning evaluate their performance on a wide range of natural language classification tasks. specifically, we conduct the evaluations on two popular benchmarks: (1) the general language understanding evaluation (glue) benchmark @xcite , and (2) the long rang arena (lra), a collection of challenging long range context tasks @xcite . for the backbone transformer, we choose bert base @xcite for the glue benchmarks, and two state-of-the-art long-text transformers big bird @xcite and performers (choromanski et al., 2020) for the lra. for dataset statistics such as the number of classes and input sequence length, and the details of the backbone transformers, see appendix c. for sequence-length configurations, we generate a set f of 30 sequence-length configurations using equation 5, the details of which are listed in for each token selection method, we show the predictive performance for 1.5x, 2x, 3x, and 3.5x speedup. similarly, we show the predictive performance for 30% and 70% space complexity reductions of the attention layer. the reason we consider the space reduction for the attention layer alone is that its quadratic complexity serves the main bottleneck for long sequences. for details on how to get the performance at different speedup and mathematical formula for computing speedup and space complexity reduction, see appendix c. ## ablation studies we conduct four ablation studies to better study pyramid-bert: (1) performance comparisons for coreset-select with different values of m, the number of centers to add per iteration. (2) justification on the token importance measured by the coresetselect. next, we conduct a study to validate the importance of tokens measured by our coreset-select strategy. we consider a trained bert that has been fine-tuned on a downstream dataset without any sequence length reduction. during inference, given a encoder j and input example consist of a sequence of tokens, we eliminate the k-th most important token measured by the core-set selection method with 1 ≤ @xmath3, and obtain a classification output (prediction label). the classification outputs for all input examples are then compared with those generated by bert without any sequence length reduction. the comparison between the two classification outputs is measured by the mutual information @xcite . larger mutual information indicates more similarity between the two classification outputs. the importance score is specifically the order of tokens (centers) added by the core-set selection method. for a batch of size m tokens that are added at the same time, their importance is determined by the maximum distance between the token and its nearest centers that have already been added. the expectation is that the importance of tokens is negatively correlated with the corresponded mutual information. for an example, reducing the least important token for each input example should make the classification outputs have the least difference from those generated without reducing any token, and thus result in the largest mutual information. a result on sst-2 dataset is presented in figure 4 . the input sequence length l is set as 64, and thus @xmath4, 64]. the encoder index j takes value {1, 3, 6, 9}. since the target variable of sst-2 is a relatively balanced binary value, the largest mutual information, which corresponds to no difference between the classification outputs, is ln(2) ≈ 0.69. for the token selection method, we choose coresetselect with @xmath5. the pattern shown in the figure aligns with our expectation that the importance of tokens is negatively correlated with the mutual information. same pattern is observed for coresetselect-k-1. see figure 5 . next, we study the difference between applying coreset-select at both fine-tuning and inference, and at only inference. finally, we study the position to insert the coreset-select method in the encoder. two choices of position have been considered. the first option is to plug-in the token selection method right after the attention layer and before the feed-forward network, and the second option is to place it at the end of the encoder. the results is shown in ## conclusion we provide pyramid-bert, a theoretically justified technique for sequence length reduction, achieving speedup and memory reduction for both training and inference of transformers, while incurring significantly less accuracy drop than competitive methods. however, this technique can be applied only for classification and ranking tasks which use single embedding from the top layer for prediction. also, our token selection approach requires finetuning the network. an interesting future study would be to eliminate the need of fine-tuning.