Early Exiting BERT for Efficient Document Ranking (2020)

Abstract: Pre-trained language models such as BERT have shown their effectiveness in various tasks. Despite their power, they are known to be computationally intensive, which hinders real-world applications. In this paper, we introduce early exiting BERT for document ranking. With a slight modification, BERT becomes a model with multiple output paths, and each inference sample can exit early from these paths. In this way, computation can be effectively allocated among samples, and overall system latency is significantly reduced while the original quality is maintained. Our experiments on two document ranking datasets demonstrate up to 2.5x inference speedup with minimal quality degradation. The source code of our implementation can be found at https://github.com/castorini/earlyexiting-monobert .

## introduction large scale pre-trained language models such as elmo @xcite , gpt @xcite , bert @xcite , roberta @xcite , and albert @xcite have brought impressive improvements to natural language processing (nlp) and information retrieval (ir) applications. however, these large-scale models bring to our community not only exciting results, but also concerns about intensive computation demands and high inference latency, especially in real-world deployments. in this paper, we study how to accelerate inference of bert-based ir models. we follow the framework of monobert @xcite , which performs binary classification on query-document pairs into relevant/non-relevant. to accelerate inference for bert, we employ the idea of early exiting as in deebert @xcite . in deebert, extra classification layers ⋮ ⋮ transformer 1 classifier 1 exit continue transformer 2 classifier 2 exit continue transformer 𝑛 classifier 𝑛 exit input figure 1: overview of early exiting bert for document ranking. blue blocks are transformer layers and orange blocks are classifiers. we conduct experiments on bert base with two document ranking datasets, ms marco passage @xcite and asnq @xcite . we compare against cascade transformer @xcite , a recently proposed technique to accelerate inference in bert-based document ranking. results show that our method can reduce inference latency by up to 2.5× with minimal effectiveness degradation. ## related work neural document ranking models have brought significant improvements to ir tasks. throughout this paper, we refer to candidates to be retrieved as documents, although they may actually be passages, answers to a question, etc. representative neural ranking models include drmm @xcite , @xcite , knrm @xcite , co-pacrr @xcite -just to name a few. in recent years, large-scale pretrained language models, especially those based on the transformer architecture @xcite , have been applied to ir tasks and have pushed the state of the art even further @xcite @xcite . the idea of early exiting for neural networks originates from branchynet @xcite , and is also applied to nlp tasks in several papers @xcite @xcite . our work differs from them by using an early exiting strategy that specializes for document ranking. another related work that focuses on retrieval is cascade transformer @xcite , where a fixed proportion of samples are dropped after each layer. in contrast, our work drops samples based on their scores, and empirically we are able to achieve higher inference speedups. ## early exiting for document ranking the task of concern is document re-ranking, i.e., to rank among a small candidate document set, which is generated by a "bag of words" ir technique such as bm25. we assume in this paper that the candidate set is provided as input. ## experimental setup we apply early exiting on bert base and conduct experiments on two datasets for document ranking, ms marco passage @xcite and asnq @xcite . ## experimental results we show the trade-offs between model quality and computation of early exiting in tables 1 and 2 for ms marco passage and ansq, respectively. different trade-offs are achieved by setting various negative confidence threshold τ n , while τ p is always set to 1; we will provide detailed analyses of these thresholds later. inference efficiency is quantified by the average exit layer of inference samples; this metric is, according to our experiments, proportional to actual wall-clock runtime, while being invariant across multiple runs. we can see that in both datasets, early exiting is able to accelerate inference by ∼2.5× while maintaining the original model effectiveness. it is worth noting that in cascade transformer (ct) @xcite , only a part of the development set is used for evaluation, and therefore the scores are not directly comparable. however, in we also compare confidence-based early exiting trade-offs with layer-wise scores of the model in figures 2 and 3 . the layer-wise score of the i th layer is obtained by forcing all inference samples to exit through the classifier at the i th layer. it provide a series of baselines: if we want to save 50% inference computation for a 12-layer model, a straightforward way is to use the 6 th layer's classifier for all samples. the first two layers are omitted from the layer-wise score curves since their scores are too low to be useful. the figures show that the early exiting idea significantly outperforms the naïve baselines. to analyze the effect of different confidence thresholds, we plot in figure 4 the comparison of confidence thresholds τ p and τ n using the ms marco passage dataset. results from asnq are very similar and therefore omitted. each curve corresponds to one τ p value, and points within a curve are plotted by choosing different τ n values. 3 we notice that the trade-off performance is monotonic with respect to the positive confidence threshold: the higher τ p is, the better the trade-offs are. we speculate the reason is that while smaller τ p improves efficiency by allowing more samples to exit earlier, the quality of predictions from earlier layers of relevant samples degrades drastically. considering the fact that relevant samples constitute only a tiny fraction of all candidates in both datasets, setting τ @xmath0, i.e., using as many transformer layers as we can on positive samples, is the optimal choice. it is worth noting that for other datasets with higher relevant candidate proportions, the optimal τ p may be smaller than 1. within one curve (a fixed τ p ), τ n controls the trade-offs between efficiency and quality: lower τ n allows more negative samples to exit earlier, thus improving efficiency with relatively small quality degradation. such asymmetry demonstrates the necessity of using two confidence thresholds instead of one: while we can safely perform early exiting on negative samples and save computation, we should allocate far more resources to positive samples for the most accurate predictions. ## conclusion we propose asymmetric early exiting bert for document ranking, an effective method to improve model efficiency in ir tasks. computation resources are allocated to samples according to their needs. experiments show that our method is able to achieve different quality-latency trade-offs by setting different thresholds, and also improves model efficiency over baselines.