W HEN IN D OUBT , S UMMON THE T ITANS : E FFICIENT I NFERENCE WITH L ARGE M ODELS
Anonymous authors
Paper under double-blind review
A BSTRACT
Scaling neural networks to “large” sizes, with billions of parameters, has been shown to yield impressive results on many challenging problems.
However, the inference cost incurred by such large models often prevent their application in most real-world settings.
In this paper, we propose a two-stage framework based on distillation that realizes the modelling beneﬁts of the large models, while largely preserving the computational beneﬁts of inference with more lightweight models.
In a nutshell, we use the large teacher models to guide the lightweight student models to only make correct predictions on a subset of “easy” examples; for the “hard” examples, we fall-back to the teacher.
Such an approach allows us to efﬁciently employ large models in practical scenarios where easy examples are much more frequent than rare hard examples.
Our proposed use of distillation to only handle easy instances allows for a more aggressive trade-off in the student size, thereby reducing the amortized cost of inference and achieving better accuracy than standard distillation.
Empirically, we demonstrate the beneﬁts of our approach on both image classiﬁcation and natural language processing benchmarks.
I NTRODUCTION
Scaling neural networks to “large” sizes has brought dramatic quality gains over a wide variety of machine learning problems, including at the tails.
In computer vision, the high performing models for image classiﬁcation (Kolesnikov et al., 2019; Xie et al., 2020; Tan & Le, 2019; Foret et al., 2021) and segmentation (Ghiasi et al., 2020) have upto 928M parameters and require up to 600G FLOPs for a prediction.
Similarly, in natural language processing, transformer-based approaches, which have several billion parameters and require up to a tera-FLOP for a prediction, are leading performance on language understanding tasks (Raffel et al., 2019; Brown et al., 2020; Fedus et al.) and neural machine translation (Bapna & Firat, 2019; Huang et al., 2018).
The immensely expensive inference cost of these large models is, however, hindering their direct widespread adoption (Jouppi et al., 2017; Ning, 2013;
Crankshaw et al., 2017; Zhang et al., 2019).
The issue is further exacerbated in deployment over resource-constrained edge devices such as mobile phones (Zhang et al., 2020).
As a workaround, many model compression techniques have been proposed to reduce the computational cost and memory footprint by trading-off accuracy, including quantization (Mozer & Smolensky, 1988; Han et al., 2015), pruning (LeCun et al., 1989; Hassibi &
Stork, 1993), and distillation (Bucilˇa et al., 2006; Romero et al., 2014; Hinton et al., 2015).
However, there is a limit to how far such model compression techniques can be pushed to reduce inference cost while retaining good performance across all inputs (cf. teacher-student accuracy gaps in (Cho & Hariharan, 2019; Menon et al., 2020b; Mirzadeh et al., 2020; Wang et al., 2017a)).
Ideally, the compute required to make predictions on an instance should depend on the hardness of the instance.
But the large models do not adapt their computational budget based on the complexity of the task at hand.
We conjecture that the full ability of a large model is needed only for a small fraction of “hard” instances.
The majority of real-inputs are “easy”, for which performing full computation of a large model is wasteful; rendering the overall ML system inefﬁcient.
Such an inefﬁcient utilization of compute gets even more pronounced for many real-world data that are heavily long-tailed (Zhu et al., 2014;
Wang et al., 2017b; Van Horn & Perona, 2017), with hard instances belonging to the tail.
[Figure] [Figure]
56 Student Size [Layers]
T o t a l C
o m p u t e
C o s t [ M
F L O P S ]
[Figure]
In this paper we focus on realizing the beneﬁts of a large model on the hard instance without incurring the unnecessary large inference cost on prevalent easy instances.
Towards this, we propose to employ a novel distillation-based two-stage inference framework in Figure 1 (left): First use a lightweight student model to make a prediction.
If the student is conﬁdent, we emit the prediction and we want the student to be conﬁdent on all the easy instances, which should be a large fraction of the test time queries.
When the student is in doubt, ideally only for a small number of hard examples, we fall-back to the large teacher.
Our main contributions for leveraging the excellent performance of large models to realize a desirable inference cost vs. performance trade-off are as follows
• The instance-aware two-stage inference mechanism crucially relies on the ability of the student model to detect the “hardness” of an input instance on the ﬂy and routing it to the large model.
To enable this routing, we propose modiﬁed distillation procedures.
In particular, we employ novel distillation loss functions (cf.
Sec. 4) such that the student gets penalized heavily for making mistakes on easy examples while for harder out-of-domain examples we encourage the student to be less conﬁdent, e.g., the prediction distribution be closer to the uniform distribution.
• We conduct a detailed empirical evaluation of the proposed distillation-based two-stage inference framework (cf.
Sec. 5) and show that it allows us to much more aggressively trade-off size of the student for multiple image classiﬁcation and natural language processing (NLP) benchmarks.
Interestingly, as summarized
in
Figure
, there is a sweet spot where we can achieve the same accuracy as the teacher with 45% less compute.
This beneﬁt is further magniﬁed when considering only in-domain examples.
Thus, we can reduce the overall computation over the data distribution and achieve better accuracy than performing inference with only the student model.
Note that, traditionally, the distillation approach aims to utilize a complex model to learn a simple model that has its overall performance as close to the complex model as possible.
This is done under the assumption that during the inference time one can ‘throw away’ the complex model and rely on only the simple model for the ﬁnal predictions.
We would like to highlight that our goal is not to train a student/simple model that will be used as a standalone model to generate predictions.
It’s worth mentioning that the proposed two-stage inference can also be useful in a modern setup like edge computing and 5G cloudlets (Fang et al., 2019), where a lightweight student model runs on a device to make most of the predictions with low latency and only once in a while a hard instance is delegated to a shared large teacher model running in the cloud.
R ELATED WORK
Techniques to reduce inference cost for deep models mainly fall under two different approaches: quantization and pruning, and adaptive computation.
Quantization and pruning.
The primary way suggested in the literature to accelerate predictions from deep neural networks has been quantization and pruning (Mozer & Smolensky, 1988; LeCun et al., 1989; Hassibi & Stork, 1993;
Li et al., 2020; Carreira-Perpin´an, 2017; Howard et al., 2019).
Signiﬁcant progress was made by introducing Huffman encoding methods for non-uniform quantization which led to a reduction in network sizes by orders of magnitude and up to 4x reduction in overall prediction cost (Han et al., 2015).
Since then pruning and quantization have been widely adopted in computer vision and more details can be found in the recent survey by Liang et al. (2021).
In the NLP domain, Gordon et al.
; Zadeh et al.
(2020) proposed pruning BERT during training, which resulted in 30%-40% reduction in model size with minimal effect on the accuracy of the ﬁnal task, however, not much compute/time savings was observed as arbitrary sparsity might not be leveraged by modern hardware accelerators.
Towards this, structured pruning is more beneﬁcial, as it removes a series of weights that correspond to an entire component of the model (Ganesh et al., 2020).
In transformers, this would correspond to pruning out entire attention heads (Kovaleva et al., 2019; Raganato et al., 2020) or encoder units (Fan et al., 2019).
Our proposed approach of two-stage inference is complementary to such techniques and can be combined with these to further reduce the inference cost.
Adaptive computation.
In line with our proposed approach, there have been works trying to adapt the amount of computation of neural model based on an input instance.
Effort in this space started in the vision community for enabling real-time object detection by Rowley et al.
and later formalized by Viola & Jones (2001).
The basic idea was to design a cascade of independent classiﬁer and reject early on and cheaply.
The idea has been generalized from a linear chain of cascaded classiﬁers to trees (Xu et al., 2014).
Instead of combining many independent classiﬁers, a similar idea to stop early has emerged in monolithic deep models.
One approach to this problem is represented by Adaptive Computation Time (ACT) (Graves, 2016; Chung et al., 2016).
ACT is a mechanism for learning a scalar halting probability, called the “ponder time”, to dynamically modulate the number of computational steps needed for each input.
An alternative approach is represented by Adaptive
Early Exit Networks (Bolukbasi et al., 2017), which gives the network the ability to exit prematurely - i.e., not computing the whole hierarchy of layers - if no more computation is needed.
A modern incarnation of this approach in NLP with transformers encoders appeared in Schwartz et al.
Liu et al.
Dabre et al.
This idea has been extended to generative tasks as well, where a number of decoder layers per time step are adapted in Elbayad et al.
As a further generalization, Bapna et al.
(2020) introduced “control symbols” to determine which components are skipped in a transformer, i.e. not all previous components need to be executed.
Similar ideas had already existed in the vision community, for example, Wang et al.
(2017a; 2018) introduced a method for dynamically skipping convolutional layers.
All of these approaches are specialized to a task and rely on designing the whole pipeline from scratch which can be expensive if we want to achieve state-of-the-art (SoTA) performance.
In contrast, we want to design efﬁcient inference techniques achieving SoTA performance by only training cheap components, like the student model using a novel distillation procedure, while leveraging existing SoTA large models without re-training or modifying them.
Moreover, our approach is a generic framework to leverage the large models independent of the underlying model architecture and problem domain.
Our proposed approach ensures that large model is only invoked on instances that necessarily beneﬁt from its large model capacity and a lite distilled model sufﬁces to predict a large portion of test instances.
B ACKGROUND 3. M ULTICLASS CLASSIFICATION
Consider a standard multiclass classiﬁcation problem where given an instance x ∈ X , the objective is to classify the instance as a member of one of the L classes, indexed by Y (cid:44)
[ L ] .
In the most common setting, given a training set comprising of n instance and label pairs or training examples S n = { ( x i , y i ) } i ∈ [ n ] , one learns a classiﬁcation model f : X → R L , where f ( x ) =
( f ( x ) 1 , . .
f ( x ) L ) represent the model scores assigned to instance x for L classes.
Based on the model scores, an instance can be predicted to belong to class ˆ y x (cid:44) arg max
i ∈
[ L ] f ( x ) i .
Accordingly, the model (top1 ) accuracy is deﬁned as
[Equation]
Ideally, one prefers a classiﬁcation model with a high accuracy.
Let (cid:96) :
Y × R L → R be a surrogate loss function such that (cid:96) ( y, f ( x )) closely approximates the misclassiﬁcation error
ˆ y x
(cid:54) =
y .
Typically, given the training set S n and the loss function (cid:96) , one selects a desired classiﬁcation model via empirical risk minimization (ERM).
Softmax cross-entropy loss is one of the most widely used surrogate loss functions for multiclass classiﬁcation: Given model scores f ( x ) =
( f ( x ) 1 , . .
f ( x ) L ) , one computes the softmax distribution
[Equation]
where τ denotes the temperature parameter of the softmax operation.
1 Further, we deﬁne p y ∈ { 0 , 1 } L to be the one-hot label distribution corresponding to the true label
y ∈
[ L ] , which has non-zero value at only y -th coordinate.
Now, softmax cross-entropy loss corresponds to the distance between the distributions p y and ˆ p f,x , as measured by the cross-entropy function.
[Equation]
3. M ODEL DISTILLATION
Distillation is a celebrated training techniques that utilizes one model’s scores to train another model (Bucil ˇ a et al., 2006; Hinton et al., 2015).
The former model is typically referred to as the ‘teacher’ model while the latter model is called the ‘student’ model.
During distillation, given a teacher model g : X → R L and an example ( x, y ) ∈
X × Y , one ﬁrst deﬁnes the teacher (softmax) distribution ˆ p g,x ( i ) as per (2).
Now, as opposed to utilizing the ‘one-hot’ label distribution p y , we treat ˆ p g,x
( i ) as the pseudo label distribution and deﬁne the distillation version of the softmax cross-entropy loss for the student model f : X → R L as (cid:96) distill ( g ( x ) , f ( x ))
=
H (ˆ p g,x , ˆ p f,x ) .
For a, b ∈ R + , distillation involves minimizing
[Equation]
Note that the objective in (3) utilized both the true labels { y i } and the teacher scores { g ( x i ) } .
When b = 0 , this reduces to the standard training.
More interestingly, when a = 0 , (3) correspond to training with solely ˆ p g,x
( i ) – a fairly common way to utilize distillation (Menon et al., 2020a).
Remark 1.
For distillation, teacher models are usually much more complex and powerful as compared to student models (Bucil ˇ a et al., 2006; Hinton et al., 2015).
However, the distillation with an equally complex, and even a simpler, teacher model has been shown to improve the quality of the student model via distillation (see, e.g., Rusu et al., 2016; Furlanello et al., 2018; Yuan et al., 2020).
D ISTILLATION FOR TWO - STAGE INFERENCE
We now propose various distillation approaches to power the proposed two-stage inference framework.
Recall that we intend to obtain a lightweight student model that can generate highly accurate predictions on easy instances and route hard instances to the large teacher model.
This raises an interesting question if one needs to modify the distillation process in any way whatsoever to aid our objective as typically distillation envisions the student model to be used as a standalone model.
Towards this, we explore a general distillation framework that partitions the training examples S into two groups: 1) easy instances S easy and 2) hard instances S hard .
Accordingly, we modify the distillation process such that the student incurs larger loss when it makes incorrect predictions S easy .
Furthermore, we penalize the student less for making mistakes on S hard , which we accomplish by carefully designed supervision during the distillation.
Now, we present two speciﬁc realizations of the above generic distillation framework based on two different strategies to partition the training examples into S easy and S hard .
4. C LASS - SPECIFIC DISTILLATION
Many real-world data distributions exhibit a long-tail behavior where most of the instances we observe belong to a small number of classes, and the remaining classes are represented by very few instances.
This has sparked a long line of work on improving model performance on the tail classes (see, e.g., Cui et al., 2019;
Cao et al., 2019; Kang et al., 2020; Menon et al., 2021).
In contrast, we explore an orthogonal direction and leverage the data imbalance to enable efﬁcient inference with large models.
In particular, for L in ⊆
[ L ] , we deﬁne
S easy = {
[Equation] [Equation]
y
[Equation]
L in } and S hard =
S \ S easy =
{ ( x, y
[Equation]
S
[Equation]
y / ∈ L in } .
Thus, we require the lite model to perform well only on a subset of classes L in and route the examples from the remaining classes to the large model.
Here, we hypothesize that the smaller model can better utilize its limited capacity to perform well on a subset of classes.
Now, setting L in to be the head classes will ensure that the lite model itself tries to predict the examples from the head classes.
Since the underlying data-distribution is long-tail, only the examples from the tail classes (and a few hard examples from the head classes depending on the exact implementation details described later) are sent to the large model during inference.
With this general approach in mind, we propose a class-speciﬁc distillation approach.
2 Now, given a large teacher model g and example ( x, y ) , we deﬁne a pseudo label distribution (cid:101) p class g,x as follows:
[Equation]
where ˆ p g,x and p y denote the teacher’s softmax distribution and the one-hot label distribution, respectively.
In addition, α ∈ [0 , 1] denotes a label-smoothing parameter.
Now, we train a lite student f : X → R L with the distillation loss
[Equation]
Note that the loss in (5) utilizes teacher softmax distribution for classes in L in and relies on labelsmoothed one-hot distribution for the remaining classes.
This loss has two desirable properties for our two-stage inference objective: 1) As a result of standard distillation, the lite student behaves as a well-calibrated model with good performance on the examples from the classes in L in . 2)
Due to standard label-smoothing (Szegedy et al., 2016; M ¨ uller et al., 2019), the lite student achieves a smaller margin on the examples belonging to the classes in [ L ] \L in .
Let f L in be the lite student model obtained by minimizing the loss in (5).
Now, given a test instance x ∈ X , we ﬁrst run inference with the lite model to obtain f L in ( x ) .
Subsequently, we decide if we need to make the ﬁnal prediction based on f L in ( x ) or delegate the example to the large teacher g to obtain the ﬁnal prediction.
We identify two useful delegation schemes, which we detail next.
Class-based delegation.
Recall that the class-speciﬁc distillation aims to utilize the lite model to classify only the examples belonging to the classes in L in .
Thus, the prediction made by f L in serves as a natural candidate for the delegation.
In particular, when
[Equation]
we declare ˆ y student ( x ) as the ﬁnal prediction.
Otherwise, x is sent to the teacher and ˆ y teacher ( x ) (cid:44) arg
max j ∈
[ L ] g ( x ) becomes the ﬁnal prediction.
Margin-based delegation.
The class-based delegation is designed under the assumption that the lite model achieves high accuracy on the examples belonging to the classes in L in and identiﬁes the examples from [ L ] \L in with high ﬁdelity.
Both of these assumptions don’t always hold in practice.
In particular, f L in may ﬁnd some instances from L in hard and incorrectly predict a wrong class in
L in when presented with those instances.
Similarly, f L in may predict a class from L in when the test instance belongs to [ L ] \L in .
Recall that the distillation loss in (5) is designed to ensure that f L in is well-calibrated on the instances from L in ; as a result, it attains small margin on hard instances from L in .
Moreover, by design, f L in realizes small margin on the instances from [ L ] \L in .
Thus, one can utilize the margin of the lite model f L in to perform delegation.
Towards this, recall that
[Equation]
denotes the margin realized by f L in on x .
Here, ˆ p f L in ,x ([ i ]) denotes the i -th largest element of the vector
ˆ p f L in ,x .
Now, given a design parameter ρ ∈ (0 , 1) , we assign the following ﬁnal prediction to a test instance x ∈ X .
[Equation]
The margin-based delegation aims to ensure that hard instances from L in as well as all instances in [ L ] \L in get routed to the teacher for the ﬁnal prediction.
Remark 2.
The distillation based on the loss in (5) constitutes only one of many possible ways to transfer the performance of a teacher over a subset of classes to a lite student.
We discuss two other class-speciﬁc distillation-based approaches in Sec.
A of the appendix and evaluate those in Sec. 5.
4. M ARGIN - BASED DISTILLATION
As discussed before, the margin assigned to an instance by a model is a natural proxy for the hardness of the instance, as viewed by the model.
In Sec. 4.1, we utilize the margins assigned by the student to delegate the examples to the teacher.
This raises an interesting question if one can utilize the margins to partition the training data into S easy and S hard example during the distillation.
Towards this, given a teacher g and a parameter ρ tr ∈ (0 , 1) , we add a training example ( x, y ) to S easy iff γ g ( x ) (cid:44) ˆ p g,x ([1]) − ˆ p g,x ([2]) >
ρ tr .
Given this data partition, for an example ( x, y ) , we deﬁne the pseudo label distribution as follow:
[Equation]
Now, we obtain a (lite) student f Lρ tr : X → R L by minimizing
[Equation]
For the two-stage inference, given a test instance x , we make the following ﬁnal prediction.
[Equation]
where ˆ
y f
Lρ tr = arg max
j f Lρ
tr
( x ) j .
Remark 3.
For a small value of ρ tr , we expect the student to make the correct prediction on almost all examples.
Thus, the margin-based distillation proposed above becomes very similar to the normal distillation discussed in Sec. 3.2.
On the other hand, for large values of ρ tr , the student is expected to do well on only a small subset of examples during training and be able to identify the rest of the examples as hard instances that need to be routed to the teacher.
Similar to class-speciﬁc distillation, there are multiple potential variants of the margin-based distillation.
We discuss one such variant that relies on an ‘abstain’ class in Sec.
C of the appendix.
E XPERIMENTS
We now conduct a comprehensive empirical study of our distillation-based two-stage inference procedure.
In particular, we evaluate various distillation frameworks introduced in Sec.
4 along with different choices of delegation methods.
On standard image classiﬁcation tasks, we establish that:
(i)
A large improvement in accuracy over the student-only approach can be realized by delegating only a small fraction of examples to the large teacher model (Sec. 5.1).
This validates our claim that one can rely on a much smaller student and achieve an accuracy similar to the teacher-only approach with a small increment in inference cost.
[Figure]
(ii) Advantages of the two-stage inference are even more pronounced if we focus on the in-domain performance, where the in-domain portion of the instance space corresponds to
a subset of classes or instances with a large margin (based on the large teacher model).
This validates the utility of two-stage inference for those real-world settings where a large data imbalance exists in favor of in-domain instances (Sec. 5.2).
(iii) Class-speciﬁc distillation deﬁned by (5) indeed achieves the desired behavior where the student enables a clear dichotomy among the in-domain and out-of-domain instances.
By varying the label-smoothing parameter, we can improve in-domain model performance and delegate a small number of instances to the teacher at the cost of performance on the entire test data (Sec. 5.2).
Furthermore, by delegating a small fraction of hard instances to a large teacher model, our proposed class-speciﬁc distillation and a variant of margin-based distillation enable efﬁcient inference on sentence classiﬁcation and reading comprehension tasks in the NLP domain, respectively (Sec. 5.3).
We mainly focus on three benchmark image datasets – CIFAR-100 (Krizhevsky, 2009), ImageNet1k (Russakovsky et al., 2015), and ImageNet-21k (Deng et al., 2009).
In addition, we also evaluate the proposed two-stage inference procedure on a sentence classiﬁcation task based on
MNLI (Williams et al., 2018) and a reading comprehension task based on SQuAD dataset (Rajpurkar et al., 2016).
On image classiﬁcation tasks, we use EfﬁcientNet-L2 (Xie et al., 2020) as the large teacher.
As for the lite student, we utilize ResNet (He et al., 2016a;b) and MobileNetV3 (Howard et al., 2019) for CIFAR-100 and ImageNet, respectively.
For the sentence classiﬁcation and reading comprehension tasks, RoBERTa-Large (Liu et al., 2019) and T5-11B (Raffel et al., 2019) serve as the teachers, respectively.
We use MobileBERT (Sun et al., 2020) as the lite student for both MNLI and SQuAD.
We provide a detailed description in Sec.
D of the appendix.
5. O VERALL ACCURACY GAINS
We begin by establishing the utility of the two-stage inference procedure for overall performance improvement.
Towards this, Fig.
2a, 3a, and 5a (in appendix) depict the two-stage overall accuracy achieved by various class-speciﬁc distillation approaches via margin-based delegation.
Besides, we also include conventional model distillation as Baseline in our evaluation.
Note that, for all approaches, allowing the lite student model to delegate hard instances (the ones with a small margin) to the large teacher signiﬁcantly increases the accuracy as compared to the performance of (one-stage) student-only inference approach.
Focusing on CIFAR-100, both Baseline and CD-I can approximate the performance of the teacher model by delegating only ∼ 40% test instance, which translates to ∼ 60% reduction in the inference cost as compare to the large teacher (cf.
Fig.
2a).
Here, we note that CD-II and CD-III train the student to make the correct prediction on only the instances belonging to the classes in L in , leading to poor overall accuracy on the entire test set.
However, unsurprisingly, as student delegates more instances to the teacher, the two-stage overall accuracy improves.
Similar conclusions also hold on ImageNet datasets (cf.
Fig.
3a and 5a).
Table 1 and(in the appendix) show that the class-speciﬁc distillation based two-stage inference leads to similar improvements in the overall accuracy when we employ class-based delegation from Sec. 4.1.
[Figure]
5. I N - DOMAIN PERFORMANCE
[Figure]
[Table]
Next, we verify that two-stage distillation is even more beneﬁcial in those real-world settings where ML models encounter heavily imbalanced data during the inference.
Ideally, the lite student model should make a prediction on frequent but easy instances and delegate rare but hard instances to the large teacher.
This would ensure that the two-stage inference procedure realizes high overall accuracy (compared to inference with only the student) while signiﬁcantly lowering the inference cost (compared to inference with only the teacher).
With this in mind, we evaluate the performance of the two-stage inference procedure on indomain (so-called easy instances).
By design, for class-speciﬁc distillation, in-domain instances belong to the classes in L in .
Similarly, for the margin-based distillation, in-domain instances are the ones where teacher assigns a large margin ρ .
For CIFAR-100, Fig.
2b shows that two-stage inference achieves much better in-domain (deﬁned by |L in | = 30 classes) accuracy as compared to overall accuracy.
This implies that in a real-world setting where in-domain instances are frequent, the two-stage inference can efﬁciently achieve very good overall performance by having the student predict most of the in-domain instances.
Note that overall performances in such a scenario (with data-imbalance) is not captured by Fig.
2a which is based on fair balanced test data.
Instead of aiming to imitate the data imbalance encountered in practice, we have separately highlighted the performance of two-stage inference on the in-domain instances.
This shows that the more imbalanced the data is the more advantageous two-stage distillation would be in terms of saving the inference cost without affecting the classiﬁcation performance.
As per Fig 2c, we essentially maintain same accuracy as the teacher but reduce latency by more than 2x.
Here, we also note that Fig.
2a implies that two-stage inference would indeed enable one to approximate the teacher performance on all instances, whether they belong to in-domain or not.
Interestingly, Fig. 2 indicates that CD-I with margin-based delegation gives a favorable tradeoff between in-domain and (balanced) overall accuracy.
The label-smoothing parameter α in (4) helps create a dichotomy between in-domain and out-of-domain instances.
The larger values of α ensure that the student assigns a smaller margin to out-of-domain examples and tries to improve its performance on in-domain instances.
This is reﬂected in Fig.
2b where CD-I with large α achieves higher
[Figure]
accuracy by predicting a larger fraction of in-domain instances at the student.
We note the similar trend in Table 1 where we combine class-speciﬁc distillation with class-based delegation
We also evaluate the in-domain two-stage performance of class-based distillation on ImageNet-1k with |L in | = 300 .
The conclusions from Fig.
3b, 3c and Table 4 (in the appendix) are similar to those observed on CIFAR-100.
Also, see Fig.
5b (in the appendix) for the identical trends on
ImageNet-21k.
Finally, we also studied the in-domain performance of two-stage inference enabled by margin-based distillation from Sec.
4.2 on both CIFAR-100 and ImageNet-1k in the appendix.
5. T WO - STAGE INFERENCE IN NLP DOMAIN
Sentence classiﬁcation task.
Fig. 4a and 4b show the performance of our proposed two-stage inference framework on MNLI.
Since it has only 3 classes, we employ margin-based distillation (with ρ tr = 8 . 0) from (9) with margin-based delegation.
Our conclusion for the image classiﬁcation also extends to the text domain and the two-stage inference framework enables a lite student to leverage high-quality teacher by delegating a small portion of (hard) instances to the teacher.
We provide inference-latency vs. in-domain performance trade-off for MNLI in the appendix (cf. Fig. 8a).
Reading comprehension task.
To showcase the generality of our approach , we apply the twostage inference method to a machine comprehension task (SQuAD), where there is a fundamental mismatch between student architecture which is based on span selection whereas teacher is based on encoder-decoder architecture.
Thus, one cannot do standard distillation by transferring logits from teacher to students, but a suitable modiﬁcation of margin-based distillation from (9) still works.
We partition the training examples into easy and hard instances by thresholding teacher’s log-likelihood of ground-truth answer given the context.
While training student, for easier examples, we use one hot labels for the start and end span as no teacher logits are available.
Whereas for hard examples, we still employ label smoothing.
We try two values α = 0 .
4 in the experiments.
Our results in
Fig.
4d and 4d, which has similar conclusions to the classiﬁcation experiments.
See Fig.
8b for the inference-latency vs. in-domain performance trade-off achieved by the two-stage inference.
D ISCUSSION
We propose a distillation-based two-stage inference framework to efﬁciently leverage large models with prohibitively large inference costs in real-world settings.
Given a large model, we distill a lite student that utilizes its limited model capacity to perform well on easy (in-domain) instances and can identify hard (out-of-domain) instances.
When deployed in tandem, the lite model generates the ﬁnal prediction on easy but frequent instances and delegates hard but rare instances to the large model.
This ensures much higher performance (compared to inference based on only the lite model) and a much smaller inference cost (compared to inference with only the large model).
We propose various distillation methods to enable such two-stage inference.
We establish the utility of these approaches for realizing efﬁcient and accurate inference on both image classiﬁcation and NLP benchmarks.
Generalizing our framework to design a multi-stage inference procedure is a natural direction for future research.
To enhance the applicability of our method to wider settings, another research direction would be to devise a principled approach for applying our method to other architectures including non-parametric models like k-NN, generalizing what we did for reading comprehension.
R EPRODUCIBILITY S TATEMENT
In Sec. 5 and Sec.
D (in the appendix) we provide detailed descriptions of the experimental setup used in this paper, including the pre-trained (teacher) models, student models, datasets and how to train and evaluate our two-stage framework.
Also, since we are only training small student models via standard and proposed distillation methods it is relatively cheap to reproduce our experimental results
A V ARIANTS OF CLASS - SPECIFIC DISTILLATION
Here, we present two additional variants of class-speciﬁc distillation that enable transferring the performance of a large teacher model to a lite student model on a subset of class.
Subsequently, the lite student can be employed in our two-stage inference framework.
A. I N - DOMAIN CLASS - DISTILLATION
Since we do not intend the student to make the ﬁnal prediction on an instance from [ L ] \L in , we can train the student to perform an |L in | =
| L (cid:48) | -way classiﬁcation.
In particular, given a teacher g and example ( x, y ) , we deﬁne a pseudo-label distribution:
[Equation]
Here, ˆ p L (cid:48) g,x denotes the softmax distribution restricted to L in :
[Equation] [Equation]
Now we get the student f L in ,L (cid:48) :
X → R L
[Equation]
by minimizing
[Equation]
A lite model f L in ,L (cid:48) trained with the loss in (13) aims to classify an instance from L in to the correct class.
At the same time, such a model is also trained to generate a non-informative uniform distribution 1 L (cid:48) · 1 on the instances from [ L ] \L in , which by deﬁnition corresponds to zero margin.
Now, given the student f L in ,L (cid:48) and teacher g , one can deﬁne a two-stage inference by using marginbased delegation.
In particular, for a test instance x ∈ X , the ﬁnal prediction becomes
[Equation]
where γ f L in ,L (cid:48) ( x ) is the margin assigned to x by the student f L in ,L (cid:48) .
A. I N - DOMAIN CLASS - DISTILLATION WITH ‘ ABSTAIN ’ OPTION .
We next explore another natural candidate for class-speciﬁc distillation, where the lite student aims to correctly classify an instance x from L in and declares to ‘abstain’ on the instances from [ L ] \L i .
To realize this, we train the student to perform an ( L (cid:48) + 1) -way classiﬁcation:
Given a teacher g and example ( x, y ) , we deﬁne a pseudo label distribution:
[Equation]
where ˆ p L (cid:48) g,x denotes the softmax distribution restricted to L (cid:48) classes, as deﬁned in (12).
Now, one can perform the distillation by utilizing (cid:101) p class ,L (cid:48)
+1 g,x , i.e., minimize
[Equation]
Note that (14) encourages the trained lite model f L in ,L (cid:48)
+1 to predict ( L (cid:48)
+ 1) -th class, i.e., ‘abstain’ class, for all instances from [ L ] \L in .
This leads to a natural abstain class-based delegation approach for two-stage inference.
Let
[Equation]
Then, given f L in ,L (cid:48) +1 and g , one makes the following ﬁnal prediction for a test instance.
[Equation]
In addition, analogous to (7), one can also deﬁne a two-stage inference procedure via margin-based delegation.
Remark 4.
Note that using an ‘abstain’ class is closely related to classiﬁcation with a reject option (see Sec. B).
However, as opposed to the traditional classiﬁcation with the reject paradigm, we intend to provide supervision for the so call reject class as well.
B C LASSIFICATION WITH A REJECT OPTION
There is a large literature on selective classiﬁcation, also known as classiﬁcation with a reject option or abstention Grandvalet et al.
(2009); Bartlett & Wegkamp (2008)
; Cortes et al.
; Geifman &
El-Yaniv (2017); Ramaswamy et al.
; Ni et al.
Here, one seeks a predictor f : X →
Y ∪{⊥} , where a prediction of ⊥ denotes the classiﬁer is uncertain.
To avoid the degenerate solution of abstaining on all samples, one assumes a ﬁxed rejection cost c ∈ (0 , 1] .
The goal is to then tradeoff the misclassiﬁcation error on non-abstained samples with the total cost incurred on abstained samples, i.e., minimize the loss
[Equation]
The Bayes-optimal classiﬁer for this objective abstains on samples with high uncertainty on the “true” label, i.e., Ramaswamy et al.
[Equation]
C V ARIANT OF MARGIN - BASED DISTILLATION : USING ABSTAIN CLASS
Another natural approach for margin-based distillation is to utilize an ‘abstain’ class to encourage student to not spend its model capacity on correctly classifying hard instances (where teacher achieves a low-margin).
Towards this, we can deﬁne the following pseudo label distribution for an example ( x, y ) .
[Equation]
g,
[Equation]
if
[Equation] [Equation]
otherwise .
Now, we distill a lite student model f L +1 ρ tr : X → R L +1 based on (cid:101) p margin g,x , i.e., we minimize
[Equation]
Note that we train the student to perform an ( L + 1) -way classiﬁcation, where it aims to classify an example in S easy to one of L original classes and the examples in S hard to the ‘abstain’ class.
Now, one may employ g and f L +1 ρ tr to enable a two-stage inference procedure that makes the ﬁnal prediction for a test instance x ∈ X as
[Equation]
otherwise , where ˆ y f L +1 ρ tr = arg max j f L +1 ρ tr ( x ) j .
D D ETAILS OF EXPERIMENTAL SETUP
Datasets.
We use three benchmark image datasets – CIFAR-100 (Krizhevsky, 2009), ImageNet
ILSVRC 2012 (a.k.a.
ImageNet-1k)
(Russakovsky et al., 2015), and ImageNet-21k (Deng et al., 2009).
CIFAR-100 contains 60k (50k train/10k test) images annotated with one of 100 object categories distributed uniformly.
ImageNet (ILSVRC 2012), on the other hand, is a much larger dataset with 1.33M (1.28M train/50k test) images annotated with one of 1000 object categories also distributed uniformly.
As for ImagenNet-21k, it originally contains images 12.8M from 21,843 classes.
We select 17,203 classes with at least 100 images and create a balanced test set with 50 images from each of the selected classes.
This remaining images from the selected classes provides us with a training set containing ∼ 11.8M images.
In addition to image datasets, we also evaluate the proposed distillation-based two-stage inference procedure on MNLI (Williams et al., 2018) and SQuAD datasets (Rajpurkar et al., 2016), which are
[Table] [Table]
standard sentence classiﬁcation and QA benchmarks, respectively.
MNLI dataset corresponds to a 3 -way classiﬁcation task with 392,702 training examples and a matched test set consisting of 9, instances.
As for the SQuAD dataset, it is a span selection task over a sequence length of 384 with 87,599 and 10,570 train and test examples, respectively.
Large teacher models.
For the image classiﬁcation tasks, we use EfﬁcientNet-L2 (Xie et al., 2020;
Tan & Le, 2019; Foret et al., 2021) which is the state-of-art for image classiﬁcation on multiple datasets.
It is an optimized convolutional neural network pretrained on both ImageNet and unlabeled
JFT-300M (Sun et al., 2017) with input resolution of 475.
EfﬁcientNet-L2 is a large model with 480M parameters which requires 478G FLOPs per inference and achieves an accuracy of 88.6% on
ImageNet.
For CIFAR-100, we used base EfﬁcientNet-L2 model with a new ﬁne-tuned classiﬁcation layer achieving an accuracy of 94.4%.
For the sentence classiﬁcation task, a RoBERTa-Large (Liu et al., 2019) model serves as a teacher, which is pre-trained on a general purpose text corpus of
GB size.
It is a transformer-based encoder model with 355M parameters that requires 155G FLOPs per inference while achieving an accuracy of 90.2% on MNLI.For the QA task requiring reading comprehension, we use T5-11B (Raffel et al., 2019) which exhibits competitive performance on a wide variety of NLP tasks.
It is a text-to-text transformer architecture pretrained on a subset of the common crawl with sequence lengths of 512.
T5-11B is a large model with 11B parameters and requires 6.6T FLOPs per inference achieving an accuracy of 90.2% on SQuAD.
Lite student models.
For CIFAR-100, we use small ResNet networks (He et al., 2016a;b) whose details are listed in Table 2.
As for ImageNet, we explore MobileNetV3 (Howard et al., 2019) of varying sizes as lite student models (see Table 3).
For NLP tasks, we use MobileBERT (Sun et al., 2020) as the lite student model which has 25M parameters and requires 13.5G FLOPs per inference.
Details of the experiment in Figure 1.
For this experiment CIFAR-100 dataset was used.
To sweep full spectrum all the way till teacher size, we used ResNet networks as detailed in Tableas the student networks and the largest ResNet-56 as the teacher network.
We used class-speciﬁc distillation from Sec 4.1.
E A DDITIONAL EXPERIMENTAL RESULTS
Class-speciﬁc distillation with class-based delegation.
Table 4 represents the two-stage performance (both in-domain and overall) realized by class-speciﬁc distillation approaches (cf.
Sec. 4.1) on ImageNet when we employ appropriate class-based delegation.
Class-speciﬁc distillation on ImageNet-21k dataset.
Figure 5 shows the performance of our proposed distillation-based two-stage inference framework on ImageNet-21k dataset.
As discussed in
Sec.
D, we work with 17,203 out of 21,843 classes originally present in the dataset.
Since we don’t have access to a high-performing teacher model on this dataset, we work with the oracle teacher (the one that knows the true label) to simulate the two-stage stage inference framework.
Also, while training a MobileNetV3-0.75 model as a student, we use the one-hot labels as the supervision for the classes in L in and utilize label-smoothing for the instances belonging to the remaining classes.
Accordingly, Baseline here corresponds to a normally trained student (MobileNetV3-0.75 model) combined with the oracle teacher.
Margin-based distillation with margin-based delegation.
Here, we study both overall and in-domain performance of two-stage inference enabled by margin-based distillation from Sec. 4.
[Table]
[Figure]
Fraction processed by student
[Figure]
As described in Sec.
D, we work with a subset corresponding to 17,203 out of 21,843 classes.
Baseline denotes the standard distillation from (3) with a = 0 and b = 1 .
CD -I denote the classspeciﬁc distillation approaches deﬁned in Sec. 4.1.
with |L in | =
L (cid:48)
= 1000 .
Accordingly, in-domain accuracy is computed via focusing on the test instances from the 1000 classes in L in .
Here, each lite student model (a MobileNetV3-0.75 model) employs margin-based delegation.
on both CIFAR-100 (cf.
Fig. 6) and ImageNet (cf.
Fig. 7).
Here, in-domain instances correspond to those instances where the student assigns a margin of at least ρ = 0 .
As evident, margin-based distillation also leads to improved in-domain two-stage accuracy.
However, Baseline and MD from (17) have very similar performance on both datasets.
[Figure]
[Figure]
based distillation approach from Sec.
4.2 on ImageNet.
Here, each lite student model (a MobileNetV3-0.75 model) employs margin-based delegation .
In-domain accuracy is computed on the test instances where student achieves a margin ρ of at least 0.4.
[Figure]