Consistent Accelerated Inference via Confident Adaptive Transformers (2021)

Abstract: We develop a novel approach for confidently accelerating inference in the large and expensive multilayer Transformers that are now ubiquitous in natural language processing (NLP). Amortized or approximate computational methods increase efficiency, but can come with unpredictable performance costs. In this work, we present CATs – Confident Adaptive Transformers – in which we simultaneously increase computational efficiency, while guaranteeing a specifiable degree of consistency with the original model with high confidence. Our method trains additional prediction heads on top of intermediate layers, and dynamically decides when to stop allocating computational effort to each input using a meta consistency classifier. To calibrate our early prediction stopping rule, we formulate a unique extension of conformal prediction. We demonstrate the effectiveness of this approach on four classification and regression tasks.

## introduction large pre-trained language models have become the de facto standard approach for solving natural language processing tasks @xcite . despite their impressive performance, however, their often massive computational burden makes them costly to run @xcite . concerns about their efficiency have kindled a large body of research in the field @xcite @xcite . for multilayered architectures such as the transformer, a popular approach is adaptive early exiting @xcite . early exiting takes advantage of the observation that task instances vary in complexity. in this setting, "early" classifiers are added on top of the simpler features of intermediate layers in the base model, and can trigger a prediction before the full model is executed. naively deciding when to preempt computation, however, can result in unpredictable decreases in model accuracy. quantifying the uncertainty in a prediction in order to decide when additional computation is needed (or not) is critical to making predictions quickly without excessively sacrificing performance. in this paper, we present confident adaptive transformers (cats), a general method for increasing transformer-based model efficiency while remaining confident in the quality of our predictions. specifically, given a fixed, expensive llayer model f(x), we create an amortized model g(x) that includes early classifiers {f 1 , . . . , f l }. 2 we then make g provably consistent with the original f with arbitrarily high probability (e.g., 95% of the time). this process is illustrated in figure 1 . our approach builds on conformal prediction (cp), a model-agnostic and distribution-free framework for creating well-calibrated predictions @xcite . concretely, suppose we have been (ex.1) claim: all airports in guyana were closed for all international passenger flights until 1 may 2020. evidence: airports in guyana are closed to all international passenger flights until 1 may 2020. evidence: the film was a success and broke box office sales record for mainland-produced romance films. given n examples, x @xmath0, @xmath1, . . . , n, as unlabeled calibration data, that have been drawn exchangeably from some underlying distribution p . let x n+1 ∈ x be a new exchangeable test example for which we would like to make a prediction. the aim of our method is to construct g such that it agrees with f with distribution-free marginal coverage at a tolerance @xmath2, 1), i.e., we consider g to be -consistent if the frequency of error, g(x n+1 ) = f(x n+1 ), does not exceed . 3 by design, this ensures that g preserves at least (1 -)-fraction of f's original performance. within these constraints, the remaining challenge is to make g relatively efficient (e.g., a consistent, but vacuous, model is simply the identity g f). in order to support an efficient g, we need a reliable signal for inferring whether or not the current prediction is likely to be stable. past work (e.g., @xcite rely on potentially poorly correlated metrics such as the early classifier's softmax response. we address this challenge by instead directly learning meta "consistency predictors" for each of the l -1 early classifiers of our l layer model, by leveraging patterns in past predictions. 4 figure 2 demonstrates the progression of meta confidence scores across layers when applied to "easy" versus "hard" instances from the vitaminc fact verification task @xcite . 3 for regression, we define equality as |g(•) -f (•)| ≤ τ . 4 we refer to the meta aspect of the classifier, not the optimization process (i.e., not to be confused with meta-learning). we pair the scores of our meta classifier for each layer with a stopping rule that is calibrated using a unique twist on standard conformal prediction. traditionally, cp is used to construct prediction sets that cover the desired target (e.g., y n+1 ) with high probability. we invert the cp problem to first infer the multi-label set of inconsistent layers, and then exit at the first layer that falls in its complement. we then demonstrate that this can be reduced to setting a simple (but well-calibrated) exit threshold for the meta classifier scores. our resulting algorithm is (1) fast to compute in parallel to the main transformer, (2) requires only unlabeled data, and (3) is statistically efficient in practice, in the sense that it finds low exit layers on average while still maintaining the required predictive consistency. we validate our method on four diverse nlp tasks-covering both classification and regression, different label space sizes, and varying amounts of training data. we find that it constitutes a simpleyet-effective approach to confident adaptive prediction with minimal interventions and desirable theoretical guarantees. in short, we provide: ## early exiting transformers in the following, we describe our dynamic early exiting model. we summarize early classification (following previous work) for convenience ( §3.1), and then present our novel meta consistency classifier ( §3.2). we focus on classification and regression tasks, given a model f(x) = y. we assume that f maps the input @xmath3. here, f is a multilayered transformer @xcite composed of l layers (although our method can be applied to any multilayer network). for all downstream tasks we follow standard practice and assume that the input contains a [cls] token whose representation is used for prediction. for classification, we use a task-specific head, softmax(w o (φ(w p h [cls] ))), where h [cls] ∈ r d is the hidden representation of the [cls] token, foot_1 φ is a nonlinear activation, and w * are linear projections, where w @xmath4. regression is treated similarly, but uses a 1-d output projection, w o • h @xcite . [cls] }, where h (k) ## conformalized early exits we now formulate the main contribution of this paper, which is a distribution-free and model-agnostic method based on cp for guaranteeing any performance bound an end-user chooses to specify. 8 our training ( §3), conformal calibration ( §4), and inference pipelines are summarized in algorithm 1. ## experimental setup for our main results, we use an albert-xlarge model @xcite with 24 transformer layers. results using an albert-base model and a roberta-large model @xcite are in appendix c. see appendix b for implementation details. we did not search across different values for the hyper-parameters of f or g as our approach is general and guarantees consistency for any f with any nonconformity measure (see appendix c.2). tuning the hyper-parameters could further improve the efficiency of g while preserving consistency. ## experimental results we present our main results. we experiment with both our meta classifier m k confidence score (meta, §3.2), and, for classification tasks, the early classifier's softmax response, p max k (sm), as a dropin replacement for m k (at no additional computational cost). appendix c reports results with other drop-in m k replacements, in addition to results using our naive development set calibration approach ( §3.3). appendix d provides qualitative examples. ## conclusion the ability to make predictions quickly without excessively degrading performance is critical to production-level machine learning systems. in fact, being capable of quantifying the uncertainty in a prediction and deciding when additional computation is needed (or not) is a key challenge for any intelligent system (e.g., see the system 1 vs. system 2 dichotomy explored in @xcite ). in this work, we addressed the crucial challenge of deciding when to sufficiently trust an early prediction of transformer-based models by learning from their past predictions. our confident adaptive transformers (cats) framework leverages meta predictors to accurately assess whether or not the prediction of a simple, early classifier trained on an intermediate transformer representation is likely to already be consistent with that of the full model f(x) (i.e., after all l layers of f are computed). importantly, we develop a new conformal prediction approach for calibrating the confidence of the meta classifier that is (1) simple to implement, (2) fast to compute alongside the transformer, (3) requires only unlabeled data, and (4) provides statistically efficient marginal guarantees on the event that the prediction of the faster, amortized cat model is consistent with that of the full f. our results on multiple tasks demonstrate the generality of our approach, and its effectiveness in consistently improving computational efficiencyall while maintaining a reliable margin of error.