Fine- and Coarse-Granularity Hybrid Self-Attention for Efficient BERT

Our strategy of scoring the informativeness of tokens is based on the self-attention map. Concretely, taking a single token vector $x_{i}$ as an example, its attention head $x_{i}^{(t)}$ is updated by: $x_{i}^{(t)}=\sum_{j=1}^{n}a_{i,j}x_{j}^{(t)}$ (Eq. 2). $a_{i,j}$ is an element in attention map $A$. Therefore, $a_{i,j}$ represents the information contribution from token vector $x_{j}$ to $x_{i}$ over $\text{head}_{t}$. Intuitively, we define the informativeness of a token by accumulating along the columns of attention map $A$:

The overall informativeness of $x_{j}$ is defined as the average over the heads:

We next analyze some properties of defined informativeness in BERT-base. The first sub-figure in Figure 1 displays the normalized variance and standard deviation of informativeness of layers from 1 to 12 on RTE (classification dataset), which supports the phenomenon that the informativeness distributions at the bottom layers are relatively uniform and the top layers are volatile. The last five sub-figures further present the informativeness distributions of some BERT-base layers, where the first token is [CLS] and its representations are used for the final prediction. We can see that as the layers deepen, the informativeness is progressively concentrated on two tokens. This means that maintaining full-length token-level representations for the classification tasks may be redundant.

A straightforward approach for reducing the sequence length of self-attention is to maintain the informative tokens and prune the rest. We argue that this approach is effortless but encounters the risk of information loss, especially for lower layers.

Instead of pruning, we propose to process the uninformative tokens with more efficient units. Figure 2 shows the architecture of FCA layer, which inserts a granularity hybrid sub-layer after MHA. At each layer, it first divides tokens into informative and uninformative ones based on their assigned informativeness. The CLS token is always divided into informative part as it is used to derive the final prediction.

Let $x_{cls}^{(l)}\oplus X^{(l)}$ be the sequence of token vectors input to $l\mbox{-}$th layer, where $X^{(l)}=[x_{1}^{(l)},...,x_{n}^{(l)}]$ and $n$ is the sequence length of $X^{(l)}$. We gather the token vectors from $X^{(l)}$ with the top-$k$ informativeness to form the informative sequence $X_{in}^{(l)}$ and the rest vectors to form the uninformative sequence $X_{un}^{(l)}$, where $X_{in}^{(l)}\in\mathbb{R}^{k\times d_{h}}$ and $X_{un}^{(l)}\in\mathbb{R}^{(n\mbox{-}k)\times d_{h}}$. The length of the uninformative sequence is reduced by performing certain type of aggregating operations along the sequence dimension, such as average pooling:

or weighted average pooling with informativeness as weights:

where $x^{(l)}_{un}$ is the token vector in $X_{un}^{(l)}$. The aggregated sequence $X_{un}^{\prime(l)}\in\mathbb{R}^{k^{\prime}\times d_{h}}$ and $k^{\prime}$ is a fixed parameter. After hybrid layer, token sequence is updated to [$x_{cls}^{(l)}\oplus X_{in}^{(l)}\oplus X_{un}^{\prime(l)}$] and sequence length is shortened by $n\mbox{-}k\mbox{-}k^{\prime}$. Therefore, in addition to the following layers, the computation cost of FFN in $l\mbox{-}$th layer is reduced as well. It should be noted that the relative position of uninformative tokens should be preserved, which contains their contextual features to a certain extent and they can be captured by aggregating operations.

The parameter $k$ is learnable and progressively shortened. Inspired by  Goyal et al. (2020), we train $n$ learnable parameters to determine the configuration of $k$, denoted $R=[r_{1},...,r_{n}]$. The parameters are constrained to be in the range, i.e., $r_{i}\in[0;1]$ and added after MHA sub-layer. Given a token vector $x_{i}$ output by MHA, it is modified by:

where $pos(x_{i})$ is the sorted position of $x_{i}$ over informativeness. Intuitively, the parameter $r_{i}$ represents the extent to which the informativeness of the token at $i\mbox{-}$th position is retained. Then, for the $l\mbox{-}$th layer, we obtain the configuration of $k_{l}$ from the sum of the above parameters, i.e.,

Our experiments are mainly conducted on GLUE (General Language Understanding Evaluation) ²²2https://gluebenchmark.com/  Wang et al. (2018) and RACE Lai et al. (2017) datasets. GLUE benchmark covers four tasks: Linguistic Acceptability, Sentiment Classification, Natural Language Inference, and Paraphrase Similarity Matching. RACE is the Machine Reading Comprehension dataset.

For experiments on RACE, we denote the input passage as $P$, the question as $q$, and the four answers as $\{a_{1},a_{2},a_{3},a_{4}\}$. We concatenate passage, question and each answer as a input sequence $[\text{CLS}]P[\text{SEP}]q[\text{SEP}]a_{i}[\text{SEP}]$, where $[\text{CLS}]$ and $[\text{SEP}]$ are the special tokens used in the original BERT. The representation of $[\text{CLS}]$ is treated as the single logit value for each $a_{i}$. Then, a softmax layer is placed on top of these four logits to obtain the normalized probability of each answer, which is used to compute the cross-entropy loss.

