Original: Towards Group Robustness in the Presence of Partial Group Labels
Revised: Towards Group Robustness in the Presence of Partial Group Labels

Original: Abstract
Revised: Abstract

Original: Learning invariant representations is a fundamental requirement for training machine learning models that are inﬂuenced by spurious correlations.
Revised: Learning invariant representations is a fundamental requirement for training machine learning models that are inﬂuenced by spurious correlations.

Original: These spurious correlations, present in the training datasets, wrongly direct the neural network predictions resulting in reduced performance on certain groups, especially the minority groups.
Revised: These spurious correlations, present in the training datasets, wrongly direct the neural network predictions resulting in reduced performance on certain groups, especially the minority groups.

Original: Robust training against such correlations requires the knowledge of group membership on every training sample.
Revised: Robust training against such correlations requires the knowledge of group membership on every training sample.

Original: This need is impractical in situations where the data labeling efforts, for minority/rare groups, are signiﬁcantly laborious or where the individuals comprising the dataset choose to conceal sensitive information pertaining to the groups.
Revised: This need is impractical in situations where the data labeling efforts, for minority/rare groups, are signiﬁcantly laborious or where the individuals comprising the dataset choose to conceal sensitive information pertaining to the groups.

Original: On the other hand, the presence of data collection efforts often results in datasets that contain partially labeled group information.
Revised: On the other hand, the presence of data collection efforts often results in datasets that contain partially labeled group information.

Original: Recent works, addressing the problem, have tackled fully unsupervised scenarios where no labels for groups are available.
Revised: Recent works, addressing the problem, have tackled fully unsupervised scenarios where no labels for groups are available.

Original: We aim to ﬁll a missing gap in the literature that addresses a more realistic setting by leveraging partially available group information during training.
Revised: We aim to ﬁll a missing gap in the literature that addresses a more realistic setting by leveraging partially available group information during training.

Original: First, we construct a constraint set and derive a high probability bound for the group assignment to belong to the set.
Revised: First, we construct a constraint set and derive a high probability bound for the group assignment to belong to the set.

Original: Second, we propose an algorithm that optimizes for a worst-off group assignment from the constraint set.
Revised: Second, we propose an algorithm that optimizes for a worst-off group assignment from the constraint set.

Original: Through experiments on image and tabular datasets, we show improvements in the minority group’s performance while preserving overall accuracy across groups.
Revised: Through experiments on image and tabular datasets, we show improvements in the minority group’s performance while preserving overall accuracy across groups.

Original: Our code is available on https://github.com/ googleinterns/fairness_ssl
Revised: Our code is available on https://github.com/ googleinterns/fairness_ssl

Original: Published at the ICML 

Revised: Published at the ICML

Original: Workshop on Spurious Correlations, Invariance
Revised: Workshop on Spurious Correlations, Invariance

Original: Copyright 2022 by
Revised: Copyright 2022 by

Original: (s).

Original: 1. Introduction
Revised: 1. Introduction

Original: Neural networks being overly biased to certain groups of the data is an increasing concern within the machine learning community (Agarwal et al., 2018).
Revised: Neural networks being overly biased to certain groups of the data is an increasing concern within the machine learning community (Agarwal et al., 2018).

Original: A primary cause for bias against speciﬁc groups is the presence of extraneous attributes in the datasets that wrongly direct the model responses (Xie et al., 2017).
Revised: A primary cause for bias against speciﬁc groups is the presence of extraneous attributes in the datasets that wrongly direct the model responses (Xie et al., 2017).

Original: An inevitable consequence of such correlations to extraneous attributes is disparities in performance across different groups.
Revised: An inevitable consequence of such correlations to extraneous attributes is disparities in performance across different groups.

Original: Speciﬁcally, if certain groups form a minority, a model can simply cheat by having a high overall aggregate accuracy but poor minority group accuracy (Oakden-Rayner et al., 2020).
Revised: Speciﬁcally, if certain groups form a minority, a model can simply cheat by having a high overall aggregate accuracy but poor minority group accuracy (Oakden-Rayner et al., 2020).

Original: Existing works (Arjovsky et al., 2019; Sagawa et al., 2019) operate in the regime where the number of groups are known apriori.
Revised: Existing works (Arjovsky et al., 2019; Sagawa et al., 2019) operate in the regime where the number of groups are known apriori.

Original: Further, most works assume a complete knowledge of the group membership of individual samples.
Revised: Further, most works assume a complete knowledge of the group membership of individual samples.

Original: While these methods have been proven effective, it is not realistic to assume complete access to the group labels.
Revised: While these methods have been proven effective, it is not realistic to assume complete access to the group labels.

Original: For example, in a lung cancer detection problem, the label class could contain many unrecognized subgroups such as solid/subsolid tumors and central/peripheral neoplasms (Oakden-Rayner et al., 2020).
Revised: For example, in a lung cancer detection problem, the label class could contain many unrecognized subgroups such as solid/subsolid tumors and central/peripheral neoplasms (Oakden-Rayner et al., 2020).

Original: These unrecognized subgroups are difﬁcult to label and form a minority in the dataset, thereby resulting in an unbalanced performance across the different subgroups (Sohoni et al., 2020).
Revised: These unrecognized subgroups are difﬁcult to label and form a minority in the dataset, thereby resulting in an unbalanced performance across the different subgroups (Sohoni et al., 2020).

Original: In this work, we consider a setting where a signiﬁcant portion of the training data is devoid of group labels.
Revised: In this work, we consider a setting where a signiﬁcant portion of the training data is devoid of group labels.

Original: We choose to ﬁll a missing gap in the literature where several works bifurcate into methods that either are fully supervised or fully unsupervised.
Revised: We choose to ﬁll a missing gap in the literature where several works bifurcate into methods that either are fully supervised or fully unsupervised.

Original: We answer the question using a framework of distributionally robust optimization (DRO) (Shapiro et al., 2021;
Revised: We answer the question using a framework of distributionally robust optimization (DRO) (Shapiro et al., 2021;

Original: Namkoong & Duchi, 2016).
Revised: Namkoong & Duchi, 2016).

Revised: Prior works that utilize a Maximum Mean Discrepancy (MMD) framework exist (Goldstein et al., 2022).

Original: DRO optimizes for the worstcase training loss over predeﬁned set of groups and is closely connected to the Rawlsian criterion (Rawls, 2001).
Revised: DRO optimizes for the worst-case training loss over predeﬁned set of groups and is closely connected to the Rawlsian criterion (Sagawa et al., 2019; Rawls, 2001).

Original: The application of DRO to the partial group label setting poses some challenges: (1) the lack of group label makes it infeasible to compute the worst-off group loss; (2) optimizing only for the high-loss samples, by considering them as a worst-off group, discards considerable portion of the training data thereby impacting the overall accuracy; and (3) inferring missing group labels with pseudo-label methods is a cause for ethical concerns.
Revised: Applying DRO to the partial group label setting poses some challenges: (1) the lack of group label makes it infeasible to compute the worst-off group loss; (2) optimizing only for the high-loss samples, by considering them as a worst-off group, discards considerable portion of the training data thereby impacting the overall performance of the method; and (3) inferring missing group labels with pseudo-label methods is a cause for ethical concerns.

Original: The third challenge above suggests a straightforward way of handling partial labels wherein we directly estimate the group label for unlabelled samples.
Revised: The third challenge above suggests a straightforward way of handling partial labels wherein we directly estimate the group label for unlabelled samples.

Original: However, this approach could be harmful in the context of fairness problems because the estimated labels are susceptible to misuse by a wrongdoer.
Revised: However, this approach could be harmful in the context of fairness problems because the estimated labels are susceptible to misuse by a wrongdoer.

Original: For example, when the groups are indicative of sensitive information such as age or gender, an incorrect estimation would wrongly designate the demographics of an individual.
Revised: For example, when the groups are indicative of sensitive information such as age or gender, an incorrect estimation would wrongly designate the demographics of an individual.

Original: Moreover, when it’s desirable to conceal such sensitive information, a direct estimation of groups would be violation of privacy.
Revised: Moreover, when it’s desirable to conceal such sensitive information, a direct estimation of groups would be violation of privacy.

Original: Thus we cautiously avoid building or utilizing pseudo-label based methods in this paper.
Revised: Thus we cautiously avoid building or utilizing pseudo-label based methods in this paper.

