Even your Teacher Needs Guidance: Ground-Truth Targets Dampen Regularization Imposed by Self-Distillation

1 Introduction . Knowledge distillation , most commonly known from Hinton et al . ( 2015 ) , is a procedure to transfer knowledge from one neural network ( teacher ) to another neural network ( student ) .1 Often the student has fewer parameters than the teacher , and the procedure can be seen as a model compression technique . Originally , the distillation procedure achieves the knowledge transfer by training the student network using the original training targets , denoted as ground-truth targets , as well as a softened distribution of logits from the ( already trained and fixed ) teacher network.2 Since the popularization of knowledge distillation by Hinton et al . ( 2015 ) , the idea of knowledge distillation has been extended to a variety of settings.3 This paper will focus on the special case where the teacher and student are of identical architecture , called self-distillation , and where the aim is to improve predictive performance , rather than compressing the model . The idea of self-distillation is to use outputs from a trained model together with the original targets as new targets for retraining the same model from scratch . We refer to this as one step of self-distillation , and one can iterate this procedure for multiple distillation steps ( see Figure 1 ) . Empirically , it has been shown that this procedure often generalizes better than the model trained merely on the 1Often knowledge distillation is also referred to under the name Teacher-Student learning . 2We will refer to the weighted outputs of the penultimate layer , i.e . pre-activation of the last layer , as logits . 3See Section 2 for a brief overview , or see Wang and Yoon ( 2020 ) for a more exhaustive survey 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . original targets , and achieves higher predictive performance on validation data , despite no additional information being provided during training ( Furlanello et al. , 2018 ; Ahn et al. , 2019 ; Yang et al. , 2018 ) . Modern deep neural networks are often trained in the over-parameterized regime , where the amount of trainable parameters highly exceed the amount of training samples . Under simple first-order methods such as gradient descent , such large networks can fit any target , but in order to generalize well , such overfitting is usually undesirable ( Zhang et al. , 2017 ; Nakkiran et al. , 2020 ) . Thus , some type of regularization is typically imposed during training , in order to avoid overfitting . A common choice is to add an ` 2-regularization4 term to our objective function , which has been shown to perform comparably to early-stopping gradient descent training ( Yao et al. , 2007 ) . However , in the theoretical study of the over-parameterized regime , regularization is often overlooked , but recent results have shown a connection between wide neural networks and kernel ridge regression through the Neural Tangent Kernel ( NTK ) ( Lee et al. , 2019 , 2020 ; Hu et al. , 2019 ) . We briefly elaborate on this connection in Section D , which motivates our problem setup and connection to deep learning in Section 5 . Our Contributions Through a theoretical analysis we show that • the solution at any distillation step can easily be calculated as a function of the initial fit , and infinitely many steps of self-distillation ( with fixed distillation weight ) correspond to solving the usual kernel ridge regression problem with a specific amplified regularization parameter when the distillation weight is non-zero , • for fixed distillation weights , self-distillation amplifies the regularization at each distillation step , and the ground-truth targets dampen the sparsification and regularization of the selfdistilled solutions , ensuring non-zero solutions for any number of distillation steps , • the optimal distillation weight has a closed form solution for kernel ridge regression , and can be estimated efficiently for neural networks compared to a grid search . Proofs of all our results can be found in Supplementary Material A , and code to reproduce our illustrative example in Section 4.5 . Experimental results in Section B can be found at github.com/Kennethborup/self_distillation . 2 Related Work . The idea of knowledge distillation dates back to Bucila et al . ( 2006 ) , and was later brought to the deep learning setting by Ba and Caruana ( 2014 ) and more recently popularized by Hinton et al . ( 2015 ) in the context of compressing neural networks . Since the original formulation , various extensions have been proposed . Some approaches focus on matching the teacher and student models on statistics other than the distribution of the logits , such as intermediate representations ( Romero et al. , 2015 ) , spacial attention maps ( Zagoruyko and Komodakis , 2019 ) , Jacobians ( Srinivas and Fleuret , 2018 ) , Gram matrices ( Yim et al. , 2017 ) , or relational information between teacher outputs ( Park et al. , 2019 ) . Other extensions focus on developing the transfer procedure , such as self-distillation ( Furlanello et al. , 2018 ) , data-free distillation ( Lopes et al. , 2017 ; Nayak et al. , 2019 ; Micaelli and Storkey , 2019 ; Chen et al. , 2019 ; Fang et al. , 2019 ) , data distillation ( Radosavovic et al. , 2018 ) , residual knowledge distillation ( Gao et al. , 2020 ) , online distillation ( Anil et al. , 2018 ) or contrastive distillation ( Ahn et al. , 2019 ; Tian et al. , 2020a ) . 4With slight differences , ` 2 regularization is often referred to as weight decay and ridge regularization in deep learning and statistical learning literature , respectively . See e.g . Loshchilov and Hutter ( 2019 ) for details . The practical benefits of knowledge distillation have been proven countless of times in a variety of settings , but the theoretical justification for knowledge distillation is still highly absent . Hinton et al . ( 2015 ) conjecture that the success of knowledge distillation should be attributed to the transfer of dark knowledge ( e.g . inter-class relationships revealed in the soft labels ) . Müller et al . ( 2019 ) ; Tang et al . ( 2020 ) support this conjecture , and argue that knowledge distillation is similar to performing adaptive label smoothing weighted by the teacher ’ s confidence in the predictions . Dong et al . ( 2019 ) show the importance of early stopping when training over-parameterized neural networks for distillation purposes by arguing that neural networks tend to fit informative and simple patterns faster than noisy signals , and knowledge distillation utilizes these simple patterns for knowledge transfer . Abnar et al . ( 2020 ) empirically investigate how knowledge distillation can transfer inductive biases between student and teacher models , and Gotmare et al . ( 2019 ) empirically shows how the dark knowledge shared by the teacher mainly is disbursed to some of the deepest layers of the teacher . To the best of our knowledge , few papers investigate knowledge distillation from a rigorous theoretical point of view , and those that do , do so with strong assumptions on the setting . Phuong and Lampert ( 2019 ) ignore the ground-truth targets during distillation and furthermore assume linear models . Mobahi et al . ( 2020 ) investigate self-distillation in a Hilbert space setting with kernel ridge regression models where the teacher is trained on the ground-truth targets , and the student ( and subsequent iterations ) is only trained on the predictions from the teacher without access to the ground-truth targets . They show that self-distillation progressively limits the number of basis functions used to represent the solutions , thus eventually causing the solutions to underfit . In this paper , we build on the theoretical results of Mobahi et al . ( 2020 ) , but we include the weighted ground-truth targets in the self-distillation procedure , where we allow the weight to depend on the self-distillation step , and show how this drastically affects the behavior and effect of self-distillation.5 3 Problem Setup . Notation Vectors and matrices are denoted by bold-faced letters ; vectors are column vectors by default , and for a vector a let [ a ] i be the i-th entry and for a matrix A let [ A ] i , j be the ( i , j ) -th entry . Let In denote the identity matrix of dimension n , [ k ] = { 1 , 2 , . . . , k } , and let ‖·‖2 and ‖·‖F denote the ` 2-norm and the Frobenius norm , respectively . Finally , for a function h : Rn → Rd and X ∈ Rm×n , we denote by h ( X ) the Rm×d matrix of outcomes , where the i ’ th row of h ( X ) is the function applied to the i ’ th row of X , i.e . [ h ( X ) ] i , · = h ( xi ) . Consider the training dataset D ⊆ Rd ×R , and let X = { x | ( x , y ) ∈ D } and Y = { y | ( x , y ) ∈ D } denote the inputs and targets , respectively . Let X = [ xi ] i∈ [ n ] ∈ Rn×d be the matrix of inputs , y = [ yi ] i∈ [ n ] the vector of targets , and X̃ ∈ Rm×d , ỹ ∈ Rm be the matrix and vector of validation inputs and targets , respectively . Given a feature map ϕ : Rd → V , where V has dimension D , we denote by K = κ ( X , X ) = [ κ ( xi , xj ) ] ni , j=1 ∈ Rn×n , where κ ( xi , xj ) = 〈ϕ ( xi ) , ϕ ( xj ) 〉 , the symmetric kernel ( Gram ) matrix associated with the feature map ϕ.6 3.1 Self-Distillation of Kernel Ridge Regressions . In order to avoid overfitting our training data , we will impose a regularization term on our weights , and thus investigate the kernel ridge regression functions f ∈ F mapping f : X → Y , to construct a solution which best approximates the true underlying data generating map and generalize well to new unseen data from this underlying map . We consider self-distillation in the kernel ridge regression setup ; i.e . consider the ( self-distillation ) objective function Ldistill ( f ( X , β ) , y1 , y2 ) = α 2 ‖f ( X , β ) − y1‖22 + 1− α 2 ‖f ( X , β ) − y2‖22 + λ 2 ‖β‖22 , ( 1 ) where α ∈ [ 0 , 1 ] , λ > 0 , y1 , y2 ∈ Rn and f ( X , β ) = ϕ ( X ) β . The objective in ( 1 ) is a weighted sum of two Mean Square Error ( MSE ) objective functions with different targets7 and an ` 2-regularization 5In Supplementary Material E we relate our problem setup to Mobahi et al . ( 2020 ) and extend some of our results to a constrained optimization setting with a regularization functional in Hilbert space . 6Since the kernel trick makes the predictions depend only on inner products in the feature space , it is not a restriction if D is infinite . However , for ease of exposition we assume D is finite . 7It is straightforward to verify that minimizing ( 1 ) and the classic MSE objective with a weighted target , i.e . L̃distill ( f ( X , β ) , y1 , y2 ) = 12 ‖f ( X , β ) − ( αy1 + ( 1− α ) y2 ) ‖ 2 2 + λ 2 ‖β‖22 , are equivalent and that the objective functions are equal up to the additive constant α ( α− 1 ) ‖y1 − y2‖22 . on the model weights . Minimization of ( 1 ) w.r.t . β is straightforward and yields the minimizer β̂ def = argmin β Ldistill ( f ( X , β ) , y1 , y2 ) = ϕ ( X ) ᵀ ( K+ λIn ) −1 ( αy1 + ( 1− α ) y2 ) ( 2 ) by Woodbury ’ s matrix identity and definition of K. This solution can also be seen as a direct application of the Representer Theorem ( Schölkopf et al. , 2001 ) . Let y ( 0 ) def= y , i.e . the original targets , and recursively define for the steps τ ≥ 1 , β̂ ( τ ) def = argmin β Ldistill ( f ( X , β ) , y , y ( τ−1 ) ) ( 3 ) = ϕ ( X ) ᵀ ( K+ λIn ) −1 ( α ( τ ) y + ( 1− α ( τ ) ) y ( τ−1 ) ) , f ( x , β̂ ( τ ) ) def = ϕ ( x ) ᵀβ̂ ( τ ) ( 4 ) = κ ( x , X ) ᵀ ( K+ λIn ) −1 ( α ( τ ) y + ( 1− α ( τ ) ) y ( τ−1 ) ) , y ( τ ) def = f ( X , β̂ ( τ ) ) , ( 5 ) for fixed α ( τ ) ∈ [ 0 , 1 ] . Notice , the initial step ( τ = 1 ) corresponds to standard training by definition and as such is independent of α ( 1 ) . Self-distillation treats the weighted average of the predictions , y ( 1 ) , from this initial model on X , and the ground-truth targets , y , as targets . This procedure is repeated as defined in ( 3 ) - ( 5 ) and we obtain the self-distillation procedure as illustrated in Figure 1 . Note , the special cases α ( τ ) = 0 and α ( τ ) = 1 correspond to merely training on the predictions from the previous step , and only training on the original targets , respectively . Thus , α ( τ ) = 1 is usually not of interest , as the solution is equal to a classical kernel ridge regression , and self-distillation plays no role in this scenario . We will often consider the special case of equal weights , α ( 2 ) = · · · = α ( τ ) = α , and if α = 0 this corresponds to the setting investigated in Mobahi et al . ( 2020 ) in a slightly different setup . Thus , some of the following results can be seen as a generalization of Mobahi et al . ( 2020 ) to step-wise and non-zero α .

Summary: This paper studies the theory of self-distillation in the kernel ridge regression setting. The authors extend the theoretical results in Mobahi et al. (2020) to further incorporate the ground-truth labels in the distillation objective and highlight that the ground-truth labels serve to dampen the sparsification and regularization of the self-distilled solutions.