Original: K NOWLEDGE
Revised: K NOWLEDGE

Original: D ISTILLATION
Revised: D ISTILLATION

Original: AS S EMIPARAMETRIC
Revised: AS S EMIPARAMETRIC

Original: I NFERENCE
Revised: I NFERENCE

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: 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.
Revised: 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.

Original: Surprisingly, this two-step knowledge distillation process often leads to higher accuracy than training the student directly on labeled data.
Revised: Surprisingly, this two-step knowledge distillation process often leads to higher accuracy than training the student directly on labeled data.

Original: 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.
Revised: 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.

Original: By adapting modern semiparametric tools, we derive several new guarantees for the prediction error of standard distillation and develop several enhancements with improved guarantees.
Revised: By adapting modern semiparametric tools, we derive several new guarantees for the prediction error of standard distillation and develop several enhancements with improved guarantees.

Original: We validate our ﬁndings empirically on both tabular data and image data and observe consistent improvements from our knowledge distillation enhancements.
Revised: We validate our ﬁndings empirically on both tabular data and image data and observe consistent improvements from our knowledge distillation enhancements.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Knowledge distillation (KD) (Craven & Shavlik, 1996; Breiman & Shang, 1996; Bucila et al., 2006;
Revised: Knowledge distillation (KD) (Craven & Shavlik, 1996; Breiman & Shang, 1996; Bucila et al., 2006;

Original: Li et al., 2014; Ba & Caruana, 2014; Hinton et al., 2015) is a widely used model compression technique that enables deployment of ML models on devices such as phones, watches, and virtual assistants (Stock et al., 2020).
Revised: Li et al., 2014; Ba & Caruana, 2014; Hinton et al., 2015) is a widely used model compression technique that enables deployment of ML models on devices such as phones, watches, and virtual assistants (Stock et al., 2020).

Original: KD works by training a large teacher model and then a smaller student model that tries to match the teacher’s predictions.
Revised: KD works by training a large teacher model and then a smaller student model that tries to match the teacher’s predictions.

Original: It is easily applied to any classiﬁcation procedure and any teacher and works remarkably well across a wide variety of domains including computer vision (Hinton et al., 2015), natural language processing (Sanh et al., 2019), and tabular data settings (Liu et al., 2018; Tan et al., 2018; Fakoor et al., 2020).
Revised: It is easily applied to any classiﬁcation procedure and any teacher and works remarkably well across a wide variety of domains including computer vision (Hinton et al., 2015), natural language processing (Sanh et al., 2019), and tabular data settings (Liu et al., 2018; Tan et al., 2018; Fakoor et al., 2020).

Original: Despite KD’s ubiquity and importance, the reasons for its success remains elusive.
Revised: Despite KD’s ubiquity and importance, the reasons for its success remains elusive.

Original: The practice of KD is often guided by heuristics.
Revised: The practice of KD is often guided by heuristics.

Original: In practice, KD is often performed in diverse ways, for different tasks and different domains.
Revised: In practice, KD is often performed in diverse ways, for different tasks and different domains.

Original: It is often tailored to particular model types (e.g., CNN, RNNs, Transformer).
Revised: It is often tailored to particular model types (e.g., CNN, RNNs, Transformer).

Original: This precludes a general theoretical understanding of KD.
Revised: This precludes a general theoretical understanding of KD.

Original: Hinton et al.
Revised: Hinton et al.

Original: (2015) remarks that KD works by transferring the “dark knowledge” from the teacher’s predicted class probabilities.
Revised: (2015) remarks that KD works by transferring the “dark knowledge” from the teacher’s predicted class probabilities.

Original: As we detail below, several works have tried to explain how this dark knowledge can improve the student’s performance beyond training from scratch.
Revised: As we detail below, several works have tried to explain how this dark knowledge can improve the student’s performance beyond training from scratch.

Original: Our work is motivated by the observation of Menon et al.
Revised: Our work is motivated by the observation of Menon et al.

Original: (2020) that the teacher probabilities 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.
Revised: (2020) that the teacher probabilities 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.

Original: Building on this observation, we cast KD as a plug-in approach to semiparametric inference : distillation aims to ﬁt a student model ˆ f in the presence of nuisance (the Bayes probabilities p 0 ) with the teacher’s probabilities ˆ p as a plug-in estimate of p 0 .
Revised: Building on this observation, we cast KD as a plug-in approach to semiparametric inference : distillation aims to ﬁt a student model ˆ f in the presence of nuisance (the Bayes probabilities p 0 ) with the teacher’s probabilities ˆ p as a plug-in estimate of p 0 .

Original: This insight allows us to adapt modern tools from semiparametric inference to analyze the error of a distilled student in Sec. 3.
Revised: This insight allows us to adapt modern tools from semiparametric inference to analyze the error of a distilled student in Sec. 3.

Original: Our analysis further suggests several enhancements for KD.
Revised: Our analysis further suggests several enhancements for KD.

Original: For example, in Sec. 4 we show that avoiding sample reuse between the teacher and student through cross-ﬁtting (see, e.g., Chernozhukov et al., 2018) can improve
Revised: For example, in Sec. 4 we show that avoiding sample reuse between the teacher and student through cross-ﬁtting (see, e.g., Chernozhukov et al., 2018) can improve

Original: KD performance.
Revised: KD performance.

Original: Moreover, a common strategy in semiparametric inference is to improve plug-in estimation by reducing plug-in bias.
Revised: Moreover, a common strategy in semiparametric inference is to improve plug-in estimation by reducing plug-in bias.

Original: The method of orthogonal machine learning (Chernozhukov et al., 2018; Foster & Syrgkanis, 2019) suggests correcting the teacher’s probabilities by stepping in the y − ˆ p direction where y is an observed label vector.
Revised: The method of orthogonal machine learning (Chernozhukov et al., 2018; Foster & Syrgkanis, 2019) suggests correcting the teacher’s probabilities by stepping in the y − ˆ p direction where y is an observed label vector.

Original: We argue in Sec.
Revised: We argue in Sec.

Original: 4 that while this minimizes the teacher bias, the variance could explode.
Revised: 4 that while this minimizes the teacher bias, the variance could explode.

Original: We instead propose a correction that balances the bias and variance, leading to a better overall generalization bound.
Revised: We instead propose a correction that balances the bias and variance, leading to a better overall generalization bound.

Original: We empirically validate our theoretical analysis in Sec. 5.
Revised: We empirically validate our theoretical analysis in Sec. 5.

Original: In a synthetic setting where we have access to the Bayes classiﬁer, our modiﬁed loss yields estimators that are closer to the Bayes classiﬁer.
Revised: In a synthetic setting where we have access to the Bayes classiﬁer, our modiﬁed loss yields estimators that are closer to the Bayes classiﬁer.

Original: Moreover, our enhanced distillation improve dramatically over vanilla distillation in failure regimes of the teacher, e.g., when the teacher is heavily biased or overﬁtted.
Revised: Moreover, our enhanced distillation improves dramatically over vanilla distillation in failure regimes of the teacher, e.g., when the teacher is heavily biased or overﬁtted.

Original: On ﬁve real tabular datasets, cross-ﬁtting and corrected distillation improve the student’s performance by up to 4% AUC point compared to vanilla KD.
Revised: On ﬁve real tabular datasets, cross-ﬁtting and corrected distillation improve the student’s performance by up to 4% AUC point compared to vanilla KD.

Original: Furthermore, on a common image classiﬁcation dataset (CIFAR-10, Krizhevsky & Hinton, 2009), when the teacher model overﬁts, cross-ﬁtting and our corrected loss improve the student accuracy by up to 1.5%.
Revised: Furthermore, on a common image classiﬁcation dataset (CIFAR-10, Krizhevsky & Hinton, 2009), when the teacher model overﬁts, cross-ﬁtting and our corrected loss improve the student accuracy by up to 1.5%.

Original: This demonstrates the wide applicability of our enhancements. 
Revised: This demonstrates the wide applicability of our enhancements. Related work.

Revised: Since we do not aim review the vast literature on KD in its entirety here, we point the interested reader to Gou et al.

Revised: for a good overview of the results.

Revised: In the following we review theory research studying why distillation works.

Revised: We further go over some empirical studies and applications of knowledge distillation in the extended literature review in App. A.

Original: Related work 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.
Revised: 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.

Original: This was hypothesized in Hinton et al.
Revised: This was hypothesized in Hinton et al.

Original: (2015) with empirical justiﬁcation.
Revised: (2015) with empirical justiﬁcation.

Original: Phuong & Lampert (2019) consider the case where the teacher is a ﬁxed linear classiﬁer and the student is a linear classiﬁer or a deep linear network.
Revised: Phuong & Lampert (2019) consider the case where the teacher is a ﬁxed linear classiﬁer and the student is a linear classiﬁer or a deep linear network.

Original: They show that the student can learn the teacher perfectly if it has access to training examples more than the ambient dimensions.
Revised: They show that the student can learn the teacher perfectly if it has access to training examples more than the ambient dimensions.

Original: Vapnik & Izmailov (2015) discusses the setting of learning with privileged information where one has additional information at training time which is not available at test time.
Revised: Vapnik & Izmailov (2015) discusses the setting of learning with privileged information where one has additional information at training time which is not available at test time.

Original: Lopez-Paz et al. (2015) draws a connection between this and distillation, arguing that distillation is effective because the teacher learns a better representation allowing the student to learn at a faster rate.
Revised: Lopez-Paz et al. (2015) draws a connection between this and distillation, arguing that distillation is effective because the teacher learns a better representation allowing the student to learn at a faster rate.

Original: They hypothesize that a teacher’s soft predictions indicating if an example is hard or easy enables this.
Revised: They hypothesize that a teacher’s soft predictions indicating if an example is hard or easy enables this.

Original: Tang et al.
Revised: Tang et al.

Original: (2020) argues using empirical evidence that label smoothing and reweighting of training examples using the teacher’s
Revised: (2020) argues using empirical evidence that label smoothing and reweighting of training examples using the teacher’s

Original: predictions are key to beneﬁts offered by knowledge distillation.
Revised: predictions are key to beneﬁts offered by knowledge distillation.

Original: Mobahi et al.
Revised: Mobahi et al.

Original: (2020) analyzed the case of self-distillation where the student and teacher function classes are the same.
Revised: (2020) analyzed the case of self-distillation where the student and teacher function classes are the same.

Original: They showed that this is equivalent to increasing regularisation strength for kernelized models.
Revised: They showed that this is equivalent to increasing regularisation strength for kernelized models.

Original: Bu et al.
Revised: Bu et al.

Original: (2020) considers 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.
Revised: (2020) considers 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.

Original: Menon et al.
Revised: Menon et al.

Original: (2020) consider the case of losses such that the population risk is linear in the Bayes class probabilities.
Revised: (2020) consider the case of losses such that the population risk is linear in the Bayes class probabilities.

Original: 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.
Revised: 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.

Original: They show that the variance of the Bayes distilled empirical risk is lower than the empirical risk.
Revised: They show that the variance of the Bayes distilled empirical risk is lower than the empirical risk.

Original: Then using analysis from Maurer & Pontil (2009); Bennett (1962), they derive the excess risk of the distilled empirical risk as a function of the (cid:96) 2 distance between the teacher’s class probabilities and the Bayes class probabilities.
Revised: Then using analysis from Maurer & Pontil (2009); Bennett (1962), they derive the excess risk of the distilled empirical risk as a function of the (cid:96) 2 distance between the teacher’s class probabilities and the Bayes class probabilities.

Original: We signiﬁcantly depart from Menon et al.
Revised: We signiﬁcantly depart from Menon et al.

Original: (2020) in multiple ways: i) our Thm. 1 allows for the common practice of data re-use, ii) our results cover the standard distillation losses (1) and (2) which are non-linear in p 0 , iii) we use localized Rademacher analysis to achieve tight fast rates for strongly-convex losses; which the majority of distillation losses satisﬁes, iv) we use techniques from semiparametric inference to improve upon vanilla distillation.
Revised: (2020) in multiple ways: i) our Thm. 1 allows for the common practice of data re-use, ii) our results cover the standard distillation losses (1) and (2) which are non-linear in p 0 , iii) we use localized Rademacher analysis to achieve tight fast rates for strongly-convex losses, which are commonly used for distillation, and iv) we use techniques from semiparametric inference to improve upon vanilla distillation.

