Coping with Label Shift via Distributionally Robust Optimisation

1 INTRODUCTION . Classical supervised learning involves learning a model from a training distribution that generalises well on test samples drawn from the same distribution . While the assumption of identical train and test distributions has given rise to useful methods , it is often violated in many practical settings ( Kouw & Loog , 2018 ) . The label shift problem is one such important setting , wherein the training distribution over the labels does not reflect what is observed during testing ( Saerens et al. , 2002 ) . For example , consider the problem of object detection in self-driving cars : a model trained in one city may see a vastly different distribution of pedestrians and cars when deployed in a different city . Such shifts in label distribution can significantly degrade model performance . As a concrete example , consider the performance of a ResNet-50 model on ImageNet . While the overall error rate is ∼ 24 % , Figure 1 reveals that certain classes suffer an error as high as ∼ 80 % . Consequently , a label shift that increases the prevalence of the more erroneous classes in the test set can significantly degrade performance . Most existing work on label shift operates in the setting where one has an unlabelled test sample that can be used to estimate the shifted label probabilities ( du Plessis & Sugiyama , 2014 ; Lipton et al. , 2018 ; Azizzadenesheli et al. , 2019 ) . Subsequently , one can retrain a classifier using these probabilities in place of the training label probabilities . While such techniques have proven effective , it is not always feasible to access an unlabelled set . Further , one may wish to deploy a learned model in multiple test environments , each one of which has its own label distribution . For example , the label distribution for a vehicle detection camera may change continuously while driving across the city . Instead of simply deploying a separate model for each scenario , deploying a single model that is robust to shifts may be more efficient and practical . Hence , we address the following question in this work : can we learn a single classifier that is robust to a family of arbitrary shifts ? We answer the above question by modeling label shift via distributionally robust optimisation ( DRO ) ( Shapiro et al. , 2014 ; Rahimian & Mehrotra , 2019 ) . DRO offers a convenient way of coping with distribution shift , and have lead to successful applications ( e.g . Faury et al . ( 2020 ) ) . Intuitively , by seeking a model that performs well on all label distributions that are “ close ” to the training data label distribution , this task can be cast as a game between the learner and an adversary , with the latter allowed to pick label distributions that maximise the learner ’ s loss . We remark that while adversarial perspectives have informed popular paradigms such as GANs , these pursue fundamentally different objectives from DRO ( see Appendix A for details ) . Although several previous works have explored DRO for tackling the problem of example shift ( e.g. , adversarial examples ) ( Namkoong & Duchi , 2016 ; 2017 ; Duchi & Namkoong , 2018 ) , an application of DRO to the label shift setting poses several challenges : ( a ) updating the adversary ’ s distribution naı̈vely requires solving a nontrivial convex optimisation subproblem with limited tractability , and also needs careful parameter tuning ; and ( b ) naı̈vely estimating gradients under the adversarial distribution on a randomly sampled minibatch can lead to unstable behaviour ( see §3.1 ) . We overcome these challenges by proposing the first algorithm that successfully optimises a DRO objective for label shift on a large scale dataset ( i.e. , ImageNet ) . Our objective encourages robustness to arbitrary label distribution shifts within a KL-divergence ball of the empirical label distribution . Importantly , we show that this choice of robustness set admits an efficient and stable update step . Summary of contributions ( 1 ) We design a gradient descent-proximal mirror ascent algorithm tailored for optimising large-scale problems with minimal computational overhead , and prove its theoretical convergence . ( 2 ) With the proposed algorithm , we implement a practical procedure to successfully optimise the robust objective on ImageNet scale for the label shift application . ( 3 ) We show through experiments on ImageNet and CIFAR-100 that our technique significantly improves over baselines when the label distribution is adversarially varied . 2 BACKGROUND AND PROBLEM FORMULATION . In this section we formalise the label shift problem and motivate its formulation as an adversarial optimisation problem . Consider a multiclass classification problem with distribution ptr over instances X and labels Y = [ L ] . The goal is to learn a classifier hθ : X → Y parameterised by θ ∈ Θ , with the aim of ensuring good predictive performance on future samples drawn from ptr . More formally , the goal is to minimise the objective minθ E ( x , y ) ∼ptr [ ` ( x , y , θ ) ] , where ` : X × Y × Θ → R+ is a loss function . In practice , we only have access to a finite sample S = { ( xi , yi ) } ni=1 ∼ pntr , which motivates us to use the empirical distribution pemp ( x , y ) = 1n ∑n i=1 1 ( x = xi , y = yi ) in place of ptr . Doing so , we arrive at the objective of minimising the empirical risk : min θ Epemp [ ` ( x , y , θ ) ] : = 1 n ∑n i=1 ` ( xi , yi , θ ) . ( 1 ) The assumption underlying the above formulation is that test samples are drawn from the same distribution ptr that is used during training . However , this assumption is violated in many practical settings . The problem of learning from a training distribution ptr , while attempting to perform well on a test distribution pte 6= ptr is referred to as domain adaptation ( Ben-David et al. , 2007 ) . In Label distribution Reference Train distribution Standard ERM Specified a-priori ( e.g. , balanced ) ( Elkan , 2001 ; Xie & Manski , 1989 ; Cao et al. , 2019 ) Estimated test distribution ( du Plessis & Sugiyama , 2014 ; Lipton et al. , 2018 ; Azizzadenesheli et al. , 2019 ; Garg et al. , 2020 ; Combes et al. , 2020 ) Worst-performing class ( Hashimoto et al. , 2018 ; Mohri et al. , 2019 ; Sagawa et al. , 2020 ) Worst k-performing classes ( Fan et al. , 2017 ; Williamson & Menon , 2019 ; Curi et al. , 2019 ; Duchi the special case of label shift , one posits that pte ( x | y ) = ptr ( x | y ) , but the label distribution pte ( y ) 6= ptr ( y ) ( Saerens et al. , 2002 ) ; i.e. , the test distribution satisfies pte ( x , y ) = pte ( y ) ptr ( x | y ) . The label shift problem admits the following three distinct settings ( see Table 1 for a summary ) : ( 1 ) Fixed label shift . Here , one assumes a-priori knowledge of pte ( y ) . One may then adjust the outputs of a probabilistic classifier post-hoc to improve test performance ( Elkan , 2001 ) . Even when the precise distribution is unknown , it is common to posit a uniform pte ( y ) . Minimising the resulting balanced error has been the subject of a large body of work ( He & Garcia , 2009 ) , with recent developments including Cui et al . ( 2019 ) ; Cao et al . ( 2019 ) ; Kang et al . ( 2020 ) ; Guo et al . ( 2020 ) . ( 2 ) Estimated label shift . Here , we assume that pte ( y ) is unknown , but that we have access to an unlabelled test sample . This sample may be used to estimate pte ( y ) , e.g. , via kernel meanmatching ( Zhang et al. , 2013 ) , minimisation of a suitable KL divergence ( du Plessis & Sugiyama , 2014 ) , or using black-box classifier outputs ( Lipton et al. , 2018 ; Azizzadenesheli et al. , 2019 ; Garg et al. , 2020 ) . One may then use these estimates to minimise a suitably re-weighted empirical risk . ( 3 ) Adversarial label shift . Here , we assume that pte ( y ) is unknown , and guard against a suitably defined worst-case choice . Observe that an extreme case of label shift involves placing all probability mass on a single y∗ ∈ Y . This choice can be problematic , as ( 1 ) may be rewritten as min θ ∑ y∈ [ L ] pemp ( y ) · { 1 ny ∑ i : yi=y ` ( xi , yi , θ ) } , where ny is the number of training samples with label y . The empirical risk is thus a weighted average of the per-class losses . Observe that if some y∗ ∈ Y has a large per-class loss , then an adversary could degrade performance by choosing a pte with pte ( y∗ ) being large . One means of guarding against such adversarial label shifts is to minimise the minimax risk ( Alaiz-Rodrı́guez et al. , 2007 ; Davenport et al. , 2010 ; Hashimoto et al. , 2018 ; Mohri et al. , 2019 ; Sagawa et al. , 2020 ) min θ max π∈∆L ∑ y∈ [ L ] π ( y ) · { 1 ny ∑ i : yi=y ` ( xi , yi , θ ) } , ( 2 ) where ∆L denotes the simplex . In ( 2 ) , we combine the per-label risks according to the worst-case label distribution . In practice , focusing on the worst-case label distribution may be overly pessimistic . One may temper this by instead constraining the label distribution . A popular choice is to enforce that ‖π‖∞ ≤ 1k for suitable k , which corresponds to minimising the average of the top-k largest per-class losses for integer k ( Williamson & Menon , 2019 ; Curi et al. , 2019 ; Duchi et al. , 2020 ) . We focus on the adversarial label shift setting , as it meets the desiderata of training a single model that is robust to multiple label distributions , and not requiring access to test samples . Adversarial robustness has been widely studied ( see Appendix A for more related work ) , but its application to label shift is much less explored . Amongst techniques in this area , Mohri et al . ( 2019 ) ; Sagawa et al . ( 2020 ) are most closely related to our work . These works optimise the worst-case loss over subgroups induced by the labels . However , both works consider settings with a relatively small ( ≤ 10 ) number of subgroups ; the resultant algorithms face many challenges when trained with many labels ( see Section 4 ) . We now detail how a suitably constrained DRO formulation , coupled with optimisation choices , can overcome this limitation . Algorithm 1 ADVSHIFT ( θ0 , γc , λ , NNOpt , pemp , ηπ ) 1 : Initialise adversary distribution as π1 = ( 1L , ... , 1 L ) . 2 : for t = 1 , . . . , T do 3 : Sample mini-batch of b examples { ( xi , yi ) } bi=1 . 4 : Evaluate stochastic gradient gθ = 1b ∑b i=1 πt ( yi ) pemp ( yi ) · ∇θ ` ( xi , yi , θt ) 5 : Update neural network parameters θt+1 = NNOpt ( gθ ) 6 : Update Lagrangian variable α = 0 if r > KL ( πt , pemp ) , α = 2γcλ if r < KL ( πt , pemp ) . 7 : Evaluate adversarial gradient gπ ( i ) = 1b ∑b j=1 1 { yj=i } pemp ( i ) · ∇π ` ( xj , yj , θt+1 ) . 8 : Update adversarial distribution πt+1 = ( πt · pαemp ) 1/ ( 1+α ) · exp ( ηπgπ ) /C

Summary: This paper tackles label shift in supervised learning via distributionally robust optimization. The main idea is to train by solving a min-max problem, where the max problem searches for the worst-case label shift in an Kullback-Leibler divergence ambiguity set. The KL ambiguity set will generate some form of adversarial reweighing of the sample training points, which gives us hope that the learned parameters will perform better in the test data (with shift). The paper proposes to solve the Lagrangian version instead of the constrained version of the DRO problem, and proposes a gradient descent-ascent type of algorithm.