Evaluating model performance under worst-case subpopulations

1 Introduction . The training population typically does not accurately represent what the model will encounter under operation . Model performance has been observed to substantially degrade under distribution shift [ 16 , 28 , 69 , 80 , 53 ] in speech recognition [ 52 ] , automated essay scoring [ 4 ] , and wildlife conservation [ 11 ] . Similar trends persist for state-of-the-art NLP and computer vision models [ 78 , 74 ] , even on new data constructed under a near-identical process [ 57 , 66 ] . Heavily engineered commercial models are no exception [ 19 ] , performing poorly on rare entities in named entity linking and examples that require abstraction and distillation in summarization tasks [ 38 ] . A particularly problematic form of distribution shift comes from embedded power structures in data collection . Data forms the infrastructure on which we build prediction models [ 30 ] , and they inherit socioeconomic and political inequities against marginalized communities . For example , out of 10,000+ cancer clinical trials the National Cancer Institute funds , less than 5 % of participants were non-white [ 21 ] . Typical models replicate and perpetuate such bias , and their performance drops significantly on underrepresented groups . Speech recognition systems work poorly for Blacks [ 52 ] ∗Authors ordered alphabetically . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . and those with minority accents [ 3 ] . More generally , model performance degrades across demographic attributes such as race , gender , or age , in facial recognition , video captioning , language identification , and academic recommender systems [ 41 , 46 , 17 , 72 , 79 , 19 ] . Model training typically relies on varied engineering practices . It is crucial to rigorously certify model robustness prior to deployment for these heuristic approaches to bear fruit and transform consequential applications . Ensuring that models perform uniformly well across subpopulations is simultaneously critical for reliability , fairness , satisfactory user experience , and long-term business goals . While a natural approach is to evaluate performance across a small set of groups , disadvantaged subpopulations are hard to define a priori because of intersectionality . The most adversely affected are often determined by a complex combination of variables such as race , income , and gender [ 19 ] . For example , performance on summarization tasks varies across demographic characteristics and document specific traits such as abstractiveness , distillation , and location and dispersion of information [ 38 ] . Motivated by these challenges , we study the worst-case subpopulation performance across all subpopulations of a given size . This conservative notion of performance evaluates robustness to unanticipated distribution shifts in Z , and automatically accounts for complex intersectionality by virtue of being agnostic to demographic groupings . Formally , let Z be a set of core attributes that we wish to guarantee uniform performance over . It may include protected demographic variables such as race , gender , income , age , or task-specific information such as length of the prompt or metadata on the input ; notably , it can contain any continuous or discrete variables . We let X ∈ X be the input / covariate , and Y ∈ Y be the label . In NLP and vision applications , X is high-dimensional and typically dim ( Z ) dim ( X ) . We use θ ( X ) to denote a fixed prediction model and consider flexible and abstract losses ` ( θ ( x ) ; y ) . Our goal is to ensure that the model θ performs well over all subpopulations defined over Z . We evaluate model losses on a mixture component , which we call a subpopulation . Postulating a lower bound α ∈ ( 0 , 1 ] on the demographic proportion ( mixture weight ) , we consider the set of subpopulations of the data-generating distribution PZ Qα : = { QZ | PZ = aQZ + ( 1− a ) Q′Z for some a ≥ α , and subpopulation Q′Z } . ( 1 ) The demographic proportion ( mixture weight ) a represents how underrepresented the subpopulation is under the data-generating distribution PZ . Before deploying the model θ , we wish to evaluate the worst-case subpopulation performance Wα ( θ ) : = sup QZ∈Qα EZ∼QZ [ E [ ` ( θ ( X ) , Y ) | Z ] ] . ( 2 ) The worst-case subpopulation performance ( 2 ) guarantees uniform performance over subpopulations ( 1 ) and has a clear interpretation that can be communicated to diverse stakeholders . The minority proportion α can often be chosen from first principles , e.g. , we wish to guarantee uniformly good performance over subpopulations comprising at least α = 20 % of the collected data . Alternatively , it is often informative to study the threshold level of α ? when α 7→ Wα ( θ ) crosses the maximum level of acceptable loss . The threshold α ? provides a certificate of robustness on the model θ ( · ) , guaranteeing that all subpopulations large than α ? enjoy good performance . We develop a principled and scalable procedure for estimating the worst-case subpopulation performance ( 2 ) and the certificate of robustness α ? . A key technical challenge is that for each data point , we observe the loss ` ( θ ( X ) ; Y ) but never observe the conditional risk evaluated at the attribute Z µ ( Z ) : = E [ ` ( θ ( X ) ; Y ) | Z ] . ( 3 ) In Section 2 , we propose a two-stage estimation approach where we compute an estimate ĥ1 ( · ) of the conditional risk µ ( · ) . Then , we compute a plug-in estimate of the worst-case subpopulation performance under ĥ1 ( · ) using a dual reformulation of the worst-case problem ( 2 ) . We show several theoretical guarantees for our estimator of the worst-case subpopulation performance ( 2 ) . Our first finite-sample result ( Section 3.1 ) shows convergence at the rate Op ( √ Compn ( H ) /n ) , where Compn denotes a notion of complexity for the model class estimating the conditional risk ( 3 ) . In some applications , it may be natural to define Z using images or natural languages describing the input and use deep networks to predict the conditional risk ( 3 ) . As the complexity term Compn ( H ) becomes prohibitively large in this case [ 10 , 86 ] , our second result ( Section 3.2 ) shows data-dependent dimension-free concentration of our two-stage estimator : our bound only depends on the complexity of the model class H through the out-of-sample error for estimating the conditional risk ( 3 ) . This error can be made small using overparameterized deep networks , allowing us to estimate the conditional risk ( 3 ) using even the largest deep networks and still obtain a theoretically principled upper confidence bound on the worst-case subpopulation performance . Leveraging these guarantees , we develop principled procedures for estimating the certificates of robustness α ? in Section 3.3 . In Section 4 , we demonstrate the effectiveness of our procedure on real data . By evaluating model robustness under subpopulation shifts , our methods allow the selection of robust models before deployment as we illustrate using the recently proposed CLIP model [ 62 ] . Related work . The long line of works on distributionally robust optimization ( DRO ) aims to train models to perform well under distribution shifts . Previous approaches considered finite-dimensional worst-case regions such as constraint sets [ 29 , 39 , 5 ] and those based on notions of distances for probability measures such as f -divergences [ 12 , 13 , 56 , 55 , 60 , 33 , 32 ] , Levy-Prokhorov [ 34 ] , Wasserstein distances [ 35 , 73 , 15 , 37 , 14 , 82 ] , and integral probability metrics based on reproducing kernels [ 77 , 87 ] . The distribution shifts considered in these approaches are often contrived and difficult to interpret and often result in overly conservative models . Furthermore , these approaches do not currently scale to modern large-scale NLP or vision applications . Our work is most closely related to Duchi et al . [ 31 ] , who proposed algorithms for training models with respect to the worst-case subpopulation performance ( 2 ) , which is a more ambitious goal than our narrower viewpoint of evaluating model performance pre-deployment . Their ( full-batch ) training procedure requires solving a convex program with n2 variables per gradient step , which is often prohibitively expensive . Furthermore , training with respect to the worst-case conditional risk E [ ` ( θ ( X ) ; Y ) | Z ] do not scale to deep networks that can overfit to the training data [ 70 ] . By contrast , our evaluation perspective aims to take advantage of the rapid progress in deep learning . We build scalable evaluation methods that apply to arbitrary models , which allows leveraging state-of-the-art engineered approaches for training θ ( · ) . Our narrower focus on evaluation allows us to provide convergence rates that scale advantageously with the dimension of Z , compared to the nonparametric Op ( n −1/d ) rates for training [ 31 ] . Recently , Jeong and Namkoong [ 48 ] studied a similar notion of worst-case subpopulation performance in causal inference . Our notion of worst-case subpopulation performance is also related to the by now vast literature on fairness in ML . We give a necessarily abridged discussion and refer readers to Barocas et al . [ 8 ] and Corbett-Davies and Goel [ 27 ] for a comprehensive treatment . A large body of work studies equalizing a notion of performance over fixed , pre-defined demographic groups for classification tasks [ 24 , 36 , 7 , 43 , 51 , 84 ] . Kearns et al . [ 49 , 50 ] , Hébert-Johnson et al . [ 45 ] consider finite subgroups defined by a structured class of functions over Z , and study methods of equalizing performance across them . By contrast , our approach instantiates Rawls ’ theory of distributive justice [ 64 , 65 ] , where we consider the allocation of the loss ` ( · ; · ) as a resource . Rawls ’ difference principle maximizes the welfare of the worst-off group and provides incentives for groups to maintain the status quo [ 64 ] . Similarly , Hashimoto et al . [ 44 ] studied negative feedback loops generated by user retention—they use a more conservative notion of worst-case loss than ours—as poor performance on a currently underrepresented user group can have long-term consequences . Our diagnostics complement the recent approaches to benchmarking under distribution shifts [ 80 , 66 , 74 , 53 , 71 , 78 , 57 ] as our procedure does not require out-of-distribution data . Since good performance on a particular distribution shift does not necessitate robustness , we evaluate models using the worst-case subpopulation performance ( 2 ) . 2 Approach . We begin by contrasting our approach to standard alternatives that consider pre-defined , fixed demographic groups [ 59 ] . Identifying disadvantaged subgroups a priori is often challenging as they are determined by intersections of multiple demographic variables . To illustrate such complex intersectionality , consider a drug dosage prediction problem for Warfarin [ 26 ] , a common anti-coagulant ( blood thinner ) . Taking the best prediction model for the optimal dosage on this dataset based on genetic , demographic and clinical factors [ 26 ] , we present the squared error on the root dosage . In Figure 1 , when age and race are considered simultaneously instead of separately , subpopulation performance vary significantly across intersectional groups . The worst-case subpopulation performance ( 2 ) automatically accounts for latent intersectionality . It is agnostic to demographic groupings and allows considering infinitely many subpopulations that represent at least α-fraction of the training population P . By allowing the modeler to select arbitrary protected attributes Z , we are able to consider potentially complex subpopulations . For example , Z can even be defined with respect to a natural language description of the input X . The choice of Z—and subsequent worst-case subpopulation performance ( 2 ) of the conditional risk µ ( Z ) = E [ ` ( θ ( X ) ; Y ) | Z ] —interpolates between the most conservative notion of subpopulations ( when Z = ( X , Y ) ) and simple counterparts defined over a single variable . The choice of the subpopulation size α should be informed by domain knowledge—desired robustness of the system—and the dataset size relative to the complexity of Z . Often , proxy groups can be used for selecting α . If we wish to ensure good performance over patients of all races aged 50 years or older , we can choose α to be the proportion of the least represented ( race , age ≥ 50 ) group—this leads to α = 5 % in the Warfarin data . The corresponding worst-case subpopulation performance ( 2 ) guarantees good performance over all groups of similar size . When it is challenging to commit to a specific subpopulation size , it may be natural to postulate a maximum level of acceptable loss ¯̀ . To measure the robustness of a model , we define the smallest subpopulation size α ? ( θ ) for which the worst-case subpopulation performance is acceptable α ? ( θ ) : = inf { α : Wα ( θ ) ≤ ¯̀ } . ( 4 ) This provides a certificate of robustness : if α ? ( θ ) is large , then θ is brittle against even majority subpopulations ; if it is sufficiently small , then θ performs well on underrepresented subpopulations . We now derive estimators for the worst-case subpopulation performance ( 2 ) and the certificate of robustness ( 4 ) , based on i.i.d . observations ( Xi , Yi , Zi ) ni=1 ∼ P . We assume our observations are independent from the data used to train the model θ ( · ) . Dual reformulation The worst-case subpopulation performance ( 2 ) is unwieldy as it involves an infinite dimensional optimization problem over probabilities . Instead , we use its dual reformulation for tractable estimation . We denote [ · ] + = max ( · , 0 ) , and abuse notation by letting Wα ( h ) be the worst-case subpopulation performance for h ( Z ) ( so that Wα ( θ ) = Wα ( µ ) ) . Lemma 1 ( Shapiro et al . [ 75 , Example 6.19 ] and Rockafellar and Uryasev [ 67 ] ) . If E [ h ( Z ) + ] < ∞ , Wα ( h ) : = sup QZ∈Qα EZ∼QZ [ h ( Z ) ] = inf η∈R { 1 α EP [ h ( Z ) − η ] + + η } . ( 5 ) The dual optimum is attained at the ( 1 − α ) -quantile of the h ( Z ) [ 68 , Theorem 10 ] . The dual ( 5 ) hence shows Wα ( θ ) is a tail-average of µ ( Z ) , a popular risk measure known as the conditional value-at-risk ( CVaR ) in portfolio optimization [ 67 ] . Algorithm 1 Two-stage procedure for estimating worst-case subpopulation performance ( 2 ) 1 : INPUT : Subpopulation size α , model classH , samples S1 and S2 2 : On S1 , solve ĥ1 ∈ argminh∈H 1|S1| ∑ i∈S1 ( ` ( θ ( Xi ) ; Yi ) − h ( Zi ) ) 2 . 3 : On S2 , compute the plug-in estimator Ŵα ( ĥ1 ) = infη { 1 α|S2| ∑ i∈S2 [ ĥ1 ( Zi ) − η ] + + η } . Two-stage procedure A key remaining challenge in estimating Wα ( θ ) is that we can only observe losses ` ( θ ( Xi ) ; Yi ) and never observe the conditional risk µ ( · ) ( 3 ) . We propose a two-stage procedure ( Algorithm 1 ) , where we split the sample into two sets S1 and S2 . On the first sample S1 , we fit an estimator ĥ1 ( Z ) of the conditional risk µ ( Z ) , using any model classH ( class of mappings Z → R ) , by solving an empirical approximation to the loss minimization problem minimize h∈H E [ ( ` ( θ ( X ) ; Y ) − h ( Z ) ) 2 ] . ( 6 ) We denote by h ? a minimizer of ( 6 ) ; for a sufficiently rich model class H , the minimizer is given by µ ( Z ) = E [ ` ( θ ( X ) ; Y ) | Z ] . The loss minimization formulation ( 6 ) allows the use of any ML estimator , as well as standard tools for model selection ( e.g . cross validation ) . In the second stage , on S2 we construct a plug-in estimator for the dual form ( 5 ) , under the estimated conditional risk ĥ1 ( · ) . In practice , we switch the roles of S1 and S2 and average the resulting estimates to leverage the entire sample . To estimate the threshold subpopulation size α ? ( θ ) , we simply take the plug-in estimator α̂ : = inf { α : Ŵα ( ĥ1 ) ≤ ¯̀ } . ( 7 ) Since α 7→ Ŵα ( ĥ1 ) is decreasing , the threshold can be efficiently found by a simple bisection search .

Summary: The authors developed the evaluation method for the worst-case loss of a given predictor over all the subpopulations. The method is motivated to evaluate how a given predictor is robust against population change of the random variable Z.  The authors theoretically investigate the estimation error of the proposed method and show two high probability bounds on the estimation error. One of the bound depends on the complexity of the model to estimate the conditional mean of loss; however, we can achieve a convergence rate faster than $1/n^{1/4}$ if the model complexity is low. Another one is the data-dependent bound and is independent of the model complexity; however, the convergence rate is $1/n^{1/4}$. They also provide some showcases to utilize the proposed methods.