SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yihe Deng* Yu Yang* Baharan Mirzasoleiman Quanquan Gu
{yihedeng,yuyang,baharan,qgu}@cs.ucla.edu
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. Codes are available at https://github.com/uclaml/PDE.
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
Figure 1: A overview of the problem, our proposed solution, and the resultant outcomes. (A) We demonstrate the data distribution and provide an example of the statistics of Waterbirds. (B) The overall procedure of PDE. (C) we use GradCAM (Selvaraju et al., 2017) to show the attention of the model trained with PDE as compared to ERM. While ERM focuses on the background, PDE successfully trains the model to capture the birds.
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 two-layer 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 \times$  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 = \{(\mathbf{x}_i, y_i)\}_{i=1}^N$ , where  $\mathbf{x}_i \in \mathbb{R}^d$  is the input and  $y \in \{\pm 1\}$  is the output label. We train a model  $f(\mathbf{x}; \mathbf{W})$  with weight  $\mathbf{W}$  to minimize the empirical loss function:
$$
\mathcal {L} (\mathbf {W}) = \frac {1}{N} \sum_ {i = 1} ^ {N} \ell \left(y _ {i} f \left(\mathbf {x} _ {i}; \mathbf {W}\right)\right), \tag {2.1}
$$
where  $\ell$  is the logistic loss defined as  $\ell(z) = \log(1 + \exp(-z))$ . The empirical risk minimizer refers to  $\mathbf{W}^*$  that minimizes the empirical loss:  $\mathbf{W}^* \coloneqq \operatorname{argmin}_{\mathbf{W}} \mathcal{L}(\mathbf{W})$ . Typically, gradient-based optimization algorithms are employed for ERM. For example, at each iteration  $t$ , gradient descent (GD) has the following update rule:
$$
\mathbf {W} ^ {(t + 1)} = \mathbf {W} ^ {(t)} - \eta \nabla \mathcal {L} (\mathbf {W} ^ {(t)}). \tag {2.2}
$$
Here,  $\eta > 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.,  $\mathbf{x} = [\mathbf{x}_{\mathrm{core}},\mathbf{x}_{\mathrm{spu}},\mathbf{x}_{\mathrm{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  $(\mathbf{x},y,a)\in (\mathbb{R}^d)^P\times \{\pm 1\} \times \{\pm 1\}$  is generated from the distribution  $\mathcal{D}$  as follows.
- Randomly generate the true label  $y \in  \{  \pm  1\}$  .
- Generate spurious label  $a \in \{\pm y\}$ , where  $a = y$  with probability  $\alpha > 0.5$ .
- Generate  $\mathbf{x}$  as a collection of  $P$  patches:  $\mathbf{x} = (\mathbf{x}^{(1)},\mathbf{x}^{(2)},\dots ,\mathbf{x}^{(P)})\in (\mathbb{R}^d)^P$ , where
- Core feature. One and only one patch is given by  $\beta_c \cdot y \cdot \mathbf{v}_c$  with  $\| \mathbf{v}_c \|_2 = 1$ .
- Spurious feature. One and only one patch is given by  $\beta_s \cdot a \cdot \mathbf{v}_s$  with  $\| \mathbf{v}_s \|_2 = 1$  and  $\langle \mathbf{v}_c, \mathbf{v}_s \rangle = 0$ .
- Random noise. The rest of the  $P - 2$  patches are Gaussian noises  $\xi$  that are independently drawn from  $N(0, (\sigma_p^2 / d) \cdot \mathbf{I}_d)$  with  $\sigma_p$  as an absolute constant.
And  $0 < \beta_{c} \ll \beta_{s} \in \mathbb{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 = \{(\mathbf{x}_i, y_i, a_i)\}_{i=1}^N$  and let  $S$  be partitioned into large group  $S_1$  and small group  $S_2$  such that  $S_1$  contains all the training
Figure 2: Visualization of the data.
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  $\widehat{\alpha} = \frac{|S_1|}{N}$  and therefore  $1 - \widehat{\alpha} = \frac{|S_2|}{N}$ .
Visualization of our data. In Figure 2, we present the visualization in 2D space of the higher-dimensional 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:
$$
f (\mathbf {x}; \mathbf {W}) = \sum_ {j \in [ J ]} \sum_ {p = 1} ^ {P} \sigma \left(\langle \mathbf {w} _ {j}, \mathbf {x} ^ {(p)} \rangle\right), \tag {2.3}
$$
where  $\mathbf{w}_j\in \mathbb{R}^d$  is the weight vector of the  $j$ -th filter,  $J$  is the number of filters (neurons) of the network, and  $\sigma (z) = z^3$  is the activation function.  $\mathbf{W} = [\mathbf{w}_1,\dots ,\mathbf{w}_J]\in \mathbb{R}^{d\times 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 = \mathrm{polylog}(d)$ . We initialize  $\mathbf{W}^{(0)}\sim \mathcal{N}(0,\sigma_0^2)$ , where  $\sigma_0^2 = \frac{\mathrm{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 frequently considered metric is the test accuracy:  $\mathrm{Acc}(\mathbf{W}) = \mathbb{P}_{(\mathbf{x},y,a)\sim \mathcal{D}}\big[\mathrm{sgn}(f(\mathbf{x};\mathbf{W})) = y\big]$ . With spurious correlations, researchers are more interested in the worst-group accuracy:
$$
\mathrm {A c c} _ {\mathrm {w g}} (\mathbf {W}) = \min _ {y \in \{\pm 1 \}, a \in \{\pm 1 \}} \mathbb {P} _ {(\mathbf {x}, y, a) \sim \mathcal {D}} [ \mathrm {s g n} (f (\mathbf {x}; \mathbf {W})) = y ],
$$
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  $\langle \mathbf{w}_i^{(t)},\mathbf{v}_s\rangle$  and  $\langle \mathbf{w}_i^{(t)},\mathbf{v}_c\rangle$  respectively:
Theorem 2.2. Consider the training dataset  $S = \{(\mathbf{x}_i, y_i)\}_{i=1}^N$  that follows the distribution in Definition 2.1. Consider the two-layer nonlinear CNN model as in (2.3) initialized with  $\mathbf{W}^{(0)} \sim \mathcal{N}(0, \sigma_0^2)$ . After training with GD in (2.2) for  $T_0 = \widetilde{\Theta}\big(1 / (\eta \beta_s^3 \sigma_0)\big)$  iterations, for all  $j \in [J]$  and  $t \in [0, T_0)$ , we have
$$
\widetilde {\Theta} (\eta) \beta_ {s} ^ {3} (2 \widehat {\alpha} - 1) \cdot \left\langle \mathbf {w} _ {j} ^ {(t)}, \mathbf {v} _ {s} \right\rangle^ {2} \leq \left\langle \mathbf {w} _ {j} ^ {(t + 1)}, \mathbf {v} _ {s} \right\rangle - \left\langle \mathbf {w} _ {j} ^ {(t)}, \mathbf {v} _ {s} \right\rangle \leq \widetilde {\Theta} (\eta) \beta_ {s} ^ {3} \widehat {\alpha} \cdot \left\langle \mathbf {w} _ {j} ^ {(t)}, \mathbf {v} _ {s} \right\rangle^ {2}, \tag {2.4}
$$
$$
\widetilde {\Theta} (\eta) \beta_ {c} ^ {3} \widehat {\alpha} \cdot \left\langle \mathbf {w} _ {j} ^ {(t)}, \mathbf {v} _ {c} \right\rangle^ {2} \leq \left\langle \mathbf {w} _ {j} ^ {(t + 1)}, \mathbf {v} _ {c} \right\rangle - \left\langle \mathbf {w} _ {j} ^ {(t)}, \mathbf {v} _ {c} \right\rangle \leq \widetilde {\Theta} (\eta) \beta_ {c} ^ {3} \cdot \left\langle \mathbf {w} _ {j} ^ {(t)}, \mathbf {v} _ {c} \right\rangle^ {2}. \tag {2.5}
$$
After training for  $T_0$  iterations, with high probability, the learned weight has the following properties: (1) it learns the spurious feature  $\mathbf{v}_s$ :  $\max_{j\in [J]}\langle \mathbf{w}_j^{(T)},\mathbf{v}_s\rangle \geq \widetilde{\Omega} (1 / \beta_s)$ ; (2) it almost does not learn the core feature  $\mathbf{v}_c$ :  $\max_{j\in [J]}\langle \mathbf{w}_j^{(T)},\mathbf{v}_c\rangle = \widetilde{\mathcal{O}} (\sigma_0)$ .
Discussion. The detailed proof is deferred to Appendix E, and we provide intuitive explanations of the theorem as follows. A larger value of  $\langle \mathbf{w}_i^{(t)},\mathbf{v}\rangle$  for  $\mathbf{v}\in \{\mathbf{v}_s,\mathbf{v}_c\}$  implies better learning of the feature vector  $\mathbf{v}$  by neuron  $\mathbf{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  $\langle \mathbf{w}_i^{(t)},\mathbf{v}_s\rangle$  is significantly faster than that of  $\langle \mathbf{w}_i^{(t)},\mathbf{v}_c\rangle$ . 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  $(\mathbf{x},y,a)\sim \mathcal{D}$ , the neural network's prediction is given by  $\mathrm{sgn}\big(f(\mathbf{x};\mathbf{W})\big) = \mathrm{sgn}\big(\sum_{j\in [J]}\big(y\beta_c^3\langle \mathbf{w}_j,\mathbf{v}_c\rangle^3 +a\beta_s^3\langle \mathbf{w}_j,\mathbf{v}_s\rangle^3 +\langle \mathbf{w}_j,\xi \rangle^3\big)\big)$ . Since the term  $\beta_s^3\max_{j\in [J]}\langle \mathbf{w}_j,\mathbf{v}_s\rangle^3$  dominates the summation, the prediction will be  $\mathrm{sgn}(f(\mathbf{x};\mathbf{W})) = a$ . Consequently, we obtain the test accuracy as  $\mathrm{Acc}(\mathbf{W}) = \alpha$ , since  $a = y$  with probability  $\alpha$ , and the model accurately classifies the large group. However, when considering the small group and examining examples where  $y\neq a$ , the models consistently make errors, resulting in  $\mathrm{Acc}_{wg}(\mathbf{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  $\langle \mathbf{w}_i^{(t)},\mathbf{v}_s\rangle$  in (2.4) and  $\langle \mathbf{w}_i^{(t)},\mathbf{v}_c\rangle$  in (2.5) follows the formula  $x_{t + 1} = x_t + \eta Ax_t^2$ , where  $x_{t}$  represents the inner product sequence with regard to iteration  $t$  and  $A$  is the coefficient containing  $\hat{\alpha}$ ,  $\beta_{c}$  or  $\beta_{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  $\hat{\alpha}$  and feature strengths  $\beta_{c},\beta_{s}$ .
- When the model learns spurious feature  $(\beta_c^3 < \beta_s^3 (2\widehat{\alpha} - 1))$ . We examine the lower bound for the growth of  $\langle \mathbf{w}_i^{(t)}, \mathbf{v}_s \rangle$  in (2.4) and the upper bound for the growth of  $\langle \mathbf{w}_i^{(t)}, \mathbf{v}_c \rangle$  in (2.5) in Theorem 2.2. If  $\beta_c^3 < \beta_s^3 (2\widehat{\alpha} - 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:  $\widehat{\alpha} > 1/2$  (groups are imbalanced) and  $\beta_c < \beta_s$  (the spurious feature is stronger). This observation is consistent with real-world datasets, such as the Waterbirds dataset, where  $\widehat{\alpha} = 0.95$  and the background is much easier to learn than the intricate features of the birds.
- When the model learns core feature  $(\beta_{c} > \beta_{s})$ . However, if we deviate from the aforementioned conditions and consider  $\beta_{c} > \beta_{s}$ , we can examine the lower bound for the growth of  $\langle \mathbf{w}_i^{(t)}, \mathbf{v}_c \rangle$  in (2.5) and the upper bound for the growth of  $\langle \mathbf{w}_i^{(t)}, \mathbf{v}_s \rangle$  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  $\hat{\alpha} > 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  $\widehat{\alpha} = 1/2$  in a group-balanced data subset. Subsequently, we utilize the momentum to amplify the core feature's coefficient, allowing for tolerance of  $\widehat{\alpha} > 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  $\eta$ ; momentum coefficient  $\gamma$ ; initialization scale  $\sigma_0$ ; training set  $S = \{(\mathbf{x}_i, y_i, a_i)\}_{i=1}^n$ ; model  $f_{\mathbf{W}}$ .
1: Initialize  $\mathbf{W}^{(0)}$
Warm-up stage
2: Divide the  $S$  into groups by values of  $y$  and  $a$ :  $S_{y,a} = \{(\mathbf{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_{y,a}^0| = \min_{y',a'} |S_{y',a'}|$  for  $y \in \{\pm 1\}$  and  $a \in \{\pm 1\}$ .
4: for  $t = 0,1,\ldots ,{T}_{0}$  do
5: Compute loss on  $S^0$ :  $\mathcal{L}_{S^0}(\mathbf{W}^{(t)}) = \frac{1}{|S^0|}\sum_{i\in S^0}\ell (y_if(x_i;\mathbf{W}^{(t)}))$
6: Update  $\mathbf{W}^{(t + 1)}$  by (3.1) and (3.2).
7: end for
Expansion stage
8: for  $k = 1, \dots, K$  do
9: Draw  $m$  examples  $(S_{[m]})$  from  $S / S^{k - 1}$  and let  $S^k = S^{k - 1}\cup S_{[m]}$ .
10: for  $t = 1,\dots ,J$  do
11: Compute loss on  $S^k$ :  $\mathcal{L}_{S^k}(\mathbf{W}^{(T)}) = \frac{1}{|S^k|}\sum_{i\in S^k}\ell (y_if(x_i;\mathbf{W}^{(T)}))$ , where  $T = T_0 + (k - 1)*J + t$
12: Update  $\mathbf{W}^{(T + 1)}$  by (3.1) and (3.2).
13: end for
14: end for
15: return  $\mathbf{W}^{(t)} = \mathrm{argmax}_{\mathbf{W}(t')} \mathrm{Acc}_{\mathrm{wg}}^{\mathrm{val}}(\mathbf{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  $(\mathrm{GD} + \mathbf{M})$ , at each iteration  $t$  and with momentum coefficient  $\gamma > 0$ , it updates as follows
$$
\mathbf {g} ^ {(t + 1)} = \gamma \mathbf {g} ^ {(t)} + (1 - \gamma) \nabla \mathcal {L} (\mathbf {W} ^ {(t)}), \tag {3.1}
$$
$$
\mathbf {W} ^ {(t + 1)} = \mathbf {W} ^ {(t)} - \eta \cdot \mathbf {g} ^ {(t + 1)}, \tag {3.2}
$$
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  $\widehat{\alpha} = 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 = \{(\mathbf{x}_i, y_i, a_i)\}_{i=1}^{N_0}$  with  $\widehat{\alpha} = 1/2$  as in Definition 2.1 and CNN as in (2.3). The gradient on  $\mathbf{v}_s$  will be 0 from the beginning of training.