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: The paper proposes to have a closer look at (logit matching) knowledge distillation, and explores whether the gains are uniform across all portions of the test data or if instead there are gains on some portions of the data but damaged performance on others. This is justified by the potential of KD to amplify biases. The authors quickly show that indeed the gains are not uniform, some classes being affected negatively by KD so that the bad performance of the teacher on some classes is amplified on the student. Then the authors conduct a number of experiments to study the impact of different aspects of training (e.g. architectures involved and aspects of data like class imbalance) over this amplification of class-specific poor performance. They are then able to identify what classes are likely to be negatively affected and follow that with the proposal of two modifications of the basic KD formulation with the aim to improve KD effect on the worse-performing classes. Experiments focus on CIFAR-100 and ImageNet, including long-tail variants.