Title: Robust Learning with Progressive Data Expansion Against Spurious Correlation
Venue: NeurIPS 2023
Accepted: yes
Keywords: spurious correlation

robustness

robust learning
Abstract: While deep learning models have shown remarkable performance in various tasks, they are susceptible to learning non-generalizable _spurious features_ rather than the core features that are genuinely correlated to the true label. In this paper, beyond existing analyses of linear models, we theoretically examine the learning process of a two-layer nonlinear convolutional neural network in the presence of spurious features. Our analysis suggests that imbalanced data groups and easily learnable spurious features can lead to the dominance of spurious features during the learning process. In light of this, we propose a new training algorithm called **PDE** that efficiently enhances the model's robustness for a better worst-group performance. PDE begins with a group-balanced subset of training data and progressively expands it to facilitate the learning of the core features. Experiments on synthetic and real-world benchmark datasets confirm the superior performance of our method on models such as ResNets and Transformers. On average, our method achieves a $2.8$ \% improvement in worst-group accuracy compared with the state-of-the-art method, while enjoying up to $10\times$ faster training efficiency.

Robust Learning with Progressive Data Expansion Against Spurious Correlation

Abstract

While deep learning models have shown remarkable performance in various tasks, they are susceptible to learning non-generalizable spurious features rather than the core features that are genuinely correlated to the true label. In this paper, beyond existing analyses of linear models, we theoretically examine the learning process of a two-layer nonlinear convolutional neural network in the presence of spurious features. Our analysis suggests that imbalanced data groups and easily learnable spurious features can lead to the dominance of spurious features during the learning process. In light of this, we propose a new training algorithm called PDE that efficiently enhances the model's robustness for a better worst-group performance. PDE begins with a group-balanced subset of training data and progressively expands it to facilitate the learning of the core features. Experiments on synthetic and real-world benchmark datasets confirm the superior performance of our method on models such as ResNets and Transformers. On average, our method achieves a 2 . 8% improvement in worst-group accuracy compared with the state-of-the-art method, while enjoying up to 10 × faster training efficiency. Codes are available at https://github.com/uclaml/PDE .

1 Introduction

Despite the remarkable performance of deep learning models, recent studies (Sagawa et al., 2019, 2020; Izmailov et al., 2022; Haghtalab et al., 2022; Yang et al., 2022, 2023b,c) have identified their vulnerability to spurious correlations in data distributions. A spurious correlation refers to an easily learned feature that, while unrelated to the task at hand, appears with high frequency within a specific class. For instance, waterbirds frequently appear with water backgrounds, and landbirds with land backgrounds. When training with empirical risk minimization (ERM), deep learning models tend to exploit such correlations and fail to learn the more subtle features genuinely correlated with the true labels, resulting in poor generalization performance on minority data (e.g., waterbirds with land backgrounds as shown in Figure 1). This observation raises a crucial question: Does the model genuinely learn to classify birds, or does it merely learn to distinguish land from water? The issue is particularly concerning because deep learning models are being deployed in critical applications such as healthcare, finance, and autonomous vehicles, where we require a reliable predictor.

Researchers formalized the problem by considering examples with various combinations of core features (e.g., landbird/waterbird) and spurious features (e.g., land/water backgrounds) as different groups . The model is more likely to make mistakes on certain groups if it learns the spurious feature. The objective therefore becomes balancing and improving performance across all groups. Under this formulation, we can divide the task into two sub-problems: (1) accurately identifying the groups, which are not always known in a dataset, and (2) effectively using the group information

Training data

Warmup stage:

balanced groups

Expansion stage: progressive expansion

y=+1

y=-1

a=+1

a=-1

...

Our method

Examples

ERM: 72.7%

Ours: 90.0%

# of data

ERM

Our method

Original

Core feature

GradCAM results

(A)

(C)

(B)

to finally improve the model's robustness. While numerous recent works (Nam et al., 2020; Liu et al., 2021; Creager et al., 2021; Ahmed et al., 2021; Taghanaki et al., 2021; Zhang et al., 2022) focus on the first sub-problem, the second sub-problem remains under studied. The pioneering work (Sagawa et al., 2019) still serves as the best guidance for utilizing accurate group information. In this paper, we focus on the second sub-problem and aim to provide a more effective and efficient algorithm to utilize the group information. It is worth noting that the theoretical understanding of spurious correlations lags behind the empirical advancements of mitigating spurious features. Existing theoretical studies (Sagawa et al., 2020; Chen et al., 2020; Yang et al., 2022; Ye et al., 2022) are limited to the setting of simple linear models and data distribution that are less reflective of real application scenarios.

We begin by theoretically examining the learning process of spurious features when training a twolayer nonlinear convolutional neural network (CNN) on a corresponding data model that captures the influence of spurious correlations. We illustrate that the learning of spurious features swiftly overshadows the learning of core features from the onset of training when groups are imbalanced and spurious features are more easily learned than core features. Based upon our theoretical understanding, we propose Progressive Data Expansion ( PDE ), a neat and novel training algorithm that efficiently uses group information to enhance model's robustness against spurious correlations. Existing approaches, such as GroupDRO (Sagawa et al., 2019) and upsampling techniques (Liu et al., 2021), aim to balance the data groups in each batch throughout the training process. In contrast, we employ a small balanced warm-up subset only at the beginning of the training. Following a brief period of balanced training, we progressively expand the warm-up subset by adding small random subsets of the remaining training data until using all of them, as shown in the top right of Figure 1. Here, we utilize the momentum from the warm-up subset to prevent the model from learning spurious features when adding new data. Empirical evaluations on both synthetic and real-world benchmark data validate our theoretical findings and confirm the effectiveness of PDE. Additional ablation studies also demonstrate the significance and impact of each component within our training scheme. In summary, our contributions are highlighted as follows:

We provide a theoretical understanding of the impact of spurious correlations beyond the linear setting by considering a two-layer nonlinear CNN.

We introduce PDE, a theory-inspired approach that effectively addresses the challenge posed by spurious correlations.

-PDE achieves the best performance on benchmark vision and language datasets for models including ResNet and Transformer. On average, it outperforms the state-of-the-art method by 2 . 8% in terms of worst-group accuracy.

-PDE enjoys superior training efficiency, being 10 × faster than the state-of-the-art methods.

2 Why is Spurious Correlation Harmful to ERM?

In this section, we simplify the intricate real-world problem of spurious correlations into a theoretical framework. We provide analysis on two-layer nonlinear CNNs, extending beyond the linear setting prevalent in existing literature on this subject. Under this framework, we formally present our theory concerning the training process of empirical risk minimization (ERM) in the presence of spurious features. These theoretical insights motivate the design of our algorithm.