Original:  K NOWLEDGE D ISTILLATION B ACKGROUND
Revised:  K NOWLEDGE D ISTILLATION B ACKGROUND

Original: In this section, we present our problem setup and an introduction to knowledge distillation.
Revised: We consider a multiclass classiﬁcation problem with k classes and n training datapoints z i =( x

Original: We are given a training dataset { ( x 1 ,y 1 ),..., ( x n ,y n ) } sampled independently from someprobability distribution Pwith w :=( x,y ) ∈X ×Y =: W . The goal is to learn some function f : X →Y from some function class F that minimizes the risk (i.e., expected loss over the population distribution P ):
Revised: i ,y i ) sampled independently from some distribution P . Each feature vector x belongs to a set X , each label vector y ∈{ e 1 ,...,e j }⊂{ 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 p 0 ( x )

Original: [Equation] where (cid:96) :
Revised: =

Original: W×Y → Ris some given loss function. In multiclass classiﬁcation with k classes, the label vector y ∈{ e 1 ,...,e j }⊂{ 0 , 1 } k is a one-hot encoding of the class label.
Revised: E [ Y | X = x ] . Our aim is to identify a scoring rule f : X → R k that minimizes a prediction loss on average under the distribution P .

Original: Knowledge distillation Knowledge distillation (KD) is a two-step training process where one ﬁrst uses a labeled dataset to train a teacher model and then trains a student model to predict the teacher’s predicted class probabilities.
Revised: Knowledgedistillation. Knowledgedistillation(KD)isatwo-steptrainingprocesswhereoneﬁrstuses alabeleddatasettotrainateachermodelandthentrainsastudentmodeltopredicttheteacher’spredictedclassprobabilities.

Original: Typically the teacher model is larger and more cumbersome, while the student is smaller and more efﬁcient.
Revised: Typicallytheteachermodelislargerandmorecumbersome,whilethestudentis smallerandmoreefﬁcient.

Original: Knowledge distillation was ﬁrst motivated by model compression (Bucila et al., 2006), to ﬁnd compact yet high-performing models to be deployed (such as on mobile devices).
Revised: Knowledgedistillationwasﬁrstmotivatedbymodelcompression(Bucila etal. ,2006),toﬁndcompactyethigh-performingmodelstobedeployed(suchasonmobiledevices).

Original: In training the student to match the teacher’s prediction probability, there are several types of loss functions that are commonly used.
Revised: In training the student to match the teacher’s prediction probability, there are several types of loss functions that are commonly used.

Original: Let ˆ p ( x ) ∈
Revised: Let ˆ p ( x ) ∈

Original: R k be the teacher’s vector of predicted class probabilities, f ( x ) ∈
Revised: R k be the teacher’s vector of predicted class probabilities, f ( x ) ∈

Original: R k be the student model’s output, and [ k ] (cid:44) { 1 , 2 ,...,k } .
Revised: R k be the student model’s output, and [ k ] (cid:44) { 1 , 2 ,...,k } .

Original: The most popular  loss functions(cid:96) ( w ; f ( x ) , ˆ p ( x )) include the squared error logit loss (Ba & Caruana, 2014)
Revised: The most popular distillation loss functions 1 (cid:96) ( z ; f ( x ) , ˆ p ( x )) include the squared error logit loss (Ba & Caruana, 2014)

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: ) for an inverse temperature β > 0 .
Revised: for an inverse temperature β > 0 .

Original: These loss functions measure the divergence between the probabilities predicted by the teacher and the student.
Revised: These loss functions measure the divergence between the probabilities predicted by the teacher and the student.

Original: 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).
Revised: 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).

Original: In Secs.
Revised: In Secs.

Original: 3 and 4, we will adapt modern tools from semiparametric inference to understand and enhance this phenomenon.
Revised: 3 and 4, we will adapt modern tools from semiparametric inference to understand and enhance this phenomenon.

Original:  D ISTILLATION AS S EMIPARAMETRIC I NFERENCE
Revised:  D ISTILLATION AS S EMIPARAMETRIC I NFERENCE

Original: In semiparametric inference (Kosorok, 2007), one aims to estimate a target parameter or function f 0 , but that estimation depends on an auxiliary nuisance function p 0 that is unknown and not of primary interest.
Revised: In semiparametric inference (Kosorok, 2007), one aims to estimate a target parameter or function f 0 , but that estimation depends on an auxiliary nuisance function p 0 that is unknown and not of primary interest.

Original: We cast the knowledge distillation process as a semiparametric inference problem, by treating the  teacher’s  probabilities as a nuisance function that approximates the Bayes probability p 0 ( x )=
Revised: We cast the knowledge distillation process as a semiparametric inference problem, by treating the unknown Bayes class probabilities p 0 as nuisance and the teacher’s predicted probabilities as a plug-in estimate of that nuisance.

Original: E [ Y | X = x ] .

Original: This perspective allows us bound the generalization of the student in terms of the MSE between the teacher and the Bayes probabilities.
Revised: 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.

Original: In the next section (Sec. 4) we use techniques from semi-parametric inference to enhance the performance of the student.
Revised: In the next section (Sec. 4) we use techniques from semiparametric inference to enhance the performance of the student.

Revised: The interested reader could consult Tsiatis (2007) for more details on semiparametric inference.

Original: Our analysis starts from taking the following perspective on distillation.
Revised: Our analysis starts from taking the following perspective on distillation.

Original: We view the goal of the student as minimizing an oracle population loss over a function class F , 
Revised: For a given pointwise loss function (cid:96) ( z ; f ( x ) ,p 0 ( x )) , we view the goal of the student as minimizing an oracle population loss over a function class F , [Equation] ( Z

Original: [Equation]
Revised: [Equation]

Revised: f ( X ) ,p 0 ( X ))] with f 0 (cid:44) argmin f ∈F L D ( f,p 0 ) .

Original: The main hurdle is that this is objective depends on the unknown Bayes probabilities p 0 .
Revised: The main hurdle is that this is objective depends on the unknown Bayes probabilities p 0 .

Original: We view the teacher’s model ˆ p as an approximate version of p 0 and bound the distillation error of the student as a function of the teacher’s estimation error.
Revised: We view the teacher’s model ˆ p as an approximate version of p 0 and bound the distillation error of the student as a function of the teacher’s estimation error.

Original: Typical semi-parametric inference considers cases where f 0 is a ﬁnite dimensional parameter, however recent work of Foster & Syrgkanis (2019) extends this framework to inﬁnite dimensional models fand proposes the statistical learning theory with a nuisance component framework.
Revised: Typical semiparametric inference considers cases where f 0 is a ﬁnite dimensional parameter; however recent work of Foster & Syrgkanis (2019) extends this framework to inﬁnite dimensional models f 0 and to develop statistical learning theory with a nuisance component framework.

Original: The distillation problem ﬁts exactly into this setup.
Revised: The distillation problem ﬁts exactly into this setup.

Original: Bounds on vanilla distillation with data re-use As a ﬁrst step we derive a vanilla bound on the error of the distilled student model without any further modiﬁcations 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  on the plug-in loss, plugging in the teacher’s model instead of p 0 , i.e.,
Revised: Bounds on vanilla distillation with data re-use As a ﬁrst step we derive a vanilla bound on the error of the distilled student model without any further modiﬁcations 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 p 0 , i.e.,

Original: [Equation]
Revised: [Equation]

Original: where E n [ X ]= 1 n (cid:80) ni
Revised: where E n [ X ]= 1 n (cid:80) ni

Original: =1 X i denotes the empirical expectation of a random variable.
Revised: =1 X i denotes the empirical expectation of a random variable.

Original: Technical deﬁnitions Before presenting our main theorem we introduce some technical notation.
Revised: Technical deﬁnitions Before presenting our main theorem we introduce some technical notation.

Original: For a vector valued function f that takes as input a random variable X , we use the shorthand notation (cid:107) f (cid:107) p,q (cid:44) (cid:107)(cid:107) f
Revised: For a vector valued function f that takes as input a random variable X , we use the shorthand notation (cid:107) f (cid:107) p,q (cid:44) (cid:107)(cid:107) f

Original: ( X ) (cid:107) p (cid:107) L q = E (cid:2) (cid:107) f
Revised: ( X ) (cid:107) p (cid:107) L q = E (cid:2) (cid:107) f

Original: ( X ) (cid:107) qp (cid:3) 1 /q .
Revised: ( X ) (cid:107) qp (cid:3) 1 /q .

Original: Let ∇ φ and ∇ π denote the partial derivatives of (cid:96) ( w ; φ, π ) , with respect to its second and third input correspondingly and ∇ φπ the Jacobian of cross partial derivatives, i.e. [ ∇ φπ (cid:96) ( w ; φ,π )] i,j = {+∂ 
+}
Revised: Let ∇ φ and ∇ π denote the partial derivatives of (cid:96) ( z ; φ,π ) , with respect to its second and third input correspondingly and ∇ φπ the Jacobian of cross partial derivatives, i.e., [ ∇ φπ (cid:96) ( z ; φ,π )] i,j = {+∂ 
+}

Original: [Equation]

Original:  j ∂π i (cid:96) ( w ; φ,π ) .
Revised: ∂φ j ∂π i (cid:96) ( z ; φ,π ) .

Original: Finally, let q f,p ( x )=
Revised: Finally, let q f,p ( x )=

Original: E [ ∇ φπ (cid:96) ( W ; f ( X ) ,p ( X )) | X = x ] and γ f,p ( x )=
Revised: E [ ∇ φπ (cid:96) ( Z ; f ( X ) ,p ( X )) | X = x ] and γ f,p ( x )=

Original: E U ∼ Unif ([0 , 1])
Revised: E U ∼ Unif ([0 , 1])

Original: [ q f,Up +(1 − U ) p 0 ( x )] .
Revised: [ q f,Up +(1 − U ) p 0 ( x )] .

Original: Critical radius Finally, we need to deﬁne the notion of the critical radius (see e.g., Wainwright (2019)) of a function class, which typically provides tight learning rates for statistical learning theory tasks.
Revised: Critical radius Finally, we need to deﬁne the notion of the critical radius (see e.g., Wainwright (2019)) of a function class, which typically provides tight learning rates for statistical learning theory tasks.

Original: For any function class F we deﬁne the localized Rademacher complexity as:
Revised: For any function class F we deﬁne the localized Rademacher complexity as:

Original: [Equation]
Revised: [Equation]

Original: where (cid:15) i are i.i.d.
Revised: where (cid:15) i are i.i.d.

Original: random variables taking values equiprobably in {− 1 , 1 } .
Revised: random variables taking values equiprobably in {− 1 , 1 } .

Original: The critical radius of a class F , taking values in [ − H,H ] , is the smallest solution δ n to the inequality R ( δ ; F ) ≤ δ 2 H .
Revised: The critical radius of a class F , taking values in [ − H,H ] , is the smallest solution δ n to the inequality R ( δ ; F ) ≤ δ 2 H .

Original: Theorem 1 (Vanilla distillation) .
Revised: Theorem 1 (Vanilla distillation analysis) .

Original: Assume that f 0 ∈ F and F is a convex set and that the teacher estimates a model ˆ p ∈ P on the same sample as the student.
Revised: Assume that f 0 ∈ F and F is a convex set and that the teacher estimates a model ˆ p ∈P on the same sample as the student.

Original: Let δ n,ζ = δ n + c 0 (cid:113) log( c 1 /ζ )
Revised: Let δ n,ζ = δ n + c 0 (cid:113) log( c 1 /ζ )

Original: n for some universal constants c 0 ,c 1 and δ n is an upper bound on the critical radius the function class:
Revised: n for some universal constants c 0 ,c 1 and δ n is an upper bound on the critical radius the function class:

Original: [Equation]
Revised: [Equation]

Original: Assume that, for all w ∈W and π , the loss (cid:96) ( w ; φ,π ) is σ -strongly convex with respect to φ , uniformly bounded in [ − H,H ] for all w,φ,π and let: µ ( w )
Revised: Assume that, for all z ∈Z and π , the loss (cid:96) ( z ; φ,π ) is σ -strongly convex with respect to φ , uniformly bounded in [ − H,H ] for all z,φ,π and let: µ ( z )

