When in Doubt, Summon the Titans: A Framework for Efficient Inference with Large Models
Authors: Ankit Singh Rawat
Manzil Zaheer
Aditya Krishna Menon
Amr Ahmed
Sanjiv Kumar
Venue: ICLR 2022 Submitted (2022)
Keywords: Distillation
Large models
Efficient inference

Abstract: 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 benefits of the large models, while largely preserving the computational benefits 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 efficiently 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 benefits of our approach on both image classification and natural language processing benchmarks.

Under review as a conference paper at ICLR 2022
WHEN IN DOUBT, SUMMON THE TITANS: EFFICIENT INFERENCE WITH LARGE MODELS
Anonymous authors Paper under double-blind review
ABSTRACT
Scaling neural networks to "large" sizes, with billions of parameters, has been shown to yield impressive results on many challenging problems. However, the in- ference 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 dis- tillation that realizes the modelling benefits of the large models, while largely pre- serving the computational benefits of inference with more lightweight models. nutshell, we use the large teacher models to guide the lightweight student models to only make correct predictions on subset of "easy" examples; for the "hard" examples, we fall-back to the teacher. Such an approach allows us to efficiently employ large models 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 stan- dard distillation. Empirically, we demonstrate the benefits of our approach on both image classification and natural language processing benchmarks.
INTRODUCTION
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 classification (Kolesnikov et al. 2019 Xie et al. 2020 Tan & Le 2019 Foret et al., 2021 and segmentation (Ghiasi al.) (020)) 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 tera-FLOP for 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 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 deployment over resource-constrainee edge devices such as mo- bile phones (Zhang al. 2020). As 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 al.
1989 Hassibi &
Stork 1993), and distillation (Bucilà et al. 2006) Romero et al. 2014]] Hinton et al. 2015). However,
there is limit to how far such model compression techniques can be pushed reduce inference
cost while retaining good performance across all inputs (cf. teacher-student accuracy gaps in (Cho)
& Hariharan 2019 Menon al. 2020b Mirzadeh et al.) 2020 Wang 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 com- plexity 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 computation of large model is wasteful; rendering the overall ML system inefficient. Such an inef- ficient 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 the tail.
"easy", for which performing full

---

Under review as a conference paper at ICLR 2022
Two-Stage Framework
Compute Cest [MLLPPS)
150
Increasing share of examples cassified by
Einil prediction
0.75
Studant now Too complex
0.75
0.75
Test
Student (Lite model)
confident
100
75
classifed
axamples
0.75
0.750.73
0.75
0.73
0.75
0.74
50-
Teachee heee
No
0.7
25
0.67
Stage
Sumtog the Taans
Teacher (Large model)
Emil prediction
1.0.00
F0.58
Studen: Only
14
32
Student Size [Layers]
Figure I: Leftt: A schematic of the proposed two-stage inference framework where after distillation retain both the lightweight student model and the large teacher model. At inference time if the student finds an instance hard, we fall back to the teacher for prediction. Right: The two- stage framework trade-off size of the student, thereby reducing the overall computation in expectation and achieving better accuracy compared performing inference based only the student. Note that the numbers annotated at each (student size, total compute cost) point on the plots denote the overall accuracy of the corresponding setup. For the two-stage inference, always achieve an accuracy of 0.75 by delegating an appropriate fraction of instances to the teacher. Compared to classifying all examples using the teacher, computation cost savings come from those instances where the student makes the final prediction.
CIFAR-100 image classification task using ResNets allows to aggressively
increase the size of the student,
can
In this paper, we focus realizing the benefits 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 novel distillation basedd two-stage inference framework in Figure [ (left): First use a lightweight student model to make prediction. If the student is confident, we emit the prediction and we want the student to be confident on all the easy instances, which should be large fraction of 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 test time queries. When the student
The instance-aware two-stage inference mechanism crucially relies on the ability of the student model detect the "hardness" of an input instance on the fly and routing it to the large model. To enable this routing, we propose modified distillation procedures. In particular, we employ novel distillation loss functions (ef. 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 confident, 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 allows us to much more aggressively trade-off size of the student for multiple image classification and natural language processing (NLP) benchmarks. Interestingly, as summarized in Figure [ there a sweet spot where we can achieve the same ac- curacy as the teacher with 45% less compute. This benefit is further magnified 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 final predictions. We would like to highlight that our goal is not to train 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 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 shared large teacher model running in the cloud.
RELATED WORK
Techniques to reduce inference cost for deep models mainly fall under two different approaches: quantization and pruning, and adaptive computation.
2

---

Under review as a conference paper at ICLR 2022
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. [989 Hassibi & Stork [993] Li et al. 2020;
2017, Howard et al. 2019).
Sigmificant progress was made by introducing Huffman encoding methods for non- non-uniform quan- tization which led to reduction in network sizes by orders of magnitude and up to 4x reduction in overall prediction cost (Han et al. 2013). 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. 2222I) In the NLP domain, Gordon et al. 2020); Zadeh et al. (2020) proposed pruning BERT during train- ing, which resulted in 30%--40% reduction in model size with minimal effect on the accuracy of the final 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 beneficial, 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 al, 2019 Raganato et al. 22200 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. (1998) and later formalized by Viola & Jones (000I). The basic idea was to design a cascade of independent classifier and reject early on and cheaply. The idea has been generalized from linear chain of cascaded classifiers to trees (Xu et al. 2014). Instead of combining many independent classifiers, 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 scalar halting probability, called the "ponder time". to dynamically modulate the number of computational steps needed for each input. An alternative approach represented by Adaptive Early Exit Networks (Bolukbasi al.] 2017), which gives the network the ability to exit prematurely i..., not computing the whole hicrarchy of layers - no more computation is needed. A modern incarnation of this approach in NLP with transformers encoders appeared in (Schwartz et al. (2020); Liu et al. 22020); Dabre et al. (2020). 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. (2020). As a further generalization, [aanna 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 a task and rely on designing the whole pipeline from scratch which can be expensive if we want to achieve state-of-the- the-ttr-tt (SoTA) performance. In contrast, 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 benefit from its large model capacity and a lite distilled model suffices to predict a large portion of test instances.
specialized to
want to design efficient inference
BACKGROUND
3.1 MULTICLASS CLASSIFICATION
Consider standard multiclass classification problem where given instance x E X, the objective is to classify the instance as a member of one of the L classes, indexed by y 2 [L]. In the most common setting, given a training set comprising of n instance and label pairs or training examples S, = (())) 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 dN - arg max [(x); Accordingly, the model (top- 1) accuracy is defined as
((Ii, y4) sejnj one learns a classification model f X - RR-, where f(x)
P Øリ- = - EExy [1 [1.-yl = 1 Ex,y [1g.4yl-
(1)
3

---

Under review as a conference paper at ICLR 2022
Ideally, one prefers classification model with a high accuracy. Let é: R RR be surrogate loss function such that Ey. f()) closely approximates the misclassification error Ly Typically, given the training set S, and the loss function f. one selects desired classification model via empir- ical risk minimization (ERM). Softmax cross-entropy loss is one of the most widely used surrogate loss functions for multiclass classification: Given model scores f(x) (f(x)....... (()))) one computes the softmax distribution
(t) c R()); E J et,,,,,, for E [L],
(2)
where denotes the temperature parameter of the softmax operation D Further, we define Py E {0, 1 to be the one-hot label distribution corresponding to the true label E [L], which has non-zero value only y-th coordinate. Now, softmax cross-entropy loss corresponds to the distance between the distributions Py and f5.a. as measured by the cross-entropy function.
f(x)) 11 Pr
- Po(i) loggogggr.. (t)
3.2 MODEL DISTILLATION
Distillation is celebrated training techniques that utilizes one model's scores to train another model (Bucilà 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 teacher model g X + H and an example (a, y) E X x y, one first defines the teacher (soft- max) distribution (t) per (2). Now, as opposed to utilizing the 'one-hot' label distribution py, we treat Py.x(() as the pseudo label distribution and define the distillation version of the softmax cross-entropy loss for the student model ʃ : X - IR- as frtistin (g()), [(x)) = (fg Pf.z). For b E R4, distillation involves minimizing
+ Σ selm) (( f(x:))
"
(1))), ni E [E|n)
b
Note that the objective in (3) utilized both the true labels {2:} and the teacher scores (z)))) When b = 0, this reduces to the standard training. More interestingly, when a = 0. (3) correspond to training with solely Pg (i) fairly common way to utilize distillation (Menon al. 2020a).
Remark 1. For distillation, teacher models are usually much more complex and powerful as com- pared to student models (Bucilà et al., 2006 Hinton et al., 2015). However, the distillation with an cqually 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
4 DISTILLATION FOR TWO-STAGE INFERENCE
We now propose various distillation approaches to power the proposed two-stage inference frame- work. 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 aid our objective as typically distillation envisions the student model to be used as a standalone model.
Towards this, we explore general distillation framework that partitions the training examples S into two groups: 1) easy instances Seasy and 2) hard instances Shard- Accordingly, we modify the distillation process such that the student incurs larger loss when it makes incorrect predictions Seasy. Furthermore, we penalize the student less for making mistakes on Shard, which we accomplish by carefully designed supervision during the distillation. Now, we present two specific realizations of the above generic distillation framework based on two different strategies to partition the training examples into Soasy and Shard-
4.1 CLASS-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
For brevity, do explicitly represent the temperature parameter in the rest of the paper.
4