Original: In light of above challenges, we make the following contributions .
Revised: In light of above challenges, we make the following contributions .

Original: We propose a method that deﬁnes a constraint set of group assignments and optimize over worst-off assignments within the set.
Revised: We propose a method that deﬁnes a constraint set of group assignments and optimize over worst-off assignments within the set.

Original: We show that the constraint set encompasses the ground-truth group labels of the unlabeled data with high probability.
Revised: We show that the constraint set encompasses the ground-truth group labels of the unlabeled data with high probability.

Original: Since worst-off assignments do not directly relate to ground-truth, our approach is safe and does not violate privacy.
Revised: Since worst-off assignments do not directly relate to ground-truth, our approach is safe and does not violate privacy.

Original: As we shall see, our method assigns high loss samples to groups with low marginal probability and does not discard any samples.
Revised: As we shall see, our method assigns high loss samples to groups with low marginal probability and does not discard any samples.

Original: We show experiments on three imaging datasets and one tabular dataset.
Revised: We show experiments on three imaging datasets and one tabular dataset.

Original: 2. Methodology
Revised: 2. Methodology

Original: We revisit the GroupDRO in Section 2.1 and detail our method, Worst-off DRO, in Section 2.2 and 2.3.
Revised: We revisit the GroupDRO in Section 2.1 and detail our method, Worst-off DRO, in Section 2.2 and 2.3.

Original: In Section 2.4, we describe a practical algorithm for optimization.
Revised: In Section 2.4, we describe a practical algorithm for optimization.

Original: 2.1. Preliminary: Group DRO
Revised: 2.1. Preliminary: Group DRO

Original: Let x ∈ X ⊂ R d be data descriptors, y ∈ Y ⊂ { 0 , 1 } be target labels, and g ∈ G ⊂ { 1 , ..., M } be group labels.
Revised: Let x ∈ X ⊂ R d be data descriptors, y ∈ Y ⊂ { 0 , 1 } be target labels, and g ∈ G ⊂ { 1 , ..., M } be group labels.

Original: We assume training a neural network parameterized by the weights w that corresponds to a per-sample loss l ( x, y ; w ) .
Revised: We assume training a neural network parameterized by the weights w that corresponds to a per-sample loss l ( x, y ; w ) .

Original: Given data triplets { ( x i , y i , g i ) } Ni =1 , we seek to optimize w for the
Revised: Given data triplets { ( x i , y i , g i ) } Ni =1 , we seek to optimize w for the

Original: Rawlsian criterion (Rawls, 2001; Hashimoto et al., 2018), which minimizes the loss of the worst-off group, as follows:
Revised: Rawlsian criterion (Rawls, 2001; Hashimoto et al., 2018), which minimizes the loss of the worst-off group, as follows:

Original: [Equation]
Revised: [Equation]

Original: Sagawa et al.
Revised: Sagawa et al.

Original: (2019) proposed a practical algorithm to solve (1), called Group DRO .
Revised: (2019) proposed a practical algorithm to solve (1), called Group DRO .

Original: This method optimizes a weighted expected loss across all groups.
Revised: This method optimizes a weighted expected loss across all groups.

Original: These weights over the groups, denoted by q , are drawn from a simplex ∆ M .
Revised: These weights over the groups, denoted by q , are drawn from a simplex ∆ M .

Original: The objective function L GDRO is as described below,
Revised: The objective function L GDRO is as described below,

Original: [Equation]
Revised: [Equation]

Original: 2.2. Our Approach: Worst-off DRO
Revised: 2.2. Our Approach: Worst-off DRO

Original: In this work, we are interested in training a robust neural network when group labels are only partially available.
Revised: In this work, we are interested in training a robust neural network when group labels are only partially available.

Original: That is, our training dataset consists of fully-labeled { ( x i , y i , g (cid:63)i ) }
Revised: That is, our training dataset consists of fully-labeled { ( x i , y i , g (cid:63)i ) }

Original: Ki =1 and task-labeled samples { ( x i , y i , − ) } Ni = K +1 , where − indicates the missing groups.
Revised: Ki =1 and task-labeled samples { ( x i , y i , − ) } Ni = K +1 , where − indicates the missing groups.

Original: As noted in (2), the Group DRO requires group labels of entire dataset.
Revised: As noted in (2), the Group DRO requires group labels of entire dataset.

Original: When some of them are missing, we propose to optimize for the following objective for L WDRO ( C ) :
Revised: When some of them are missing, we propose to optimize for the following objective for L WDRO ( C ) :

Original: [Equation]
Revised: [Equation]

Original: where C is a set of group assignments { ˆ g i } Ni =1 satisfying ˆ
Revised: where C is a set of group assignments { ˆ g i } Ni =1 satisfying ˆ

Original: g i =
Revised: g i =

Original: g (cid:63)i , ∀ i ≤ K .
Revised: g (cid:63)i , ∀ i ≤ K .

Original: We call the objective in (3) Worst-off
Revised: We call the objective in (3) Worst-off

Original: DRO method; as it optimizes over the worst-off assignments in a certain constraint set C (more details on C soon).
Revised: DRO method; as it optimizes over the worst-off assignments in a certain constraint set C (more details on C soon).

Original: Note that the Worst-off DRO objective forms an upper bound to the Group DRO objective evaluated at the groundtruth group labels if { g (cid:63)i } Ni =1 ∈ C .
Revised: Note that the Worst-off DRO objective forms an upper bound to the Group DRO objective evaluated at the groundtruth group labels if { g (cid:63)i } Ni =1 ∈ C .

Original: Under identical parameters w and q , this is rather a straightforward consequence from the fact that the ground-truths { g (cid:63)i } Ni =1 falls within the constraint set C .
Revised: Under identical parameters w and q , this is rather a straightforward consequence from the fact that the ground-truths { g (cid:63)i } Ni =1 falls within the constraint set C .

Original: However, the following lemma generalizes the upper bound relationship between Worst-off DRO and
Revised: However, the following lemma generalizes the upper bound relationship between Worst-off DRO and

Original: Group DRO objectives for all w and q parameters.
Revised: Group DRO objectives for all w and q parameters.

Original: Lemma 2.1.
Revised: Lemma 2.1.

Original: Denote L GDRO at a speciﬁc w and q as
Revised: Denote L GDRO at a speciﬁc w and q as

Original: L GDRO( w,q ) .
Revised: L GDRO( w,q ) .

Original: Similarly L WDRO ( C ) at a ﬁxed w and q is denoted by L WDRO( w,q ) ( C ) .
Revised: Similarly L WDRO ( C ) at a ﬁxed w and q is denoted by L WDRO( w,q ) ( C ) .

Original: When the ground-truth group assignment { g (cid:63)i } Ni =1 ∈ C , we have
Revised: When the ground-truth group assignment { g (cid:63)i } Ni =1 ∈ C , we have

Original: [Equation]
Revised: [Equation]

Original: The proof is in Appendix A.3.
Revised: The proof is in Appendix A.3.

Original: For safety-critical applications, such as learning a fair classiﬁer, it is important that the optimal objective (i.e., Group DRO with a groundtruth group assignment) is bounded by the objective used in optimization as in Lemma 2.1.
Revised: For safety-critical applications, such as learning a fair classiﬁer, it is important that the optimal objective (i.e., Group DRO with a groundtruth group assignment) is bounded by the objective used in optimization as in Lemma 2.1.

Original: This is simply because optimizing the proposed objective guarantees that the corresponding lower bound Group DRO is also optimized.
Revised: This is simply because optimizing the proposed objective guarantees that the corresponding lower bound Group DRO is also optimized.

Original: 2.3. Reducing Constraint Set with Marginal Distribution Constraint
Revised: 2.3. Reducing Constraint Set with Marginal Distribution Constraint

Original: Clearly the constraint set plays an important role in connecting Worst-off DRO to Group DRO.
Revised: Clearly the constraint set plays an important role in connecting Worst-off DRO to Group DRO.

Original: We see that Worst-off
Revised: We see that Worst-off

Original: DRO can be made closer to Group DRO by further restricting C as long as C contains the ground-truths.
Revised: DRO can be made closer to Group DRO by further restricting C as long as C contains the ground-truths.

Original: To achieve this goal, we utilize marginal distribution constraints.
Revised: To achieve this goal, we utilize marginal distribution constraints.

