C ee BERT : Cross-Domain Inference in Early Exit BERT (2024)

Abstract: Pre-trained Language Models (PLMs), like BERT, with self-supervision objectives exhibit remarkable performance and generalization across various tasks. However, they suffer in inference latency due to their large size. To address this issue, side branches are attached at intermediate layers, enabling early inference of samples without requiring them to pass through all layers. However, the challenge is to decide which layer to infer and exit each sample so that the accuracy and latency are balanced. Moreover, the distribution of the samples to be inferred may differ from that used for training necessitating cross-domain adaptation. We propose an online learning algorithm named Cross-Domain Inference in Early Exit BERT (CeeBERT) that dynamically determines early exits of samples based on the level of confidence at each exit point. CeeBERT learns optimal thresholds from domain-specific confidence observed at intermediate layers on the fly, eliminating the need for labeled data. Experimental results on five distinct datasets with BERT and ALBERT models demonstrate CeeBERT’s ability to improve latency by reducing unnecessary computations with minimal drop in performance. By adapting to the threshold values, CeeBERT can speed up the BERT/ALBERT models by 2\times - 3.1\times with minimal drop in accuracy. The anonymized source code is available at https://github.com/Div290/CeeBERT.

## introduction in recent years, pre-trained language models (plms) have demonstrated substantial advancements in enhancing natural language processing (nlp) tasks. these models, such as elmo @xcite , bert @xcite , albert @xcite , gpt @xcite , xl-net @xcite , and roberta @xcite , encapsulate extensive knowledge transferable to diverse downstream tasks. despite their remarkable efficacy, large-scale plms suffer from inference latencies attributed to their substantial size. with millions or even billions of parameters, these models become computationally expensive, inefficient in terms of memory consumption, and exhibit latency challenges. moreover, past approaches like @xcite @xcite have highlighted the concern that plms tend to be overparameterized, leading to the 'overthinking' issue. this problem arises when shallow representations at early layers are often adequate for accurate predictions on many input samples, while the final layer's representations may suffer distractions due to over-complicated or irrelevant features that lack generalization. overthinking in plms wastes the computation, degrades model performance and hinders model generalization. to circumvent this, many variants of the bert, like deebert @xcite , elasticbert @xcite , fastbert @xcite , pabee @xcite , etc. facilitate inference at the intermediary layers of plms through early exits. in this setup, each sample must ascertain whether the inference can be completed at intermediary layers or the last layer. the decisions of early exit are based on the confidence at the intermediary layers being above a threshold. in the following, we consider the cost quantified as inference time (latency). however, depending on the application, the cost can also be present as other factors like power and computational resources. the threshold used to compare the confidence levels significantly impacts the amount of latency and accuracy: with a lower threshold, more samples exit early, but with a lower confidence value, leading to lower accuracy and lower latency. with a higher threshold, fewer samples exit early, leading to higher latency but improved accuracy. hence, one has to set the threshold that optimally trades-off between latency and accuracy. the threshold is often determined using a labeled dataset @xcite @xcite or by using some fraction of training data @xcite @xcite during training and serves as a crucial reference point for decision-making during inference. however, a significant challenge arises when deploying pre-trained models that are later tested on samples whose latent distribution can be different from the training samples in a zero-shot setting @xcite . for instance, a sentiment analysis model pre-trained on electronic product reviews (source data) to analyze sentiment in a distinct domain, like movie reviews (target data), poses a challenge. the language and sentiment expressions in movie-related content may significantly differ from electronic reviews as illustrated in fig 1 . plms exhibit strong generalization across domains on datasets with similar task types @xcite . nevertheless, the distribution of confidence in the prediction of classifiers attached to the intermediary layers can change when transitioning between domains. this can render early exit methods less efficient if the thresholds are not adjusted as per the latent distribution of input samples. this real-world challenge prompts the question: how to adapt the threshold in deployed early exit plms to maintain efficiency and robustness to domain shifts in a zero-shot setting? also, in post-deployment scenarios, samples are fed sequentially, in an online fashion, where inference for each sample needs to be performed before the next sample is fed, this leads to the challenge of learning the optimal threshold in an online and unsupervised manner. to tackle this issue, we introduce an online learning algorithm using the multi-armed bandit framework. @xcite . our algorithm, cross domain inference in early exits bert (ceebert), learns a threshold from a set of thresholds that achieves optimal trade-off between accuracy and latency. we extensively evaluate the performance of ceebert on five datasets viz. imdb, mrpc, scitail, snli, yelp, and qqp cover different classification tasks -sentiment, entailment, and natural language inference. in our evaluation, we trained the plm with exits on a source dataset and assessed its performance on a target dataset which exhibits variation in the latent data distribution from the source dataset but shares a similar task type. ceebert achieves speedup ranging from 2× to 3.5× in inference time, while maintaining a minimal accuracy loss of 0.1% to 3% compared to the naive bert classification. notably, instead of using pre-fixed thresholds learned on the source dataset, ceebert learns the thresholds on the fly for incoming samples from the target dataset. this sets it apart from previous early exit plms that rely on validation datasets representative of the target dataset to determine threshold values. our primary contributions are as follows: ## related work early exit methods are applied for various tasks such as image classification, image captioning and nlp tasks to reduce the computational resources and inference latency. early exits in image tasks: for image classification tasks, branchynet @xcite uses classification entropy at each attached exit to decide whether to infer the sample at the side branch based on the entropy of prediction. shallow-deep @xcite and msdnet @xcite improve upon branchynet by effectively choosing the thresholds based on the confidence distribution. similar architectures @xcite @xcite split the nn to be deployed on edge and cloud. see @xcite work in service outage scenarios. flexdnn (fang et al., 2020) and edgent (li et al., 2019) focus mainly on the most appropriate neural network (nn) depth. other works such as dynexit (wang et al., 2019b) focus on deploying the multi-exit nn in hardware. it trains and deploys the nn on field programmable gate array (fpga) hardware while paul et al. (kim and park, 2020) explains that implementing a multi-exit nn on fpga board reduces inference time and energy consumption. in a parallel vein, the mue and deecap (tang et al., 2023; fei et al., 2022) model employs a distinctive approach to apply early exits to the image captioning. deecap only applies to to decoder while mue applies to the encoder as well as the decoder. early exit in plms: multiple approaches have been proposed to effectively apply early exits to plms and solve multiple nlp tasks @xcite @xcite @xcite @xcite @xcite . dee-bert @xcite , elasticbert @xcite and berxit @xcite are based on the transformer-based @xcite bert model. berxit proposes an efficient finetuning strategy for the bert model with attached exits. deebert is obtained by training the exit points attached before the last module to the bert backbone separately. in contrast, elasticbert is obtained by training all the exit points attached to the bert backbone jointly. pabee @xcite ) is another multi-exit model that makes the exit decision based on the stability of the predic-tions after different exits. leebert @xcite proposed a self-distillation framework that has similar exiting criteria as pabee. etfee @xcite adds an adapter on top of the transformer layers and an (entangled frame) etf classifier to make intermediate exits learn better. multi-armed bandits in early exit nn: lee @xcite , dee @xcite and uee-ucb @xcite leverage the mab framework to learn the optimal exit in early exit nns. lee and dee mainly focus on learning optimal depth in image classification tasks, while uee-ucb finds optimal depth for nlp tasks employing a pre-trained elasticbert @xcite model. uee-ucb does not need any labels but works under the assumption that the prediction of the intermediary layers follows strong dominance property @xcite . @xcite utilize mabs to find the optimal splitting layers in a mobile-cloud co-inference setup. our approach differs from past works as 1) unlike previous studies, our work is primarily concerned with making unsupervised cross-domain inference efficient in early exit plms. 2) our work is focused on adapting the threshold values based on the underlying distribution of a dataset. 3) we use the multi-armed bandits framework to solve the problem of threshold learning. we compare against different early exiting models in ## problem setup we start with a pre-trained plm like bert or al-bert and attach exit classifiers at each layer. in the following, we discuss how the exit layers are trained and used for early inference. ## algorithm we develop an algorithm named cross domain inference in early exits in bert (ceebert). its pseudo-code is given in algorithm 1. the inputs to the algorithm are exploration constant γ and latency factor o i for each layer i. for the first |a| samples, the algorithm plays each arm once. in the subsequent rounds, it plays the arm with the highest upper confidence bound (ucb) index denoted as β t . ucb indices are obtained by taking the weighted sum of the empirical average of the rewards q t (α) and the confidence bonuses with γ as the weight factor. if c i at the ith layer is larger than β t then the sample exits, otherwise, the sample is passed to the next layer in the backbone while adding latency. if the sample does not exit at any intermediate classifier then it is inferred at the final layer. finally, the algorithm updates the number of pulls (n (β t )) and empirical mean (q(β t )) of the played arm. note that the algorithm is applied directly to the inference (target) dataset. observe an instance x t 5: apply threshold β t and observe c i 9: process and infer at the last layer. ## conclusion in this work, we introduced the concept of crossdomain inference in early exiting plms. we pro-