The Right Tool for the Job: Matching Model and Instance Complexities

TL;DR

The paper addresses the growing cost of NLP inference by proposing a dynamic, multi-layer exit strategy for BERT-large. By attaching multiple layer-wise classifiers and calibrating their confidence, the model exits early for easy instances and proceeds to deeper layers for harder ones, reusing intermediate computations to save time. Across five datasets, the method delivers faster inference with comparable accuracy, and can be combined with distillation without retraining for each budget point. A detailed analysis of difficulty notions and calibration highlights both the potential and limitations, while showing practical applicability and public code release.

Abstract

As NLP models become larger, executing a trained model requires significant computational resources incurring monetary and environmental costs. To better respect a given inference budget, we propose a modification to contextual representation fine-tuning which, during inference, allows for an early (and fast) "exit" from neural network calculations for simple instances, and late (and accurate) exit for hard instances. To achieve this, we add classifiers to different layers of BERT and use their calibrated confidence scores to make early exit decisions. We test our proposed modification on five different datasets in two tasks: three text classification datasets and two natural language inference benchmarks. Our method presents a favorable speed/accuracy tradeoff in almost all cases, producing models which are up to five times faster than the state of the art, while preserving their accuracy. Our method also requires almost no additional training resources (in either time or parameters) compared to the baseline BERT model. Finally, our method alleviates the need for costly retraining of multiple models at different levels of efficiency; we allow users to control the inference speed/accuracy tradeoff using a single trained model, by setting a single variable at inference time. We publicly release our code.

Figures (7)

• Figure 1: An illustration of our approach. Some layers of a BERT-large model are attached to output classifiers, which make their respective predictions. The confidence of each layer-wise prediction is computed. If high enough, the model takes an early exit, avoiding the computation associated with successive (higher) layers (grayed out). Otherwise, the model continues to the next layer/classifier.
• Figure 2: Illustration of our baselines. (\ref{['fig:partial-baseline']}) Efficient baseline: adding a single output layer to an intermediate layer, while not processing the remaining BERT layers. (\ref{['fig:standard-baseline']}) The standard model: adding a single output layer to the final BERT layer. (\ref{['fig:our-approach']}) Our approach: adding multiple output layers to intermediate BERT layers; running the corresponding classifiers sequentially, while taking early exits based on their confidence scores.
• Figure 3: Test accuracy and processing time of our approach (blue squares, each point representing a different confidence threshold), our standard baseline (std., green diamond), efficient baselines (eff., red dots), and oracle baseline (orange star). Left and higher is better. Our method presents similar or better speed/accuracy tradeoff in almost all cases.
• Figure 4: Experiments with tinyBERT. Our method (light-blue pentagons) provides a better speed-accuracy tradeoff compared to the standard (light-green diamonds) and efficient (small light-red dots) baselines. For comparison, we also show the results of our method (blue squares) and our efficient baselines (large red dots) with BERT-large. Our method applied to BERT-large provides the overall best tradeoff.
• Figure 5: Instances with different labels are predicted with different degrees of confidence.