BERT Loses Patience: Fast and Robust Inference with Early Exit

TL;DR

The paper introduces Patience-based Early Exit (PABEE), a plug-in mechanism that appends an internal classifier after every layer of a pretrained language model and halts inference when intermediate predictions stabilize for a predefined patience t. By leveraging a cross-layer consensus among multiple classifiers, PABEE aims to reduce latency while maintaining or improving accuracy and robustness, addressing the overthinking phenomenon in deep PLMs. The authors provide theoretical conditions for potential accuracy gains, and validate the approach on GLUE with ALBERT, showing improved speed-accuracy trade-offs and enhanced adversarial robustness compared with prediction-score based exits and static compression baselines. They also demonstrate that PABEE scales with deeper models and can adapt to different resource constraints, with potential applicability beyond NLP, such as image classification.

Abstract

In this paper, we propose Patience-based Early Exit, a straightforward yet effective inference method that can be used as a plug-and-play technique to simultaneously improve the efficiency and robustness of a pretrained language model (PLM). To achieve this, our approach couples an internal-classifier with each layer of a PLM and dynamically stops inference when the intermediate predictions of the internal classifiers remain unchanged for a pre-defined number of steps. Our approach improves inference efficiency as it allows the model to make a prediction with fewer layers. Meanwhile, experimental results with an ALBERT model show that our method can improve the accuracy and robustness of the model by preventing it from overthinking and exploiting multiple classifiers for prediction, yielding a better accuracy-speed trade-off compared to existing early exit methods.

Figures (5)

• Figure 1: Comparison between Shallow-Deep Net, a prediction score based early exit (threshold is set to $0.9$), and our Patience-based Early Exit (patience $t=1$). A classifier is denoted by $C_i$, and $n$ is the number of layers in a model. In this figure, Shallow-Deep incorrectly exits based on the prediction score while PABEE considers multiple classifiers and exits with a correct prediction.
• Figure 2: Analogy between overfitting in training and overthinking in inference. (a) In training, the error rate keeps going down on the training set but goes up later on the development set. (b) We insert a classifier after every layer. Similarly, the predicted entropy keeps dropping when more layers are added to inference but the error rate goes up after 10 layers. The results are obtained with ALBERT-base on MRPC.
• Figure 3: Speed-accuracy curves of BranchyNet, Shallow-Deep and PABEE on MNLI and SST-2 with ALBERT-base model.
• Figure 4: Accuracy scores and speed-up ratios under different patience with ALBERT-base model. The baseline is denoted with gray dash lines.
• Figure 5: Monte Carlo simulation of per PABEE classifier's accuracy vs. the original inference accuracy under different patience settings.

Theorems & Definitions (2)

• Theorem 1
• Proof B.1