ELBERT: FAST ALBERT WITH CONFIDENCE-WINDOW BASED EARLY EXIT

As shown in Fig. 1, ELBERT uses ALBERT as the backbone model, which is composed of an encoder and a classifier. Additionally, ELBERT is designed to put a early exit decision after each propagation processed by the encoder and the classifier.

To fit the early exit mechanism in inference, the losses of inputs exiting at different depths of ELBERT are calculated during the training. For classification, the early exit loss at the $i$-th layer $\mathcal{L}_{i}$ is calculated with Cross-Entropy

$$\mathcal{L}_{i}=-\sum_{c\in C}\left[\mathbb{I}\left[\mathbf{\hat{y}}_{i}=c\right]\cdot\log P\left(\mathbf{\hat{y}}_{i}=c\mid\mathbf{h}_{i}\right)\right],$$

(1)

where $c$ and $C$ denote one class label and the set of class labels, respectively. The common practice is to simply add up $\mathcal{L}_{i}$ as the total loss $\mathcal{L}$ [21, 23]. For better training under various combinations of losses, we assign a trainable variable $t_{i}$ with an initial value of 4 to each layer, inspired by Wang et al.[25]. The weight of the $i$-th layer $w_{i}$ is calculated by

$$w_{i}=\left\{\begin{array}[]{ll}\sigma(t_{i})&0<i\leq M-1\\ M-\sum_{i=1}^{M-1}\sigma(t_{i})&i=M\end{array}\right.,$$

(2)

where $M$ denotes the depth of ELBERT and $\sigma(\cdot)$ denotes sigmoid funciton $\sigma\left(t_{i}\right)=1/\left(1+\exp({-t_{i})}\right)$. Then the total loss $\mathcal{L}$ is calculated by a weighted sum

$$\mathcal{L}=\sum_{i=1}^{M}w_{i}\cdot\mathcal{L}_{i}.$$

(3)

In this way, the cases that inputs may exit at different depths are well considered, which helps to bridge the gap between training and inference of ELBERT.

ELBERT first introduces a two-stage early exit mechanism, which focuses on both the intermediate state and the historical trend of classifier output to decide whether an early exit of computation is needed.

Formally, the input $\mathbf{x}$ goes through the encoder iteratively. The hidden state $\mathbf{h}_{i}$ after the $i$-th forward propagation of the encoder is calculated by

$$\mathbf{h}_{i}=\left\{\begin{array}[]{ll}Encoder(\mathbf{h}_{i-1})&0<i\leq M\\ Embedding(\mathbf{x})&i=0\end{array}\right..$$

(4)

After each forward propagation in the encoder, the hidden state $\mathbf{h}_{i}$ is sent to the classifier that outputs a prediction probability distribution $p_{i}=Classifier(\mathbf{h}_{i})$ via fully-connected layer and softmax function for classification. Then we can get the predicted label $\mathbf{\hat{y}}_{i}=argmax(p_{i})$.

The first stage of the early exit focuses on confidence, or intermediate state, of the classifier. Given a probability distribution $p_{i}$, we take its normalized entropy as the $Puzzlement$ of the current classifier, which is calculated by

$$Puzzlement(i)=\frac{\sum_{j=1}^{C}p_{i}(j)\log p_{i}(j)}{\log(1/C)},$$

(5)

where $C$ denotes the number of labeled classes. The model will stop the inference in advance and take $\mathbf{\hat{y}}_{i}$ as the final prediction to skip further computations when $Puzzlement(i)<{\delta}$, where $\delta$ is a user-defined threshold. When a faster model is needed and some accuracy degradation is tolerable, we can set a higher $\delta$, while the opposite situation leads to a lower $\delta$.

The second stage traces the historical trend of the classifier output in a time window, whose size $N$ is defined based on user demands. We propose three criteria for triggering the second stage early exit in a time window: 1) The prediction probability $p_{i}$ of a certain class varies monotonically. 2) The range of $max(p_{i})$ is less than a set value. 3) The predicted label $\mathbf{\hat{y}}_{i}$ stays the same. Experimental results show that the first criterion outperforms others. In subsequent experiments, we will use the first criterion for the second stage by default, and the window size $N$ is set to 8.