The input length of BERT is set to 512 by default. However, the instances in these datasets are relatively short, rarely reaching 512. If we keep the default length settings, most of the input tokens are [PAD] tokens. In this way, our model can easily save computational resources by discriminating [PAD] tokens as the uninformative ones, which is meaningless. To avoid this, we constrained the length of the datasets. The statistic information of the datasets is summarized in Table  1.

For accuracy evaluation, we adopt Matthew’s Correlation for CoLA, F1-score for QQP, and Accuracy for the rest datasets. For efficiency evaluation, we use the number of floating operations (FLOPs) to measure the inference efficiency, as it is agnostic to the choice of the underlying hardware.

We compare our model with both distillation and pruning methods. Distillation methods contain four models DistilBERT Sanh et al. (2019), BERT-PKD Sun et al. (2019), Tiny-BERT Jiao et al. (2019), and Mobile-BERT Sun et al. (2020). All four models are distillation from BERT-base and have the same structure (6 transformer layers, 12 attention heads, dimension of the hidden vectors is 768). Pruning methods contain FLOP Wang et al. (2020), SNIP Lin et al. (2020), and PoWER-BERT Goyal et al. (2020). PoWER-BERT Goyal et al. (2020) is the state-of-the-art pruning method which reduces sequence length by eliminating word-vectors. To make fair comparisons, we set the length of our informative tokens at each layer same to the sequence length of PoWER-BERT.

We deploy BERT-base as the standard model in which transformer layers $L\mbox{=}12$, hidden size $d_{h}\mbox{=}512$, and number of heads $h\mbox{=}12$. All models are trained with 3 epochs. The batch size is selected in list 16,32,64. The model is optimized using Adam Kingma and Ba (2014) with learning rate in range [2e-5,6e-5] for the BERT parameters $\Theta$, [1e-3,3e-3] for $R$. Hyper-parameter  $\lambda$ that controls the trade-off between accuracy and FLOPs is set in range [1e-3,7e-3]. We conducted experiments with a V100 GPU. The FLOPs for our model and the baselines were calculated with Tensorflow and batch size=1. The detailed hyper-parameters setting for each dataset are provided in the Appendix.

Table 3 and Table 2 display the accuracy and inference FLOPs of each model on GLUE benchmark respectively. As the FLOPs of PoWER-BERT is almost the same as that of FCA-BERT and the number of coarse units has little affect on FLOPs, Table 2 only lists the FLOPs of FCA-BERT.

Among our models, the weighted average pooling operation raises the better performance than the average pooling operation. The number of coarse units contributes the model accuracy for both two operations, especially for pooling operation. This is reasonable as when the number of coarse units increases, the information stored in each FCA gradually approaches the standard BERT. But overmuch coarse units grow FLOPs. Therefore, it is necessary to balance impact on FLOPS and performance brought by the coarse units.

We next compare our model to the SOTA pruning model PoWER-BERT. Their acceleration effects are comparable and we focus on comparing their accuracy. The results on Table 3 show that our models achieve better accuracy than PoWER-BERT on all datasets. This is because PoWER-BERT discards the computing units, which inevitably causes information loss. Instead of pruning, FCA layer stockpiles the information of uninformative tokens in a coarse fashion (aggregating operations). Moreover, we noticed that coarse units are not always classified as uninformative. In other words, they sometimes participate in the calculation of self-attention as informative tokens. This shows the total informativeness contained in uninformative tokens can not be directly negligible and can be automatically learned by self-attention.

In order to visually demonstrate the advantages of our model, Figure 3 draws curves of trade-off between accuracy and efficiency on three datasets. The results of FCA-BERT and PoWER-BERT are obtained by tuning the hyper-parameter $\lambda$. For DistilBERT, the points correspond to the distillation version with 4 and 6 Transformer layers. It can be seen that with the decrease of FLOPs, (1) PoWER-BERT and our model outperform Distil-BERT by a large margin; (2) our model exhibits the superiority over all the prior methods consistently; (3) more importantly, the advantage of our model over PoWER-BERT gradually becomes apparent. This is because PoWER-BERT prunes plenty of computation units to save FLOPs, which results in the dilemma of information loss. In contrast, our model preserves all information to a certain extent.

In this section, we explore that can we not differentiate between tokens and perform the average pooling on all tokens to reduce the computation cost. To make fair comparisons, we set the length of pooled sequence at each layer equal to the FCA-BERT-Pool₅. The results show that pooling all tokens decreases the model accuracy from 75.0 to 73.8. This is because the pooling operation weakens the semantic features learned by the informative tokens, which are often decisive for the final prediction. On the contrary, our model does not conduct pooling on informative tokens and instead delegates the burden of saving computational overhead to uninformative tokens. And this does not cause serious damage to the representative features learned by the model.

In this section, we further investigate the extent to which these compressed models can retain the essential information of the original BERT. Concretely, we adopt the Euclidean distance of the CLS representation between BERT and the compressed models as the evaluation metric, which is proportional to the information loss caused by model compression, formally:

where $M$ is the number of the instances in corresponding dataset. Table 4 shows the distance of baselines and our models with standard BERT. Combining the results in Table 3, it can be found that the distance is consistent with the test accuracy. Large distance leads to low accuracy and vice versa. This provides an inspiration, that is, we can add a distance regulation term to the objective function to forcibly shorten the distance between the compression model and the original BERT, i.e.,

However, the experimental results show that the accuracy has not been significantly improved. This may be because the information learned by the compressed model has reached the limit of approaching the BERT, and the regulation term can not further improve the potential of the compressed model.