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 attacks the issue of mismatch in distribution of labels between train and test samples. The authors propose a DRO-based approach which amounts to solving a modified ERM problem. Compared to classical approaches, the proposed method doesn’t entail fitting many different models: just a single model. The method builds on recent progress on solving nonconvex-concave games for approximate stationary points. The resulting algorithm is a mirror-descent scheme with explicit convergence rates, under Setting : train and test label distribution do not match