PaperlayerView exact revision on arXiv

Source-linked AI summary

BERT Loses Patience: Fast and Robust Inference with Early Exit

Wangchunshu Zhou, Canwen Xu, Tao Ge, Julian McAuley, Ke Xu, Furu Wei

arXiv:2006.04152v3cs.CLcs.LG

TL;DR

Pretrained language models are computationally expensive and can suffer from overthinking during inference. The paper proposes Patience-based Early Exit, which dynamically stops inference when intermediate predictions remain unchanged, and reports simultaneous gains in efficiency, accuracy, and robustness on ALBERT, including 1.57× faster inference.

  • Problem

    Pretrained language models are computationally expensive and may overthink by worsening predictions after earlier layers already suffice.

  • Method

    Patience-based Early Exit attaches internal classifiers to model layers and exits when their intermediate predictions remain unchanged for a predefined patience period.

  • Results

    1.57× faster inference while simultaneously improving ALBERT’s efficiency, accuracy, and adversarial robustness on GLUE.

  • Takeaways & Limitations

    PABEE provides an accuracy-speed trade-off that can be adjusted for different devices and resource constraints by tuning patience without retraining.

  • Takeaways & Limitations

    The theoretical analysis assumes internal classifiers have equal and independent accuracy, an assumption generally unattainable in practice.

Abstract

from arXiv · show

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.

1 Introduction

The paper introduces Patience-based Early Exit (PABEE), which dynamically stops pretrained language model inference using layerwise internal classifiers whose predictions remain unchanged for a preset patience period. It targets the computational costs of deep PLMs while improving efficiency, accuracy, and robustness without retraining for different resource constraints.

  • PLMs contain many layers and parameters, making inference expensive in memory and latency and limiting deployment where speed and computational cost matter.
  • PABEE couples an internal classifier with each PLM layer and stops inference dynamically when intermediate predictions remain unchanged for a predefined number of steps.
  • Tuning the patience hyperparameter adjusts the accuracy–efficiency trade-off for different devices and resource constraints without retraining the model.
  • Experiments on GLUE show that PABEE can simultaneously improve the accuracy and robustness of ALBERT while speeding inference across tasks with trivial additional training resources.

2 Related Work

Related work divides efficient neural-network inference into static approaches, which use fixed computation, and dynamic approaches, which adapt computation to each input. Prior adaptive methods face training, parameter, performance, or task-coverage limitations that motivate PABEE as a distinct early-exit criterion for pretrained language models.

  • Static and Dynamic Approaches: Efficient neural-network methods comprise static approaches with fixed computation and dynamic approaches that select input-dependent computational paths.Static models send every input through the same layers, whereas dynamic methods can use fewer calculations for simpler inputs.
  • Static Approaches: Compact Network Design and Model Compression: Static approaches include compact architectures such as MobileNet, ShuffleNet, EfficientNet, and ALBERT, alongside sparsification, weight quantization, and knowledge distillation.These methods target resource-constrained applications or compress heavy models while retaining a static inference structure.
  • Dynamic Approaches: Input-Adaptive Inference: Dynamic approaches enable input-adaptive inference through mechanisms such as Adaptive Computation Time and entropy-based early-exit methods including BranchyNet.Adaptive Computation Time uses a trainable halting mechanism, while BranchyNet calculates prediction entropy to guide inference.
  • Limitations of Existing Methods: Prediction-probability-based methods can substantially underperform an oracle selecting the smallest sufficient model and support only classification, excluding regression.These limitations reduce their applicability across tasks.
  • PABEE: PABEE differs from recent work by introducing a novel early-exit criterion rather than directly applying existing efficient-inference methods on pretrained language models.The passage presents this distinction as a response to limitations in prior prediction-probability-based approaches.

3 Patience-based Early Exit

PABEE is a plug-and-play early-exit method that attaches internal classifiers to intermediate layers and stops inference when predictions remain stable for a preset patience period. It is motivated by overthinking in inference and is designed to reduce latency while potentially improving accuracy.

  • Motivation: PABEE addresses overthinking by using an early-stopping-like inference mechanism when later-layer predictions become unnecessarily confident but less accurate.The paper reports that prediction entropy decreases as more layers participate, while the model’s actual error rate increases.
  • Inference mechanism: PABEE couples an internal classifier or regressor with each intermediate layer to produce per-layer predictions.For layer Li, the hidden state hi is passed to its internal classifier or regressor to obtain yi.
  • Inference mechanism: Inference exits at layer Lj when the prediction-stability counter reaches patience t; otherwise, the final classifier Cn produces the prediction.For classification, stability is based on unchanged argmax predictions; for regression, it is based on changes smaller than a predefined threshold τ.
  • Training: Training minimizes a weighted average of internal-classifier losses so every inference branch is covered and training weights can reflect relative inference cost.Classification uses cross-entropy loss, while regression uses mean squared error.
  • Theoretical analysis: Theorem 1 states that PABEE improves conventional-inference accuracy when its patience and internal-classifier error conditions satisfy the derived inequality.For n = 12, q = 0.2, and p = 0.1, the condition holds when patience t ≥4, although equal and independent internal-classifier accuracy is generally unattainable in practice.

4 Experiments

