Year: 2021

Abstract:
Learning invariant representations is an important requirement when training machine learning models that are driven by spurious correlations in the datasets. These spurious correlations, between input samples and the target labels, wrongly direct the neural network predictions resulting in poor performance on certain groups, especially the minority groups. Robust training against these spurious correlations requires the knowledge of group membership for every sample. Such a requirement is impractical in situations where the data labelling efforts for minority or rare groups is significantly laborious or where the individuals comprising the dataset choose to conceal sensitive information. On the other hand, the presence of such data collection efforts result in datasets that contain partially labelled group information. Recent works have tackled the fully unsupervised scenario where no labels for groups are available. Thus, we aim to fill the missing gap in the literature by tackling a more realistic setting that can leverage partially available sensitive or group information during training. First, we construct a constraint set and derive a high probability bound for the group assignment to belong to the set. Second, we propose an algorithm that optimizes for the worst-off group assignments from the constraint set. Through experiments on image and tabular datasets, we show improvements in the minority group’s performance while preserving overall aggregate accuracy across groups.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 METHODOLOGY
text: 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.

heading: 3.1 PRELIMINARY: GROUPDRO
text: 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)

heading: 3.2 WORST-OFF DRO
text: 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.

heading: 3.3 REDUCING CONSTRAINT SET WITH MARGINAL DISTRIBUTION CONSTRAINT
text: 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̄,δ+ )

heading: 3.4 A PRACTICAL OPTIMIZATION ALGORITHM
text: We are interested in solving the optimization problem LWDRO(Cp, ). Unfortunately, the inner maximization problem with respect to the group assignments {ĝ} in Equation 3 is challenging as variables are discrete and the objective cannot be decomposed due to the marginal distribution constraint. In this section, we describe an optimization recipe with a few approximations.
First, we propose to use a soft group assignments. This not only converts the problem into continuous optimization problem, but also accommodates inherent uncertainties in group assignment
Algorithm 1 Worst-off DRO Algorithm
1: Input: Fully-labelled dataset {(xi, yi, g?i )}Ki=1, task-labeled dataset {(xi, yi,−)}Ni=K+1 2: Initialization: learning rates ηw and ηq , Marginal distribution p̄, 3: Parameters: Group Weights qj , Worst-off DRO group assignments ĝ,
Neural network parameter w 4: for t = 0, 1, 2, ..., T do 5: {ĝt} ← max{ĝ}∈Cp̄, ∑M j=1 q t j ∑ i ĝij l(xi,yi;w t)∑ i ĝij
where, Cp̄, as defined in Equation 9. 6: Gradient descent on w:
wt+1 ← wt − ηw∇w ∑M j=1 q t j ∑ i ĝ t ij l(xi,yi;w)∑ i ĝ t ij
7: Exponential ascent on q: qt+1 ← qt exp (ηq∇q ∑M j=1 qj ∑ i ĝ t ij l(xi,yi;w t+1)∑ i ĝ t ij ) 8: end for 9: Output: Trained neural network parameters wT
for data with missing group labels. Specifically, for each data, we retain a soft group assignment ĝi ∈∆M , and optimize the Worst-off DRO objective over the constraint set Cp̄, as defined below:
min w max q∈∆M max {ĝ}∈Cp, M∑ j=1 qj [∑N i=1 ĝij l(xi, yi, w)∑N i=1 ĝij ] (8)
where the constraint set is defined as:
Cp̄, = { {ĝi}Ni=1 ∣∣∣ ĝi ∈∆M ,∀i≤N,ĝi(g?i ) = 1,∀i≤K, | 1N ∑N i=1 ĝij − p̄j | ≤ ,∀j≤M } (9)
The first condition ensures assignments in the probability simplex, second one ensures assignments are consistent with ground-truth for labeled data, and the third one ensures the data marginal distribution follows the provided distribution. This resulting optimization problem is solved using offthe-shelf CVXPY solver [8]. More details on this and an example are provided in the Appendix A.2.
Finally, we alternate optimization over w, q and {ĝ} as in Algorithm 1. That is, we first solve inner maximization over {ĝ}, and conduct gradient descent on w and the exponential gradient ascent on q, and iterate. An exponential ascent on q achieves smaller losses for linear predictors (like q) [17].