Original: = sup φ (cid:107)∇ φ (cid:96) ( w ; φ, ˆ p ( x )) (cid:107) 2 .
Revised: = sup φ (cid:107)∇ φ (cid:96) ( z ; φ, ˆ p ( x )) (cid:107) 2 .

Original: Moreover, suppose that the function class f satisﬁes the (cid:96) 2 /(cid:96) 4 ratio condition: sup f ∈F (cid:107)
Revised: Moreover, suppose that the function class f satisﬁes the (cid:96) 2 /(cid:96) 4 ratio condition: sup f ∈F (cid:107)

Original: f − f 0 (cid:107) 2 , 4 (cid:107) f − f 0 (cid:107) 2 , 2 ≤ C.
Revised: f − f 0 (cid:107) 2 , 4 (cid:107) f − f 0 (cid:107) 2 , 2 ≤ C.

Original: Then ˆ f (3) satisﬁes (cid:107) ˆ f − f 0 (cid:107) 22 , 2 ≤ 1 σ 2 O ( δ 2 n,ζ C 2 H 2 (cid:107) µ (cid:107) 24 + (cid:107) γ (cid:62) f 

Revised: Then ˆ f (3) satisﬁes (cid:107) ˆ f − f 0 (cid:107) 22 , 2 ≤ 1 σ 2 O ( δ 2 n,ζ C 2 H 2 (cid:107) µ (cid:107) 24 + (cid:107) γ (cid:62) f

Original: ˆ p (ˆ p − p 0 )
Revised: ˆ p (ˆ p − p 0 )

Original: (cid:107) 22 , 2 ) with probability 1 − ζ.
Revised: (cid:107) 22 , 2 ) with probability 1 − ζ.

Original: Remark
Revised: Remark

Original: The proof is in App.
Revised: The proof is in App.

Original: B. We note that the (cid:96) 2 /(cid:96) 4 ratio requirement can be removed, albeit the ﬁnal bound will depend on a uniform bound (cid:107) µ (cid:107) ∞
Revised: D. We note that the (cid:96) 2 /(cid:96) 4 ratio requirement can be removed, albeit the ﬁnal bound will depend on a uniform bound (cid:107) µ (cid:107) ∞

Original: =
Revised: =

Original: sup w | µ ( w ) | , instead of (cid:107) µ (cid:107) 4 .
Revised: sup z | µ ( z ) | , instead of (cid:107) µ (cid:107) 4 .

Original: Moreover, we note thatthestrongconvexityrequirementfor (cid:96) issatisﬁedbyallstandarddistillationobjectives, asitisstrong convexity with respect to the output of f and not the parameters of f .
Revised: Moreover, we note thatthestrongconvexityrequirementfor (cid:96) issatisﬁedbyallstandarddistillationobjectives, asitisstrong convexity with respect to the output of f and not the parameters of f .

Original: Even this requirement could be removed, but this would yield slow rate bounds of the form: (cid:107) ˆ f − f 0 (cid:107) 22 , 

Revised: Even this requirement could be removed, but this would yield slow rate bounds of the form: (cid:107) ˆ f − f 0 (cid:107) 22 ,

Original: =
Revised: =

Original: O ( δ n,ζ + (cid:107) γ (cid:62) f 0 , ˆ p (ˆ p − p 0 )
Revised: O ( δ n,ζ + (cid:107) γ (cid:62) f 0 , ˆ p (ˆ p − p 0 )

Original: (cid:107) 22 , 2 ) .
Revised: (cid:107) 22 , 2 ) .

Original: Thm. 1 shows that the standard approach to distillation can provably work to provide an accurate student and potentially improve upon training from scratch.
Revised: Thm. 1 shows that the standard approach to distillation can provably work to provide an accurate student and potentially improve upon training from scratch.

Original: For instance, if the teacher is very accurate, i.e. ˆ p ≈ p 0 , then the term that remains contains the fourth moment of the loss gradient (cid:107) µ (cid:107) 4 .
Revised: For instance, if the teacher is very accurate, i.e., ˆ p ≈ p 0 , then the term that remains contains the fourth moment of the loss gradient (cid:107) µ (cid:107) 4 .

Original: If we were to replace the student ˆ p ( x ), with the raw labels y , then we would have gotten a dependence on the fourth moment of the loss gradient ˜ µ ( w ) (cid:44) sup f ∈F (cid:107)∇ φ (cid:96) ( w ; f ( x ) ,y ) (cid:107) 2 , which now contains extra variation due to the variance of the labels y .
Revised: If we were to replace the teacher ˆ p ( x ) with the raw labels y , then we would have gotten a dependence on the fourth moment of the loss gradient ˜ µ ( z ) (cid:44) sup f ∈F (cid:107)∇ φ (cid:96) ( z ; f ( x ) ,y ) (cid:107) 2 , which now contains extra variation due to the variance of the labels y .

Original: Thus in practice it can be that (cid:107) µ (cid:107) 4 (cid:28)(cid:107) ˜
Revised: Thus in practice it can be that (cid:107) µ (cid:107) 4 (cid:28)(cid:107) ˜

Original: µ (cid:107) 4 .
Revised: µ (cid:107) 4 .

Original: Drawbacks of vanilla distillation There are two major drawbacks of Thm. 1, which are inherent drawbacks of vanilla distillation.
Revised: Drawbacks of vanilla distillation There are two major drawbacks of Thm. 1, corresponding to two distinct sources of KD student error: error due to teacher overﬁtting and error due to teacher underﬁtting.

Original: First, the function class G contains the variation of the student’s model over the whole P .
Revised: First, the function class G contains the variation of the student’s model over the whole P , the function class of the teacher.

Original: If P is a large model space, with a large critical radius, then G, will inherit this complexity and the ﬁnal bound will depend on the sample complexity of the teacher.
Revised: If P is a large model space, with a large critical radius, then G will inherit this complexity and the ﬁnal bound will depend on the sample complexity of the teacher.

Original: Second, if the teacher is not very accurate, i.e. (cid:107) ˆ p − p 0 (cid:107) 2 , 2 is large, then the student’s model will be signiﬁcantly impacted.
Revised: Second, if the teacher is not very accurate, i.e., (cid:107) ˆ p − p 0 (cid:107) 2 , 2 is large, then the student’s model will be signiﬁcantly impacted.

Original: In the next section, we use techniques from semi-parametric inference to tackle both of these problems.
Revised: In the next section, we use techniques from semiparametric inference to tackle both of these problems.

Original: E NHANCING K NOWLEDGE D ISTILLATION
Revised: We show two examples to illustrate how the problems of teacher overﬁtting/underﬁtting could manifest themselves in knowledge distillation. These examples serve to lower bound student performance, in the worst case, by the teacher’s critical radius and the teacher’s 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.

Revised: Example 1 (Impact of teacher overﬁtting on vanilla distillation) .

Revised: Suppose that the Bayes probability function p 0 has Lipschitz gradient and, for known (cid:15)> 0 , belongs to the set

Revised: [Equation]

Revised: Suppose moreover that X ∈ R d has Lebesgue density bounded away from 0 and ∞ and that (cid:15) < 14 E [ p 0 ,j ( X )(1 − p 0 ,j ( X )) 2 ] E [(1 − p 0 ,j ( X )) /p 0 ,j ( X )] for each j .

Revised: Consider the teacher estimates ˆ p j ( x ) = max( (cid:15), ˜ p

Revised: j ( x )) for ˜ p the Nadaraya-Watson kernel smoothing estimator (Nadaraya, 1964; Watson, 1964)

Revised: [Equation]

Revised: if

Revised: [Equation]

Original: {+) otherwise with kernel K ( x ) = (cid:107) x (cid:107) − a 
+}
Revised: {+) otherwise with kernel K ( x ) = (cid:107) x (cid:107) − a 
+}

Revised: I [ (cid:107) x

Revised: (cid:107) 2 ≤ 1] , a ∈ (0 , d/ 2) , and h = n − 1 / (4+ d ) .

Revised: Notably, ˜ p exactly interpolates the observed labels (i.e., ˜ p ( x i ) =

Revised: y i )

Revised: but, by Belkin et al.

Revised: (2019, Thm. 1), also satisﬁes

Revised: E [ (cid:107) p 0 − ˆ p (cid:107) 22 , 2 ]

Revised: =

Revised: O ( n − 4 / (4+ d ) ) .

Revised: Consider now a student that learns a constant prediction rule via distilled ERM (3) with the squared error logit loss (1) and

Revised: [Equation]

Revised: [log( (cid:15) ) , 0] k for all x,x (cid:48) ∈X} .

Revised: Since the student only has access to the teacher’s training set probabilities, its estimate ˆ f ( x )= 1 n (cid:80) ni =1 log(max( y i ,(cid:15) )) is inconsistent for the optimal constant rule f 0 ( x )=

Revised: E [log( p 0 ( X ))] as f 0 ,j ( x ) − E ˆ f j ( x )=

Revised: E [log

Revised: [Equation]

Revised: p 0 ,j ( X ))

Revised: − log

Revised: [Equation]

Revised: max

Revised: [Equation]

Revised: Y j ,

Revised: [Equation]

Revised: ))] ≥ E [ p

Revised: [Equation]

Revised: j

Revised: [Equation]

Revised: X

Revised: [Equation]

Revised: ] by Taylor’s theorem with Lagrange remainder.

Revised: This non-vanishing student error reﬂects the nonvanishing critical radius δ n of the composite student-teacher function class G deﬁned in Thm. 1; since the student function class F has low complexity, the complexity of G is driven by the highly ﬂexible interpolating teacher.

Revised: We will demonstrate how to avoid this dependence on teacher complexity in Sec. 4.2.

Revised: Example 2 (Impact of teacher underﬁtting on vanilla distillation) .

Revised: Suppose that p 0 belongs to the set (4) and does not vary with x .

Revised: Consider the constant teacher estimate ˆ p ( x )=

Revised: ¯ y 1+ λ obtained via ridge regression with regularization strength λ and ¯ y (cid:44) 1 n (cid:80) ni

Original: {+=
+}
Revised: {+=
+}

Revised: y i .

Revised: A constant student prediction rule trained via distilled ERM (3) with the squared error logit loss (1) yields ˆ f ( x )=log( ¯ y 1+ λ ) for all x ∈X and satisﬁes

Revised: [Equation]

Revised: by the concavity of the logarithm.

Revised: Since P (¯ y j − p 0 ,j ( x ) ≥ θp 0 ,j ( x )) ≤ ζk for θ ≥ (cid:113) 2(1 − p 0 ( x ))

Original: {+np 
+}
Revised: {+np 
+}

Revised: ( x ) log( kζ ) +1 np 0 ( x ) log( kζ ) by Bernstein’s inequality (Bernstein, 1946), we have

Revised: P ( (cid:107) f 0 − ˆ f (cid:107) 22 , 2 ≥ (cid:107) γ (cid:62) f 0 ,p 0 ( p 0 − ˆ p ) (cid:107) 22 , 2 ) ≥ 1 − ζ matching the teacher MSE dependence of the Thm.upper bound whenever λ ≥ (cid:113) 2 n(cid:15) log( kζ )+ 43 1 n(cid:15) log( kζ ) .

Revised: Moreover, since limsup n

Revised: [Equation]

Revised: with probability 1 by the law of the iterated logarithm, (cid:107) f

Revised: [Equation]

Revised: (min(1 ,λ 2 )) with probability 1 whenever λ ≥ (cid:113)

Revised: [Equation]

Revised:  E NHANCING K NOWLEDGE D ISTILLATION

Revised: To address the two distinct inefﬁciencies of vanilla distillation revealed in Sec. 3, we will adapt and generalize two distinct techniques from semiparametric inference: orthogonal correction and cross-ﬁtting.

Original: We can view the plug-in distillation loss (cid:96) ( w ; f ( x ) , ˆ p ( x )) as a zeroth order Taylor approximation to the ideal loss (cid:96) ( w ; f ( x ) ,p 0 ( x )) around ˆ p .
Revised: We can view the plug-in distillation loss (cid:96) ( z ; f ( x ) , ˆ p ( x )) as a zeroth order Taylor approximation to the ideal loss (cid:96) ( z ; f ( x ) ,p 0 ( x )) around ˆ p .

Original: An ideal ﬁrst-orderTaylor approximationto the loss would take the form
Revised: An ideal ﬁrst-order approximation would take the form

Original: [Equation]
Revised: [Equation]