2.1 Empirical Risk Minimization

We begin with the formal definition of the ERM-based training objective for a binary classification problem. Consider a training dataset S = { ( x i , y i ) } N i =1 , where x i ∈ R d is the input and y ∈ {± 1 } is the output label. We train a model f ( x ; W ) with weight W to minimize the empirical loss function:

where ℓ is the logistic loss defined as ℓ ( z ) = log(1 + exp( -z )) . The empirical risk minimizer refers to W ∗ that minimizes the empirical loss: W ∗ := argmin W L ( W ) . Typically, gradient-based optimization algorithms are employed for ERM. For example, at each iteration t , gradient descent (GD) has the following update rule:

Here, η > 0 is the learning rate. In the next subsection, we will show that even for a relatively simple data model which consists of core features and spurious features, vanilla ERM will fail to learn the core features that are correlated to the true label.

2.2 Data Distribution with Spurious Correlation Fails ERM

Previous work such as (Sagawa et al., 2020) considers a data model where the input consists of core feature, spurious feature and noise patches at fixed positions, i.e., x = [ x core , x spu , x noise ] . In real-world applications, however, features in an image do not always appear at the same pixels. Hence, we consider a more realistic data model where the patches do not appear at fixed positions.

Definition 2.1 (Data model) . A data point ( x , y, a ) ∈ ( R d ) P ×{± 1 } × {± 1 } is generated from the distribution D as follows.

Randomly generate the true label y ∈ {± 1 } .

Generate spurious label a ∈ {± y } , where a = y with probability α > 0 . 5 .

Generate x as a collection of P patches: x = ( x (1) , x (2) , . . . , x ( P ) ) ∈ ( R d ) P , where

-Core feature. One and only one patch is given by β c · y · v c with ∥ v c ∥ 2 = 1 .

-Spurious feature. One and only one patch is given by β s · a · v s with ∥ v s ∥ 2 = 1 and ⟨ v c , v s ⟩ = 0 .

Random noise. The rest of the P -2 patches are Gaussian noises ξ that are independently drawn from N (0 , ( σ 2 p /d ) · I d ) with σ p as an absolute constant.

And

0

< β

c

≪

β

s

∈

R

.

Similar data models have also been considered in recent works on feature learning (Allen-Zhu & Li, 2020; Zou et al., 2021; Chen et al., 2022; Jelassi & Li, 2022), where the input data is partitioned into feature and noise patches. We extend their data models by further positing that certain feature patches might be associated with the spurious label instead of the true label. In the rest of the paper, we assume P = 3 for simplicity. With the given data model, we consider the training dataset S = { ( x i , y i , a i ) } N i =1 and let S be partitioned into large group S 1 and small group S 2 such that S 1 contains all the training

Decision

boundary

a=-1

a=+1

data that can be correctly classified by the spurious feature, i.e., a i = y i , and S 2 contains all the training data that can only be correctly classified by the core feature, i.e., a i = -y i . We denote ̂ α = | S 1 | N and therefore 1 -̂ α = | S 2 | N .

Visualization of our data. In Figure 2, we present the visualization in 2 D space of the higherdimensional data generated from our data model using t-SNE (Van der Maaten & Hinton, 2008), where data within each class naturally segregate into large and small groups. The spurious feature is sufficient for accurate classification of the larger group data, but will lead to misclassification of the small group data.

2.3 Beyond Linear Models

We consider a two-layer nonlinear CNN defined as follows:

where w j ∈ R d is the weight vector of the j -th filter, J is the number of filters (neurons) of the network, and σ ( z ) = z 3 is the activation function. W = [ w 1 , . . . , w J ] ∈ R d × J denotes the weight matrix of the CNN. Similar two-layer CNN architectures are analyzed in in (Chen et al., 2022; Jelassi & Li, 2022) but for different problems, where the cubic activation serves a simple function that provides non-linearity. Similar to Jelassi & Li (2022); Cao et al. (2022), we assume a mild overparameterization of the CNN with J = polylog ( d ) . We initialize W (0) ∼ N (0 , σ 2 0 ) , where σ 2 0 = polylog ( d ) d . Due to the CNN structure, our analysis can handle data models where each data can have arbitrary order of patches while linear models fail to do so.

2.4 Understanding the Training Process with Spurious Correlation

In this subsection, we formally introduce our theoretical result on the training process of the two-layer CNN using gradient descent in the presence of spurious features. We first define the performance metrics. A frequenly considered metric is the test accuracy: Acc ( W ) = P ( x ,y,a ) ∼D [ sgn ( f ( x ; W )) = y ] . With spurious correlations, researchers are more interested in the worst-group accuracy:

which accesses the worst accuracy of a model among all groups defined by combinations of y and a . We then summarize the learning process of ERM in the following theorem. Our analysis focuses on the learning of spurious and core features, represented by the growth of ⟨ w ( t ) i , v s ⟩ and ⟨ w ( t ) i , v c ⟩ respectively:

Theorem 2.2. Consider the training dataset S = { ( x i , y i ) } N i =1 that follows the distribution in Definition 2.1. Consider the two-layer nonlinear CNN model as in (2.3) initialized with W (0) ∼ N (0 , σ 2 0 ) . After training with GD in (2.2) for T 0 = ˜ Θ ( 1 / ( ηβ 3 s σ 0 ) ) iterations, for all j ∈ [ J ] and t ∈ [0 , T 0 ) , we have

After training for T 0 iterations, with high probability, the learned weight has the following properties: (1) it learns the spurious feature v s : max j ∈ [ J ] ⟨ w ( T ) j , v s ⟩ ≥ ˜ Ω(1 /β s ) ; (2) it almost does not learn the core feature v c : max j ∈ [ J ] ⟨ w ( T ) j , v c ⟩ = ˜ O ( σ 0 ) .

Discussion. The detailed proof is deferred to Appendix E, and we provide intuitive explanations of the theorem as follows. A larger value of ⟨ w ( t ) i , v ⟩ for v ∈ { v s , v c } implies better learning of the feature vector v by neuron w i at iteration t . As illustrated in (2.4) and (2.5), the updates for both spurious and core features are non-zero, as they depend on the squared terms of themselves with non-zero coefficients, while the growth rate of ⟨ w ( t ) i , v s ⟩ is significantly faster than that of ⟨ w ( t ) i , v c ⟩ . Consequently, the neural network rapidly learns the spurious feature but barely learns the core feature, as it remains almost unchanged from initialization.

̸

