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 formulates knowledge distillation as a semi-parametric inference problem. Then, the paper adapts techniques from semi-parametric inference to analyze the error of a student model and improve the performance of knowledge distillation. The presentation is overall of good quality and it is relatively easy for me to follow the logic flow. That being said, there is still some room of improving the clarity of the paper.