Original: However, its computation also requires knowledge of p 0 .
Revised: However, its computation also requires knowledge of p 0 .

Original: Nevertheless, since p 0 ( x )=
Revised: Nevertheless, since p 0 ( x )=

Original: E [ Y | X = x ] , we can always construct an unbiased estimate of the ideal ﬁrst order term by replacing p 0 ( x ) with y :  (cid:96)ortho ( w ; f ( x ) , ˆ p ( x ))=
Revised: E [ Y | X = x ] , we can always construct an unbiased estimate of the ideal ﬁrst order term by replacing p 0 ( x ) with y : [Equation] For standard distillation base losses like (1) and (2), the orthogonal loss (5) has an especially simple form, as ∇ π (cid:96) ( z ; f ( x ) , ˆ p ( x )) is linear in f .

Original:  (cid:96) ( w ; f ( x ) , ˆ p ( x ))+ (cid:104)
Revised: Indeed, this is true more generally for the following class of Bregman divergence losses. Deﬁnition 1 (Bregman divergence losses) . Any Bregman divergence loss function of the form (cid:96) ( z ; f ( x ) ,p ( x )) (cid:44) Ψ( f ( x ))

Original: y − ˆ p ( x  ) , E [ ∇ π (cid:96)  ( w ; f ( x ) , ˆ p ( x ))| x ] (cid:105)
Revised: − Ψ( g ( p ( x ))) −(cid:104)∇ g Ψ( g ( p ( x ))) ,f ( x ) − g ( p ( x )) (cid:105) has (cid:96) ortho ( z ; f ( x ) ,p ( x ))

Original: This correction in the direction of y − ˆ p can potentially improve approximation error.

Original: To allow more ﬂexibility so as to balance bias and variance, we consider a more general form of a correction term, which we call the γ -corrected loss for γ : X → R k × R k a matrix valued function: (cid:96)γ ( w ; f ( x ) ,p ( x )) 
Revised: = (cid:96) ( z ; f ( x ) ,p ( x )) +( y − p ( x )) (cid:62) ∇ p g ( p ( x )) (cid:62) ∇ 2 gg Ψ( g ( p ( x ))) f ( x )+const (6) with the second term bilinear in f ( x ) and y − p ( x ) . For the squared error logit loss (1) , Ψ( s )=

Original: (cid:44) (cid:96) ( w ; f ( x ) ,p ( x ))
Revised: 12 (cid:107) s (cid:107) 22 , g ( p ) =

Original: +( y − p ( x )) (cid:62) γ ( x ) f ( x )  .We denote associated population and empirical risks by
Revised: log( p ) , and the correction matrix ∇ p g ( p ( x )) (cid:62) ∇ 2 gg Ψ( g ( p ( x ))) = diag( 1 p ( x ) ) .

Original: L
Revised: Similarly, the annealed cross-entropy loss (2) falls into the class of Bregman divergence losses. We will show that orthogonal correction (5) can signiﬁcantly improve student bias due to teacher underﬁtting; however, for our standard distillation losses (1) and (2), the same orthogonal correction term often introduces unreasonably large variance due to division by small probabilities appearing in the correction matrix (see Deﬁnition 1). To grant ourselves more ﬂexibility in balancing bias and variance, we propose and analyze a family of γ -corrected losses , parameterized by a matrix valued function γ : X → R k × R k :

Original: [Equation]
Revised: [Equation]

Original: and
Revised: (cid:96) (

Original: [Equation]
Revised: [Equation]

Original: p
Revised: f

Original: [Equation]
Revised: [Equation]

Original: f
Revised: x

Original: [Equation]
Revised: [Equation]

Original: X
Revised: ,p [Equation] +

Original: [Equation]
Revised: [Equation]

Original: p
Revised: p

Original: [Equation]
Revised: [Equation]

Original: X
Revised: ) (cid:62) γ ( x

Original: [Equation]
Revised: [Equation]

Revised: f ( x ) to mimic the bilinear structure of Bregman orthogonal losses (6).

Revised: Note that we can always recover the vanilla distillation loss by taking γ ≡ 0 .

Revised: We denote the associated population and empirical risks by

Revised: L D ( f,p,γ ) (cid:44) E

Revised: [ (cid:96) γ ( Z ; f ( X ) ,p ( X ))] and L n ( f,p,γ ) (cid:44) E n [ (cid:96) γ

Original:  Observe that at p 0 the correction term is mean-zero and hence L D ( f,p 0 ,γ ) is independent of γ
Revised: ( Z ; f ( X ) ,p ( X ))] . Observe that at p 0 the correction term is mean-zero and hence L D ( f,p 0 ,γ ) is independent of γ

Original: L D ( f,p 0 ) (cid:44) E
Revised: L D ( f,p 0 ) (cid:44) E

Original: [ (cid:96) ( W ; f ( X ) ,p 0 ( X ))
Revised: [ (cid:96) ( Z ; f ( X ) ,p 0 ( X ))

Original: ]=
Revised: ]=

Original: L D ( f,p 0 ,γ ) for all γ.
Revised: L D ( f,p 0 ,γ ) for all γ.

Original: 4. R OBUSTNESS OF γ -C ORRECTED L OSS TO T EACHER E RROR

Original: The γ -corrected loss has strong connections to the literature on Neyman orthogonality (Chernozhukov et al., 2018;
Revised: The γ -corrected loss has strong connections to the literature on Neyman orthogonality (Chernozhukov et al., 2018;

Original: Chernozhukov et al., 2016; Nekipelov et al., 2018; Chernozhukov et al., 2018; Foster & Syrgkanis, 2019).
Revised: Chernozhukov et al., 2016; Nekipelov et al., 2018; Chernozhukov et al., 2018; Foster & Syrgkanis, 2019).

Original: In particular, if the function γ is set appropriately, then one can show that the γ -corrected loss function satisﬁes the condition of a Neyman orthogonal loss deﬁned by Foster &
Revised: In particular, if the function γ is set appropriately, then one can show that the γ -corrected loss function satisﬁes the condition of a Neyman orthogonal loss deﬁned by Foster &

Original: Syrgkanis (2019).
Revised: Syrgkanis (2019).

Original: We begin our analysis by showing a general lemma for any estimator ˆ f , which adapts the main theorem of Foster & Syrgkanis (2019) to account for approximate orthogonality; the proof can be found in App. A.
Revised: We begin our analysis by showing a general lemma for any estimator ˆ f , which adapts the main theorem of Foster & Syrgkanis (2019) to account for approximate orthogonality; the proof can be found in App. C.

Original: Lemma 2 (Algorithm-agnostic analysis) .
Revised: Lemma 2 (Algorithm-agnostic analysis) .

Original: Consider any estimation algorithm that produces an estimate ˆ f with small plug-in excess risk, i.e.,
Revised: Consider any estimation algorithm that produces an estimate ˆ f with small plug-in excess risk, i.e.,

Original: [Equation]
Revised: [Equation]

Original: and that the loss L D is σ -strongly convex with respect to f and F is a convex set.
Revised: and that the loss L D is σ -strongly convex with respect to f and F is a convex set.

Original: Then:σ
Revised: Then:

Original: [Equation]
Revised: [Equation]

Original: If, further, sup w,φ,π,i ∈ [ d ] (cid:107)∇ φ i ππ (cid:96) ( w ; φ,π ) (cid:107) op ≤ M , then:
Revised: If, further, sup z,φ,π,i ∈ [ d ] (cid:107)∇ φ i ππ (cid:96) ( z ; φ,π ) (cid:107) op ≤ M , then:

Original: [Equation]
Revised: [Equation]

Original: Connection to Neyman orthogonality Note that if we set γ = q f 0 , ˆ p , then we obtain that the MSE error of the student depends on the MSE error of the teacher only in a second order manner!
Revised: Connection to Neyman orthogonality Note that if we set γ = q f 0 , ˆ p , then we obtain that the MSE error of the student depends on the MSE error of the teacher only in a second-order manner!

Original: Morever, as we show in the Example 3, q f 0 , ˆ p is typically independent of f and hence a known closed form function of ˆ p .

Original: Setting γ = q f 0 , ˆ p renders the loss function Neyman orthogonal (Foster & Syrgkanis, 2019).
Revised: Setting γ = q f 0 , ˆ p renders the loss function Neyman orthogonal (Foster & Syrgkanis, 2019).

Original: However, we note that this setting of the γ can lead to larger variance, i.e., the achievable excess risk can be much larger than the excess risk without the γ -correction.

Original: For instance, in the case of the squared error logit loss q f 0 , ˆ p ( x )= 1ˆ p ( x ) , which can be really large when ˆ p is close to 0 , leading to a large increase in the variance of our loss.

Original: Thus unlike the standard approach in semiparametric inference, we will be choosing γ in practice to balance bias and variance.

Original: Example 3 (Bregman Divergence Losses) . For any Bregman divergence loss function of the form [Equation] − g
Revised: Moreover, q f 0 , ˆ p is an observable quantity for any Bregman divergence loss (Deﬁnition 1) as q f 0 , ˆ p is independent of f 0 .

Original: [Equation] π ) [Equation] we have
Revised: However, we note that this setting of the γ can lead to larger variance, i.e., the achievable excess risk can be much larger than the excess risk without the γ -correction.

Original: [Equation] For such losses, the quantities q f,p and γ f,p are independent of  f  , since:{+ˆ p ( x ) = 
+}
Revised: For instance, in the case of the squared error logit loss (1) q f 0 , {+ˆ p ( x ) = 
+}

Original: [Equation] Thus by choosing γ = q ˆ p  , which  is  a known closed form function of the teacher’s model, we achieve a second order dependence on the teacher’s estimation error. For the squared error logit loss (1) , Ψ( s ) =
Revised: ˆ p ( x ) , which can be excessively large when ˆ p is close to 0 , leading to a large increase in the variance of our loss.

Original: 12 (cid:107) s (cid:107) 22 , g ( p ) = log( p ) , and q p ( x ) = diag (cid:16) 1 p 1 ( x ) ,..., 1 p K ( x ) (cid:17) . Similarly, the annealed cross-entropy loss (2) falls into the class of Bregman divergence losses.
Revised: Thus, in a departure from the standard approach in semiparametric inference, we will be choosing γ in practice to balance bias and variance.

Original: Example instantiation of student’s estimation algorithm
Revised: Example instantiation of student’s estimation algorithm

Original: If we use plug-in empirical risk minimization , i.e., ˆ f = argmin f ∈F L n ( f, ˆ p,γ ) , to estimate f 0 with ˆ p estimated on an independent sample, then the results of Maurer & Pontil (2009) directly imply that as long as the loss function (cid:96) ( w ; φ,π ) is uniformly bounded in [ − H,H ] , then w.p. 1 − δ :
Revised: If we use plug-in empirical risk minimization , i.e., ˆ f = argmin f ∈F L n ( f, ˆ p,γ ) , to estimate f 0 with ˆ p estimated on an independent sample, then the results of Maurer & Pontil (2009) directly imply that as long as the loss function (cid:96) ( z ; φ,π ) is uniformly bounded in [ − H,H ] , then w.p. 1 − δ :

Original: [Equation]
Revised: [Equation]

Original: where τ ( n )= N ∞ (1 /n, F , 2 n ) and N ∞ ( (cid:15), F ,m ) is the (cid:96) ∞ empirical covering number of function class
Revised: where τ ( n )= N ∞ (1 /n, F , 2 n ) and N ∞ ( (cid:15), F ,m ) is the (cid:96) ∞ empirical covering number of function class

Original: F in the worst-case over all realizations of m samples and at approximation level (cid:15) .
Revised: F in the worst-case over all realizations of m samples and at approximation level (cid:15) .

Original: This result has two drawbacks: it is a slow rate result that scales as 1 / √ n for parametric of bounded VC-dimension classes and it requires the student to be ﬁt on a completely separate sample than the teacher.
Revised: This result has two drawbacks: it is a slow rate result that scales as 1 / √ n for parametric of bounded VC-dimension classes and it requires the student to be ﬁt on a completely separate sample than the teacher.

Original: In the next theorem, we address both of these drawbacks: i) we invoke localized Rademacher complexity analysis to provide a fast rate result which would be of the order of 1 /n for VC or parametric function classes, ii) we use a more sophisticated sample-splitting technique called cross-ﬁtting, which allows the student to be trained on all the available data.
Revised: In the next theorem, we address both of these drawbacks: i) we invoke localized Rademacher complexity analysis to provide a fast rate result which would be of the order of 1 /n for VC or parametric function classes, ii) we use a more sophisticated sample-splitting technique called cross-ﬁtting, which allows the student to be trained on all the available data.

