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

We thank the reviewer for the valuable feedback, and acknowledging our well-motivated work and our comprehensive experiments and ablations.

1. When and why one should choose the last output activation vector to define the clustering instead of intermediate activation vector?

• Our theorems 3.2 and 3.3 show that examples with at least one fast-learnable feature are learned in early training iterations, and this is well reflected in the model's final output (normalized logit vector). Intuitively, regardless of the layer in which different fast-learnable features are learned, the model’s output for examples containing any fast-learnable feature become similar to its one-hot encoded label. Hence, examples with at least one fast-learnable feature can be well separated based on the model’s final output. While one may be able to separate examples with the same fast-learnable feature by clustering model’s activation at different layers, this is not necessary for our purpose, which is reducing the simplicity bias of early training. In our experiments, we used the model’s final output (output of the softmax), without further tuning the layer number.

2. How to find the separating epoch?

• In Fig 14a, we showed that the epoch that can best separate the examples is when the reduction in training error starts to diminish. Intuitively, the first part of the plot with a high negative slope is the time that fast-learnable features are learned due to the simplicity bias. When the training loss curve diminishes the examples with fast learnable features can be best separated from the remaining examples. Figure 14b shows our ablation and confirms that separating examples at multiple epochs around this time (when the slope of the training loss curve diminishes) and upsampling the remaining examples outperform SGD. This observation can help find the separating epoch relatively quickly.

3. How robust is the scaling factor of 2 on long-tail distribution?

• The reason for the scaling factor of 2 is to reduce the simplicity bias of early training without further modifying the training data distribution. We conducted new experiments on long tail CIFAR10 with an imbalanced ratio of 10. Our results show that our method with a scaling factor of 2 can indeed improve the generalization performance. Notably, USEFUL outperforms class balancing to address long tail distribution. As can be seen in Fig 1 of the PDF, interestingly our method may upsample more examples from some of the larger classes (not smaller classes). This confirms that, instead of balancing the training data distribution, it benefits the performance by reducing the simplicity bias and learning features of a class at a more uniform speed. Empirically, as we confirmed in Fig 13 and our new experiments on long-tail data, a scaling factor of 2 works best to reduce the simplicity bias. We note that our method can be stacked with existing methods for long-tail data to balance the training data distribution and further improve the performance, as we confirmed in our new experiments.

  Table 3: Test error (avg of 3 runs) on long-tailed CIFAR10.

  Ratio	SGD	SGD + USEFUL	SAM	SAM + USEFUL
  1:10	10.01 ± 0.21	9.53 ± 0.13	8.85 ± 0.08	8.22 ± 0.04
  Balancing	9.77 ± 0.17	9.25 ± 0.11	8.31 ± 0.11	7.93 ± 0.02

4. The authors claim that their method is generalizing to OOD tasks while providing experiments on only the WaterBird dataset.

• Prior works have shown the benefits of reduced simplicity bias to improving the OOD performance [Evading the Simplicity Bias, CVPR’22, Overcoming Simplicity Bias, ICML’23, Identifying Spurious Biases Early, AISTATS’24]. Our method reduces the simplicity bias and hence is expected to benefit the OOD performance too. We confirmed this on one of the standard OOD benchmark datasets to not distract the reader from the main contribution of the paper, i.e. improved ID performance. We thank the reviewer and will revise our language as suggested.

We hope our rebuttal has effectively addressed your concerns. As the discussion phase is nearing its conclusion, we are wondering if there's anything further we can clarify or elaborate on. Your feedback is greatly appreciated.

I would like to thank the authors for addressing my concerns. I really appreciate the new long tail experiment and the fact that the authors will revise some of the paper's language. I am maintaining my score.

We thank the reviewer for their valuable feedback and recognizing the originality and novelty of our work, and our comprehensive experiments.

1. Our work (ID) vs [1-5] (OOD).

• As the reviewers correctly mentioned and we discussed in our general comment, our work studies ID, while prior work including [1-5] studied OOD with spurious features (CelebA, Waterbirds, CMNIST and CIFAR-MNIST are all benchmark datasets with spurious features in training data but not on the test data. Thus spurious features are not predictive on test). In the OOD setting, simple spurious (non-predictive) features are learned from the training data instead of the predictive features, due to simplicity bias. Training with SAM [1] or methods for alleviating the simplicity bias [2-5] suppresses learning the spurious features and allows learning more/other predictive features. In the ID setting (CIFAR10-100, TinyImageNet), there is no spurious feature and all features are predictive on the test set. Hence, it is not clear why reducing the simplicity bias should help. We proved that reducing the simplicity bias allows learning the same set of predictive features (not new features) at a more uniform speed. This benefits the ID performance. To our knowledge, our study is the first to show the benefits of reduced simplicity bias to ID, hence is a major contribution. Our results on Waterbirds shows that our method also benefits OOD, but is not the main contribution of our work (as we mentioned in lines 338-339).