Original: These constraints may be given as a side information or could be estimated from the small set of partial group labels.
Revised: These constraints may be given as a side information or could be estimated from the small set of partial group labels.

Original: Let C p ,(cid:15) be subset of C whose elements { g i } Ni =1 satisfy, do 5: { ˆ g t } ←
Revised: Let C p ,(cid:15) be subset of C whose elements { g i } Ni =1 satisfy, do 5: { ˆ g t } ←

Original: [Equation]
Revised: [Equation]

Original: where, C ¯ p ,(cid:15) as deﬁned in (10).
Revised: where, C ¯ p ,(cid:15) as deﬁned in (10).

Original: Gradient descent on w : w
Revised: Gradient descent on w : w

Original: [Equation]
Revised: [Equation]

Original: Exponential ascent on q :
Revised: Exponential ascent on q :

Original: [Equation]
Revised: [Equation]

Original: 8: end for 9: Output: Trained neural parameters
Revised: 8: end for 9: Output: Trained neural parameters

Original: [Equation]
Revised: [Equation]

Original: T
Revised: T

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: where (5) implies that the true group labels are assigned whenever available, and (6) implies that the data marginal distribution should be close to the marginal distribution p .
Revised: where (5) implies that the true group labels are assigned whenever available, and (6) implies that the data marginal distribution should be close to the marginal distribution p .

Original: Then, for any marginal distribution p and (cid:15) > 0 , it is easy to show L WDRO ( C p ,(cid:15) )
Revised: Then, for any marginal distribution p and (cid:15) > 0 , it is easy to show L WDRO ( C p ,(cid:15) )

Original: ≤ L WDRO ( C ) as C p ,(cid:15)
Revised: ≤ L WDRO ( C ) as C p ,(cid:15)

Original: ⊂ C .
Revised: ⊂ C .

Original: Moreover, we will see in Lemma 2.2 that, with high probability, the constraint set C p (cid:63) ,(cid:15) with the true marginal distribution p (cid:63) contains the true group assignment { g (cid:63) } .
Revised: Moreover, we will see in Lemma 2.2 that, with high probability, the constraint set C p (cid:63) ,(cid:15) with the true marginal distribution p (cid:63) contains the true group assignment { g (cid:63) } .

Original: Lemma 2.2.
Revised: Lemma 2.2.

Original: The constraint set C p (cid:63) ,(cid:15) contains the true group labels { g (cid:63)i } Ni =1 with high probability:
Revised: The constraint set C p (cid:63) ,(cid:15) contains the true group labels { g (cid:63)i } Ni =1 with high probability:

Original: [Equation]
Revised: [Equation]

Original: The proof is in Appendix A.3.
Revised: The proof is in Appendix A.3.

Original: As in (7), the probability of the constraint set containing the true group labels gets closer to 1 by allowing a larger variance ( (cid:15) ) from the true marginals.
Revised: As in (7), the probability of the constraint set containing the true group labels gets closer to 1 by allowing a larger variance ( (cid:15) ) from the true marginals.

Original: For ﬁxed (cid:15) > 0 , the probability gets closer to 1 as we increase the number of unlabeled data ( N ).
Revised: For ﬁxed (cid:15) > 0 , the probability gets closer to 1 as we increase the number of unlabeled data ( N ).

Original: Finally, this implies that L WDRO ( C p (cid:63) ,(cid:15) ) is an upper bound to L GDRO :
Revised: Finally, this implies that L WDRO ( C p (cid:63) ,(cid:15) ) is an upper bound to L GDRO :

Original: [Equation]
Revised: [Equation]

Original: In practice, however, the true marginal distribution p (cid:63) may not be available.
Revised: In practice, however, the true marginal distribution p (cid:63) may not be available.

Original: For our setting where group labels are partially available, with an assumption that group labels are missing completely at random (MCAR) (Rubin, 1976), the true marginal distribution could be estimated from the labelled subset.
Revised: For our setting where group labels are partially available, with an assumption that group labels are missing completely at random (MCAR) (Rubin, 1976), the true marginal distribution could be estimated from the labelled subset.

Original: This again allows us to formulate a constraint set that contains the ground-truth group assignment with high probability.
Revised: This again allows us to formulate a constraint set that contains the ground-truth group assignment with high probability.

Original: Let ¯ p be the estimate of the marginal distribution from { ( x i , y i , g (cid:63)i ) }
Revised: Let ¯ p be the estimate of the marginal distribution from { ( x i , y i , g (cid:63)i ) }

Original: Ki =1 .
Revised: Ki =1 .

Original: Lemma 2.3.
Revised: Lemma 2.3.

Original: The constraint set C ¯ p ,δ + (cid:15) contains the true group labels { g (cid:63)i } Ni =1 with high probability:
Revised: The constraint set C ¯ p ,δ + (cid:15) contains the true group labels { g (cid:63)i } Ni =1 with high probability:

Original: [Equation]
Revised: [Equation]

Original: Proof in Appendix A.3.
Revised: Proof in Appendix A.3.

Original: Here, δ accounts for estimation error from the true marginals p (cid:63) .
Revised: Here, δ accounts for estimation error from the true marginals p (cid:63) .

Original: 2.4. A Practical Optimization Algorithm
Revised: 2.4. A Practical Optimization Algorithm

Original: We are interested in solving the optimization problem
Revised: We are interested in solving the optimization problem

Original: L WDRO ( C p ,(cid:15) ) .
Revised: L WDRO ( C p ,(cid:15) ) .

Original: Unfortunately, the inner maximization problem with respect to the group assignments { ˆ g } in (3) is challenging as variables are discrete and the objective cannot be decomposed due to the marginal distribution constraint.
Revised: Unfortunately, the inner maximization problem with respect to the group assignments { ˆ g } in (3) is challenging as variables are discrete and the objective cannot be decomposed due to the marginal distribution constraint.

Original: Thus, we propose to use a soft group assignments.
Revised: Thus, we propose to use a soft group assignments.

Original: Specifically, for each sample, we retain a soft group assignment ˆ g i ∈
Revised: Specifically, for each sample, we retain a soft group assignment ˆ g i ∈

Original: ∆ M , and optimize the Worst-off DRO objective over the constraint set C ¯ p,(cid:15) as deﬁned below:
Revised: ∆ M , and optimize the Worst-off DRO objective over the constraint set C ¯ p,(cid:15) as deﬁned below:

Original: [Equation]
Revised: [Equation]

Original: The ﬁrst condition ensures that the assignments form a probability simplex, second one assures consistency with labeled data, and the third one validates the data marginal distribution follows the provided distribution.
Revised: The ﬁrst condition ensures that the assignments form a probability simplex, second one assures consistency with labeled data, and the third one validates the data marginal distribution follows the provided distribution.

Original: The third constraint also provides for a mitigation strategy when ¯p is misspeciﬁed (likely when data is not MCAR).
Revised: The third constraint also provides for a mitigation strategy when ¯p is misspeciﬁed (likely when data is not MCAR).

Original: We alternate optimization over w , q and { ˆ g } as shown in Algorithm 1.
Revised: We alternate optimization over w , q and { ˆ g } as shown in Algorithm 1.

Original: Next, we will see how the worst-off assignments computed by the algorithm look to be.
Revised: Next, we will see how the worst-off assignments computed by the algorithm look to be.

Original: For simplicity, consider the case (cid:15) = 0 and K = 0 , (i.e., no group-labelled samples).
Revised: For simplicity, consider the case (cid:15) = 0 and K = 0 , (i.e., no group-labelled samples).

Original: Denoting
Revised: Denoting

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: we can re-write the maximization over { ˆ g ij } as,
Revised: we can re-write the maximization over { ˆ g ij } as,

Original: [Equation]
Revised: [Equation]

Original: The constraints ensure (cid:80) Ni
Revised: The constraints ensure (cid:80) Ni

Original: =

Revised: =

Original: ˆ g ij =
Revised: ˆ g ij =

Original: N ¯p j and (cid:80) Mi
Revised: N ¯p j and (cid:80) Mi

Original: =
Revised: =

Original: j ˆ g ij = 1 , ˆ g ij ≥ 0 for all i ≤ N and j ≤ M respectively.
Revised: j ˆ g ij = 1 , ˆ g ij ≥ 0 for all i ≤ N and j ≤ M respectively.

