Don't be picky, all students in the right family can learn from good teachers

1 INTRODUCTION . Recently-developed deep learning models have achieved remarkable performance in a variety of tasks . However , breakthroughs leading to state-of-the-art ( SOTA ) results often rely on very large models : GPipe , Big Transfer and GPT-3 use 556 million , 928 million and 175 billion parameters , respectively ( Huang et al. , 2019 ; Kolesnikov et al. , 2020 ; Brown et al. , 2020 ) . Deploying these models on user devices ( e.g . smartphones ) is currently impractical as they require large amounts of memory and computation ; and even when large devices are an option ( e.g . GPU clusters ) , the cost of large-scale deployment ( e.g . continual inference ) can be very high ( Cheng et al. , 2017 ) . Additionally , target hardware does not always natively or efficiently support all operations used by SOTA architectures . The applicability of these architectures is , therefore , severely limited , and workarounds using smaller or simplified models lead to a performance gap between the technology available at the frontier of deep learning research and that usable in industry applications . In order to bridge this gap , Knowledge Distillation ( KD ) emerges as a potential solution , allowing small student models to learn from , and emulate the performance of , large teacher models ( Hinton et al. , 2015a ) . The student model can be constrained in its size and type of operations used , so that it will satisfy the requirements of the target computational environment . Unfortunately , successfully achieving this in practice is extremely challenging , requiring extensive human expertise . For example , while we know that the architecture of the student is important for distillation ( Liu et al. , 2019b ) , it remains unclear how to design the optimal network given some hardware constraints . With Neural Architecture Search ( NAS ) it is possible to discover an optimal student architecture . NAS automates the choice of neural network architecture for a specific task and dataset , given a search space of architectures and a search strategy to navigate that space ( Pham et al. , 2018 ; Real et al. , 2017 ; Liu et al. , 2019a ; Carlucci et al. , 2019 ; Zela et al. , 2018 ; Ru et al. , 2020 ) . One im- portant limitation of most NAS approaches is that the search space is very restricted , with a high proportion of resources spent on evaluating very similar architectures , thus rendering the approach limited in its effectiveness ( Yang et al. , 2020 ) . This is because traditional NAS approaches have no tools for distinguishing between architectures that are similar and architectures that are very different ; as a consequence , computational resources are needed to compare even insignificant changes in the model . Conversely , properly exploring a large space requires huge computational resources : for example , recent work by Liu et al . ( 2019b ) investigating how to find the optimal student requires evaluating 10 , 000 models . By focusing on the comparison between distributions we ensure to use computational resources only on meaningful differences , thus performing significantly more efficiently : we evaluate 33× less architectures than the most related work to ours ( Liu et al. , 2019b ) . To overcome these limitations , we propose an automated approach to knowledge distillation , in which we look for a family of good students rather than a specific model . We find that even though our method , AutoKD , does not output one specific architecture , all architectures sampled from the optimal family of students perform well when trained with KD . This reformulation of the NAS problem provides a more expressive search space containing very diverse architectures , thus increasing the effectiveness of the search procedure in finding good student networks . Our contributions are as follows : ( A ) a framework for combining KD with NAS and effectively emulate large models while using a fraction of the memory and of the parameters ; ( B ) By searching for an optimal student family , rather than for specific architectures , our algorithm is up to 20x more sample efficient than alternative NAS-based KD solutions ; ( C ) We significantly outperform advanced KD methods on a benchmark of vision datasets , despite using the traditional KD loss , showcasing the efficacy of our found students . 2 RELATED WORK . Model compression has been studied since the beginning of the machine learning era , with multiple solutions being proposed ( Choudhary et al. , 2020 ; Cheng et al. , 2017 ) . Pruning based methods allow the removal of non-essential parameters from the model , with little-to-none drop in final performance . The primary motive of these approaches was to reduce the storage requirement , but they can also be used to speed up the model ( LeCun et al. , 1990 ; Han et al. , 2015 ; Li et al. , 2016a ) . The idea behind quantization methods is to reduce the number of bits used to represent the weights and the activations in a model ; depending on the specific implementation this can lead to reduced storage , reduced memory consumption and a general speed-up of the network ( Fiesler et al. , 1990 ; Soudry et al. , 2014 ; Rastegari et al. , 2016 ; Zhu et al. , 2016 ) . In low rank factorization approaches , a given weight matrix is decomposed into the product of smaller ones , for example using singular value decomposition . When applied to fully connected layers this leads to reduced storage , while when applied to convolutional filters , it leads to faster inference ( Choudhary et al. , 2020 ) . All the above mentioned techniques can successfully reduce the complexity of a given model , but are not designed to substitute specific operations . For example , specialized hardware devices might only support a small subset of all the operations offered by modern deep learning frameworks . In Knowledge Distillation approaches , a large model ( the teacher ) distills its knowledge into a smaller student architecture ( Hinton et al. , 2015b ) . This knowledge is assumed to be represented in the neural network ’ s output distribution , hence in the standard KD framework , the output distribution of a student ’ s network is optimized to match the teacher ’ s output distribution for all the training data ( Yun et al. , 2020 ; Ahn et al. , 2019 ; Yuan et al. , 2020 ; Tian et al. , 2020 ; Tung & Mori , 2019 ) . The work of Liu et al . ( 2019b ) shows that the architecture of a student network is a contributing factor in its ability to learn from a given teacher . The authors propose combining KD with a traditional NAS pipeline , based on Reinforcement Learning , to find the optimal student . While this setup leads to good results , it does so at a huge computational cost , requiring over 5 days on 200 TPUs . Similarly , Gu & Tresp ( 2020 ) also look for the optimal student architecture , but do so by searching for a subgraph of the original teacher ; therefore , it can not be used to substitute unsupported operations . Orthogonal approaches , looking at how KD can improve NAS , are explored by Trofimov et al . ( 2020 ) and Li et al . ( 2020 ) . The first establishes that KD improves the correlation between different budgets in multi-fidelity methods , while the second uses the teacher supervision to search the architecture in a blockwise fashion . 3 SEARCHING FOR THE OPTIMAL STUDENT NETWORK GENERATOR . The AutoKD framework ( Fig . 1 ) combines Bayesian Optimization ( BO ) , Neural Architecture Search ( NAS ) and Knowledge Distillation ( KD ) . AutoKD defines a family of random network generators G ( θ ) parameterized by a hyperparameter θ , from where student networks are sampled . BO uses a surrogate model to propose generator hyperparameters , while students from these generators are trained with KD using a state-of-the-art teacher network . The student performances are evaluated and provided as feedback to update the BO surrogate model . To improve our BO surrogate model , the search procedure is iterated , until the best family of student networks G ( θ∗ ) is selected . In this section we specify all components of AutoKD . See also Fig . 1 and Algorithm 1 for an overview . 3.1 KNOWLEDGE DISTILLATION . Knowledge Distillation ( KD ; Hinton et al. , 2015b ) is a method to transfer , or distill , knowledge from one model to another—usually from a large model to small one—such that the small student model learns to emulate the performance of the large teacher model . KD can be formalized as minimizing the objective function : LKD = ∑ xi∈X l ( fT ( xi ) , fS ( xi ) ) ( 1 ) where l is the loss function that measures the difference in performance between the teacher fT and the student fS , xi is the ith input , yi is the ith target . The conventional loss function l used in practice is a linear combination of the traditional cross entropy loss LCE and the Kullback–Leibler divergence LKL of the pre-softmax outputs for fT and fS : l = ( 1− α ) LCE + αLKL ( σ ( fT ( xi ) /τ ) , σ ( fS ( xi ) /τ ) ) ( 2 ) where σ is the softmax function σ ( x ) = 1/ ( 1 + exp ( −x ) ) , and τ is the softmax temperature . Hinton et al . ( 2015b ) propose “ softening ” the probabilities using temperature scaling with τ ≥ 1 . The parameter α represents the weight trade-off between the KL loss and the cross entropy loss LCE . The LKD loss is characterized by the hyper-parameters : α and τ ; popular choices are τ ∈ { 3 , 4 , 5 } and α = 0.9 ( Huang & Wang , 2017 ; Zagoruyko & Komodakis , 2016 ; Zhu et al. , 2018 ) . Numerous other methods ( Polino et al. , 2018 ; Huang & Wang , 2017 ; Tung & Mori , 2019 ) can be formulated as a form of Equation ( 2 ) , but in this paper we use the conventional loss function l. Traditionally in KD , both the teacher and the student network have predefined architectures . In contrast , AutoKD defines a search space of student network architectures and finds the optimal student by leveraging neural architecture search , as detailed below . 3.2 STUDENT SEARCH VIA GENERATOR OPTIMIZATION . Most NAS method for vision tasks employ a cell-based search space , where networks are built by stacking building blocks ( cells ) and the operations inside the cell are searched ( Pham et al. , 2018 ; Real et al. , 2017 ; Liu et al. , 2019a ) . This results in a single architecture being output by the NAS procedure . In contrast , more flexible search spaces have recently been proposed that are based on Algorithm 1 : AutoKD 1 : Input : Network generator G , BOHB hyperparameters ( η , training budget bmin and bmax ) , Evaluation function fKD ( θ , b ) which assesses the validation performance of a generator hyperparameterθ by sampling an architecture from G ( θ ) and training it with the KD loss LKD ( equations 1 and 2 ) for b epochs . 2 : smax = blogη bmaxbmin c ; 3 : for s ∈ { smax , smax − 1 , . . . , 0 } do 4 : Sample M = d smax+1s+1 · η se generator hyperparameters Θ = { θj } Mj=1 which maximises the raito of kernel density estimators ; . ( Falkner et al. , 2018 , Algorithm 2 ) 5 : Initialise b = ηs · bmax ; . Run Successive Halving ( Li et al. , 2016b ) 6 : while b ≤ bmax do 7 : L = { fKD ( θ , b ) : θ ∈ Θ } ; 8 : Θ = top k ( Θ , L , b|Θ|/ηc ) ; 9 : b = η · b ; 10 : end while 11 : end for 12 : Obtain the best performing configuration θ∗ for the student network generator . 13 : Sample k architectures from G ( θ∗ ) , train them to completion , and obtain test performance . neural network generators ( Xie et al. , 2019 ; Ru et al. , 2020 ) . The generator hyperparameters define the characteristics of the family of networks being generated . NAGO optimizes an architecture generator instead of a single architecture and proposes a hierarchical graph-based space which is highly expressive yet low-dimensional ( Ru et al. , 2020 ) . Specifically , the search space of NAGO comprises three levels of graphs ( where the node in the higher level is a lower-level graph ) . The top level is a graph of cells ( Gtop ) and each cell is itself a graph of middlelevel modules ( Gmid ) . Each module further corresponds to a graph of bottom-level operation units ( Gbottom ) such as a relu-conv3×3-bn triplet . NAGO adopts three random graph generators to define the connectivity/topology of Gtop , Gmid and Gbottom respectively , and thus is able to produce a wide variety of architectures with only a few generator hyperparameters . AutoKD employs NAGO as the NAS backbone for finding the optimal student family . Our pipeline consists of two phases . In the first phase ( search ) , a multi-fidelity Bayesian optimisation technique , BOHB ( Falkner et al. , 2018 ) , is employed to optimise the low-dimensional search space . BOHB uses partial evaluations with smaller-than-full budget to exclude bad configurations early in the search process , thus saving resources to evaluate more promising configurations . Given the same time constraint , BOHB evaluates many more configurations than conventional BO which evaluates all configurations with full budget . As Ru et al . ( 2020 ) empirically observe that good generator hyperparameters lead to a tight distribution of well-performing architectures ( small performance standard deviation ) , we similarly assess the performance of a particular generator hyperparameter value with only one architecture sample . In the second phase ( retrainA ) , AutoKD uniformly samples multiple architectures from the optimal generator found during the search phase and evaluates them with longer training budgets to obtain the best architecture performance . Instead of the traditionally used cross-entropy loss , AutoKD uses the KD loss in equation 2 to allow the sampled architecture to distill knowledge from its teacher . The KD hyperparameters temperature τ and loss weight α are included in the search space and optimized simultaneously with the architecture to ensure that the student architectures can efficiently distill knowledge both from the designated teacher and the data distribution . A full overview of the framework is shown in Fig . 1 .

Summary: This paper proposes searching for an architecture generator that outputs good student architectures for a given teacher. The authors claim that by learning the parameters of the generator instead of relying directly on the search space, it is possible to explore the search space of architectures more effectively, increasing the diversity of the architectures explored. They show that this approach combined with the standard knowledge distillation loss is able to learn good student architectures requiring substantially less samples and achieving competitive performances when comparing to other knowledge distillation algorithms.