heading: 4 EXPERIMENTS
text: We test the efficacy of our method on image and tabular datasets, each of which consists of samples from mutually exclusive groups or environments. These groups are indicative of the background or an RGB identification for image datasets, and attributes such as gender or race for tabular datasets. As discussed in Section 1, one or more of the available groups form a minority in terms of the sample size and demographics. The presence of minority groups results in a possible scenario where the aggregate performance is (falsely) remarkable, because evaluations are dominated by larger groups, even though the performance on the minority groups is poor. In our experiments we assume group numbers to be known but group labels are missing completely at random at a fixed rate at train time.
In Section 4.1, we outline our baselines for comparison, our metrics of evaluation and the model selection strategy. In Section 4.2, we describe each dataset in detail and highlight the differences across the groups within the dataset. All the quantitative results are available in Table 2 and per-group summary statistics are present in Table 1. More analysis of our method is provided in Section 4.3.

heading: 4.1 EXPERIMENTAL SETTINGS
text: Baselines. We contrast the performance of our method with respect to a few well-known baselines.
1. ERM: Empirical Risk Minimization that optimizes aggregate average loss over all the samples in the training dataset.
2. Unsup DRO [14]: Samples with losses exceeding a threshold η are considered as a group whose average loss is optimized. Since the method doesn’t require group labels, it is an unsupervised algorithm. 3. Group DRO [30]: A method that optimizes the Rawlsian criterion by assigning simplex weights to the groups. The group labels for individual samples are assumed to be available, hence this method is fully-supervised in terms of the group label. 4. Group DRO (Partial): We consider another variant of Group DRO that only uses samples with group labels at train time. We call the method Group DRO (Partial), to contrast with the above baseline, Group DRO (Oracle).
We compare above methods with our proposal, Worst-off DRO. Note that our approach requires marginal probabilities as an input to the algorithm, which are computed from the training dataset in our experiments. For baselines and Worst-off DRO implementations, samples are drawn randomly for every batch ensuring an unbiased comparison to the ERM baseline. This is unlike [30] who adopt a weighted sampling procedure which could be noisy when group labels are uncertain.
Evaluation Metrics. We set aside a test set whose group labels are fully available. Since all of our datasets characterize a classification task, we evaluate overall accuracies and per-group accuracies in our experiments. Specifically, we highlight the accuracy of the minority group (min) together with the overall (avg) accuracy where individual samples are equally weighted regardless of their group.
Model Selection. Model selection plays a crucial role when distributional differences are observed in a dataset [13]. In our problem setting, individual groups may differ from each other in the joint distribution over the data and the label space, however, the testing set resembles the training set. That is, there are no out of distribution samples and the focus is to improve robustness over a predefined set of groups common to both training and testing datasets. Consequently, among the recommendations made in [13], a training domain validation set is a feasible strategy for our problem. In our algorithm, learning rate and weight decay are important hyper-parameters. Prior works [30, 31] noted that ERM fails to optimize for the minority group’s performance under high regularization regime thus necessitating an alternative. Hence, we compare our methods in this regime. We also tune for algorithmic specific hyper-parameters for each baseline. These hyper-parameters are the loss threshold in Unsup DRO (ηUDRO) and the step size for the group weights in Group DRO (ηGDRO) and in Worst-off DRO (ηWDRO). A list of all hyper-parameter choices used in the experiments is provided in the Appendix A.3. All the numbers reported in the paper were averaged over three random seeds.
We adopt NVP (novel validation procedure) [9] in our experiments. In this procedure, we first search for hyper-parameters with the best overall accuracy. Then, from the top five best performing hyper-parameters, we select the model that achieves the highest minority group accuracy. Such a procedure offers robustness to hyper-parameters in the reported numbers.

heading: 4.2 QUANTITATIVE RESULTS
text: We describe key results on four datasets, Waterbirds [30], CMNIST [2], Adult [10], and CelebA [21].