Original: The linear program (11) sets the highest mass on ˆ g ij for i and j that maximize θ
Revised: The linear program (11) sets the highest mass on ˆ g ij for i and j that maximize θ

Original: j and sample loss l i .
Revised: j and sample loss l i .

Original: A large θ j represents groups with a high group weight q j and low marginal probability ¯p j , characteristic of a worst-off group.
Revised: A large θ j represents groups with a high group weight q j and low marginal probability ¯p j , characteristic of a worst-off group.

Original: In summary,
Revised: In summary,

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: beyond 0 .
Revised: beyond 0 .

Original: high loss samples are assigned to groups with high group weights and low marginal probabilities .
Revised: high loss samples are assigned to groups with high group weights and low marginal probabilities .

Original: 3. Experiments
Revised: 3. Experiments

Original: We test the efﬁcacy of our method on image and tabular datasets 1 .
Revised: We test the efﬁcacy of our method on image and tabular datasets 1 .

Original: We use Waterbirds (Sagawa et al., 2019), Group
Revised: We use Waterbirds (Sagawa et al., 2019), Group

Original: CMNIST (Arjovsky et al., 2019), Group Adult (Dua et al., 2017), and CelebA (Liu et al., 2015).
Revised: CMNIST (Arjovsky et al., 2019), Group Adult (Dua et al., 2017), and CelebA (Liu et al., 2015).

Original: We compare our method against ERM, Unsup DRO (Hashimoto et al., 2018) and Group DRO.
Revised: We compare our method against ERM, Unsup DRO (Hashimoto et al., 2018) and Group DRO.

Original: We consider a variant of Group DRO that only uses samples with group labels at train time.
Revised: We consider a variant of Group DRO that only uses samples with group labels at train time.

Original: We call the method Group DRO (Partial) , to contrast with fullylabelled counterpart, Group DRO (Oracle) .
Revised: We call the method Group DRO (Partial) , to contrast with fullylabelled counterpart, Group DRO (Oracle) .

Original: Our quantitative results in Table 1 show that the minority group accuracy is higher relative to the baselines for Worst-off DRO method 2 .
Revised: Our quantitative results in Table 1 show that the minority group accuracy is higher relative to the baselines for Worst-off DRO method 2 .

Original: Moreover, the average accuracy across all the groups is within a 2% window for all the methods.
Revised: Moreover, the average accuracy across all the groups is within a 2% window for all the methods.

Original: Ablation Experiments.
Revised: Ablation Experiments.

Original: In Figure 1a, we show that the minority group accuracy increases with more labelled samples.
Revised: In Figure 1a, we show that the minority group accuracy increases with more labelled samples.

Original: More labels provides two beneﬁts.
Revised: More labels provides two beneﬁts.

