Feature Kernel Distillation

1 INTRODUCTION & BACKGROUND . A prevailing belief in the Deep Learning community is that feature learning , where data-dependent features are acquired during training , is crucial to explaining the empirical success of Neural Networks ( NNs ) ( Fort et al. , 2020 ; Baratin et al. , 2021 ) . A comparison in this regard is often made to kernel methods ( Jacot et al. , 2018 ) , which can be thought of as feature selection methods from a fixed data-independent set of features . This separation has been caricaturised as a distinction between feature learning and kernel learning regimes ( Chizat et al. , 2019 ; Yang & Hu , 2020 ; Woodworth et al. , 2020 ) of NN training . Though less amenable to theoretical analysis compared to kernel regimes , feature learning regimes have the capability to capture more of the complex empirical phenomenon that one can observe in NNs due to parameter-space non-convexity , such as : i ) how ensembling trained NNs differing solely in their independent parameter initialisations can lead to improvements in predictive accuracy & uncertainty ( Lakshminarayanan et al. , 2017 ; Allen-Zhu & Li , 2020 ) , or ii ) the effectiveness of knowledge distillation with both single & ensemble teacher models ( Buciluǎ et al. , 2006 ; Hinton et al. , 2015 ) . This implies that in order to understand ensembling & knowledge distillation ( KD ) in NNs , we need to understand the mechanisms of NN feature learning regimes . Ensembling can be loosely summarised as aggregating predictions from multiple models , & is used widely across machine learning ( ML ) to improve performance ( Dietterich , 2000 ; Breiman , 2001 ) . Conversely , knowledge distillation , the idea of transferring knowledge from a teacher model to a student model , is more elaborate as described below , & has garnered the most attention with NNs . Suppose we have a set of input-target ( x , y ) ∈ Rd×RCS data pairs D̂ = { ( xi , yi ) } Ni=1 sampled i.i.d . ( identically and independently distributed ) from some distribution D , and a student NN architecture fS ( x ) = WS ·hS ( x , θS ) . Here hS ( x , θS ) ∈ Rm×1 is a student-specific feature extractor model ( e.g . MLP , CNN , ResNet , or Transformer ) that takes input x & parameters θS , and WS ∈ RCS×m is a matrix of parameters for a linear output layer . Assume also that we have a loss function : L ( θS , WS ) = 1 N N∑ i=1 L ( fS ( xi ) , yi ) , which we seek to minimise over θS & WS in the hope that fS can generalise to unseen ( x , y ) pairs . L is typically the cross-entropy loss in the classification setting . In vanilla KD ( Hinton et al. , 2015 ) , we distill knowledge from a trained teacher network fT ( x ) =WT ·hT ( x , θT ) ∈ RCT to a student by augmenting the loss with a regularisation term towards the teacher ’ s temperature-scaled logits : L̃ ( θS , WS ) = L ( θS , WS ) + λKD 1 N N∑ i=1 L ( fS ( xi ) τ , fT ( xi ) τ ) , ( 1 ) for temperature τ > 0 & regularisation λKD > 0 hyperparameters . Note , this is only valid if CT =CS . Following Hinton et al . ( 2015 ) , many methods have been proposed using different quirks of NNs to distill knowledge from teacher to student . A relevant line of work involves encouraging the student to match how similar/related the teacher views two inputs x , x′ to be ( Tung & Mori , 2019 ; Park et al. , 2019 ) . These approaches have the benefit of being agnostic to teacher/student architectures & prediction spaces CT & CS , but as of yet remain heuristically motivated . In this work , we explore such approaches under the more general framework of NN feature kernel ( the kernel induced by the inner product of last-layer features h ) learning , allowing us to provide the missing theoretical justification . Moreover , our feature kernel learning perspective enables us to derive implementation considerations in Section 3.1 , which we show outperform these previous works in Section 4 . A remarkable property one observes in NN Knowledge Distillation is that it is possible to distill and improve a single student ’ s generalisation with a more powerful teacher , or an ensemble of teachers . This means that the single student ’ s model has enough flexibility to generalise well ( relative to the teacher ) , so that one must factor in the optimisation process ( such as the parameter initialisation ) in order to explain the mechanisms of ensembling and KD in NNs . Allen-Zhu & Li ( 2020 ) provided the first theoretical exposition of the mechanisms by which vanilla KD and ensembling improve generalisation in NNs . To this end , the authors introduce the idea of multi-view data , which is when a class in a multi-class classification problem has multiple identifying features/attributes . For example , an image of a car can be discerned by i ) wheels , ii ) windows , or iii ) headlights . A key idea from Allen-Zhu & Li ( 2020 ) is that the parameter initialisation of an NN , and its random correlations with certain attributes , will bias the NN to learn only a subset of the entire set of attributes pertaining to a given class . When presented with single-view data lacking the classidentifying attribute that the NN has learnt , the NN will not generalise . For example , an NN that has learnt to classify cars based on if they have headlights or not will not generalise to a side-on image of a car that occludes headlights . Allen-Zhu & Li ( 2020 ) suggest then that ensembling NNs works in part because independent parameter initialisations learn independent sets of attributes , so more data features will be learnt by the ensemble model . This is fundamentally different to ensembling in strongly convex feature selection problems , such as linear or random features ( Rahimi & Recht , 2007 ) models with ` 2 regularisation . In such cases , different initialisations reach the same unique optimum , and additional noise must be added to ensure predictive diversity in the ensemble ( Matthews et al. , 2017 ) . Moreover , Allen-Zhu & Li ( 2020 ) claim that vanilla KD in NNs works because the features learnt by the teacher model ( or models ) are imparted to the student via soft teacher labels that capture ambiguity in a given data input ( such as an image of a car whose headlights look like the eyes of a cat ) . These analyses highlight the fact that without feature learning , it is not possible to fully explain how vanilla KD , variants of KD , or ensembling work in NNs . Our contributions Feature learning can be thought of as when the feature kernel , induced by the inner product of last-layer representations in a NN , changes during training ( Yang & Hu , 2020 ) , and kernel learning in NNs can be thought of as when this kernel is constant . In this work , we take a feature learning perspective of knowledge distillation ( KD ) . We first highlight the importance of the feature kernel by viewing trained NNs as data-dependent kernel machines , & use this to motivate Feature Kernel Distillation ( FKD ) . In FKD , our objective is to ensure that the student ’ s feature kernel is well suited for improved generalisation , using both the teacher ’ s data-dependent feature kernel as well as an understanding of the student NN ’ s optimisation process . In Section 3 , we extend the framework of Allen-Zhu & Li ( 2020 ) to show that FKD offers the same generalisation benefits as found in vanilla KD in a multi-view data setting , and is further amenable to ensemble distillation . We then derive practical considerations from our theoretical insights , highlighting the benefit of motivating FKD through an understanding of an NN ’ s feature learning optimisation process , compared to previous methods which implicitly considered the feature kernel for KD . Finally , in Section 4 , we provide experimental support that our theoretical claims extend to standard image classification settings , by : verifying that FKD is amenable to ensemble distillation ; can transfer knowledge across datasets with different prediction spaces ( unlike vanilla KD ) ; and outperforms vanilla KD & previous feature kernel based distillation methods over a range of architectures on CIFAR-100 . 2 MOTIVATION FOR FEATURE KERNEL DISTILLATION . One obvious limitation of vanilla KD is that we require student fS and teacher fT to share prediction spaces , i.e . CS=CT . In many situations , we may have a teacher network trained on a dataset with a different number of classes than the student ’ s dataset , and it is not clear how one could apply vanilla knowledge distillation . One possibility could be to regularise directly in feature space by comparing hS and hT element-wise , but again this either requires same teacher-student last-layer sizes or additional projection layers . To eschew such unnecessary complications , we take the perspective of NNs as data-dependent kernel machines . Define an NN ’ s feature kernel to be : Definition 1 ( Feature Kernel ) . Suppose we have parameters θ and last-layer NN feature extractor h ( · , θ ) : Rd 7→ Rm . For two inputs xi , xj ∈ Rm , the feature kernel , k , is the kernel induced by the inner product of h ( xi , θ ) and h ( xj , θ ) : k ( xi , xj ) def = 〈h ( xi , θ ) , h ( xj , θ ) 〉 . At initialisation , it is well known that in the infinite NN-width limit , with appropriate scaling , the feature kernel k converges almost surely to a deterministic kernel known as the Neural Network Gaussian Process ( NNGP ) kernel ( Neal , 2012 ; Lee et al. , 2018 ; Matthews et al. , 2018 ; Yang , 2019 ) . Yang & Hu ( 2020 ) show that there is a parameterisation-dependent dichotomy between kernel & feature learning regimes for infinite-width NNs , where the feature kernel k is constant or changes during training , respectively . It has been widely demonstrated that a crucial component of the success of finite-width NNs is their ability to flexibly learn features , and indeed the feature kernel , from data during training ( Fort et al. , 2020 ; Aitchison , 2020 ; Chen et al. , 2020 ; Maddox et al. , 2021 ) . To see the importance of the feature kernel , note that for a fixed θ with many common loss functions L , and some mild assumptions on strong convexity ( which could be enforced e.g . with standard ` 2 regularisation ) , the optimal W is uniquely determined and k determines the entire predictive function f∗ ( · ) . For example , with squared error , L ( f ( x ) , y ) = ‖f ( x ) − y‖22 , and ` 2 regularisation strength λ > 0 , a trained NN is precisely kernel ridge regression with the data-dependent feature kernel k , whose job is to measure how similar different inputs are . Thus , all teacher knowledge is contained in its feature kernel , kT , so the feature kernel can act as our primary distillation target , as depicted in Fig . 1 . We show a corresponding result for cross-entropy loss in Appendix A . Fig . 2 corroborates our claims . For a given ResNet20v1 ( He et al. , 2016 ) ‘ reference ’ model trained on CIFAR10 with cross entropy , we plot test class prediction confusion matrices in log-scale between said model and : i ) a retrained version where we fix all but the last-layer parameters ( hence fully determined by the reference model ’ s feature kernel as per Appendix A ) & ii ) an independent model trained from a different initialisation . We see , as expected , that there is significantly more disagree- ment across test predictions for models with different initialisations than those which share feature kernels . This highlights the fact that : a ) different initialisations bias the same architecture to learn different features , and b ) the feature kernel ( largely ) determines a model ’ s test predictions . Experimental details and a breakdown of the predictive disagreements can be found in Appendix F. Having described the feature kernel as a central object in any NN , we use this to motivate our proposed FKD , where we treat the student ’ s feature kernel , kS , as a key distillation target using the teacher ’ s feature kernel , kT . Encouraging similarity across feature kernels shares useful features that the teacher has learnt with the student , which we theoretically show in Section 3 . We define the FKD student loss function via an additive regularisation term between feature kernels:1 LλKD ( θ , W ) def= L ( θ , W ) + λKD · D ( kθ , kT ) ( 2 ) where λKD > 0 is the regularisation strength , D is some ( pseudo- ) distance over kernels , and the student feature kernel k = kθ is written to make explicit the dependence on student feature extractor parameters θ . We stress that Eq . ( 2 ) does not require matching prediction spaces between teacher and student , allowing us to apply FKD across tasks and datasets . We train θ using a stochastic optimisation algorithm such as stochastic gradient descent . For that reason , we require D to be amenable to stochastic mini-batching , therefore , in this work we choose D ( kθ , kT ) = E x1 , x2 i.i.d.∼D̂ [ ( kθ ( x1 , x2 ) − kT ( x1 , x2 ) ) 2 ] , ( 3 ) where we can approximate the expectation with an average of pairwise differences in a minibatch .

Summary: In this paper the authors consider extending the knowledge distillation (KD) framework to cases where the student and the teacher do not share prediction spaces. To this end they propose Feature Kernel Distillation (FKD), where one views the last layer weights of a neural network as a data dependent feature kernel. Instead of regularising the distance between temperature scaled logits corresponding to each data point of the student and the teacher in vanilla Knowledge Distillation, the authors instead regularise based on the distance between feature kernels of a pair of data points (inner product of weights of the pre-final layer of the NN).   Under a simplified setting, viewing data as being multi-view the authors show a setting where classical knowledge distillation would not do well, but a FKD based knowledge distillation using an ensemble of teachers does well (the teacher kernel between two data points is the average teacher kernel across the ensemble). The authors then propose using a correlation kernel instead of raw kernels and feature regularisation to spread out the kernel values.  The authors proceed to empirically demonstrate  that this approach helps in dataset transfer between datasets of different number of classes. They also show that using an ensemble of teachers enables FKD to improve performance (but KD also improves performance).