We derive the neural network's prediction after training for T 0 iterations. For a randomly generated data example ( x , y, a ) ∼ D , the neural network's prediction is given by sgn ( f ( x ; W ) ) = sgn ( ∑ j ∈ [ J ] ( yβ 3 c ⟨ w j , v c ⟩ 3 + aβ 3 s ⟨ w j , v s ⟩ 3 + ⟨ w j , ξ ⟩ 3 )) . Since the term β 3 s max j ∈ [ J ] ⟨ w j , v s ⟩ 3 dominates the summation, the prediction will be sgn ( f ( x ; W )) = a . Consequently, we obtain the test accuracy as Acc ( W ) = α , since a = y with probability α , and the model accurately classifies the large group. However, when considering the small group and examining examples where y = a , the models consistently make errors, resulting in Acc wg ( W ) = 0 . To circumvent this poor performance on worst-group accuracy, an algorithm that can avoid learning the spurious feature is in demand.

3 Theory-Inspired Two-Stage Training Algorithm

In this section, we introduce Progressive Data Expansion (PDE), a novel two-stage training algorithm inspired by our analysis to enhance robustness against spurious correlations. We begin with illustrating

the implications of our theory, where we provide insights into the data distributions that lead to the rapid learning of spurious features and clarify scenarios under which the model remains unaffected.

3.1 Theoretical Implications

Notably in Theorem 2.2, the growth of the two sequences ⟨ w ( t ) i , v s ⟩ in (2.4) and ⟨ w ( t ) i , v c ⟩ in (2.5) follows the formula x t +1 = x t + ηAx 2 t , where x t represents the inner product sequence with regard to iteration t and A is the coefficient containing ̂ α , β c or β s . This formula is closely related to the analysis of tensor power methods (Allen-Zhu & Li, 2020). In simple terms, when two sequences have slightly different growth rates, one of them will experience much faster growth in later times. As we will show below, the key factors that determine the drastic difference of spurious and core features in later times are the group size ̂ α and feature strengths β c , β s .

When the model learns spurious feature ( β 3 c < β 3 s (2 ̂ α -1) ). We examine the lower bound for the growth of ⟨ w ( t ) i , v s ⟩ in (2.4) and the upper bound for the growth of ⟨ w ( t ) i , v c ⟩ in (2.5) in Theorem 2.2. If β 3 c < β 3 s (2 ̂ α -1) , we can employ the tensor power method and deduce that the spurious feature will be learned first and rapidly. The condition on data distribution imposes two necessary conditions: ̂ α > 1 / 2 (groups are imbalanced) and β c < β s (the spurious feature is stronger). This observation is consistent with real-world datasets, such as the Waterbirds dataset, where α = 0 . 95 and the background is much easier to learn than the intricate features of the birds.

̂ · When the model learns core feature ( β c > β s ). However, if we deviate from the aforementioned conditions and consider β c > β s , we can examine the lower bound for the growth of ⟨ w ( t ) i , v c ⟩ in (2.5) and the upper bound for the growth of ⟨ w ( t ) i , v s ⟩ in (2.4). Once again, we apply the tensor power method and determine that the model will learn the core feature rapidly. In real-world datasets, this scenario corresponds to cases where the core feature is not only significant but also easier to learn than the spurious feature. Even for imbalanced groups with ̂ α > 1 / 2 , the model accurately learns the core feature. Consequently, enhancing the coefficients of the growth of the core feature allows the model to tolerate imbalanced groups. We present verification through synthetic experiments in the next section.

As we will show in the following subsection, we initially break the conditions of learning the spurious feature by letting ̂ α = 1 / 2 in a group-balanced data subset. Subsequently, we utilize the momentum to amplify the core feature's coefficient, allowing for tolerance of ̂ α > 1 / 2 when adding new data.

3.2 PDE: A Two-Stage Training Algorithm

We present a new algorithm named Progressive Data Expansion (PDE) in Algorithm 1 and explain the details below, which consist of (1) warm-up and (2) expansion stages.

Algorithm 1 Progressive Data Expansion (PDE)

Require: Number of iterations T 0 for warm-up training; number of times K for dataset expansion; number of iterations J for expansion training; number of data m for each expansion; learning rate η ; momentum coefficient γ ; initialization scale σ 0 ; training set S = { ( x i , y i , a i ) } n ; model f W .

i =1 1: Initialize W (0) . Warm-up stage 2: Divide the S into groups by values of y and a : S y,a = { ( x i , y i , a i ) } y i = y,a i = a . 3: Generate warm-up set S 0 from S by randomly subsampling from each group of S such that | S 0 y,a | = min y ′ ,a ′ | S y ′ ,a ′ | for y ∈ {± 1 } and a ∈ {± 1 } . 4: for t = 0 , 1 , . . . , T 0 do 5: Compute loss on S 0 : L S 0 ( W ( t ) ) = 1 | S 0 | ∑ i ∈ S 0 ℓ ( y i f ( x i ; W ( t ) )) . 6: Update W ( t +1) by (3.1) and (3.2). 7: end for Expansion stage 8: for k = 1 , . . . , K do 9: Draw m examples ( S [ m ] ) from S/S k -1 and let S k = S k -1 ∪ S [ m ] . 10: for t = 1 , . . . , J do 11: Compute loss on S k : L S k ( W ( T ) ) = 1 | S k | ∑ i ∈ S k ℓ ( y i f ( x i ; W ( T ) )) , where T = T 0 +( k -1) ∗ J + t . 12: Update W ( T +1) by (3.1) and (3.2). 13: end for 14: end for 15: return W ( t ) = argmax W ( t ′ ) Acc val wg ( W ( t ′ ) ) .

As accelerated gradient methods are most commonly used in applications, we jointly consider the property of momentum and our theoretical insights when designing the algorithm. For gradient descent with momentum (GD+M), at each iteration t and with momentum coefficient γ > 0 , it updates as follows

Warm-up Stage. In this stage, we create a fully balanced dataset S 0 , in which each group is randomly subsampled to match the size of the smallest group, and consider it as a warm-up dataset. We train the model on the warm-up dataset for a fixed number of epochs. During this phase, the model is anticipated to accurately learn the core feature without being influenced by the spurious feature. Note that, under our data model, a completely balanced dataset will have ̂ α = 1 / 2 . We present the following lemma as a theoretical basis for the warm-up stage.

Lemma 3.1. Given the balanced training dataset S 0 = { ( x i , y i , a i ) } N 0 i =1 with ̂ α = 1 / 2 as in Definition 2.1 and CNN as in (2.3). The gradient on v s will be 0 from the beginning of training.

In particular, with ̂ α = 1 / 2 we have | S 0 1 | = | S 0 2 | : an equal amount of data is positively correlated with the spurious feature as the data negatively correlated with the spurious feature. In each update, both groups contribute nearly the same amount of spurious feature gradient with different signs, resulting in cancellation. Ultimately, this prevents the model from learning the spurious feature. Detailed proofs can be found in Appendix F.