2. How is the separating epoch chosen?

• In Fig 14a, we showed that the epoch that can best separate the examples is when the reduction in training error starts to diminish. Intuitively, the first part of the plot with a high negative slope is the time that fast-learnable features are learned due to the simplicity bias. When the training loss curve diminishes the examples with fast learnable features can be best separated from the remaining examples. Figure 14b shows that separating examples at multiple epochs around this time (when the slope of the training loss curve diminishes) and upsampling the remaining examples outperform SGD. We report the best-performing separating epoch (red points in the figures) in our experiments.
• Is there a separating epoch number that works across various datasets?
  • There is no universal separating epoch. This is reasonable because each data has a different data distribution, i.e. slow- and fast-learnable features, thus, a different theoretical time T in Theorems 3.2 and 3.3.
• The average gain on most of the datasets with USEFUL is less than 1% with additional cost for training.
  • Fig 4, 5 show that SAM’s improvement over SGD is also around 1%. While SAM increases the training time by 2x, and requires tuning the inner step size $\rho$ for different model architectures and datasets, it attracted a lot of attention and is considered a very important contribution as it improves the SOTA performance. Our method is much cheaper (~1.3x cost) than SAM (2x cost) and can be easily stacked with SAM and TA to further improve the SOTA performance. Thus, we believe our contribution is novel and important and can lead to a line of follow up work.

Questions:

1. Separating epoch used for all datasets

• In Figure 4, the best separating epochs for STL10, CINIC10, and Tiny-ImageNet are 11, 4, and 10, respectively. The separating epoch for Waterbirds is 5. We will add these details into our revision.

2. What is the takeaway from Figure 1? What kind of examples are usually clustered in fast learnable cluster vs slow learnable cluster?

• The bottom row in Fig 1 shows examples with fast learnable features that are learned in the first few epochs and are found by our method. We see that these examples are not ambiguous and are clear representatives of their corresponding class, hence are very easy to visually classify (the entire object is in the image and the area associated with the background is small). Such examples are learned early in training due to the simplicity bias. In contrast, the top row shows that examples without any fast learnable feature are harder, to identify visually and look more ambiguous (part of the object is in the image or the object is smaller and the area associated with the background is larger). These examples still contain features that are predictive on the test set, but are learned later during the training (require more iterations to be learned). Our method enables learning both types of examples containing fast-learnable and easy-learnable features at a more uniform speed, which we showed benefits the ID performance.

Thank you for reading our rebuttal, and we’re glad that you found our findings interesting and striking. We also believe that this is an important result and we are quite excited about it.

As shown by our theory, we expect prior methods for reducing the simplicity bias to also benefit ID, and our preliminary results with JTT and EIIL showed their promise. Such methods often have several hyperparameters that need to be carefully tuned via a grid search, as we listed for JTT and EIIL in our rebuttal. Our method is directly motivated by our theory and requires tuning only one hyperparameter within a small range (guided by the training loss curve). Hence, we expect it to be more effective in the ID setting. Nevertheless, we agree that adding prior simplicity bias methods would be a nice addition to our work and further support the finding and generality of our result. We thank the reviewer for their valuable suggestion, and we will add experiments with more simplicity bias methods to our revised version.

Thanks for your feedback!

1. Comparison with simplicity bias baselines

• Prior work, including papers referred to by the reviewer, showed the benefits of reducing the simplicity bias to out-of-distribution (OOD), where there is a shift between training and test distribution (c.f. general comment for detailed discussion).
• Our main contribution is to show that reducing the simplicity bias benefits ID. USEFUL is directly motivated by our Theorem 3, and we don’t see alternative methods for reducing the simplicity bias as baselines for our contribution. We conducted new experiments on CIFAR10 with EIIL & JTT. Our theoretically-motivated method outperforms EIIL and JTT in the ID setting (Table 5 in the PDF).

2. Example difficulty metrics vs USEFUL clusters & training on the k examples with highest difficulty scores per class