Original: 4. C ROSS - FITTED ERM ANALYSIS
Revised: 4. C OMBATING TEACHER OVERFITTING WITH CROSS - FITTING

Original: We now provide a more sophisticated version of sample splitting to make use of all datapoints in our ﬁnal estimation, while at the same time not taking a hit by the sample complexity of the teacher’s function space.
Revised: We now provide a more sophisticated version of sample splitting to make use of all datapoints in our ﬁnal estimation, while at the same time not taking a hit by the sample complexity of the teacher’s function space.

Original: This approach is referred to as cross-ﬁtting in the semiparametric inference literature (see e.g., Chernozhukov et al.
Revised: This approach is referred to as cross-ﬁtting in the semiparametric inference literature (see e.g., Chernozhukov et al.

Original: Partition the dataset into B equally sized partitions P 1 ,...,P B 2.
Revised: Partition the dataset into B equally sized partitions P 1 ,...,P B 2.

Original: For each fold t ∈ [ B ] estimate ˆ p ( t ) and ˆ γ ( t ) using all the out-of-fold samples 3.
Revised: For each fold t ∈ [ B ] estimate ˆ p ( t ) and ˆ γ ( t ) using all the out-of-fold samples 3.

Original: Estimate ˆ f by minimizing the empirical loss:
Revised: Estimate ˆ f by minimizing the empirical loss:

Original: [Equation]
Revised: [Equation]

Original: In other words, the nuisance estimates ˆ γ ( t ) , ˆ p ( t ) that are evaluated on the samples in fold t , are estimated based only on samples outside of P t .
Revised: In other words, the nuisance estimates ˆ γ ( t ) , ˆ p ( t ) that are evaluated on the samples in fold t , are estimated based only on samples outside of P t .

Original: Theorem 4 (Cross-ﬁtted ERM analysis) .
Revised: Theorem 3 (Cross-ﬁtted ERM analysis) .

Original: Assume that f 0 ∈ F and F is a convex set.
Revised: Assume that f 0 ∈ F and F is a convex set.

Original: Let
Revised: Let

Original: [Equation]
Revised: [Equation]

Original: n
Revised: n

Original: [Equation]
Revised: [Equation]

Original: for some universal constants c 0 ,c 1 and δ n is an upper bound on the critical radii of the classes:
Revised: for some universal constants c 0 ,c 1 and δ n is an upper bound on the critical radii of the classes:

Original: [Equation]
Revised: [Equation]

Original: for each t ∈ [ B ] .
Revised: for each t ∈ [ B ] .

Original: Assume that, for all w ∈ W and π , the loss (cid:96) ( w ; φ, π ) is σstrongly convex with respect to φ , uniformly bounded in [ − H, H ] for all w, φ, π and let: µ ( w )
Revised: Assume that, for all z ∈Z and π , the loss (cid:96) ( z ; φ,π ) is σ -strongly convex with respect to φ , uniformly bounded in [ − H,H ] for all z,φ,π and let: µ ( z )

Original: =sup f ∈F ,t ∈
Revised: =sup f ∈F ,t ∈

Original: [ B ] (cid:13)(cid:13) ∇ φ (cid:96) ( w ; f ( X ) , ˆ p ( t ) ( x ))
Revised: [ B ] (cid:13)(cid:13) ∇ φ (cid:96) ( z ; f ( X ) , ˆ p ( t ) ( x ))

Original: (cid:13)(cid:13) 2 .
Revised: (cid:13)(cid:13) 2 .

Original: Moreover, suppose that the function class f satisﬁes the (cid:96) 2 /(cid:96) 4 ratio condition: sup f ∈F (cid:107)
Revised: Moreover, suppose that the function class f satisﬁes the (cid:96) 2 /(cid:96) 4 ratio condition: sup f ∈F (cid:107)

Original: f − f 0 (cid:107) 2 , 4 (cid:107) f − f 0 (cid:107) 2 , 2 ≤ C.
Revised: f − f 0 (cid:107) 2 , 4 (cid:107) f − f 0 (cid:107) 2 , 2 ≤ C.

Original: If ˆ f is the output of the cross-ﬁtted empirical risk minimization algorithm on a fresh set of n samples, then, w.p. 1 − ζ :
Revised: If ˆ f is the output of the cross-ﬁtted empirical risk minimization algorithm on a fresh set of n samples, then, w.p. 1 − ζ :

Original: [Equation]
Revised: [Equation]

Original: The proof is found in App.
Revised: The proof is found in App.

Original: C. Observe that unlike Thm. 1, the function classes G (ˆ p ( t ) , ˆ γ ( t ) ) in the latter theorem do not vary the teacher’s model over P , but rather only vary f ∈F and evaluate p at the out-of-fold estimate p ( t ) .
Revised: E. Observe that unlike Thm. 1, the function classes G (ˆ p ( t ) , ˆ γ ( t ) ) in the latter theorem do not vary the teacher’s model over P , but rather only vary f ∈F and evaluate p at the out-of-fold estimate p ( t ) .

Original: Since in practice the teacher’s model can be quite complex, getting rid of this dependence on the sample complexity of the teacher’s function space can bring immense improvement.
Revised: Since in practice the teacher’s model can be quite complex, getting rid of this dependence on the sample complexity of the teacher’s function space can bring immense improvement.

Original: Thus the critical radius of G (ˆ p ( t ) , ˆ γ ( t ) ) , can be signiﬁcantly smaller than G .
Revised: Thus the critical radius of G (ˆ p ( t ) , ˆ γ ( t ) ) , can be signiﬁcantly smaller than G .

Original: For instance, if P and F are VC-subgraph classes with VC dimensions d P (cid:29) d F , correspondingly, then the critical radius of
Revised: For instance, if P and F are VC-subgraph classes with VC dimensions d P (cid:29) d F , correspondingly, then the critical radius of

Original: G (ˆ p ( t ) , ˆ γ ( t ) ) , will be of the order (cid:112) d F log( n ) /n , while that of G of the order of (cid:112) d P log( n ) /n .
Revised: G (ˆ p ( t ) , ˆ γ ( t ) ) , will be of the order (cid:112) d F log( n ) /n , while that of G of the order of (cid:112) d P log( n ) /n .

Original: Moreover, we can see in the bound of Thm.
Revised: Moreover, we can see in the bound of Thm.

Original: 4 the interplay between bias and variance introduced by γ .
Revised: 3 the interplay between bias and variance introduced by γ .

Original: In particular, the part of the bound that depends on  γ  can be further simpliﬁed as
Revised: In particular, the part of the bound that depends on ˆ γ ( t ) can be further simpliﬁed as:

Original: [Equation]
Revised: [Equation]

Original: where the ﬁrst part corresponds to the bias increase due to γ , while the second part the bias decrease.
Revised: where the ﬁrst part corresponds to the bias increase due to γ , while the second part the bias decrease.

Original: 4. B IASED STOCHASTIC GRADIENT DESCENT ANALYSIS
Revised: 4. B IASED STOCHASTIC GRADIENT DESCENT ANALYSIS

Original: When the set of candidate prediction rules f θ is parameterized by a vector θ ∈ R d , we may alternatively ﬁt θ via stochastic gradient descent (SGD) (Robbins & Monro, 1951; Bottou & Bousquet, 2008) on the γ -corrected objective L D ( f θ , ˆ p,γ ) .
Revised: When the set of candidate prediction rules f θ is parameterized by a vector θ ∈ R d , we may alternatively ﬁt θ via stochastic gradient descent (SGD) (Robbins & Monro, 1951; Bottou & Bousquet, 2008) on the γ -corrected objective L D ( f θ , ˆ p, ˆ γ ) .

Original: With a minibatch size of 1 and a starting point θ 0 , the parameter updates take the form
Revised: With a minibatch size of 1 and a starting point θ 0 , the parameter updates take the form

Original: [Equation]
Revised: [Equation]

Original: Ideally, these updates would converge to a minimizer of the ideal risk L ( θ ; p 0 ) =
Revised: Ideally, these updates would converge to a minimizer of the ideal risk L ( θ ; p 0 ) =

Original: L D ( f θ ,p 0 ) .
Revised: L D ( f θ ,p 0 ) .

Original: Our next result shows that, if the teacher ˆ p is independent of ( W t ) t ∈ [ n ] , then the SGD updates (4) have excess ideal risk governed by a bias term ζ ( γ ) and a variance term σ ( γ ) 2 /n .
Revised: Our next result shows that, if the teacher ˆ p is independent of ( W t ) t ∈ [ n ] , then the SGD updates (8) have excess ideal risk governed by a bias term ζ (ˆ γ ) and a variance term σ (ˆ γ ) 2 /n .

Original: Here, σ 20 ( θ ) represents the baseline stochastic gradient variance that would be incurred if SGD were run directly on the ideal risk L ( θ ; p 0 ) rather than our surrogate risk.
Revised: Here, σ 20 ( θ ) represents the baseline stochastic gradient variance that would be incurred if SGD were run directly on the ideal risk L ( θ ; p 0 ) rather than our surrogate risk.

Original: Our proof in App.
Revised: Our proof in App.

Original: D builds upon the biased SGD bounds of Ajalloeian & Stich (2020).
Revised: F builds upon the biased SGD bounds of Ajalloeian & Stich (2020).

Original: Theorem 5 (Biased SGD analysis) .
Revised: Theorem 4 (Biased SGD analysis) .

Original: Suppose that the loss L ( θ ; p 0 ) is λ -strongly smooth in θ .
Revised: Suppose that the loss L ( θ ; p 0 ) is λ -strongly smooth in θ .

Original: Deﬁne the bias and root-variance parameters
Revised: Deﬁne the bias and root-variance parameters

Original: [Equation]
Revised: [Equation]

Original: ∇ [Equation] (cid:62) ( 
Revised: for σ 20 ( θ ) (cid:44) (cid:80)

Original: [Equation] p
Revised: i ∈

Original: [Equation] [Equation] 2 ,σ ( γ ) (cid:44) sup [Equation]
Revised: [ d ] Var[ ∇ θ i (cid:96) ( W ; f θ ( X ) , p 0 ( X ))] the unbiased SGD variance.

Original: [Equation] )+ (cid:113) E [ (cid:107)∇ [Equation]
Revised: If

Original: f{+F 
+}
Revised: {+F 
+}

Original: [-[Equation] ( X ) (cid:62) γ ( X ) (cid:62) ( Y − p 
 ( X )) (cid:107) 22 ]+ (cid:13)(cid:13) ∇ θ f (cid:62)-]
Revised: [-[Equation] ( X ) (cid:62) γ ( X ) (cid:62) ( Y − p 
 ( X )) (cid:107) 22 ]+ (cid:13)(cid:13) ∇ θ f (cid:62)-]= L

Original: θ ( γ f θ , ˆ p − γ ) (cid:62) (ˆ p − p 0 ) (cid:13)(cid:13) 2 2 ,for σ 20 ( θ ) (cid:44) (cid:80) i ∈ [ d ] Var[ ∇ θ i (cid:96) ( W ; f θ ( X ) , p 0 ( X ))] the unbiased SGD variance.
Revised: [Equation]

Original: If
Revised: θ

Original: [-F 
 =-]
Revised: [-F 
 =-][Equation]

Original: L ( θ 0 ; p 0 ) − min θ ∈ R d L ( θ ; p 0 ) , then the iterates { θ t } nt =1 of the γ -corrected SGD algorithm satisfy
Revised: p 0 ) − [Equation] θ ∈ R d L ( θ ; p 0 ) , then the iterates { θ t } n t =1 of the γ -corrected SGD algorithm satisfy

Original: [Equation]
Revised: [Equation]

Original: If, in addition, L ( θ ; p 0 ) is µ -strongly convex in θ , then the iterates satisfy
Revised: If, in addition, L ( θ ; p 0 ) is µ -strongly convex in θ , then the iterates satisfy

Original: [Equation]
Revised: [Equation]

