SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In particular, with  $\widehat{\alpha} = 1/2$  we have  $|S_1^0| = |S_2^0|$ : 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  $\mathbf{v}$ , it retains a substantial historical gradient in the momentum containing  $\mathbf{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  $\mathbf{v}_s$  and momentum does not help, as we will show in synthetic experiments, PDE avoids learning  $\mathbf{v}_s$  and learns  $\mathbf{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  $\mathbf{v}_c$  in the expansion stage. For a specific illustration, the learning of the core feature by GD+M will be
$$
\langle \mathbf {w} _ {j} ^ {(t + 1)}, \mathbf {v} _ {c} \rangle = \langle \mathbf {w} _ {j} ^ {(t)} - \eta \big (\gamma g ^ {(t)} + (1 - \gamma) \nabla_ {\mathbf {w} _ {j}} \mathcal {L} (\mathbf {W} ^ {(t)}) \big), \mathbf {v} _ {c} \rangle ,
$$
where  $g^{(t)}$  is the additional momentum as compared to GD with  $\gamma = 0$ . While the current gradient along  $\mathbf{v}_c$  might be small (i.e.,  $\beta_c$ ), we can benefit from the historical gradient in  $g^{(t)}$  to amplify the growth of  $\langle \mathbf{w}_j^{(t + 1)},\mathbf{v}_c\rangle$  and make it larger than that of the spurious feature (i.e.,  $\beta_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.
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.
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 section 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  $\mathrm{GD + M}$  is unable to accurately predict the small group in our specified data distribution where  $\hat{\alpha} = 0.98$  and  $\beta_{c} < \beta_{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  $\beta_{c} > \beta_{s}$  instead. Even with the same  $\hat{\alpha}$  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.
Table 1: Synthetic data experiments. We report the worst-group accuracy and the gap (i.e., overall - worst). We further consider several variations of PDE to demonstrate the importance of each component of our method. Reset: we reset the momentum to zero after the warm-up stage. Warmup+All: we let PDE to incorporate all of the new training data at once after the warm-up stage.
Worst-group (%)	Gap (%)
ERM (GD)	0.00	97.71
ERM (GD+M)	0.00	97.71
Warmup+All (Reset)	67.69	31.18
Warmup+All	74.24	24.76
PDE (Reset)	92.51	2.29
PDE	93.01	0.03
(a) ERM (case 1)
(b) ERM (case 2)
(c) PDE
Figure 3: Training process of ERM vs. PDE. We consider the same dataset generated from the distribution as in Definition 2.1 for ERM (case 1) and PDE. On the same training data, ERM learns the spurious feature while PDE successfully learns the core feature. We further consider ERM (case 2) when training on the data distribution where  $\beta_{c} > \beta_{s}$  and  $\widehat{\alpha} = 0.98$ . We show the growth of the max inner product between the model's neuron and core/spurious signal vector and the decrease of training loss with regard to the number of iterations  $t$ .
We conduct experiments on real benchmark datasets to: (1) compare our approach with state-of-the-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) CivilComments-WILDS (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  $\mathrm{DFR}^{\mathrm{Tr}}$  for fair comparison as all methods only use the training dataset to train and finetune the model, while  $\mathrm{DFR}^{\mathrm{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
Table 2: The worst-group and average accuracy (%) of PDE compared with state-of-the-art methods. The bold numbers indicate the best results among the methods that require group information, while the underscored numbers represent methods that only train once. All methods use validation data for early stopping and model selection, while  $\sqrt{\sqrt{}}$  indicates that the method also re-trains the last layer using the validation data
Method	Group info	Train once	Val info	Waterbirds	CelebA	CivilComments
Worst	Average	Worst	Average	Worst	Average
ERM	×	√	√	70.0±2.3	97.1±0.1	45.0±1.5	94.8±0.2	58.2±2.8	92.2±0.1
LfF	×	×	√	78.0N/A	91.2N/A	77.2N/A	85.1N/A	58.8N/A	92.5N/A
EIIL	×	×	√	77.2±1.0	96.5±0.2	81.7±0.8	85.7±0.1	67.0±2.4	90.5±0.2
JTT	×	×	√	86.7N/A	93.3N/A	81.1N/A	88.0N/A	69.3N/A	91.1N/A
Subsample	√	√	√	86.9±2.3	89.2±1.2	86.1±1.9	91.3±0.2	64.7±7.8	83.7±3.4
DFRTr	√	×	√	90.2±0.8	97.0±0.3	80.7±2.4	90.6±0.7	58.0±1.3	92.0±0.1
DFRVal	√	×	√√	92.9±0.2	94.2±0.4	88.3±1.1	91.3±0.3	70.1±0.8	87.2±0.3
GroupDRO	√	√	√	86.7±0.6	93.2±0.5	86.3±1.1	92.9±0.3	69.4±0.9	89.6±0.5
PDE	√	√	√	90.3±0.3	92.4±0.8	91.0±0.4	92.0±0.6	71.5±0.5	86.3±1.7
Table 3: Training efficiency of PDE and GroupDRO on Waterbirds. We compare with GroupDRO at their learning rate and weight decay, as well as at ours. We report the worst-group accuracy, average accuracy and the number of epochs till early stopping as the model reached the best performance on validation data. Note: for a fair comparison, we consider one training epoch as training over the  $N$  data as the size of the training dataset.
Method	Learning rate	Weight decay	Worst	Average	Early-stopping epoch*
GroupDRO	1e-5	1e-0	86.7±0.6	93.2±0.5	92±4
GroupDRO	1e-2	1e-2	77.3±2.0	97.1±0.5	15±15
PDE	1e-2	1e-2	90.3±0.3	92.4±0.8	8.9±1.8
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 \times$  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 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.
Table 4: Performance of PDE after each stage. We report the worst-group and average accuracy.
Dataset	Warm-up	Addition
Worst	Avg	Worst	Avg
Waterbirds	86.0	91.9	90.3	92.4
CelebA	87.8	92.1	91.0	92.0
CivilComm	67.7	78.8	71.5	86.3
(a) Worst-group accuracy.
(b) Average accuracy.
Figure 4: The effect of reset the momentum to zero after the warm-up stage for PDE on Waterbirds.
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.
(a) Worst-group accuracy.
Figure 5: The variations in both worst-group and average accuracy on the test set of Waterbirds during expansion stage under different expansion learning rates. Each vertical dashed line denotes an expansion and the arrow denotes the early stopping.
(b) Average accuracy.
Table 5: Ablation study on Waterbirds. Exp. size: number of data points added in each expansion.
Exp. size	Exp. lr	Worst	Average
5	1e-4	89.9±0.5	92.1±0.3
10	1e-4	90.3±0.3	92.4±0.8
50	1e-4	88.1±0.8	93.4±0.4
Table 6: Ablation study on Waterbirds. Exp. Ir: the learning rate in expansion stage.
Exp. size	Exp. lr	Worst	Average
10	1e-2	85.4±3.1	92.1±2.0
10	1e-3	89.4±0.7	92.6±0.3
10	1e-5	89.5±0.2	92.1±0.1
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 self-supervised 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.
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.
We sincerely thank Dongruo Zhou for the constructive suggestions on the structure and writings of the paper. We also thank the anonymous reviewers for their helpful comments. YD and QG are supported in part by the National Science Foundation CAREER Award 1906169 and IIS-2008981, and the Sloan Research Fellowship. BM is supported by the National Science Foundation CAREER Award 2146492. The views and conclusions contained in this paper are those of the authors and should not be interpreted as representing any funding agencies.
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., Chaumont, 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-the-art 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.