• Methods to calculate example difficulty require either training the model to convergence (forgetting score, Toneva et al, NeurIPS’18) or training several models partially (El2N, Paul et al, NeurIPS’22). USEFUL only requires training one model for as few as 4 epochs. The fast-learnable cluster that we find is correlated with the easy examples, as we confirm in our PDF. But, USEFUL does not need a fine grained calculation of difficulty for all examples. In the method suggested by the reviewer, the optimal choice of K may be different in each class (since the learning difficulty of classes are often different), and finding the best K per class requires extensive hyperparameter tuning. USEFUL automatically finds the optimal choice for K in each class, and is cheap to apply.

3. Novelty w.r.t [64]

• As discussed in our general comment, prior work including [64] studied OOD. [64] showed that SAM suppresses the spurious/redundant (non-predictive) features, and enables learning other/more diverse features which can benefit OOD. They studied toy/real dataset with a known spurious feature and group labels.
• In contrast, we consider the in-distribution (ID) setting and prove that SAM learns the same set of predictive features at a more uniform speed, which benefits ID. Unlike [64], we do not require any group labels. To our knowledge, our study is the first to show the benefits of reduced simplicity bias to ID, hence is pretty novel.

4. Would USEFUL fail if:

• Most training examples have one or more slow-learnable features
  • USEFUL finds examples with at least one fast-learnable feature (and arbitrary number of slow-learnable features), and upsample the remaining examples. If most examples have at least one slow-learnable feature, nothing can be concluded. If there is no example with fast-learnable features, there is no simplicity bias to alleviate and there won't be any cluster with low training loss. But, this is unlikely for real-world datasets.
  • Implicit assumption of one-to-one relation between examples and features. What if all examples contain an easy and a hard feature (e.g. CIFAR)? For simplicity of theoretical analysis, we assumed that every example contains one slow-learnable and at most one fast-learnable. But, such assumption is not required in practice, as we confirmed empirically on several benchmark datasets containing multiple slow and fast learnable features, including CIFAR10/100, TinyImageNet and CINIC10 where there is no one-to-one relation between examples and features.
• Noisy labeled data. Our method and analysis consider a clean dataset. But, as we confirmed in our new experiments in Table 4 in the PDF, USEFUL can easily stack with robust methods for learning against noisy labels to reduce the simplicity bias and improve their performance.

5. Comparing {SGD, SAM, SAM+USEFUL} w/ and w/o TA

• We compared (i) SGD + USEFUL with SAM, and (ii) SGD, SAM, SAM + TA when they stack with USEFUL to show that USEFUL can stack with different gradient-based optimizers and data augmentations. Our goal is not to compare methods w/ and w/o TA. Our evaluations are fair and confirm the effectiveness of USEFUL.

6. Larger datasets

• Please see Q5 of Reviewer Y69v. We do not believe there is anything specific to ImageNet that does not work with our method.

Questions:

1. If #slow-learnable features >> #number of fast-learnable features

• In this case, the dataset is very difficult to learn and learning is not affected much by the simplicity bias. In the extreme case where there is no fast learnable feature in the data, there won't be a low-loss cluster early in training and our method is not applicable (for the right reason). In practice, however, most datasets have several fast-learnable features and our method effectively improves the performance as we confirmed on several benchmark datasets.

2. Why is max alignment between weight vector and ground-truth feature v_e the right way to evaluate the feature’s contribution?

• Since the model output of example $x_i$ at iteration t is given by $f(x_i; W) = \sum_{j \in [J]} ( y \beta_d^3 \langle w_j^{(t)}, v_d \rangle^3 + y \beta_e^3 \langle w_j^{(t)}, v_e \rangle^3 + \langle w_j^{(t)}, \xi_i \rangle^3 )$, the term $\beta_e^3 \max_{j \in [J]} \langle w_j^{(t)}, v_e \rangle^3$, i.e., max alignment between weight vector and the ground-truth feature $v_e$, greatly affects the prediction when the fast-learnable feature is present.

3. Isn’t it possible that SAM solutions rely more on the simpler feature if more weight vectors rely on v_e instead of v_d?

• Indeed, SAM relies more on $v_e$ than $v_d$ in early training as proved in our Theorem 3.3. We showed that (1) SAM learns features at a more uniform speed as indicated by a smaller gap between features’ contribution in Theorem 3.4; and (2) SAM relies on $v_d$ relatively (w.r.t $v_e$) more than SGD in Theorem 3.5.

We hope our rebuttal has effectively addressed your concerns. As the discussion phase is nearing its conclusion, we are wondering if there's anything further we can clarify or elaborate on. Your feedback is greatly appreciated.

We thank the reviewer for the valuable feedback and acknowledging our theoretical results and comprehensive experiments. We discuss the questions below.

