Knowledge Distillation as Semiparametric Inference

A popular approach to model compression is to train an inexpensive student model to mimic the class probabilities of a highly accurate but cumbersome teacher model . Surprisingly , this two-step knowledge distillation process often leads to higher accuracy than training the student directly on labeled data . To explain and enhance this phenomenon , we cast knowledge distillation as a semiparametric inference problem with the optimal student model as the target , the unknown Bayes class probabilities as nuisance , and the teacher probabilities as a plug-in nuisance estimate . By adapting modern semiparametric tools , we derive new guarantees for the prediction error of standard distillation and develop two enhancements—cross-fitting and loss correction—to mitigate the impact of teacher overfitting and underfitting on student performance . We validate our findings empirically on both tabular and image data and observe consistent improvements from our knowledge distillation enhancements . 1 INTRODUCTION . Knowledge distillation ( KD ) ( Craven & Shavlik , 1996 ; Breiman & Shang , 1996 ; Bucila et al. , 2006 ; Li et al. , 2014 ; Ba & Caruana , 2014 ; Hinton et al. , 2015 ) is a widely used model compression technique that enables the deployment of highly accurate predictive models on devices such as phones , watches , and virtual assistants ( Stock et al. , 2020 ) . KD operates by training a compressed student model to mimic the predicted class probabilities of an expensive , high-quality teacher model . Remarkably and across a wide variety of domains ( Hinton et al. , 2015 ; Sanh et al. , 2019 ; Jiao et al. , 2019 ; Liu et al. , 2018 ; Tan et al. , 2018 ; Fakoor et al. , 2020 ) , this two-step process often leads to higher accuracy than training the student directly on the raw labeled dataset . While the practice of KD is now well developed , a general theoretical understanding of its successes and failures is still lacking . As we detail below , a number of authors have argued that the success of KD lies in the more precise “ soft labels ” provided by the teacher ’ s predicted class probabilities . Recently , Menon et al . ( 2020 ) observed that these teacher probabilities can serve as a proxy for the Bayes probabilities ( i.e. , the true class probabilities ) and that the closer the teacher and Bayes probabilities , the better the student ’ s performance should be . Building on this observation , we cast KD as a plug-in approach to semiparametric inference ( Kosorok , 2007 ) : that is , we view KD as fitting a student model f̂ in the presence of nuisance ( the Bayes probabilities p0 ) with the teacher ’ s probabilities p̂ as a plug-in estimate of p0 . This insight allows us to adapt modern tools from semiparametric inference to analyze the error of a distilled student in Sec . 3 . Our analysis also reveals two distinct failure modes of KD : one due to teacher overfitting and data reuse and the other due to teacher underfitting from model misspecification or insufficient training . In Sec . 4 , we introduce and analyze two complementary KD enhancements that correct for these failures : cross-fitting—a popular technique from semiparametric inference ( see , e.g. , Chernozhukov et al. , 2018 ) —mitigates teacher overfitting through data partitioning while loss correction mitigates teacher underfitting by reducing the bias of the plug-in estimate p̂ . The latter enhancement was inspired by the orthogonal machine learning ( Chernozhukov et al. , 2018 ; Foster & Syrgkanis , 2019 ) approach to semiparametric inference which suggests a particular adjustment for the teacher ’ s log probabilities . We argue in Sec . 4 that this orthogonal correction minimizes the teacher bias but often at the cost of unacceptably large variance . Our proposed correction avoids this variance explosion by balancing the bias and variance terms in our generalization bounds . In Sec . 5 , we complement our theoretical analysis with a pair of experiments demonstrating the value of our enhancements on six real classification problems . On five real tabular datasets , cross-fitting and loss correction improve student performance by up to 4 % AUC over vanilla KD . Furthermore , on CIFAR-10 ( Krizhevsky & Hinton , 2009 ) , a benchmark image classification dataset , our enhancements improve vanilla KD accuracy by up to 1.5 % when the teacher model overfits . Related work . Since we can not review the vast literature on KD in its entirety , we point the interested reader to Gou et al . ( 2020 ) for a recent overview of the field . We devote this section to reviewing theoretical advances in the understanding of KD and summarize complementary empirical studies and applications of in the extended literature review in App . A . A number of papers have argued that the availability of soft class probabilities from the teacher rather than hard labels enables us to improve training of the student model . This was hypothesized in Hinton et al . ( 2015 ) with empirical justification . Phuong & Lampert ( 2019 ) consider the case in which the teacher is a fixed linear classifier and the student is either a linear model or a deep linear network . They show that the student can learn the teacher perfectly if the number of training examples exceeds the ambient dimension . Vapnik & Izmailov ( 2015 ) discuss the setting of learning with privileged information where one has additional information at training time which is not available at test time . Lopez-Paz et al . ( 2015 ) draw a connection between this and KD , arguing that KD is effective because the teacher learns a better representation allowing the student to learn at a faster rate . They hypothesize that a teacher ’ s class probabilities enable student improvement by indicating how difficult each point is to classify . Tang et al . ( 2020 ) argue using empirical evidence that label smoothing and reweighting of training examples using the teacher ’ s predictions are key to the success of KD . Mobahi et al . ( 2020 ) analyzed the case of self-distillation in which the student and teacher function classes are identical . Focusing on kernel ridge regression models , they proved that self-distillation can act as increased regularization strength . Bu et al . ( 2020 ) considers more generic model compression in a rate-distortion framework , where the rate is the size of the student model and distortion is the difference in excess risk between the teacher and the student . Menon et al . ( 2020 ) consider the case of losses such that the population risk is linear in the Bayes class probabilities . They consider distilled empirical risk and Bayes distilled empirical risk which are the risk computed using the teacher class probabilities and Bayes class probabilities respectively rather than the observed label . They show that the variance of the Bayes distilled empirical risk is lower than the empirical risk . Then using analysis from Maurer & Pontil ( 2009 ) ; Bennett ( 1962 ) , they derive the excess risk of the distilled empirical risk as a function of the ` 2 distance between the teacher ’ s class probabilities and the Bayes class probabilities . We significantly depart from Menon et al . ( 2020 ) in multiple ways : i ) our Thm . 1 allows for the common practice of data re-use , ii ) our results cover the standard KD losses SEL and ACE which are non-linear in p0 , iii ) we use localized Rademacher analysis to achieve tight fast rates for standard KD losses , and iv ) we use techniques from semiparametric inference to improve upon vanilla KD . 2 KNOWLEDGE DISTILLATION BACKGROUND . We consider a multiclass classification problem with k classes and n training datapoints zi= ( xi , yi ) sampled independently from some distribution P. Each feature vector x belongs to a setX , each label vector y∈ { e1 , ... , ek } ⊂ { 0,1 } k is a one-hot encoding of the class label , and the conditional probability of observing each label is the Bayes class probability function p0 ( x ) =E [ Y |X =x ] . Our aim is to identify a scoring rule f : X →Rk that minimizes a prediction loss on average under the distribution P. Knowledge distillation . Knowledge distillation ( KD ) is a two-step training process where one first uses a labeled dataset to train a teacher model and then trains a student model to predict the teacher ’ s predicted class probabilities . Typically the teacher model is larger and more cumbersome , while the student is smaller and more efficient . Knowledge distillation was first motivated by model compression ( Bucila et al. , 2006 ) , to find compact yet high-performing models to be deployed ( such as on mobile devices ) . In training the student to match the teacher ’ s prediction probability , there are several types of loss functions that are commonly used . Let p̂ ( x ) ∈Rk be the teacher ’ s vector of predicted class probabilities , f ( x ) ∈Rk be the student model ’ s output , and [ k ] , { 1,2 , ... , k } . The most popular distillation loss functions1 ` ( z ; f ( x ) , p̂ ( x ) ) include the squared error logit ( SEL ) loss ( Ba & Caruana , 2014 ) ` se ( z ; f ( x ) , p̂ ( x ) ) , ∑ j∈ [ k ] 1 2 ( fj ( x ) −log ( p̂j ( x ) ) ) 2 ( SEL ) and the annealed cross-entropy ( ACE ) loss ( Hinton et al. , 2015 ) ` β ( z ; f ( x ) , p̂ ( x ) ) =− ∑ j∈ [ k ] p̂j ( x ) β∑ l∈ [ k ] p̂l ( x ) β log ( exp ( βfj ( x ) ) ∑ l∈ [ k ] exp ( βfl ( x ) ) ) ( ACE ) for an inverse temperatureβ > 0 . These loss functions measure the divergence between the probabilities predicted by the teacher and the student . A student model trained with knowledge distillation often performs better than the same model trained from scratch ( Bucila et al. , 2006 ; Hinton et al. , 2015 ) . In Secs . 3 and 4 , we will adapt modern tools from semiparametric inference to understand and enhance this phenomenon . 3 DISTILLATION AS SEMIPARAMETRIC INFERENCE . In semiparametric inference ( Kosorok , 2007 ) , one aims to estimate a target parameter or function f0 , but that estimation depends on an auxiliary nuisance function p0 that is unknown and not of primary interest . We cast the knowledge distillation process as a semiparametric inference problem , by treating the unknown Bayes class probabilities p0 as nuisance and the teacher ’ s predicted probabilities as a plug-in estimate of that nuisance . This perspective allows us bound the generalization of the student in terms of the mean squared error ( MSE ) between the teacher and the Bayes probabilities . In the next section ( Sec . 4 ) we use techniques from semiparametric inference to enhance the performance of the student . The interested reader could consult Tsiatis ( 2007 ) for more details on semiparametric inference . Our analysis starts from taking the following perspective on distillation . For a given pointwise loss function ` ( z ; f ( x ) , p0 ( x ) ) , we view the goal of the student as minimizing an oracle population loss over a function classF , LD ( f , p0 ) =E [ ` ( Z ; f ( X ) , p0 ( X ) ) ] with f0 , argminf∈FLD ( f , p0 ) . The main hurdle is that this is objective depends on the unknown Bayes probabilities p0 . We view the teacher ’ s model p̂ as an approximate version of p0 and bound the distillation error of the student as a function of the teacher ’ s estimation error . Typical semiparametric inference considers cases where f0 is a finite dimensional parameter ; however recent work of Foster & Syrgkanis ( 2019 ) extends this framework to infinite dimensional models f0 and to develop statistical learning theory with a nuisance component framework . The distillation problem fits exactly into this setup . Bounds on vanilla KD As a first step we derive a vanilla bound on the error of the distilled student model without any further modifications of the distillation process , i.e. , we assume that the student is trained on the same data as the teacher and is trained by running empirical risk minimization ( ERM ) on the plug-in loss , plugging in the teacher ’ s model instead of p0 , i.e. , f̂=argminf∈FLn ( f , p̂ ) for Ln ( f , p̂ ) , En [ ` ( Z ; f ( X ) , p̂ ( X ) ) ] ( Vanilla KD ) where En [ X ] = 1n ∑n i=1Xi denotes the empirical expectation of a random variable . Technical definitions Before presenting our main theorem we introduce some technical notation . For a vector valued function f that takes as input a random variableX , we use the shorthand notation ‖f‖p , q , ‖‖f ( X ) ‖p‖Lq =E [ ‖f ( X ) ‖qp ] 1/q . Let∇φ and∇π denote the partial derivatives of ` ( z ; φ , π ) , with respect to its second and third input correspondingly and ∇φπ the Jacobian of cross partial derivatives , i.e. , [ ∇φπ ` ( z ; φ , π ) ] i , j= ∂ 2 ∂φj∂πi ` ( z ; φ , π ) . Finally , let qf , p ( x ) =E [ ∇φπ ` ( Z ; f ( X ) , p ( X ) ) |X=x ] and γf , p ( x ) =EU∼Unif ( [ 0,1 ] ) [ qf , Up+ ( 1−U ) p0 ( x ) ] . Critical radius Finally , we need to define the notion of the critical radius ( see , e.g. , Wainwright ( 2019 , 14.1.1 ) ) of a function class , which typically provides tight learning rates for statistical learning theory tasks . For any function classF we define the localized Rademacher complexity as : R ( δ ; F ) =EX1 : n , 1 : n [ supf∈F : ‖f‖2≤δ 1 n ∑n i=1 if ( Xi ) ] where i are i.i.d . random variables taking values equiprobably in { −1,1 } . The critical radius of a classF , taking values in [ −H , H ] , is the smallest positive solution δn to the inequalityR ( δ ; F ) ≤ δ 2 H . 1These loss functions do not depend on the ground-truth label y , but we use the augmented notation ` ( z ; f ( x ) , p̂ ( x ) ) to accommodate the enhanced distillation losses presented in Sec . 4 . Theorem 1 ( Vanilla KD analysis ) . Suppose f0 belongs to a convex set F satisfying the ` 2/ ` 4 ratio condition supf∈F ‖f−f0‖2,4/‖f−f0‖2,2 ≤C and that the teacher estimates p̂∈P from the same dataset used to train the student . Let δn , ζ = δn+c0 √ log ( c1/ζ ) n for universal constants c0 , c1 and δn an upper bound on the critical radius of the function class G , { z→r ( ` ( z ; f ( x ) , p ( x ) ) − ` ( z ; f0 ( x ) , p ( x ) ) ) : f ∈F , p∈P , r∈ [ 0,1 ] } . Let µ ( z ) = supφ‖∇φ ` ( z ; φ , p̂ ( x ) ) ‖2 , and assume that the loss ` ( z ; φ , π ) is σ-strongly convex in φ for each z and that each g∈G is uniformly bounded in [ −H , H ] . Then the Vanilla KD f̂ satisfies ‖f̂−f0‖22,2 = 1σ2O ( δ 2 n , ζC 2H2‖µ‖24+‖γ > f0 , p̂ ( p̂−p0 ) ‖ 2 2,2 ) with probability at least 1−ζ . Thm . 1 , proved in App . C , shows that vanilla distillation yields an accurate student whenever the teacher generalizes well ( i.e. , ‖p̂−p0‖2,2 is small ) and the student and teacher model classesF andP are not too complex . The ` 2/ ` 4 ratio requirement can be removed at the expense of replacing ‖µ‖4 by ‖µ‖∞=supz|µ ( z ) | in the final bound . Moreover , we highlight that the strong convexity requirement for ` is satisfied by all standard distillation objectives including SEL and ACE , as it is strong convexity with respect to the output of f and not the parameters of f . Even this requirement could be removed , but this would yield slow rate bounds of the form : ‖f̂−f0‖22,2 =O ( δn , ζ+‖γ > f0 , p̂ ( p̂−p0 ) ‖ 2 2,2 ) . Failure modes of vanilla KD Thm . 1 also hints at two distinct ways in which vanilla distillation could fail . First , since the student only learns from the teacher and does not have access to the original labels , we would expect the student to be erroneous when the teacher probabilities are inaccurate due to model misspecification , an overly restrictive teacher function class , or insufficient training . Prop . 2 , proved in App . D , confirms that , in the worst case , student error suffers from inaccuracy due to this teacher underfitting even when both the student and teacher belong to low complexity model classes . Proposition 2 ( Impact of teacher underfitting on vanilla KD ) . There exists a classification problem in which the following properties all hold simultaneously with high probability for f0 =log ( p0 ) : • The teacher learns p̂ ( x ) = 1n ( 1+λ ) ∑n i=1yi for all x∈X via ridge regression with λ=Θ ( 1/n1/4 ) . • Vanilla KD with SEL loss and constant f̂ satisfies ‖f̂−f0‖22,2 ≥‖γ > f0 , p0 ( p̂−p0 ) ‖ 2 2,2 = Ω ( 1√ n ) , matching the dependence of the Thm . 1 upper bound up to a constant factor . • Enhanced KD with SEL loss , γ̂ ( t ) =diag ( 1 p̂ ( t ) ) , and constant f̂ satisfies ‖f̂−f0‖22,2 =O ( 1n ) . Second , the critical radius in Thm . 1 depends on the complexity of the teacher model class P . If P has a large critical radius , then the student error bound suffers due to potential teacher overfitting even if the teacher generalizes well . Prop . 3 , proved in App . E , shows that , in the worst case , this teacher overfitting penalty is unavoidable and does in fact lead to increased student error . This occurs as the student only has access to the teacher ’ s training set probabilities which , due to overfitting , need not reflect its test set probabilities . Proposition 3 ( Impact of teacher overfitting on vanilla KD ) . There exists a classification problem in which the following properties all hold simultaneously with high probability for f0 =E [ log ( p0 ( X ) ) ] : • The critical radius δn of the teacher-student function class G in Thm . 1 is a non-vanishing constant , due to the complexity of the teacher ’ s function class . • The Vanilla KD error ‖f̂−f0‖22,2 for constant f̂ with SEL loss is lower bounded by a non-vanishing constant , matching the δn dependence of the Thm . 1 upper bound up to a constant factor . • Enhanced KD with SEL loss , γ̂ ( t ) =0 , and constant f̂ satisfies ‖f̂−f0‖22,2 =O ( n−4/ ( 4+d ) ) . These examples serve to lower bound student performance in the worst case by the teacher ’ s critical radius and class probability MSE , matching the upper bounds given in Thm . 1 . However , we note that in other better-case scenarios vanilla distillation can perform better than the upper-bounding Thm . 1 would imply . In the next section , we adapt and generalize techniques from semiparametric inference to mitigate the effects of teacher overfitting and underfitting in all cases .

Summary: This paper gives generalization guarantees for vanilla knowledge distillation, where it identifies a large variance issue due to the teacher’s complexity in this guarantee, and then proposes alternative approaches to eliminate this issue. The framework it uses is inspired by semi-parametric methods. The main resulting algorithm has two key components: cross-fitting and a corrective surrogate loss. An SGD variant is also presented. Experiments illustrate some of the benefits of the cross-fitting approach.