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: This paper focuses on evaluating the performance of a supervised learning model $\theta: X \to Y$ over a set of subpopulations $Z \sim Q$. The authors consider a setting with "subpopulation" shift -- i.e., where the training data under/over samples data from a particular group. In this setting, they propose to bound the ``worst-case" group performance defined as: $$W_{\alpha}(\theta) := \max_{Q \in \mathcal{Q}_alpha} \mathbb{E}_Q[l(\theta(X),Y|Z)].$$ The authors present a new technique to estimate this quantity, and characterize its convergence in finite-sample settings. The authors apply their measures to determine the "worst-case" group-specific performance of models in 2 real-world prediction tasks: Warfarin dosage prediction, and land use from satellite images.