Original: Similar to Thm. 4, the bound in Thm. 5, portrays the interplay of bias and variance as  γ ranges from 0 to q f θ , ˆ p (which we remind that for any Bregman loss is independent of f θ ).
Revised: Similar to Thm. 3, the bound in Thm. 4, portrays the interplay of bias and variance as ˆ γ ranges from 0 to q f θ , ˆ p (which we remind that for any Bregman loss is independent of f θ ).

Original: In particular, the part of the bound for strongly convex losses that depends on  γ can be further simpliﬁed to:
Revised: In particular, the part of the bound for strongly convex losses that depends on ˆ γ can be further simpliﬁed to:

Original: [Equation]
Revised: [Equation]

Original: This has a very intuitive form: the ﬁrst term is the impact of  γ ( X ) on the variance, which is also related to the square of the noise of y , divided by the standard error scaling.
Revised: This has a very intuitive form: the ﬁrst term is the impact of ˆ γ ( X ) on the variance, which is also related to the square of the noise of y , divided by the standard error scaling.

Original: The second controls how  γ improves the bias introduced by the error in the teacher’s ˆ p .
Revised: The second controls how ˆ γ improves the bias introduced by the error in the teacher’s ˆ p .

Original: Both the variance and the bias are re-weighted based on the gradient of θ , which expresses that some regions of X are more important in identifying the parameters and therefore, when we minimize the average bias-variane trade-off, we should be re-weighting the samples based on how informative they are about the parameters θ .
Revised: Both the variance and the bias are re-weighted based on the gradient of θ , which expresses that some regions of X are more important in identifying the parameters and therefore, when we minimize the average bias-variane trade-off, we should be re-weighting the samples based on how informative they are about the parameters θ .

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We empirically validate that the two components of our solution, cross-ﬁtting and corrected loss function, can improve the student’s performance, in a variety of settings: (i) a synthetic setting that validates the theory in Secs.
Revised: We empirically validate that the two components of our solution, cross-ﬁtting and γ -correction, can address the two failure modes of KD (teacher overﬁtting/underﬁtting), in a variety of settings: (i) a synthetic setting that validates the theory in Secs.

Original: 3 and 4, (ii) tabular data showing that cross-ﬁtting and γ -correction yield improvement, and (iii) image data to show the wide applicability of our methods.
Revised: 3 and 4, (ii) tabular data showing that cross-ﬁtting and γ -correction yield improvement, and (iii) image data to show the wide applicability of our methods.

Original: Throughout, we use the squared error logit loss (1).
Revised: Throughout, we use the squared error logit loss (1).

Original: Additionalexperiment details and results aredescribed in App.
Revised: Additional details and results are in App.

Original: E.
Revised: G.

Original: Objective for picking γ  in correction.
Revised: Selecting the correction matrix ˆ γ Motivated by the analyses in Sec. 4, for each training point ( x,y ) , we will select our correction matrix ˆ γ ( x ) to balance bias and variance by minimizing a pointwise upper bound on the γ -correction error (9) (ideally with a closed-form solution to avoid excessive computational overhead). 2 To eliminate dependence on the unobserved p 0 , we observe that the bias term (cid:107) ( γ f θ , ˆ p ( x ) − ˆ γ ( x )) (cid:62) (ˆ p ( x ) − p 0 ( x ))

Original: Motivated by the analysis in Secs.{+(cid:107) 
+}
Revised: {+(cid:107) 
+}

Original: 4.1 and 4.3, we propose to ﬁnd the correction γ by balancing the bias  and  variance terms. In particular, at each point ( x,y ) , we choose γ ( x )=diag( v ( x )) to minimize 
Revised: = O ( (cid:107) q f θ , ˆ p ( x ) − ˆ γ ( x ) (cid:107) 2op ) up to additive terms independent of ˆ γ . We introduce a tunable hyperparameter α> 0 to trade off between this bias bound and the variance term in (9) and select ˆ γ ( x )= diag( v ( x )) to minimize (cid:107) ˆ γ ( x ) (cid:107) 2op

Original:  [Equation]
Revised: E [ (cid:107) y − ˆ p ( x ) (cid:107) 22 | x ]+ α (cid:107) q f θ [Equation] ˆ [Equation]

Original: where the hyperparameter α controls the exact tradeoff between bias and variance.
Revised: ( x ) − [Equation] x ) (cid:107) 2op [Equation]

