Pseudo Knowledge Distillation: Towards Learning Optimal Instance-specific Label Smoothing Regularization

1 INTRODUCTION . Neural networks , predominantly , pre-trained language models have grown massively in scale over the past few years ( Fedus et al. , 2021 ; Brown et al. , 2020 ; Zeng et al. , 2021 ) . Some researchers have gone so far as to term these over-parameterized models as Foundation Models ( Bommasani et al. , 2021 ) . Although larger models trained on more data tend to dominate various benchmarks ( Wang et al. , 2019b ; a ; Deng et al. , 2009 ) , deploying them on edge devices is a challenge . Training smaller networks is an option but empirical evidence suggests that training larger models and then compressing them can help achieve a higher accuracy per unit computation ( Li et al. ) . Knowledge distillation ( KD ) ( Hinton et al. , 2015 ; Buciluǎ et al. , 2006 ) which is a knowledge transfer technique , from a single or an ensemble of neural networks to another , is a widely applied neural model compression technique across domains ranging from computer vision ( He et al. , 2019 ; Xu et al. , 2020 ; Jafari et al. , 2021 ) to natural language processing ( Jiao et al. , 2020 ; Sanh et al. , 2019 ) . This empirical success is attributed to a few factors such as a ) deeper teacher networks learn better representations , b ) the teacher is able to transfer dark knowledge ( Hinton et al. , 2014 ) in its predictions and c ) the regularization effect of the soft labels in the KD loss ( Yuan et al. , 2020 ) . Nonetheless , it is unclear how exactly student networks benefit from these soft labels . Recent work suggests that with KD , the student is unable to match the teacher on the training distribution but improves generalization on the test set ( Stanton et al. , 2021 ) . In trying to understand KD , a series of investigations have looked at regularization and have demonstrated that the success of both KD and label smoothing is due to a similar regularization effect of soft targets ( Yuan et al. , 2020 ; Zhang & Sabuncu , 2020 ) . Label smoothing ( LS ) ( Szegedy et al. , 2016 ) is a popular technique for enhancing the generalization of a wide range of models without incurring additional computational costs . It encourages the model to treat each non-target class as equally likely for classification by using a uniform distribution to smooth the distribution of the hard labels . Although combining the uniform distribution with the original hard label is beneficial for regularisation , LS does not take into account the true relationships between different label categories like the way KD does . Zhang & Sabuncu ( 2020 ) demonstrate the importance of an instance specific LS regularization . They demonstrate better performance compared to LS , but use a trained model to infer prior knowledge on the label space and thereby sacrifice some of the efficiency of LS . In this work , we investigate whether we can bridge the performance gap between LS and KD , while maintaining the efficiency of LS . Motivated by this , we first revisit the conventional LS technique and generalize the uniform LS method to an instance-specific smoothing regularization framework . Within this framework , we demonstrate that both LS and KD can be interpreted as instances of a smoothing distribution . Next , we propose to learn the optimal smoothing function along with the model training , by devising a two-stage optimization problem . We solve the first-stage problem by giving a deterministic and interpretable solution and apply gradient based optimization to solve the second stage . Our contributions can be summarized as follows : • Our method , Pseudo-KD , provides a general framework , that improves upon conventional LS without an additional training cost . • We justify our framework by demonstrating the new formulation unifies the conventional LS and KD objectives , which confirm their relationship from a different perspective . • Based on our framework , we further propose to learn the smoothing distribution along with training by formulating it as a bi-level optimization problem and we derive a deterministic and interpretable solution . • We conducted extensive experiments on image classification ( CIFAR10 and CIFAR100 ) and natural language understanding ( the GLUE benchmark ) and show that our method outperforms label smoothing , and performs on par with KD while being significantly faster in training ( approximately up to 3x for some settings ) . 2 RELATED WORK . We will provide a brief overview of KD and LS and teacher-free methods which reconcile these two . Knowledge Distillation . KD came into prominence after the work of Hinton et al . ( 2015 ) . Given access to a trained teacher model , assume that we want to train a student . Denote by PT ( x ) and pθ ( x ) the teacher ’ s and student ’ s predictions respectively . For a classification problem , the total KD loss is defined as : L = ( 1− α ) H ( q ( x ) , pθ ( x ) ) + αDKL ( PT ( x ) , pθ ( x ) ) ( 1 ) where q ( x ) is the one-hot label , H ( q ( x ) , p ( x ) ) is the cross-entropy loss between the student ’ s output and the labels , DKL is the KL divergence and α is the scaling parameter . Note that we assume a temperature of 1 and omit it without loss of generality . KD training is equivalent to training with a smoothed label P̂ ( x ) such that : P̂T ( x ) = ( 1− α ) q ( x ) + αPT ( x ) ( 2 ) Label Smoothing We observe that the loss in equation 2 is similar in form to the LS loss ( Yuan et al. , 2020 ) . In LS the KL divergence is between the student output and a uniform distribution U ( k ) = 1K where K is the number of classes . Therefore we can define label smoothened label as : P̂U ( x ) = ( 1− α ) q ( x ) + αU ( 3 ) Both LS and KD incorporate training a model with a smoothened label . KD tends to perform better as it leads to a more accurate student ( Müller et al. , 2019 ; Shen et al. , 2021 ) . However , LS is more efficient since it does not need a pre-trained teacher network . We note from equation 3 that a higher α can lead to a smoother label distribution and a lower α to a peakier label . Yuan et al . ( 2020 ) exploited this to propose a better LS loss to bridge the performance gap between KD and LS . Chen et al . ( 2021 ) note that training on high entropy label distributions can also lead to better robustness under adversarial attacks . Self Distillation We can take advantage of distillation without necessarily using a cumbersome teacher . A student can improve generalization by learning for the output of a trained , identically parameterized model Furlanello et al . ( 2018 ) ; Yuan et al . ( 2020 ) . Zhang et al . ( 2018 ) propose to train an ensemble of students without a static teacher , which benefits from learning and teaching among different students collaboratively . Another self-boosting direction is to apply regularization on the intermediate layers instead of predictive distributions . During the training stage , Zhang et al . ( 2019 ) introduce auxiliary classifiers on top of intermediate layers and penalize the difference between the predictions of these classifiers and the penultimate layer . Along this direction , Liu et al . ( 2020 ) further propose a top-down feature fusion process to handle the capacity gap between the intermediate layers with different depths . However , all these methods require extra parameterized modules for training which are abandoned at inference . In contrast , some recent works focus on boosting the performance with explicitly regularized soft targets . Zhang & Sabuncu ( 2020 ) interpret student-teacher training as an amortized maximum a posteriori estimation and derive an equivalence between self-distillation and instance-specific label smoothing . This analysis helped them to devise a regularization scheme , Beta Smoothing . However , they still use an extra model to infer a prior distribution on their smoothing technique during the training . Yun et al . ( 2020 ) introduce an additional regularization to penalize the predictive output between different samples of the same class . Ding et al . ( 2021 ) propose to capture the relation of classes by introducing decoupled labels table which increases O ( K ×K ) space complexity . The concurrent work ( Kim et al. , 2021 ) utilizes the model trained at i-th epoch as the teacher to regularize the training at ( i+1 ) -th epoch along with annealing techniques . The difference emerges as it still requires a separated model as teacher , or saving all soft labels generated at i-th epoch for the next epoch training which introduces O ( N × K ) space complexity . Our work has a different basis : we develop Pseudo-KD motivated by a principled starting point : to generalize the smoothing distribution to a general form with neither time nor space complexity increases during the training and inference . 3 METHODOLOGY . In this section , we derive our Pseudo-KD framework from the two-stage optimization perspective . First , we revisit the conventional uniform LS method and introduce a closed-form solution to our optimal instance-specific label smoothing . Then , we describe a teacher-free KD implementation under this formulation . 3.1 REVISITING LABEL SMOOTHING . Suppose that we have a K-class classification task to be learned by a neural network , Sθ ( · ) . Given a training set of { xi , yi } Ni=1 samples where xi is a data sample and yi is the ground-truth label , the model Sθ ( · ) outputs the probability pθ ( j|xi ) of each class j ∈ { 1 , . . . , K } : pθ ( j|xi ) = exp ( zi , j ) ∑K j′=1 exp ( zi , j′ ) , ( 4 ) where zi , j = [ Sθ ( xi ) ] j is the logit for j-th class of input xi . A general label smoothing can be formally written as : P̂ls ( j|xi ) = { 1− α+ α · Pls ( j|xi ) , j = k , where k is the ground-truth class ; α · Pls ( j|xi ) , j 6= k , ( 5 ) where P̂ls is a smoothed label , Pls is a smoothing distribution , and α is a hyperparameter that determines the amount of smoothing . If α = 0 , we obtain the original one-hot encoded label . For the original label smoothing method , the probability Pls ( ·|x ) is independent on the sample x and is taken to be the uniform distribution Pls ( j ) = U ( j ) , However , the training can benefit from instancespecific regularization . ( Yuan et al. , 2020 ; Zhang & Sabuncu , 2020 ) . In this work , we consider the general form of LS P̂ls ( ·|xi ) which is instance-specific and not necessarily uniform . Let us consider the Cross Entropy ( CE ) loss with the smoothed labels : Lθ ( Pls ) = 1 N N∑ i=1 − K∑ j=1 P̂ls ( j|xi ) log pθ ( j|xi ) = 1 N N∑ i=1 EP̂ls [ − log pθ ( j|xi ) ] . ( 6 ) Note that computing the weighted sum of negative log-likelihood of each probability of label can be viewed as taking expectation of the negative log-likelihood over label space under a certain distribution P̂ls . We modify this loss by adding a Kullback–Leibler ( KL ) divergence term DKL ( Pls ( ·|xi ) ‖U ( · ) ) into Eq . 6 which encourages the sample-wise smoothing distribution Pls ( ·|xi ) to be close to the uniform distribution . Rθ ( Pls ) = 1 N N∑ i=1 [ EP̂ls [ − log pθ ( j|xi ) ] + βDKL ( Pls ( ·|xi ) ‖U ( · ) ) ] , ( 7 ) where β is a hyper-parameter . The instance-specific smoothing Pls ( ·|xi ) can be viewed as the prior knowledge over the label space for a sample xi . This KL divergence term can be understood as a measure of the ‘ smoothness ’ or ‘ overconfidence ’ of each training label . We choose a Uniform distribution to constrain Pls in the KL term because we would like Pls contain more information on non-target classes . It plays a role similar to temperature in KD and controls the sharpness of the distribution Pls . On the one hand side , in case of KD , it regulates the amount of prior knowledge we inject in the soft label . On the other hand side , it reduce the overconfidence of the trained network by increasing the entropy of soft labels ( a phenomenon studied by Zhang & Sabuncu ( 2020 ) ) . Next we show the post-hoc interpretability of our new formulation.The following two Remarks discuss the relationship of our objective with LS and KD . Remark . When Pls ( ·|xi ) is taken to be the uniform distribution U ( · ) for any xi , the objective in Eq . 7 reduces back to the one in Eq . 6 since DKL ( U ( · ) ‖U ( · ) ) = 0 . Remark . When Pls ( ·|xi ) is taken to be PT ( ·|xi ) which is the output probability of a teacher model T ( · ) for any xi , our objective can be converted to the conventional KD objective with temperature τ = 1 and plus a constant α logK Rθ ( PT ) = 1 N N∑ i=1 [ EP̂T [ − log pθ ( j|xi ) ] + αDKL ( PT ( ·|xi ) ‖U ( · ) ) ] , ( 8 ) = 1 N N∑ i=1 [ ( 1− α ) LCE ( yi , pθ ( ·|xi ) ) + αDKL ( PT ( ·|xi ) ‖pθ ( ·|xi ) ) ] + α logK , where P̂T ( j|xi ) = ( 1−α+α ·PT ( k|xi ) ) for j = k the ground truth , and P̂T ( j|xi ) = α ·PT ( j|xi ) for j 6= k. Our objective function Eq . 7 demonstrates that the both objective of KD and LS can be rewritten as an expectation of negative log-likelihood w.r.t . a transformed distribution P̂T or U ( · ) plus a KL term between the smoothing distribution and uniform distribution . This is not only to confirm previous findings from a different perspective but also justify why our framework is reasonable . Our objective gives a unified formulation for instance-specific label smoothing in a general form . It is easy to bridge the objective of label smoothing and knowledge distillation by converting our objective to KD or LS via replacing the Pls ( ·|xi ) in Eq . 7 with PT ( ·|xi ) or U ( · ) , respectively . There is an inherent relation among these methods . KD and LS , in particular , can be interpreted as instances of our method with a fixed smoothing distribution .

Summary: This paper proposes a new formulation of an objective function for label smoothing regularization to unify label smoothing (LS) regularization and knowledge distillation (KD). The objective function learns an instance-specific label smoothing regularization while training a model towards its prediction target. When the instance-specific label smoothing regularization uses a simple uniform distribution for all training samples, the objective function falls back to traditional LS. When the instance-specific label smoothing regularization uses the output distribution of a teacher model, the objective function learns a KD model. The paper formulates the learning of the optimal instance-specific label smoothing regularization as a bi-level optimization problem and derived a closed-form solution for the inner level of the problem. Experimental results on both image classification and natural language understanding tasks show that the proposed objective and learning algorithm are effective and yield results comparable to the baseline LS and KD models.