Usually, we prefer the moment when we get enough confidence. Only when the first stage condition isn’t satisfied will we consider the second stage early exit.

We select three baselines. 1) Original model: We choose ALBERT-large (depth=24). 2) Plain compression: We evaluate several models with smaller depths based on ALBERT-large. 3) Early exit approach: The methods in DeeBERT and FastBERT are applied to ALBERT for comparison.

To test the generalization ability of ELBERT, widely used GLUE benchmark [26], AG-news [27] and IMDB [28] are evaluated in our experiments. These datasets include various NLP tasks such as Natural Language Inference, Sentiment Analysis and News Classification.

Training For GLUE we use the corresponding hyperparameters in ALBERT original paper for a fair comparison, while for other datasets, we use a default learning rate of 3e-5 and a batch size of 32.
Inference In practical scenes, the user requests often arrive one by one. Our batch size of inference is set to 1, following prior work[21, 19]. The experiments are done on an NVIDIA 2080Ti GPU.

Efficient inference acceleration We evaluate ELBERT on the above datasets and report the median of 5 runs in Fig. 2 and Fig. 3. The curves are drawn by interpolating several points that correspond to different $\delta$, which changes from 0.1 to 1.0 with a step size of 0.1 in the first stage of early exit. For all datasets, ELBERT achieves at least two times inference speedup while keeping or even improving the accuracy. When a little accuracy degradation is tolerable, the inference acceleration ratio can be up to ten. This demonstrates ELBERT’s superiority of inference acceleration.

Task-related trends An interesting observation is that there are different trends of curves in Fig. 2 on different kinds of tasks. For News Classification (AG), ELBERT gets the best acceleration performance, followed by Sentiment Analysis (SST-2, IMDB), the curves of which drop a little faster. NLI (QNLI, RTE) is the case with the lowest performance. This indicates that different tasks may have different internal characteristics and acceleration difficulty. Early exit may help us understand tasks better. We will do some discussions about this in Section 4.5.
Flexible and better accuracy-speed tradeoffs We compare different models on several datasets. The results are shown in Fig. 3, where the red star-shaped points represent different models obtained by plain compression. Our first observation is that ELBERT significantly outperforms plain compression models. Also, compared to other early exit based methods, ELBERT obtains higher accuracy than both DeeBERT and FastBERT under the same computation cost, which shows ELBERT’s great advantages over other approaches.

To visualize the decision-making process of the early exit in ELBERT, we make some changes to BertViz[29], a tool for visualizing attention in Transformer. We use the attention-scores of each layer to get the cumulative attention-scores, which allows us to see the attention relationships between tokens clearly as the input passes through different depths of ELBERT. Since ELBERT only takes the [cls] token as the representation of one sentence to do classification, we only show the cumulative attention-scores of [cls] to other tokens in the figures. We take SST, a Sentiment Analysis dataset for example, and find two main patterns of early exit. ^††footnotetext: $\dagger$The results are based on our implementation on ALBERT-large.
Simple input, simple exit For the most common inputs without emotional turns or negative words, as shown in Fig. 4, the attention of [cls] to the emotional keywords (i.e., hampered) tends to increase monotonously. Early exit is triggered when such attention exceeds a certain limit determined by the $\delta$, thus reducing redundant computations. Actually the prediction remains unchanged after the exit layer 11.
Mitigating overthinking As Fig. 3 shows, ELBERT sometimes achieves even higher accuracy than that of the original model, indicating that the early exit mechanism corrects some wrong predictions of the final layer. As shown in Fig. 5, the model first pays attention to the commendatory word (i.e., benign) and predicts $Positive$. Next, an irrelevant negative word (i.e., rarely) is noticed, seen as the negation of commendatory words. Then the model predicts $Negative$. This is exactly an example of overthinking. In correct cases, the negation and the corresponding word are often noticed simultaneously.

The above patterns demonstrate that ELBERT’s prediction for classification is mainly determined by some key words, such as negatives and those words with strong emotional orientation. The early exit mechanism helps to establish appropriate attention to these words, which enables the model to exit from simple inputs in advance and avoid overthinking.