When in Doubt, Summon the Titans: A Framework for Efficient Inference with Large Models

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 . 1 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 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 ( Bucilǎ 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 inefficient . Such an inefficient 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 . In this paper , we focus on 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 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 confident , we emit the prediction and we want the student to be confident 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 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 ( 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 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 it 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 1 , there is a sweet spot where we can achieve the same accuracy 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 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 . 2 RELATED 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-Perpinán , 2017 ; Howard et al. , 2019 ) . Significant 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 . ( 2020 ) ; 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 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 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 . ( 1998 ) and later formalized by Viola & Jones ( 2001 ) . The basic idea was to design a cascade of independent classifier and reject early on and cheaply . The idea has been generalized from a 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 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 . ( 2020 ) ; Liu et al . ( 2020 ) ; 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 , 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 efficient 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 benefit from its large model capacity and a lite distilled model suffices to predict a large portion of test instances .

Summary: 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.   The 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).