1. USEFUL vs resampling methods for long-tail data & example difficulty.

• As discussed in the general comment, long-tail data is an instance of OOD. Long-tail methods resample the data at the class or subclass level to match the training and test distribution. For example, [A Re-Balancing Strategy, CVPR’22] identifies instances with slow learning speed as more difficult instances and dynamically increases their weights during training to effectively change the data distribution to match the test distribution [page 2]. In contrast, we showed in the ID setting that the simplicity bias of gradient-based methods makes them find solutions with suboptimal performance. USEFUL reduces the simplicity bias by finding examples that are not learned during the first few training epochs, upsamples the rest of examples once, and restart training. Hence, USEFUL does not require calculating the difficulty of examples or do dynamic instance-wise reweighting or resampling during the training. Hence, the effectiveness of our method is attributed to learning features at a more uniform speed, and not matching the training and test data distributions. To our knowledge, the benefit of reduced simplicity bias to ID is studied for the first time by our work.
• We conducted new experiments on long-tail CIFAR10 with an imbalance ratio of 10. Table 3 in the PDF shows that our method can also improve the performance of SGD and SAM, by reducing the simplicity bias on the long-tail data. Fig 1 in the PDF shows the distribution of classes before and after upsampling by USEFUL. Interestingly, we see that USEFUL upsamples more examples from some of the larger classes and still improves the accuracy on the balanced test set. This improvement is attributed to the more uniform speed of feature learning, and not balancing the training data distribution. Notably, USEFUL outperforms class balancing to address long tail distribution. Besides, USEFUL can be stacked with methods to address long tail data to further improve the performance, as we confirmed in our new experiments (please refer to Q3 of Reviewer F4yx).

2. USEFUL’s solution is not as flat as SAM.

• The motivation of USEFUL is not to get a flatter minima like SAM. In fact, a flatter solution has been only conjectured (not proven) to have a better performance, and recent studies have shown conflicting evidence on the relationship between flatness and generalization, suggesting that flatness does not fully explain SAM’s success [16 in paper, When do flat minima optimizers work?, NeurIPS’22, A modern look at the relationship between sharpness and generalization, ICML’23]. A key contribution of our work is identifying an orthogonal effect of SAM that is beneficial in-distribution (ID): we proved that SAM has less simplicity bias early in training and thus learns features at a more uniform speed.
• Our method, USEFUL, is inspired by SAM in reducing the simplicity bias to improve the ID generalization performance. Due to the same reason, USEFUL can effectively improve the SAM’s performance itself, as we confirmed empirically. While our goal is not to directly find a flatter minima or the same solution as SAM, we showed that USEFUL can find flatter minima. Note that to capture the sharpness/flatness, multiple different criteria have been proposed (largest Hessian eigenvalue and bulk of Hessian), and one criteria is not enough to accurately capture the sharpness. Table 1 shows that while the solution of SGD + USEFUL has a higher largest Hessian eigenvalue, it achieves the smallest bulk. Notably, the solution found by both SGD + USEFUL and SAM have lower sharpness in both indicators than SGD, proving that USEFUL successfully reduces the sharpness of the solution.

3. Formulation of the Data distribution.

• The data distribution is defined in Definition 3.1 in Section 3.1. Each data point consists of 3 different patches, a slow-learnable, a fast-learnable, and random noise. Our data model is consistent with many prior theoretical works [2, 7, 10, 11, 15, 29, 37 in paper].

4. "patch" meaning in Definition 3.1.

• It is similar to the patch in ViT. Each patch is a small region in the image and each image consists of P patches.

5. More experiments on popular models and big datasets, e.g., Transformer ImageNet-1k

• Both CINIC10 and Tiny ImageNet contain images from ImageNet datasets and have a relatively large number of examples. CINIC10 dataset consists of 270,000 images, which is 4.5 times larger than CIFAR10. Tiny-ImageNet comprises 100,000 images distributed across 200 classes of ImageNet. Due to the computational constraints, we cannot conduct experiments on ImageNet (on our NVIDIA RTX A5000, a single training of ResNet18 with SGD on ImageNet takes around 1068 GPU hours). In Fig 5, we confirmed the effectiveness of USEFUL on other architectures such as ViT-small, which is a popular architecture for computer vision tasks [69 in paper, When vision transformers outperform resnets without pre-training or strong data augmentations, ICLR’22].

We hope our rebuttal has effectively addressed your concerns. As the discussion phase is nearing its conclusion, we are wondering if there's anything further we can clarify or elaborate on. Your feedback is greatly appreciated.