heading: 4.2.1 WATERBIRDS
text: The dateset, used in [30], comprises of 4795 images of birds from the CUB dataset [37] and the backgrounds taken from the Places dataset [41]. Each image in the dataset has a background of land or water. The target labels are either “landbirds” or “waterbirds”. The authors in [30] create four groups with each target label and a background class considered as a group. In this dataset the groups “landbirds” on water and “waterbirds” on land form a minority. Our results in Table 2 firstly shows that the ERM method attains a small minority group accuracy of 60%. All the invariant learning baselines, except for Group DRO (Partial), improve the minority group’s accuracy. Next, we observe that in comparison to Group DRO (Partial), our proposed Worst-off DRO improves the
minority group’s performance by a significant margin of 21%. Due to this improvement, the allgroup accuracy also improves by 8%. Minority group’s performance on fully-supervised method Group DRO (Oracle) is at 83% accuracy with a window of 18% difference from Worst-off DRO.

heading: 4.2.2 COLORED MNIST (CMNIST)
text: CMNIST, derived from an MNIST [19], is a digit recognition dataset where each image is colored either red or green. Digits < 5/ ≥ 5 are considered as label 0/1. We consider three groups in our experiments. In the first two groups, label 0 images are predominantly colored red and vice versa. In the third group, which forms a minority, we switch coloring such that the label 1 images are predominantly colored red. Specifically, for the first two groups, the color id is sampled by flipping the target label with probabilities 0.2 and 0.1 respectively, while the third group with probability 0.9. Both training and testing sets contain three groups. The overall setup for generating a given group is similar to [2]. We show the results on CMNIST in Table 2. Similar to the Waterbirds dataset, Worstoff DRO improves the minority group’s accuracy compared to the ERM method. Relative to Group DRO (Partial), Worst-off DRO improves the accuracy of the minority group by 4% and 1% in the overall accuracy. The margin between Worst-off DRO and Group DRO (Oracle) is 11%. Among all the baseline, Unsup DRO attains lowest minority group accuracy of 10%. A large trade-off between the minority group accuracy and the all-group accuracy was seen for Unsup DRO in this dataset.

heading: 4.2.3 ADULT DATASET
text: We use a semi-synthetic version of the Adult dataset [10] for this experiment. Similar to [18], we consider race and sex as the four demographic groups. The target label is income > 50K$ and is treated as label 1. Similar to the CMNIST dataset, each group has a different correlation strength to the target label. For the purposes of the experiment, we exaggerate these spurious correlations caused by group membership close to [7]. Particularly, for samples with group label as AficanAmerican, we undersample examples with probability P (y = 1 | group) = 0.06 whereas for the non African-American group labels, we oversample examples with probability P (y = 1 | group) = 0.94. Table 2 indicates a 5% improvement in the minority group’s accuracy while maintaining the similar overall accuracy of 90% compared to Group DRO (Partial). The Group DRO (Oracle) method reaches an accuracy of 82% for the minority group compared to Worst-off DRO which achieves 71%. Evidently, ERM underperforms in terms of the minority group’s accuracy and attains about 68% accuracy.

heading: 4.2.4 CELEBA DATASET
text: CelebA [21] is a dataset containing about 200k celebrity faces curated from the internet. There are 40 labels available in this dataset which are annotated by a group of paid adult participants [5]. Similar to [30], we aim to predict the target attribute Blond Hair that is spuriously correlated to the Gender attribute. Specifically, having blond hair correlates with the female attribute. The minority group in this dataset are the images with attributes (blond, male). The proportion of samples in the minority and the majority group is show in Table 1. The quantiative results in Table 2 indicate an improvement of 9% over the Group DRO (Partial) method for the proposed Worst-off DRO algorithm. The Group DRO (Oracle) method achieves the highest minority group accuracy of 80%. The minority group performance for the ERM method, with an accuracy of 45%, is comparable to Group DRO (Partial). All the methods are similar in terms of the average group accuracy with values > 90%.

heading: 4.3 ABLATION STUDIES
text: In this section, we discuss different components of our algorithm that influences it’s performance.

