DE$^3$-BERT: Distance-Enhanced Early Exiting for BERT based on Prototypical Networks

TL;DR

This work tackles the latency of BERT-style inference by highlighting that exiting indicators based only on local information can misestimate prediction correctness. It introduces DE$^3$-BERT, which attaches prototypical networks to every internal layer and trains them with distance-aware regularization to learn a global metric space of class prototypes; its inference uses a hybrid exit indicator, EDR, that fuses local entropy with a distance ratio derived from prototype distances via $DR(r_1^{(m)},r_2^{(m)})=0.5\times\left(1+\frac{r_1^{(m)}-r_2^{(m)}}{\max\{r_1^{(m)},r_2^{(m)}\}}\right)$ and $EDR=\frac{\lambda+1}{\frac{\lambda}{DR}+\frac{1}{Entropy}}$. Experiments on GLUE show that DE$^3$-BERT achieves a superior speed-accuracy trade-off with minimal overhead and generalizes across backbones and languages, while ablations confirm the critical roles of prototypical networks and DAR in learning discriminative metric spaces. The approach provides a more reliable exiting mechanism by leveraging global information and offers practical benefits for edge and real-time NLP deployments, with avenues for extension to regression and OOD scenarios.

Abstract

Early exiting has demonstrated its effectiveness in accelerating the inference of pre-trained language models like BERT by dynamically adjusting the number of layers executed. However, most existing early exiting methods only consider local information from an individual test sample to determine their exiting indicators, failing to leverage the global information offered by sample population. This leads to suboptimal estimation of prediction correctness, resulting in erroneous exiting decisions. To bridge the gap, we explore the necessity of effectively combining both local and global information to ensure reliable early exiting during inference. Purposefully, we leverage prototypical networks to learn class prototypes and devise a distance metric between samples and class prototypes. This enables us to utilize global information for estimating the correctness of early predictions. On this basis, we propose a novel Distance-Enhanced Early Exiting framework for BERT (DE$^3$-BERT). DE$^3$-BERT implements a hybrid exiting strategy that supplements classic entropy-based local information with distance-based global information to enhance the estimation of prediction correctness for more reliable early exiting decisions. Extensive experiments on the GLUE benchmark demonstrate that DE$^3$-BERT consistently outperforms state-of-the-art models under different speed-up ratios with minimal storage or computational overhead, yielding a better trade-off between model performance and inference efficiency. Additionally, an in-depth analysis further validates the generality and interpretability of our method.

Figures (10)

• Figure 1: The discrepancy between model confidence and prediction correctness. On a sample in QQP, the first internal classifier makes a wrong prediction with a low entropy value (high confidence), while the sixth classifier makes a correct prediction with a high entropy value (low confidence).
• Figure 2: The performance-efficiency trade-offs of existing exiting strategies, the oracle, and the BERT-base backbone on the QNLI development set. Right and higher is better. The significant gap between the existing exiting strategies and the oracle indicates considerable potential for improving the estimation of prediction correctness, which is crucial to reliable exiting decisions.
• Figure 3: Comparison between the classic early exiting framework and our DE$^3$-BERT. Most existing frameworks rely on local information from individual test samples to estimate the correctness of early predictions. Instead, our DE$^3$-BERT considers both local information (entropy) and global information (distance ratio) to enhance the estimation of prediction correctness for more reliable exiting decisions. For the exit decision-making process of a sample $x$ shown in (b), (1) compute the entropy of the early prediction $p_x$ suggested by the internal classifier; (2) compute the distance ratio between the sample representation $h'_{x}$ and the two class prototypes ${cp}_{k1}$ and ${cp}_{k2}$ as illustrated in Equation (\ref{['dr']}), where ${cp}_{k1}$ and ${cp}_{k2}$ correspond to the classes $k1$ and $k2$ with the highest two probability scores in $p_x$, and the prototypical network $\rm PN$ is used to map the hidden states $h_x$ into the metric space; (3) output the hybrid exiting indicator $\mathrm{EDR}$ according to Equation (\ref{['edr']}), and terminate the inference process once the $\mathrm{EDR}$ value falls below the threshold $\tau$.
• Figure 4: Performance-efficiency trade-off curves using different exiting strategies on the development sets of four GLUE tasks. DE$^3$-BERT is our hybrid exiting strategy. Entropy denotes the classic entropy-based exiting strategy. Each point on the curve corresponds to a selected threshold. Its horizontal and vertical coordinates represent the corresponding speed-up ratio and task performance, respectively. Models are trained as illustrated in Section \ref{['section:training_scheme']}.
• Figure 5: Impact of prototypical networks and DAR on the trade-off between model performance and efficiency on the development sets of four GLUE tasks. We employ the proposed hybrid exiting strategy for all models during inference. PN denotes the prototypical networks.