Original: Firstly, the standard deviation of errors in estimating the marginal probabilities reduces (Wasserman, 2004) ( ≈ √
Revised: Firstly, the standard deviation of errors in estimating the marginal probabilities reduces (Wasserman, 2004) ( ≈ √

Original: # samples rate).
Revised: # samples rate).

Original: Secondly, labelled groups reinforce an accurate evaluation of the Rawlsian objective in (2).
Revised: Secondly, labelled groups reinforce an accurate evaluation of the Rawlsian objective in (2).

Original: In Figure 1b, we plot the minority/average group accuracies for different hyper-parameter choices.
Revised: In Figure 1b, we plot the minority/average group accuracies for different hyper-parameter choices.

Original: Due to distributional differences amongst the groups, a mild tradeoff exists between the two accuracy measures.
Revised: Due to distributional differences amongst the groups, a mild tradeoff exists between the two accuracy measures.

Original: Evidently, the top-right regions in the plot are desirable, at which the Worst-off DRO models are positioned.
Revised: Evidently, the top-right regions in the plot are desirable, at which the Worst-off DRO models are positioned.

Original: Next, in
Revised: Next, in

Original: Figure 1c, we plot group weights ( q − values) in the training phase.
Revised: Figure 1c, we plot group weights ( q − values) in the training phase.

Original: The plots show a high q − value on the minor- ity groups indicating that the loss values on them are upweighted relative to other groups.
Revised: The plots show a high q − value on the minor- ity groups indicating that the loss values on them are upweighted relative to other groups.

Original: Lastly, Figure 1d describes experiments where the constraint set size is gradually increased by varying the (cid:15) parameter of C ¯p ,(cid:15) .
Revised: Lastly, Figure 1d describes experiments where the constraint set size is gradually increased by varying the (cid:15) parameter of C ¯p ,(cid:15) .

Original: Increasing (cid:15) parameter accommodates the case where ¯p is misspeciﬁed.
Revised: Increasing (cid:15) parameter accommodates the case where ¯p is misspeciﬁed.

Original: 4. Conclusion
Revised: 4. Conclusion

Original: We present Worst-off DRO, an invariant learning method for partially labelled datasets.
Revised: We present Worst-off DRO, an invariant learning method for partially labelled datasets.

Original: Worst-off DRO extends Group
Revised: Worst-off DRO extends Group

Original: DRO by optimizing the loss against the worst-off group assignments in a constraint set.
Revised: DRO by optimizing the loss against the worst-off group assignments in a constraint set.

Original: By reducing the constraint set with the marginal distribution, we reduce the optimization parameter space while maintaining the objective to be an upper bound to that of the Group DRO with true group assignments.
Revised: By reducing the constraint set with the marginal distribution, we reduce the optimization parameter space while maintaining the objective to be an upper bound to that of the Group DRO with true group assignments.

Original: By harnessing both labeled and unlabeled data in terms of group, we demonstrate in experiments
Revised: By harnessing both labeled and unlabeled data in terms of group, we demonstrate in experiments

Original: that
Revised: that

Original: the
Revised: the

Original: Worst-off DRO outperforms both ERM, UnsupDRO, which do not make use of available group labels, as well as the
Revised: Worst-off DRO outperforms both ERM, UnsupDRO, which do not make use of available group labels, as well as the

Original: Group DRO (Partial), which does not use unlabeled data.
Revised: Group DRO (Partial), which does not use unlabeled data.

Original: Acknowledgements The authors are grateful to Prof. Vikas
Revised: Acknowledgements The authors are grateful to Prof. Vikas

Original: Singh (University of Wisconsin - Madison) for discussions discussion on this project.
Revised: Singh (University of Wisconsin - Madison) for discussions discussion on this project.

Original: Special thanks to Sercan Arik,
Revised: Special thanks to Sercan Arik,

Original: Jeremy Martin and Alex Beutel for reviewing and providing feedback on the manuscript.
Revised: Jeremy Martin and Alex Beutel for reviewing and providing feedback on the manuscript.

Original: A. Appendix A.1. Ethics Statement
Revised: A. Appendix A.1. Ethics Statement

Original: 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.
Revised: 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.

Original: To democratize ML in real world, learning ML models that perform robustly across groups or environments has become an important venue of research.
Revised: To democratize ML in real world, learning ML models that perform robustly across groups or environments has become an important venue of research.

Original: The proposed Worst-off DRO is a versatile method that could be employed to train an invariant classiﬁer across groups even when the group information is available only for the portion of the data.
Revised: The proposed Worst-off DRO is a versatile method that could be employed to train an invariant classiﬁer across groups even when the group information is available only for the portion of the data.

Original: This is a rather practical scenario as the group information could be missing for various reasons during the data collection.
Revised: This is a rather practical scenario as the group information could be missing for various reasons during the data collection.

Original: 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.
Revised: 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.

Original: A.2. Reproducibility
Revised: A.2. Reproducibility

Original: We write our experimental code from scratch using PyTorch library (Paszke et al., 2019).
Revised: We write our experimental code from scratch using PyTorch library (Paszke et al., 2019).

Original: Due to its similarity, our implementation may closely follow that of Group DRO (Sagawa et al., 2019).
Revised: Due to its similarity, our implementation may closely follow that of Group DRO (Sagawa et al., 2019).

Original: 3 One of the key differentiation of Worst-off
Revised: 3 One of the key differentiation of Worst-off

Original: DRO is the inner maximization solver for the worst-off group assignments { ˆ g } , which we elaborate the exact code using
Revised: DRO is the inner maximization solver for the worst-off group assignments { ˆ g } , which we elaborate the exact code using

Original: CVXPY solver (Diamond & Boyd, 2016) in Algorithm 2 of Appendix.
Revised: CVXPY solver (Diamond & Boyd, 2016) in Algorithm 2 of Appendix.

Original: 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 algorithm-speciﬁc parameters are summarized in Table 2 and Section A.8 and A.11.
Revised: 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 algorithm-speciﬁc parameters are summarized in Table 2 and Section A.8 and A.11.

Original: A.3. Proof of Lemmas
Revised: A.3. Proof of Lemmas

Original: Lemma A.1.
Revised: Lemma A.1.

Original: Denote L GDRO at a given w and q parameters as L GDRO( w,q ) .
Revised: Denote L GDRO at a given w and q parameters as L GDRO( w,q ) .

Original: Similarly L WDRO ( C ) at a ﬁxed w and q as
Revised: Similarly L WDRO ( C ) at a ﬁxed w and q as

Original: L WDRO( w,q ) ( C ) .
Revised: L WDRO( w,q ) ( C ) .

Original: When the ground-truth group assignment { g (cid:63)i } Ni =1 ∈ C , we have
Revised: When the ground-truth group assignment { g (cid:63)i } Ni =1 ∈ C , we have

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Under the case { g (cid:63)i } Ni =1 ∈ C , due to the max over C , we have
Revised: Under the case { g (cid:63)i } Ni =1 ∈ C , due to the max over C , we have

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Deﬁne q ∗ WDRO =
Revised: Deﬁne q ∗ WDRO =

Original: arg max q ∈ ∆ M L WDRO( w,q ) ( C ) and q ∗ GDRO = arg max q ∈ ∆ M L GDRO( w,q ) .
Revised: arg max q ∈ ∆ M L WDRO( w,q ) ( C ) and q ∗ GDRO = arg max q ∈ ∆ M L GDRO( w,q ) .

Original: From the above deﬁnitions, we have,
Revised: From the above deﬁnitions, we have,

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Moreover,
Revised: Moreover,

Original: [Equation]
Revised: [Equation]

Original: Minimizing (17) over w , we obtain,
Revised: Minimizing (17) over w , we obtain,

Original: [Equation]
Revised: [Equation]

Original: Lemma A.2.
Revised: Lemma A.2.

Original: The constraint set C p (cid:63) ,(cid:15) contains the true group labels { g (cid:63)i } Ni =1 with high probability:
Revised: The constraint set C p (cid:63) ,(cid:15) contains the true group labels { g (cid:63)i } Ni =1 with high probability:

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: The probability of the true group assignment { g (cid:63)i } Ni =1 in the constraint set C p (cid:63) ,
Revised: The probability of the true group assignment { g (cid:63)i } Ni =1 in the constraint set C p (cid:63) ,

Original: [Equation]
Revised: [Equation]

Original: is written as follows:
Revised: is written as follows:

Original: [Equation]
Revised: [Equation]

Original: where (18) holds true from the Hoeffding’s inequality.
Revised: where (18) holds true from the Hoeffding’s inequality.

Original: Lemma A.3.
Revised: Lemma A.3.

Original: The constraint set C ¯ p ,δ + (cid:15) contains the true group labels { g (cid:63)i } Ni =1 with high probability:
Revised: The constraint set C ¯ p ,δ + (cid:15) contains the true group labels { g (cid:63)i } Ni =1 with high probability:

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Using Hoeffding’s inequality, we can show that the estimation error of the marginal distribution is bounded by δ with high probability as follows:
Revised: Using Hoeffding’s inequality, we can show that the estimation error of the marginal distribution is bounded by δ with high probability as follows:

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Furthermore, we show using Hoeffding’s inequality that
Revised: Furthermore, we show using Hoeffding’s inequality that

Original: [Equation]
Revised: [Equation]

Original: where (22) is due to that the intersection of events and is a subset of an event in (21).
Revised: where (22) is due to that the intersection of events and is a subset of an event in (21).

Original: Then, (23) is derived using union bound.
Revised: Then, (23) is derived using union bound.

Original: Now, the probability of the true group assignment { g (cid:63)i } Ni =1 in the constraint set C ¯ p ,δ + (cid:15) is written as follows:
Revised: Now, the probability of the true group assignment { g (cid:63)i } Ni =1 in the constraint set C ¯ p ,δ + (cid:15) is written as follows:

Original: [Equation]
Revised: [Equation]

Original: P (cid:16)
Revised: P (cid:16)

Original: M (cid:91) j
Revised: M (cid:91) j

Original: = (cid:12)(cid:12)
Revised: = (cid:12)(cid:12)

Original: ¯ p j − 1 N (cid:88) N i =1 1 { g i = j } (cid:12)(cid:12) >
Revised: ¯ p j − 1 N (cid:88) N i =1 1 { g i = j } (cid:12)(cid:12) >

Original: δ + (cid:15) (cid:17) (26) from union bound, we get, (27) ≥ 1 −
Revised: δ + (cid:15) (cid:17) (26) from union bound, we get, (27) ≥ 1 −

Original: M (cid:88) j
Revised: M (cid:88) j

Original: =
Revised: =

Original: P (cid:16)(cid:12)(cid:12) ¯ p j − 1 N (cid:88) N i =1 1 {
Revised: P (cid:16)(cid:12)(cid:12) ¯ p j − 1 N (cid:88) N i =1 1 {

Original: g i =
Revised: g i =

Original: j } (cid:12)(cid:12) >
Revised: j } (cid:12)(cid:12) >

Original: δ + (cid:15) (cid:17) (28) from (23), we have (29)
Revised: δ + (cid:15) (cid:17) (28) from (23), we have (29)

Original: [Equation]
Revised: [Equation]

Original: A.4. Notes on the Optimization Procedure
Revised: A.4. Notes on the Optimization Procedure

Original: When using CVXPY to solve for the Worst-off DRO assignments, we simplify the problem by replacing the data marginal distribution (cid:80) Ni
Revised: When using CVXPY to solve for the Worst-off DRO assignments, we simplify the problem by replacing the data marginal distribution (cid:80) Ni

Original: =

Revised: =

Original: ˆ g ij in the denominator of (10) to ¯ p j , thus providing us with a convex optimization problem.
Revised: ˆ g ij in the denominator of (10) to ¯ p j , thus providing us with a convex optimization problem.

Original: The code for the solver is available in Algorithm 2.
Revised: The code for the solver is available in Algorithm 2.

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: A.
Revised: A.

Original: . An example of worst
Revised: . An example of worst

Original: off assignments
Revised: off assignments

Original: Using three samples, we provide an example of the worst-off assignments made by our algorithm,
Revised: Using three samples, we provide an example of the worst-off assignments made by our algorithm,

Original: Example A.4 .
Revised: Example A.4 .

Original: Consider three samples with loss values l 1 > l 2 > l 3 and two predeﬁned groups.
Revised: Consider three samples with loss values l 1 > l 2 > l 3 and two predeﬁned groups.

Original: Assume the marginal probabilities ¯ p 1 = 0 .
Revised: Assume the marginal probabilities ¯ p 1 = 0 .

Original: 6 and ¯ p 2 = 0 .
Revised: 6 and ¯ p 2 = 0 .

Original: Without loss in generality, assume q 1 ¯ p 1 > q 2 ¯ p 2 .
Revised: Without loss in generality, assume q 1 ¯ p 1 > q 2 ¯ p 2 .

Original: With constraint C ¯p ,(cid:15) =0 and solving for
Revised: With constraint C ¯p ,(cid:15) =0 and solving for

Original: Worst-off DRO objective results in the following group assignments, { ˆ g t } =
Revised: Worst-off DRO objective results in the following group assignments, { ˆ g t } =

Original: (cid:18)
Revised: (cid:18)

Original: [Table]
Revised: [Table]

Original: (cid:19) .
Revised: (cid:19) .

Original: Here, the i th row indicates the assignment given to sample l i .
Revised: Here, the i th row indicates the assignment given to sample l i .

Original: The group assignments can be derived by identifying a { ˆ g } that satisﬁes the constraints (cid:80) Ni
Revised: The group assignments can be derived by identifying a { ˆ g } that satisﬁes the constraints (cid:80) Ni

Original: =

Revised: =

Original: ˆ g i 1 ≤ N ¯ p 1 and (cid:80) Ni
Revised: ˆ g i 1 ≤ N ¯ p 1 and (cid:80) Ni

Original: =

Revised: =

Original: ˆ g i 2 ≤
Revised: ˆ g i 2 ≤

Original: N ¯ p 2 , where N = 3 , ¯ p 1 = 0 .
Revised: N ¯ p 2 , where N = 3 , ¯ p 1 = 0 .

Original: 6 and ¯ p 1 = 0 .
Revised: 6 and ¯ p 1 = 0 .

Original: 4 , and correspondingly maximizes Worst-off DRO objective.
Revised: 4 , and correspondingly maximizes Worst-off DRO objective.

Original: The above example informs us that group assignments depend on the magnitude of loss values in addition to the group weights and marginal probabilities.
Revised: The above example informs us that group assignments depend on the magnitude of loss values in addition to the group weights and marginal probabilities.

Original: As indicated in the paper, we ﬁnd that high loss samples are assigned to groups with high group weights and low marginal probabilities , characteristic of a worst-off group.
Revised: As indicated in the paper, we ﬁnd that high loss samples are assigned to groups with high group weights and low marginal probabilities , characteristic of a worst-off group.

Original: Marginal constraints form a key ingredient of our algorithm as per the above example.
Revised: Marginal constraints form a key ingredient of our algorithm as per the above example.

Original: Without the marginal constraints, the group assignments { ˆ g t } =
Revised: Without the marginal constraints, the group assignments { ˆ g t } =

Original: (cid:18) 1 0 1 0 1 0 (cid:19) .
Revised: (cid:18) 1 0 1 0 1 0 (cid:19) .

Original: That is, the assignments would have been made independent of the loss values and sparsely restricted to the group with large q j ¯ p j value.
Revised: That is, the assignments would have been made independent of the loss values and sparsely restricted to the group with large q j ¯ p j value.

Original: [Table]
Revised: [Table]

Original: A.6. Discussion on Unsupervised DRO methods
Revised: A.6. Discussion on Unsupervised DRO methods

Original: In this section
Revised: In this section

Original: [Table]
Revised: [Table]

Original: we contrast Worst-off DRO method against Unsup DRO (Hashimoto et al., 2018) and CVaR DRO (Levy et al., 2020).
Revised: we contrast Worst-off DRO method against Unsup DRO (Hashimoto et al., 2018) and CVaR DRO (Levy et al., 2020).

Original: CVaR DRO (Levy et al., 2020) is a coherent risk measure (Rahimian & Mehrotra, 2019) that optimizes over a certain ﬁxed-sized sub-populations within the training dataset.
Revised: CVaR DRO (Levy et al., 2020) is a coherent risk measure (Rahimian & Mehrotra, 2019) that optimizes over a certain ﬁxed-sized sub-populations within the training dataset.

Original: In essence, CVaR DRO is alike Unsup DRO where the size of the sub-population is controlled by a threshold on the loss value.
Revised: In essence, CVaR DRO is alike Unsup DRO where the size of the sub-population is controlled by a threshold on the loss value.

Original: In both CVaR DRO and Unsup DRO, the size of the selected sub-population needs to be close to the size of the smallest group as identiﬁed in Section 3.2.2 of (Liu et al., 2021).
Revised: In both CVaR DRO and Unsup DRO, the size of the selected sub-population needs to be close to the size of the smallest group as identiﬁed in Section 3.2.2 of (Liu et al., 2021).

Original: Such a requirement demands wider hyper-parameter search space for α/η parameters that control the size of the sub-populations.
Revised: Such a requirement demands wider hyper-parameter search space for α/η parameters that control the size of the sub-populations.

Original: Our experiments justify this need, Table 2 of Appendix A.8 shows that the search space of Unsup DRO is twice relative to Worst-off DRO in order to attain comparable average group accuracies.
Revised: Our experiments justify this need, Table 2 of Appendix A.8 shows that the search space of Unsup DRO is twice relative to Worst-off DRO in order to attain comparable average group accuracies.

Original: Clearly, a wider search space contributes to a harder model selection procedure.
Revised: Clearly, a wider search space contributes to a harder model selection procedure.

Original: Moreover, scenarios where extensive search is not possible (eg, small validation set/dataset regimes) could result in incorrect/unstable model selection.
Revised: Moreover, scenarios where extensive search is not possible (eg, small validation set/dataset regimes) could result in incorrect/unstable model selection.

Original: From the perspective of the methodology,
Revised: From the perspective of the methodology,

Original: CVaR DRO / Unsup DRO train only on the highest loss samples while discarding the remaining samples.
Revised: CVaR DRO / Unsup DRO train only on the highest loss samples while discarding the remaining samples.

Original: In contrast,
Revised: In contrast,

Original: Worst-off DRO does not discard any sample rather downweights/upweights as per the worst-off group assignment.
Revised: Worst-off DRO does not discard any sample rather downweights/upweights as per the worst-off group assignment.

Original: This property aids in maintaining a high overall accuracy besides reaching good minority group accuracy.
Revised: This property aids in maintaining a high overall accuracy besides reaching good minority group accuracy.

Original: A.7. Discussion on MAR case
Revised: A.7. Discussion on MAR case

Original: The δ gap in Lemma 2.3 captures the error in misspeciﬁcation of ¯ p in relation to p (cid:63) .
Revised: The δ gap in Lemma 2.3 captures the error in misspeciﬁcation of ¯ p in relation to p (cid:63) .

Original: When ¯ p is misspeciﬁed due to the data being Missing at Random (MAR) rather than MCAR (Missing Completely at Random), a solution could be to estimate the propensity of missingness from other features; then use inverse propensity weighting to get a consistent estimate of the fraction of samples in each group as discussed in (Zhao & Udell, 2020).
Revised: When ¯ p is misspeciﬁed due to the data being Missing at Random (MAR) rather than MCAR (Missing Completely at Random), a solution could be to estimate the propensity of missingness from other features; then use inverse propensity weighting to get a consistent estimate of the fraction of samples in each group as discussed in (Zhao & Udell, 2020).

Original: Alternatively, if provided with the knowledge of the data-generation process, the core effort in extending our method simply involves using off-the-shelf estimators to characterize the probability distributions (see (Mohan & Pearl, 2014) for example.
Revised: Alternatively, if provided with the knowledge of the data-generation process, the core effort in extending our method simply involves using off-the-shelf estimators to characterize the probability distributions (see (Mohan & Pearl, 2014) for example.

Original: import cvxpy as cp import numpy as np class Solver(object): def __init__(self, n_controls, bsize, marginals, epsilon, labeled=
Revised: import cvxpy as cp import numpy as np class Solver(object): def __init__(self, n_controls, bsize, marginals, epsilon, labeled=

Original: None): """Group assignment solver.
Revised: None): """Group assignment solver.

Original: Arguments:n_controls: An integer for the number of groups.
Revised: Arguments:n_controls: An integer for the number of groups.

Original: bsize: An integer for the batch size.
Revised: bsize: An integer for the batch size.

Original: marginals: A 2D array for the marginal distribution.
Revised: marginals: A 2D array for the marginal distribution.

Original: epsilon: A float for the variance.
Revised: epsilon: A float for the variance.

Original: labeled: A tuple for labeled data indices and their value. """ self.
Revised: labeled: A tuple for labeled data indices and their value. """ self.

Original: X =
Revised: X =

Original: cp.
Revised: cp.

Original: Variable((bsize, n_controls)) self.l
Revised: Variable((bsize, n_controls)) self.l

Original: = cp.
Revised: = cp.

Original: Parameter((bsize, 1)) self.p
Revised: Parameter((bsize, 1)) self.p

Original: = cp.Parameter((n_controls, 1), value=marginals) self.q = cp.
Revised: = cp.Parameter((n_controls, 1), value=marginals) self.q = cp.

Original: Parameter(n_controls) if labeled is not None: labeled_idx, labeled_value = labeled counts = cp.sum(self.
Revised: Parameter(n_controls) if labeled is not None: labeled_idx, labeled_value = labeled counts = cp.sum(self.

Original: X, axis=0, keepdims=
Revised: X, axis=0, keepdims=

Original: True)
Revised: True)

Original: obj
Revised: obj

Original: =
Revised: =

Original: ((self.l.
Revised: ((self.l.

Original: T @ self.
Revised: T @ self.

Original: X) / self.p.T) @ self.q constraints =
Revised: X) / self.p.T) @ self.q constraints =

Original: [self.
Revised: [self.

Original: X >= 0, cp.sum(self.
Revised: X >= 0, cp.sum(self.

Original: X, axis=1, keepdims=
Revised: X, axis=1, keepdims=

Original: True) =
Revised: True) =

Original: =
Revised: =

Original: np.ones((bsize, 1)), cp.abs(cp.sum(self.
Revised: np.ones((bsize, 1)), cp.abs(cp.sum(self.

Original: X, axis=0, keepdims=
Revised: X, axis=0, keepdims=

Original: True) / bsize - self.p.
Revised: True) / bsize - self.p.

Original: T) <
Revised: T) <

Original: = epsilon] if labeled is not None: constraints += [self.
Revised: = epsilon] if labeled is not None: constraints += [self.

Original: X[labeled_idx]
Revised: X[labeled_idx]

Original: ==
Revised: ==

Original: labeled_value] self.prob =
Revised: labeled_value] self.prob =

Original: cp.
Revised: cp.

Original: Problem(cp.
Revised: Problem(cp.

Original: Maximize(obj), constraints) def cvxsolve(self, losses, weights): """Solver.
Revised: Maximize(obj), constraints) def cvxsolve(self, losses, weights): """Solver.

Original: Arguments:losses: A 2D array for loss values.
Revised: Arguments:losses: A 2D array for loss values.

Original: weights: A 1D array for group weights q.
Revised: weights: A 1D array for group weights q.

Original: Returns:A2Darray for soft group assignments. "
Revised: Returns:A2Darray for soft group assignments. "

Original: "" self.l.value = losses self.q.value = weights self.prob.solve()returnself.
Revised: "" self.l.value = losses self.q.value = weights self.prob.solve()returnself.

Original: X.value
Revised: X.value

Original: A.8. Hyper-parameter Tuning
Revised: A.8. Hyper-parameter Tuning

Original: Hyper-parameters were selected for each algorithm by performing an NVP procedure (see Section 3).
Revised: Hyper-parameters were selected for each algorithm by performing an NVP procedure (see Section 3).

Original: The best performing model was identiﬁed on the validation set associated with each dataset.
Revised: The best performing model was identiﬁed on the validation set associated with each dataset.

Original: All the measures were computed and averaged over three random runs.
Revised: All the measures were computed and averaged over three random runs.

Original: A list of all the hyper-parameters that were tuned for are available in Table 2.
Revised: A list of all the hyper-parameters that were tuned for are available in Table 2.

Original: The ﬁnal hyper-parameters selected for each method can be viewed from Table 3.
Revised: The ﬁnal hyper-parameters selected for each method can be viewed from Table 3.

Original: A.9. Additional ablation experiments
Revised: A.9. Additional ablation experiments

Original: In this section, we extend the ablation experiments presented in the main paper.
Revised: In this section, we extend the ablation experiments presented in the main paper.

Original: Speciﬁcally, results over all four datasets is provided in Figures 3,4,5.
Revised: Speciﬁcally, results over all four datasets is provided in Figures 3,4,5.

Original: As we see in the plots, all the datasets exhibit similar behaviour across the ablation experiments.
Revised: As we see in the plots, all the datasets exhibit similar behaviour across the ablation experiments.

Original: A.10. Ablation study on increasing the constraint set size.
Revised: A.10. Ablation study on increasing the constraint set size.

Original: We conduct experiments on Worst-off DRO method for different values of the (cid:15) parameter in the set { 0 , 0 .
Revised: We conduct experiments on Worst-off DRO method for different values of the (cid:15) parameter in the set { 0 , 0 .

Original: The test set accuracies on the minority group and average group are reported in Figure 6.
Revised: The test set accuracies on the minority group and average group are reported in Figure 6.

Original: Increasing the (cid:15) value also increases the constraint set size because the marginal constraint is gradually relaxed.
Revised: Increasing the (cid:15) value also increases the constraint set size because the marginal constraint is gradually relaxed.

Original: Figure 6 shows that the both minority group accuracy and average group accuracy values reduce with increase in (cid:15) value beyond 0 .
Revised: Figure 6 shows that the both minority group accuracy and average group accuracy values reduce with increase in (cid:15) value beyond 0 .

Original: 1 threshold.
Revised: 1 threshold.

Original: The accuracy values for (cid:15) ≤ 0 .
Revised: The accuracy values for (cid:15) ≤ 0 .

Original: 01 are comparable.
Revised: 01 are comparable.

Original: A similar trend holds on other datasets as well.
Revised: A similar trend holds on other datasets as well.

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Table]
Revised: [Table]

Original: A.11. More details on the datasets A.11.1. W ATERBIRDS
Revised: A.11. More details on the datasets A.11.1. W ATERBIRDS

Original: The dateset, used in (Sagawa et al., 2019), comprises of 4795 images of birds from the CUB dataset (Welinder et al., 2010) and the backgrounds taken from the Places dataset (Zhou et al., 2017).
Revised: The dateset, used in (Sagawa et al., 2019), comprises of 4795 images of birds from the CUB dataset (Welinder et al., 2010) and the backgrounds taken from the Places dataset (Zhou et al., 2017).

Original: Each image in the dataset has a background of land or water.
Revised: Each image in the dataset has a background of land or water.

Original: The target labels are either “landbirds” or “waterbirds”.
Revised: The target labels are either “landbirds” or “waterbirds”.

Original: In this dataset the groups “landbirds” on water and “waterbirds” on land form a minority.
Revised: In this dataset the groups “landbirds” on water and “waterbirds” on land form a minority.

Original: A ResNet50 model, pre-trained with ImageNet weights, has been used for training in experiments on this dataset.
Revised: A ResNet50 model, pre-trained with ImageNet weights, has been used for training in experiments on this dataset.

Original: No data augmentation has been applied for any of the Algorithms.
Revised: No data augmentation has been applied for any of the Algorithms.

Original: [Figure]
Revised: [Figure]

Original: A.11.2. G ROUP CMNIST
Revised: A.11.2. G ROUP CMNIST

Original: CMNIST, derived from an MNIST (LeCun et al., 1998), is a digit recognition dataset where each image is colored either red or green.
Revised: CMNIST, derived from an MNIST (LeCun et al., 1998), is a digit recognition dataset where each image is colored either red or green.

Original: Digits < 5 / ≥ 5 are considered as label 0 / 1 .
Revised: Digits < 5 / ≥ 5 are considered as label 0 / 1 .

Original: We consider three groups in our experiments.
Revised: We consider three groups in our experiments.

Original: In the ﬁrst two groups, label 0 images are predominantly colored red and vice versa.
Revised: In the ﬁrst two groups, label 0 images are predominantly colored red and vice versa.

Original: In the third group, which forms a minority, we switch coloring such that the label 1 images are predominantly colored red.
Revised: In the third group, which forms a minority, we switch coloring such that the label 1 images are predominantly colored red.

Original: Speciﬁcally, for the ﬁrst two groups, the color id is sampled by ﬂipping the target label with probabilities 0 .
Revised: Speciﬁcally, for the ﬁrst two groups, the color id is sampled by ﬂipping the target label with probabilities 0 .

Original: 2 and 0 .
Revised: 2 and 0 .

Original: 1 respectively, while the third group with probability 0 .
Revised: 1 respectively, while the third group with probability 0 .

Original: Both training and testing sets contain three groups.
Revised: Both training and testing sets contain three groups.

Original: The overall setup for generating a given group is similar to (Arjovsky et al., 2019).
Revised: The overall setup for generating a given group is similar to (Arjovsky et al., 2019).

Original: However, unlike (Arjovsky et al., 2019) the training and testing phases share the set of three pre-deﬁned groups so as to evaluate for group-robustness criterion.
Revised: However, unlike (Arjovsky et al., 2019) the training and testing phases share the set of three pre-deﬁned groups so as to evaluate for group-robustness criterion.

Original: A.11.3. G ROUP A DULT
Revised: A.11.3. G ROUP A DULT

Original: We use a semi-synthetic version of the Adult dataset (Dua et al., 2017) for this experiment.
Revised: We use a semi-synthetic version of the Adult dataset (Dua et al., 2017) for this experiment.

Original: Similar to (Lahoti et al., 2020), we consider race and sex as the four demographic groups.
Revised: Similar to (Lahoti et al., 2020), we consider race and sex as the four demographic groups.

Original: The target label is income > 50 K $ .
Revised: The target label is income > 50 K $ .

Original: Similar to the CMNIST dataset, each group has a different correlation strength to the target label.
Revised: Similar to the CMNIST dataset, each group has a different correlation strength to the target label.

Original: For the purposes of the experiment, we exaggerate these spurious correlations caused by group membership close to (Creager et al., 2021).
Revised: For the purposes of the experiment, we exaggerate these spurious correlations caused by group membership close to (Creager et al., 2021).

Original: However, distinct from (Creager et al., 2021), the training and testing phases share same set of groups.
Revised: However, distinct from (Creager et al., 2021), the training and testing phases share same set of groups.

Original: For samples with group label as Aﬁcan-American, we undersample examples with probability P
Revised: For samples with group label as Aﬁcan-American, we undersample examples with probability P

Original: ( y = 1 | group )
Revised: ( y = 1 | group )

Original: = 0 .
Revised: = 0 .

Original: 06 whereas for the non African-American group labels, we oversample examples with probability P
Revised: 06 whereas for the non African-American group labels, we oversample examples with probability P

Original: ( y = 1 | group )
Revised: ( y = 1 | group )

Original: = 0 .
Revised: = 0 .

Original: A.11.4. C ELEB A
Revised: A.11.4. C ELEB A

Original: CelebA (Liu et al., 2015) is a dataset containing about 200k celebrity faces curated from the internet.
Revised: CelebA (Liu et al., 2015) is a dataset containing about 200k celebrity faces curated from the internet.

Original: Similar to (Sagawa et al., 2019), we aim to predict the target attribute Blond Hair that is spuriously correlated to the Gender attribute.
Revised: Similar to (Sagawa et al., 2019), we aim to predict the target attribute Blond Hair that is spuriously correlated to the Gender attribute.

Original: The minority group in this dataset are the images with attributes (blond, male).
Revised: The minority group in this dataset are the images with attributes (blond, male).

Original: For this dataset, the ofﬁcial train-val-test splits as recommended by (Liu et al., 2015) has been used.
Revised: For this dataset, the ofﬁcial train-val-test splits as recommended by (Liu et al., 2015) has been used.

Original: Similar to the Waterbirds experiments, a pre-trained ImageNet-based
Revised: Similar to the Waterbirds experiments, a pre-trained ImageNet-based

Original: ResNet50 model has been used for the implementations.
Revised: ResNet50 model has been used for the implementations.

Original: A.12. Related work
Revised: A.12. Related work

Original: Group Robust Optimization.
Revised: Group Robust Optimization.

Original: Methods in the literature handling robustness to extraneous attributes can be broadly categorized into two classes.
Revised: Methods in the literature handling robustness to extraneous attributes can be broadly categorized into two classes.

Original: The ﬁrst class, domain generalization methods (Arjovsky et al., 2019; Mahajan et al., 2021;
Revised: The ﬁrst class, domain generalization methods (Arjovsky et al., 2019; Mahajan et al., 2021;

Original: Moyer et al., 2018), aim at learning representations invariant to a predeﬁned set of extraneous attributes or groups.
Revised: Moyer et al., 2018), aim at learning representations invariant to a predeﬁned set of extraneous attributes or groups.

Original: The goal is to be able to generalize to unseen domains or environments in the testing phase.
Revised: The goal is to be able to generalize to unseen domains or environments in the testing phase.

Original: On the other hand, the second class of methods, called as the group robust methods (Oakden-Rayner et al., 2020; Sagawa et al., 2019; Liu et al., 2021), seek to improve the worst-off or the minority group performance within the set of pre-deﬁned groups.
Revised: On the other hand, the second class of methods, called as the group robust methods (Oakden-Rayner et al., 2020; Sagawa et al., 2019; Liu et al., 2021), seek to improve the worst-off or the minority group performance within the set of pre-deﬁned groups.

Original: Here the training and testing phases share the same set of groups.
Revised: Here the training and testing phases share the same set of groups.

Original: Our approach falls into the second class of methods.
Revised: Our approach falls into the second class of methods.

Original: Robust Optimization with Demographics.
Revised: Robust Optimization with Demographics.

Original: When group information is known at train time, Group DRO (Hu et al., 2018;
Revised: When group information is known at train time, Group DRO (Hu et al., 2018;

Original: Sagawa et al., 2019) or Invariant Risk Minimization (IRM) (Arjovsky et al., 2019) could be employed to improve the performance over multiple groups.
Revised: Sagawa et al., 2019) or Invariant Risk Minimization (IRM) (Arjovsky et al., 2019) could be employed to improve the performance over multiple groups.

Original: Speciﬁcally, 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.
Revised: Speciﬁcally, 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.

Original: Robust Optimization without Demographics.
Revised: Robust Optimization without Demographics.

Original: Several works have focused on developing methods that remove the dependence on the group labels.
Revised: Several works have focused on developing methods that remove the dependence on the group labels.

Original: (Hashimoto et al., 2018) proposed to minimize the loss of the samples with losses larger than a certain threshold.
Revised: (Hashimoto et al., 2018) proposed to minimize the loss of the samples with losses larger than a certain threshold.

Original: (Lahoti et al., 2020) developed an adversarial re-weighting scheme that assigns large weight to high loss samples.
Revised: (Lahoti et al., 2020) developed an adversarial re-weighting scheme that assigns large weight to high loss samples.

Original: Recently, (Liu et al., 2021) proposed a simple yet effective two-stage approach called Just-Train-Twice (JTT) that trains a model by upweighting samples with high losses from the initial ERM model.
Revised: Recently, (Liu et al., 2021) proposed a simple yet effective two-stage approach called Just-Train-Twice (JTT) that trains a model by upweighting samples with high losses from the initial ERM model.

Original: Lastly, EIIL (Creager et al., 2021) also proposed a two-stage method where in the ﬁrst stage the group or environment labels are inferred and in the second stage a Group DRO or an IRM optimization is employed on inferred labels.
Revised: Lastly, EIIL (Creager et al., 2021) also proposed a two-stage method where in the ﬁrst stage the group or environment labels are inferred and in the second stage a Group DRO or an IRM optimization is employed on inferred labels.

Original: Drawbacks of two-stage methods.
Revised: Drawbacks of two-stage methods.

Original: While two-stage methods, like JTT (Liu et al., 2021) and EIIL (Creager et al., 2021), have demonstrated signiﬁcant improvements in minority group accurcies, they bear few drawbacks in relation to single-stage methods.
Revised: While two-stage methods, like JTT (Liu et al., 2021) and EIIL (Creager et al., 2021), have demonstrated signiﬁcant improvements in minority group accurcies, they bear few drawbacks in relation to single-stage methods.

Original: Firstly, two-stage methods introduce additional set of hyper-parameters.
Revised: Firstly, two-stage methods introduce additional set of hyper-parameters.

Original: For example, it’s crucial for JTT to tune for the number of epochs in its ﬁrst stage.
Revised: For example, it’s crucial for JTT to tune for the number of epochs in its ﬁrst stage.

Original: Similarly, several hyper-parameters are introduced in EIIL method especially in optimizing the EI objective and in identifying a pre-trained reference model.
Revised: Similarly, several hyper-parameters are introduced in EIIL method especially in optimizing the EI objective and in identifying a pre-trained reference model.

Original: Secondly, in a two-stage model, a failed ﬁrst stage leads to an unsuccessful second stage.
Revised: Secondly, in a two-stage model, a failed ﬁrst stage leads to an unsuccessful second stage.

Original: This is because the errors from the ﬁrst stage are propagated to the later stages.
Revised: This is because the errors from the ﬁrst stage are propagated to the later stages.

Original: For example, a ﬁrst stage model could fail in JTT due to model overﬁtting, and likewise an inaccurate group inference in the
Revised: For example, a ﬁrst stage model could fail in JTT due to model overﬁtting, and likewise an inaccurate group inference in the

Original: EIIL method may block second-stage invariant learning besides raising ethical issues on pseudo-label misuse.
Revised: EIIL method may block second-stage invariant learning besides raising ethical issues on pseudo-label misuse.

Original: In summary, efforts to reduce a two-stage model to a single-stage method are beneﬁcial and, as we shall see shortly, our proposal ﬁts in class of single stage methods.
Revised: In summary, efforts to reduce a two-stage model to a single-stage method are beneﬁcial and, as we shall see shortly, our proposal ﬁts in class of single stage methods.