Expansion Stage. In this stage, we proceed to train the model by incrementally incorporating new data into the training dataset. The rationale for this stage is grounded in the theoretical result by the previous work (Jelassi & Li, 2022) on GD with momentum, which demonstrates that once gradient descent with momentum initially increases its correlation with a feature v , it retains a substantial historical gradient in the momentum containing v . Put it briefly, the initial learning phase has a considerable influence on subsequent training for widely-used accelerated training algorithms. While ERM learns the spurious feature v s and momentum does not help, as we will show in synthetic experiments, PDE avoids learning v s and learns v c in the warm-up stage. This momentum from warm-up, in turn, amplifies the core feature that is present in the gradients of newly added data, facilitating the continued learning of v c in the expansion stage. For a specific illustration, the learning of the core feature by GD+M will be

where g ( t ) is the additional momentum as compared to GD with γ = 0 . While the current gradient along v c might be small (i.e., β c ), we can benefit from the historical gradient in g ( t ) to amplify the growth of ⟨ w ( t +1) j , v c ⟩ and make it larger than that of the spurious feature (i.e., β s ). This learning process will then correspond to the case when the model learns the core feature discussed in Subsection 3.1. Practically, we consider randomly selecting m new examples for expansion every J epochs by attempting to draw a similar number of examples from each group. During the last few epochs of the expansion stage, we expect the newly incorporated data exclusively from the larger group, as the smaller groups have been entirely integrated into the warm-up dataset.

It is worth noting that while many works address the issue of identifying groups from datasets containing spurious correlations, we assume the group information is known and our algorithm focuses on the crucial subsequent question of optimizing group information utilization. Aiming to prevent the learning of spurious features, PDE distinguishes itself by employing a rapid and lightweight warm-up stage and ensuring continuous improvement during the expansion stage with the momentum acquired from the warm-up dataset. Our training framework is both concise and effective, resulting in computational efficiency and ease of implementation.

4 Experiments

In this section, we present the experiment results from both synthetic and real datasets. Notably, we report the worst-group accuracy, which assesses the minimum accuracy across all groups and is commonly used to evaluate the model's robustness against spurious correlations.

4.1 Synthetic Data

In this section, we present synthetic experiment results in verification to our theoretical findings. In Appendix A, we illustrate the detailed data distribution, hyper-parameters of the experiments and more extensive experiment results. The data used in this seciton is generated following Definition 2.1. We consider the worst-group and overall test accuracy. As illustrated in Table 1, ERM, whether trained with GD or GD+M, is unable to accurately predict the small group in our specified data distribution where ̂ α = 0 . 98 and β c < β s . In contrast, our method significantly improves worst-group accuracy while maintaining

overall test accuracy comparable to ERM. Furthermore, as depicted in Figure 3(a), ERM rapidly learns the spurious feature as it minimizes the training loss, while barely learning the core feature. Meanwhile, in Figure 3(b) we show the learning of ERM when the data distribution breaks the conditions of our theory and has β c > β s instead. Even with the same ̂ α as in Figure 3(a), ERM correctly learns the core feature despite the imbalanced group size. These two figures support the theoretical results we discussed to motivate our method. Consequently, on the same training dataset as in Figure 3(a), Figure 3(c) shows that our approach allows the model to initially learn the core feature using the warm-up dataset and continue learning when incorporating new data.

0

100

200

300

400

0.0

0.5

1.0

1.5

j

[

J

]

w

(

t

)

j

,

v

s

j

[

J

]

w

(

t

)

j

,

v

c

training loss

(a) ERM (case 1)

0

100

200

300

400

0.0

0.5

1.0

1.5

(b) ERM (case 2)

Expansion stage

(c) PDE

4.2 Real Data

We conduct experiments on real benchmark datasets to: (1) compare our approach with state-ofthe-art methods, highlighting its superior performance and efficiency, and (2) offer insights into the design of our method through ablation studies.

Datasets. We evaluate on three wildly used datasets across vision and language tasks for spurious correlation: (1) Waterbirds (Sagawa et al., 2019) contains bird images labeled as waterbird or landbird, placed against a water or land background, where the smallest subgroup is waterbirds on land background. (2) CelebA (Liu et al., 2015) is used to study gender as the spurious feature for hair color classification, and the smallest group in this task is blond-haired males. (3) CivilCommentsWILDS (Koh et al., 2021b) classifies toxic and non-toxic online comments while dealing with demographic information. It creates 16 overlapping groups for each of the 8 demographic identities.

Baselines. We compare our proposed algorithm against several state-of-the-art methods. Apart from standard ERM, we include GroupDRO (Sagawa et al., 2019) and DFR (Kirichenko et al., 2023) that assume access to the group labels. We only include DFR Tr for fair comparison as all methods only use the training dataset to train and finetune the model, while DFR Val also finetunes on validation data. We also design a baseline called subsample that simply trains the model on the warm-up dataset only. Additionally, we evaluate three recent methods that address spurious correlations without the need for group labels: LfF (Nam et al., 2020), EIIL (Creager et al., 2021), and JTT (Liu et al., 2021). We

report results for ERM, Subsample, GroupDRO and PDE based on our own runs using the WILDS library Koh et al. (2021a); for others, we directly reuse their reported numbers.

We present the experiment details including dataset statistics and hyperparameters as well as comprehensive additional experiments in Appendix B.

4.2.1 Consistent Superior Worst-group Performance

We assess PDE on the mentioned datasets with the state-of-the-art methods. Importantly, we emphasize the comparison with GroupDRO, as it represents the best-performing method that utilizes group information. As shown in Table 2, PDE considerably enhances the worst-performing group's performance across all datasets, while maintaining the average accuracy comparable to GroupDRO. Remarkably, although GroupDRO occasionally fails to surpass other methods, PDE's performance consistently exceeds them in worst-group accuracy.

4.2.2 Efficient Training

In this subsection, we show that our method is more efficient as it does not train a model twice (as in JTT) and more importantly avoids the necessity for a small learning rate (as in GroupDRO). Specifically, methods employing group-balanced batches like GroupDRO require a very small learning rate coupled with a large weight decay in practice. We provide an intuitive explanation as follows. When sampling to achieve balanced groups in each batch, smaller groups appear more frequently than larger ones. If training progresses rapidly, the loss on smaller groups will be minimized quickly, while the majority of the large group data remains unseen and contributes to most of the gradients in later batches. Therefore, these methods necessitate slow training to ensure the model encounters diverse data from larger groups before completely learning the smaller groups. We validate this observation in Table 3, where GroupDRO trained faster than the default results in significantly poorer performance similar to ERM. Conversely, PDE can be trained to converge rapidly on the warm-up set and reaches better worst-group accuracy 10 × faster than GroupDRO at default. Note that methods which only finetune the last layer (Kirichenko et al., 2023; Wei et al., 2023) are also efficient. However, they still require training a model first using ERM on the entire training data till convergence. In contrast, PDE does not require to further finetune the model.

