Changing the Training Data Distribution to Reduce Simplicity Bias Improves In-distribution Generalization

Summary

It is known that usually deep neural networks will learn “easy examples" that contain fast-learnable features first while learning more complex examples in a second time. The authors argue that mitigating such simplicity bias is the reason method like SAM are outperforming SGD. Based on such analysis, the authors introduce their methods coined as USEFUL that consists in two setups: 1) Identifying the examples with fast-learnable features using a clustering method based on layer output similarity 2) Upsampling by a constant factor the remaining examples with slow-learning features. By doing so, the authors can significantly increase model performances and training time on different classification tasks using different optimizers. They assess their methods across a wide range of dataset and different hyper-parameters and outperform random clustering baseline.

Strengths

This paper is well motivated and written. The method seems to be sounded and I really appreciate that the authors assess their method using different hyper-parameters such as optimizer, batch size, datasets, upsampling factor, architectures, and data augmentation. It is also great that they ran a baseline with random clustering.

Weaknesses

It is not clear when and why one should choose the last output activation vector to define the clustering instead of intermediate activation vector. It is also not clear at which epoch one should decide to do the clustering since for a dataset like CIFAR10 the optimal performances are achieved at epoch 8 while for CIFAR100 it is epoch 20. So, finding the correct hyper-parameters for the clustering might be costly and thus impact how fast convergence can really be (if we consider this needed additional ablation on clustering epoch). In addition, the authors mention that they are using an upscaling factor of 2, but I am wondering how robust this is when using long-tail distribution. For example, I am not sure that on something like ImageNet-LT or Inaturalist, we will get the best performances by using a constant factor. I would also be a bit more cautious about some of the claims made in the papers. For example, the authors claim that their method is generalizing to OOD tasks while providing experiments on only the WaterBird dataset. So, it would be better to write about promising preliminary results than claiming generalization on OOD.

Questions

1) Do you think your method will also generalize on Long-Tail dataset while keeping an upscaling factor constant? 2)) Any ideas or heuristics about how to find the optimal epoch/layer to perform the clustering without running an expensive ablation?

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors did not really discuss any limitations (outside the fact that their theoretical result does not extend to CNN) or societal impact. I think that one limitation that could have been highlighted is the smaller scale of the experiment and the focus on classification tasks. Another limitation is the lack of results on OOD or long-tail benchmarks which would seem to be well suited for this type of work.

Summary

This work aims to modify the training data distribution to improve in-distribution generalization. First, the authors theoretically analyse a 2-layer CNN and compare the feature learning dynamics (fast learnable and slow-learnable features) of Gradient Descent (GD) and Sharpness-Aware Minimization (SAM). It is then shown that SAM mitigates simplicity bias compared to GD. The authors then propose USEFUL (UpSample Early For Uniform Learning), a method that upsamples the examples in the training set that contains slow-learnable features. USEFUL first clusters the examples with similar outputs early in the training and then upsamples the slow-learnable clusters. The main idea behind USEFUL is to learn features at a uniform speed (similar to SAM) by changing the training data distribution. USEFUL can be trained with SGD, SAM and SAM + Trivial Augment. Results on CIFAR-10, CIFAR-100, STL10, TinyImageNet indicate that USEFUL is across datasets and architectures. Additonal ablation and analysis show that USEFUL learns similar properties to SAM (for e.g less sharp solutions).

Strengths

1. Originality: The question posed by the authors “Can we change the training data distribution such that the model trained on it has similar properties to SAM?” is interesting and novel. The proposed method is also well-motivated. 2. Results: The authors perform a comprehensive set of ablations and analysis on the proposed method USEFUL. Section 5.4 that shows that USEFUL’s solution has similar properties to SAM, which answers the question raised in the motivation of the paper. I also particularly like the ablations with upweighting loss and data selection method in Appendix D.6. 3. Overall, the paper is fairly well written. One minor point to address here is that the paper covers multiple concepts like SAM, simplicity bias, flat minima and uniform feature learning. It would be good to explain the relationship between these more clearly.

Weaknesses

