Invariant Learning with Partial Group Labels

1 INTRODUCTION . Neural networks being overly biased to certain groups of the data is an increasing concern within the machine learning community [ 1 ] . A primary cause for bias against specific groups is the presence of extraneous attributes in the datasets that wrongly direct the model responses [ 38 ] . Such extraneous attributes are features that need to be controlled for . For example , in computer vision tasks such as image classification or object detection , an extraneous attribute could correspond to the background in an image or a co-occurring object irrelevant to the task , e.g . a person making a speech in a football field could be predicted as playing football [ 6 ] . The presence of such extraneous attributes warrant a model to derive the predictions by making spurious correlations to extraneous features in an image rather than an actual object of interest . An inevitable consequence of such correlations to extraneous attributes is disparities in performance across different groups within the dataset . Specifically , if certain groups form a minority , a model can simply cheat by having a high overall aggregate accuracy but poor minority group accuracy [ 24 ] . Existing works for this problem [ 2 , 30 ] operate in the regime where the number of groups , likely to be adversely impacted through spurious correlations , are known apriori . Further , they assume a complete knowledge of the group membership of individual samples in the training dataset . While these methods have been proven effective , it is not realistic to assume access to the group labels for every sample . Consider the scenario where the minority / majority groups could be defined by demographic information such as gender or race . An individual can simply choose not to reveal this information due to privacy considerations [ 16 ] . Alternatively , in medical image settings a label class could contain unrecognized subgroups that demand significant burden on the data labelling efforts [ 33 ] . An example of such unrecognized subgroups could be a lung cancer detection problem , where the class label could comprise of groups such as solid/subsolid tumors and central/peripheral neoplasms [ 24 ] , with many of these groups naturally forming a minority in the dataset . In this work , we consider a setting where a significant portion of the training data is devoid of group labels . We choose to fill a missing gap in the literature where several works bifurcate into methods that either are fully supervised or fully unsupervised in terms of the groups labels . Knowledge on the number of groups in the data makes it convenient to obtain group labels for a tiny portion of the data or take advantage of an existing labelled samples if available . Hence , we address the following research question : Can we train a model that is invariant to group membership using partially labelled data ? We answer the question using a framework of distributionally robust optimization ( DRO ) [ 23 , 32 ] . DRO allows for a training routine that optimizes for the worst-case training loss over predefined set of groups closely connected to the Rawlsian fairness measure [ 27 ] . When the group membership is fully known , the method simply upweights/downweights average training loss of different groups through the course of training [ 30 ] . The application of DRO to the partial group label setting poses significant challenges : ( 1 ) the lack of group label makes it infeasible to compute the worst-off group loss ; and ( 2 ) optimizing only for the high-loss samples , by considering them as a worst-off group , discards considerable portion of the training data adversely impacting the overall accuracy . In light of these challenges , we make the following contributions in this paper . We construct a constraint set that can with high probability encompass the group labels of the unlabeled portion of the data . We optimize for the a worst-off soft group assignment from the constraint set which upper bounds the DRO objective with true group membership . Approximately , high loss samples are assigned to groups with a higher weight and lower marginal probability improving the worstoff group accuracies . Further , the low loss samples are not discarded thus preserving the aggregate accuracy across groups . We show experiments on two images datasets and one tabular datasets and contrast the performance against several baselines . 2 RELATED WORK . Distributionally robust optimization [ 4 , 11 , 23 ] has been studied as a way to train robust ML models across multiple groups and environments . Group Robust Optimization . When group information is known at train time , Group DRO [ 15 , 22 , 25 , 30 , 40 ] or Invariant Risk Minimization ( IRM ) [ 2 ] could be employed to improve the performance over multiple groups . Specifically , Group DRO proceeds by minimizing the loss of the group with the largest loss , while IRM enforces a shared predictor across multiple environments to be optimal in the form of a Lagrangian multiplier . Robust Optimization without Demographics . As the group information may not be always available reliably , several studies have been focusing on developing methods that remove or reduce their dependence on the group information . [ 14 ] has developed a method based on the distributionally robust optimization that minimizes the loss of the samples with losses larger than a certain threshold . [ 18 ] has proposed to reweight the samples in an adversarial way so that the high loss sample could receive more attention over the course of training . Moreover , [ 20 ] has proposed a simple yet effective two-stage approach called Just-Train-Twice that trains a model by upweighting samples with high losses from the initial ERM model . While above methods are based on optimizing high loss samples based on sample reweighting , GEORGE [ 33 ] and EIIL [ 7 ] are two-stage optimization methods , where in the first stage the group or environment labels are inferred and in the second stage Group DRO or IRM optimization are employed with inferred group labels from the first stage . Our method shares a similar idea to GEORGE or EIIL in that we infer group labels for unlabeled data and apply a ( variant of ) Group DRO for optimization . However , we infer group labels in a way that the objective becomes the lower bound of that of Group DRO with true group labels with high probability . Such theoretical guarantee is important for safety-critical ML applications . 3 METHODOLOGY . We introduce our robust optimization framework , Worst-off DRO , with partial group information . We revisit the GroupDRO in Section 3.1 and detail our method in Section 3.2 and 3.3 . In Section 3.4 , we describe a practical method for optimization . 3.1 PRELIMINARY : GROUPDRO . Let x∈X ⊂Rd be data descriptors , y ∈Y ⊂ { 0 , 1 } be target labels , and g ∈G ⊂ { 1 , ... , M } be group labels.We assume training a neural network parameterized by the weights w that corresponds to a per-sample loss l ( x , y ; w ) . Given data triplets { ( xi , yi , gi ) } Ni=1 , we seek to optimize w for the Rawlsian criterion [ 14 , 27 , 39 ] , which minimizes the loss of the worst-off group , as follows : min w max g∈G E [ l ( x , y ; w ) |g ] . ( 1 ) [ 30 ] proposed a practical algorithm to optimize Equation 1 , called Group DRO . This method optimizes a weighted expected loss across all groups . These weights over the groups , denoted by q , are drawn from a simplex ∆M . The objective function is as described below , LGDRO = min w max q∈∆M M∑ j=1 qj [ ∑N i=1 1 { gi = j } l ( xi , yi , w ) ∑N i=1 1 { gi = j } ] ( 2 ) 3.2 WORST-OFF DRO . In this work , we are interested in training a distributionally robust neural network when group labels are only partially available in the entire dataset . That is , our training dataset constitutes of the fullylabeled dataset { ( xi , yi , g ? i ) } Ki=1 and the task-labeled dataset { ( xi , yi , − ) } Ni=K+1 , where− indicates the missing group labels . As noted in Equation 2 , the Group DRO requires group labels of entire dataset . When some of them are missing , we propose to optimize for the following objective : : LWDRO ( C ) = min w max q∈∆M max { ĝ } ∈C M∑ j=1 qj [ ∑N i=1 1 { ĝi = j } l ( xi , yi , w ) ∑N i=1 1 { ĝi = j } ] ( 3 ) where C is a set of group assignments { ĝi } Ni=1 satisfying ĝi = g ? i , ∀i ≤ K. We call the objective in Equation 3 a Worst-off DRO as it optimizes neural network parameters with respect to the worst-off group assignment in some constraint set C. Observe that the Worst-off DRO objective forms an upper bound to the Group DRO objective when evaluated at the ground-truth group labels . This is rather a straightforward consequence from the fact that the ground-truth group assignment { g ? i } Ni=1 falls within the constraint set C. However , for safety-critical applications , such as learning a fair classifier , it is important that the optimal objective ( i.e. , Group DRO with a ground-truth group assignment ) is bounded by the learning objective used in optimization . This is simply because optimizing the proposed learning objective guarantees that the corresponding lower bound ground-truth Group DRO is also optimized ( assuming same parameters for both the methods ) . Conversely , objectives of methods such as EIIL [ 7 ] or GEORGE [ 33 ] , which optimize the Group DRO or IRM objectives using pseudo group labels , would be difficult to compare with LGDRO as it depends on a heuristic to obtain a single set of pseudo group labels . 3.3 REDUCING CONSTRAINT SET WITH MARGINAL DISTRIBUTION CONSTRAINT . It is clear that the constraint set C plays an important role that connects Worst-off DRO to Group DRO . Specifically , the Worst-off DRO objective can be made a tighter bound to that of Group DRO by further constraining C so long as it contains the ground-truth group assignment { g ? i } Ni=1 . In the subsequent paragraph , we describe how we reduce the constraint set while retaining the ground-truth group assignment using a marginal distribution constraint . These constraints may be given as a side information or could be estimated from the small set of partial group labels under certain conditions . Let Cp , ⊂C is a subset of C whose elements { gi } Ni=1 satisfy the following condition : gi = g ? i , ∀i ≤ K , ( 4 ) | 1 N ∑N i=1 1 { gi = j } − pj | ≤ , ∀j ≤M , ( 5 ) where Equation 4 implies that the true group labels are assigned whenever available , and Equation 5 implies that the data marginal distribution should be close to the marginal distribution p. Then , for any marginal distribution p and > 0 , it is easy to show LWDRO ( Cp , ) ≤LWDRO ( C ) as Cp , ⊂C . Moreover , we will see in Lemma 1 that , with high probability , the constraint set Cp ? , with the true marginal distribution p ? contains the true group assignment { g ? } . Lemma 1 . The constraint set Cp ? , contains the true group labels { g ? i } Ni=1 with high probability : P ( { g ? i } Ni=1 ∈Cp ? , ) ≥ 1− 2e−2N 2 ( 6 ) The proof is in Appendix A.1 . As in Equation 6 , the probability of the constraint set containing the true group labels gets closer to 1 by allowing a larger variance ( ) from the true marginal distribution . For fixed > 0 , the probability gets closer to 1 as we increase the number of unlabeled data ( N ) . Finally , this implies that LWDRO ( Cp ? , ) is an upper bound to that of Group DRO : LGDRO ≤ w.h.p LWDRO ( Cp ? , ) ≤ LWDRO ( C ) In practice , however , the true marginal distribution p ? may not be available . For our setting where group labels are partially available , with an assumption that group labels are missing completely at random ( MCAR ) [ 29 ] , the true marginal distribution could be estimated from the subset of data with group labels . This again allows us to formulate a constraint set that contains the ground-truth group assignment with high probability . To be more specific , let p̄ be the estimate of the marginal distribution from { ( xi , yi , g ? i ) } Ki=1 . Lemma 2 . The constraint set Cp̄ , δ+ contains the true group labels { g ? i } Ni=1 with high probability : P ( { g ? i } Ni=1 ∈Cp̄ , δ+ ) ≥ 1− 2e−2N 2 − 2e−2Kδ 2 ( 7 ) We provide a proof in Appendix A.1 . Here , δ is introduced to take into account the estimation error of the true marginal distribution p ? . When K , the number of labeled data , is large , the bound in Equation 7 is close to 1 . Assuming MCAR , for large K : LGDRO ≤ w.h.p LWDRO ( Cp̄ , δ+ )

Summary: The paper proposes a new variant of DRO, called Worst-Off DRO, to address a ubiquitous real-world setting where group labels are only partially available over the training set. The core idea is to introduce a nested optimization to maximize the worst-off group assignment over the entire training set, given some sensible constraints, hence enabling to alternate the standard DRO optimization. Empirical results are reported to support the proposed methods.