4.2.3 Understanding the Two Stages

We examine each component and demonstrate their effect in PDE. In Table 4, we present the worst-group and average accuracy of the model trained following the warm-up and expansion stages. Indeed, the majority of learning occurs in the warm-up stage, during which a satisfactory worst-group accuracy is established. In the expansion stage, the model persists in learning new data along the established trajectory, leading to continued performance improvement. In Figure 5,

we corroborate and emphasize that the model has acquired the appropriate features and maintains learning based on its historical gradient stored in the momentum by presenting the following evidence. As shown, if the optimizer is reset after the warm-up stage and loses all its historical gradients, it rapidly acquires spurious features, resulting in a swift decline in performance accuracy as shown in the blue line.

0.800

0.825

0.850

0.875

0.900

0.925

0.950

0.975

1.000

Worst-group accuracy

Expansion stage

With reinitialization

Without reinitialization

(a) Worst-group accuracy.

0.86

0.88

0.90

0.92

0.94

0.96

0.98

1.00

Average accuracy

Expansion stage

With reinitialization

Without reinitialization

(b) Average accuracy.

4.2.4 Ablation Study on the Hyper-parameters of PDE

PDE is robust within a reasonable range of hyperparemeter choices, although some configurations outperform others. As shown in Table 5, it is necessary to limit the number of data points introduced during each expansion to prevent performance degradation. Similarly, in Appendix A, we emphasize the importance of gradual data expansion. In Table 6, we show that post-warmup learning rate decay is essential, though PDE exhibits a tolerance to the degree of this decay. Lastly, as illustrated in Figure 5, adopting a smaller learning rate often necessitates increased data expansions. Nonetheless, a reduced learning rate does not necessarily lead to improved performance.

140

160

180

200

220

240

260

280

300

Epoch

0.6

0.7

0.8

0.9

1.0

Worst-group accuracy

Expansion LR = 1E-3

Expansion LR = 1E-4

Expansion LR = 1E-5

140

160

180

200

220

240

260

280

300

Epoch

0.80

0.85

0.90

0.95

1.00

Average accuracy

Expansion LR = 1E-3

Expansion LR = 1E-4

Expansion LR = 1E-5

(a) Worst-group accuracy.

(b) Average accuracy.

5 Related Work

Existing approaches for improving robustness against spurious correlations can be categorized into two lines of research based on the tackled subproblems. A line of research focuses on the same subproblem we tackle: effectively using the group information to improve robustness. With group information, one can use the distributionally robust optimization (DRO) framework and dynamically increase the weight of the worst-group loss in minimization (Hu et al., 2018; Oren et al., 2019; Sagawa et al., 2019; Zhang et al., 2021). Within this line of work, GroupDRO (Sagawa et al., 2019) achieves state-of-the-art performances across multiple benchmarks. Other approaches use importance weighting to reweight the groups (Shimodaira, 2000; Byrd & Lipton, 2019; Xu et al., 2021) and class balancing to downsample the majority or upsample the minority (He & Garcia, 2009; Cui et al., 2019; Sagawa et al., 2020). Alternatively, Goel et al. (2021) leverage group information to augment the minority groups with synthetic examples generated using GAN. Another strategy (Cao et al., 2019, 2020) involves imposing Lipschitz regularization around minority data points. Most recently, methods that train a model using ERM first and then only finetune the last layer on balanced data from training or validation (Kirichenko et al., 2023) or learn post-doc scaling adjustments (Wei et al., 2023) are shown to be effective.

The other line of research focuses on the setting where group information is not available during training and tackles the first subproblem we identified as accurately finding the groups. Recent notable works (Nam et al., 2020; Liu et al., 2021; Creager et al., 2021; Zhang et al., 2022; Yang et al., 2023a) mostly involve training two models, one of which is used to find group information. To finally use the found groups, many approaches (Namkoong & Duchi, 2017; Duchi et al., 2019; Oren et al., 2019; Sohoni et al., 2020) still follow the DRO framework.

The first theoretical analysis of spurious correlation is provided by Sagawa et al. (2020). For selfsupervised learning, Chen et al. (2020) shows that fine-tuning with pre-trained models can reduce the harmful effects of spurious features. Ye et al. (2022) provides guarantees in the presence of label noise that core features are learned well only when less noisy than spurious features. These theoretical work only provides analyses of linear models. Meanwhile, a parallel line of work has established theoretical analysis of nonlinear CNNs in the more realistic setting Allen-Zhu & Li (2020); Zou et al. (2021); Wen & Li (2021); Chen et al. (2022); Jelassi & Li (2022). Our work builds on this line of research and generalizes it to the study of spurious features. Lastly, we notice that a concurrent work (Chen et al., 2023) also uses tensor power method (Allen-Zhu & Li, 2020) to analyze the learning of spurious features v.s. invariant features, but in the setting of out-of-distribution generalization.

6 Conclusion

In conclusion, this paper addressed the challenge of spurious correlations in training deep learning models and focused on the most effective use of group information to improve robustness. We provided a theoretical analysis based on a simplified data model and a two-layer nonlinear CNN. Building upon this understanding, we proposed PDE, a novel training algorithm that effectively and efficiently enhances model robustness against spurious correlations. This work contributes to both the theoretical understanding and practical application of mitigating spurious correlations, paving the way for more reliable and robust deep learning models.

Limitations and future work. Although beyond the linear setting, our analysis still focuses on a relatively simplified binary classification data model. To better represent real-world application scenarios, future work could involve extending to multi-class classification problems and examining the training of transformer architectures. Practically, our proposed method requires the tuning of additional hyperparameters, including the number of warm-up epochs, the number of times for dataset expansion and the number of data to be added in each expansion.

References

Ahmed, F., Bengio, Y., Van Seijen, H., and Courville, A. Systematic generalisation with group invariant predictions. In International Conference on Learning Representations , 2021.

Allen-Zhu, Z. and Li, Y. Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. arXiv preprint arXiv:2012.09816 , 2020.

Byrd, J. and Lipton, Z. What is the effect of importance weighting in deep learning? In International conference on machine learning , pp. 872-881. PMLR, 2019.

Cao, K., Wei, C., Gaidon, A., Arechiga, N., and Ma, T. Learning imbalanced datasets with labeldistribution-aware margin loss. Advances in neural information processing systems , 32, 2019.

