B ARACK : Partially Supervised Group Robustness With Guarantees
Abstract
While neural networks have shown remarkable success on classification tasks in terms of average-case performance, they often fail to perform well on certain groups of the data, for instance when spurious correlations are present.
Unfortunately, group information may be expensive to obtain; thus, recent works in robustness and fairness have proposed ways to improve worst-group performance even when group labels are unavailable.
However, these methods generally underperform methods that utilize group information at training time.
In this work, we assume access to a small number of group labels alongside a larger dataset without group labels.
We propose B ARACK , a simple two-step framework to utilize this partial group information to improve worst-group performance: train a model to predict the missing group labels for the training data, and then use these predicted group labels in a robust optimization objective.
Theoretically, we provide generalization bounds for our approach in terms of the worst-group performance, which scale with respect to both the total number of training points and the number of training points with group labels .
Empirically, across four spurious correlation and robustness benchmark tasks, our method outperforms the baselines that do not use group information, even when only 1-33% of points have group labels.
1. Introduction
On classification tasks, deep neural networks can often underperform on certain groups of the data.
For example, on datasets with spurious correlations , standard neural networks have been shown to achieve high average accuracy, yet drastically lower accuracy on groups that violate the
2022 by the author spurious correlation (Sagawa et al., 2019).
Similarly, when certain groups are underrepresented in the training data, models tend to perform poorly on these rare groups (Sohoni et al., 2020).
In many settings, such as applications where fairness or safety are important, this behavior is undesirable; for example, gender classification systems have been shown to underperform for non-white faces (Buolamwini & Gebru, 2018), and medical triage systems have been shown to miss certain abnormality subtypes (Oakden-Rayner et al., 2020).
To avoid this, we want to ensure group robustness , i.e., high accuracy on the worst-performing group.
Unfortunately, group annotations are often unavailable.
Many datasets only have labels for the task , not groups.
Group labels may also be expensive to obtain; for instance, in the common case where the group labels are finer-grained than the class labels, it may require higher annotation cost to obtain group labels than class labels (Gebru et al., 2017).
This makes ensuring group robustness more challenging.
Existing works to address the issue of group robustness fall into two main categories: those that assume access to the group labels for all of the training data, and those that assume no access to the group labels for the training data.
For instance, in the first category, Sagawa et al.
(2019) propose group distributionally robust optimization (GDRO), an efficient algorithm for minimizing the worst-group loss when the groups are known.
More recently, several approaches have been proposed to improve group robustness when group labels are unavailable .
A common approach is to first estimate the group labels, then train a robust classifier using these estimated group labels (Sohoni et al., 2020;
Liu et al., 2021).
However, in terms of worst-group performance, the methods that require group labels unsurprisingly (and often substantially) outperform those that do not.
A fundamental question is: can we close this gap if we have partial group information ?
Specifically, we seek to understand the intermediate regime in which group labels are available for a (small) subset of the training data, while the remainder has class labels only.
The distinction between this setting and the aforementioned prior work is akin to the difference between semi-supervised learning vs. supervised or unsupervised learning.
From an application standpoint, when the identities of the groups are known, it is often feasible to obtain group labels for a small subset of the data.
From a theoretical standpoint, the relative value of class vs. group labels for ensuring group robustness is still unknown.
To address this question, we propose B ARACK , 1 a simple two-stage approach to improve group robustness for the setting wherein group labels are only known for a subset of datapoints.
In the first stage of B ARACK , we use the available group labels to train a model to predict the group labels on the data without group annotations.
In the second stage, we use these predicted group labels in the GDRO objective (Sagawa et al., 2019) to train a robust model.
Theoretically, we show how worst-group performance scales with the number of total points and the number of points with group labels .
Empirically, we show that even if only a small fraction (1-33%) of points have group labels, B ARACK improves over approaches that do not use group labels.
Contributions.
In summary, our main contributions are: - We propose a simple framework, B ARACK , that can improve group robustness with only a small number of group labels: we train a model to predict the missing group labels, then use these group labels in a robust training objective.
- On four benchmark tasks, we show that with only 1-33% of data having group labels, our method empirically outperforms baselines that do not use group information, approaching the performance of GDRO trained on the full dataset.
- We prove a generalization bound on the worst-group performance of our method, showing that it scales with the inverse square root of the total number of points with group labels in the smallest group.
(In Appendix C, we show how to tighten our bound under additional assumptions.)
- We conduct ablation studies to better understand the importance of the different components of B ARACK .
2. Background 2.1. Problem Setup
We consider a similar setting to Sagawa et al.
(2019): we have n training points sampled IID from a distribution P : { ( x i , y i , z i ) } ni =1 ∈ X ×Y ×G .
x i denotes a datapoint, y i its class label, and z i its group label.
However, unlike Sagawa et al.
(2019), we do not assume we know all the z
i ’s; we only assume knowledge of z 1 , . . .
, z m , where m < n .
We denote the group-labeled subset as D 1 :=
{ ( x i , y i , z i ) } mi =1 , and D 2 :
= { ( x i , y i ) } ni = m
+1 as the group-unlabeled subset.
Note that we assume D 1 and D 2 are samples from the same distribution P ; but the z i ’s are unobserved on D 2 .
Our end goal is to maximize the worst-group accuracy on the task of predicting the correct class label for each datapoint.
In other words, given a function class of classifiers F (where each f ∈ F is a function X → ∆ |Y| , i.e., a function that outputs probabilities for each class), we wish to find f ∈ F that maximizes min g ∈G E ( x,y ) | z
[Equation]
(argmax
[Equation]
In practice, we instead seek the f ∈ F that minimizes the worst-group loss over the training data:
[Equation]
When the z i ’s are known, the latter problem can be solved with group DRO (GDRO) (Sagawa et al., 2019).
GDRO is a stochastic optimization method designed for minimax problems of exactly the form of Eq.
However, in our setting, solving this problem is challenging because we only know a subset of the z
i ’s, so we cannot compute Eq.
directly.
3. Method
To address the problem of improving group robustness when only some group labels are available, we propose B ARACK , a two-stage framework which leverages the group-labeled examples to generate group “pseudolabels” for the remaining datapoints, and then uses these pseudolabels to train a robust model on the target task.
This two-stage approach is inspired by prior works such as J TT and G EORGE (Liu et al., 2021;
Sohoni et al., 2020).
However, unlike these methods which assume all group labels are unknown, B ARACK is capable of exploiting the additional information in the group labels that are known for some datapoints.
Pseudocode for
B ARACK ’s overall workflow is given in Alg.
1 (App. B).
Stage 1: Predicting group labels (via a “class-conditional” classifier).
First, we train a model f group ∈ F 1 to predict the group labels for the training and validation datapoints that do not have provided group labels.
To do so, we train a supervised classifier on the training points with known group labels.
Despite the small number of these points, we show that this simple approach can perform surprisingly well with a key modification: we use the class label (assumed known for all train data) as an input to the group classifier, as the probabilities of each group can vary conditional on the class.
Specifically, we compute the empirical probabilities and corresponding logits of each group conditioned on the class.
For each example fed into the group classifier, these logits for its class are summed with the output of the last layer of the network.
In the datasets we evaluate on, the groups are subsets of the classes, so this means that we take the softmax over the logits output by the network over all groups belonging to the known class to get the predicted per-group probabilities (assigning 0 probability to groups in different classes).
In this way, the class information helps the model learn to distinguish the groups, offsetting the dearth of data.
We train f group with GDRO to encourage good performance at predicting each group.
To select the best group classifier model over the course of training, we use a group-labeled subset of the validation set with the same size as the group-labeled training set (ensuring that the total number of group labels required is small across both training and validation splits).
This group classifier is then used to generate “pseudo-group-labels” for all training datapoints without a known group label.
We term this approach (together with Stage 2) B ARACK -Base; where unspecified, B ARACK refers to B ARACK -Base.
We find that this simple approach works well, and thus focus on it for the main paper.
A more complex method can also be used in Stage 1: for instance, one could use semi-supervised learning to leverage the group-unlabeled points to train an improved group classifier.
We term this B ARACK -SSL; our preliminary experiments in Appendix D.4 and analysis in Appendix C show that this can indeed further improve performance, at the cost of added complexity and runtime.
Stage 2: Training a robust model.
Intuitively, if the predicted group labels ˆ z i from Stage 1 align well with the (unobserved) true group labels z
i , then training a model to be robust with respect to the predicted groups should also induce good robustness with respect to the true groups.
Thus, we train a model f robust ∈ F 2 on the original task using
GDRO (Sagawa et al., 2019), with the groups defined by the predicted group labels from Stage 1 (except for datapoints with known ground-truth group labels, for which we use this ground-truth).
The same small group-labeled validation subset as in Stage 1 is used to validate this model.
4. Analysis
We now analyze the theoretical worst-group performance of B ARACK .
For brevity, much of the theory is deferred to
Appendix C.
For notation, we use F 1 , F 2 to denote the spaces of possible classifiers used in Stage 1 and 2 respectively.
We assume
F 2 is parametrized by θ ∈ R d , so the B ARACK model ˆ f robust has parameters ˆ θ robust .
We define L robust ( θ ) :
= max g ∈G E ( x,y ) | z = g [ ℓ ( f ( x ; θ ) , y )]
(the worst-group loss), and
L ∗ robust :=
min θ ∈ R
d
L robust ( θ ) .
Finally, q is the population proportion of the rarest group.
We assume ℓ is either the squared loss between probabilities, or the truncated crossentropy loss, so that ℓ is bounded and Lipschitz.
First, we relate the performance of the group classifier in the first stage to the excess worst-group risk of the end model.
Theorem 4.1.
Suppose that on each group, the error rate of the group classifier from “Stage 1” is ≤ r .
Then with high probability, L robust (ˆ θ robust )
≤ L ∗ robust + ˜ O (cid:16) rq + 1 √ qn (cid:17) .
Theorem 4.1 says that the excess worst-group risk scales linearly in the error rate of the group classifier, plus an additional O ( 1 √ qn ) term which is small if the total number of datapoints is large.
In particular, if we use standard learning-theoretic results to bound the error rate of the group classifier, then under the assumption that the group classification problem is realizable , we obtain the following corollary.
(Realizability means there exists f ∗ ∈ F 1 such E ( x,y,z ) ∼P
[ ℓ ( f ( x, y ) , z )]
= 0 .
We show how to relax the realizability assumption in Theorem C.8 (Appendix C).
Corollary 4.2.
With high probability, for B ARACK -Base we have L robust (ˆ θ robust )
≤ L ∗ robust + ˜ O (cid:16) 1 q √ m (cid:17) .
A strength of Theorem 4.1 and Corollary 4.2 are that they do not require assumptions on the data (unlike prior work such as (Sohoni et al., 2020) which requires specific distributional assumptions to obtain generalization bounds).
We discuss how to further improve our bounds in Appendix C.
5. Experiments
We empirically validate that B ARACK improves group robustness on four different image classification tasks.
In
Section 5.2, we study how the worst-group performance of B ARACK scales with the number of group labels, and compare it with several baselines.
We show that with as few as 1-33% of points having group labels, B ARACK attains better worst-group performance than the baselines that do not use this group information, and approaches the worst-group performance of GDRO on the full dataset as the number of group labels increases.
We also show that B ARACK outperforms GDRO trained on only the subset of points with group labels.
In Section 5.3, we confirm that the worst-group accuracy of the final model increases with the accuracy of stage 1 of B ARACK .
In Appendix D, we present extensive ablations to better understand the impact of B ARACK ’s design choices.
In our experiments, we study how performance varies as we increase the number of group-labeled examples per group .
In other words, we pick a fixed budget of (training and validation) examples to label for each group, sampled randomly from the appropriate groups in the original dataset.
5.1. Datasets / Tasks
We evaluate on four image classification tasks: two spurious correlation tasks, Waterbirds and CelebA (Sagawa et al., 2019), and two tasks with a rare group , U-MNIST (Sohoni et al., 2020) and U-CIFAR10.
Further dataset details are in
Appendix D.1. [Figure] [Table]
5.2. Results: Worst-Group Performance
Across the four tasks in Section 5.1, B ARACK matches or improves worst-group accuracy compared to baselines that don’t use group information; results are in Table 1 & Figure 1.
The baselines we study are ERM, G EORGE
(Sohoni et al., 2020) (which use no group information), and E IIL (Creager et al., 2021) and J TT (Liu et al., 2021) (which use group information on the val.
set).
Additional results are in
Appendix D.
With as few as 1%-33% of training datapoints having group labels, B ARACK improves over these baselines (Figure 1).
For instance, CelebA has 1387 training points in the smallest group, and B ARACK outperforms the baselines on CelebA with 16 group-labeled training examples per group.
As the number of group-labeled points increases, the worst-group performance of B ARACK approaches that of full GDRO.
In terms of average accuracy, B ARACK and full-dataset
GDRO are typically similar, while ERM is usually higher.
This is expected, since on these tasks there is a tradeoff between optimizing for average-case and worst-case performance, as observed in the literature (Sagawa et al., 2019).
B ARACK also substantially outperforms SubsetGDRO (GDRO trained only on the subset of group-labeled points) on both metrics.
Subset-GDRO fails to generalize well due to the limited amount of training data it uses.
5.3. Results: Group Prediction Accuracy
In this section, we study the performance of Stage 1 of
B ARACK (accuracy at predicting group labels), in order to better understand the performance of Stage 2 (worst-group accuracy on the target classification task).
In Table 2, we report the Stage 1 group prediction accuracies per dataset.
Table 2 shows that the group prediction models are far from perfect.
Indeed, on U-CIFAR10 the accuracy at predicting the ‘airplane’ group is below 70%, and the average accuracy [Table] over all groups is only ≈ 84% .
Nevertheless, this only causes a modest drop in performance for the final robust model, as seen in Table 1.
Acknowledgments
NS performed this work during an internship with Facebook
AI (now Meta AI)
A. Related Work
Beyond average performance.
Our work primarily builds on prior work in the area of group robustness.
This line of work has a long history in the literature; for example, see (Mohri et al., 2019; Zhang et al., 2021) and references therein.
Several methods have been proposed to improve group robustness when the group labels are known at train time.
In our algorithms, we focus on GDRO (Sagawa et al., 2019), a stochastic algorithm for minimizing the worst-group loss, which has been shown to substantially improve worst-group performance on tasks with spurious correlations.
We use GDRO as a component of our method.
While we mainly focus on worst-group accuracy, average and worst-group accuracy are not the only measures of performance of interest to ML practitioners.
Several other works have used different approaches to strike a balance between important performance measures, for example through the lenses of distributional robustness (Duchi et al., 2019;
Wang et al., 2020; Zhang et al., 2021; Ben-Tal et al., 2013), fairness (Hardt et al., 2016; Agarwal et al., 2018;
Li et al., 2019; 2020), or outlier/noisy sample detection (Huber, 1992; Bhatia et al., 2015; Menon et al., 2019; Li et al., 2020).
Balashankar et al.
and Martinez et al.
(2020) propose different methods for ensuring group Pareto fairness , which seeks to find Pareto-efficient solutions in terms of the accuracies on each group.
Group robustness without group labels.
When the group labels are not known, alternative methods exist that still attempt to improve group robustness and robustness to spurious correlations.
Several of these works aim to first estimate the group labels, then train a robust classifier using these estimated group labels (Sohoni et al., 2020; Nam et al., 2020; Liu et al., 2021;
Zhang et al., 2022; Ahmed et al., 2021).
Others make no assumptions on the structure of the groups, and simply try to perform well on “all possible” data subsets above a specified size (Levy et al., 2020; Martinez et al., 2021).
Other categories of methods to improve group robustness without group labels include methods that aim to learn representations invariant to spurious features (similar to the ideas behind Invariant Risk Minimization (Arjovsky et al., 2019), but without having pre-provided environment labels), such as (Creager et al., 2021;
Goldstein et al., 2022; Ahmed et al., 2021); and
Pareto-fairness inspired methods, such as (Lahoti et al., 2020; Martinez et al., 2021).
Unsurprisingly, all of these approaches typically underperform methods that do utilize ground-truth group labels.
Our method, B ARACK , involves training two models sequentially.
This is commonly used in different ways as an approach to increasing model robustness in the literature (Yaghoobzadeh et al., 2021; Utama et al., 2021).
Among these works, (Liu et al., 2021; Goel et al., 2020; Creager et al., 2021;
Nam et al., 2020; Sohoni et al., 2020;
Zhang et al., 2022) are most relevant to our work, where a model is trained first and then the outputs of this model are used in some manner (for example, in the GDRO objective as done in (Sohoni et al., 2020)) to train the second model to be robust.
Our key point of difference is that none of these works are designed to actually utilize possible group labels when they are known for some samples.
B ARACK can yield superior performance to these methods by utilizing such additional group information, even if it is limited.
Semi-supervised learning.
Semi-supervised learning (SSL) is a rich field with several recent developments.
In the appendix, we explore the use of FixMatch (Sohn et al., 2020), a recent state-of-the-art method for SSL, for group classification.
Other recent successful approaches to semi-supervised learning involve learning self-supervised representations (without using labels), and then using the labeled examples for fine-tuning; examples include (Xie et al., 2020a;
Chen et al., 2020;
Caron et al., 2020).
Our work also has connections to self-training.
In standard self-training, a labeled dataset is used to train a model to generate pseudolabels for a separate unlabeled dataset; the labeled and pseudolabeled data are then used together to train a downstream model (Zoph et al., 2020; Xie et al., 2020b; Lee et al., 2013; Rosenberg et al., 2005).
In our work, we instead generate pseudolabels for the task of classifying the groups , which are then used to train a robust model for the original task. Concurrent work.
We would also like to acknowledge the following important concurrent work: Spread Spurious
Attribute (SSA) (Nam et al., 2021), which considers a very similar problem to ours: specifically, they consider the problem of group robustness when there are spurious attributes that are known for a subset of the training data.
Their proposed algorithm can be viewed as a special case of our general two-step framework (B ARACK ), in which they use semi-supervised learning techniques for the group prediction stage. Compared to (Nam et al., 2021), our work is more focused on understanding and analyzing the effectiveness of B ARACK through theoretical and ablation analyses. We believe that these analyses provide valuable insight regardless of the precise method used to estimate group labels. As we show in our experiments and analysis, while the basic version of B ARACK performs quite well, better training methods for the first (group classification) stage, such as with SSL, can translate to better results (possibly at the cost of increased computational complexity).
(Lokhande et al., 2022) also address a similar version of our partial group robustness problem, using a different approach based on minimizing an upper bound to the GDRO loss.
However, unlike our work (and that of (Nam et al., 2021)), they avoid any estimation of missing group labels due to privacy considerations, which unfortunately results in substantially lower worst-group performance (albeit still better than ERM) due to their upper bound function possibly being quite loose.
B. Algorithm Pseudocodes
On the following pages, we include algorithm boxes for B ARACK , as well as the subroutines used therein.
We also provide a schematic illustrating B ARACK (Figure 7) for explanatory purposes.
Additionally, to demonstrate the flexibility of our B ARACK framework, we present Algorithm 5 (B ARACK -SSL), an instantiation of the general B ARACK framework (Algorithm 1) which uses semi-supervised learning to train the group classifier.
Our preliminary empirical results on B ARACK -SSL (Appendix D.4) and theoretical analysis of B ARACK -SSL (Appendix C.10) suggest that B ARACK -SSL is a promising approach to further improve performance in our setting.
Algorithm 1 B ARACK (General)
Require: Group-labeled data D 1 = { ( x i , y i , z i ) } mi =1 , group-unlabeled data D 2 = { ( x i , y i ) } ni
= m +1 .
1: ˆ f group ← T RAIN ( D 1 , D 2 ) 2: (ˆ z m +1 , . . .
, ˆ z n ) ← P REDICT ( ˆ f group , D 2 ) 3: (ˆ z 1 , . . .
, ˆ z m ) ← ( z 1 , . . .
z m ) 4: ˆ f robust ← T RAIN R OBUST ( { ( x i , y i , ˆ z i ) } ni =1 ) output : Final model
ˆ f robust .
Algorithm 2).
Require: Group-labeled data D 1 = { ( x i , y i , z i ) } mi =1 , group-labeled validation data D val 1: Initialize model f θ ∈ F2:
Initialize acc best =3: for epoch in 1 , . . .
, T do 4: for i in 1 , . . .
, m do 5: θ ← GDRO Update ( ℓ, f θ ( x i , y i ) , z i ) 6: end for 7: if ValidationSubsetWorstGroupAcc( f θ ) >
acc best then 8: ˆ f group ← f θ 9: acc best ← ValidationSubsetWorstGroupAcc( f θ ) 10: end if 11: end for output
Require: Group prediction model ˆ f group , group-unlabeled training data D 2 = { ( x i , y i ) } ni
= m +1 .
1: for i = 1 , . . .
, m do 2: ˆ z i ← argmax g ∈G ˆ f group ( x i , y i ) 3: end for output : { ˆ z i } mi =
Require: Group-labeled data D 1 = { ( x i , y i , z i ) } mi =1 , group-
C. Theoretical Analysis and Proofs C.1. Overview
First, we state an upper bound on the worst-group loss of GDRO, and a lower bound on that of ERM.
Next, we restate a generalization bound on the worst-group loss of B ARACK (Theorem 4.1), and discuss how this result relates to the GDRO and ERM bounds.
Proofs are provided in the following subsections.
First, we show that if GDRO is trained on the dataset of the m group-labeled points, the excess worst-group risk of the resulting model (compared to the worst-group-optimal model, i.e. the model in F 2 with the lowest population worst-group performance on the task) scales as ˜ O
(1 / √ qm ) .
Lemma C.1.
Let ˆ f gdro,labeled ∈ F 2 (with associated parameters ˆ θ gdro,labeled ) denote the GDRO classifier trained on Donly.
Then with high probability,
L
[Equation]
By contrast, ERM (even trained on the full dataset) can result in a worst-group risk multiple times higher than that of the optimal worst-group model:
Lemma C.2.
Let ˆ f erm,full ∈ F 2 (with associated parameters ˆ θ erm,full ) denote the ERM classifier trained on the full dataset.
There exists a distribution P such that with high probability, L robust ( ˆ θ erm,full ) ≥ log(1 /q ) log( |Y| )
L ∗ robust − ˜ O
(1 / √ n ) .
For comparison, recall our bound (Theorem 4.1) on the generalization risk of B ARACK based on the group classifier accuracy, and the resulting Corollary 4.2 under the assumption of realizability of the group classification problem:
Theorem 4.1.
Suppose that on each group, the error rate of the group classifier from “Stage 1” is ≤ r .
Then with high probability, L robust (ˆ θ robust )
≤ L ∗ robust + ˜ O (cid:16) rq + 1 √ qn (cid:17) .
Corollary C.3.
With high probability, for B ARACK -Base we have L robust (ˆ θ robust )
≤ L ∗ robust + ˜ O (cid:16) 1 q √ m (cid:17) .
(Note: Realizability of the group classifier is not a fundamentally necessary assumption, and can be relaxed; in Appendix
C.6 we provide Theorem C.8, an analogue to Theorem 4.1 with an improved bound in the non-realizable case.
For ease of exposition, we defer further description of Theorem C.8 to Appendix C.6.)
As mentioned, Theorem 4.1 and Corollary 4.2 do not require data distribution assumptions.
However, the downside of
Theorem 4.1 is that the bound is relatively weak unless the group classifier is known to perform near-optimally; indeed,
Corollary 4.2 yields a slightly weaker asymptotic bound than GDRO on the labeled data alone.
This result can be improved when further assumptions are made to guarantee a stronger bound on the group classification error.
For instance, in Corollary C.13 (Appendix C.10), we show that the excess worst-group generalization error rate bound can be improved to O (1 / √ n ) when semi-supervised learning is used to train the group classifier.
With an appropriate choice of semi-supervised learning method (such as FixMatch (Sohn et al., 2020)), if m = Ω( √ n ) , then under appropriate conditions, the worst-group generalization error rate bound of B ARACK -SSL is O (1 / √ n ) .
This result is based on the PAC-learning based results of (Balcan & Blum, 2009), and requires the assumptions therein as well as realizability (the exact conditions are somewhat technical, and are deferred to the discussion in Appendix C.10).
This translates to an excess robust risk of
O (1 / √ n ) for the final model, as stated in Corollary C.13.
Next, another natural question is how the performance of B ARACK -Base compares to that of ERM.
We study this in
Corollary C.4. Corollary C.4.
In addition to the assumptions of Theorem 4.1, suppose that for all f ∈ F 2 , ℓ ( f ( x ) , y ) ⊥ argmax { ˆ f group ( x, y ) } (cid:12)(cid:12) z .
Let θ ∗ avg be the minimizer of the population average loss E ( x,y ) ∼ P [ ℓ ( x, y ; θ )] .
Then for B ARACKBase we have L robust
[Equation]
ˆ θ robust [Equation] For comparison, the robust loss of the ERM model is upper bounded by L robust ( θ ∗ avg ) + ˜ O
( 1 √ qn ) with high probability.
Corollary C.4 says that if we can assume that the errors made by f group are “random” conditioned on the true group identity—i.e., they do not affect the distribution of the loss on the target task—then in addition to the bound of Theorem 4.1, we can also guarantee that the worst-group loss is at least as good as that of the ERM model (plus
small ˜ O
(1 / √ qn ) noise).
Of course, this “random error” assumption is very strong; nevertheless, in Appendix D.3.1 we compare the performance of
B ARACK -Base and simulated “group predictions” with the same confusion matrix as those of B ARACK -Base but random errors, and find that the B ARACK predictions do not substantially degrade performance compared to these randomized predictions.
Thus, we hypothesize that the errors made by B ARACK -Base are “sufficiently random” to make the conditions of Corollary C.4 hold approximately.
We now proceed to prove the results stated in this section.
C.2. Notation
The notation “ i ∈ [ b ] ” means i ∈ { 1 , 2 , . . .
, b } .
log denotes the natural logarithm.
Where unspecified, ∥·∥ denotes the
Euclidean norm ∥·∥ 2 .
∆ d denotes the d -coordinate ( d − 1 dimensional) simplex.
We have a training dataset { ( x i , y i , z i ) } ni =1 ∈ X × Y × G .
Here x i is a datapoint, y i is a discrete class label, and z i is a discrete group label.
Let F 1 be a function class { f θ : X × Y → ∆ |G| } , where each f θ is a member of F 1 parameterized by the vector θ ∈ R d .
(The inputs to f θ are the features x and class label y , and the output is a vector of predicted probabilities for each group.)
Define F R 1 :=
{ f θ ∈ F 1 : ∥ θ ∥ 2 ≤ R } .
Similarly let F 2 be a function class { f θ : X → ∆ |Y| } and F R 2 :=
{ f θ ∈ F 2 : ∥ θ ∥ 2 ≤ R } .
Let ℓ ( · , · ) be a nonnegative loss function globally bounded by B ; ℓ is either in R |Y| × Y
→ R or R |G| × G → R depending on whether it is for the group classification task or target classification task, and will be clear from context.
We will use the notations ℓ ( f θ ( x ) , y ) and ℓ ( x, y ; θ ) interchangeably (or even ℓ ( f ( x ) , y ) or ℓ ( x, y ) when clear from context).
We assume that the datapoints are IID samples from a distribution P , i.e., ( x i , y i , z i ) ∼ P .
We can write P as a mixture of distributions p 1 , . . .
, p |G| , where p g is the distribution conditioned on the group label being equal to g .
By overloading notation, we will also write E ( x,y ) ∼ p
g to denote the expectation of a quantity conditioned on the group label being g (i.e., we will not write ( x, y, z ) ∼ p g because z = g always for points sampled from p g ).
Let q 1 , . .
, q |G| be the corresponding mixture weights, i.e., q g = P ( z = g ) , and let q :=
min g ∈G q g .
Recall that we denote D 1 :=
{ ( x i , y i , z i ) } mi =1 to be the group-labeled dataset of m points, and D 2 :
= { ( x i , y i ) } ni = m
+1 to be the dataset of n − m group-unlabeled points.
We are interested in the regime where m is small compared to n , so we shall implicitly assume that n − m is Ω( n ) .
In our theoretical results we assume for simplicity that D 1 and D 2 are samples from the same distribution (although the z i ’s are unobserved on D 2 ).
However, our analysis easily extends to the case where
D 1 , D 2 are sampled from distributions on which the group proportions q 1 , . . .
, q |G| differ, but the per-group distributions p g
= P ( x, y |
z = g ) are the same.
(This only complicates notation, as we need to distinguish the different q =
min
g ∈G q g between D 1 and D 2 .)
We will use the notation L robust ( θ ) to denote either max g ∈G E
[ ℓ (( x, y ) , z ; θ )
| z = g ] (worst-group population loss for the task of classifying the groups ) or max g ∈G E [ ℓ ( x, y ; θ )
| z = g ] (worst-group population loss for the task of classifying the classes ), which will be clear in context.
We will also use L ∗ robust as shorthand for min θ L robust ( θ ) .
C.3. “Helper” Results
We will use the following standard result from learning theory (Liang & Ma, 2019):
Theorem
C.5.
Suppose ℓ ( x, y ; θ ) is nonnegative, globally bounded by
B and
LLipschitz continuous.
Define
L [Equation]
] and ˆ L
[Equation]
θ
[Equation]
:= n [Equation] i = ℓ ( x i , y i ; θ ) , where { ( x i , y i ) } ni =1 are sampled IID from distribution P .
Let p > 0 .
Then, with probability ≥ 1 − O ( e − p ) , for all θ such that ∥ θ ∥ 2 ≤ R we have | ˆ L ( θ ) − L ( θ ) | ≤ O
[Equation]
We will also use the following simple lemma relating the minimizer of a “perturbed” GDRO-style objective to the minimizer of the unperturbed version.
Lemma C.6.
Define f ( θ ) =
max k ∈ [ G ] f k ( θ ) and ˜ f
( θ ) = max k ∈
[ G ] ( f k ( θ ) + ϵ k ) where | ϵ k | ≤ ϵ .
Let the minimizers of f, ˜ f be θ ∗ , ˜ θ ∗ respectively.
Then f (˜ θ ∗ ) − f ( θ ∗ ) ≤ 2 ϵ .
Proof.
Note that ˜ f
( θ ∗ ) ≤ f ( θ ∗ ) + ϵ .
Similarly f (˜ θ ∗ ) ≤ ˜ f
(˜ θ ∗ )
+ ϵ .
So f (˜ θ ∗ ) ≤ ˜ f
(˜ θ ∗ ) + ϵ ≤ ˜ f
( θ ∗ ) + ϵ ≤ f ( θ ∗ ) + 2 ϵ .
C.4. Warm-up: Proof of Group DRO Generalization Bound (Lemma C.1)
We restate Lemma C.1 below:
Lemma C.7.
Let ˆ f gdro,labeled ∈ F 2 (with associated parameters ˆ θ gdro,labeled ) denote the GDRO classifier trained on Donly.
Then with high probability,
[Equation]
Proof .
The population group DRO loss is L [Equation] ( θ ) :
[Equation]
max
[Equation]
E
[Equation] x [Equation] ) ∼ p g [ ℓ ( x, y ; θ )] (cid:1) , i.e. the maximum of the average per-group losses.
Here p g denotes the conditional distribution P ( x, y | z = g ) .
For g ∈ G , denote S g to be the subset of points on D 1 such that the group label is g .
The empirical GDRO loss on D 1 is L robust, D 1 ( θ ) : = max g ∈ [ G ] [Equation]
E
[Equation]
x
[Equation]
) ∼ S g [ ℓ ( x, y
[Equation]
θ )] (cid:1) .
Note that each of the per-group losses is simply the empirical estimate of the corresponding population per-group loss, over the set S g .
We can apply Theorem C.5 to each of the groups individually, since each set S g is an IID sample from p g .
Thus, we obtain that for a given group g , (cid:12) (cid:12) E ( x,y ) ∼ S g [ ℓ ( x, y ; θ )] − E ( x,y ) ∼ p
g
[ ℓ ( x, y ; θ )] (cid:12) (cid:12) ≤ O ( (cid:113) p log | S g | | S g | ) for all θ
such that ∥ θ ∥ 2 ≤ R , with probability ≥ 1 − O
( e − p ) .
( | S g | is the size of set S g .)
Thus by union bound, max g ∈G (cid:12)(cid:12) E ( x,y ) ∼ S g [ ℓ ( x, y ; θ )] − E ( x,y ) ∼ p
g
[ ℓ ( x, y ; θ )] (cid:12)(cid:12) ≤
O (cid:18) max S g (cid:113) p log | S g | | S g | (cid:19) w.p. ≥ 1 − O
( |G| e − p ) .
(Here and henceforth we omit the constants B, L, R from the notation.)
Let q g :=
P ( z = g ) , i.e., the population fraction of group g .
By Hoeffding’s inequality, for a given group g we have | S g − q g m | ≤ √ m log m with probability ≥ 1 − 2 /m 2 .
Thus by union bound and the fact that q g ≥ q , we have
[Equation]
with probability greater than 1 − 2 |G| /m 2 .
Thus max g ∈G log | S g | | S g | ≤ log m qm ·
[Equation]
log m q √ m with probability greater than 1 − 2 |G| /m 2 .
As 11 − x ≤ 1 + 2 x for all x ∈ [0 , 1 / 2] , we thus have
[Equation]
for sufficiently large m (relative to q ).
Thus
[Equation]
with probability greater than 1 − 2 |G| /m 2 .
To recap, the training loss is max g ∈G E ( x,y ) ∼ S g E
[ ℓ ( x, y ; θ )] and we showed that
[Equation]
for all g ∈ G and all θ with ∥ θ ∥ 2 ≤ R with high probability, so applying Lemma C.6 and union bound, we have that
[Equation]
with high probability, as desired.
Note that as long as we optimize over a bounded domain that is sufficiently large to contain θ ∗ , there exists a valid norm constraint R such that all θ under consideration have ∥ θ ∥ 2 ≤ R .
C.4.1. R ELATING LOSS TO CLASSIFICATION ERROR
For example, we could let ℓ be the truncated cross-entropy loss (i.e., the cross-entropy loss clipped to remain in [0 , B ] for some large constant B in order to ensure boundedness).
Observe that the classification error is upper bounded by the cross-entropy loss divided by log( |Y| ) , where |Y| is the number of classes: as the class prediction is the class with highest predicted probability, if the model makes an error then the predicted probability of the correct class is at most 1 / |Y| , which means that the cross-entropy loss for that example is ≥ − log(1 / Y ) =
log |Y| .
In other words, if argmax { f ( x ) } ̸ = y then ℓ ( f ( x ) , y ) ≥ log |Y| .
Thus 1 (argmax { f ( x ) } ̸ = y ) ≤ ℓ ( f ( x ) ,y ) log |Y| , and so E [ 1 (argmax { f ( x ) } ̸ = y )] ≤ 1 log |Y| E [ ℓ ( f ( x ) , y )] .
(Thus, the clipping constant B just needs to be ≥ log |Y| .)
We could also let ℓ be the squared loss (the square of 1 minus the predicted probability of the correct class), which is 2-Lipschitz and bounded by 1.
The classification error is upper bounded by |Y| |Y|− 1 times the squared loss (since when the model makes an error, the loss on that example must be at least (1 − 1 / |Y| )
+( |Y|− 1) / |Y| 2 = 1 − 2 / |Y| +1 / |Y| =
|Y|− 1 |Y| ).
C.5. Proof of Theorem 4.
We restate Theorem 4.1 below, and prove it in the following two subsubsections.
Theorem 4.1.
Suppose that on each group, the error rate of the group classifier from “Stage 1” is ≤ r .
Then with high probability, L robust (ˆ θ robust )
≤ L ∗ robust + ˜ O (cid:16) rq + 1 √ qn (cid:17) .
C.5.1. P ER - GROUP POPULATION LOSS VS . PER - ESTIMATED - GROUP POPULATION LOSS .
Let ˜ z denote the prediction of the
Stage group classifier, i.e.
˜ z
= argmax { ˆ f group ( x, y ) } .
The assumption on the error rate on each group is equivalent to assuming
P (˜ z ̸ =
z | z = g ) ≤ r for each g ∈
G .
Notice that
P ( z ̸ = ˜
z | ˜ z = g )
=
(cid:80) g ′ ∈|G| ,
g ′ ̸
= g P ( z = g ′ | ˜
z = g )
= (cid:80) g ′ ∈|G| ,
g ′ ̸
= g P
(˜ z = g | z = g ′ )
P ( z = g ′ )
P (˜ z = g ) .
Now, P (˜ z = g ) ≥ P (˜ z
= g | z = g ) P ( z = g ) ≥ (1 − r ) q , and P (˜ z
= g | z = g ′ )
≤ r for g ̸ = g ′ , and (cid:80) g ′ ̸
=
g P ( z = g ′ ) ≤ 1 − q , so P ( z ̸ = ˜
z | ˜ z = g ) ≤ (1 − q ) r
(1 − r ) q ≤ r (1 − r ) q .
Also, of course P ( z ̸ = ˜
z | ˜ z = g ) ≤ 1 as well, so
P ( z ̸ = ˜
z | ˜ z = g ) ≤ min { r q (1 − r ) , 1 } .
| E [ ℓ ( x, y ) | ˜
z = g ] − E [ ℓ ( x, y ) |
z = g ] | =
| E [ ℓ ( x, y ) | ˜
z
= g, z = g ] P ( z = g | ˜
z = g )
+
E [ ℓ ( x, y ) | ˜
z
= g, z ̸ = g ] P ( z ̸ = g | ˜
z = g ) − E [ ℓ ( x, y ) |
z = g ] | ≤ | E [ ℓ ( x, y ) | ˜
z
= g, z = g ] P ( z = g | ˜
z = g ) − E [ ℓ ( x, y ) |
z = g ] | +
E [ ℓ ( x, y ) | ˜
z
= g, z ̸ = g ] P ( z ̸ = g | ˜
z = g ) ≤ | E [ ℓ ( x, y ) | ˜
z
= g, z = g ] P ( z = g | ˜
z = g ) − E [ ℓ ( x, y ) |
z = g ] | + B min { r q (1 − r ) , 1 } .
Similarly,
E [ ℓ ( x, y ) |
z = g ] = (cid:16) E
[ ℓ ( x, y ) | ˜ z
= g, z = g ] P (˜ z
= g | z = g ) +
E [ ℓ ( x, y ) | ˜ z ̸ = g, z = g ] P (˜ z
̸ = g | z = g ) (cid:17) , so | E [ ℓ ( x, y ) | ˜
z
= g, z = g ] P ( z = g | ˜
z = g ) − E [ ℓ ( x, y ) |
z = g ] | = (cid:16)
| E [ ℓ ( x, y ) | ˜
z
= g, z = g ] P ( z = g | ˜
z = g ) − E [ ℓ ( x, y ) | ˜
z
= g, z = g ] P (˜ z
= g | z = g ) −
E [ ℓ ( x, y ) | z ̸ = g, z = g ] P (˜ z
̸ = g | z = g ) | (cid:17) ≤
E [ ℓ ( x, y ) | ˜
z
= g, z = g ] | P ( z = g | ˜
z = g ) − P (˜ z
= g | z = g ) | + Br.
Finally, we have ≥
[Equation]
≥ max { 1 − r q (1 − r ) , 0 } , and 1 ≥ P (˜ z
= g | z = g ) ≥ 1 − r , so
| P ( z = g | ˜
z = g ) − P (˜ z
= g | z = g ) | ≤ min { r q (1 − r ) , 1 } . Thus, altogether we have | E [ ℓ ( x, y ) | ˜ z = g ] − E [ ℓ ( x, y ) |
z = g ] | ≤
B min { r q (
[Equation]
r ) , 1 } + B min { r q (
[Equation]
r ) , 1 } + Br ≤ 3 B min { r q (1 − r ) , 1 } , which is O ( r/q ) in terms of r and q .
C.5.2. P ER - ESTIMATED - GROUP POPULATION LOSS VS . TRAINING LOSS
Let ˜ S
g denote the set of training points with predicted group label g .
The datapoints in ˜ S g are independent samples from P ( x, y | ˜
z = g ) .
So by Theorem C.5, for a given g we have that the difference between the per-group training and population losses on the Stage 2 classification task is | E (
x,y ) ∼ ˜ S
g
[ ℓ ( x, y ; θ )] − E [ ℓ ( x, y ; θ ) | ˜
z = g ] | ≤ O (cid:18)(cid:114) log |
˜ S g | | ˜
S g | (cid:19) with high probability.
By Hoeffding’s inequality we have O (cid:18)(cid:114) log | ˜
S g | | ˜
S g | (cid:19) ≤ O (cid:16)(cid:113) log n q (1 − r ) n (cid:17) =
O (cid:16)(cid:113) log n qn (cid:17) .
Thus by triangle inequality and union bound, we have | E ( x,y ) ∼ ˜ S
g
[ ℓ ( x, y ; θ )] − E [ ℓ ( x, y ; θ )
| z = g ] | ≤ O ( r/q ) + O (cid:16)(cid:113) log n qn (cid:17) =
˜ O (cid:16) rq + 1 √ qn (cid:17) for all g ∈ G and all θ with ∥ θ ∥ 2 ≤ R with high probability.
Finally, applying Lemma C.6 yields the desired result, as the training loss is max g ∈G E ( x,y ) ∼ ˜ S g [ ℓ ( x, y ; θ )] .
□
C.6. Extension of Theorem 4.1: Non-realizable case
In this section, we show that if the group classifier makes randomized predictions according to the predicted probabilities (rather than classifying groups by picking the group with maximum predicted probability), we can yield a more general bound than Theorem 4.1 that is better in the non-realizable case (i.e., when there does not exist a perfect group classifier).
We relax the realizability requirement stated in Section 4 to the requirement that there must exist some f ∗ ∈ F 1 such that f ( x, y ) is the vector of true probabilities P ( z | x, y ) .
This theorem is stated below.
Theorem C.8.
Suppose there exists f ∗ ∈ F 1 such that f ( x, y ) is the vector of true probabilities P ( z | x, y ) .
Let ˆ P ( z | x, y ) denote the probabilities output by the group classifier from ”Stage 1,” and suppose for each datapoint ( x, y ) the predicted group label ˜ z is sampled from ˆ P ( z | x, y ) .
Suppose that the total variation between ˆ P ( z | x, y ) and P ( z | x, y ) is bounded by r , for all x, y in the support of P .
Then with high probability, L robust (ˆ θ robust )
≤ L ∗ robust + ˜ O (cid:16) rq + 1 √ qn (cid:17) .
Theorem C.8 shows that even if it is impossible to perfectly distinguish the groups, this is not necessarily an obstacle to the downstream robust performance.
To prove Theorem C.8, we first prove the following lemma.
Lemma
C.9.
Suppose
[Equation]
Then
[Equation] , y ; θ ) | ˜
z = g ] , for all g ∈ G .
Proof.
By Bayes’ rule P ( x, y | ˜ z = g ) =
P (˜ z
= g | x, y ) P ( x, y ) /P
(˜ z = g ) .
By assumption, for any ( x, y ) , P (˜ z = g | x, y ) =
P ( z = g | x, y ) .
Therefore P ( x, y | ˜
z = g )
=
P (˜ z
= g | x, y ) P ( x, y ) /P
(˜ z = g ) =
P ( z = g | x, y ) P ( x, y ) /P
( z = g ) =
P ( x, y | z = g ) by applying Bayes’ rule again.
The claim follows.
Lemma C.9 implies that if we use the predicted group labels ˜ z rather than the “true” group labels z , there is essentially no difference since
˜ z
and z have the same distribution conditioned on x, y .
Thus, samples from ( x, y, ˜ z ) and ( x, y, z ) are equivalent for our purposes, and applying Lemma C.1 shows that the minimizer of ˆ˜ θ of max g E [ ℓ ( x, y ; θ ) | ˜
z = g ] satisfies
L robust (ˆ˜ θ ) ≤ L ∗ robust + ˜ O
( 1 √ qn ) with high probability (since in this case we have n total datapoints).
Proof of Theorem C.8 .
Note that if r = 1 the statement follows trivially from boundedness.
Similarly if r = 0 the statement follows from the argument above.
So assume 0 < r < 1 .
Given x, y , suppose we sample z ′ ,
˜ z ′ in the following “coupled” fashion.
Flip a biased coin with probability of heads being (cid:80) g ∈G min (cid:8) ˆ P ( z = g | x, y ) ,
P (˜ z
= g | x, y ) (cid:9) .
If the coin is heads, then sample z ′′ from the distribution where
[Equation]
and set z ′ =
˜ z ′ =
z ′′ .
Note that by the assumption r < 1 , the denominator is nonzero.
If the coin is tails, then sample z ′ from the distribution where
[Equation]
and independently sample ˜
z ′ from the distribution where
[Equation]
Also note that by the assumption r > 0 , the denominator is nonzero.
Notice that using this sampling procedure, P (˜ z ′
| x, y ) =
P ( z | x, y ) .
Similarly, P ( z ′ | x, y )
=
ˆ P ( z ′ | x, y ) .
Thus by Lemma C.9,
E [ ℓ ( x, y ; θ ) | z ′ =
g ]
=
E [ ℓ ( x, y ; θ )
| z = g ] =
L robust ( θ ) , and E [ ℓ ( x, y ; θ ) | ˜ z ′ =
g ]
=
E [ ℓ ( x, y ; θ ) | ˜
z = g ] .
Also observe that the probability that z ′ ̸ =
˜ z ′ is ≤ 1 − (cid:80)
g ∈G min { ˆ P ( z = g | x, y ) , P (˜ z
= g | x, y ) } which is precisely the total variation between ˆ P ( z = g | x, y ) and P (˜ z
= g | x, y ) .
From here, the remainder of the proof is essentially identical to the proof of
Theorem 4.1, as the sampled group labels z ′ from the true conditional distribution are “equivalent” to “correct” group labels, and the disagreement rate between z ′ and the estimated group labels ˜
z ′ is ≤ r by the total variation assumption.
C.7. Proof of Corollary 4.
Corollary C.10.
With high probability, for B ARACK -Base we have L robust (ˆ θ robust )
≤ L ∗ robust + ˜ O (cid:16) 1 q √ m (cid:17) .
Proof .
First, by Lemma C.1, we have that if we train a classifier with group DRO on D 1 to classify the group labels , the worst-group population loss is ˜ O
(1 / √ qm ) with high probability (since in the realizable case, there exists a group classifier with 0 population loss).
As discussed in Appendix C.4.1, for both the cross-entropy and the squared loss this translates to a (population) worst-group misclassification error of ˜ O
(1 / √ qm ) as well (when the classifier prediction is the group with maximum predicted probability).
In fact, by inspecting the proof of Lemma C.1, we can make the slightly stronger statement that the population classification error on group g is ˜ O (
1 / √ q g m ) with high probability.
So,
[Equation]
g
[Equation]
√ m (cid:80) √ q g ′ (cid:33) =
˜ O ( 1 q √ m )
since P (˜ z = g ) ≥ P (˜ z
= g | z = g ) P ( z = g ) ≥ (1 − ˜ O
(1 / √ q g m ))
q g with high probability and (cid:80) (cid:112) q ′ g ≤ (cid:80) 1 / (cid:112) |G| =
(cid:112) |G| .
Thus by a similar argument to the proof of Theorem 4.1, we have | E [ ℓ ( x, y ) | ˜
z = g ] − E [ ℓ ( x, y ) |
z = g ] | ≤ ˜ O
( 1 q √ m ) , and in turn can obtain the desired result (by an analogous argument as Appendix C.5.2).
C.8. Proof of Corollary C.
Corollary C.4 is restated below.
Corollary C.11.
In addition to the assumptions of Theorem 4.1, suppose that for all f ∈ F 2 , ℓ ( f ( x ) , y ) ⊥ argmax { ˆ f group ( x, y ) } (cid:12)(cid:12) z .
Let θ ∗ avg be the minimizer of the population average loss E ( x,y ) ∼ P [ ℓ ( x, y ; θ )] .
Then for B ARACKBase we have L
[Table]
[Table]
[Table]
Proof .
In words, the assumption that ℓ ( f ( x ) , y ) ⊥ argmax { ˆ f group ( x, y ) } (cid:12)(cid:12) z for all f ∈ F 2 says that for all classifiers in
F 2 , the loss on the target task (of classifying the class labels) is independent of the prediction of the group classifier, when conditioned on the actual group label.
Using this assumption, we have that E ( x,y )
[ ℓ ( x, y
[Equation]
| ˜ z = g ] =
E
[Equation]
Denote L avg ( θ ) =
E ( x,y ) [ ℓ ( x, y ; θ )] . Note that E ( x,y )
[ ℓ ( x, y ; θ
[Equation]
| z ̸
[Equation]
g
[Equation]
z [Equation] [Equation] y
[Equation]
g
[Equation]
P ( z = g ) =
L avg ( θ ) .
Thus,
[Equation]
As a result,
[Equation]
where a = P ( z ̸ = g | ˜
z = g )
P ( z ̸ = g ) ≤ ˜ O ( 1 √ qm ) with high probability (as shown in previous sections).
Let ˜ θ pop denote the minimizer of E ( x,y )
[ ℓ ( x, y ; θ ) | ˜
z = g ] .
Suppose for contradiction that L robust (˜ θ pop ) >
L robust ( θ ∗ avg ) .
Let
[Equation]
Then we have
[Equation]
>
E
[Equation]
[Equation]
i.e. that
[Equation]
and thus that
[Equation]
<
E ( x,y ) ∼ P [ ℓ ( x, y ; θ )] .
But there must be some k ∈ [ G ] such that
[Equation] ) ∼ P [ ℓ ( x, y ; θ )] , which contradicts the definition of k max .
Thus, L robust (˜ θ pop ) ≤ L robust ( θ ∗ avg ) .
Using the fact that the datapoints in ˜ S g are independent samples from P ( x, y | ˜
z = g ) , we have for all g
that | E ( x,y ) ∼ ˜ S
g
[ ℓ ( x, y ; θ )] − E [ ℓ ( x, y ; θ ) | ˜
z = g ] | ≤ O (cid:18)(cid:114) log |
˜ S g | | ˜
S g | (cid:19) with high probability (as argued in Appendix C.5.2).
By combining this with the result of Corollary 4.2, we obtain the desired result.
□
C.9. Proof of Lemma C.
Lemma C.12.
Let ˆ f erm,full ∈ F 2 (with associated parameters ˆ θ erm,full ) denote the ERM classifier trained on the full dataset.
There exists a distribution P such that with high probability, L robust (ˆ θ erm,full ) ≥ log(1 /q ) log( |Y| )
L ∗ robust − ˜ O
(1 / √ n ) .
Proof.
Consider the following simple distribution.
Suppose that the distribution of x is a point mass on a single point.
Suppose that there are k classes and the classes are identical to the groups (that is, y =
z always).
If cross-entropy loss is used, then clearly to minimize worst-group loss one should predict a uniform distribution over each class (since the point x gives no information about the class or group).
Thus, the loss on each point would be log k , so L ∗ robust =
log k in this case.
Recall that the cross-entropy loss has the property that the average population cross-entropy loss is minimized when the predicted probability of each class is simply the true probability of that class conditioned on the features x .
In our case, the latter is simply the probability of the class.
Thus, for the minimizer of the population average cross entropy loss, the predicted class probabilities for any point are the true class probabilities P ( c ) , and so the loss given that the true class is c is − log P ( c ) by definition.
So the worst-group loss is max c − log P ( c )
= − log min c P ( c )
= log(1 /q )
[note that 1 /q ≥ k ].
The lemma now follows by a simple application of Hoeffding’s inequality, since the training dataset has n points.
C.10. Theoretical Analysis of B ARACK -SSL
In Corollary C.13 (below), we analyze B ARACK -SSL (Algorithm 5), which uses semi-supervised learning to train the group classifier and is described further in Appendix B.
For this corollary, we assume that the group classification problem is realizable (i.e., there exists a classifier with 0 training loss on the group classification task) and that m = Ω( √ n ) .
Corollary C.13.
Under appropriate conditions, for B ARACK -SSL we have
L robust (ˆ θ robust ) ≤ L ∗ robust + ˜ O (cid:16)
1 q √ n + 1 qm (cid:17) with high probability.
In B ARACK -SSL, semi-supervised learning is used to train the group classifier, leveraging the group-unlabeled points.
With an appropriate choice of semi-supervised learning method (such as FixMatch (Sohn et al., 2020)), if m = Ω( √ n ) , then under appropriate conditions, the worst-group generalization error rate bound of B ARACK -SSL is O (1 / √ n ) .
This result is based on the PAC-learning based results of (Balcan & Blum, 2009), and requires the assumptions therein as well as realizability (the exact conditions are somewhat technical, and are deferred to the discussion in Appendix C).
This translates to an excess robust risk of O (1 / √ n ) for the final model, as stated in Corollary C.13.
To prove Corollary C.13, we make use of the following theorem (see Theorem C.14 below, a restated version of Theorem 21.
from (Balcan & Blum, 2009)).
First, we need some additional notation: let χ : F 1 × X
→ { 0 , 1 } be a function and define the overloaded notation χ ( f, P )
=
E x ∼ P [ χ ( f, x )] .
Let V C ( C ) denote the VC-dimension of a function class C .
For function classes that output probabilities rather than labels directly, we overload notation so that V C ( F ) =
V C ( { argmax { f } | f ∈
F} ) .
Recall that D 2 denotes the training set of group-unlabeled points.
Theorem C.14.
Given δ ∈ (0 , 1) , if n = Ω (cid:0) 1 ϵ 2 (cid:0) V C ( χ ( F 1 )) log 1 ϵ + log 1 δ (cid:1)(cid:1) and m ≥ 2 ϵ (cid:0) log(2 s (2 m, t + 2 ϵ ))
+ log 2 δ (cid:1) , then with probability ≥ 1 − δ it holds that all f ∈ F 1 with zero training error and 1 − χ ( f, D 2 ) ≤ t + ϵ have population error ≤
ϵ .
Here, s (2 m, t + 2 ϵ ) denotes the expected number of splits when 2 m points are drawn IID from P with concepts f ∈ F 1 having χ ( f, P ) ≤ t + 2 ϵ .
FixMatch (Sohn et al., 2020) (which we use for Stage 1 of B ARACK - SSL , as described in 5) minimizes a weighted sum of the supervised loss (computed on the training points with group labels) and an unsupervised loss (consistency of predictions between examples and augmented versions of the same example, computed on all training points).
Concretely, in our case the FixMatch loss for a group classifier f ∈ F 1 is L fixmatch =
λ L sup + (1 − λ )
L consistency , where λ ∈ [0 , 1] , L sup =
E ( x,y,z ) ∼D 1 [ ℓ ( f ( x, y ) , z )] , and L consistency =
E ( x,y ) ∼D
2 [ 1 (max { f ( x, y ) } ≥ τ ) · ℓ ( f ( aug ( x ) , y ) , f ( x, y ))] , where τ ∈ [0 , 1] is a predefined constant.
Here, aug ( · ) denotes the augmentation function; for simplicity assume it is a fixed (non-random) function.
Let χ ( f, x ) beif max { f ( x ) }≥ τ and argmax { f ( x ) } ̸ = argmax { f ( aug ( x )) } , and 1 otherwise (in other words, χ ( f, x ) is 1 unless f makes a confident prediction on x but makes a different prediction on the augmented version of x ).
By the realizability assumption, there exists f ∈ F 1 with zero loss (and therefore zero training loss) on the supervised task of classifying the groups.
Note that if aug is the identity function, then f also attains zero consistency loss; more generally, given ξ ≥
we can choose aug to be a “weak enough” augmentation such that there exists f ∈ F 1 such that L consistency ≤ ξ .
If we choose ξ to be sufficiently small, we can guarantee that L consistency ≤ δ implies χ ( f, D 2 ) ≤ t + ϵ , since t, ϵ ≥ 0 .
Because there exists f ∈ F 1 with zero loss, this implies that if we set the weight λ on the supervised part of the loss to be large enough, the function F 1 we learn (corresponding to the minimizer of L fixmatch ) will have zero training error.
(If we find a function with nonzero training error, we can increase the weight of the supervised loss and rerun.)
Then, invoking Theorem C.14, we get that the population error of the returned group classification model will be ≤ O ( ϵ ) with high probability as long as m, n satisfy m = Ω(1 /ϵ 2 ) and n = Ω(1 /ϵ ) .
The value of the term t in Theorem C.14 depends on the consistency error χ
( f, D 2 ) on the training data, which in turn will depend on the augmentation chosen, the threshold τ , and the weight 1 − λ (smaller λ will encourage lower consistency loss).
Similarly, the constants in the preceding Ω( · ) terms depend on the choice of augmentation function aug ( · ) ; intuitively, we would like to choose an augmentation function that is as strong as possible while still being label-preserving (so as not to make the consistency loss large).
In summary, given the required conditions, if m = Ω(1 /ϵ 2 ) and n = Ω(1 /ϵ ) then the population error of the returned group classification model is ≤ O ( ϵ ) with high probability.
The population error of the group classifier on each group g will then be ≤ O ( ϵ/q g ) with high probability.
Equivalently, with high probability, for all groups g the population error of the group classifier on that group will be O
( 1 q g m + 1 q g √ n ) [as long as m = Ω( √ n ) ].
Finally, the desired result now follows by invoking Theorem 4.1.
D. Experimental Details and Additional Results
We first describe the datasets, baselines, and experimental details more thoroughly.
We then present extensive ablation experiments in Appendix D.3 and results for B ARACK -SSL in Appendix D.4 to support the robustness and extensibility of our B ARACK framework.
D.1. Dataset Details D.1.1. D ATASETS
Waterbirds.
Waterbirds (Sagawa et al., 2019), is a popular robustness benchmark that consists of images from ‘landbird’ and ‘waterbird’ species on either land or water backgrounds.
The task is to classify images as ‘landbird’ vs. ‘waterbird’, and the groups are defined by background.
5 95% of landbirds are on land backgrounds and similarly for waterbirds; this spurious correlation makes landbirds on water and waterbirds on land harder to classify.
CelebA. CelebA
(Liu et al., 2015a) is a popular face classification dataset which is also often used to evaluate robustness to spurious correlations.
The task is to classify faces as ‘blond’ or ‘not blond’, and the groups are defined by gender.
Blondness is spuriously correlated with gender: the female faces are disproportionately blond and the male faces disproportionately non-blond.
For instance, only 6% of blond examples are male, leading to poor performance on this group.
U-MNIST.
U-MNIST (Sohoni et al., 2020) is a modified version of MNIST (LeCun et al., 2010), where the task is to classify digits as ‘ < 5 ’ or ‘ ≥ 5 ’, the groups are the individual digits, and only 5% of images in the ‘8’ group are retained from the individual dataset.
This rarity makes ‘8’ images more difficult to classify.
U-CIFAR10.
We introduce U-CIFAR10 as a modification of the CIFAR-10 dataset (Krizhevsky, 2009), where the task is to classify the image as ‘animal’ or ‘vehicle’, the groups are the 10 original CIFAR-10 classes, and we undersample the ‘airplane’ class to 5%.
Though similar to U-MNIST, this task is much more challenging.
D.1.2. D ATASET S ETUP
Recall that for our experiments, we pick a fixed budget of (training and validation) examples to label for each group, which are randomly sampled from the appropriate group in the original dataset, and evaluate how our approach fares as this budget is increased.
We use the same group label budget for both the training and validation sets.
Note that it is not fundamentally necessary for the group-labeled training and validation subsets to be the same size; this is merely a simple heuristic to trade off the amount of data for training the group classifier, and the amount of data for model selection.
Note that in reality, a group-labeled dataset may not have the same number of labeled points for each group, but rather reflect the training/population distribution, i.e. have the same group proportions as the population.
In this case, however, we can still convert our group-labeled dataset to one with balanced numbers of points per group, by simply subsampling the group-labeled points by group to get a (smaller) balanced dataset; the performance we achieve after this subsampling can be interpreted as a rough lower bound on the performance without such subsampling, since it simply throws away data.
In fact, we empirically observed that this subsampling does not meaningfully degrade the final performance, compared to using the un-subsampled version (containing more group-labeled points for the larger groups).
D.2. Training Details
For all datasets and methods, we use a fixed training/validation/test split.
Models are trained on the training set, and the validation set is used for model selection (both selecting the best model during training, and for hyperparameter selection).
All results reported in plots and tables are on the test set; the test set is not used for any model selection or tuning purposes.
For consistency and direct comparability with prior works, we do not use data augmentation for any of the baselines or while training the “Stage 2” robust models, except when explicitly specified otherwise.
For all experimental settings, reported means and standard deviations are over 5 trials with different random seeds.
(We also average over 5 trials for all ablation experiments in Appendix D.3.)
U-MNIST.
This task is based on the MNIST dataset (LeCun et al., 2010) (available under the Creative Commons
Attribution-Share Alike 3.0 license).
The U-MNIST task is to classify digits between ≤ 4 and ≥ 5 ; the groups are the individual digits.
We use a fixed training-validation split for all methods, to set aside 20% of the original MNIST training set (12,000 points) for validation.
On the training set, the ‘8’ digits are subsampled such that only 5% of them are kept.
The total number of points with each digit label in the training set are 0 : 4769 , 1 : 5382 , 2 : 4845 , 3 : 4950 , 4 : 4662 , 5 : 4300 , 6 : 4728 , 7 : 4980 , 8 : 234 , 9 : 4720 .
The validation set is approximately balanced.
Without modification, this can actually help methods that do not require group labels on the full validation set (such as ERM, B ARACK , G EORGE ) to more easily select models with good worst-group performance, because the rare ‘8’ group is overrepresented in the validation set compared to the training set, so it is easier to detect poor performance on that group even when looking at overall performance or performance using noisy group labels.
Thus, when computing an average metric (loss or accuracy) on a subset of the validation data, we compute a weighted average which is the sum over each group g of: the average of that metric for all points in the subset whose true group label is g , times the proportion of group g in the training dataset.
This reweighting procedure is the same as what is done in (Sagawa et al., 2019; Sohoni et al., 2020).
We use a 4 layer LeNet (LeCun et al., 1998) and the Adam optimizer for all methods.
For U-MNIST, for training the robust model we train for 100 epochs with a batch size of 128, the Adam optimizer (Kingma & Ba, 2015), and decay the learning rate by a factor of 0.1 at epochs 50 and 75.
These hyperparameters were taken from (Sohoni et al., 2020).
We tune all methods over the cross product of learning rates [2e-3, 2e-4] and weight decays [1e-4, 3e-4, 1e-5].
For methods using GDRO for the second stage, we also tune the GDRO group adjustment parameter in the set { 0, 3 } , and use uniform per-group sampling, as described in (Sagawa et al., 2019).
U-CIFAR10.
This task is based on the publicly available CIFAR-10 dataset (Krizhevsky, 2009) (license unknown).
The
U-CIFAR10 task is to classify images as “vehicle” or “animal”; the groups are the original CIFAR-10 classes.
We use a fixed training-validation split for all methods, to set aside 20% of the original MNIST training set (10,000 points) for validation.
On the training set, the ‘airplane’ images are subsampled such that only 5% of them are kept.
The total number of points with each group label in the training set are airplane: 204, automobile: 4004, bird: 3976, cat: 4017, deer: 3997, dog: 3999, frog: 4000, horse: 3976, ship: 3957, truck: 4003.
The validation set is approximately balanced, so as described for U-MNIST we compute reweighted metrics where appropriate.
We use a ResNet-50 (He et al., 2016) model and train for 200 epochs with a batch size of 128, SGD with momentum 0.9, and a cosine learning rate schedule.
These hyperparameters were taken from the implementation at https://github.com/ kuangliu/pytorch-cifar .
We tune all methods over the cross product of learning rates [1e-2, 1e-3] and weight decays [1e-3, 3e-3, 1e-2, 3e-2, 1e-1].
For methods using GDRO for the second stage, we also tune the GDRO group adjustment parameter in the set { 0, 3 } , and use uniform per-group sampling, as described in (Sagawa et al., 2019).
Waterbirds.
The Waterbirds dataset was created by (Sagawa et al., 2019) as a modification of the CUB dataset (Wah et al., 2011) (license unknown).
It consists of different bird species (with class labels either “waterbird” or “landbird”) on either a land or water background.
There are 3498 training images of landbirds on land, 184 of landbirds on water, 56 of waterbirds on land, and 1057 of landbirds on land (these are the four groups).
The validation set is more balanced, so as described for
U-MNIST we compute reweighted metrics where appropriate.
We use a ResNet-50 (He et al., 2016) model and train for 300 epochs with a batch size of 128 and SGD with momentum 0.9.
These hyperparameters were taken from (Sagawa et al., 2019).
We tune all methods over the (learning rate, weight decay) pairs (1e-4, 1e-1), (1e-3, 1e-4), and (1e-5, 1.0), as done in (Sagawa et al., 2019; Liu et al., 2021).
For methods using
GDRO for the second stage, we set the GDRO group adjustment parameter to 2 as in (Sagawa et al., 2019), and use uniform per-group sampling.
CelebA.
The CelebA dataset (Liu et al., 2015b) (license unknown) is a dataset of celebrity faces annotated with several descriptors (such as gender, hair color, wearing glasses).
It is often used as a benchmark for robustness to spurious correlations.
The task we consider is classifying the hair color of the person in the image as blond or non-blond, as in prior works such as (Sagawa et al., 2019;
Sohoni et al., 2020; Liu et al., 2021).
In this dataset, hair color is spuriously correlated with gender: there are 71629 images in the “female, non-blond” group, 66874 in the “female, blond”,“male, non-blond”, and just 1387 “male, blond”.
We use a ResNet-50 (He et al., 2016) model and train for 50 epochs with a batch size of 128 and SGD with momentum 0.9.
These hyperparameters were taken from (Sagawa et al., 2019).
We tune all methods over the (learning rate, weight decay) pairs (1e-4, 1e-2), (1e-4, 1e-4), and (1e-5, 0.1), as done in (Sagawa et al., 2019; Liu et al., 2021).
For methods using
GDRO for the second stage, we set the GDRO group adjustment parameter to 3 as in (Sagawa et al., 2019), and use uniform per-group sampling.
D.2.1. B ASELINE DETAILS
We reimplemented ERM and GDRO (and subset-GDRO) ourselves, along with B ARACK .
For the other baseline methods (George (Sohoni et al., 2020), JTT (Liu et al., 2021), and EIIL (Creager et al., 2021)), we use the authors’ publicly available repositories, adapting the code (such as to plug in our dataloaders) where necessary.
All three methods require first training an ERM model.
For G EORGE , this model is either a standard ERM model or one trained with high regularization, and is selected based on the Silhouette score of the clustered activations.
For U-CIFAR (the only dataset which George did not evaluate on originally), we used a learning rate of 1e-3 for this model and tuned the weight decay in [1e-3, 1e-1] based on this Silhouette score criterion.
For JTT and EIIL, the model is an ERM model trained with high regularization for a shorter number of epochs.
For U-MNIST and U-CIFAR10 (which the JTT paper did not evaluate on originally) we used a learning rate of 2e-3 and tuned the weight decay in [1e-1, 1e-3] for this ERM model, and tuned the number of training epochs in { 1, 50 } (the tuning criterion in this case was the validation worst-group accuracy of the final robust model), for both EIIL and JTT.
EIIL also did not evaluate on CelebA, so we use the same ERM model as in JTT.
For JTT and EIIL, we assume the group labels on the entire validation set are known, and use these for model selection (as done in the respective papers).
This gives them a slight advantage for model selection compared to B ARACK , in which we only use a small number of group-labeled validation examples (the same number as we use for training).
On the other hand,
George does not assume validation set group labels, but rather estimates them the same way the training group labels are estiated (via clustering the activations of the last layer).
D.2.2. B ARACK DETAILS
For B ARACK -base, we train a supervised group classifier model (with GDRO) to predict group pseudolabels, as described in
Section 3.
This is challenging due to the low number of points with known group labels: for instance, with 8 group-labeled points per group on Waterbirds or CelebA, the entire training dataset for this stage is 32 examples.
Thus, we generally need to train for more epochs to obtain reasonable results.
We also use mild data augmentation (random crops and flips) for the Stage 1 group classifier only to help deal with this lack of data (although, similarly to prior works, we do not use data augmentation in Stage 2, except for the experiment with RotNet with data augmentation).
For the learning rate and weight decay for the group classifier, we tune over the same pairs of hyperparameters as described above.
For the number of epochs, we train for 500 / ( n lab / 64) epochs, where n lab is the number of group-labeled examples in the smallest group (in our experiments, we evaluate the settings n lab ∈ { 8 , 16 , 32 , 64 } ).
(We evaluate the validation accuracy every n lab / 64 epochs to keep a constant number of validation evaluations.)
We set the batch size for training the group classifier to be the minimum of 128 or the total number of group-labeled examples.
We note that the group classifier and end model need not have the same architecture.
For instance, one could potentially select a smaller model for the group classifier since it is trained with a small amount of data.
In our preliminary experiments, however, we found that using the larger ResNet-50 model (the same architecture as the Stage 2 model) performed better than using a smaller ResNet-18 (by 5-10% group classification accuracy on average across all datasets).
For B ARACK (and Subset-GDRO), the points for which we know the group label are selected randomly from the training and validation datasets.
We select different sets of these points depending on the random seed (but for a fixed seed, these sets are the same, to facilitate direct comparisons).
We do this to avoid over-indexing interpretation of results to a particularly
[Figure]
“easy” or “difficult” set chosen by happenstance.
In the real world, of course, the points for which the group labels are known would generally be a fixed set.
We note that hyperparameters should not be tuned by looking at averaged metrics over the different seed results, because in a way this “shares information” between different trials with different group-labeled points, so in some sense is using more group label information than it should.
Thus, we indeed select the best hyperparameters separately for each seed (based on the appropriate validation metric for that seed).
For consistency, we do this for all methods (e.g., tune hyperparameters on a per-seed basis based on the validation metric) [Figure]
D.2.3. I NITIALIZATION
For training both the group classifier and the robust model, we typically start from a pretrained model:
ResNet-50 pretrained on ImageNet.
The exception is that for U-MNIST, we use a LeNet-5.
In Appendix D.3.3, we evaluate the impact of using different pretrained models (for the tasks other than U-MNIST).
D.3. Ablation Experiments
In this section, we present ablation experiments to study the reasons behind the worst-group accuracy gains offered by
B ARACK .
First, we run a synthetic experiment in which we run GDRO with randomly generated noisy group labels at different noise levels, to better understand how group prediction errors affect the final robust performance.
Next, we ablate the importance of using the class label as an input to the group prediction model (as described in Section 3).
Finally, we explore using models that are pretrained on ImageNet in a self-supervised manner (instead of supervised) as the starting model for B ARACK (and all the baselines), to understand how the worst-group accuracy trends from Section 5.2 generalize when different pretrained models are used.
D.3.1. R ANDOM F LIPPING .
To better understand why B ARACK can achieve worst-group accuracy close to that of GDRO even with fairly inaccurate predicted group labels, we run a synthetic experiment: we run GDRO with a varying fraction of the group labels randomly flipped, and compare the performance of this to that of B ARACK at an equivalent group prediction error rate.
Results are in
Figure 8; GDRO with the randomly perturbed group labels performs similarly to B ARACK .
Specifically, we take the ground-truth group labels, randomly flip them to get the same error rate and confusion matrix as the predicted group labels from B ARACK ’s group prediction model, and then use these “noisy group labels” as the groups for
GDRO.
For most settings, the final worst-group error is quite similar to that of B ARACK , which suggests that the errors made by our group prediction model are indeed “sufficiently random” to not adversely affect the downstream worst-group accuracy too much.
However, on U-MNIST, B ARACK underperforms the random flipping version with a small number of
[Table]
group-labeled points, although this difference decreases as the number of group-labeled points increases.
Additional Details.
As described above, for this experiment we created random “synthetic group pseudolabels” to have the same confusion matrix with respect to the true group labels as the confusion matrix of the B ARACK group pseudolabels with respect to the true group labels.
Specifically, for each setting (dataset, seed
[Table] and number of known group labels) we computed the confusion matrix of the group predictions output by the corresponding B ARACK model (i [Table] e., the one selected as the “best model” based on the criteria described in Appendix D.2)
[Table]
and then created the “synthetic pseudolabels” by taking the true group labels and flipping randomly chosen ones to match the confusion matrix.
We then used these synthetic pseudolabels in the GDRO objective, with the same hyperparameters as those of the B ARACK model.
(Thus, one possible explanation for the fact that the B ARACK models generally slightly outperform the “randomly flipped” models with equivalent error rates is the fact that we did not perform a separate hyperparameter search for the “randomly flipped” model, instead using the same ones as those of the B ARACK model with the same confusion matrix.)
D.3.2. C LASS L ABEL I NPUT .
To evaluate the importance of using the class label as input in the first stage of B ARACK , we run the standard B ARACK procedure except without using the class label as input.
This substantially decreases the final worst-group accuracy on all datasets (often by 10% or more, see Figure 8), showing that usage of the class label is indeed an important part of
B ARACK procedure.
On most datasets, this drop can largely be explained by the reduced group prediction accuracy when the class feature is not used (the worst-group accuracy of the group classifier is up to 20 points worse when the class is not used).
Interestingly, on U-CIFAR10, the group classifier’s worst-group accuracy drops only 1-5% when it does not use the class label, but this still results in a drop of 8-11% in the worst-group accuracy of the final B ARACK model.
Note that for all datasets considered, the classes are disjoint unions of the groups, meaning that knowing the class narrows down the possibilities for the group label.
Thus, it is unsurprising that using the class label in the group prediction model significantly improves the worst-group accuracy of both the group prediction model and the final model.
We hypothesize that the class label is more essential for group prediction in the spurious correlation setting (as suggested by the results in
Fig.
8), because it enables the group prediction model to disambiguate between examples with the same spurious attribute but different classes, allowing it to focus on identifying the spurious attribute itself.
By contrast, there is no clear “spurious attribute” on U-MNIST and U-CIFAR10; knowing the class label only reduces the number of candidate group labels for each example from 10 to 5.
D.3.3. P RETRAINED M ODEL C HOICE .
As described in Section 3, on all tasks except U-MNIST, in our default experiments we start from a pretrained model trained on the supervised ImageNet task for all methods (and for both the group classifier and final model in B ARACK ).
This is the standard approach when training on Waterbirds and CelebA, as in previous works (Sagawa et al., 2019; Sohoni et al., 2020; Levy et al., 2020; Liu et al., 2021).
In this subsection, we investigate how the choice of pretrained model affects performance (Table 3).
One motivation for this experiment is the potential for overlap or “leakage” between ImageNet and other image classification datasets (Kolesnikov et al., 2020).
Moreover, we seek to confirm that the observations in previous sections regarding the performance of B ARACK , compared to the baselines, are robust to the choice of pretrained model initialization.
Thus, in this section we use pretrained models that were trained only with self-supervision (i.e., not using any labels).
To be specific, in this section we use RotNet (Gidaris et al., 2018) from the VISSL library (Goyal et al., 2021) as opposed to the ResNet-50 pretrained on supervised ImageNet (from PyTorch) used in previous sections and previous works.
Results are reported in Table 3 and plotted in Figure 9.
[Figure]
of
[Figure]
Overall, we observe that using the RotNet model achieves somewhat worse accuracies (both worst-group and average) for all methods, although this gap can be reduced or eliminated by using data augmentation.
This unsurprisingly aligns with the findings of the original RotNet paper (Gidaris et al., 2018) that RotNet has somewhat worse transfer performance compared to supervised pretrained models.
Despite the gap, the key takeaway is that B ARACK still outperforms the baselines and remains competitive with full-dataset GDRO with the same initialization.
We additionally find that using data augmentation for the RotNet-based model can boost its performance closer to that of the supervised pretrained model.
We observe that data augmentation does not seem to benefit the RotNet model much on
U-CIFAR10, but it does on the other tasks (on CelebA, augmentation even boosts the performance using RotNet past the performance using the supervised ImageNet model, although of course the performance of the latter could likely also be boosted by using data augmentation).
For the RotNet training with augmentation, the augmentations used in both Stageand Stage 2 of B ARACK are random crops, random flips, and random rotations (of up to 15 degrees).
The reason we added the random rotations is because rotation is a key part of pretraining the RotNet model itself.
D.4. B ARACK -SSL: Semi-Supervised Learning for Group Prediction
As an extension, we investigate the use of semi-supervised learning (SSL) using FixMatch (Sohn et al., 2020), for Stage 1 of B ARACK .
(We still use GDRO for Stage 2.)
We refer to this procedure as B ARACK -SSL.
On the U-CIFAR10 task,
B ARACK -SSL can improve group prediction accuracy and, correspondingly, final robust performance.
For example, with only 8 group-labeled examples per group, the worst-group accuracy of the final B ARACK -SSL model is 94.0 %, compared to 90.4% for B ARACK -Base.
Correspondingly, the worst-group prediction accuracy of the group classifier is 83.7% when trained using FixMatch, while it is substantially lower at 42.9% when trained using simple supervised learning as in
B ARACK -Base, which helps explain these results.
Thus, while B ARACK -Base is simple and attains good worst-group accuracy, these results highlight the exciting potential of using more advanced SSL techniques to further boost worst-group performance of B ARACK at the cost of more complexity.
(Note: When training the group classifier using FixMatch, we use the class label as an input to the prediction head, just as in B ARACK -Base.)
Additional Details.
For B ARACK - SSL , we use FixMatch (Sohn et al., 2020) to train the semi-supervised group classifier.
We adapt the PyTorch implementation at https://github.com/kekmodel/FixMatch-pytorch , modifying it to use the class label the same way as described in Section 3 in order to assign zero probability to the groups that do not belong to the given class.
Other than that, we use the default FixMatch hyperparameters.
The group classifier used is simply the model at the end of FixMatch training (so we do not use the validation set at all for Stage 1).
For the Stage 2 GDRO model, we used the same hyperparameter search and model selection approach as described in Appendix D.2.
[Table]
D.5. Additional Results D.5.1. W ORST -G ROUP P ERFORMANCE .
In Table 4, we provide a more complete version of Table 1.
In this table, all results are rerun by us except for G EORGE on U-MNIST, CelebA and Waterbirds (since the original
G EORGE paper already reports results over 5 seeds on those datasets).
We ran 5 seeds for all methods except 10 seeds for
B ARACK , due to the increased variability from selecting different subsets of points with known group labels.
We note that our averaged results were somewhat better on Waterbirds and worse on CelebA than those reported in the J TT paper (which reports results from one trial).
Similarly, our results for EIIL on Waterbirds are also somewhat better than those reported in the EIIL paper (Creager et al., 2021).
Separately, the recent work (Zhang et al., 2022) proposes a method (CnC) based on contrastive learning for improving robustness to spurious correlations.
This method does very well on the spurious correlation datasets (Waterbirds and
CelebA), even exceeding the performance of B ARACK on CelebA. However, we found that CnC did not work well on U-MNIST and {+U-CIFAR
+} (attaining worst-group accuracies lower than those of ERM), which are tasks without explicit spurious correlations (even though CnC can in principle be applied to such tasks).
D.5.2. G ROUP P REDICTION A CCURACY .
In Figure 10 we provide more results on the performance of the “Stage 1” group classification model.
We plot both worst-group and average accuracies for the group classifier, corresponding to the same settings as in Table 4.
E. Conclusion
In this paper, our focus is to gain understanding into the practically important, yet so far understudied, version of the popular group robustness problem in which only a few group labels are known.
We show simple techniques can be highly effective for this novel setting, and support our work with theoretical results and extensive ablations.
Specifically, to address this problem, we present B ARACK , a two-stage approach to improve group robustness in the setting when only a small number of group labels are known.
We empirically validate that B ARACK outperforms methods that do not use training group labels, even with just a small number of group-labeled examples.
We theoretically provide generalization bounds on the worst-group performance of B ARACK .
Our results indicate that even a small number of group labels can be helpful for substantially improving worst-group performance.
[Figure]