Experiments on GLUE show that PABEE improves ALBERT inference speed and performance while retaining higher accuracy than competing early-exit methods at matched speed-ups. Additional results demonstrate versatility across BERT and deeper ALBERT models, favorable training cost, and improved adversarial robustness.

  • GLUE Results: PABEE consistently improves original ALBERT performance while accelerating inference by 1.57×.The paper presents this as improving both speed and performance of a fine-tuned PLM.
  • GLUE Results: PABEE retains higher accuracy than BranchyNet and Shallow-Deep under the same speed-up ratio.The comparison uses speed-accuracy curves with different trade-off hyperparameters.
  • GLUE Results: On BERT, PABEE significantly outperforms methods using knowledge distillation or prediction-probability-based adaptive inference, though it is slightly less accurate than original BERT.In contrast, PABEE improves accuracy when ALBERT is the backbone.
  • Training Cost: PABEE training is no slower, and sometimes slightly faster, than conventional fine-tuning despite adding internal classifiers.The paper attributes this result partly to the additional loss functions of the internal classifiers.
  • Impact of Model Depth: On 24-layer ALBERT-large, PABEE consistently improves accuracy as layers and classifiers increase while producing a larger speed-up ratio.This supports applying PABEE to deeper pretrained language models.
  • Adversarial Robustness: PABEE defends more than 3× as many attacks as original ALBERT on NLI tasks and 2× as many on Yelp sentiment analysis.It also substantially increases the number of queries required by TextFooler and improves robustness over Shallow-Deep.

5 Discussion

PABEE is an efficient inference method that achieves a better accuracy-speed trade-off than existing methods. Experiments and theoretical analysis show simultaneous gains in efficiency, accuracy, and adversarial robustness on ALBERT, while the method currently applies only to single-branch models.

  • Contributions: PABEE yields a better accuracy-speed trade-off than existing inference methods.The paper presents PABEE as a novel efficient inference method.
  • Evaluation: Experiments on GLUE and theoretical analysis verify PABEE’s effectiveness and efficiency.The evaluation uses GLUE, alongside a theoretical analysis of the method.
  • Empirical results: PABEE simultaneously improves efficiency, accuracy, and adversarial robustness over a competitive ALBERT model.The reported improvements are measured against ALBERT.
  • Limitations: PABEE currently works only on models with a single branch, such as ResNet and Transformer.The paper identifies this single-branch restriction as a limitation.

Broader Impact

PABEE’s efficient, plug-in inference can expand mobile and edge applications while reducing energy use and carbon emissions, though its effects on biases encoded in pretrained language models require further study.

  • PABEE can facilitate more applications on mobile and edge computing while reducing energy use and carbon emissions.The paper characterizes PABEE as an efficient inference technique.
  • As a plug-in for existing pretrained language models, PABEE does not introduce significant new ethical concerns.
  • Further work is needed to determine PABEE’s effects on biases already encoded in pretrained language models, such as gender bias.

A Image Classification

On CIFAR-10 and CIFAR-100, PABEE uses internal classifiers in a ResNet-56 backbone and improves accuracy and inference speed over the original ResNet while outperforming or matching baseline efficiency.

  • A Image Classification: Experiments use ResNet-56 on CIFAR-10 and CIFAR-100, adding an internal classifier after every two convolutional layers.The setup follows Shallow-Deep and compares PABEE with BranchyNet and Shallow-Deep.
  • A Image Classification: 0.2 and 0.5 accuracy gains and 1.26× and 1.22× inference speedups are achieved on CIFAR-10 and CIFAR-100, respectively.These results compare PABEE with the original ResNet model.
  • A Image Classification: PABEE demonstrates better performance and a similar speed-up ratio than both BranchyNet and Shallow-Deep.The comparison is reported for the CIFAR experiments.

B Proof of Theorem 1

The proof analyzes binary classification by comparing early-stopped misclassification probability with early-stopping probability. It derives bounds involving the number of internal classifiers, patience, and classifier error rates, including a special case where q = p.

  • Setup: The proof denotes patience as t, the number of internal classifiers as n, internal-classifier error as q, and final/original-classifier error as p.It seeks a condition on n and t under which PABEE improves conventional inference accuracy.
  • Special case: When q = p, the sufficient accuracy condition specializes to the stated n −t inequality.The supplied passage gives the specialization but truncates its right-hand side.
  • Probability comparison: For samples that do not early-stop, PABEE and conventional inference have the same misclassification probability.Thus, the proof focuses on the ratio between early-stopped misclassified samples, pmisc, and early-stopped samples, pstop.
  • Early-stopping bound: The lower bound for pstop is minimized at q = 0.5, corresponding to random guessing and the greatest instability of consecutive internal predictions.This establishes a worst-case setting for the early-stopping probability bound.
  • Misclassification bound: The proof upper-bounds pmisc using consecutive misclassification from the first internal classifier and from later classifiers after prior correct classification.The later contribution is represented by (n −t −1)(1 −q)q^(t+1), without accounting for inference that may already have stopped.

C Monte Carlo Simulation

Monte Carlo simulation tests Patience-based Early Exit in a simplified 12-layer binary classifier model. Across 10,000 Bernoulli-sampled runs, the method reduces the accuracy required of each individual classifier to match original inference accuracy.

  • Simulation setup: The simulation uses a 12-layer binary classification model whose classifiers C1–C12 share the same correct-prediction probability.This setup is designed to verify the theoretical feasibility of Patience-based Early Exit.
  • Simulation setup: 10,000 runs sample random Bernoulli distributions across original accuracies from 0.5 to 1.0 and patience settings t ∈ [1, 11].The simulation evaluates every 0.01 increment of original accuracy in that range.
  • Simulation findings: Patience-based Early Exit reduces the accuracy lower bound required for each individual classifier to achieve the original inference accuracy.Figure 5a reports the per-classifier accuracy requirement under different patience settings.
Loading 2006.04152v3…