heading: 4.3.1 INCREASING THE LABELLED SAMPLES.
text: Recall that for the quantitative results in Table 2, the number of labelled samples were around 10% of the total training samples. In this section, we investigate the effects of increasing the number of labelled samples provided to the training algorithm. Although obtaining annotations for groups is an arduous task [18], having more labelled groups provides two benefits for the algorithm. Firstly, the standard deviation of errors in estimating the marginal probabilities from the labelled portion of the data reduces [36] ( ≈ √ # samples rate). Secondly, labelled groups reinforce an accurate evaluation of the Rawlsian objective in Equation 2 and appropriate weight updates for the groups. The results shown in Figure 2 depict the minority group accuracy at different labelled percent thresholds. The corresponding plots for average group accuracies are provided in the Appendix Figure 5. Worst-off DRO method is compared with Group DRO (Partial). It is observed that for both the methods, the minority group’s accuracy increases with more labelled data. Furthermore, the accuracy values for Worst-off DRO method are better than Group DRO at several thresholds. The methods converge at a threshold specific to the datasets. Increasing the labelled counts beyond such a threshold saturates the Worst-off DRO performance, however, Group DRO (Partial) consistently improves until Oracle performance is attained.

heading: 4.3.2 MINORITY GROUP VS. OVERALL ACCURACY.
text: As discussed in Section 1, several groups in the training dataset, especially the minority groups, could be distributionally different from the majority group samples. Consequently, a mild tradeoff surfaces between the minority group accuracy values and the aggregate accuracies. Addressing this issue, recall that we leverage a robust model selection criterion such as NVP (see Section 4, model selection paragraph) that balances both the minority and aggregate group accuracies. We extend the results in this section, by plotting evaluations at different hyper-parameter choices for our algorithm. Figure 3 contrasts minority group and overall accuracy across all the datasets. Evidently, the topright corners are desirable regions for the models to be present with maximum performance across both the metrics. The Adult dataset in Figure 3 shows a clear envelope on the Worst-off DRO models that surpass the corresponding Group DRO models. Similar trend exists on the remaining datasets with more Worst-off DRO models concentrated in the top-right corner.

heading: 4.3.3 PROGRESSION OF GROUP WEIGHTS
text: Our algorithm 1, proceeds by assigning weights, the q−value’s, to every group. These q−values are updated through exponential ascent similar to [30]. Noticeably, the updates on q depend on the worstoff group assignments determined from the constraint set Cp̄,δ+ . In this section, we investigate on how these group weights evolve through the several iteration of the proposed algorithm. Figure 4 plots this evolution across different datasets with the number of curves in a given plot matching the group count of that dataset. Since the q− weights are initialized uniformly at random, the curves begin at the same value. As the training progresses, it is observed that the weights on the minority
groups gradually increase and those on the majority groups reduce. The plots indicate a high q− value on the minority groups towards the end of training. This is desirable because the empirical risk on the minority groups getting upweighted relative to the majority groups.

heading: 5 CONCLUSION
text: We present Worst-off DRO, an invariant learning method across groups when partial group labels are available. The formulation of Worst-off DRO extends that of Group DRO by optimizing the loss against the worst-off group assignments in the constraint set. By reducing the constraint set with the marginal distribution, we reduce the optimization parameter space while keeping the objective to be an upper bound to that of the Group DRO with true group assignments with high probability. By harnessing both labeled and unlabeled data in terms of group, we demonstrate in experiments that the Worst-off DRO outperforms both ERM, UnsupDRO, which do not make use of available group labels, as well as the Group DRO (Partial), which does not use unlabeled data.
One future direction, when marginal distribution is not available, is to relax our missing completely at random assumption and bring in different but more realistic modeling assumptions on the missingness of group labels, such as missing at random (MAR), where missing values depend on other observed attributes [12]. In addition, it would be valuable to design a reduced constraint set containing the true group assignment to reduce the performance gap to the Group DRO (Oracle).

heading: ETHICS STATEMENT
text: Machine Learning (ML) models that perform poorly on a minority group or environment have raised a lot of concerns within the AI community and broader society in recent years. To democratize ML in real world, learning ML models that perform robustly across groups or environments has become an important venue of research. The proposed Worst-off DRO is a versatile method that could be employed to train an invariant classifier across groups even when the group information is available only for the portion of the data. This is a rather practical scenario as the group information could be missing for various reasons during the data collection. We further emphasize the importance of theoretical result showing the objective of Worst-off DRO being an upper bound to that of Group DRO with complete group information for safety-critical ML applications.

heading: REPRODUCIBILITY
text: We write our experimental code from scratch using PyTorch library [26]. Due to its similarity, our implementation may closely follow that of Group DRO [30].1 One of the key differentiation of Worst-off DRO is the inner maximization solver for the worst-off group assignments {ĝ}, which we elaborate the exact code using CVXPY solver [8] in Algorithm 2 of Appendix. Additional implementation details, including the neural network architectures, as well as value for hyperparameters including the learning rate, weight decay, batch size, number of training epochs, and algorithmspecific parameters are summarized in Table 3 of Appendix and Section A.3 and A.5.

heading: A APPENDIX

heading: A.1 PROOF OF LEMMAS
text: 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
Proof. The probability of the true group assignment {g?i }Ni=1 in the constraint set Cp?, is written as follows:
P ({g?i }Ni=1 ∈Cp?, ) = P (∣∣p?j − 1N ∑Ni=1 1{g?i = j}∣∣≤ ) ≥ 1− 2e−2N 2 (10)
where Equation 10 holds true from Hoeffding’s inequality.
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
Proof. Using Hoeffding’s inequality, we can show that the estimation error of the marginal distribution is bounded by δ with high probability as follows:
P (|p?j − p̄j | ≤ δ) = P (∣∣p?j − 1K∑Ki=1 1{g?i = j}∣∣ ≤ δ) ≥ 1− 2e−2Kδ2 (11)
Furthermore, we show using Hoeffding’s inequality that P (∣∣p?j − 1N ∑Ni=1 1{gi = j}∣∣ ≤ ) ≥ 1− 2e−2N 2 (12)
Now, the probability of the true group assignment {g?i }Ni=1 in the constraint set Cp̄,δ+ is written as follows:
P ({g?i }Ni=1 ∈Cp̄,δ+ ) = P (∣∣p̄j − 1
N ∑N i=1 1{gi = j} ∣∣≤ δ + ) (13)
≥ P ({∣∣p?j − p̄j∣∣ ≤ δ} ∩ {∣∣p?j − 1N ∑Ni=1 1{gi = j}∣∣ ≤ }) (14)
≥ P (∣∣p?j − p̄j∣∣ ≤ δ)+ P(∣∣p?j − 1N ∑Ni=1 1{gi = j}∣∣ ≤ )− 1 (15) ≥ 1− 2e−2Kδ 2 − 2e−2N 2 (16)
where Equation 14 is due to that the intersection of events in Equation 14 is a subset of an event in Equation 13, and Equation 15 is derived using union bound.

heading: A.2 NOTES ON OPTIMIZATION
text: When using CVXPY to solve for the Worst-off DRO assignments, we simplify the problem by replacing the data marginal distribution ∑N i=1 ĝij in the denominator of Equation 9 to p̄j , thus providing us with a convex optimization problem. The code for the solver is available in Algorithm 2.
Next, Using three sample, we provide an example of the worst-off assignments made by our algorithm,
Example 3. Consider three samples with loss values l1 > l2 > l3 and two predefined groups. Assume the marginal probabilities p̄1 = 0.6 and p̄2 = 0.4. Without loss in generality, assume q1 p̄1 > q2p̄2 . Solving for Worst-off DRO objective results in the following group assignments. With
Algorithm 2 Group Assignment Solver using CVXPY library
import cvxpy as cp import numpy as np
class Solver(object): def __init__(self, n_controls, bsize, marginals, epsilon, labeled=None): """Group assignment solver.
Arguments: n_controls: An integer for the number of groups. bsize: An integer for the batch size. marginals: A 2D array for the marginal distribution. epsilon: A float for the variance. labeled: A tuple for labeled data indices and their value. """ self.X = cp.Variable((bsize, n_controls)) self.l = cp.Parameter((bsize, 1)) self.p = cp.Parameter((n_controls, 1), value=marginals) self.q = cp.Parameter(n_controls) if labeled is not None: labeled_idx, labeled_value = labeled counts = cp.sum(self.X, axis=0, keepdims=True)
obj = ((self.l.T @ self.X) / self.p.T) @ self.q constraints = [self.X >= 0,
cp.sum(self.X, axis=1, keepdims=True) == np.ones((bsize, 1)), cp.abs(cp.sum(self.X, axis=0, keepdims=True) / bsize - self.p.T) <=
epsilon] if labeled is not None: constraints += [self.X[labeled_idx] == labeled_value]
self.prob = cp.Problem(cp.Maximize(obj), constraints)
def cvxsolve(self, losses, weights): """Solver.
Arguments: losses: A 2D array for loss values. weights: A 1D array for group weights q.
Returns: A 2D array for soft group assignments. """ self.l.value = losses self.q.value = weights self.prob.solve() return self.X.value
constraint Cp̄, =0, we obtain {ĝt} =
( 1 0
0.8 0.2 0 1
) . Here, the ith row indicates the assignment given
to sample li.
The above example can be derived by identifying a {ĝ} that satisfies the constraints ∑N i=1 ĝi1 ≤
N p̄1 and ∑N i=1 ĝi2 ≤ N p̄2, where N = 3, p̄1 = 0.6 and p̄1 = 0.4 and correspondingly maximizes Worst-off DRO objective. The above example informs us that group assignments depend on the magnitude of loss values in addition to the group weights and marginal probabilities. Importantly, we find that high loss samples are assigned to groups with high group weights and low marginal probabilities.
Observe that the marginal constraints form a key ingredient of our algorithm as per the above example. Without the marginal constraints, we would have obtained the group assignments
{ĝt} = ( 1 0 1 0 1 0 ) . That is, the assignments would have been made independent of the loss values and sparsely restricted to the group with large qjp̄j value.

heading: A.3 HYPER-PARAMETER TUNING
text: Hyper-parameters were selected for each algorithm by performing an NVP procedure (see Section 4). The best performing model was identified on the validation set associated with each dataset.
All the measures were computed and averaged over three random runs. A list of all the hyperparameters that were tuned for are available in Table 3. The final hyper-parameters selected for each method can be viewed from Table 4.

heading: A.4 ADDITIONAL EXPERIMENTAL RESULTS
text: We provide the following additional results, first, in Figure 5, we show average group accuracies as a function of labelled sample counts. The average group accuracies of the Worst-off DRO method are closely similar across various labelled sample counts. The Group DRO method shows a slight increasing trend in the average accuracies as the number of labelled samples increase. Next, corresponding to the quantitative results of Table 2 in the paper, we provide standard deviations of those results in Table 5. The standard deviations for all the methods are comparable.

heading: A.5 MORE DETAILS ON THE DATASETS

heading: A.5.1 WATERBIRDS
text: This dataset was first introduced in [30] and has been developed by cropping images of birds from the CUB dataset [34] and pasting them on the backgrounds from the Places dataset [41]. A ResNet50 model, pre-trained with ImageNet weights, has been used for training in experiments on this dataset. No data augmentation has been applied for any of the Algorithms.

heading: A.5.2 CMNIST
text: CMNIST dataset comprised of two groups of MNIST images each with a specific color. As per the description in the main paper, the target label is flipped with a specific correlation to the color. Following the implementation of [7], the digit images contain two channels and were downsampled to 14× 14 pixels.

heading: A.5.3 ADULT
text: The Adult dataset used in the paper was obtained from the UCI repository [10]. It contains 44, 842 samples. The features that were used in the experiments include “age”, “workclass”, “fnlwgt”, “education”, “education-num”, “marital-status”, “occupation”, “relationship”, “race”, “sex”, “capitalgain”, “capital-loss”, “hours-per-week”, “native-country”, “income”. A positve target label in this dataset is indicated by the attribute “income-bracket” being above 50K$.

heading: A.5.4 CELEBA
text: For this dataset, the official train-val-test splits as recommended by [21] has been used. Similar to the Waterbirds experiments, a pre-trained ImageNet-based ResNet50 model has been used for the implementations.

heading: A.6 ABLATION STUDY ON INCREASING THE CONSTRAINT SET SIZE.
text: We conduct experiments on Worst-off DRO method for different values of the parameter in the set {0, 0.001, 0.01, 0.1, 1}. The test set accuracies on the minority group and average group are
reported in Figure 6. Increasing the value also increases the constraint set size because the marginal constraint is gradually relaxed. Figure 6 shows that the both minority group accuracy and average group accuracy values reduce with increase in value beyond 0.1 threshold. The accuracy values for ≤ 0.01 are comparable. A similar trend hold on other datasets as well.

Review summaries:
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.