Cao, K., Chen, Y., Lu, J., Arechiga, N., Gaidon, A., and Ma, T. Heteroskedastic and imbalanced deep learning with adaptive regularization. arXiv preprint arXiv:2006.15766 , 2020.

Cao, Y., Chen, Z., Belkin, M., and Gu, Q. Benign overfitting in two-layer convolutional neural networks. Advances in neural information processing systems , 35:25237-25250, 2022.

Chen, Y., Wei, C., Kumar, A., and Ma, T. Self-training avoids using spurious features under domain shift. Advances in Neural Information Processing Systems , 33:21061-21071, 2020.

Chen, Y., Huang, W., Zhou, K., Bian, Y., Han, B., and Cheng, J. Towards understanding feature learning in out-of-distribution generalization. arXiv preprint arXiv:2304.11327 , 2023.

Chen, Z., Deng, Y., Wu, Y., Gu, Q., and Li, Y. Towards understanding the mixture-of-experts layer in deep learning. Advances in neural information processing systems , 2022.

Creager, E., Jacobsen, J.-H., and Zemel, R. Environment inference for invariant learning. In International Conference on Machine Learning , pp. 2189-2200. PMLR, 2021.

Cui, Y., Jia, M., Lin, T.-Y., Song, Y., and Belongie, S. Class-balanced loss based on effective number of samples. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pp. 9268-9277, 2019.

Duchi, J. C., Hashimoto, T., and Namkoong, H. Distributionally robust losses against mixture covariate shifts. Under review , 2:1, 2019.

Goel, K., Gu, A., Li, Y., and Re, C. Model patching: Closing the subgroup performance gap with data augmentation. In International Conference on Learning Representations , 2021.

Haghtalab, N., Jordan, M., and Zhao, E. On-demand sampling: Learning optimally from multiple distributions. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), Advances in Neural Information Processing Systems , 2022.

He, H. and Garcia, E. A. Learning from imbalanced data. IEEE Transactions on knowledge and data engineering , 21(9):1263-1284, 2009.

He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition , pp. 770-778, 2016.

Hu, W., Niu, G., Sato, I., and Sugiyama, M. Does distributionally robust supervised learning give robust classifiers? In International Conference on Machine Learning , pp. 2029-2037. PMLR, 2018.

Izmailov, P., Kirichenko, P., Gruver, N., and Wilson, A. G. On feature learning in the presence of spurious correlations. Advances in Neural Information Processing Systems , 35:38516-38532, 2022.

Jelassi, S. and Li, Y. Towards understanding how momentum improves generalization in deep learning. In International Conference on Machine Learning , pp. 9965-10040. PMLR, 2022.

Kirichenko, P., Izmailov, P., and Wilson, A. G. Last layer re-training is sufficient for robustness to spurious correlations. In International Conference on Learning Representations , 2023.

Koh, P. W., Sagawa, S., Marklund, H., Xie, S. M., Zhang, M., Balsubramani, A., Hu, W., Yasunaga, M., Phillips, R. L., Gao, I., Lee, T., David, E., Stavness, I., Guo, W., Earnshaw, B. A., Haque, I. S., Beery, S., Leskovec, J., Kundaje, A., Pierson, E., Levine, S., Finn, C., and Liang, P. WILDS: A benchmark of in-the-wild distribution shifts. In International Conference on Machine Learning (ICML) , 2021a.

Koh, P. W., Sagawa, S., Marklund, H., Xie, S. M., Zhang, M., Balsubramani, A., Hu, W., Yasunaga, M., Phillips, R. L., Gao, I., et al. Wilds: A benchmark of in-the-wild distribution shifts. In International Conference on Machine Learning , pp. 5637-5664. PMLR, 2021b.

Liu, E. Z., Haghgoo, B., Chen, A. S., Raghunathan, A., Koh, P. W., Sagawa, S., Liang, P., and Finn, C. Just train twice: Improving group robustness without training group information. In International Conference on Machine Learning , pp. 6781-6792. PMLR, 2021.

Liu, Z., Luo, P., Wang, X., and Tang, X. Deep learning face attributes in the wild. In Proceedings of the IEEE international conference on computer vision , pp. 3730-3738, 2015.

Nam, J., Cha, H., Ahn, S., Lee, J., and Shin, J. Learning from failure: De-biasing classifier from biased classifier. Advances in Neural Information Processing Systems , 33:20673-20684, 2020.

Namkoong, H. and Duchi, J. C. Variance-based regularization with convex objectives. Advances in neural information processing systems , 30, 2017.

Oren, Y., Sagawa, S., Hashimoto, T. B., and Liang, P. Distributionally robust language modeling. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP) , pp. 4227-4237, Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10.18653/v1/D19-1432.

Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S. Pytorch: An imperative style, high-performance deep learning library. In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, F., Fox, E., and Garnett, R. (eds.), Advances in Neural Information Processing Systems 32 , pp. 8024-8035. 2019.

Sagawa, S., Koh, P. W., Hashimoto, T. B., and Liang, P. Distributionally robust neural networks. In International Conference on Learning Representations , 2019.

Sagawa, S., Raghunathan, A., Koh, P. W., and Liang, P. An investigation of why overparameterization exacerbates spurious correlations. In International Conference on Machine Learning , pp. 83468356. PMLR, 2020.

Selvaraju, R. R., Cogswell, M., Das, A., Vedantam, R., Parikh, D., and Batra, D. Grad-cam: Visual explanations from deep networks via gradient-based localization. In Proceedings of the IEEE international conference on computer vision , pp. 618-626, 2017.

Shimodaira, H. Improving predictive inference under covariate shift by weighting the log-likelihood function. Journal of statistical planning and inference , 90(2):227-244, 2000.

Sohoni, N., Dunnmon, J., Angus, G., Gu, A., and Ré, C. No subclass left behind: Fine-grained robustness in coarse-grained classification problems. Advances in Neural Information Processing Systems , 33:19339-19352, 2020.

Taghanaki, S. A., Choi, K., Khasahmadi, A. H., and Goyal, A. Robust representation learning via perceptual similarity metrics. In International Conference on Machine Learning , pp. 10043-10053. PMLR, 2021.

Van der Maaten, L. and Hinton, G. Visualizing data using t-sne. Journal of machine learning research , 9(11), 2008.

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. Attention is all you need. In Advances in neural information processing systems , pp. 5998-6008, 2017.

Wah, C., Branson, S., Welinder, P., Perona, P., and Belongie, S. The caltech-ucsd birds-200-2011 dataset. 2011.

