Teacher's pet: understanding and mitigating biases in distillation

Knowledge distillation is widely used as a means of improving the performance of a relatively simple “ student ” model using the predictions from a complex “ teacher ” model . Several works have shown that distillation significantly boosts the student ’ s overall performance ; however , are these gains uniform across all data subgroups ? In this paper , we show that distillation can harm performance on certain subgroups , e.g. , classes with few associated samples , compared to the vanilla student trained using the one-hot labels . We trace this behaviour to errors made by the teacher distribution being transferred to and amplified by the student model . To mitigate this problem , we present techniques which soften the teacher influence for subgroups where it is less reliable . Experiments on several image classification benchmarks show that these modifications of distillation maintain boost in overall accuracy , while additionally ensuring improvement in subgroup performance . 1 INTRODUCTION . Knowledge distillation is a technique for improving the performance of a “ student ” model using the predictions from a “ teacher ” model . At its core , distillation involves replacing the one-hot training labels with the teacher ’ s predicted label distribution . Empirically , distillation has proven successful for model compression ( Bucilǎ et al. , 2006 ; Hinton et al. , 2015 ) , improving the performance of a fixed model architecture ( Anil et al. , 2018 ; Furlanello et al. , 2018 ) , and semi-supervised learning ( Radosavovic et al. , 2018 ) . Theoretically , several works ( Lopez-Paz et al. , 2016 ; Mobahi et al. , 2020 ; Tang et al. , 2020 ; Menon et al. , 2020 ; Zhang & Sabuncu , 2020 ; Ji & Zhu , 2020 ; Allen-Zhu & Li , 2020 ; Zhou et al. , 2021 ; Dao et al. , 2021 ) have studied how distillation affects learning . Put together , both strands of work further the understanding of when and why distillation helps . In this paper , we are similarly motivated to better understand the mechanics of distillation , but pose a slightly different question : does distillation help all data subgroups uniformly ? Or , do its overall gains come at the expense of degradation of performance on certain subgroups ? To our knowledge , there has been no systematic study ( empirical or otherwise ) of this question . This consideration is topical given the study of fairness of machine learning algorithms on under-represented subgroups ( Hardt et al. , 2016 ; Buolamwini & Gebru , 2018 ; Chzhen et al. , 2019 ; Sagawa et al. , 2020a ) . Our first finding is that even in standard settings — e.g. , on image classification benchmarks such as CIFAR — distillation can disproportionately harm performance on subgroups defined by the individual classes ( see Figure 1 ) . To discern the source of this behaviour , we ablate the teacher and student architectures ( §3.2 ) , dataset complexity ( §3.3 ) , label frequencies ( §3.4 ) . These point to the potential harms of distillation when the teacher makes confident mispredictions on a subgroup . We emphasise that this phenomenon is not simply the teacher ’ s worst class accuracy translating to student ; instead , distillation amplifies worst class imbalance much more than the teacher . Having identified a potential limitation of distillation , we present two simple techniques to remedy it . These apply per-subgroup mixing weights between the teacher and one-hot labels , and per-subgroup margins respectively ( §4 ) . Intuitively , these limit the influence of teacher predictions on subgroups it models poorly . Experiments on image classification benchmarks show that these methods typically maintain a boost in overall accuracy , while ensuring a more equitable improvement across subgroups . In sum , this work provides novel insights into distillation performance , with three main contributions : ( i ) we identify a hitherto unexplored issue with distillation , namely , that its improvements in overall accuracy may come at the expense of harming accuracy on certain subgroups ( §3.1 ) . Such a finding is topical given the widespread practical use of generic learning paradigms such as distillation , and the increasing societal applications of learning systems more broadly . ( ii ) we ablate potential sources for the above phenomenon ( §3.2 , §3.3 , §3.4 ) , and in the process identify certain characteristics of data ( e.g. , skewed label distributions ) where it can manifest . This systematic empirical analysis aims at understanding a non-trivial phenomenon , rather than empirically comparing different learning methods . ( iii ) we propose two simple modifications of distillation that mitigate the above problem , based on applying per-subgroup mixing weights and margins ( §4 ) ; these perform well empirically ( §5 ) . While the paper ’ s analysis is empirical rather than theoretical , we note that our focus is on systematic analysis aiming at understanding a non-trivial phenomenon , rather than merely empirical comparison : Section §3 is devoted to carefully understanding the extent of , and causes for , the non-uniform gains of distillation . This is in line with works which empirically analyse neural network phenomena , e.g . Zhang et al . ( 2017 ) ; Müller et al . ( 2019 ) ; Nakkiran et al . ( 2020 ) ; Neyshabur et al . ( 2020 ) . 2 BACKGROUND AND RELATED WORK . Knowledge distillation . Consider a multi-class classification problem over instances X and labels Y = [ L ] . = { 1 , . . . , L } . Given a training set S = { ( xn , yn ) } Nn=1 drawn from some distribution P , we seek a classifier h : X → Y that minimises the misclassification error Eavg ( h ) . = P ( h ( x ) 6= y ) . In practice , one may learn logits f : X→ RL to minimise R̂ ( f ) = 1N ∑N n=1 ` ( yn , f ( xn ) ) , where ` is a loss function such as the softmax cross-entropy , which for softmax probabilities py ( x ) ∝ exp ( fy ( x ) ) is ` ( y , f ( x ) ) .= − log py ( x ) . One may then classify the sample via h ( x ) = arg maxy∈ [ L ] fy ( x ) . Knowledge distillation ( Bucilǎ et al. , 2006 ; Hinton et al. , 2015 ) employs the logits f t : X→ RL of a “ teacher ” model to train a “ student ” model . The latter learns logits f s : X→ RL to minimise : R̂dist ( f ) = 1 N N∑ n=1 [ ( 1− α ) · ` ( yn , f ( xn ) ) + α · ∑ y′∈ [ L ] pty′ ( xn ) · ` ( y′ , f ( xn ) ) ] , ( 1 ) where α ∈ [ 0 , 1 ] . Here , one converts the teacher logits to probabilities pt : X→ ∆L for simplex ∆ , e.g . via a softmax transformation pty′ ( x ) ∝ exp ( f ty′ ( x ) ) . The second term smooths the student labels based on the teacher ’ s confidence that they explain the sample . The first term includes the original label , so as to prevent incorrect teacher predictions from overwhelming the student . One further important trick is temperature scaling of the teacher logits , so that pty′ ( x ) ∝ exp ( T−1 · f ty′ ( x ) ) . Setting T 0 makes pt more uniform , thus preventing overconfident predictions ( Guo et al. , 2017 ) . Average versus subgroup performance . The above exposition treats the misclassification error Eavg ( h ) as the fundamental performance measure of interest . However , suppose the data contains subgroups G = { 1 , . . . , G } . Defining the per-subgroup errors errg ( h ) . = P ( h ( x ) 6= y | g ) , we have Eavg ( h ) = ∑ g∈G P ( g ) · errg ( h ) , which may mask errors on samples with P ( g ) ∼ 0 ( Sagawa et al. , 2020a ; b ; Sohoni et al. , 2020 ) . To this end , one may instead measure the balanced error ( Menon et al. , 2013 ) Ebal ( h ) . = ∑ g∈G 1 |G| · errg ( h ) which treats the subgroup distribution as uniform , or the worst-subgroup error ( Sagawa et al. , 2020a ; b ; Sohoni et al. , 2020 ) Emax ( h ) . = maxg∈G errg ( h ) , which focusses on the worst-performing subgroup . An intermediary is the average of the k worstperforming subgroups ( Williamson & Menon , 2019 ) : for ith largest per-subgroup error err [ i ] ( h ) , Etop−k ( h ) . = 1k ∑k i=1 err [ i ] ( h ) . The definition of subgroups is a domain-specific consideration . One special case is where each label defines a subgroup ( i.e. , G = Y ) , and P ( y ) is skewed . In such long-tail settings ( Buda et al. , 2017 ) , classifiers with good average performance can perform poorly on “ tail ” labels where P ( y ) ∼ 0 . Related work . There is limited prior study that dissects distillation ’ s overall gains per subgroup . Zhao et al . ( 2020 ) showed that in incremental learning settings , distillation can be biased towards recently observed classes . We show that even in offline settings , distillation can harm certain classes . Recently , Zhou et al . ( 2021 ) studied the standard aggregate performance Eavg of distillation , which was tied to a certain subset of “ regularisation samples ” . By contrast , our primary concern is to understand the subgroup performance of distillation . Study of the fairness of machine learning algorithms on under-represented data subgroups has received recent attention ( Dwork et al. , 2012 ; Hardt et al. , 2016 ; Buolamwini & Gebru , 2018 ; Chzhen et al. , 2019 ) . This has prompted dissection of the performance of techniquess such as dimensionality reduction ( Samadi et al. , 2018 ) , increasing model capacity ( Sagawa et al. , 2020a ) , and selective classification ( Jones et al. , 2021 ) . We follow the general spirit of such works , studying a more delicate setting involving two separate models ( the student and teacher ) , each with their own inductive biases . We present more discussion of related directions in §6 . In a related effort , Hooker et al . ( 2019 ) explore how model compression may harm subgroup performance compared to the original model . 3 ARE DISTILLATION ’ S GAINS UNIFORM ? . We demonstrate that the gains of distillation are not uniform across subgroups : specifically , considering subgroups defined by classes , distillation can harm the student ’ s performance on the “ hardest ” few classes ( §3.1 ) . To understand the genesis of this problem , we perform ablations ( cf . Table 1 ) that establish its existence in settings where there are insufficient samples to model certain classes , either due to the number of classes being large ( §3.3 ) , or the class distribution being skewed ( §3.4 ) . We then identify that the student may amplify the teacher ’ s errors ( §3.6 ) . Finally , we corroborate these results for a more general notion of subgroup in a fairness dataset ( §3.5 ) .

Summary: In this work the authors examine the impact of distillation on certain sub-groups of the data. More specifically, they demonstrate that even though distillation can increase the overall accuracy of a model, it can harm the accuracy on certain subclasses. This includes both self-distillation, as well as regular distillation setups. To overcome this, the authors propose using class-specific weights, that have been tuned according to a holdout set, to tune the importance of each class during distillation, as well as using softmax cross-entropy with margins tuned according to the difficulty of each class. Furthermore, it is also demonstrated that these phenomena can also impact the fairness of the model.