---

Under review as a conference paper at ICLR 2022
few instances. This has sparked a long line of work improving model performance the tail classes (see, c.g.. Cui et al. 2019 Cao et al.] 2019 Kang et al., 2020; Menon et al., (021). ). In contrast, explore an orthogonal direction and leverage the data imbalance to enable efficient inference with large models. In particular, for Lin _ [L], we define
Seasy = {(2.)) y) E y:ye E Cin} and Shard = Si,Seayy = {(x,)) E S & C.n..
Thus, we require the lite model to perform well only on a subset of classes Lin and route the exam- ples 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 subset of classes. Now, setting Lin 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 the exact implementation details described later) are sent to the large model during inference.
With this general approach in mind, we propose class-specific distillation approach Now, given a large teacher model and example (z, y), we define a pseudo label distribution gg.
prlasss as follows:
[pg.r p,
if y E Lins if E [L]
glaass Pgg.
(4)
where a and Py denote the teacher's softmax distribution and the one-hot label distribution, spectively, In addition, E [0, 1] denotes a label-smoothing parameter. Now, we train lite student f: X + RR' with the distillation loss
elem f()) H (Pylars .Pj.)).
(5)
Note that the loss in 5 utilizes teacher softmax distribution for classes in Cin and relies label- smoothed one-hot distribution for the remaining classes. This loss has two desirable properties for 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 Cin- 2) Due to standard label-smoothing (Szegedy et al. 2016 Müller et al.) 2019), the lite student achieves smaller margin on the examples belonging to the classes in Lin.
Let f.... be the lite student model obtained by minimizing the loss in (5). Now, given a test instance a X, we first run inference with the lite model to obtain fc,, (z). Subsequently, we decide if we need to make the final prediction based on Se.. (x) or delegate the example to the large teacher g to obtain the final prediction. We identify two useful delegation schemes, which we detail next.
Class-based delegation. Recall that the class-specific distillation aims to utilize the lite model to classify only the examples belonging to the classes in Lin. Thus, the prediction made by fc.. serves as a natural candidate for the delegation. In particular, when
Matudent (x) A arg max Sea (2) E Lin;
we declare Istudent (I) as the final prediction. Otherwise, x is sent to the teacher and Dteacher (x) arg max g(x) becomes the final prediction.
Margin-based delegation. The class-based delegation designed under the assumption that the lite model achieves high accuracy on the examples belonging to the classes in Lin and identifies the examples from Cin with high fidelity. Both of these assumptions don't always hold in practice. In particular, fe... may find some instances from Lin hard and incorrectly predict a wrong class in Li when presented with those instances. Similarly, f'.,, may predict class from Lin when the test instance belongs to Lij.
Recall that the distillation loss in 55) is designed to ensure that fc.. cell--caibbrtted on the instances from Lin; as a result, it attains small margin on hard instances from Li- Moreover, by design, fc... realizes small margin on the instances from [L], Thus, one can utilize the margin of the lite model fc.. to perform delegation. Towards this, recall that
Me. (x) 2 PSec s ([4) - -Pfe ((2))
(6)
In what follows, without loss of generality, assume Lin [T']] for T.' S T.. One can easily express approach for general Lin with slightly cumbersome notation.
5

---

Under review as a conference paper at ICLR 2022
denotes the margin realized by fei. Here, PIc., .z ((i) denotes the i-th largest element of the vector Pfa,, Now, given design parameter p (0, 1), assign the following final prediction to test instance E X.
Ustudenee (x) Meam (x) 2 P. if Mean (x) <
((x) -
Breacher (x)
The margin- bargin-based delegation aims to ensure that hard instances from Lin as well as all instances in [L]) get routed the teacher for the final prediction.
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 lite student. We discuss two other class-specific distillation-basee approaches in Sec. Afof the appendix and evaluate those in Sec. [
Remark 2.
4.2 MARGIN-BASEE BASED DISTILLATION
As discussed before, the margin assigned to an instance by a model is a natural proxy for the hardness of the instance, viewed by the model. In Sec. Sec.14.1| to delegate the examples to the teacher. This raises an interesting question one can utilize the margins to partition the training data into Seasy and Shard example during the distillation. Towards this, given a teacher and a parameter E (0, 1). we add a training example (xr, y) to Seasy iff
utilize the margins assigned by the student
(o()) - p ((1) - Po. (12) >
Ptr-
Given this data partition, for an example (x, y). we define the pseudo label distribution as follow:
(pg. (1), .I).....) (L)) if (a. y) Seasy otherwise.
Imargin.L
(9)
(11 aa) .ay p +
Now, we obtain a (lite) student f/e X RR by minimizing
gmargin.T. (g(x), I (xx) = ( ,pf.)) distill
(10)
For the two-stage inference, given a test instance we make the following final prediction.
Ösz = max, f/r (x); Bteacher (x)
if Ysk (x) 2 P. otherwise,
(xx)
where Úsa. arg max, f/r. (z)j-
Remark 3. For small value of Per. 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 Sec. 3.2, On the other hand, for large values of Ptz. the student is expected to do well on only small subset of examples during training and be able identify the rest of the examples as hard instances that need to be routed to the teacher.
Similar to class- distillation, there are multiple potential variants of the margin-based distil- lation. We discuss one such variant that relies on an 'abstain" class in Sec. C of the appendix.
EXPERIMENTS
We now conduct comprehensive empirical study of our distillation- based two-stage inference pro- cedure. In particular, evaluate various distillation frameworks introduced in Sec. along with different choices of delegation methods. On standard image classification tasks, we establish that:
(i) large improvement in accuracy over the student-only approach can be realized by delegating only small fraction of examples the large teacher model (Sec. 5.T). This validates claim that one can rely on a much smaller student and achieve an accuracy similar to the teacher-only approach with small increment in inference cost.

---

Under review as a conference paper at ICLR 2022
0.9
0.9
0.9
0.8-
0.8
0.8
Seashre
Jeae
Seasihe
2--1 c1
CD (D-1am00.31
CD: (ymo.c)
(b) (2-0
(b-1
0.6
0.6
0.6
0.0
0.2
0.4
0.6
0.8
1.0
0.2
0.6
0.8
1.0
0.5
1.0
1.
2.0
2.5
Fraction processed by student
Fraction processed by student
Total Latency (5)
(a) Overall accuracy
(b) In- in-domain accuracy
(c) Latency vs. in-domain accuracy
Figure 2: Comparison of various class- distillation methods on CIFAR-100. Baseline denotes the standard distillation from (3) CD-1, CD-II, and CD-III denote the class-specific distillation ap- proaches defined in Sec. 4.1 Sec. A.1 and Sec. A.2 respectively, with Liin I' 30. Accord- ingly, we compute in-domain accuracy on test instances from the 30 classes in Lin- Here, each lite student (ResNet-32) employs margin-based delegation the teacher. The right-most plot depicts the (inference) latency vs. in-domain accuracy trade-off for the two-stage inference procedure.
(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 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-worke settings where a large data imbalance exists in favor of in-domain instances (Sec. (.2)).
(iii) Class-specific distillation defined by (5) indeed achieves the desired behavior where the student enables 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 large teacher model, proposed class-specific distillation and a variant of margin- based distillation enable efficient inference on sentence classification and reading comprehension tasks in the NLP domain, respectively (Sec. (.3
We mainly focus on three benchmark image datasets CIFAR-100 (Krizhevsky) 2009), ImageNet- Ik (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 classification task based on MNLI (Williams et al. 2018) and reading comprehension task based SQuAD dataset (Rajpurkar) et al. 2016). On image classification tasks, we use EfficientNet-L. (Xie et al. 2020) as the large teacher. As for the lite student, utilize ResNet (He et al.) 2016a by and MobileNetV3 (Howard
ct al. 2019) for CIFAR- and ImageNet, respectively. For the sentence classification and reading (Liu al. (2099) and T5-1 IB (Raffel al.) 2019) serve as
comprehension tasks, RoBERTa- the teachers, respectively. We use MobileBERT (Sun et al. 2020) as the lite student for both MNLI and SQuAD. We provide detailed description in Sec. D of the appendix.
5.1 OVERALL ACCURACY GAINS
We begin by establishing the utility of the two-stage inference procedure for overall performance improvement. Towards this, Fig. 2a Bal and (in appendix) depict the two-stage overall accuracy achieved by various class-specift distillation approaches via margin-basse based delegation. Besides, also include conventional model distillation as Baseline in our evaluation. Note that, for all ap- proaches, allowing the lite student model delegate hard instances (the gin) to the large teacher significantly 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 ~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 Lin. leading to poor overall accuracy unsurprisingly, as student delegates more instances to the teacher, the two-stage overall accuracy improves. Similar conclusions also hold ImageNet datasets (cf. Fig. 3a and 5a Table 1 and (in the appendix) show that the class-specific distillation based two-stage inference leads to similar improvements in the overall accuracy when employ class-based delegation from Sec. 4.1
with small mar-
the entire test set. However,

---

Under review as a conference paper at ICLR 2022
0.90
accurscy
0.80
0.75
Eanskie
CD-1 LOnO 05
Reasirg
CD
(cm0.))
Seasihe
2D: (ymo.c)
CD Jeet
(b: (y-0-))
0.70
(sm0.5)
0.65
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.5
1.0
1.5
2.0
2.5
Fraction processed by student
Fracbon processed by student
Total Latency [s]
(a) Overall accuracy
(b) In-domain accuracy
(c) Latency vs. in-domain accuracy
Figure 3: Comparison of various class-specific distillation methods ImageNet- Ik. Baseline de- notes the standard distillation from (3) CD-I and CD-III denote the class-specific distillation ap- proaches defined in Sec. 4.1 and Sec. A.2 respectively, with Linl L' 300. Accordingly, compute in-domain accuracy on the test instances from the 300 classes in Lin- Here, each lite student (MobileNetV3-0075) employs margin-based delegation. The right-most plot depicts the (inference) latency vs. in-domain accuracy trade-off for the two-stage inference procedure.
5.2 IN-DOMAIN PERFORMANCE
In-domain
Overall
Approuch
Next, we verify that two-stage dis- tillation is more beneficial in those real-world settings where MI. models encounter heavily imbal- anced data during the inference. Ide- ally, the lite student model should make a prediction on frequent but easy instances and delegate rare but
Accuracy
Fraction
Accuracy
Fraction
Buseline
0.71
1.00
0.72
1.00
ResNet-32
CD-1 (a 0.0)
0.74
0.88
CD-I (a 0.6)
0.88
0.84
0.86
0.32
CD-II
0.78
1.00
0.24
1.00
CD-III
0.69
0,90
Baseline
0.75
1.00
0.75
1.00
ResNet-56
CD-I 0.0)
0.89
0.77
0.90
0.27
CD-1 (a 0.63
0.90
0.82
0.88
0.29
hard instances to the large teacher. This would ensure that the two-stage inference procedure realizes high overall accuracy (compared to infer- ence with only the student) while significantly 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 in- domain (so-called easy instances). By design, for class-specific distilla- tion, in-domain instances belong to the classes in Lin- Similarly, for the margin-based distillation, in-domain instances are the ones where teacher assigns a large margin p.
CD-II
0.80
1.00
0.24
1.00
CD-III
Table 1: Performance of two-stage inference procedure on CIFAR-100, The student employs class-specific distillation with | L' 30, with in-domain referring the stances from the classes in Lin- During inference we, use an appropriate flss-based delegation method. See Fig. for the identity of the distillation approaches. Fraction denotes the fraction of test instances where the student model makes the final prediction. Unlike margin-based delegation, for given teacher and student models, we obtain a single value of (Ac- curacy, Fraction) tuple with class-based delegation.
For C1FAR-100, Fig. 2b shows that two-stage inference achieves much better in-domain (defined by I = 30 classes) accuracy as compared to overall accuracy. This implies that a real-world setting where in-domain instances are frequent, the two-stage inference 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, have separately highlighted the performance of two-stage inference the in-domain instances. This shows that the more imbalanced the data is the more advantageous two-stage distilla- tion would be in terms of saving the inference cost without affecting the classification performance. As per Fig 2c we essentially maintain same accuracy as the teacher but reduce latency by more than 2x. Here, 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.
efficiently achieve very
Interestingly, Fig. indicates that CD-I with margin- based delegation gives a favorable tradeoff be- tween in-domain and (balanced) overall accuracy. The label-smoothing parameter a in (4) helps create dichotomy between in-domain and out-of-dommin instances. The larger values of a ensure that the student assigns a smaller margin to out-of-domait examples and tries to improve its perfor- mance in-domain instances. This is reflected in Fig. 2b where CD-I with large achieves higher
8

---

Under review as a conference paper at ICLR 2022
0.950
0.90
Accurgyy
0.925
0.88
0.50
Bxoe the
MNLI
0.500
Saoeliee te
MNU
0.06
SQUAD
Baseine tee
0.96
SQUAD
Sa ne
0.54
No Lomo
0.94
(om)
5.00 0.25 0.50 0,75 1.00
0.00
0.25 0.50 0.75 1.00
0.00
0.25 0,50 0.75 1.00
0.00 0.25 0.50 0.75 1.00
Fraction processed at stusent
Fraction processed at student
Fraction processed at stacent
Praction processed at student
(a) Overall accuracy
(b) In- nn-domain accuracy
(c) Overall accuracy
(d) In-domain accuracy
Figure 4: Performance of the proposed two-stage inference on NLP tasks. Left half: Overall and in-
domain accuracy of the margin- based distillation coupled with margin- baggid delegation on MNLI.
Here baseline corresponds to the two-stage inference enabled by the standard distillation (cf. (3) Right half: Overall and in-domain accuracy (exact match) comparison of margin- bargin-based based distillation on SQuAD task. Here, Baseline denotes a normally trained student used in combination with the teacher. Note that MD denotes a margin-based distillation with label smoothing.
accuracy by predicting a larger fraction of in-domain instances at the student. We note the similar trend in Table where we combine class-specific distillation with class-based delegation
We also evaluate the in-domain two-stage performance of class-based distillation on ImageNet- Ik with Lin| 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. k. Finally, we also studied the in-domain performance of two-stage inference enabled by margin-basee based distillation from Sec. 4.2 on both CIFAR- and ImageNet- Ik in the appendix.
5.3 TWO-STAGE INFERENCE IN NLP DOMAIN
Sentence classification task. Fig, 4a and 4b show the performance of proposed two-stage in- ference framework on MNLI. Since it has only 3 classes, we employ margin-based distillation (with Per 8.0) from ( with margin- based delegation. Our conclusion for the image classification also extends to the text domain and the two-stage inference framework enables a lite student to leverage high-quality teacher by delegating 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. Ba).
Reading comprehension task. To showease the generality of our approach. we apply the two- stage 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 suitable modification 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, hot labels for the start and end span as no teacher logits are available. Whereas for hard examples, still employ label smoothing. We try two values a 0.2, 0.4 in the experiments. Our results in Fig. [d) and 4d] which has similar conclusions to the classification experiments. See Fig. [8b] for the inference- latency vs. in in-domiin performance trade-off achieved by the two-stage inference.
one
DISCUSSION
We propose a distillaaionnbbased two-stage inference framework efficiently 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 easy (in-domain) instances and can identify hard (out-of-domain instances. When deployed in tandem, the lite model generates the final 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 efficient and accurate inference on both image classification 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 wider settings, another research di- rection would be to devise a principled approach for applying our method to other architectures including par-meraettric models like k-NN, generalizing what we did for reading comprehension.

---

Under review as a conference paper at ICLR 2022
REPRODUCIBILITY STATEMENT
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 two-stage framework. Also, since we only training small student models via standard and proposed distillation methods, it is relatively cheap to reproduce our experimental results.
REFERENCES
Ankur Bapna and Orhan Firat. Exploring massively multilingual, massive neural machine trans- lation. In Google AI Blog, Oct 2019. URL htt.:httts//// htts:///a...ooo.... goog lellog.lloo...oo..... com www.c0mmmm100cc/ 201 9/10/ i html
sxploring-masssiellyy
Ankur Bapna, Naveen Arivazhagan, and Orhan Firat. Controlling computation versus quality for neural sequence models. arXiv preprini arXiv:2002.07106, 2020.
Peter L.. Bartlett and Marten H. Wegkamp. Classification with reject option using hinge loss. J. Mach. Learn. Res., 1211823-8444 June 2008. ISSN 1532-4435.
Tolga Bolukbasi, Joseph Wang. Ofer Dekel, and Venkatesh Saligrama. Adaptive neural networks for efficient inference. In Proceedings of the 34th International Conference on Machine Learning Volume 70, ICML'11,, pp. 527-536. JMLR.org, 2017.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Ma- teusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCan- dlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models learners. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Ad- vances in Neural Information Processing Systems, volume 33, pp. 1877-1901. Curran Asso- ciates, Inc., 2020. URL httpss:/pproceeeddgss..ee wwtpeeewwwr...eeeeeeeeeeeeeeeeee ccc/paper//22///1eeee pdf
few-shot
Cristian Bucilă, Rich Caruana, and Alexandru Niculescu-Mizil Model compression. In Proceedings of the 12th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, New York, NY, USA, 2006. ACM.
KDD '06, pp. 535-535-11
Kaidi Cao, Colin Wei, Adrien Gaidon, Nikos Arechiga, and Tengyu Ma. Learning imbalanced datasets with label-diisributtionaaaaee margin loss. In Advances in Neural Information Processing Systems, 2019.
Miguel A Carreira-Perpinae Model compression constrained optimization, with application to neural nets, part i: General framework. arXiv preprint arXiv:1707.01209, 2017.
Jang Hyun Cho and Bharath Hariharan. On the efficacy of knowledge distillation. In Proceedings of the IEEE/CVF International Conference on Computer Vision. pp. 4794 4802, 2019.
Junyoung Chung, Sungjin Ahn, and Yoshua Bengio. Hierarchical multiscale recurrent neural net- works. arXiv preprint 160.v0 704, 2016.
Corinna Cortes, Giulia DeSalvo, and Mchryar Mohri. Boosting with abstention. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 29, pp. 1660-1668. Curran Associates, Inc., 2016.
Daniel Crankshaw, Xin Wang, Giulio Zhou, Michael J. Joseehh E. Gonzalez, and Ion Stoica. Clipper: A low-latency online prediction serving system. In Proceedings of the 14th USENIX Conference on Networked Systems Design and Implementation, NSDI' 17, pp. 613-627, USA, 2017. USENIX Association. ISBN 97819311 7781931971379
Yin Cui, Menglin Jia, Tsung-Yr Yi Lin, Yang Song, and Serge Belongie. effective number of samples. In CVPR, 2019.
loss based on
10

---

Under review as a conference paper at ICLR 2022
Raj Dabre, Raphael Rubino, and Atsushi Fujita. Balancing cost and benefit with tied-multi trans- formers. In Proceedings of the Fourth Workshop on Neural Generation and Translation, pp. 24- 34, Online, July 2020. Association for Computational Linguisties doi: 10. 10.18653/w1/22200mgg-- aattps://ww.aclveett http://www.....ooo//////ooooo///
2gg.:ahnoo..og222....gg.1 ngt-1.3
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hier- archical image database. In 2009 IEEE Conference Computer Vision and Pattern Recognition, pp. 245,255, 2009. doi: 10.1109/CVPR.2200..52 2206548
Maha Elbayad, Jiatao Gu, Edouard Grave, and Michael Auli. Depth-adaptive transformer. In In- ternational Conference Learning Representations, 2020. URL https::t////.. Dtctps://openreviee
Angela Fan, Edouard Grave, and Armand Joulin. Reducing transformer depth on demand with structured dropout. arXiv preprint arXiv: /909. 1556, 2019.
Zhou Fang, Dezhi Hong, and Rajesh K. Gupta. Serving deep neural networks at the cloud edge for vision applications on mobile platforms. In Proceedings of the 10th ACM Multimedia Systems Conference, MMSys '19, pp. 36-47, New York, NY, USA, 2019. Association for Computing Machinery. ISBN 978 450362979. doi: 10.1145/33041099330222
URL https://doi.
hxg/10.1145/20.
3306221
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. arXiv preprint arXiv:2101.03961
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimiza- tion for efficiently improving generalization. In International Conference on Learning Represen- tations, 2021. URL httpttttt///prreeeeeee
net forum 1eniidd-mmm wtw.idddmmmmmmss.ommmmmmmmm 1rM
Tommaso Furlanello, Zachary Chase Lipton, Michael Tschannen, Laurent Itti, and Anima Anand- kumar. Born-again neural networks. In Proceedings of the 35th International Conference on Machine Learning. ICML 2018, pp. 1602-1611, 2018.
Prakhar Ganesh, Yao Chen, Xin Lou, Mohammad Ali Khan, Yin Yang, Deming Chen, Marianne Winslett, Hassan Sajjad, and Preslav Nakov. Compressing large-scale transformer-based models: A case study on bert. arXiv preprint arXiv:2002./1985 2020.
Yonatan Geifman and Ran Ea-ia..v. Selective classification for deep neural networks. In Proceed- ings of the 31st International Conference on Neural Information Processing Systems, NIPS' 17, pp. 4885 4894, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 9781510860961
Golnaz Ghiasi, Yin Cui, Aravind Srinivas, Rui Qian, Tsung-Yi Yi Lin, Ekin D Cubuk, Quoc V Le, and Barret Zoph. Simple copy-paste is strong data augmentation method for instance segmentation. arXiv preprint 27X7::201 2020.
Mitchell Gordon, Kevin Duh, and Nicholas Andrews. Compressing BERT:
Studying the effects of weight pruning on transfer learning. Proceedings of the 5th Workshop on Representation Learn- ing for NLP, pp. 143-155, Online, July 2020. Association for Computational Linguistics. doi: 10. 10.18653/1120222rrp44.l. 1.18. URL https: http://wwccccccoccccc aclweb org/ owtp///ootoooooog///22.......
rep14nip-1.12
Yves Grandvalet, Alain Rakotomamonjy, Joseph Keshet, and Stéphane Canu. Support vector ma- chines with a reject option. D. Koller, D. Schuurmans, Y. Bengio, and L. Bottou (eds.), Ad- vances in Neural Information Processing Systems, volume 21, pp. 537-544. Curran Associates, Inc., 2009.
Alex Graves. Adaptive computation time for recurrent neural networks. arXiv preprint wwXiv:1603.08983 2016.
Song Han, Huizi Mao, and William Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. arXiv preprint 1rxiv:1510.00144 2015.
Babak Hassibi and David G Stork. Second order derivatives for network pruning: Optimal brain surgeon. Morgan Kaufmann, 1993.
11

---

Under review as a conference paper at ICLR 2022
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2016a.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In Bastian Leibe, Jiri Matas, Nicu Sebe, and Max Welling (eds.), Computer Vision ECCV 2016, pp. 630-645, Cham, 2016b. Springer International Publishing. ISBN 978-3-319. 46493-0.
Geoffrey E. Hinton, Oriol Vinyals, and Jeffrey Dean. Distilling the knowledge in neural network. CoRR, 503//1503,0231,, 2015.
Andrew Howard, Mark Sandler, Grace Chu, Liang-Chieh Chen, Bo Chen, Mingxing Tan, Weijun Wang, Yukun Zhu, Ruoming Pang, Vijay Vasudevan, Quoc V. Le, and Hartwig Adam. Searching for mobilenetv3. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), October 2019.
Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Mia Xu Chen, Dehao Chen, Hy- oukJoong Lee, Jiquan Ngiam, Quoe V Le. Yonghui Wu, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. arXiv preprint arXiv: 1e111111100055 2018.
N. P. Jouppi, C. Young, N. Patil, D. Patterson, G. Agrawal, R. Bajwa, S. Bates, S. Bhatia, and N. Boden et al. In-datacenter performance analysis of tensor processing unit. In 2017 ACM/IEEE 44th Anmual International Symposium on Computer Architecture (ISCA), pp. 1-12, Los Alamitos,
CA, USA, jun 2017. IEEE Computer Society. doi: 10.1145/309985663000888 3080246.. URL https: 3080246
7/doi. www.computeeee.................e www..eeeeettttesssstett.t..... wrw.ccett.ororooooooooooo
Bingyi Kang, Saining Xie, Marcus Rohrbach, Zhicheng Yan, Albert Gordo, Jiashi Feng, and Yan- nis Kalantidis. Decoupling representation and classifier for long-tailed recognition. In Fighth International Conference on Learning Representations (ICLR), 2020.
Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Joan Puigcerver, Jessica Yung, Sylvain Gelly, and Neil Houlsby. Big transfer (bit): General visual representation learning. arXiv preprint 6(2):8, 2019.
Olga Kovaleva, Alexey Romanov, Anna Rogers, and Anna Rumshisky. Revealing the dark secrets of bert. arXiv preprint 1rXiv:1908.0859; 2019.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
Yann LeCun, John S Denker, Sara A Solla, Richard E Howard, and Lawrence D Jackel. Optimal brain damage. In NIPs, volume 2, pp. 198-605. Citeseer, 1989.
Zhuohan Li, Eric Wallace, Sheng Shen, Kevin Lin, Kurt Keutzer, Dan Klein, and Joey Gonzalez. Train big, then compress: Rethinking model size for efficient training and inference of transform- ers. Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. PMLR, 13-18 Jul 2020. URL htRl.ttt::/// htttp::/pproceeiiiiei ngs. ml r www.r.peess/ve httpss/vvv1///1 Om. html
Tailin Liang, John Glossner, Lei Wang, and Shaobo Shi. Pruning and quantization for deep neural network acceleration: A survey. arXiv preprint wxXiv:2101.0967. 2021.
Weijie Liu, Peng Zhou, Zhiruo Wang, Zhe Zhao, Haotang Deng, and Qi Ju. FastBERT: a self- distilling BERT with adaptive inference time. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 6035-6044, Online, July 2020. Association for Computational Linguistics. doi: 10. 10.18666 11183/6//2222..ccl mm.n........m URL https::/wwwg
aclweb. orq/ http:///athhoooyy2200 aal.020.cccl1o manl--mainn... 537
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. RoBERTa: A robustly optimized bert pretraining approach. arXiv preprint arXiv: /907. 1192,, 2019.
Aditya Krishna Menon, Ankit Singh Rawat, Sashank J. Reddi, Seungyeon Kim, and Sanjiv Kumar. Why distillation helps: a statistical perspective. CoRR, abs/2005. 10419, 2020a. URL https: 2001/
w/arxiv.org/aana wwt.:/orggabaaae.. wttps//wwww//ba///......o v.
12

---

Under review as a conference paper at ICLR 2022
Aditya Krishna Menon, Ankit Singh Rawat, Sashank J Reddi, Seungyeon Kim, and Sanjiv Kumar. Why distillation helps: a statistical perspective. arXiv preprint arXiv:2005. 10419, 2020b.
Aditya Krishna Menon, Sadeep Jayasumana, Ankit Singh Rawat, Himanshu Jain, Andreas Veit, and Sanjiv Kumar. Long-tail learning via logit adjustment. In International Conference on Learning Representations, 2021. URL hetps: http///pppeeeiiieeeeeeeee net/
Seyed Iman Mirzadeh, Mehrdad Farajtabar, Ang Li, Nir Levine, Akihiro Matsukawa, and Hassan Ghasemzadeh. Improved knowledge distillation via teacher assistant. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 5191-5198, 2020.
Michael C. Mozer and Paul Smolensky. Skeletonization: A technique for trimming the fat from network via relevance assessment. Proceedings of the 1st International Conference on Neural Information Processing Systems, pp. 107-115, 1988.
Rafael Müller, Simon Kornblith, and Geoffrey E Hinton. When does label smoothing help? In H. Wallach, H. Larochelle, A. Beygelzimer, F. dAlché Buc, E. Fox, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Asso- ciates, Inc., 2019. URL http:: http::::s//prroceeeengge.. www.:/eeeddiin............e...ee neurips wwtprriss.cppppppeeeeeeeeeeee pttppappe//119//1111 (1746600854443311151011102223333
Chenri Ni, Nontawat Charoenphakdee, Junya Honda, and Masashi Sugiyama. On the calibration of multiclass classification with rejection. In H. Wallach, H. Larochelle, Beygelzimer, F. dAlché Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, vol- ume 32, pp. 2586-2282 25966. Curran Associates, Inc., 2019.
Emma Ning. Microsoft open sources breakthrouug optimiiiionn for transformer inference and CPU. ht
on on GPU hatrroottooons-eeeeeooorrreeeete gpu-cpu/ 2013. Online; accessed March 2021.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint a9Xiv: 10683, 2019.
Alessandro Raganato, Yves Scherrer, and Jörg Tiedemann. Fixed encoder self attention patterns in transformer-batt machine translation. arXiv preprint arXiv:2002. 10260, 2020.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100,000+ questions for machine comprehension of text. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pp. 2383-2392, Austin, Texas, November 2016. Association for
Computational Linguistics. doi: 10. (0.18653/v1/D166 1264. URL https: www ac 1 web.. orq/
aathology/D16-1222
Harish G. Ramaswamy, Ambuj Tewari, and Shivani Agarwal. Consistent algorithms for multiclass classification with an abstain option. Electron. J. Statist. 12(1):530-555 554, 2018. doi: 10.121// EJS1388. URL
Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and Yoshua Bengio. Fitnets: Hints for thin deep nets. arXiv preprint 1412.6541 2014.
H. A. Rowley, S. Baluja, and T. Kanade. Neural network-based face detection. IEEE Transactions on Pattern Analysis and Machine Intelligence, 20(1):23-38, 1998. doi: 10. 10.1100 00.1109/34.6556477
Olga Russakovsky, Jia Deng. Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang. Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. Im- ageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJC), 115(3):111-25 2015. doi: 10. 1007/s1 1233-015-08166y,
Andrei A. Rusu, Sergio Gomez Colmenarejo, Çaglar Gülçehre, Guillaume Desjardins, James Kirk- patrick, Razvan Pascanu, Volodymyr Mnih, Koray Kavukcuoglu, and Raia Hadsell. Policy distil- lation. In 4th International Conference Learning Representations, ICLR 2016, 2016.
13

---

Under review as a conference paper at ICLR 2022
Roy Schwartz, Gabriel Stanovsky, Swabha Swayamdipta, Jesse Dodge, and Noah A. Smith. The right tool for the job: Matching model and instance complexities. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 6640-6655 Online, July 2020. Association for Computational Linguistics. doi: 20.18863/v1/2220.0 aall-main.593, URL
https::/www.aclwwbb.c org/ http:///w..............oo/////.. main. 593.
Cun, A. Shrivastava, S. Singh, and A. Gupta. Revisiting unreasonable effectiveness of data in deep learning era. In 20/7 IEEE International Conference on Computer Vision (ICCV), pp. 843-852, 2017. doi: 11.1 1,.110//CCV.2017.99
Zhiqing Sun, Hongkun Yu, Xiaodan Song, Renjie Liu, Yiming Yang, and Denny Zhou. Mobile- BERT: a compact aank-sgi BERT for lemmree devices. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 2158 2170, Online, July 2020. Association for Computational Linguistics. doi: 10.18653/v1/2020.. acl 195n. URL wtwp:/gg/y/22..020 2020 awwaaallmanann. 195.
httpps:/www. wwtp:/awwwww..//./aae....
Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. Rethink- ing the inception architecture for computer vision. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2016.
Mingxing Tan and Quoc Le. EfficientNet: Rethinking model scaling for convolutional neural net- works. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.). Proceedings of the 36th In- ternational Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 6105-6114. PMLR, 09-15 Jun 2019. URL httt:://proceeei/i/t ngs. www.mls......................... .mlr.
press/v97/tanl wtw.tannnaaah...
Grant Van Horn and Pietro Perona. The devil is in the tails: Fine-grained classification in the wild. 2017.
arXiv preprint
P. Viola and M. Jones. Rapid object detection using a boosted cascade of simple features. In Proceedings of the 2001 IEEE Computer Society Conference on Computer Vision and Pattern Recognition. CVPR 2001, volume 1, pp. I-1, 2001. doi: 10.1 N0.1109/CVPP.22 2000CCV7
Xin Wang, Yujia Luo, Daniel Crankshaw, Alexey Tumanov, Fisher Yu, and Joseph E Gonzalez. Idk cascades: Fast deep learning by learning not to overthink. arXiv preprint arXiv:1706.000:: 2017a.
Xin Wang, Fisher Yu, Zi-Yi Dou, Trevor Darrell, and Joseph E Gonzalez. Skipnet: Learning dynamic routing in convolutional networks. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 409-424, 2018.
Yu-Xion Ya-Xiong Wang, Deva Ramanan, and Martial Hebert. Learning to model the tail. In Proceedings of the 31st International Conference on Neural Information Processing Systems, pp. 1032-7042. 2017b.
Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sen- tence understanding through inference. In Proceedings of the 2018 Conference of the North Amer- ican Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers), pp. 1112-1122, New Orleans, Louisiana, June 2018 Association for Computational Linguistics. doi: 10. 10.11853//1/11--111 1101. URL
org/N18-1107
Qizhe Xie, Minh-Thang Luong, Eduard Hovy, and Quoe V. Le. Self-training with noisy student improves imagenet classification. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2020.
Zhixiang (Eddic) Xu, Matt J. Kusner, Kilian Q. Weinberger, Minmin Chen, and Olivier Chapelle Classifier cascades and trees for minimizing feature evaluation cost. Journal of Machine Learn- ing Research, (162))2113-2144 2014. URL http: http:://mmrr../// tttp:/ppppp/////////ea/
html
Li Yuan, Francis EH Tay, Guilin Li, Tao Wang, and Jiashi Feng. Revisiting knowledge distillation via label smoothing regularization. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2020.
14

---

Under review as a conference paper at ICLR 2022
Ali Hadi Zadeh, Isak Edo, Omar Mohamed Awad, and Andreas Moshovos. Gobo: Quantizing attention-bate based nlp models for low latency and energy efficient inference. 2020 53rd Anmual IEEE/ACM International Symposium on Microarchitecture (MICRO), pp. 811-824. IEEE, 2020.
Mi Zhang, Faen Zhang, Nicholas D Lane, Yuanchao Shu, Xiao Zeng, Biyi Fang, Shen Yan, and Hui Xu. Deep learning in the era of edge computing: Challenges and opportunities. Fog Computing: Theory and Practice, pp. 67-78, 2020.
Minjia Zhang, Samyam Rajbandari. Wenhan Wang, Elton Zheng, Olatunji Ruwase, Jeff Rasley, Jason Li, Junhua Wang, and Yuxiong He. Accelerating large scale deep learning infer- ence through deepcpu at microsoft. In 2019 USENIX Conference on Operational Ma- chine Learning (OpML 19), pp. 5-7, Santa Clara, CA, May 2019. USENIX Association. ISBN 978-1-939133-00-7. URL https::/www/ uswnix.e........... httpa///.//e/oooeeee///////mnn// presentation/ wtt.i/ooon//aag-nniiiiaaa
Xiangxin Zhu, Dragomir Anguelov, and Deva Ramanan. Capturing long-tail distributions of object subcategories. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recogni- tion, pp. 915-922, 2014.
15

---

Under review as a conference paper at ICLR 2022
A VARIANTS OF CLASS-SPECIHLE DISTILLATION
Here, we present two additional variants of class-specific distillation that enable transferring the performance of large teacher model to lite student model subset of class. Subsequently, the lite student can be employed in our two-stage inference framework.
A.1 IN-DOMAIN CLASS DISTILLATION
Since we do not intend the student to make the final prediction on an instance from [Z]] Cin, we can train the student to perform iiil = -way classification. particular, given teacher g and example (zr, y), we define a pseudo-lake label distribution:
[ 10,11 1e RR
if fin:
pelass, L'' Pg.
(11)
if [4]|Cin-
Here, 14/ denotes the softmax distribution restricted to Lin:
1) =c((t), EsEco
g(a)s for j E Lin-
(12)
SEC.
Now we get the student f : : + R² by minimizing
pclass, (g(x), f(x)) = H1 (Fgla .,,.)
distill
(13)
lite model fc.. trained with the loss (33 aims to classify an instance from Li to the cor- rect class. At the same time, such model is also trained to generate non-informative uniform distribution L/ on the instances from Lm, which by definition corresponds to zero margin.
Now, given the student fc... and teacher g. based delegation. In particular, for a test instance x E X, the final prediction becomes
define two-stage inference by using margin-
arg maxje(L1 L (z); if if (x) 2 P. arg maX 9(z);
((x)
if (x) <
where Ye us (x) is the margin assigned to x by the student L'..
A.2 IN-DOMAIN CLASS-DISTILLAN DISTILLATION WITH 'ABSTAIN' OPTION.
We next explore another natural candidate for class-specific distillation, where the lite student aims to correctly classify instance from Cin and declares to 'abstain' on the instances from [L]\C To realize this, we train the student to perform an (L' + 1)-way classification: Given a teacher g and example (x, y), we define a pseudo label distribution:
(W/'e 0) E [0, 1 (0.
if Lin: if E [L]
relass,
where P denotes the softmax distribution restricted to classes, as defined in (12) Now, one can perform the distillation by utilizing Polu pelase,L'+1 i.e., minimize
class, I.' I'+1
(g(x), f(x) = H (p .f.w).
(14)
distill
Note that (14) encourages the trained lite model J241 to predict (L' + 1)-th class, i.e., 'abstain' class, for all instances from [i). This leads natural abstain class-based delegation approach for two-stage inference. Let
gabstain
(x) arg max jE[[+1]
(Cm,
Then, given L'+1 and g, one makes the following final prediction for test instance.
((x) Ustatein (x) if Oahasein (x) < I',
Treacher (x) if gahasain (z) = I'
In addition, analogous to 7 one can also define a two-stage inference procedure via margin-based delegation.
16

---

Under review as a conference paper at ICLR 2022
Remark 4. Note that using 'abstain' class is closely related to classification with reject option (see Sec. B However, as opposed to the traditional classification with the reject paradigm, we intend to provide supervision for the call reject class well.
B CLASSIFICATION WITH REJECT OPTION
There is a large literature on selective classification, also known as classification with a reject option
abstention Grandvalet et al. (2009): Bartlett & Wegkamp (2008); Cortes et al. (2016); 6);; Geifman & El-Yaniv (2017); Ramaswamy et al. 2018); Ni et al. (2019). Here, one seeks predictor f: X ]u(L). where prediction of I denotes the classifier is uncertain. To avoid the degenerate solution of abstaining on all samples, one assumes fixed rejection cost E (0, 1J. The goal is to then trade- off the misclassification error on non-abstained samples with the total cost incurred abstained
samples, i.e., minimize the loss
{(y, fx)) = 1 (()))(()))) t f(x)=1.
(15)
The Bayes-optimal classifier for this objective abstains samples with high uncertainty the "true" label, i.e., Ramaswamy et al. (2018)
if max Ply x) 11-
f())=
yEY
(16)
argmax Pay x))
else.
yEy
C VARIANT OF MARGIN- DISTILLATION, USING ABSTAIN CLASS
Another natural approach for margin-based based distillation is to utilize an 'abstain' class to encour- age student to not spend its model capacity correctly classifying hard instances (where teacher achieves a low-margin). Towards this, we can define the following pseudo label distribution for an example (xr, y)-
(Bg (1), Pa.er (L),0) if (2,z)= v) E Seasy (0, ..., 0, E (0, 1] otherwise.
gmargin Pg.
Now, we distill lite student model f :x' RL+1 based on Prarginn i.e., we minimize
(17) Note that we train the student to perform an (L + 1)- classification, where it aims to classify an example in Seasy to one of L original classes and the examples in Shard to the 'abstain' class. Now, one may employ g and f4+11 to enable two-stage inference procedure that makes the final prediction for test instance x as
eaaarii (g(x), f(x)) H .Pf=)
i I (z) 2 P.
((x) =
Steacher (x)
otherwise,
where 0 = arg max, 1 (
D DETAILS OF EXPERIMENTAL SETUP
Datasets. We three benchmark image datasets CIFAR-100 (Krizhevsky 2009). ImageNet ILSVCC 2012 (a. ImageNet-1k) (Russakovsky et al. 2015). and ImageNet-214 (Degg et al.) 2009). CIFAR-100 contains 60k (50k train/TOk test) images annotated with one of 100 object cate- gories distributed uniformly. ImageNet (ILSVRC 2012), on the other hand, is much larger dataset with 1.33M (1.28M train/50k test) images annotated with of 1000 object categories also dis- tributed uniformly. As for ImagenNet-2111 k, it originally contains images 12.8M from 21,843 classes. We select 17,203 classes with least 100 images and create balanced test with 50 images from each of the selected classes. This remaining images from the selected classes provides us with a training set containing 11.8MM images.
In addition to image datasets, we also evaluate the proposed distillation-based two-stage inference procedure MNLI (Williams et al. 2018) and SQuAD datasets (Rajpurkar et al. 2016), which
17

---

Under review as a conference paper at ICLR 2022
Model
Parameters
FLOPs
Model
Parameters
FLOPS
ResNet-8
0.09 M
16 M
MobileNerv3-0.35
2.14 M
40 M
ResNet-14
0.19 M
30 M
MobileNerv3.0.50
2.70 M
70 M
ResNet-20
0.27 M
44 M
MobileNerv3-0.75
4.01 M
156 M
ResNet-32
0.46 M
72 M
MobileNeev3-L.00
5.50 M
218 M
ResNer-44
0.66 M
100 M
MobileNetv3-1.25
8.29 M
358 M
ResNet-56
0.85 M
128 M
Table (Lite) ResNet models for CIFAR-106 100.
Table 3: (Lite) MobileNetV3 models for ImageNet.
standard sentence classification and QA benchmarks, respectively. MNLI dataset corresponds to a 3-was 3-way classification task with 392,702 training examples and matched test set consisting of 9,815 instances, As for the SQuAD dataset, a span selection task over a sequence length of 384 with 87,599 and 10,570 train and test examples, respectively.
EfficientNet-L2 (Xie et al, Foret et al. 202I) which is the state-of-art for image classification on multiple datasets. is an optimized convolutional neural network pretrained both ImageNet and unlabeled JFT-300M (Sun et al. 2017) with input resolution
Large teacher models. For the image classification tasks, 2020 Tan & (0192 Foret
475. EfficientNer-L22 is a large model with 480M parameters which requires 478G FLOPs per inference and achieves an accuracy of 88.6% ImageNet. For CIFAR-106 100, we used base EfficientNet-L2 model with new fine-tuned classification layer achieving accuracy of 94.4%. For the sentence classification task, RoBERTa- et al. 2019) model serves as a teacher, which is pre- vre-traine on a general purpose text corpus of 160 encoder model with 355M parameters that requires 155G FLOPs
(Liu
GB is transformer-basee per inference while achieving an accuracy of 90.2% on MNLLFor the QA task requiring reading comprchension,
T5-11B (Rafell et al.. 2019) which exhibits competitive performance 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 large model with 1 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. 201 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 100 dataset was used. To sweep full spectrum all the way till teacher size, used ResNet networks detailed in Table 2 as the student networks and the largest ResNet-56 as the teacher network. We used class- specific distillation from See 4.1
E ADDITIONAL EXPERIMENTAL RESULTS
Class-specific distillation with class-based delegation. Table 4 represents the two-stage perfor- mance (both in-domain and overall) realized by class-specific distillation approaches (cf. Sec. 4.1 on ImageNet when we employ appropriate class-based delegation.
Class-specific distillation on ImageNet-21k dataset. Figure 5 the performance of our pro- posed distillation-basee two-stage inference framework on ImageNet-21k dataset. As discussed in Sec. D 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 that knows the true label) to simulate the two-stage stage inference framework. Also, while training a MobileNetV3-00755 model as a student, we use the one-hot labels as the supervision for the classes in Cin and utilize label-smoothing for the instances belonging the remaining classes. Accordingly, Baseline here corresponds to a normally trained student (MobileNetV3-00.7 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.
18

---

Under review as a conference paper at ICLR 2022
In-domain Fraction
Overall
Approach
Accuracy
Accuracy
Fraction
Baseline
0.75
1.00
0.68
1.00
MobsieNee-11 1.25 MobilNNet-0.5
CD)- (a = 0.0)
0.85
0.76
0.84
0.23
CD-I (a 0.4)
0.84
0,82
0.81
0.32
CD-I (a 0.6)
0.83
0,86
0.78
0.37
CD-III
0.87
0.63
0.85
0.23
Baseline
0.79
1.00
0.72
1.00
CD-I (a = 0.0)
0.86
0,79
0.85
0.28
CDD-I (a 0.4)
0.85
0,83
0.83
0.31
CD-I (a 0.6)
0.85
0.86
0.81
0.36
CD-I1I
0.87
0.70
0.85
0.25
Table 4: Performance of two-stage inference procedure on ImageNet-Ik lk. The student model is ob- tained by a class-specific distillation approach with SLinl L' 300, with in-domain referring to the instances belonging to the classes in Lin- CD-I and CD-III denote the class-bases distillation approaches defined in Sec.44.1 and Sec. A.2 Baseline refers the standard distillation from (3) with and 1. The inference procedure employs an appropriate class-based delegation for each distillation approach. Fraction denotes the fraction of test instances where the student model makes the final prediction. Note that, for standard distillation (Baseline), the student cannot delegate any examples to the teacher via bass- delegation.
1.0
1.0
Top-1 accuracy
0.8
Top-1 accuracy
0.8
0.6
Baseline
0.6
Baseline
0.4
CD-1 (a=0.4)
CD-I (a=0.4)
CD-I (o=0.66)
CD-I (a=0.66)
0.2
(a=0.8)
0.4
(a=0.8)
CD-I (a=0.9)
CDH (a=0.9)
0.0
0.2
0.4 0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Fraction processed by student
Fraction processed by student
(a) Overall accuracy
(b) In-domain accuracy
Figure 5: Comparison of various class-specific distillation methods (subset of) ImageNet-2Ik. As described in Sec. D we work with subset corresponding to 17,203 out of 21,843 classes. Baseline denotes the standard distillation from with = 0 and 1. CD-I denote the class- specific distillation approaches defined in Sec. 4.1 with Cin L' = 1000. Accordingly, in-domain accuracy is computed via focusing on the test instances from the 1000 classes in Cin- Here, each lite student model (a MobileNetV3-0.77 model) employs margin-bart delegation.
on both CIFAR-100 (cf. Fig. 6 and ImageNet (ef. Fig. Here, in-domain instances correspond to those instances where the student assigns margin of at least 0.4. As evident, margin-based distillation also leads to improved in-domain two-stage accuracy. However, Baseline and MD from have very similar performance on both datasets.
19

---

Under review as a conference paper at ICLR 2022
0.95
1.00
Top-1 accuracy
0.90
Top-1 accuracy
0.95
0.85
0.90
0.80
0.85
Baseline
Baseline
0.80
0.75
(Pram =0.4)
(Prran 00.4)
MD (Pran =07)
MD (pran =07)
0.75
0.70
0.0
0.2
0.4
0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
Fraction processed by student
Fraction processed by student
(a) Overall accuracy
(b) In-domain accuracy
Figure 6: Performance of two-stage inference procedure enabled by the margin-based distillation approach from Sec. 4.2 on CIFAR-100. Again, Baseline denotes the standard distillation from (3) with a and 1. MD represent the margin-based distillation approach from (7). Here, each lite student model (a ResNet-32 model) employs margin-bast based delegation. In-domain accuracy is computed on the test instances where student achieves margin p of least 0.4.
0.90
0.95
Top-1 accuracy
0.85
Top-1 accuracy
0.90
0.80
0.85
0.75
0.80
Baseline
Baseline
MD (Pran
MD (Pram =0.4)
0.70
MD (Pram -0.6)
0.75
MD (Pram -06)
0.2
0.4
0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
Fraction processed by student
Fraction processed by student
(a) Overall accuracy
(b) In-domain accuracy
Figure 7: Performance of two-stage inference procedure enabled by the margin-based distillation approach from Sec. 4.2 on ImageNet. Here, each lite student model (a MobileNetV3-0.77 model) employs margin-based delegation. In-domain accuracy is computed the test instances where student achieves a margin p of at least 0.4.
20

---

Under review as a conference paper at ICLR 2022
0.950
1.00
Accuracy
0.925
Exact Match
0.98
0.900
MNLI
0.96
SQUAD
Baseline
Baseline
0.875
MD (a=0.11
0.94
MD (a=0.2)
(a- 0.4)
MD (a 0.4)
0.850
0.25
0.50
0.75
1.00
25
50
75
Latency [s]
Latency [s]
(a) Latency vs. in-domain accuracy (MNLI)
(b) Latency vs. in-domain accuracy (SQuAD)
in-domain performance trade-off realized by the two-stage in- ference procedure on the NLP tasks. MD denotes the margin-based distillation with label smoothing parameter Left: On MNLI dataset, RoBERTa-Large and MobileBERT are used the teacher and student models, respectively. Here, Baseline corresponds to the two-stage inference enabled by the standard distillation (cf. ( Right: On SQuAD dataset, T5 and MobileBERT are used as the teacher and student models, respectively. Here, Baseline denotes a normally trained MobileBERT used in combination with T5 model.
Figure 8: The inference-latenee
21


Reviews:
[{"review_id": "mWJVwCSLTgK", "reviewer": "Reviewer_nhWf", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Traditionally model distillation usually sacrifice performance for improved efficiency. In this paper, the authors proposed a two stage framework based on distillation which can achieve both the modeling benefits of the large models and preserve the efficiency with lightweight models. The key component is a well designed loss function for the student model which can help in the delegation process. Overall, I believe the paper is well motivated, the method is described clearly, and thorough experiments are conducted. ", "summary_of_the_paper": "This paper proposes a two stage distillation framework. By falling back hard samples to the large models and use small models for easy ones, such an approach could achieve both efficiency and performance. ", "summary_of_the_review": "I only have one major concerns towards the experiment. The authors propose several different approaches in the paper, including class-specific distillation and margin based distillation, and corresponding delegation methods. However in the experiment section it looks like they only focus on class-specific distillation and some related methods which are introduced in the appendix. \n\nI would actually suggest the authors to make the paper more self-contained, if the methods are introduced in the main text they should be evaluated as well. And some alternative methods can be included in the appendix. Also, I feel like it would be better to have a high level summary in terms of the pros and cons for each method, along with the recommendation for users in practice. \n\nMinor issues: line 4 of Introduction, upto -> up tp", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "dceYPB9L8HX", "reviewer": "Reviewer_sKmp", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths\n- It is a well written and clear paper providing a two-stage inference framework that can be used to reduce the dependency on large models and thus improve inference efficiency. \n- The relevance of the framework is validated empirically on benchmark datasets in Computer Vision and NLP. \n\nWeaknesses\n- Even though the framework is clear and easy to use. The results are not very convincing as the performance of the two stage framework is very close to distillation which is also efficient and does not rely on the teacher model during inference. In particular in Figure 2, 3 and 4 curves of the proposed approach seem to be very close to the baseline standard distillation. The gains are clearer in Table 1 and Table 4. Could the authors add a similar table for Figure 4? \n- It is not clear which approach (CD-I, CD-II, CD-III) and which parameters (alpha) to use in practice. Could the authors provide a summary of the findings regarding recommendations on which approach to use in which cases based on the experiences? \n\nQuestions\n- In 4.1, the authors mention that « As a result of standard distillation, the lite student behaves as a well-calibrated model ». Could the authors justify this statement? I am not sure to understand why it is necessarily the case. \n- In Table 1, the CD-II approach obtain a much lower Accuracy score than the other approaches (0.24), Could the authors explain that?\n\nTypos\n- Introduction, third paragraph : « real-world data »\n- 3.2, 2nd paragraph : « the teacher scores »\n", "summary_of_the_paper": "The paper proposes a two-stage distillation framework to improve inference efficiency and reduce the dependency on large teacher models. The goal of this framework is to only use the large/teacher model for difficult and rare examples and to use the student, smaller model for the more frequent easy examples. The procedure is composed of a training phase and an inference phase. During the training phase, the dataset is separated into two subsets: one containing the hard/difficult examples and one containing the easy examples. Using a loss incorporating a combination of label-smoothing and distillation, the student model is taught to be certain on the easy examples (by using distillation of the large model) and to be less certain on the hard/difficult examples. During the inference phase, in order to route the hard examples to the larger teacher model, several possible methods are proposed relying in particular on whether the margin of the student’s softmax is higher or lower than a threshold. \n\nThe authors validate their two-stage framework empirically on three benchmark image datasets (CIFAR-100, ImageNet-1k and ImageNet-21k) and two benchmark NLP datasets (SQuAD and MNLI). ", "summary_of_the_review": "Even though the performance of the two-stage framework seems to be higher than the standard distillation, I am not sure yet if the gains of the proposed approach compared to standard distillation justify the acceptance of this paper. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "VFmzfCRmC2e", "reviewer": "Reviewer_npvb", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength: The paper is fairy well written and easy to follow.\n\nWeaknesses:\nThere are multiple weaknesses. The idea of combining a small student model and the original large model for inference is pretty intuitive (I am not saying intuitive is a bad thing) but does not demonstrate well. It is not clear why such a method can be interesting and be better than existing methods. For example, Section 2 discusses some related work on efficient inference. However, these are never compare with the proposal method, neither theoretically nor empirically. Especially, in the experiments, only the standard distillation baseline is compared with the proposed method, which is quite limited and fail to demonstrate the advantages of the proposed method. Thus, the motivation of the proposed method is quite unclear. In addition, since the focus of this paper is on inference efficiency, it is necessary to compare the inference time, in addition to the current accuracy.", "summary_of_the_paper": "This paper studies efficient inference problem for large models. It proposes to train a small student model, and performs inference for easy data on the student model, and for hard data on the original large model. Experiments show the proposed method performs better than the simple baseline of standard distillation model.", "summary_of_the_review": "This paper fails to demonstrate the motivation and advantage of the proposed method for efficient inference of large model. It does not meet the acceptance bar of the conference.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "CQbmhcKSeVh", "reviewer": "Reviewer_wETK", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper tackles an important and practical problem: So-called \"edge\" devices such as mobile phones cannot process heavy duty models. However, requirements on bandwidth and latency may prevent one from sending every example to the cloud for processing on a big GPU capable machine. So, this paper presents a hybrid approach where most of the cases can be processed on the \"edge\" but those edge devices can elect to send rare cases to the cloud for a heavy duty model.\n\nDespite being an important problem, the solutions in the paper \"label delegation\" and \"margin delegation\" are very heuristic and arbitrary. Solutions like early-exit networks seem much more principled in comparison. They are also crucially dependent on hyper-parameters, which the paper does not address tuning.\n\nHeuristic solutions are still noteworthy for publication if they beat all other known approaches in a rigorous benchmark. However, in this paper, most of the solutions benchmark the method against variations on itself.\n\nThe paper would be stronger if it had more of these comparisons.", "summary_of_the_paper": "This paper presents a new method to do efficient inference with a teacher-student setup for model distillation. It contrasts itself to traditional model distillation because it is optimized for the situation where rare and difficult cases can still be sent to the original teacher network.", "summary_of_the_review": "The paper attempts to solve and important problem, but is very heuristic and does not validate that it is better than the prior art.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper a distillation framework where a light-weight student model is trained to handle easy (frequent) instances, while the large teacher model is still used to handle the more difficult (rare) inputs. The models are trained to perform well in this two-stage inference setting. Experiments are conducted on computer vision and NLP tasks. While the idea is potentially interesting, the experimental results are fairly weak and not very convincing.
Decision: Reject