Wei, J., Narasimhan, H., Amid, E., Chu, W.-S., Liu, Y., and Kumar, A. Distributionally robust post-hoc classifiers under prior shifts. In International Conference on Learning Representations (ICLR) , 2023.

Wen, Z. and Li, Y. Toward understanding the feature learning process of self-supervised contrastive learning. In International Conference on Machine Learning , pp. 11112-11122. PMLR, 2021.

Wolf, T., Debut, L., Sanh, V., Chaumond, J., Delangue, C., Moi, A., Cistac, P., Rault, T., Louf, R., Funtowicz, M., Davison, J., Shleifer, S., von Platen, P., Ma, C., Jernite, Y., Plu, J., Xu, C., Scao, T. L., Gugger, S., Drame, M., Lhoest, Q., and Rush, A. M. Transformers: State-of-theart natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations , pp. 38-45, Online, October 2020. Association for Computational Linguistics. URL https://www.aclweb.org/anthology/2020. emnlp-demos.6 .

Xu, D., Ye, Y., and Ruan, C. Understanding the role of importance weighting for deep learning. In International Conference on Learning Representations , 2021.

Yang, Y., Gan, E., Dziugaite, G. K., and Mirzasoleiman, B. Identifying spurious biases early in training through the lens of simplicity bias. arXiv preprint arXiv:2305.18761 , 2023a.

Yang, Y., Nushi, B., Palangi, H., and Mirzasoleiman, B. Mitigating spurious correlations in multimodal models during fine-tuning. In International Conference on Machine Learning , 2023b.

Yang, Y., Zhang, H., Katabi, D., and Ghassemi, M. Change is hard: A closer look at subpopulation shift. arXiv preprint arXiv:2302.12254 , 2023c.

Yang, Y.-Y., Chou, C.-N., and Chaudhuri, K. Understanding rare spurious correlations in neural networks. arXiv preprint arXiv:2202.05189 , 2022.

Ye, H., Zou, J., and Zhang, L. Freeze then train: Towards provable representation learning under spurious correlations and feature noise. arXiv preprint arXiv:2210.11075 , 2022.

Zhang, J., Menon, A. K., Veit, A., Bhojanapalli, S., Kumar, S., and Sra, S. Coping with label shift via distributionally robust optimisation. In International Conference on Learning Representations , 2021.

Zhang, M., Sohoni, N. S., Zhang, H. R., Finn, C., and Ré, C. Correct-n-contrast: A contrastive approach for improving robustness to spurious correlations. arXiv preprint arXiv:2203.01517 , 2022.

Zhou, B., Lapedriza, A., Khosla, A., Oliva, A., and Torralba, A. Places: A 10 million image database for scene recognition. IEEE Transactions on Pattern Analysis and Machine Intelligence , 2017.