Original: Since the conditional expectation involves the unknown quantity p 0 , we approximate E [ (cid:107) y − ˆ p ( x ) (cid:107) 22 | x ] by its sample (cid:107) y − ˆ p ( x ) (cid:107) [-22 . 
-]
Revised: (cid:107) [Equation] ( x ) (cid:107) 22 E [ (cid:107) [-22 . 
-][Equation]

Revised: −

Revised: [Equation]

Original: {+( x ) (cid:107) 
+}
Revised: {+( x ) (cid:107) 
+}

Revised: | x ]+ α (cid:107) 1ˆ p ( x

Original: These two terms are simpliﬁed versions of the terms in Thm. 5.
Revised: [Equation] − [Equation]

Original: Note that this objective is a quadratic function in v ( x ) , and thus has closed-form solution that is fast to compute. In the case of squared error logit-loss, q ˆ p ( x ) = 1 / ˆ p ( x ) and{+(cid:107) 2 2 . 
 This objective is quadratic+} in the case of squared loss q ˆ p ( x ) = 1 .
Revised: ( x ) (cid:107) 22 . Since the conditional expectation involves the unknown quantity p 0 , we approximate E [ (cid:107) y − ˆ p ( x ) (cid:107) 22 | x ] by its sample (cid:107) y − ˆ p ( x ) {+(cid:107) 2 2 . 
 This objective is quadratic+} in v ( x ) and thus has a closed-form solution.

Original: Moreover, once we solve for  γ ( x ) , the student’s γ -corrected objective can be shown to be equivalent to a square loss with labelsequal to log( p ( x ))
Revised: Moreover, once we solve for ˆ γ ( x ) , the student’s γ -corrected objective can be shown to be equivalent to a square loss with labels log( p ( x ))

Original: + γ ( x ) (cid:62) ( y − p ( x )) .
Revised: +ˆ γ ( x ) (cid:62) ( y − p ( x )) .

Original: Validation of theoretical analysis When the Bayes probability p 0 ( x ) is known, we validate that the corrected loss function yields a student model whose accuracy is closer to the performance of the
Revised: Validation of theoretical analysis. When the Bayes probability p 0 ( x ) is known, we validate that the corrected loss function yields a student model whose accuracy is closer to the performance of the

Original: Bayes classiﬁer, compared to a student model trained from scratch.
Revised: Bayes classiﬁer, compared to a student model trained from scratch.

Original: The synthetic data is generated from a mixture of two Gaussians with equal weights, where the Bayes classiﬁer is known.
Revised: The synthetic data is generated from a mixture of two Gaussians with equal weights, where the Bayes classiﬁer is known.

Original: The student model is a linear logistic regression model, which can capture the
Revised: The student model is a linear logistic regression model, which can capture the

Original: Bayes classiﬁer.
Revised: Bayes classiﬁer.

Original: To simulate the teacher probabilities, we take the true probability p 0 ( x ) and move it closer to either 0 or 1, whichever is closer to p 0 ( x ) .
Revised: To simulate the teacher probabilities, we take the true probability p 0 ( x ) and move it closer to either 0 or 1, whichever is closer to p 0 ( x ) .

Original: We compare 4 training objectives for the student model: (i) trained from scratch (ii) trained with vanilla KD (iii) trained with γ -corrected KD loss (iv) trained to match the Bayes probability p 0 ( x ) .
Revised: We compare 4 training objectives for the student model: (i) trained from scratch (ii) trained with vanilla KD (iii) trained with γ -corrected KD loss (iv) trained to match the Bayes probability p 0 ( x ) .

Original: Figure 1 shows that the corrected KD loss brings the student’s accuracy close to the case where it is trained on the true probabilities.
Revised: Figure 1 shows that the corrected KD loss brings the student’s accuracy close to the case where it is trained on the true probabilities.

Original: [Figure]
Revised: [Figure]

Original: Tabular data We validate cross-ﬁtting and corrected KD loss on ﬁve real-world tabular datasets datasets: Adult (Dheeru & Karra Taniskidou, 2017), FICO (FIC), Higgs (Dheeru & Karra Taniskidou, 2017), MAGIC (Dheeru & Karra Taniskidou, 2017), and StumbleUpon (Eve; Liu et al., 2017).
Revised: Tabular data. We validate cross-ﬁtting and corrected KD loss on ﬁve real-world tabular datasets datasets: Adult (Dheeru & Karra Taniskidou, 2017), FICO (FIC), Higgs (Dheeru & Karra Taniskidou, 2017), MAGIC (Dheeru & Karra Taniskidou, 2017), and StumbleUpon (Eve; Liu et al., 2017).

Original: Cross-ﬁtting can improve the student performance by up to 0.04 AUC points when the teacher model likely overﬁts, while γ -corrected KD loss is most helpful when the teacher has high bias.
Revised: Cross-ﬁtting can improve the student performance by up to 4% AUC points when the teacher model likely overﬁts, while γ -corrected KD loss is most helpful when the teacher has high bias.

Original: We validate that cross-ﬁtting can improve KD when the teacher overﬁts, and γ -correction is helpful when the teacher has high bias.
Revised: We validate that cross-ﬁtting can improve KD when the teacher overﬁts, and γ -correction is helpful when the teacher has high bias.

Original: Both teacher and student models are random forest.
Revised: Both teacher and student models are random forest.

Original: In the ﬁrst case (Fig. 2a), we use a large teacher (500 trees) and vary the student’s capacity, while in the second case (Fig. 2b) we vary the teacher’s bias by limiting the teacher’s max tree depth.
Revised: In the ﬁrst case (Fig. 2a), we use a large teacher (500 trees) and vary the student’s capacity, while in the second case (Fig. 2b) we vary the teacher’s bias by limiting the teacher’s max tree depth.

Original: Figure 2 shows the result of the ﬁrst case on the FICO dataset, and the results of the second case on the Adult dataset.
Revised: Figure 2 shows the result of the ﬁrst case on the FICO dataset, and the results of the second case on the Adult dataset.

Original: The results on the rest of the datasets are similar (App. E.2).
Revised: The results on the rest of the datasets are similar (App. G.2).

Original: When the teacher overﬁts, cross-ﬁtting shows signiﬁcant improvement over using the same teacher on the training points.
Revised: When the teacher overﬁts, cross-ﬁtting shows signiﬁcant improvement over using the same teacher on the training points.

Original: The γ correction successfully reduces the teacher’s bias and thus improves the student’s performance.
Revised: The γ correction successfully reduces the teacher’s bias and thus improves the student’s performance.

Original: The effect is most pronounced when the teacher has large bias (low tree depth).
Revised: The effect is most pronounced when the teacher has large bias (low tree depth).

Original: [Figure]
Revised: [Figure]

Original: Image data On the image classiﬁcation dataset CIFAR-10 (Krizhevsky & Hinton, 2009), we show that cross-ﬁtting can improve the student’s performance in the case where the teacher overﬁts.
Revised: Image data. On the image classiﬁcation dataset CIFAR-10 (Krizhevsky & Hinton, 2009), we show that cross-ﬁtting can improve the student’s performance in the case where the teacher overﬁts.

Original: We use ResNet-8 as the student and ResNet-14/20/32/44/56 (He et al., 2016) as the teacher.
Revised: We use ResNet-8 as the student and ResNet-14/20/32/44/56 (He et al., 2016) as the teacher.

Original: It has been observed that larger or deeper teachers might not yield better students, as the teacher might overﬁt to the training set (Cho & Hariharan, 2019; M ¨ uller et al., 2019).
Revised: It has been observed that larger or deeper teachers might not yield better students, as the teacher might overﬁt to the training set (Cho & Hariharan, 2019; M¨uller et al., 2019).

Original: To simulate overﬁtting, we turn off data augmentation (random horizontal ﬂipping and cropping).
Revised: To simulate overﬁtting, we turn off data augmentation (random horizontal ﬂipping and cropping).

Original: We compare students trained with vanilla
Revised: We compare students trained with vanilla

Original: KD objective with and without cross-ﬁtting and γ -correction in Fig.
Revised: KD objective with and without cross-ﬁtting and γ -correction in Fig.

Original: Cross-ﬁtting can reduce the effect of teacher’s overﬁtting, especially for deep teacher models.
Revised: Cross-ﬁtting can reduce the effect of teacher’s overﬁtting, especially for deep teacher models.

Original: This effect is most evident in the test loss, where the student that is trained with vanilla KD has much higher loss than the student trained with KD and cross-ﬁtting.
Revised: This effect is most evident in the test loss, where the student that is trained with vanilla KD has much higher loss than the student trained with KD and cross-ﬁtting.

Original: The γ -correction consistently provides a small performance boost.
Revised: The γ -correction consistently provides a small performance boost.

Original: [Figure]
Revised: [Figure]

Revised: App.

Revised: G.3 shows the effect of varying the hyperparameter α .

Revised: Large α (closer to orthogonal loss) leads to high variance and lower test accuracy.

Revised: Intermediate α improves on the vanilla KD objective ( α =0 ).

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We provide a new analysis of knowledge distillation under the lens of semiparametric inference.
Revised: We provide a new analysis of knowledge distillation under the lens of semiparametric inference.

Original: By viewing the KD process as learning a student model with the true class probabilities as a nuisance function and the teacher’s probabilities as a plug-in estimator, we obtain generalization bounds on the student model.
Revised: By viewing the KD process as learning a student model with the true class probabilities as a nuisance function and the teacher’s probabilities as a plug-in estimator, we obtain generalization bounds on the student model.

Original: This perspective also suggests two techniques to improve KD, namely cross-ﬁtting and γ -correction, which we show leads to better KD performance, both theoretically and empirically.
Revised: This perspective also suggests two techniques to improve KD, namely cross-ﬁtting and γ -correction, which we show leads to better KD performance, both theoretically and empirically.

Original: This also opens the door to combining KD with other techniques and corrections in the semiparametric literature, such as those based on targeted maximum likelihood framework (Van Der Laan & Rubin, 2006).
Revised: This also opens the door to combining KD with other techniques and corrections in the semiparametric literature, such as those based on targeted maximum likelihood framework (Van Der Laan & Rubin, 2006).

Original: We are excited about future work on unraveling many surprising phenomena in KD, such as the success of self-distillation (Furlanello et al., 2018) and noisy student training (Xie et al., 2020). 
Revised: We are excited about future work on unraveling many surprising phenomena in KD, such as the success of self-distillation (Furlanello et al., 2018) and noisy student training (Xie et al., 2020). A A DDITIONAL LITERATURE REVIEW There has been a lot of work on knowledge distillation in the last half of a decade. Gou et al.

Revised: (2020) gives a survey of the various directions that have been studied in this regard.

Original: A P ROOF OF L EMMA 2: A LGORITHM - AGNOSTIC ANALYSIS
Revised: In this appendix, we sketch some directions. In addition to references discussing theoretic aspects discussed in Sec. 1, here we mention some work empirically studying why distillation works. Cho & Hariharan (2019) show that larger teacher models do not necessarily improve the performance of student models as parsimonious student models are not able to mimic the teacher model. They suggest early stopping in training large teacher neural networks as means of regularizing.

Revised: Cheng et al.

Revised: (2020) show that when applied to image data, distillation allows the student neural net to learn multiple visual concepts simultaneously, while in knowledge distillation neural networks learn concepts sequentially.

Revised: Knowledge distillation has been used in various other settings such as adversarial attacks (Papernot et al., 2016b; Ross &

Revised: Doshi-Velez, 2017; Gil et al., 2019; Goldblum et al., 2020), data security (Papernot et al., 2016a; Lopes et al., 2017; Wang et al., 2019), image processing (Li & Hoiem, 2017; Wang et al., 2017; Chen et al., 2018; Li et al., 2017), natural language processing (Nakashole & Flauger, 2017; Mou et al., 2016;

Revised: Hu et al., 2018; Freitag et al., 2017), speech processing (Chebotar & Waters, 2016;

Revised: Lu et al., 2017; Watanabe et al., 2017; Oord et al., 2018;

Revised: Shen et al., 2018).

Revised: B G LOSSARY

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: derivative

Revised: [Table]

Revised: [Table]

Revised: to the second

Revised: [Table]

Revised: derivative

Revised: [Table]

Revised: (cid:96) (

Revised: [Table]

Revised: [Table]

Revised: to

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: Localized Rademacher complexity of function class F

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: C P ROOF OF L EMMA 2: A LGORITHM - AGNOSTIC ANALYSIS

Original: First we deﬁne for any functional L ( f ) the Frechet derivative as:
Revised: First we deﬁne for any functional L ( f ) the Frechet derivative as:

Original: [Equation]
Revised: [Equation]

Original: When L is an operator of the form: E [ g ( f ( X ))] , then: D f L ( f )
Revised: When L is an operator of the form: E [ g ( f ( X ))] , then: D f L ( f )

Original: [ ν ]= E [ ∇ g ( f ( X )) (cid:62) ν ( X )] .
Revised: [ ν ]= E [ ∇ g ( f ( X )) (cid:62) ν ( X )] .

Original: By σ -strong convexity of L D [Equation] 3 we have that[Equation]
Revised: By σ -strong convexity of L D , 4 we have that

Original: Furthermore

Original: [Equation]
Revised: [Equation]

Original: our excess risk assumption and the optimality of f 0 give us σ

Original: [Equation]
Revised: [Equation]

Original: By Taylor’s theorem with integral remainder,
Revised: By Taylor’s theorem with integral remainder

Original: [Equation]
Revised: [Equation]

Original: whenever ∇ φπ (cid:96) is well-deﬁned.
Revised: whenever ∇ φπ (cid:96) is well-deﬁned.

Original: We can now invoke the expansion (5) and Cauchy-Schwarz to obtain the bound
Revised: We can now invoke the expansion (10) and Cauchy-Schwarz to obtain the bound

Original: [Equation]
Revised: [Equation]

Original: Thus combining all the above inequalities: σ
Revised: Thus combining all the above inequalities: σ

Original: [Equation]
Revised: [Equation]

Original: By an AM-GM inequality, for all a,b ≥ 0 : a · b ≤ 12 ( 2 σ a 2 + σ 2 b 2 ) .
Revised: By an AM-GM inequality, for all a,b ≥ 0 : a · b ≤ 12 ( 2 σ a 2 + σ 2 b 2 ) .

Original: Applying this to the product of norms on the RHS and re-arranging yields
Revised: Applying this to the product of norms on the RHS and re-arranging yields

Original: [Equation]
Revised: [Equation]

Original: To get the ﬁnal ineguality [Equation] observe that:
Revised: To get the ﬁnal ineguality, observe that:

Original: [Equation]
Revised: [Equation]

Original: Moreover, by the boundedness of the third derivative, we have:
Revised: Moreover, by the boundedness of the third derivative, we have:

Original: [Equation]
Revised: [Equation]

Original: Combining all the above yields the ﬁnal bound.
Revised: Combining all the above yields the ﬁnal bound.

Original: B P ROOF OF T HEOREM
Revised: D P ROOF OF T HEOREM

Original: Proof
Revised: Proof

Original: First by localized Rademacher analysis (see Lemma 11 of Foster & Syrgkanis (2019)) if δ n upper bounds the critical radius of the function class G and if we denote with (cid:96) ˆ p,f ( w )= (cid:96) ( w ; f ( x ) , ˆ p ( x )) , then w.p. 1 − ζ :(cid:12)(cid:12)(cid:12) L n (
Revised: First by localized Rademacher analysis (see Lemma 11 of Foster & Syrgkanis (2019)) if δ n upper bounds the critical radius of the function class G and if we denote with (cid:96) ˆ p,f ( w )= (cid:96) ( w ; f ( x ) , ˆ p ( x )) , then w.p. 1 − ζ :

Original: [Equation]
Revised: [Equation]

Original: f, [Equation] p [Equation]
Revised: Moreover, if we let µ ( w )=

Original: L n [Equation] [Equation] p 
Revised: sup φ (cid:107)∇ φ (cid:96) ( w ; φ, ˆ p ( x )) (cid:107) 2 , then we have by Cauchy-Scwarz inequality:

Original: [Equation]
Revised: [Equation]

Original: Moreover, if we let µ ( w )= sup φ (cid:107)∇ φ (cid:96) ( w ; φ, ˆ p ( x )) (cid:107) 2 , then we have by Cauchy-Scwarz inequality:
Revised: −

Original: [Equation]
Revised: [Equation]

Original: If we further assume that the function class F satisﬁes an (cid:96) 2 /(cid:96) 4 condition that:
Revised: If we further assume that the function class F satisﬁes an (cid:96) 2 /(cid:96) 4 condition that:

Original: [Equation]
Revised: [Equation]

Original: then:
Revised: then:

Original: [Equation]
Revised: [Equation]

Original: Plugging in the bound above to Lemma 2 (which holds irrespective of data re-use or sample splitting or cross-ﬁtting) and applying again the AM-GM inequality yields:
Revised: Plugging in the bound above to Lemma 2 (which holds irrespective of data re-use or sample splitting or cross-ﬁtting) and applying again the AM-GM inequality yields:

Original: [Equation]
Revised: [Equation]

Original: C P ROOF OF T HM . 4: C ROSS - FITTED ERM ANALYSIS
Revised: E P ROOF OF T HM . 3: C ROSS - FITTED ERM ANALYSIS

Original: Let L n,t denote the empirical loss over the samples in the t -th fold and ˆ p ( t ) , ˆ γ ( t ) the nuisance functions used on the samples in the k -th fold.
Revised: Let L n,t denote the empirical loss over the samples in the t -th fold and ˆ p ( t ) , ˆ γ ( t ) the nuisance functions used on the samples in the k -th fold.

Original: For any t ∈ [ K ] and conditional on ˆ p ( t ) , ˆ γ ( t ) , suppose that δ n upper bounds the critical radius of the function class
Revised: For any t ∈ [ K ] and conditional on ˆ p ( t ) , ˆ γ ( t ) , suppose that δ n upper bounds the critical radius of the function class

Original: G (ˆ p ( t ) , ˆ γ ( t ) ) , then by Lemma 11 of Foster & Syrgkanis (2019), 5 if we denote with (cid:96) t,f ( w )= (cid:96) ˆ γ ( t ) ( w ; f ( x ) , ˆ p ( t ) ( x )) , w.p. 1 − ζ : 
Revised: G (ˆ p ( t ) , ˆ γ ( t ) ) , then by Lemma 11 of Foster & Syrgkanis (2019), 6 if we denote with (cid:96) t,f ( w )= (cid:96) ˆ γ ( t ) ( w ; f ( x ) , ˆ p ( t ) ( x )) , w.p. 1 − ζ : (cid:12)(cid:12)(cid:12) L n,t ( ˆ f, ˆ p ( t ) , ˆ γ ( t ) )

Original: [Equation]
Revised: − L n,t ( f 0 , ˆ p ( t ) , ˆ γ ( t ) ) − ( L D ( ˆ f, ˆ p ( t ) , ˆ γ ( t ) ) − L D ( f 0 , ˆ p ( t ) , ˆ γ ( t ) )) (cid:12)(cid:12)(cid:12) ≤ O (cid:16) δ n/B,ζ (cid:107) (cid:96) t,f − (cid:96) t,f 0 (cid:107) 2 , 2 + δ 2 n/B,ζ (cid:17)

Original: Moreover, we have that by the deﬁnition of cross-ﬁtted ERM:
Revised: Moreover, we have that by the deﬁnition of cross-ﬁtted ERM:

Original: [Equation]
Revised: [Equation]

Original: Thus we have that w.p. 1 − ζB : 
Revised: Thus we have that w.p. 1 − ζB : B

Revised: B (cid:88) t =

Revised: L D ( ˆ f, ˆ p ( t ) , ˆ γ ( t ) )

Revised: − L D

Original: [Equation]
Revised: [Equation]

Revised: f 0 , ˆ p ( t ) , ˆ γ ( t ) )

Revised: [Equation]

Revised: Moreover

Revised: [Equation]

Original: Moreover, if we let µ ( w )=sup φ,t (cid:13)(cid:13) ∇ φ (cid:96) ( w ; φ, ˆ p ( t ) ( x ))
Revised: if we let µ ( w )=sup [Equation] [Equation] ∇ φ (cid:96) ( w ; φ, ˆ p ( t ) ( x ))

Original: (cid:13)(cid:13) 2 , then we have by Cauchy-Schwarz inequality:
Revised: (cid:13)(cid:13) 2 , then we have by Cauchy-Schwarz inequality:

Original: [Equation]
Revised: [Equation]

Original: If we further assume that the function class F satisﬁes an (cid:96) 2 /(cid:96) 4 condition that:
Revised: If we further assume that the function class F satisﬁes an (cid:96) 2 /(cid:96) 4 condition that:

Original: [Equation]
Revised: [Equation]

Original: then w.p. 1 − ζ :
Revised: then w.p. 1 − ζ :

Original: B B (cid:88) t = (cid:15)

Original: ( ˆ f, ˆ p ( t

Original: [Equation]
Revised: [Equation]

Original: Applying Lemma 2 for any ˆ p ( t ) , ˆ γ ( t ) and averaging the ﬁnal inequality we get:
Revised: Applying Lemma 2 for any ˆ p ( t ) , ˆ γ ( t ) and averaging the ﬁnal inequality we get:

Original: [Equation]
Revised: [Equation]

Original: Plugging in the bound above to Lemma 2 and applying the AM-GM inequality and Jensen’s inequality, yields:
Revised: Plugging in the bound above to Lemma 2 and applying the AM-GM inequality and Jensen’s inequality, yields:

Original: [Equation]
Revised: [Equation]

Revised: t = (cid:114) E (cid:104)(cid:13)(cid:13)

Revised: ( Y − ˆ p ( t ) ( X )) (cid:62) ˆ γ ( t ) ( X ) (cid:13)(cid:13) 4 2 (cid:105) (cid:33)(cid:33) +σ O (cid:32)

Revised: B

Revised: B (cid:88) t

Revised: = (cid:107) ( γ f 0 , ˆ p ( t ) − ˆ γ ( t ) ) (cid:62) (ˆ p ( t ) − p 0 )

Revised: (cid:107) 22 ,(cid:33) .

Original: D P ROOF OF T HM . 5: B IASED SGD ANALYSIS
Revised: F P ROOF OF T HM . 4: B IASED SGD ANALYSIS

Original: Below, for any integer s [Equation] we deﬁne the operator norm of any vector v ∈ R s and any tensor T operating on R s as [Equation]
Revised: Below, for any integer s , we deﬁne the operator norm of any vector v ∈ R s and any tensor T operating on R s as (cid:107) v (cid:107) op (cid:44) (cid:107) v (cid:107)and (cid:107) T (cid:107) op

Original: and (cid:107) T (cid:107) [Equation] (cid:44)
Revised: (cid:44) sup v : (cid:107) v (cid:107) 2 =1 (cid:107) T

Original: [Equation]
Revised: [ v ] (cid:107) op .

Original: Recall the deﬁnition
Revised: Recall the deﬁnition

Original: [Equation]
Revised: [Equation]

Original: Observe that since E [ Y | X = x ]= p 0 ( x ) , we can write for any γ :
Revised: Observe that since E [ Y | X = x ]= p 0 ( x ) , we can write for any γ :

Original: [Equation]
Revised: [Equation]

Original: Thus we also have that:
Revised: Thus we also have that:

Original: [Equation]
Revised: [Equation]

Original: Given this observation, we can decompose the gradient that is used in our SGD algorithm into a bias and variance component, when viewed from the perspective of a biased SGD algorithm for the population oracle loss:
Revised: Given this observation, we can decompose the gradient that is used in our SGD algorithm into a bias and variance component, when viewed from the perspective of a biased SGD algorithm for the population oracle loss:

Original: [Equation]
Revised: [Equation]

Original: The following two lemmas bound the gradient bias and noise terms.
Revised: The following two lemmas bound the gradient bias and noise terms.

Original: Lemma 6 (Gradient bias) .
Revised: Lemma 5 (Gradient bias) .

Original: If sup x,φ,π (cid:107) E [ ∇ ππφ (cid:96) ( W ; φ,π ) | X
Revised: If sup x,φ,π (cid:107) E [ ∇ ππφ (cid:96) ( W ; φ,π ) | X

Original: = x ] (cid:107) op ≤ M , then for any parameter vector θ and functions p and γ , we have:[Equation] p
Revised: = x ] (cid:107) op ≤ M , then for any parameter vector θ and functions p and γ , we have:

Original: [Equation]
Revised: [Equation]

Original: Proof
Revised: Proof

Original: By Taylor’s theorem with integral remainder and Lagrange remainder respectively the SGD bias for each parameter i takes the form
Revised: By Taylor’s theorem with integral remainder and Lagrange remainder respectively the SGD bias for each parameter i takes the form

Original: [Equation]
Revised: [Equation]

Original: Furthermore, our operator norm assumption and Cauchy-Schwarz imply
Revised: Furthermore, our operator norm assumption and Cauchy-Schwarz imply

Original: [Equation]
Revised: [Equation]

Original: Thus
Revised: Thus

Original: [Equation]
Revised: [Equation]

Original: by the triangle inequality and Jensen’s inequality we ﬁnd that
Revised: by the triangle inequality and Jensen’s inequality we ﬁnd that

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: Lemma 7 (Gradient Variance) .
Revised: Lemma 6 (Gradient Variance) .

Original: Deﬁne For any parameter θ and functions p and γ ,(cid:113)
Revised: Deﬁne For any parameter θ and functions p and γ ,

Original: [Equation]
Revised: [Equation]

Original: Proof
Revised: Proof

Original: For each i ∈ [ d ] , deﬁne the shorthand 
Revised: For each i ∈ [ d ] , deﬁne the shorthand [Equation] X

Original: [Equation]
Revised: [Equation]

Original: for some convex combination ¯ p ( X ) of p ( X ) and p 0 ( X ) .
Revised: p [Equation] X )) and [Equation]

Original: We begin by bounding the target expectation using Cauchy-Schwarz
Revised: [Table] E

Original: [Equation]
Revised: [Equation]

Original: We next employ the law of total variance to rewrite the variance terms:
Revised: We next employ the law of total variance to rewrite the variance terms:

Original: [Equation]
Revised: [Equation]

Original:  Finally, we control Var[ Z i ] using Cauchy-Schwarz (cid:115)(cid:88)
Revised: [Table] Var [Table] Finally, we control Var[ Z i ] using Cauchy-Schwarz (cid:115)(cid:88)

Original: [Equation]
Revised: [Table]

Original: The two claims of Thm. 5 now follow from Theorems 2 and 3 of Ajalloeian & Stich (2020) respectively, with the parameters σand ζ instantiated with quantities σ 2 ( γ ) and ζ ( γ ) of Lemmas 6 and 7.
Revised: The two claims of Thm. 4 now follow from Theorems 2 and 3 of Ajalloeian & Stich (2020) respectively, with the parameters σand ζ instantiated with quantities σ 2 ( γ ) and ζ ( γ ) of Lemmas 5 and 6.

Original: E E XPERIMENT D ETAILS AND A DDITIONAL R ESULTS E. S YNTHETIC EXPERIMENT
Revised: G E XPERIMENT D ETAILS AND A DDITIONAL R ESULTS G. S YNTHETIC EXPERIMENT

Original: The synthetic data is generated from a distribution P that is a mixture of two Gaussians with equal weights.
Revised: The synthetic data is generated from a distribution P that is a mixture of two Gaussians with equal weights.

Original: with means at ± ( r,r,...,r ) ∈ R 10 , where r ∈ R controls the separation between the two Gaussian distributions.
Revised: with means at ± ( r,r,...,r ) ∈ R 10 , where r ∈ R controls the separation between the two Gaussian distributions.

Original: For each x , the label y is index of the
Revised: For each x , the label y is index of the

Original: Gaussian from which x was drawn.
Revised: Gaussian from which x was drawn.

Original: Given this generative process, the Bayes classiﬁer has the form p 0 ( x )= p ( y =1 | x )= 1 1+ e θTx for some vector θ ∈ R 10 .
Revised: Given this generative process, the Bayes classiﬁer has the form p 0 ( x )= p ( y =1 | x )= 1 1+ e θTx for some vector θ ∈ R 10 .

Original: We set r =0 .
Revised: We set r =0 .

Original: 25 , and train with gradient descent for 40 epochs, with learning rate 0 .
Revised: 25 , and train with gradient descent for 40 epochs, with learning rate 0 .

Original: The λ hyperparameter of the γ -correction was chosen by cross-validation.
Revised: The λ hyperparameter of the γ -correction was chosen by cross-validation.

Original: We repeat the experiments 5 times to measure the mean and standard deviation.
Revised: We repeat the experiments 5 times to measure the mean and standard deviation.

Original: E. T ABULAR DATA
Revised: G. T ABULAR DATA

Original: We use cross-ﬁtting with 10 folds.
Revised: We use cross-ﬁtting with 10 folds.

Original: The λ hyperparameter of the γ -correction was chosen by cross-validation.
Revised: The λ hyperparameter of the γ -correction was chosen by cross-validation.

Original: We repeat the experiments 5 times to measure the mean and standard deviation.
Revised: We repeat the experiments 5 times to measure the mean and standard deviation.

Original: For the overﬁtting experiment, we use a random forest with 500 trees as the teacher, and a random forest with 1-40 trees as the student.
Revised: For the overﬁtting experiment, we use a random forest with 500 trees as the teacher, and a random forest with 1-40 trees as the student.

Original: The student is trained to match the class log probability logˆ p ( x ) of the teacher.
Revised: The student is trained to match the class log probability log(ˆ p ( x )) of the teacher.

Revised: When ˆ p ( x )=0 , we clip its value to 10 − 3 .

Original: We also evaluate the impact of γ correction compared to without correction (i.e. training on just the teacher’s probabilities), especially when the teacher has high bias.
Revised: We also evaluate the impact of γ correction compared to without correction (i.e., training on just the teacher’s probabilities), especially when the teacher has high bias.

Original: We vary the teacher’s bias by limited the tree depth in the teacher (from 1 to 20).
Revised: We vary the teacher’s bias by limited the tree depth in the teacher (from 1 to 20).

Original: Lower depth corresponds to more bias.
Revised: Lower depth corresponds to more bias.

Original: The teacher has 100 trees and the student has 10 trees.
Revised: The teacher has 100 trees and the student has 10 trees.

Original: For all of the dataset, γ correction successfully reduces the teacher’s bias and thus improves the student’s performance.
Revised: For all of the dataset, γ correction successfully reduces the teacher’s bias and thus improves the student’s performance.

Original: The effect is most pronounced when the teacher has large bias (low tree depth).
Revised: The effect is most pronounced when the teacher has large bias (low tree depth).

Original: We show the full results for all 5 of the datasets in Figs.
Revised: We show the full results for all 5 of the datasets in Figs.

Original: 4 and 5.
Revised: 4 and 5.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: E. E XPERIMENT ON IMAGE DATA (CIFAR-10)
Revised: G. E XPERIMENT ON IMAGE DATA (CIFAR-10)

Original: We use SGD with initial learning rate 0.1 and momentum 0.9, batch size 128, to train for 200 epochs.
Revised: We use SGD with initial learning rate 0.1 and momentum 0.9, batch size 128, to train for 200 epochs.

Original: We use the standard learning rate decay schedule, where the learning rate is divided by 5 at epoch 60, 120, and 160. 
Revised: We use the standard learning rate decay schedule, where the learning rate is divided by 5 at epoch 60, 120, and 160. Ablation on the effect of hyperparameter α

Revised: The hyperparameter α controls the tradeoff between bias and variance.

Revised: When α is very small, the objective is close to the vanilla KD objective.

Revised: When α is large, the objective is closer to the Neyman-orthogonal loss.

Revised: In Figure 6, we show the effect of varying α , with ResNet-8 as the student and ResNet-20 as the teacher, on the CIFARdataset.

Revised: Large values of α lead to high variance and thus lower test accuracy.

Revised: Intermediate values of α improves on the vanilla

Revised: KD objective, which corresponds to α =0 .

Original: {+The test accuracy drops sharply beyond some threshold of α as the variance becomes too high (due to the terms q ˆ p ( x )=diag (cid:16) 
+}
Revised: {+The test accuracy drops sharply beyond some threshold of α as the variance becomes too high (due to the terms q ˆ p ( x )=diag (cid:16) 
+}

Revised: ˆ p 1 ( x ) ,..., 1 ˆ p K

Revised: ( x ) (cid:17) ), causing training to become unstable.

Revised: [Figure]