1. The authors explicitly mention that their focus in this paper is only on “in-distribution generalization”. I am a bit confused by this given the motivation of simplicity bias and learning features uniformly. To elaborate more on this point, - Springer et al [1] also show that SAM implicitly balances the quality of diverse features (similar to the observations made in Section 3 of this paper. The experimental results in [1] is focused more on datasets with multiple predictive features like CelebA, CIFAR-MNIST. - Past work on simplicity bias and shortcut learning [2, 3, 4, 5] has focused on similar datasets like CelebA, Waterbirds, CIFAR-MNIST, Colored-MNIST to name a few. - While the authors have shown encouraging results on Waterbirds dataset in Appendix D5, it would be good to show the complete results on various groups and on other datasets as well. 2. Connection to [1]. Springer et al [1] made a very similar observation as to Section 3 in this paper. It would be great if the authors can clarify the differences with the observations in [1] and this work. Particularly, [1] also shows that SAM mitigates simplicity bias and that SAM learns higher quality representations of hard-to-learn features. The authors briefly discuss this in Related Works section but a more detailed answer would be helpful. 3. I just wanted to understand the practical usefulness of the proposed method. This method has one additional hyperparameter i.e the separating epoch. The authors have reported the best separating epoch for all the datasets which is epoch 8 for CIFAR-10 and epoch 20 for CIFAR-100 (Appendix C.2). How is this hyperparameter chosen? Is there a separating epoch number that works across various datasets? This is especially relevant given that that the average gain on most of the datasets with USEFUL is less than 1% with additional cost for training. [1] Springer, Jacob Mitchell, Vaishnavh Nagarajan, and Aditi Raghunathan. "Sharpness-Aware Minimization Enhances Feature Quality via Balanced Learning." The Twelfth International Conference on Learning Representations. [2] Shah, Harshay, et al. "The pitfalls of simplicity bias in neural networks." Advances in Neural Information Processing Systems 33 (2020): 9573-9585. [3] Geirhos, Robert, et al. "Shortcut learning in deep neural networks." Nature Machine Intelligence 2.11 (2020): 665-673. [4] Kirichenko, Polina, Pavel Izmailov, and Andrew Gordon Wilson. "Last layer re-training is sufficient for robustness to spurious correlations." arXiv preprint arXiv:2204.02937 (2022). [5] Teney, Damien, et al. "Evading the simplicity bias: Training a diverse set of models discovers solutions with superior ood generalization." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

Questions

1. I did not find the separating epoch used for most of the datasets (except CIFAR-10 and CIFAR-100). Could you please point me to that? 2. What is the takeaway from Figure 1? Please mention the observations regarding the Figure involving the fast and slow learnable features. What kind of examples are usually clustered in fast learnable cluster vs slow learnable cluster, Please refer to the Weakness section for remaining questions.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

Yes, the authors have addressed limitations.

Summary

- Proves for a 2-layer CNN with fixed second layer weighsts trained on a toy dataset, SAM learns slow-learnable and fast-learnable features more uniformly in the early epochs compared to SGD - Based on this analysis, proposes a simple clustering-based upsampling strategy for reducing simplicity bias / excessive reliance on fast-learnable features. The results show that this improves in-distribution generalization of standard small-scale image classification tasks.

Strengths

- Simple easy-to-implement method that uses SAM and upsampling to improve in-distribution generalization - The method is well justified with theoretical analysis comparing SAM and SGD on a toy data distribution. This analysis indicates that SAM is less sensitive to simplicity bias.

Weaknesses

- No baselines. There are several papers now that try to reduce simplicity bias in order to improve performance: - https://arxiv.org/abs/2105.05612 - https://arxiv.org/abs/2301.13293 - https://arxiv.org/abs/2107.09044 (does not focus on simplicity bias explicitly, but similar to method proposed in paper) - simpler baselines: there are several papers that propose “example difficulty” metrics (https://arxiv.org/abs/2106.09647). How well do this correlate with the clusters found in your method? If you just train on the k examples with the highest difficulty scores (per class), does this fare worse than the proposed method? - Limited novelty due to findings in [64] (Sharpness-aware minimization enhances feature quality via balanced learning). This paper also shows that SAM improves feature diversity (on real datasets + backed up with analysis on a toy dataset) and improves performance on transfer-learning tasks. - Lacking discussion about when this method would fail. I can imagine two scenarios where the method would not work: 1. Most training examples have one or more slow-learnable features. In this case, the clustering approach would “remove” most of the points in the dataset, and train on very points for multiple epochs. This could result in overfitting and performance that is worse than training. There’s an implicit assumption that there is some sort of one-to-one relation between examples and features. In the case where all examples contain an “easy” (e.g. patch) and a “hard” feature (e.g. CIFAR), would this method improve performance over SGD? 2. In noisy datasets, low-quality examples or mislabeled examples would require more time to learn, and this method would cluster them and train on them for longer. That is, it would group examples that are “high-quality” and hard-to-learn with “low-quality” points. In this case, would the proposed method improve performance over SGD? - “SB of SGD has been long conjectured to be the reason for the superior generalization performance of overparameterized models, by providing capacity control or implicit regularization” This incorrectly cites https://arxiv.org/abs/2006.07710v2, which shows that too much simplicity bias can lead to robustness and in-distribution generalization issues. - Unfair evaluation. The experiments compare SAM+TA augmentation and SAM+USEFUL+TA to SGD (no TA). I think there should be two plots, complaring {SGD, SAM, SAM+USEFUL} w/ and w/o TA. - Experiments on larger datasets. The image classification used here are fairly small-scale. I would like to how well this method scales to ImageNet-scale datasets (TinyImageNet is not a good proxy..) - Writing is repetitive at times, especially the theory section (3.3)

Questions

- How specific is the analysis to the toy distribution setting in which there is just 1 slow-learnable and 1 fast-learnable feature? How do things change if the number of slow-learnable features >> number of fast-learnable features? - Why is max alignment between weight vector and ground-truth feature v_e the right way to evaluate the feature’s contribution to the model? Isn’t it hypothetically possible that SAM solutions rely more on the simpler feature if more weight vectors rely on v_e instead of v_d, even when max-alignment with slow feature is higher for SAM? Some discussion connecting this metric to “feature reliance” vis-a-vis model outputs would be great.

Rating

4

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

Please see strengths and weaknesses.

Summary

This paper proposes an algorithm for changing the distribution of training data to improve the generalization of the model on origin data distribution. The paper is inspired by Sharpness Aware Minimization, which aims at finding a flat minimum meaning that it has a good generalization capability. This paper divides features into two categories: fast-learnable features and slow-learnable features and derives some observations like "SGD and SAM only learn fast-learnable or easy features early in training" and "SAM learns slow-learnable and fast-learnable features at a more uniform speed". The authors propose the method dubbed as USEFUL to train the model on some slow-learnable features repeatedly. The experiments show the effectiveness of USEFUL on CIFAR10 and CIFAR100 datasets.

Strengths

- The paper is well-written and easy to follow. - The paper has a theoretical analysis to analyze the learning progress and derive the proposed method. - The experiments are abundant and comprehensive.

Weaknesses

There are some questions based on the presentation of this paper, I will not hesitate to improve my score if the following question are solved. - Difference between this paper and methods for long-tailed data distribution or measuring the difficulty of learning examples. Algorithms for long-tailed data distribution are usually based on resampling training data or reweighing loss value. The proposed USEFUL is similar to the resampling methods except that USEFUL focuses on the features that are hard/slow to learn. Some references for understanding: [Shi, Jiang-Xin, et al. "How re-sampling helps for long-tail learning?." Advances in Neural Information Processing Systems 36 (2023).](https://arxiv.org/pdf/2310.18236), [Shrivastava, Abhinav, Abhinav Gupta, and Ross Girshick. "Training region-based object detectors with online hard example mining." Proceedings of the IEEE conference on computer vision and pattern recognition. 2016.](https://arxiv.org/pdf/1604.03540v1) and some references based on it, [A Re-Balancing Strategy for Class-Imbalanced Classification Based on Instance Difficulty](https://openaccess.thecvf.com/content/CVPR2022/papers/Yu_A_Re-Balancing_Strategy_for_Class-Imbalanced_Classification_Based_on_Instance_Difficulty_CVPR_2022_paper.pdf), [Active Teacher for Semi-Supervised Object Detection](https://openaccess.thecvf.com/content/CVPR2022/papers/Mi_Active_Teacher_for_Semi-Supervised_Object_Detection_CVPR_2022_paper.pdf), I believe a discussion of these references in paper should be helpful. - The relation between the proposed USEFUL and SAM? It seems like the motivation of USEFUL is changing the data distribution to get a flat minimum like SAM. But the results in Appendix D.2, *i.e.*, 53.8 for SGD 41.8 for SGD+USEFUL 12.4 for SAM in Table 1($\lambda_{max})$, do not show effectiveness compared with SAM. It could show the effectiveness on SGD but it's far from being comparable to SAM. Some small questions: - What's the exact formulation of the Data distribution? - What's the "patch" meaning in Definition 3.1? Is that the same as the patch in ViT or the channel of the image? It's a little confusing. - The experiments mainly focus on traditional architecture, e.g., n-Layer CNN, ResNet. More experiments on popular models and big datasets, e.g., Transformer ImageNet-1k, would be better.

Questions

See Weakness.

Rating

4

Confidence

3

Soundness

4

Presentation

3

Contribution

3

Limitations

N.A.

Response to Rebuttal

I thank the authors for the detailed rebuttal. The authors have clearly answered most of the questions. > reducing the simplicity bias benefits in-distribution (ID). This is an interesting result and a bold claim. Does this mean that past methods that focus on mitigating simplicity bias (and show results in the OOD scenarios) also improve the performance in IID setup by similar margins? The additional simplicity bias (JTT and EIIL) baselines discussed in the rebuttal pdf seem to indicate so. I would love to hear the authors' comments on this. In the future, the authors can consider running experiments with a couple of other methods that mitigate simplicity to verify the generality of the claim. Most of my concerns have been resolved and thus, I have increased my score to 6.