Zou, D., Cao, Y., Li, Y., and Gu, Q. Understanding the generalization of adam in learning neural networks with proper regularization. arXiv preprint arXiv:2108.11371 , 2021.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper describes Progressive Data Expansion (PDE), a new method for training models which are robust to spurious features. The idea of the method is to split the training into two phases: warmup and expansion. In the warmup stage, the dataset is balanced, and the model learns a classifier that ignores the spurious feature. In the expansion stage, the remaining data is used to finetune. The authors also provide a theoretical argument about feature learning, and show strong empirical performance on benchmark datasets.\n\n**W** &mdash; weakness,  **S** &mdash; strength, **Q** &mdash; question.\n\nSTRENGTHS:\n**S1.** The authors provide a theoretical analysis, that provides some intuition for the method\n**S2.** The proposed method is simple and computationally cheap\n**S3.** The authors achieve strong empirical performance across three benchmark datasets\n\nWEAKNESSES:\n**Q1/W1. Relationship to DFR**\n\nThe works [1, 2] advocated for an approach that is almost opposite to PDE: they train a model with standard ERM, and then retrain only the last layer on a group-balanced dataset. PDE on the other hand pretrains the model on a group-balanced dataset, and then finetunes with standard ERM. The methods achieve similar performance, if we consider DFR^Val, which is the main method in [1] (the authors argue that DFR^Val uses the validation data, but I believe PDE also uses the same amount of validation data to tune the hyper-parameters?). In fact, PDE achieves better performance on CelebA and CivilComments, while DFR achieves better performance on waterbirds.\n\nHowever, conceptually [1, 2] argue that in fact standard ERM learns the core features as well as GroupDRO on the standard benchmarks, and once the last layer is retrained we can recover the performance of GroupDRO. This intuition is quite the opposite of the intuition presented in this paper, which is that ERM does not learn the core features without intervention. I think it would be good if the authors could comment on this distinction.\n\nMy intuition is that the pretraining in the warmup stage prevents the model from learning the spurious feature, and makes it harder for the model to use the spurious features in the expansion stage. Alternatively, it could be that the model still represents the spurious features, but because of the momentum, it is unable to switch to using this feature in the expansion phase. It would be interesting to explore more deeply which of these options is happening.\n\nLastly, in line 287 you mention that the method is more efficient than DFR, because it does not train the model twice. However, DFR only retrains the last layer of a pretrained ERM model, which is computationally cheap, so I don't believe this statement is correct. \n\n **Q2/W2. Robustness to the choice of hyper-parameters**\n\nIt is important to understand how sensitive the method is to the choice of hyper-parameters. In particular, I would be interested in the impact of\n- Length of the expansion training phase\n- Learning rate in the expansion phase\n\nMy intuition is that the method continues going in a good direction for some time during the expansion phase, but then eventually it should converge to an optimal solution on the full data distribution, i.e. maximize the average accuracy and not the worst group accuracy. However, due to early stopping, the authors are able to achieve good performance, before the optimization learns to use the spurious feature. Is this intuition correct?\n\n**W3. Theoretical analysis**\n\nThe theoretical analysis is quite toy. The authors use a model which is basically a bag of features, i.e. the result is a sum of cubed dot-products of features and \"filters\". This model is not very non-linear. However, it is understandable, given that theoretical analysis of non-linear models is challenging.\n\nQUESTIONS:\n**Q3. Expansion phase:** In line 225 you mention that during the expansion phase you try to sample data uniformly across groups. What do you do exactly in this phase? Isn't all of the minority data used in the warmup phase?\n\n**Q4. Intuition for expansion phase:** Related to Q3, why is performance on the training data distribution go down in Fig 4 (b), even when you reset the momentum variables? Is the training loss going up then? Why would this happen? Or is the data not just sampled from the training distribution early in the expansion phase?\n\nLIMITATIONS:\n**References**\n\n[1] [_Last Layer Re-Training is Sufficient for Robustness to Spurious Correlations_](https://openreview.net/forum?id=Zb6c8A-Fghk);\nP. Kirichenko, P. Izmailov, A. G. Wilson;\nICLR 2023\n\n[2] [_On Feature Learning in the Presence of Spurious Correlations_](https://proceedings.neurips.cc/paper_files/paper/2022/hash/fb64a552feda3d981dbe43527a80a07e-Abstract-Conference.html);\nP. Izmailov, P. Kirichenko, N. Gruver, A. G. Wilson;\nNeurIPS 2022", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the author propose a learning framework for improving the performance of classifiers on the worst group in the presence of spurious features. The authors focus on a two-layer convolutional neural network. They first show that imbalanced data and easy-to-learn spurious features can lead to bias in the classifier. Then, they propose a learning procedure that focuses on balanced data and gradually progresses how the model learns from core features to all features. The new method is evaluated on several datasets demonstrating that the new approach is fast and can improve accuracy on worst groups.\n\nSTRENGTHS:\nThe paper is mostly well-written and easy to follow, and the idea makes sense. The proposed approach is motivated theoretically using a two-layer CNN model. The method is simple and relies on a two-stage training procedure, therefore can be easily implemented and adapted to other algorithms. \nThe proposed method improves the performance of the worst group in all evaluated examples; in some cases, it leads to significant improvements.\n\nWEAKNESSES:\nThe theoretical analysis in the paper is focused on binary classification. The method requires tuning several hyperparameters. While the worst group performance is improved in most cases the overall accuracy degradation is non negligible. The method is only evaluated on a small number of datasets.\n\nQUESTIONS:\nHow is the method applied to non-binary data? \nThe analysis focuses on the activation z^3. Can this be replaced with more commonly used activations?\nWhy is there a significant drop in average accuracy in some examples?\nSome parts are not well explained; for example, the caption in Figure 1 is unclear.\n\nLIMITATIONS:\nLimitations are discussed in the last section of the paper. Addressing some of the limitations presented by the authors is indeed important to increase the usability of the proposed approach.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper works on understanding and mitigating the impacts on spurious features. Specifically, the authors provide theoretical analysis on the learning process of a non-linear two-layer CNN, under spurious features. Theoretical insights reveal the need to start with balanced data, and progressively expand the train set. The proposed PDE is right based on such insights, and is demonstrated to be efficient and effective on certain datasets.\n\nSTRENGTHS:\n* (1) The authors provide theoretical insights on the learning process of a two-layer non-linear CNN under spurious features, which aligns with empirical observations.\n\n* (2) The proposed method starts with group-balanced training set and then gradually injects more data samples to capture the core features. And its effectiveness and efficiency are demonstrated on several popular datasets when compared with exiting methods.\n\n* (3) The paper is clearly presented and well organized, and insights/motivations are easy to follow.\n\nWEAKNESSES:\nAlthough the proposed PDE is super efficient by referring to existing learning/training based approaches, it would be better to mention/discuss the efficiency of some fine-tuning or post-hoc approaches during the related work section, for example, fine-tune the last layer is sufficient to spurious features [R1], post-hoc adjustment on model prediction improves robustness to spurious features [R2].\n\nReferences:\n\nR1: Last Layer Re-Training is Sufficient for Robustness to Spurious Correlations. [ICLR'23]\n\nR2: Distributionally Robust Post-hoc Classifiers under Prior Shifts. [ICLR'23]\n\nQUESTIONS:\n* (1) The results on the synthetic data is a bit confusing: in table 1, the performance of ERM on worst-group is 0%; besides, PDE has a better worst-group performance than the overall one, which is kind of conterintuitive.\n\n* (2) It would be much better if authors could add more discussions on the role and selection/suggestions of certain parameters, i.e., the number of times for dataset expansion and the number of data to be added in each expansion.\n\nLIMITATIONS:\nThe authors clearly stated their limitations at the end of paper, \n\n* The theoretical analysis rely on simplified model;\n\n* Certain hyper-parameters might be crucial in model training: i.e., the number of warmup epochs, the number of times for dataset expansion and the number of data to be added in each expansion.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a new method called PDE to address spurious features and improve generalization. Building on the existing literature, the authors consider an imbalanced binary classification problem where two types of features coexist: core features and spurious features. The core features can lead to good generalization performance but a classification model may pick up spurious features. To avoid that, the proposed method (1) starts with a downsampled balanced classification as a warm-up, which will not pick up the spurious features, and (2) then gradually increase the imbalance ratio until the full dataset is used for training.\n\nThe main justification for this two-stage algorithm is that gradient descent with momentum can benefit from the warm-up stage where no spurious features are picked up. Numerous experiments on synthetic data and real data are used to support the proposed method.\n\nSTRENGTHS:\nClassification with spurious features is a very practical and relevant problem. The proposed method is easy to implement and new to my knowledge. The experiments result show promising performance compared with alternatives.\n\nOverall, the warm-up stage in optimization seems to be an interesting idea. It starts with an \"easy\" downsampled dataset without data imbalance---which is known to be the culprit of picking up spurious features. Then, in the second stage, we hope to use the full datasets without picking up spurious features by using momentum.\n\nThe presentation of this paper is clear and theoretical/experimental results are easy to follow.\n\nWEAKNESSES:\nIn my opinion, this paper misses a crucial element in classification with spurious features---that is, the effect of overparametrization. It is known, for example in [1], that there is a very important difference between underparameterization and overparametrization when spurious features are present. It is, therefore, expected to have analyses on the inductive bias of the proposed method and how it may select a superior solution compared with ERM and recent alternatives. Unfortunately, the effect of overparametrization is absent in the analyses, as the number of neurons $J$ does not play a role in the analysis.\n\nBesides, while the two-stage optimization idea looks interesting, the analysis is a bit handwaving, so it is unconvincing to me why momentum method is able to keep avoiding spurious features in the second stage. An investigation of inductive bias would make this paper much stronger; perhaps even some in-depth experiments on synthetic datasets would shed lights on the behavior of momentum in terms on avoiding spurious features.\n\nLastly, I think the experimental results look fine, but I am not convinced that the proposed method is better than the recent alternatives. For example, in Table 2, while PDE achieves better accuracy on worst classes, it usually incurs a worse accuracy on average. Since there is a clear tradeoff between majority and minority classes in binary classification, better accuracy for the minority class can be often trivially improved by shifting the decision boundary. It would be more convincing if empirical results show improvements for *all* classes.\n\n\n[1] Sagawa, S., Raghunathan, A., Koh, P. W., and Liang, P. An investigation of why overparameterization exacerbates spurious correlations. In International Conference on Machine Learning, pp. 8346– 8356. PMLR, 2020.\n\nQUESTIONS:\nSee the weaknesses section.\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-07"}