Changing the Training Data Distribution to Reduce Simplicity Bias Improves In-distribution Generalization
Authors: Dang Nguyen
Paymon Haddad
Eric Gan
Baharan Mirzasoleiman
Venue: NeurIPS 2024 poster (2024)
Keywords: In-distribution generalization
Simplicity bias
Data modification
Sharpness-aware minimization
Primary area: optimization_for_deep_networks
TL;DR: We propose a sharpness-aware motivated data modification to improve in-distribution generalization performance.

Abstract: Can we modify the training data distribution to encourage the underlying optimization method toward finding solutions with superior generalization performance on in-distribution data? In this work, we approach this question for the first time by comparing the inductive bias of gradient descent (GD) with that of sharpness-aware minimization (SAM). By studying a two-layer CNN, we rigorously prove that SAM learns different features more uniformly, particularly in early epochs. That is, SAM is less susceptible to simplicity bias compared to GD. We also show that examples constraining features that are learned early are separable from the rest based on the model’s output. Based on this observation, we propose a method that (i) clusters examples based on the network output early in training, (ii) identifies a cluster of examples with similar network output, and (iii) upsamples the rest of examples only once to alleviate the simplicity bias. We show empirically that USEFUL effectively improves the generalization performance on the original data distribution when training with various gradient methods, including (S)GD and SAM. Notably, we demonstrate that our method can be combined with SAM variants and existing data augmentation strategies to achieve, to the best of our knowledge, state-of-the-art performance for training ResNet18 on CIFAR10, STL10, CINIC10, Tiny-ImageNet; ResNet34 on CIFAR100; and VGG19 and DenseNet121 on CIFAR10.

Changing the Training Data Distribution to Reduce Simplicity Bias Improves In-distribution Generalization
Tuan Hai Dang Nguyen
Paymon Haddad
Eric Gan Baharan Mirzasoleiman
Department of Computer Science, UCLA
Abstract
Can we modify the training data distribution to encourage the underlying optimiza- tion method toward finding solutions with superior generalization performance on in-distribution data? In this work, we approach this question for the first time by comparing the inductive bias of gradient descent (GD) with that of sharpness-aware minimization (SAM). By studying two-layer CNN, we rigorously prove that SAM learns different features more uniformly, particularly in early epochs. That is, SAM is less susceptible to simplicity bias compared to GD. We also show that examples containing features that are learned early are separable from the rest based on the model's s output. Based on this observation, we propose a method that (i) clusters examples based on the network output early in training, (ii) identifies a cluster of examples with similar network output, and (iii) upsamples the rest of examples only once to alleviate the simplicity bias. We show empirically that USEFUL effectively improves the generalization performance training with various gradient methods, including (S)GI and SAM. Notably, we demonstrate that our method can be combined with SAM variants and existing data augmentation strategies to achieve, to the best of our knowledge, state-of- of-the-art per- formance for training R on CIFARIO 10, STLI0, CINICIO, Tiny-ImageNel ResNet34 on CIFAR100; and VGG19 and DenseNet121 on CIFAR10.
the original data distribution when
Introduction
Training data is a key component of machine learning pipelines and directly impacts its performance. Over the last decade, there has been a large body of efforts concerned with improving learning from a given training dataset by designing more effective optimization methods
[22
38 76 or
neural networks with improved structures [7 56 82]
higher-capacity 49 51). More recently,
improving the quality of the training data has emerged as a popular avenue to improve generalization performance. Interestingly, higher-quality data further improve the performance when larger models and better optimization methods are unable to do so (23 27). Recent efforts to improve the data quality have mainly focused on filtering irrelevant, noisy, harmful examples 23 45 66).
Nevertheless, remains an open question if one can change the distribution of a clean training data further improve the in in-distribution generalization performance of models trained on
At first glance, the above question may seem unnatural, as it disputes a fundamental assumption that training and test data should come from the same distribution 22) Under this assumption, minimizing the training loss generalizes well on the test data 77. Nevertheless, for overparameterized neural networks with more parameters than training data, there many training error solutions, all global minima of the training objective, with different generalization performance [25). Thus, one may still hope to carefully change the data distribution drive the optimization algorithms towards finding more generalizable solutions on the original data distribution.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

In this work, we take the first steps towards addressing the above problem. To do so, we rely on recent results in non-convex optimization, showing the superior generalization performance of sharpness- aware-minimization (SAM) [22] over (stochastic) gradient descent (GD). SAM finds flatter local minima by simultaneously minimizing the loss value and loss sharpness. doing so, it outperforms (S)GD and obtains state-of-the the- performance, at the expense of doubling the training time [20] 81]. Our key idea is that if one can change the training data distribution such that learning shares similar properties to that of training with SAM, then the new distribution can drive (S)GD and even SAM toward finding more generalizable solutions.
To address the above question, we first theoretically analyze the dynamics of training a two-layer convolutional neural network (CNN) with SAM and compare with that of GD. We rigorously prove that SAM learns different features in a more uniform speed compared to GD, particularly early in training. other words, we show that SAM is less susceptible to simplicity bias than GD. Simplicity bias of SGD makes the model learn simple solutions with minimum norm (25 and has long been conjectured to be the reason for the superior generalization performance of overparameterized models
by providing implicit regularization 7 25 28
51
52 70). Nevertheless, the minimum-none solution
found by GD can have a suboptimal performance [64].
Following our theoretical results, we formulate changing the distribution of a training dataset such that different features are learned at more uniform speed. First, we prove that the model output for examples containing features that are learned early by GD is separable from the rest of examples in their class. Then, we propose changing the data distribution by (i) identifying a cluster of examples with similar model output early in training, (ii) upsampling the remaining examples once to speed up their learning, and (iii) restarting training on the modified training distribution. Our method, UpSample Early For Uniform Learning (USEFUL), effectively alleviates the simplicity bias and consequently improves the generalization performance. Intuitively. learning features in a more uniform speed prevents the model to overfit underrepresented but useful features that otherwise are learned in late training stages. When the model overfits an example, it cannot learn its features in a generalizable manner. This harms the generalization performance on the original data distribution.
We show the effectiveness of USEFUL in alleviating the simplicity bias and improving the general- ization via extensive experiments. First, we show that despite being relatively lightweight, USEFUL effectively improves the generalization performance of SGD and SAM. Additionally, we show that USEFUL. can be easily applied with various optimizers and data augmentation methods to improve in- distribution generalization performance even further. For example, applying USEFUL with SAM and TrivialAugment (TA) [50] achieves, to the best of our knowledge, state-of-the-art accuracy for image classification for training ResNet18 on CIFAR 10, STL CINICIO, Tiny- Imageemt,,Net ResNet34 on CIFAR FEAR110: 10; and VGGI9 and DenseNet121 on CIFARIO. We also empirically confirm the benefits of USEFUL to out-of-distribttion performance, but we emphasize that this is not the focus of our work.
2 Related Works
ceappness-awarr-miiine
(SAM). Motivated by the generalization advantages of flat local
minima, sharpness-ane aware minimization (SAM) was concurrently proposed in [22] 81 to minimize the training loss at the worst perturbed direction from the current parameters. SAM has been shown to
obtain rate-of-the-an on a variety of tasks 22
Additionally, SAM has been shown to be beneficial
in other settings, including label noise (22 81 and domain generalization 9 [2).
There have been recent efforts to understand the generalization benefits of SAM. The most popular ex- planation is based on the Hessian spectra, empirically (22] 36) and theoretically [6] 73). Other works showed that SAM finds a sparser solution in diagonal linear networks 3. and exhibits benign overfit- ting under much weaker signal strength compared to (S)GD [1]. More recently. SAM is shown to also benefit -ii-----ddttiitiionn (OOD). In particular, (65 suggested that SAM promotes diverse feature learning by empirically studying a simplified version of SAM which only perturbs the last layer. They showed that SAM upscales the last layer's weights to induce feature diversity, which benefits OOD, In contrast, we rigorously analyze a 2-layer non-linear CNN and prove that SAM learns (the same set of) features at a more uniform speed, which benefits the in-distribution (ID) settings. Our results reveal an orthogonal effect of SAM that benefits the ID generalization by reducing the simplicity bias, and provides a complementary view to prior works explaining superior ID generalization performance of SAM. We then propose method to learn features more evenly by changing the data distribution.

---

Simplicity bias (SB). (S)GD has an inductive bias towards learning simpler solutions with minimum norm [25] It is empirically observed [4] and theoretically proved [28] that SGD learns linear
functions in the early training phase and more complex functions later in training. SB of SGD has been long conjectured to be the reason for the superior in-distribution generalization performance of overparameterized models, by providing capacity control or implicit regularization [26] 52 55 [3]
On the other hand, in the OOD setting, simplicity bias is known to contribute to shortcut learning by causing models to exclusively rely on the simplest spurious feature and remain invariant to the
complex but more predictive features [63 67
751. Prior works on mitigating simplicity bias have been shown effective in the OOD settings 67 68 In contrast, our work shows, for the first time, that reducing the simplicity bias also benefits the ID settings. By studying the mechanism of feature learning in two-layer nonlinear CNN, we prove that SAM is less susceptible to simplicity bias than GD, in particular early in training, which contributes to its superior performance. Then, we show that training data distribution can be modified to reduce the SB and improve the in-distribution generalization. In Appendix [D77] we empirically confirm that existing simplicity bias mitigation methods also improve the in-distribution performance, but smaller extent than ours.
Distinction from Existing Settings. Our work is distinct from the following literature:
(1) Distribution Shift. Unlike distribution shift and shortcut learning 18 39 57 61]. we do not assume existence of domain- dependentt (non-gen generallizzabl) features or strong spurious correlations in the training data, or shift between training and test distribution. We focus on in di-striiibtionn generalization, where training and test distributions are the same and all the features in the training data are relevant for generalization. In Appendix D.5 we empirically show the benefits of our method to distribution shift, but we emphasize that this is not the focus of our study and we leave this direction to future work.
(2) Long-tail distribution. Loigetailed data is studied as a special case of distribution shift in which (sub) classes are highly imbalanced in training but are (more) balanced in test data [1 71. Long-tail methods resample the data at the class or subclass level to match the training and test distribution. In contrast, in our settings, training and test data follow the same distribution. Nevertheless, our method can be applied to improve the performance of long-tail tail datasets, as we confirm in Appendix D.5
(3) Improving Convergence. A body of work speeds up convergence of (S)GD to find the same solution faster. Such methods iteratively sample or reweight examples based on loss or gradient norm during training (21 33 35 80) In contrast, our work does not intend to speed up training to find the same solution faster, but intends to find a more generalizable solution on the original data distribution.
(4) Data Filtering Methods. Filtering methods identify and discard or downweight noisy labeled 45 domain mismatched [23], redundant [4] (9]. or adversarial examples crafted by data poisoning attacks (66). In contrast, we assume a clean training data and no mismatch between training and test distribution. Our work can be applied to a filtered training data to further improve the performance.
3 Theoretical Analysis: SAM Learns Different Features More Evenly
In this section, we analyze and compare feature learning mechanism of SAM. First, we introduce our theoretical settings including data distribution and neural network model in Sec. We then revisit the update rules of GD and SAM in Sec. 3.2 before presenting our theoretical results in Sec. 3.3
3.1 Theoretical Settings
Notation. We use lowercase letters, lowercase boldface letters, and uppercase boldface letters to denote scalars (a), vectors (v), and matrices (W). For vector v, we use ||||| to denote its Euclidean norm. Given two sequence } and {yn }. we denote In = Ozy,) if [xn]] < Cilya/ for some absolute positive constant C₁, In = 22y,) if [x, Z C2 yn| for some absolute positive constant C2. and x,, = (g,.) if Czlval S |z,,| < Calya, for some absolute constant Cz,, C2 > O. Besides, we use Ō(-), (()), and @ to hide logarithmic factors in these notations. Furthermore, denote Ii, D) for some positive constant D, and FT = polylog(yn) if x., = poly(log(y.)))
poly(ya) if xr, = O
Data distribution. We use popular data distribution used in recent works on feature learning 2 g 12 18 32 40 to represent data as a combination of two features and noise patches. Additionally, we introduce a probability a to control the frequency of fast-learnable features in the data distribution.
3

---

Definition 3.1 (Data distribution). A data point (x, y) E (R"))*
Px {=1} is generated from the
distribution (Be,, Sa, a:) follows. We uniformly generate the label y E {+1}. We generate æ as
collection of P patches: x (x(1), x(2) ....., x(P) E (Id)), where
Slow-learnable Feature. One and only one patch is given by Ba y va with ||eall2 (We, va) = 0, and 0 I Ba < 8e E R.
Fast-learnable feature. One and only one patch is given by 3e y U. with | = 1 with probability 1. With probability of a, this patch is masked, i.e. O.
Random noise. The rest of drawn from (0, (o)//)) Ia) with op as an absolute constant.
patches are Gaussian noise & that are independently
For simplicity, we assume P 3, and the noisy patch together with two features form an orthogonal set. Coefficients Be and 3a characterize the feature strength in our data model. A larger coefficient means that the corresponding feature is learned faster.
Two-layer nonlinear CNN. To model modern tfateoof-the-ar architectures, we analyze two-layer nonlinear CNN which also used in [8 [1 data distribution that does not require a fixed position of patches as defined above. Formally,
18
32 40.. Unlike linear models, CNN can handle
(<:W) = Σ JE(J p=1
a ((w),x(()))
(1)
where w Red is the weight vector of the j-th filter, J is the number of filters (neurons) of the network, and (=) = zá is the activation function, i.e., the main source of non-linearity. W = [w).... (w.....w..) E Rdx is the weight matrix the CNN. Following ( 18 32]. we assume mild over- parameterization with J = polylog(d). We initialize (0))) ~ (O, of) where of = polylog(d)/d
3.2 Empirical Risk Minimization: GD vs SAM
Consider N samplele training dataset D (zx W in which each data point is generated from the data distribution in Definition 3.1
The empirical loss function of model f(x: W) reads
N Σ (u.s(ar; N
1
= C(W)
where / is the logistic loss defined as ((z) = log(1 + exp(-2)). The solution W of the empirical risk minimization (ERM) minimizes the above loss, i.e., W* = arg minw L(W).
GD. Typically, ERM is solved using gradient descent (GD). The update rule at iteration of GD with learning rate 7 0 reads
( = ( - v ((W())
(3)
SAM. To find solutions with better generalization performance, Foret et al. 22 proposed the N-SAM algorithm that minimizes both loss and curvature. SAM's update rule iteration reads
((+1) ( ( + VLCC((()))
(4)
where ((0) p the inner step size that is usually normalized by gradient norm, i.c. (t) p/ [
3.3 Comparing Learning Between fast-learnable & slow-learnable Features for GD SAM
Next, we present our theoretical results training dynamics of the two-layer nonlinear CNN using GD and SAM. We characterize the learning speed of features by studying the growth of the model outputs before the activation function, i.e., (w0 v.) and (w!)) vz). We first prove that early in training, both GD and SAM only learn fast-learnable feature. 'The, we show SAM learns slow- learnalreelle and fast-lcarnable features at more uniform speed.

---

Theorem 3.2 (GD Feature Learning). Consider training two-layer nonlinear CNN model initial- ized with ( (0, og) on the training dataset D {z yr
y:) i with distribution D( 34. a)
with a o//3 Be > Sa. For small-enouge learning rate TJ. after training for Tep iterations, w.h.p., the model: (1) learns the fast-learnable feature Ve: [e: (w)) (Ten)
Ue) ( (113e))) does not learn
(2mm) the slow-learnable feature Uu :maxje(J] (wu) .ud) = (
Theorem 3.3 (SAM Feature Learning). Consider training a two-layer nonlinear CNN model initialized with W (O, on the training dataset D = {( (:) W:)] with distribution D(Be: a) with a 1/133, Be > .a-- For small-enough learning rate dj and perturbation radius p. after training for TSAM > Tan iterations, w.h.p., the model: (1) learns the fast-learnable fea- ture Ue : maxje[J] (w) (Tuu) .ve) 2 St(L/Be): (2) does not learn the slow-leamable feature va maxje(J] (w)) (Tuu), vua ((o)).
The detailed proof of Theorems 3.2 and 3.3 are deferred to Appendices A.1 and A.2
Discussion. Note that larger value of (t) v) for E (v., indicates better learning of the feature vector v by neuron w, at iteration t. From the above two theorems, the growth rate of the fast-learnable feature is significantly faster than that of the slow-leamable feature. As small portion (1 a) of the dataset does not have the fast-learnable feature, the model needs to learn the slow-lcarnable feature to improve the performance.
Next, we show that SAM learns last-leae and slow-learnable features more evenly. We denote by Gl = max (w)) (t) .u)) and G = maxje[J] (w)) va) the alignment of model weights with fast-learnable and doa-lace features, when training with GD. Similarly, we denote by s(() and { the alignment of model weights with lest-learnable and slow- leamnabb,,, when training with SAM.
Theorem 3.4 (SAM learns features more evenly than GD). Consider the same model and training dataset as Theorems 3.2 and 3.3 Assume that the learning rate n and the perturbation radius p are sufficiently small. Starting from the same initialization, the growth of fast-leurnable and slow- learnable features in SAM is more balanced than that SGD, i.e., for every iteration E [1, To]:
$(00 S(
G G
We prove Theorem by induction in Appendix A.2 and back it by toy experiments in Scction 3.1
Discussion, Intuitively, our proof is based on the fact that the difference between the growth of fast-learnable and slow-learnable features in SAM is smaller than that of GD. Thus, starting from the same initialization, the slow-learnable feature contributes relatively more to the model prediction in SAM than it does in SGD. Thus, the alow-learnable feature benefits SAM, by reducing its overreliance on the fast-learnable features. We note that neural networks are nonlinear, small change in the output can actually result in a big change in the model and its performance. Even in the extreme setting when two features have identical strength and the fast-learnabee feature exists in all examples, i... B. = 3r = a = 1, the gap in Eq. 5 is significant as we confirm in Figure S in Appendix D
Remark. The network often overfits slow-learnbbe features that are learned late during the training and do not learn them in generalizable manner. This harms the generalization performance on the test set sampled from the original data distribution.
Theorems 3.2 and 3.3 show that we can make the model learn more from the slow-learnable feature by increasing the value of Ba. Based on this intuition, we have the following theorem.
Theorem 3.5 (One-shot upsampling). Under the assumptions of Theorems 3.2 3.3) for a sufficiently small noise, from any iteration L during early training, we have the following results:
1. The slow-learnable featurr has larger contribution to the normalized gradient of the 1-step SAM update, compared to that of GD.
2. Amplifying the strength of the slow-learnable feature increases its contribution to the normalized gradients of GD and SAM.
3. There exists an upsampling factor s.1. the normalized gradient of the I-step GD update on
D (Be, kBd, a) recovers the normalized gradient of the 1-step SAM update on D(3e, Ba, a)).
5

---

Algorithm UpSample Early For Uniform Learning (USEFUL)
Input: Original dataset D, Model f(", W Separating epoch Total epochs T. Train the model f(-, W()) on D for epochs.
for every class c E D do
(C1, & E--mean([[(ee W()))
== C2. where C2 is the cluster with higher average loss
end for
Train f(:: w()) on D for T epochs
Output: Model f(; W(T)
airplane
automabile
deer
dog
frog
horse
ship
Taa matee Sanplamaabee
Figure 1: Examples of slow-learnable (top) and fast-learnable (bottom) in CIFAR-10 found by method. Examples in the top row (slow-Icarnable) are harder to identify visually and look more ambiguous (part of the object is in the image the object is smaller and the area associated with the background is larger). contrast, examples in the bottom row (fast-learnable) are not ambiguous and clear representatives of their corresponding class, hence very easy to visually classify (the entire object is the image and the area associated with the background small).
given in Appendix A.4 We see that we can learn features at more uniform speed by training on new dataset D( d ax) with a larger strength 8 > Bg. But, the value of coefficient 8y varies by the model weights and gradient at each iteration t.
Discussion. Proof of Theorem 3.5
Remark. Intuitively, Theorem 3.5 implies that by descending over a flatter trajectory, SAM learns slow-learnable features relatively earlier in training. compared to GD. While the largest difference between feature learning of SAM and (SJGD is attributed to early training dynamics (due to the simplicity bias of (S)GD and its largest contribution early in training), SAM learns features at more uniform speed during the entire training. This effect is. however, difficult to theoretically characterize exactly. Therefore, learning features at a more uniform speed via SAM help yield flatter minima with better generalization performance. We note that while SAM learns fast-learnable and slow-learnable features more uniform speed, still suffers from simplicity bias and learns fast-learnable features carlier (although less so than GD) evidenced in our Theorem 3.3
4
Method: UpSample Early For Uniform Learning (USEFUL)
Motivated by our theoretical results, we aim to speed up learning the dlow-learnable features in the training data. This drive the network to learn fast-learnable and slow-learnable features a more uniformly speed, and ultimately improves the in-distribution generalization performance.
Step 1: Identifying examples with fast-learnable features. As shown in Theorems 3.2 and 3.3
fast-learnable features are learned early in training, and the model output for examples containing fast-learnable features are highly separable from the rest of examples in their class, early in training. This is illustrated for one class seek to find cluster of examples with similar model outputs early training. To do so, apply k-means clustering to the last-layer activation vectors of examples in every class, to separate examples with fast-learnable features from the rest of examples. Formally, for examples in every class with = we find:
toy example and CIFAR- in Fig. 2 Motivated by our theory,
arg min
Σ (E.2.) ECC
Σ
[| (()) - 12,
(6)

---

y=1
slow-learnble
slow larnble
fast-learablle
2-
fast-leamable
PC2
0.50
0.25
-2
0.00
PC1
PC1
Figure 2: TSNE visualization of output vectors. (left) ResNer1 10 epoch 8. (right) CNN/toy data generated based on Definition 3.1 with Ba = 0.2, B. = a = 0.9, iteration 200.
2.0
0.6
& 2. 51.1%
96.6%
-E-AAA 56505
5/1
0.4
1.0
0.5
0.5
0.0
0.0
100
230
300
400
500
600
600
100
300
(a) Bg 0.2
(b) 3g 0.4
(c) 88 0.2 vs. Bg = 0.4
Figure 3: GD (blue) vs. SAM (orange) on toy datasets. Data is generated based on Definition 3.1 with different 3a and fixed B. 1, 0.9. and -- lines denote the alignment (i.e., inner product) of fast-learnable (v.) and slow-learnable (v,) features with the model weight (w
(w))). ()), (a). (b) GD and SAM first learn the fast-learnable feature. Notably, GD learns the fast-learnable feature very early,
(c) Test accuracy GD & SAM improves by increasing the strength of the slow-learnable feature.
where μi is the center cluster S,. The cluster with lower average loss will contain the majority of ex- amples containing fast-learnable features, whereas the remaining examples contain slow-lcarnable fea- tures in the training data. Examples of images in fast-learnable and slow-learnable clusters of CIFAR- 10 found by USEFUL are illustrated in Fig. I
The choice of clustering. Our choice of clustering motivated by our Theorems 3.2 and 3.3 which show that examples with fast-learnable features are separable based on model output from the rest of examples in their class. While examples with fast-learnable features are expected to have lower loss, loss of examples may oscillate during the training and makes it difficult to find an accurate cut-off for separating the examples. Besides, as fast-lea features may not be fully learned early in training, examples containing fast-learnable features may not necessarily have the right prediction, thus misclassification cannot separate examples accurately. In contrast, clustering does not require hy- perparameter tuning and performs well for separating examples, as we confirm in our ablation studies.
One-shot upsampling of slow-learuable features. Next, we upsample examples that are not in the cluster of points containing fast-learnable features. This speeds up learning slow- learnable features and encourages the model to learn different features it improves the in-distribution performance based on Theorem 3.5 As discussed earlier, the number of times we upsample these examples should change based on the model weight at each iteration. Hence, a multi-stage clustering and sampling can yield the best results. Nevertheless, we empirically confirm that a I-shot algorithm that finds fast-learnable examples at an early training iteration and upsample the remaining examples by factor of & = 2 effectively improves the performance. Notably, contrast to dynamic sampling or reweighting. USEFUL upsamples examples only once and restart training on the modified but fix distribution.
Step 2:
more uniform speed. Thus,
When to separate the examples. is crucial to separate examples early in training, to accurately identify examples that contribute the most to simplicity bias. We empirically verify the intuition that the optimal epoch to separate examples when the change in training error starts to shrink as visualized in Figure 15al More details can be found in Appendices C.2 and D.8
The pseudocode of USEFUL is illustrated in Alg. T and the workflow is shown in Appendix Fig.
7

---

CIFAR10
CIFAR100
SGD
5.0
22-46
SGD SGD+USEFUUL
4.5
SAM
SAM+USEFUL
6.0
SAM+TA
SAM+USEFUL+
2.54
2.6K
1.54
STL10
CINIC10
Tiry-ImageNNt
12.0
35
11.0
13
334L
30.5
10.0
14
12
30
29.21
1x
1.5x
2x
7x
7x
1x
1.4
2x
2,5x
24
2.5g
1x
3.Fz
2x
3.2
7X
3.24
Timg Tra W.T. 5GD
Time to Train MirL 50D
Tio 25 Tan MLTE
500
Figure 4: Test classification error of ResNet18 on CIFARIO, STL10, ImnylmaeeNet and ResNet34 on CIFAR R100. The numbers below bars indicate the approximate training cost and the tick on top shows the std over three runs. USEFUL enhances the performance of SGD and SAM on all datasets. TrivialAugment (TA) further boosts SAM's 's performance (except for CINICI0). Remarkably, USE- FUL consistently boosts the performance across all scenarios and achieves (to our knowledge) SOTA performance for ResNet1 8 and ResNet34 on the selected datasets when combined with SAM and TA.
5 Experiments
Outline. In Sec. 5.1 we empirically validate our theoretical results toy datasets. We then evaluate the performance of USEFUL on several real-world datasets in Sec. [5.2 and different model architectures in Sec. 5.3] In addition, Sec. 5.3 highlights the advantages of USEFUL over random upsampling. Furthermore, we show that USEFUL shares several properties with SAM in Sec. Sec.55.4 Additional experimental results are deferred to Appendix D where show that USEFUL also boosts the performance of other SAM variants, and present promising results for USEFUL applied to the OOD setting (spurious correlation, long-tail distribution), transfer learning. and label noise settings. We further conduct ablation studies on the effect of our data selection strategy for upsampling, training batch size, learning rate, upsampling factor, and separating epoch in Appendix [D8]
Settings. We used common datasets for image classification including C0, CIFARI00 [41], STL10 [3] CINIC1O [16), and Tiny-ImageNet (43] Both CINIC10 and Tiny ImageNet are large-scale datasets containing images from the ImageNet dataset [17] We trained ResNet18 on all datasets except for CIFAR 100 on which we trained ResNet34. We closely followed the setting from (3) in which our models are trained for 200 epochs with a batch size of 128. We used SGD with the momentum parameter of 0.9 and weight decay 0.0005. We also fixed p 0.1 for SAM in all experiments unless explicitly stated. We used linear learning rate schedule starting at 0.1 and decay by a factor of 10 once epoch 100 and again at epoch 150. More details are given in Appendix
5.1 Toy Datasets
our toy dataset consists of training and test sets, each containing 10K examples generated from the data distribution defined in 3.1 with dimension d0 and P P-3. We set -,=1,,a=000 a= 0.9, and ap. /vd = 0.125. We also consider a scenario with larger 3x 8x=00 0.4. We shuffle the order of patches randomly to confirm that our theory holds with arbitrary order of patches.
Datasets. Following 18
Training. We used the two-layer nonlinear CNN in Section 3.1 with J 40 filters. For GD, we set the learning rate to 7 = 0.1 and did not use momentum. For SAM, we used the same base GD optimizer and chose smaller value of the inner step, P 0.02, than other experiments satisfy the constraint in Theorem 3.4 We trained the model for 600 iterations till convergence for GD and SAM.
Results. Figure 3a illustrates that both GD (blue) and SAM (orange) first lear the fast-learnable fea- ture. particular, the blue dotted line (G.) accelerates quickly at around epoch 250 while the orange dotted line (S.) increases drastically much later at around epoch 450, That is: (I) GD learns the fast-
8

---

DonsetNet121
VIT-S
7.0
11.0
21
Error 1h))
10.5
1545
100
20
20:2
13:37
28.37
SAM
SAM +USEFUL
USEFUL
1809
6.0
9.5
19-
9.0
5.5
18
8.0
17
3x
1.3x
2x
2.5x
2x
2.64
1x
1.5x
2%
3x
25
3x
1x
1.45x 28
2.0x
2x. 2.9%
Time to Train w.t.
SGD
Time to Train 12,
SGD
Time to Train W.1.. 5GD
Figure 5: Test classification errors of different architectures on CIFARIO. USEFUL improves the performance of SGD and SAM when training different architectures. TrivialAugment (TA) further boosts SAM's capabilities. The results for 3-layer MLP can be found in Figure
learnable feature very early in training. This is well-aligned with our Theorems 3.2 and 3.3 and their discussion. Furthermore, the gap between contribution of fast-learnable and slow-learnable features towards the model output in SAM (S! - ) is much smaller than that of GD (G - G. That is: (2) fast-learnable and slow-learnable features are learned more evenly in SAM. This validates our Theorem 3.4 From around epoch 500 onwards, the contribution of the slow-learnable feature in SAM surpasses the level of that in GD while the contribution of the fast-learnable feature in SAM is still lower than the counterpart in GD. When increasing the slow-lcarnable feature strength Ba from 0.2 to 0.4 in Figure 3b the same conelusion for the growth speed of fast-learnable and slow- learnable features holds. Notably, there is a clear increase in the classification accuracy of the model trained with either GD or SAM by increasing 3d, as can be seen in Figure Be That is: (3) amplifying the strength of the slow-learnable feature improves the generalization performance. Effectively. this enables the model successfully predict examples in which the fast-learnable feature is missing.
5.2 USEFUL is Effective across Datasets
Figure 4 illustrates the performance of models trained with SGD and SAM on original vs modified data distribution by USEFUL. We see that USEFUL effectively reduces the test classification error of both SGD and SAM. Interestingly, USEFUL further improves SAM's generalization performance by reducing its simplicity bias. Notably, on the STL10 dataset, USEFUL boosts the performance of SGD to surpass that SAM. The percentages of examples found for upsampling by USEFUL for roughly 30%, 50%, 50% 45%,
CIFAR 10, CIFAR 100, STL10, CINIC10, and Tit and 60%, respectively. Thus, training SGD on the modified data distribution only incurs a cost of 1.3x, 1.5x, 1.5x, 1.45x, and 1.6x compared to 2x of SAM.
USEFUL+TA is particularly effective. Stacking strong augmentation methods e.g. TrivialAug- ment (50) further improves the performance, achieving state-of-the-art for ResNet on all datasets. When strong augmentation is combined with USEFUL, it makes more variations of the (upsampled) llonalearnable features and enhances their learning. Hence, it further boost the performance.
5.3 USEFUL is Effective across Architectures & Settings
Model architectures: CNN, ViT, MLP. Next, confirm the versatility of our method, by applying it to different model architectures including 3-layer MLP, CNNs (ResNer18, VGG19, DenseNet and Transformers (ViT-S). Figure 5 shows that USEFUL is effective across different model architectures. Remarkably, when applying to non-CNN architectures, it reduces the test error of SGD to lower level than that of SAM alone. Detailed results for 3-layer MLP is given in Appendix D.2
batch-size, learning rate, and SAM variants. In Appendix D we confirm the effectiveness of USEFUL for different batch sizes of 128. 256, 512, and different initial learning rates of 0. 0.2. 0.4. In Appendix D.4, we confirm that USEFUL applied to ASAM 422 a SAM variant which uses scale-invariant sharpness measure- wreasre-further reduces the test error.
Settings:
USEFUL vs Random Upsampling. Fig. 6 shows that USEFUL considerably outperforms SGD
and SAM on randomly upsampled CIFARTO 10 & CIFAR 100. This confirms that the main benefit of USEFUL due to the modified distribution and not longer training time. In Appendix [D88 we also confirm that upsampling outperforms upweighting for SAM & SGD.

---

CIFAR10
CIAR100
25
5.2
Drgg
Rand
USEPUL
23
21
4.0
3.8
20
Figure 6: USEFUL vs, Random Upsampling, when training ResNet18 on CIFARIO and CIFARI00)
5.4 USEFUL's Solution has Similar Properties to SAM
SAM & USEFUL Find Sparser Solutions than SGD. [ showed that SAM's solution has better sparsity-induee inducing property indicated by the L1 norm than the standard ERM. Fig. [0 shows the L1 norm of ResNet18 trained on CIFARIO and ResNet34 trained on CIFARI00 at the end of training. We see that USEFUL drives both SGD and SAM to find solutions with smaller LI norms.
SAM& & USEFUL Find Less Sharp Solutions than SGD. While our goal is not to directly find a flatter minimum or the same solution as SAM, we showed that USEFUL finds flatter minima. Following [3. we used the maximum Hessian eigenvalue (Amaa) ) and the bulk of the spectrum (Amaz/ /As) [30. which are commonly used metries for sharpness [T0 31 37 74] Table illustrates that SGD+USEFUL on CIFAR16 10 reduces sharpness metrics significantly compared to SGD, proving that USEFUL successfully reduces the sharpness of the solution. We note that to capture the sharpness/flatness, multiple different criteria have been proposed (largest Hessian eigenvalue and bulk of Hessian), and one criterion not enough to accurately capture the sharpness. While the solution of SGD+USEFUL has a higher largest Hessian eigenvalue than SAM, it achieves the smallest bulk.
669 count the number of times an example is misclassified after being correctly classified during training and is an indicator of the learning speed and difficulty of examples. We show Appendix (D.3 both SAM and USEFUL suc- cessfully reduce the forgetting scores, thus learn dow-llcae features faster than SGD. This aligns with our Theorem 3.4 and results on the toy datasets. By upsampling slow-learnable examples in the dataset, they contribute more to learning and hence SGD+USEFUL learns them faster than SGD.
SAM & USEFUL Reduce Forgetting Scores. Forgetting scores
USEFUL also Benefits Distribution Shift. While our main contribution is providing a novel and effective method to improve the in-distribution generalization performance, we conduct experiments confirming the benefits of our method distribution shift. We discuss this experiment and its results in Appendix D.5 On Waterbirds dataset [61] with strong spurious correlation (95%). both SAM and USEFUL successfully improve the performance on the balanced test set by 6.21% and 5.8%, respec- tively. We also show the applicability of USEFUL to fine-tuning ResNet50 pre-trained on ImageNet.
Conclusion
In this paper, we made the first attempt to improve the in-distribution generalization performance of machine learning methods by modifying the distribution of training data. We first analyzed learning dynamics of sharpness-award minimization (SAM), and attributed its superior performance over GD to mitigating the simplicity bias, and learning features at more speed. Inspired by SAM, we upsampled the examples that contain how-learnbble features to alleviate the simplicity bias. This allows learning features more uniformly, thus improving the performance. Our method boosts the performance image classifiers trained with SGD or SAM and casily stacks with data augmentation.
Acknowledgments
This research was partially supported by the National Science Foundation CAREER Award 2146492, National Science Foundation 2421782 and Simons Foundation, Cisco Systems, Optum AI, and UCLA Hellman Fellowship.
10

---

References
[1] Amro Abbas, Kushal Tirumala, Dániel Simig, Surya Ganguli, and Ari S Morcos. Semd- edup: Data-efficient learning at web-scale through semantic deduplication. arXiv preprint arXiv:2303.0955 09540, 2023.
[2] Zeyuan Allen-Zhu and Yuanzhi Li. Towards understanding ensemble, knowledge distillation and self-distillation deep learning. arXiv preprini urXiv:20/2.09816 2020.
[3] Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware minimization. In International Conference on Machine Learning, pages 639-668. PMLR, 2022.
[4] Maksym Andriushchenko, Dara Bahri, Hossein Mobahi, and Nicolas Flammarion. Sharpness- aware minimization leads to low-rank features. arXiv preprint arXiv:2330. 2023.
[5] Robert Baldock, Hartmut Maennel, and Behnam Neyshabur. Deep learning through the lens of example difficulty. Advances in Neural Information Processing Systems, 14:76666-,888,
2021.
[6] Peter L Bartlett, Philip M Long, and Olivier Bousquet. The dynamics of sharpness-awar minimization: Bouncing across ravines and drifting towards wide minima. Journal of Machine Learning Research, 24(316):1-36, 2023.
[7] Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconciling modern machine- learning practice and the classical bias varianee trade-off. Proceedings of the National Academy of Sciences, 116(32)15849-15585
2019.
[8] Yuan Cao, Zixiang Chen, Misha Belkin, and Quanquan Gu. Benign overfitting in two-layer convolutional neural networks. Advances in neural information processing systems, 35:25237- 25250, 2022.
[9] Junbum Cha, Sanghyuk Chun, Kyungjae Lee, Han-Cheol Cho, Seunghyun Park, Yunsung Lee, and Sungrae Park. Swad: Domain generalization by seeking flat minima. Advances in Neural Information Processing Systems, 34:22405-22418,, 2021.
[10] Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-sgd: Biasing gradient descent into wide valleys. Journal of Statistical Mechanics: Theory and Experiment, 2019(12): 124018, 2019.
[11] Zixiang Chen, Yihe Deng, Yue Wu, Quanquan Gu, and Yuanzhi Li. Towards understanding the mixture-of-experts layer in deep learning. Advances in neural information processing systems, 35:23049-2306.. 2022.
[12] Zixiang Chen, Junkai Zhang, Yiwen Kou, Xiangning Chen, Cho-Jui Hsieh, and Quanquan Gu. Why does sharpness-aware minimization generalize better than sgd? arXiv preprint arXiv:2310.07269, 2023.
[13] Adam Coates, Andrew Ng, and Honglak Lee. An analysis of single-layer networks in unsuper- vised feature learning. In Proceedings of the fourteenth international conference on artificial intelligence and statistics, pages 215-223. JMLR Workshop and Conference Proceedings. 2011.
[14] Elliot Creager, Henrik Jacobsen, and Richard Zemel. Environment inference for invariant learning. In International Conference on Machine Learning, pages 2189-2200. PMLR, 2021.
(15) Yin Cui, Menglin Jia, Tsung- Yi Lin, Yang Song, and Serge Belongie. Class-balanced loss based on effective number of samples. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 9268 9277, 2019.
[16] Luke N Darlow, Elliot J Crowley, Antreas Antoniou, and Amos J Storkey. Cinic- 10 is not imagenet or 10. arXiv preprint arXiv: 2018.
[17] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248-255 Icee, 2009.
[18] Yihe Deng, Yu Yang, Baharan Mirzasoleiman, and Quanquan Gu. Robust learning with progressive data expansion against spurious correlation. arXiv preprint arXiv:2306.044: 2023.
11

---

[19] Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In International Conference on Machine Learning, pages 1019-1028. PMLR, 2017.
[20] Jiawei Du, Hanshu Yan, Jiashi Feng, Joey Tianyi Zhou, Liangli Zhen, Rick Siow Mong Goh, and Vincent YF Tan. Efficient sharpness-aware minimization for improved training of neural networks. arXiv preprint 2rXiv:2/10.03/4/) 2021.
[21] Ayoub El Hanchi, David Stephens, and Chris Maddison. Stochastic reweighted gradient descent. In International Conference on Machine Learning, pages 8359-8374. PMLR, 2022.
[22] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware mini- mization for efficiently improving generalization. arXiv preprint arXiv:2010.01412, 2020.
[23] Samir Yitzhak Gadre, Gabriel Ilharco, Alex Fang. Jonathan Hayase, Georgios Smyrnis, Thao Nguyen, Ryan Marten, Mitchell Wortsman, Dhruba Ghosh, Jieyu Zhang, et al. Datacomp: In search of the next generation of multimodal datasets. arXiv preprint arXiv:2304. 11004.14. 14108, 2023,
[24] Behrooz Ghorbani. Shankar Krishnan, and Ying Xiao. An investigation into neural net opti- mization via hessian eigenvalue density. In International Conference on Machine Learning, pages 2232-2244 PMLR, 2019.
[25] Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Sre- bro. Implicit regularization in matrix factorization. Advances in neural information processing systems, 30, 2017.
[26] Katherine Hermann and Andrew Lampinen. What shapes feature representations? exploring datasets, architectures, and training. Advances in Neural Information Processing Systems, 33: 9995-10006, 2020.
[27] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 2022.
[28] Wei Hu, Lechao Xiao, Ben Adlam, and Jeffrey Pennington. The surprising simplicity of the early-time learning dynamics of neural networks. Advances in Neural Information Processing Systems, 33:17116-17122 2020.
[29] Gareth James, Daniela Witten, Trevor Hastie, Robert Tibshirani, al. An introduction to statistical learning, volume 112. Springer, 2013.
[30] Stanislaw Jastrzebski, Maciej Szymczak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, and Krzysztof Geras. The break-even point on optimization trajectories of deep neural networks. arXiv preprint arXiv:2002.09572. 2020.
[31] Stanistaw Jastrzkebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in sgd. arXiv preprint arXiv:1711.04625 2017.
[32] Samy Jelassi and Yuanzhi Li. Towards understanding how momentum improves generalization in deep learning. In International Conference on Machine Learning, pages 9965-10040. PMLR, 2022.
[33] Tyler B Johnson and Carlos Guestrin. Training deep models faster with robust, approximate importance sampling. Advances in Neural Information Processing Systems, 31, 2018.
[34] Dimitris Kalimeris, Gal Kaplun, Preetum Nakkiran, Benjamin Edelman, Tristan Yang. Boaz Barak, and Haofeng Zhang. Sgd on neural networks learns functions of increasing complexity. Advances neural information processing systems, 32, 2019.
[35] Angelos Katharopoulos and François Fleuret. Not all samples are created equal: Deep learning with importance sampling. In International conference on machine learning, pages 2525-2534
PMLR, 2018
[36] Simran Kaur, Jeremy Cohen, and Zachary Chase Lipton. On the maximum hessian eigenvalue and generalization. In Proceedings on, pages 51-65. PMLR, 2023.
[37] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On batch training for deep learning: Generalization gap and sharp minima. arXiv preprint 1e0.v:1009.044400 04836, 2016.
12

---

[38] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[39] Polina Kirichenko, Pavel Izmailov, and Andrew Gordon Wilson. Last layer re-training is sufficient for robustness to spurious correlations. arXiv preprint urXiv:2244.02937 2022.
[40] Yiwen Kou, Zixiang Chen, Yuanzhou Chen, and Quanquan Gu. Benign overfitting in two-layer relu convolutional neural networks. In International Conference on Machine Learning, pages 17615-17659. PMLR, 2023.
[41] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[42] Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. Asam: Adaptive sharpness- aware minimization for scale-invariant learning of deep neural networks. In International Conference on Machine Learning, pages 5905-5914. PMLR, 2021.
[43] Ya Le and Xuan Yang. Tiny imagenet visual recognition challenge. CS 23IN, 7(7):3, 2015.
[44] Katherine Lee, Daphne Ippolito, Andrew Nystrom, Chiyuan Zhang, Douglas Eck, Chris Callison- Burch, and Nicholas Carlini. Deduplicating training data makes language models 2021.
better. arXiv preprint arXiv:2107.0645
[45] Junnan Li, Richard Socher, and Steven CH Hoi. Dividemix: Learning with noisy labels as semi-supervised learning. arXiv preprint arXiv:2002.07394, 2020.
[46] Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. Advances in neural information processing systems, 32, 2019.
[47] Chenxi Liu, Barret Zoph, Maxim Neumann, Jonathon Shlens, Wei Hua, Li-Jia Li, Li Fei-Fei, Alan Yuille, Jonathan Huang, and Kevin Murphy. Progressive neural architecture search. In Proceedings of the European conference on computer vision (ECCV), pages 19-34, 2018.
[48] Evan Z Liu, Behzad Haghgoo, Annie S Chen, Aditi Raghunathan, Pang Wei Koh, Shiori Sagawa, Percy Liang, and Chelsea Finn. Just train twice: Improving group robustness without training group information. In International Conference on Machine Learning, pages 6781-6792 PMLR, 2021.
[49] Song Mei and Andrea Montanari. The generalization error of random features regression: Precise asymptotics and the double descent curve. Communications on Pure and Applied Mathematics, 75(4):667-766, 2022.
[50] Samuel G Müller and Frank Hutter. Trivialaugment: Tuning- fining-free yet state-of-the the--ar-rttt art data augmentation. In Proceedings of the IEEE/CVF international conference on computer vision, pages 774-782, 2021.
[51] Preetum Nakkiran, Gal Kaplun, Yamini Bansal, Tristan Yang, Boaz Barak, and Ilya Sutskever. Deep double descent: Where bigger models and more data hurt. Journal of Statistical Mechanics: 2021.
Theory and Experiment,
[52] Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. arXiv preprint arXiv:1412.66 2014.
[53] Vardan Papyan. The full spectrum of deepnet hessians at scale: Dynamics with sgd training and sample size. arXiv preprint htt::1811.00066 2018.
[54] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, perf-meeformance deep learning library. Advances in neural information processing systems, 32, 2019.
[55] Mohammad Pezeshki, Oumar Kaba, Yoshua Bengio, Aaron C Courville, Doina Precup, and Guillaume Lajoie. Gradient starvation: A learning proclivity in neural networks, Advances in Neural Information Processing Systems, 14:62112 2021.
[56] Hieu Pham, Melody Guan, Barret Zoph, Quoe Le, and Jeff Dean. Efficient neural architecture search via parameters sharing. In International conference on machine learning, pages 4095- 4104. PMLR, 2018
13

---

[57] Aahlad Manas Puli, Lily Zhang, Yoav Wald, and Rajesh Ranganath. Don't blame dataset shift! shortcut learning due gradients and cross entropy. Advances in Neural Information Processing Systems, 36:71874-71916 2023.
[58] Aahlad Manas Puli, Lily Zhang, Yoav Wald, and Rajesh Ranganath. Don't blame dataset shift! shortcut learning due to gradients and cross entropy. Advances in Neural Information Processing Systems, 36, 2024.
[59] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140):1-61 2020.
[60] Nicolas Roux, Mark Schmidt, and Francis Bach. A stochastic gradient method with an expo- nential convergence -rate for finite training sets. Advances in neural information processing systems, 25, 2012.
[61] Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. arXiv preprint arXiv: /9/ I. 1.08731, 2019.
[62] Mark Schmidt, Nicolas Le Roux, and Francis Bach. Minimizing finite sums with the stochastic average gradient. Mathematical Programming, 162:83-112, 12, 2017.
[63] Harshay Shah, Kaustav Tamuly, Aditi Raghunathan, Prateek Jain, and Pranecth Netrapalli. The pitfalls of simplicity bias in neural networks. Advances in Neural Information Processing Systems, 33:9573- -9755, 2020.
[64] Vatsal Shah, Anastasios Kyrillidis, and Sujay Sanghavi. Minimum norm solutions do not always generalize well for over-parameterized problems. stat, 1050:16, 2018.
[65] Jacob Mitchell Springer, Vaishnavh Nagarajan, and Aditi Raghunathan. Sharpness- aware minimization enhances feature quality via balanced learning. In The Twelfth International Conference on Learning Representations, 2024.
[66] Jacob Steinhardt, Pang Wei W Koh, and Percy S Liang. Certified defenses for data poisoning attacks. Advances in neural information processing systems, 30, 2017.
[67] Damien Teney, Ehsan Abbasnejad, Simon Lucey, and Anton Van den Hengel. Evading the simplicity bias: Training diverse set of models discovers solutions with superior ood general- ization. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 16761-16772, 2022.
[68] Rishabh Tiwari and Pradeep Shenoy. Overcoming simplicity bias in deep networks using feature sieve. In International Conference on Machine Learning, pages 34330-34343, PMLR, 2023.
[69] Mariya Toneva, Alessandro Sordoni, Remi Tachet des Combes, Adam Trischler, Yoshua Bengio, and Geoffrey J Gordon. An empirical study of example forgetting during deep neural network learning. arXiv preprint arXiv:1812.000
2018.
[70] Guillermo Valle-Perez, Chico Q Camargo, and Ard A Louis. Deep learning generalizes because the furnniter-fuunttonn map is biased towards simple functions. In International Conference on Learning Representations, 2018
[71] Grant Van Horn, Oisin Mac Aodha, Yang Song, Yin Cui, Chen Sun, Alex Shepard, Hartwig Adam, Pietro Perona, and Serge Belongie. The inaturalist species classification and detection dataset. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 8769 8769-8778, 8778, 2018
[72] Pengfei Wang. Zhaoxiang Zhang, Zhen Lei, and Lei Zhang. Sharpness-aware gradient matching for domain generalization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 3769-3778, 2023.
[73] Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How does sharpness-awart minimization minimize sharpness? arXiv preprint arXiv:2211.05729, 2022.
[74] Yeming Wen, Kevin Luk, Maxime Gazeau, Guodong Zhang, Harris Chan, and Jimmy Ba. An empirical study of latch-batch stochastic gradient descent with structured covariance noise. arXiv preprint arXiv: 090.99992.8224 2019.
14

---

[75] Yu Yang, Eric Gan, Gintare Karolina Dziugaite, and Baharan Mirzasoleiman. Identifying spuri- ous biases early in training through the lens of simplicity bias. arXiv preprint arXiv:2305.1876. 2023.
[76] Zhewei Yao, Amir Gholami, Sheng Shen, Mustafa Mustafa, Kurt Keutzer, and Michael Mahoney. Adahessian: An adaptive second order optimizer for machine learning. In proceedings of the AAAI conference on artificial intelligence, volume 35, pages 10675, 2021.
[77] Li Yuan, Yunpeng Chen, Tao Wang. Weihao Yu, Yujun Shi, Zi-H Zi-Hang Jiang, Francis EH Tay, Jiashi Feng, and Shuicheng Yan. Tokens-to-token vit: Training vision transformers from scratch on imagenet. In Proceedings of the IEEE/CVF international conference on computer vision, pages 558-567, 2021.
[78] Runtian Zhai, Chen Dan, Zico Kolter, and Pradeep Ravikumar, Understanding why generalized reweighting does not improve over erm. arXiv preprint arXiv:2201.12293, 2022.
[79] Hongyi Zhang. mixup: Beyond empirical risk minimization. arXiv preprint arXiv:1710.09412. 2017.
[80] Peilin Zhao and Tong Zhang. Stochastic optimization with importance sampling for regularized loss minimization. In international conference machine learning, pages 1-9. PMLR, 2015.
[81] Yaowei Zheng, Richong Zhang, and Yongyi Mao. Regularizing neural networks via adversarial model perturbation. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 8156-8165, 2021.
[82] Barret Zoph and Quoe V Le. Neural architecture search with reinforcement learning. arXiv preprint arXiv: 1,01,1778, 2016.
15

---

A Formal Proofs
A.1 Proof of Theorem 3.2
Notation, In this paper, we use lowercase letters, lowercase boldface letters, and uppercase boldface letters to respectively denote scalars (a), vectors (v). and matrices (W). For vector v, we use ||l|| to denote its Euclidean norm. Given two sequence } and {yn J. we denote In = (ya) if [Tll S C1 |n]
Glyya] for some absolute positive constant C₁, Zn (yyn) if |æ,l| 2 C2ly21 for some absolute positive constant C2, and 2, = ((yn) if Csl (al S |xn I Calunl for some absolute constant C3, C4 0. In addition, we use (()),9()) and ( to hide logarithmic factors in these notations. Furthermore, we denote Zn = poly(Un) T. = (UR)) for some positive constant D, and In (Vlylog(Un) if x poly(log(V*)).
First, we have the following assumption for the model weight initialization.
Assumption A.1 (Weight initialization). Assume that we initialize W NNO,, of)) such that for all j E [J], ( ve), (w) (0) va) 2 p > O. (0)
The above assumption is reasonable because we later show that both sequences (w0)" (1))
vz) and (w) vu)) are non-decreasing, So, we can obtain the above initialization by training the model for several iterations. For simplicity of the notation, we assume that o:N is an integer and the first a data examples have the fast-learnable feature while the rest do not. Before going into the analysis, denote the derivative of a data example i at iteration to be
v:f (())
(1) 1 exp(- y.J (2,; W(O) sigmoid(- y (x W())))
Lemma A.2 (Gradient). Let the loss function L be as defined in Equation For 1. 2 and j E [J].
the gradient of the loss L(W()) with regard to neuron w is
oN
w/O ((W())
(1)
(2
(9) Ud)
U Us + 3 (w) (1)
v u v. + vi (,)
N
1=1
N
3
100 (BAtw)) va) bu + y: ( EEPEE)
(8)
N
1=a
Proof. We have the following gradient
exp(-Me W(O))
((W())
(0)
+ exp( (xp(-yij(( fa;; Wir)) yef (())
N
im1
P
3
(t)
Σ
(wo))) ,(()))² x ((0)
Vi
N
1=1
p=l
Σ
(1)
(2)( ( ua) + (: (1) _ v ve +
N
i=1
N
3
Σ
(())
() + 1: (0)
(1) .0.0 ..)
N 1=o
With the above formula of gradient, we have the following equations: Fast-learnable feature gradient. The projection of the gradient on U. is
3 oN
w0 ((W()))
(9)
N
16

---

Slow-learnable feature gradient. The projection of the gradient on is
3
10 ((wi)))
(4)
vdl
(10)
N
Noise gradient. The projection of the gradient on &, is
3 wi ((W()) = - 4.0 uiiu) (9) . 11
Σ k=1,bfi
( yn
(
(x)) 51)
(11)
Derivative of data example i. For 1 Sis NN, (0 can be rewritten as
( sigmoid Σ
3y (w)) (1) ve)² - ( .u)) 8 / ..)3
(12)
while for aN +1 < i N, 1 can be rewritten
((0)
sigmoid
Σ
3g (w)
(1) _ua)² - ve _.))
sigmoid
Σ
(t) 32 { v2)² (w0 v.) y/ E,)3
(13)
Note that ( due the property of the sigmoid function. Furthermore, similarly consider that the sum of the sigmoid terms for all time steps is bounded up to logarithmic dependence [I] The sigmoid term is considered small for x such that
Σ
I exp(x) - EO)), t=
(14)
which implies & 2 (1)).
We present the detailed proofs that build up to Theorem A.8 We begin by considering the update for the fast-learnable and slow-learnable features.
Lemma A.3 (Famallea feature update.). For all 2 0 and j E [J], the the fastlearaable feature update is
(1+1) Ue)) = -ww) v.) + (()) (w) v.)2 (u0)
(15)
where gi (1) = sigmoid
- (10
(1) vatz - 3 (w .u0)
Proof.
Plugging the update rule of GD, we have
(w (1+1) ve))
(w)) (t)
v w
L(W(()), vo)
(1) Ue) + 3n52 aN
= -w
(w ve)
(w) v.) + (m)) (t) (1) v2)2
where the last equality holds due to Lemma B.5
Similarly, obtain the following update rule for slow-learnable features.

---

Lemma A.4 (Slow-learnalle feature update.). For all t 2 0 and j E [J], the fast-learnable feature update is
3n N
N
(w) (+1) vva) = ((w))). va)
Σ
(1)
(w)) vd)?
(16)
which gives
(t) va)² S (w0) (w) Ud) (w)) vu)) S
va)} (17)
where g1 (1()) = sigmoid (Ei=1 - (w0 va) - (2 ve)3).
Proof. Plugging the update rule of GD, we have
(w) l++1) va) = (w0) v w L(Wi()), va)
3n5
(w0) (2) va) +
Σ 1 (w vu)²
N
From Lemma B.5 we have for 1 S / aN,, = (1))g (1) and for aN N, S 4 < 1. Combining with the above equality, we obtain the desired inequalities.
Next, we simplify the two above update rules in the early training stage. Lemma A.5 (Fast-learnable feature update in early iterations). Let To be such that
maxje(J] (w (Ta) .ve) 2 (((/Be)) For t E [0, the fast-learnable feature update has the fol-
lowing rule
(w) (w .V)) = (w0) ve) + 6(n)a vo)²
(18)
(To) Proof. Let To > O be such that either max (w)) ( v.) 2 2(1/3a) or max ( (w) (To) Ua) 2
((1/8.). We will show later that the first condition will be met and we have maxjelJ] (10) (26) va) S
((11Ba)) for all j E [J] and E [0. To).
Recall that 91 (1) sigmoid (E/-1 -B (1) paa)3 133-B33 w have
(0) ( Then, for t E [o, we voo)))
(1)
+ -8
(t) vu))
8 (w)) (2 vv,)3)
2 + exp(x + k)
1 + exp(f((1))
where the first inequality holds due to (w), v,) S x/(J1/33.) and (w (0 vu) I kk(J1/33g) for
t E [0. Tol [18][L E.3]. Therefore, similar to 18 32), we have g (t) = O11) the early iterations. This implies the result in Lemma A.3
ww (++1) ve) = ( v.) + (w)) v.)2
(19)
Similarly, we obtain the following simplified update rule for slow-learnable features in the early iterations.
Lemma A.6 (Slow-learnable feature update in early iterations). Let To > O be such that
maxje(J] (w) (Tp) .ve) 2 ((1/Be). For t E [0, To). the fast-learnable feature update has the fol-
lowing rule"
(w)) (++1) vu)) (w))) .vd) + (()) vva).
(20)
18

---

We next show that GD will learn the fast-learnable feature quicker than learning the slow-learnable fea- ture.
Lemma A.7. Assume = 6(Baoo). Let To be the iteration number that max (w) (Ta) Me) reaches
Q(1/B.)= = 61). Then, we have for all t. S To, it holds that max (To) _uA) = Ō(a)).
Proof. Among all the possible indices j E [J], we focus on the index j = arg max (w) (0)
Ue).
Therefore, for C2 = aB²² ((1), we apply Lemma B.2 with two positive sequences (w /' Ue) and
(w))) vs) defined Lemmas A.5 and A.6 and get
(w)) va) S )((w)) va)) = Ö(
(21)
for all j [J].
Theorem A.8 (Restatement of Theorem 3.2g. We consider training a two-layer nonlinear CNN model
initialized with ( ~ (O0 oī) on the training dataset D = { w)) that follows the data
distribution D( Ba, 0x) with a a1/33e > BB. After training with GD in Equation 3 for TeD terations where
6(1) + 6(1) TOD na.300
log(on3a))
(22)
log(2)
for all E [J] and E [0, TaD), we have
(w j (++1) Ue) (w (1) ve) (())a uv)2
(23)
(w) (1+1) va) (w, (1) va) (3)) (1) vd)2
(24)
After training for TeD iterations, with high probability, the learned weight has the following properties: (I) it learns the fast- lessaaler feature v mmaxjelj] (w)) (Tam)
(To)) ,ue) 2 ((11/3)) (2) it does not learn
the slow-learnable feature Ud max (Ta)) ua) = Ō(oo).
Proof. From the results of Lemmas
A.5 A.7 it remains calculate the time TED. Plugging
(2(1/B)), M (
( into Lemma B.3 have TeD
Teo = 6(1)
log(0n;9c)
+ +
(25)
n05.00
log(2)
A.2 Proof of Theorem 3.3
Before going into the analysis, we denote the derivative of a data example i at iteration to be
(2;; ( (t))) d ff (1)) + +(t))) sigmoid(- yf ( (()),
(26)
where (!)) , V is the weighted ascent direction at the current parameter W We
denote the weight vector of the j-th filter after being perturbed by SAM as
(e) (1) j.e w te, w!) + 10 L(W()). w
(27)
where p(t) = =0>>=p//|V/ [V
19

---

First, we have the following inequalities regarding the gradient norm:
V 3
L(W(O))
(28)
V (
((W(1))))/2
(29)
3,9
2
3,9 oN
2
(t) va)²
+
(1)
w
(t)
N
:=1
N
1/2
+
3
yi (w
3 5.
(8)
(30)
N
Thus,
3 N
/c(W
(
vua)"
(31)
F
333
aN
c
Σ
(w) 00)
(32)
N
Lemma A.9 (Gradient). Let the loss function L be defined in Equation (1)
For t 2 0 and
[J]
the gradient of the loss L(W(1) + =(()) with regard to neuron w (t)
j. is
aN
(9) (WW (0)
3
(t)
va) 20 + (. j. ueluee + vi Se.
( j.e,
N
1.6
-=1
3
Σ
10
(2 va) u0 + v: (0) (..
(33)
N
i=o++1
Proof. We have the following gradient
1 L(W
000000++0)0)) +) =
1
N
+ exp(-yiffee (I;; W(O + +(() y.f'(z;; ))
exp(-yff(I;; W + ((()))
N
N
P
N
3
(1)
1.e
V/
w0
(2) J.c
x (1)))) z(P)
1=1
aN
3
(88 f. va) va + (w fo uu + 52
N
ic
N
3
Σ
(0
(1itw) vu) + vi (e 50.0 5:)
N
i=o++1
With the above formula of gradient, we have the projection of perturbed weight on V. is
(w) j. ve) (ww) Ue. + ( v0)
(1)
=ww .u.) + ( (e) ((W())),
(w0 (0) Ue) ( 3p (1) 8 aN
(w)) _.,)
(34)
N
From Equations
32 and 34 we have
0 (w))
(1) .e) - p S j. ve) (w0))) (2) (ve)
(2)
(35)
20

---

Similarly, the projection of perturbed weight on va is
30 .8
(w (1) j.e: va) ( vu))
Σ
(1) (wo) vv)?
(36)
N
S (w (t) Ua) p (e) va) S (ww)), va)
(37)
Fast-learnable feature gradient. The projection of the gradient on Ue is
3
L(W
v.))
(1) je: ue)²
(38)
N
Slow-learnable feature gradient. The projection of the gradient is
3,33
(V (s) C(W ve)
Σ
(0)
(w
j.e: va)
(0)
(39)
N
Noise gradient. The projection of the gradient on E; is
10 ((W(1) ). -
3 10 (w 6i) 11::11
10 k=1ket
Σ
(Exx
(Ex-
(40)
Derivative of data example i. For 1 sis a., ((62 can be rewritten as
sigmoid Σ
3 va) 8e (w (2 v.)
(41)
while for aN + sis N, 190 can be rewritten as
sigmoid
3 w je: val - 3e (w
sigmoid
Σ
33 m (1)
je v2)²
18 (w0 y1 (1 Se 3
(42)
We present the detailed proofs that build up to Theorem A.15 We begin by considering the update for the last-learnable and deaanlereeble features.
Lemma A.10 (Fast-learnable feature update.). For all t 0 and E [J], the fast-learnable feature update is (w)) ve) + (())a)2 ve) - p) < (t++)) ..) = (w) ve) (m))) v2)²
I (w))
(t) ve) +0
(43)
where g2(()= = sigmoid (E/11 - 32 (wu j.e va)²
3d (w v2)3))
Proof.
Plugging the update rule of SAM, we have
(w (++1) ve) (w) v
((W(()) + +e(1),
3
N
(w)) v.) +
e
(1)
(w jec v2)2
N
= (w)) v.) + (
jet v2)2
where the last equality holds due Lemma B.10 Combining with Equation 35 we obtain the desired inequalities.
D
21

---

Similarly, we obtain the following update rule for slow-lcarnable features.
Lemma A.11 (Slow-learnable feature update.). For all 2 0 and j E [J], the slow-learnable feature update is
3n
N
(wo)), vd) + ((0))82((((()))))))))))))))))))) vn) - -0) S (ww) vn) = (w)) va) +
N
Σ
W
jie
va)
==1
S (wol) va) + (1)) (1) + 11-10)((() (ww (1)
-va))
(44)
where 92(t)= sigmoid - 11-119991
je: val - (w Dc)
(1)
Proof. Plugging the update rule of GD, we have
(w) 11+1) va) = {w)) v (1) ((W(()), va)
(w0) (t) va) + 3m
j.:
N
From Lemma B.5
we have for 1 < aN, (0) = (1)g (t) and for aN N. ((1g, <
( < 1. Combining with the above equality and Equation 37 we obtain the desired inequalities.
Next, we simplify the two above update rules in the early training stage.
Lemma A.12 (Fast-learnable feature update in early iterations). Let To > be such that [a (w (Ta) ve) 2 ((1/Bc). For t E [0, Tol. the fast-learnable feature update has the fol-
lowing rule
vv)) (w)) v.) + (())) (ww)) .v))) ((1) (1) ve) v0)) p)² < (w) Uo) (w))
(45)
Proof. Let To > 0 be such that either maxje(J] (w) (To) v.) > 2((13e) or max (To) va) 2
(((1/aa) We will show later that the first condition will be met and we have maxje[J] (w) (7)) vu) S
((1/Ba) for all j E [J] and E [o, To]
Recall that g2 sigmoid (E/=1 -B va)² (2 (+) ve))))))))))) ve)²). Then, for t E [0, Tol. we have
1
g2(t)=
1 + exp(E)
-B3 (w (2)
jee v2)3 (2(ww) (1) j.e. v2)
+ exp(x + xx)
1 exp()((1))
where the first inequality holds due to (w) ( va) S (/(3/333z) for E [0, To]. Therefore, we have g2(4) = O(1) in the early iterations. Replacing 92 (t) = ((1) into the results of Lemma A.10 we obtain the desired results.
( he: v.) S (wo")) va) S k/(J/33..) and (w va) S
Similarly, we obtain the following simplified update rule for slow-learnable features in the early iterations.
Lemma A.13 leaowlle feature update in early iterations). Let To 0 be such that (Ta) maxje(J] (w j ve) 2 ((1/Be). For E [0, Tol, the fast-learnable feature update has the fol- lowing rule
(1) ((m))(((W))) ,va) p)² < (w (1+1) vua) (", (0())3)(ww)) vva)) )va)))))
(46)
22

---

We next show that SAM will learn the fast-learnable feature quicker than the slow-learnable one. Lemma A.14. Assume ll = o(B2oo). Let To be the iteration number that maXje[[] (w))T)) ve) reaches (1) (1). Then, we have for all S To, it holds that max (w) (Ta) va) = (
Proof. Among all the possible indices j E [J], we focus on the index j = arg max jE[J] (w) (0) Ue).
Therefore, for C. = a (1), we apply Lemma B.2 with two positive sequences (w) jr Ue) and
(ww))) va) defined in Lemmas A.12 and A.13 and get
(t) (w) j va) ((w vv)) Ö(oo)
(0)
(47)
for all j [J]-
Theorem A.15 (Restatement of Theorem 3.3). We consider training a two-layer nonlinear CNN model initialized with ( (0, a on the training datasei D {(a., Us)] that follows the data distribution D(Bes Sa, a) with a 3/333 > Ba. After training with SAM in Equation for Tsam
iterations where
((oo)
O(og)
log 3
TSAM 1j0:32( 0o - p)2
(48)
(00 p)²
log(2)
for all e [J] and E [0, TrMM), we have
(w)") ve) + @(1)a.82 Ue) p) S (w) (1+1)
j Ue) = (w)) (+) ve) + (())as ( j.e: v.)a
S (w) ve) + (1)
v2))
(49)
(w)) va) + +( (9 va) - p) s (w) (1+1) val = (w)) va) (() w (t) va)?
s (w))) va) + ①
v2))
(50)
After training for TSAM iterations, with high probability, the learned weight has the following proper- ties: (1) it learns the fast-learnable feature Ve maxje([]] (w) (Taur) uue 2 (11/Be): (2) it does not learn the slow-learnable feature va mmxjee (w) (Tau) va) = Ö(o)).
it remains to calculate the time TSAM- Plugging v (211/B)), m = M = (, 20 = ( into Lemma B.3 we have TSAM as
Proof. With the results of Lemmas
A.12
14
@(Ov)
log(ouBe)
Tsam 3 p)²
(51)
(00 p)²
log(2)
we can see that SAM learns the fast-learnable features later than GD. Particularly, if remove the approximate notations, have the following inequality
Comparing Eq. 22 and
48
1og(0.8.)
On
of
log(on,.))
(52)
log(2)
1) p)²
(00 p)a log(2)
which holds due to Assumption A.1 about weight initialization in Appendix A.I (ao 2 p2 0).
i...
A.3 Proof of Theorem 3.4
In this section, we show that SAM learns fast-learmable and slow-learnable features at more uniform speed. To ease the notation, we denote G(!) max (w) .ve) and G maxje (w)) va)
(t)
for model weights trained with GD. Similarly, we denote s! and s(t) for model weights trained
with SAM. We use 3(() and ((()) to denote the inner products with perturbed weights. We simplify Equation 34 and 36 for early iterations S To as
$(()) s(() (0
(53)
((t) = s(t) @(1) (
(54)
23

---

Before introducing the theorem, we assume that the model is initialized in favor of the fast- learnable feature, i.c. Ge _ P. This is reasonable as a consequence of Theorem A.8 because we can just train the model for several iterations to achieve this initialization (similar argument for Assumption A.I.
Theorem A.16 (Restatement of Theorem 3.4p. Consider the training dataset D { yo H that follows the data distribution D(Se, 9a, oi) using the two-layer nonlinear CNN model initialized with ( ~ NN, og). Assume that the fast-learnable feature strength is significantly larger a1/3 > Sg. Training the same model initialization, we have that for every iteration t S To
p 5/ 2 5
(55)
d < $(1)
(56)
( < G
(57)
S < G
(58)
5 5 < G G
(59)
Proof. We prove this by induction. For / = 0, the above hypotheses immediately hold because we use train two methods from the same initialization. Particularly, we have < ( -0
=
and ((0) - - 2 S p S N (G! -G( p 0.
Assume that the induction hypotheses hold for t, i.e.
p (1) S ((
(60)
$ < 8(0)
(61)
5 I G!!)
(62)
5 < G
(63)
S(0 s(t < G( G
(64)
We need to prove that they also hold for t + 1. From Lemma A.15 and the first two induction hypotheses,
p+ S4 p ((1) + (0)))
s((+1)
(500 < 50 + (3) = (((-1)
(65)
Then, S( g((+1)
5d > S (++1) p- p-s 0. From Equation
5((+1)
35 and Lemma A.12
s(++1) < 5 + ((1)) < G + 6()), < G(++1)
(66)
5(1 ((++))
< Glltt! (1+1) From Equations 53 and 54
Similarly, we have
8(1) (t) ((1) B3(( () < 6(1)) (( =(()) 5(0
(67)
< 5 - $ < < 5 < c. -G
(68)
Combining with Equations 35 and 37 we have
( - (5 < (S{0)² (5 < (G -G())
(69)
(o(9)² (00))=(((()))2 < (G (8!0)2
(70)
(()))((()))) -(5())) < )
(71)
G - ( + (G) (5())) < G!!) 500 + +0()) (5(0))2)
(72)
((+1)) - s((+1) < G((+1) - S((-1)
(73)
s(++1) s((+1) (t++) < G(++1) (+1)
(74)
Therefore, the induction hypotheses hold for t + 1.
24

---

A.4 Proof of Theorem 3.5
From Theorem A.16 we have the following result for switching between SAM and GD during training.
Lemma A.17. Consider the training dataset D = {(r,, that follows the data distribution
D(Be, Ba, a) using the two-layer nonlinear CNN model initialized with W N o.) Assume that the noise is sufficiently small (ref. Lemmas B.4 B.4and B.9 and the fast-learnable feature strength is significantly larger a 2/3 Be > 3g. From any iteration during early training, the normalized gradient of the one-step SAM update has larger weight on the slow-learnable feature compared to that of GD.
Proof. First, recall the gradients of GD and SAM are as follows.
aN
V ((W()))
3
(0)
(1
(1) va + (2 (1) Vu
to 9
N
1=1
3
Σ
1 (Aitw)) (1) va)2ua (
(w
( PE:
(75)
N
k=N++1
V (2)
L(W(O)
+)(
3
aN
(t) j,e
1 uu + ue (ww S.
(1)
N
1.6
3
(1)
(f ua + y ( (.) E.)
(76)
N
ima++1
Because the noise is sufficiently small, the above equations can be simplified as
N
aN
(n) ( ((W())=- --
(
(w
val2
va
(1) 33 (w (w)))
Ve (77)
N
N
i=1
aN
V ((W(() +c(0)= = 2
3
(0) 3d
(w))
va)a
vd
(1) f.
g3 (0) ,u))²
(w
Ve (78)
N
N
1=1
Note that in early training, we have an approximation for the logit terms as (.0) = (()) = O(1), we can further simplify the gradients
with ((W())) = (w val?)
va - (3a (2o vo))) Ue
(79)
( CW (0) + +) = = (383 (30((w() vul2) va- (30/32 (w) U )2)
(80)
Both gradients of GD and SAM can be decomposed into the linear combination of fast-learnable and slow-learnable features. To prove that the normalized gradient of SAM favors the slow-learnable fea- ture compared GD, it sufficient show the ratio of coefficients SAM larger than GD. In other words, we need to verify that
3,99 (w
Ua)
3,3 (1) va)
(81)
3a 33 Uey 2
w
30,33 (w (0,
w vd)
(1) va)
(82)
(w (0)
(wl ve)
ve))
(t) va) 3p(t) 33 v.a2
(w) v.)
(83)
(wi1)) (1), v=)) 3p(t) 33 1 v.y
(w(1), Ua)
1 30 (+) 3 va) 2 1 3p (1) a 23 (w (1) vv)
(84)
3 ww ve) 2 (w ().) va)
(855)
The last inequality holds due to ag33
83 and (1) (0,v)) > (w va) from Theorem A.16
25

---

From Equation 81 in the above proof. it can be seen clearly that amplifying the slow-learnable feature strength either GD or SAM, i.c., increasing 3a. places a larger weight on the lla-leernable feature. Thus, we have the next theorem.
Theorem A.18 (Restatement of Theorem 3.5). Consider the training dataset D (x,,, Mi )i that follows the data distribution D(3e, 3a, Ox) using the two-layer nonlinear CNN model initialized with ( ~ (00 ot)- Assume that the noise is sufficiently small (ref. Lemmaas R.4 and B.9 and the fast-learnable feature strength is significantly larger a 3/333 > -a-- We have the following results for one-step upsampling, i.e. increasing Ba. from any iteration during early training
1. The normalized gradient of the one-step SAM update has a larger weight on the slow- learnable feature compared to that of GD.
2. Amplifying the llow-learnable feature strength puts a larger weight on the slow-leamable fea- ture in the normalized gradients of GD and SAM.
3. There exists an upsampling factor k such that the normalized gradient of the one-step GD update on D(Be, kBa, a) recovers the normalized gradient of the one-step SAM update on D(Be, (1Be,,a,0)) a.)
Proof. The first result has already been proved in Lemma A.17 Now, consider increasing the slow- learnable feature strength from 34 to 3y Similar to the proof of Corollary A.17 to verify that the normalized of the new normalized gradient of GD favors the lla-llare feature, it sufficient to show
(3')) (w)()) v 3 (w(), v.)² (2
$9 (2)) va)²
(86)
which is trivial because 8 > Bg. Similarly, we can verify the result for SAM. Now, let's find the
new coefficient 8 = kBa(k 1) such that training one-step GD on D(Be, 3 a) can recover the
normalized gradient of the one-step SAM update on the original data distribution (Bee, Ba, a). Using Equation 81 we have
3 (6)
Ua)
(3))
.aa²
(87)
32 (w (t)
(2(((1) ,v,))
ve)²
( .Va)
3/3/2 (w0)) Va)
(88)
(we .Ve)
(w((), (v)
(w(1), vu) - -3p(0) /3 (w (1), v2)²
1/22 (w011, Ua)
(89)
(w((), vve) 3p(t) 33 w (t), vee²
(w(()),))
1 (1) 89 (t) va)
3//2
(90)
3 (w ve) 1
(p 3 (w) V&) 3p10 (w)()
2/3
(91)
Therefore, with 8y= 1-3ppt) (3 (9 ma)
2/3
Sd. the normalized gradient of the one-step GD update on D(e,, 82 a:) is similar to that of the one-step SAM update on D(5,, a).
113p(t)) 8 (w(t)) v.))
B Auxiliary Lemmas
Lemma B.1 (Claim D.20, (2). Considering an increasing sequence it > 0 defined as Z1+1 =
I,+ (2, 2)² for some C. = @(1), S p I Fo. then we have every A > Ty. every ó E (0, 1). and every 1) E (0,1):
1+8 8
O(n(A 2) log(1 1 + +8)
log( A/ A/zo)
Σ
C, <
(92)
To
2
S20.
1
(1-5)m0 A
O(T(A- p)²)
log(A (A/zo) log (1 + 8)
Σ
(93)
20
(1
8)
x/
120.5A
26

---

Proof. For every 00...... let Tg be the first iteration such that 2 (1 + 80 Let b be the smallest integer such that (1 + x0 2 A. Suppose for notation simplicity that we replace 2, with exactly A whenever 2t 2 A. By the definition of Tg. we have
Σ
Ce + x) S 1T TT, S 8(1 + 8) + +O((AA - -))²)
(94)
TT)))
Σ
ni[(1 + 1)0++)=000 =ol2 21212
TTa 2 8(1 + 8) O(n(A p))))
(95)
(96)
These imply that
8
O p))
Σ AE(T))To-1)
n S
(97)
(1 55%0
z/
O - ))
Σ
C 2
(98)
(1 8)0+2.0
30
RE 1)
Recall b is the smallest integer such that (1 + 8)020=
x0
2 A, so we can calculate
b-1
5
O(1(A - p)3)
Σ
n
S
+
b
(99)
(1 + 8)920
2
120,2,5A
9=0
1
O(n(A p)?)
+
(100)
145 40
1
O(n(A )) (o/o0)
(101)
30
z
log(1 +6)
8
Onn(A -))²
Σ 120,2
1
(102)
(1 + 8) g
z
0
8(1 10-1( 1(1
1 Onn(A- -)
(1+6)0-1)
(103)
I+5
200
z
(1+6)
Onn(A 2)
log( (A/zo) log(1 [1 8)
(104)
[1
Thus, the two desired inequalities proved.
Lemma B.2 (Lemma D.19, (2).. 2. Let {xt, V= be two positive sequences that satisfy
- Jz+1 x2 + -,(x, p)?,
Vt+1 S y1 Sq. C.y?:
for some Ct = (1). Suppose To 2 yo 1-3G where S E (0, 10,1),C G E (0,1/3) and 0 < S
min log(A/2o)) log(I/C] },0 S p < O(zo), and all A E (æo, O()), let T2 be the first iteration such that x4 2 A. Then, we have yr, & 1
G2an G² no
Proof. Let Te be the first iteration in which x, 2 A. Apply Lemma B.I for the x, sequence with C. = C2 and threshold A, we have
Tc
1 1+8
O(n(A - ))2) (og(A/z))
Σ 0
C <
(105)
30
z/
log(1 (1 + 5)
1 8 8
11(A 0)2 log( (A/zo)
0
(106)
3o
d.r 11 1/log( / A/zo)
1+0 8
0
(107)
PP
27

---

Let Ty be the first iteration in which 2 A. Apply Lemma B.1 for the sequence with
St,, C2 = C₁ p = 0 and threshold A' = G yo. we have
Tv Σ
1 - (1+
O(Sn(A))
log (A'/vo)
SnC,
(108)
yo
yo
log(1 + +6)
1 0 + G) -0
Sn(AA) log (1/G)
(109)
yo
1 O + G) 0
SSn (og((1/)))
(110)
y/o
Compare Equation 107 and 110
Choosing 8 and r) S minf log(A/za) log(1/G)
G
Go
}. together with
3o 2 you 1-3G
L42G have T2 S Ty.
Lemma B.3 (Lemma K.15, 32 ). Let {zilrmo be a positive sequence defined by the following recursions
2++1 2 + m(= - p)
2+41 S 2t + (zx)
where > p Z 0 is the initialization and m, M > 0 are some constants. Let v >> 20, then the time
Ty such that 21, 2 v for all 2 Tg is
220
4M log(v/zo)
T.
(111)
m(=0 p)2
m(=0 p)2 log(2)
Proof. Let E N Let Ta be the first time that 21 2 2" We want to find an upper bound of T. We start with the case n = 1. By summing the recursion, we have:
Ti +
27, 2 20
(zt p)2
(112)
1=0
Because 2, 2 20, we obtain
Ti < 2T: 20
(113)
m(2u p)2
Now, we want to bound 27, zg- Using again the recursion and 27, S 2zo, we have
Zr < z11-1 + (2n < 250 + =.
(114)
Combining Equation I13 and
we get a bound on Ti as
20 + 4M
20
4M 2
Ti <
(115)
m(z0 p)2
m(z0 p)²
m(20 - p)²
Now, let's find a bound for Tn Starting from the recursion and using the fact that 2 2n t T, Th-1 have
20 for
T,
2T, 2 2T, + m
(z - p)2
(116)
t=
2T.-1 + (22-1) 1 p)2 Tn-1)
(117)
On the other hand, by using 2r, -1 < 2 20 upper bound 2r,
2T. < 2T, -1 M (2T, 2" 2n 20
(118)
Besides, we know that PT, > 2n-1z0 Therefore, we upper bound žτ, - 2T. as
2T. 2T.-1 < 22112 + 22
20
(119)
28

---

Combining Equations 166 and 119 yields
2n- + M2 z/
Ta T
(120)
2m--122(222 -p)
20
2
Tn-1
+
(121)
2 3n-1m(Zo 20 p)²
m(20
p)²
Summing Equation 120 for
...... we have
20
4M p)²
2z0
4M
Ta m 20 p)2
+
(122)
m(=0
m(z0
p)²
m(00 p)²
Lastly, we know that 20 2 this implies that we can set log(e/20) [2g
in Equation 122
We make the following assumptions for every S as the same
Lemma B.4 (Induction hypothesis D.1. (32). Throughout the training process using GD for t $ T,
maintain that, for every and j E [J],
(w (1) ..)) S (
(123)
For every i, we have 1 ((1)g (t), where
Lemma B.5 (Lemma G.4, [8
91 (t) sigmoid Σ
(W
Vai 82 (1) (w)) (0 .
(124)
Lemma B.6 (Lemma K.5, 32
E R2d be a Gaussian random vector, ~ NO, La). Then with probability at least o(1). we have X = O
Lemma B.7 (Lemma K.7, 332 Let X and Y be independent Gaussian random vectors Ied and
X NO, Ia). Y (N o Ia). Assume that oog I Then, with probability least
have
2
(X,Y)) S / 2d
Lemma B.8 (Bound on noise inner products). Let N O(poly(d)). The following hold with probability at least o(1):
max (1w je: E.)) = Ō(oovvd)
(Ex, Eiil = =Ō o6d N + =2vdd)
max
Proof. For the first inequality, Lemma B.7 implies that with probability at least 1 - 1
2N.
(w (0) je Ex) S o 2d log(
2
Ō(ooooa
(125)
dN
Taking union bound over = 1....,
N gives the result.
The second statement is proved similarly.
Lemma B.9 (Bound on the noise component for SAM). Assume that o(on) and w(1) S N < (poly(d)). Throughout the truining process using SAM for S T, we maintain that, for every and j [J,
(0 (.. E.)l S (lorood)
(126)
29

---

(4) Proof. Let Xt = max m ( j.c. }. max, + (.)))
Combined with 12 S 1, the
noise gradient update rule can be bounded as
X++1 $ Xt + 3Tax?
Suppose that alI) satisfies the differential equation
a' 3arja² ?
a(0) = Xo
Observe that a(I) is increasing so, by the Mean Value Theorem there exists T E (1, L + 1) such that
a(t + 1) at) dr)
3ama (T)2
2 3oma
So an fast-learnable induction shows that a()) 2 Xe.
Now solving for a(t).
(1)
1 S 30ma(0)
3amt
a(O]
Using the high probability tail bounds B.8
(0) = Xu = V((ooVV))
a = Ō o²d
N + +zVVd)
where
Substituting these bounds gives vd
Sp.
(tt)
C2( ( Ö(nt( a²d N + o2vd))
Now Theorem
A.15 and p (()) implies that 1/T ô ) Combined with the assumption
that N = (1),, the second term in the denominator is of lower order than the first term, so
a(Tb) Vdooo
We conclude that
( ; E < Ö(oor,Va). Kw
(127)
Lemma B.10. For every i, we have ( = (1))) and 1 = ((1)g2(t) where
(1 sigmoid
(W (1) vu)² (w u2)s
(128)
g2 sigmoid
38 (w (t) j.: va)
2w (
je v3)3
(129)
The proof is the same as 18 Lemma G.4].
C Additional Experimental Settings
C.1 Datasets and Training Details
Datasets. The CIFAR10 dataset consists of 60,000 32 x 32 color images in classes, with
6000 images per class. The CIFAR 100 dataset 41 is just like the CIFARIO, except has 100 classes
30

---

Original Dataset S
USEFUL
Class-Wiss
Selected Samples S'
Y/1
y/2
Bi
Broz
g/+
Model
Model
11
IF
2'29
aF
Figure 7: USEFUL first trains the model for a few epochs t, which in practice is around 10% of the total training epochs. It then clusters examples in every class into 2 groups and upsamples the cluster with higher average loss. Finally, the base model is retrained from scratch on the modified data distribution.
containing 600 images each. For both of these datasets, the training set has 50,000 images (5,000 per class for CIFAR1O and 500 per class for CIFAR 100) with the test set having 10,000 images. CINIC10 [16] represents an image classification dataset consisting of 270,000 images, which is 4.5 times larger than CIFARIO. The dataset is created by merging CIFARIO 10 with images extracted from the ImageNet database, specifically selecting and downsampling images from the same 10 classes present in CIFARIO. Timy-ImageNet 43 comprises 100,000 images distributed across 200 classes of ImageNet [77 with each class containing 500 images. These images have been resized to 64x64 dimensions and are in color. The dataset consists of 500 training images, 50 validation images, and 50 test images per class, The STL10 dataset [13 includes 5000 96x96 training labeled images, 500 per CIFARIO class. The test set consists of 800 images per class, this counts up to 8,000 images in total.
Training on different datasets. Follow the setting from 3 we trained Pre-Activation ResNet18 on all datasets except for CIFAR 100 which was trained with ResNet34. We trained our models for 200 epochs with batch size of 128 and used basic data augmentations such random mirroring and random crop. We used SGD with the momentum parameter of 0.9 and set weight decay to 0.0005, We also fixed p 0.1 for SAM unless further specified. For all datasets, where we set the initial learning rate to 0. The learning rate is decayed by a factor of 10 after 50% and 75% epochs, i.e., we set the learning rate to 0,01 after 100 epochs and to 0.001 after 150 epochs.
used learning rate schedule
Training with different architectures. We used the same training procedures for Pre-Activation ResNet18, VGG19, and 1enseNet121. We directly used the official Pytorch (54) implementation for VGG19 and DenseNet121. For 3-layer MLPs, we used a hidden size of 512 with a dropout of 0.1 to avoid overfitting and p 0.01. For ViT-S (77), adopted Pytorch implementation https://github.com/ cww/ http///ddraiiiiiiiyooooo
In particular, the hidden size, the depth, the number of attention heads, and the MLP size are set to 768, 8, and 2304, respectively. We adjusted the patch size to 4 to fit the resolution of CIFAR 10 and set both the initial learning rate and to 0.01.
Computational resources. Each model is trained on NVIDIA RTX A5000 GPU,
C.2 Other Implementation Details
When to separate the examples? We selected the best-separating epoch t in the set of (4,5, 5,6, 6, 7, 8, 10) for CIFAR1O and {12, 14, 16, 18, 20, 22} for CIFAR 100. Particularly, we sepa- rated examples of CIFAR10 at around epoch 8 while that of CIFAR 100 is near epoch 20. Near this point, the gain in training error diminishes significantly as shown in Figure 15a] which shows sign that the model successfully learns fast-learnable features. In addition, we reported the results for differ- ent separating epochs in Appendix D In Figure 4
the best separating epochs for STL10, CINIC1O, and Tiny-ImageNet are 11, 4, and I0, respectively. The separating epoch for Waterbirds is 5.
Forgetting score. To compute forgetting scores of training examples in each dataset, we collected the same statistics as in [69 but computed at the end of each epoch. The reason is to make the statistics consistent between two versions of the same slow-leamable example which is repeated in the upsampled dataset.
Hessian spectra. We approximated the density of the Hessian spectrum using the Lanczos algo- rithm 24 53 The Hessian matrix is approximated by 1000 examples (100 per class of CIFAR10), Then we extract the top eigenvalues to calculate the maximum Hessian eigenvalue (Amax) and the bulk of spectra (Amax/ />s) [30].
31

---

0.5
G
0.4
5
0.3
0.2
0.1
0.0
100
200
400
500
600
Iteration
Figure 8: The gap between contribution of fast-learnable and slow-learnalle features towards the model output in SAM and GD. The toy datasets is generated from the distribution in Definition 3.1
with 3a = Be = a 1.
MLP
43
42,59
SGD
SGD+USEFUL
SAM
SAM+USEFUL
40
SAM+TA
39
18-76
SAM+USEFUL+TA
Time Train wrt.. SGD
Figure 9: Test classification errors of 3-layer MIP on CIFAR10. The number below each bar indicates the estimated cost to train the model and the tick on top shows the standard deviation over three runs. USEFUL improves the performance of SGD and SAM when training with 3-layer MLP.
CIFAR10
CIFAR100
SGD
Norm
109010
SGD +SEEULL
37000
SAM
36000
102022
20000
303020
23020
Figure 10: L1 norm of ResNet18 trained on CIFAR10 and ResNet34 trained CIFAR100. Lower LI norm indicates a sparser solution and stronger implicit regularization properties B Section 4.2]. SAM has lower L.1 norm than SGD, and USEFUL further reduces the L. norm of SGD and SAM.
D Additional Results
D.I An Extreme Case of Toy Datasets
We consider an extreme setting when two features have identical strength and no missing fast- learnable features, i.e., 3. Ba 1, the gap between LHS and RHS in Equation i not small as shown figure The gap is consistently around 0.2-0.3 from epoch 250 onwards.
D.2 USEFUL is Useful for MLP
Figure 9 shows that for 3-layer MI.P, USEFUL successfully reduces the test error of both SGD and SAM by nearly 1%. Additionally, SGD+USEFUL yields better performance than SAM alone.
32

---

Table Sharpness of solution convergence. We train R on CIFARIO and measure the maximum Hessian eigenvalue Amax and the bulk spectra measured as Amax/ A55-
METRIC
SGD
SGD+USEFUL
SAM
Amax
53.8
41.8
12.4
Amax A5
3.8
1.5
2.4
Forget Score [0-19)
Forget Score [20-39]
Number of Examples les
6117
SGD
42000
41392
6000
+USEFUL
5325
40000
39675
5000
38883
3606
36000
3000
Figure 11: Forgetting scores for training ResNet18 on CIFARIO. Forgetting scores measure the learning speed of examples training data. USEFUL approaches the training dynamics of SAM, with more examples being forgotten infrequently and fewer examples being forgotten frequently.
Table 2: Average scores for two clusters on CIFAR1O.
METRIC
FAST-LEARNABLE
SLOW-LEARNABLE
FORGETTING SCORE
3.8 + 6.1
14.7 9.0
FIRST LEARNED ITERATION
0.9 1.2
3.7 8.0
ITERATION LEARNED
45.6 50.0
105.8 41.0
Table 3: Average scores for two clusters on CIFAR
METRIC
FAST-LEARNABLE
SLOWWLEERNABBEF
FORGETTING SCORE
10.2 8.6
16.8 7.4
FIRST LEARNED ITERATION
4.7 6.8
9.8 + 13.7
ITERATION LEARNED
86.6 + 46.1
115.7 + 31.0
D.3 SAM & USEFUL Reduce Forgetting Scores
We used forgetting scores 69 to partition examples in CIFARIO into different groups. Forget- ting scores count the number of times an example misclassified after being correctly classi- fied during training and is an indicator of the learning-speed of examples. Figure TT illustrates that SGD+USEFUL and SAM have fewer examples with high forgetting scores than SGD does. This aligns with our theoretical analysis in Theorem 3.4 and results on the toy datasets. By upsampling slow-learnable examples in the dataset, they contribute more to learning and hence SGD+USEFUL learns slow-learnable features faster than SGD.
Furthermore, Tables and B jilustrate the average forgetting score, first learned iteration (i.e., at this epoch, the model predicts correctly for the first time) and iteration learned (i.e., after this epoch, prediction is always correct) of examples in fast-learnable and slow-learnable clusters. Iteration learned is highly correlated with prediction depth S. which is another notion of data difficulty. It can be seen clearly that examples fast-learnable clusters have a lower difficulty score for every metric, indicating that USEFUL successfully identifies last-lea examples early in training.
33

---

Table 4: Test classification errors for training SAM and ASAM on the original CIFARIO and modified datasets by USEFUL. Results are averaged over 3 seeds.
SAM
ASAM
4.23 + 0.08
4.33 0.19
USEFUL
4.04 + 0.06
4.09 + 0.10
TA
4.06 + 0.08
3.93
0.11
+ TA + USEFUL
3.49 + 0.09
3.46 0.01
Anneunn pre-trained
23p2
SGD
SGD+USEFU +USEFUL
SAM
SAM+
2017
SAM+ TA
2412
73
4/s
SAM USEFUL+TA
22445
1x
2,30
2x 2.3x
Time Taain W.n.l SCD
Figure 12: Comparing test classification errors on Waterbirds. The number below each bar indicates the approximated cost to train the model and the tick on top shows the standard deviation over three runs. USEFUL boosts the performance of SGD and SAM the balanced test set, showing its generalization to the OOD setting. In addition, the success of USEFUL in fine-tuning reveals its new application to the transfer learning setting.
Table 5: Comparison between USEFUL and upweighting loss regarding test classification errors. Upweighting loss doubled the loss for all examples in the slow-learnable clusters found by USEFUL, which different from dynamically upweighting examples during the training [78). Results are averaged over 3 seeds.
METHOD
SGD
SAM
CIFARIO
CIFARI00
CIFAR10
CIFAR100
UPWEIGHTING LOSS
5.33 = 0.09
26.70 + 3.25
4.28 + 0.02
21.75 + 0.25
USEFUL
4.79 = 0.05
22.58 + 0.08
4.04 0.06
20.40 + 0.36
D.4 USEFUL generalizes to other SAM variants
In this experiment, we show that SAM can also generalize to other variants of SAM. We chose
ASAM, which is proposed by Kwon et al. to address the sensitivity of parameter re-scaling [1]
Following the recommended settings in ASAM, we trained with a perturbation radius which is 10 times that of SAM. Other settings are identical the standard settings in Appendix C
1.0.
Table demonnstatss the results for training ResNet18 on CIFARIO I0. Both SAM and ASAM can be combined with USEFULto improve the test classification error. When using TA, ASAM shows slightly better performance than SAM.
D.5 USEFUL shows promising results for the OOD settings
While our main contribution is providing novel and effective method to improve the in-distribution generalization performance, we conducted new experiments confirming the benefits our method to improving diit----iitiittioon (OOD) generalization performance. As few very recent works 72 72 showed the benefit of SAM in improving OOD performance, it is expected that USEFUL also extend this setting.
Spurious correlation (Waterbirds). As can seen in Figure 12 both SAM and USEFUL effeetivlye improve the performance the balanced test set though the model is trained on the spurious training
34

---

Table Test error on long-tailed CIFARIO. BALANCING means that we upsampled small classes to make the classes balanced. Results are averaged over seeds.
RATIO
SGD
SGD+USEFUL
SAM
SAM+USEFUL
1:10
10.01 + 0.21
9.53
0.13
8.85 + 0.08
8.22 1 0.04
BALANCING
9.77 + 0.17
9.25 + 0.11
8.31 + 0.11
7.93 + 0.02
Table 7: Test classification errors of SGD for different partition methods. Results are averaged over
3 seeds.
PARTITION METHOD
CIFARI0
CIFARIO0
QUANTILE
5.27 0.10
23.49 0.82
MISCLASSIFICATION
4.98 0.17
23.86 + 0.70
USEFUL
4.79 + 0.05
22.58 + 0.08
set. When training ResNet18 with SGD from scratch, USEFUL decreases the classification errors by 5.8%. In addition, it can successfully applied to the pre-trained ResNet50 (on ImageNet), opening up promising application to the transfer learning setting.
Long-tail distribution (Long-tail CIFAR10). We conducted new experiments on long-tail CI- FARIO [S] with an imbalance ratio of 10. Table 6 shows that USEFUL can also improve the performance SGD and SAM, by reducing the simplicity bias on the long-tail data. Figure [3 visualizes the distribution of classes before and after upsampling by USEFUL. Interestingly, we see that USEFUL upsamples more examples from some the larger classes and still improves the accuracy 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 performance, as we confirmed in the last row.
Long-tailed CIFAR10 dataset
Original Upsample
6000
Count
4000
2000
0
0.0
2.5
5.0
7.5
Class
Figure 13: Class count before and after upsampling by USEFUL on long-tail CIFAR10 10 dataset.
D.6 USEFUL is also effective for noisy label data
Our method and analysis consider a clean dataset. But, as we confirmed in Table 8 USEFUL can easily stack with MixUp 79], a robust method for learning against noisy labels, to reduce the simplicity bias and improve their performance. Applying USEFUL on top of MixUp to CIFARI0 with 1 10-20% random label flip successfully boosts the performance both SGD and SAM when training on data with corrupt labels.
35

---

Table 8: Test error on label noise CIFARIO (all methods are with MixUp). Results are averaged over 3 seeds.
RATE
SGD
SGD+USEFUL
SAM
SAM+USEFUL
10%
7.20 + 0.17
6.64 + 0.10
5.15 + 0.05
4.75 0.09
20%
9.26 + 0.23
8.88 + 0.07
6.08 + 0.06
5.82 + 0.05
Table 9: Test errors of different simplicity bias reduction methods on CIFARIO. Results are averaged over seeds.
SGD
EIIL
JTT
SGD+USEFUL
5.07 + 0.04 5.04 + 0.04 4.89 + 0.03
4.79 + 0.05
7.0
6.5
To5t Ever
6.0
5.0
5.5
4.9
SGD U
5.0
4.8
128
256
512
0.1
0.2
0.4
Batch aize
Leariin rate
Upsamoling factor
Figure 14: Ablation studies of training ResNet18 on CIFAR1O. In each experiment, we used the standard training settings while (left) varying training batch size or (middle) varying learning rate, (right) varying upsampling factor.
24.5
Epech B
Epoch 20
Tran Ei1o
15
24.0
23.0
25
SCD-USETUL
SCD
25 50 75 700-12119017359
175200
25 50
75 101225100775200
7
9
10
14
16
20
22
Epoch
Ebocc
Epoch
Epoch
(a) Training error trajectories
(b) Test classification
Figure 15: Separating epoch analysis. (left) Red lines indicate our optimal choice of t to separate examples at and restart training. The early epoch that can best separate the examples is when the change in training error starts to shrink. (right) Red points indicate our optimal choice of t.
D.7 Comparison with simplicity bias mitigation methods
While previous works show that reducing simplicity bias benefits the OOD settings, we show that reducing the simplicity bias also benefits the ID settings. To confirm our hypothesis on our simplicity bias mitigation baselines, we applied EIIL [14] and JTT [48] to train ResNet18 on CIFARIO. The choice of the baselines is because they have publicly available code and fewer hyperparameters to tune in our limited rebuttal time. For EIIL, we tuned Ir (lee-, le-2, Ie-3, 5e-4, lc-4), number of epochs 1, 2, 4, 8) for training the reference model, and the weight decay E 1e-3, 5e-4, le-4) for training GroupDRO. For JTT, we tuned the separating epoch E44, 6, 7. 8, 10) and upsampling factor (2, 3, 4, 5], while Ir and weight decay follow the standard training of ResNet18 on CIFAR10. Table shows that all methods successfully reduce the simplicity bias, yielding an improvement over SGD. While EIIL requires tuning 3 hyperparameters with total of 60 combinations, USEFUL only requires one hyperparameter, which the separating epoch within a small range (around the time when the slope of the training loss curve diminishes).
36

---

D.8 Ablation studies
USEFUL vs Upweighting loss. We compare USEFUL with upweighting loss of examples in the slow-learnable clusters. As can be seen in Table when coupling with either SGD or SAM, USEFUL clearly outperforms upweighting loss on both CIFAR datasets. is worth mentioning that upweighting loss is different from iteratively importance sampling methods such as GRW (78). which dynamically upweights examples during the training by factors that depend on the loss value. In addition, GRW is dedicated to the distribution shift setting while our paper considers the in-distribution setting.
Data selection method. In this experiment, we compare clustering with other methods for partitioning data. The first baseline to upsample misclassified examples (MISCLASSITICATION) while the second baseline is to upsample all examples whose training errors are larger than the median value (QUANTILE). All the three methods are performed the same epoch Table 7 shows that USEFUL selects better set of upsampling examples. leading to the best model performance.
Training batch size. Figure 14 left shows the gap between USEFULand SGD when changing the batch size. Our method consistently improves the performance, proving its effectiveness is not simply captured by the gradient variance caused by the training batch size.
Learning rate. The small learning rate a nuance in our theoretical results to guarantee that fast-learnable and slow-learnable features are separable in early training. In general, a small learning rate is required for most theoretical results on gradient descent and its convergence and a standard theoretical assumption 12
60 62 73 In practice, both for separating fast-learnable vs slow- learnable examples and for training on the upsampled data, we used the standard learning rate that results in the best generalization performance for both SGD and SAM following prior work [3 42 81..
While the theoretical requirement for the learning rate is always smaller than the one that is used in practice, empirically larger learning rate does not yield better generalization for the problems we considered in contrast to other settings [46] (88.. As shown in Figure [4] middle, increasing the learning rate has an adverse effect on the model performance. Indeed, for fair comparison, the algorithms should be trained with hyperparameters that empirically yield the best performance; otherwise, the conclusions are not valid. Nevertheless, USEFUL always improves the test error across different learning rates.
We empirically found the upsampling factor of 2 to consistently work well across different datasets and architectures. Using larger upsampling factors results in a too-large discrepancy between the training and test distribution and does not work better, as is expected and discussed in Section [ As illustrated in Figure 14] right, while all factors from to 5 bring performance improvement, the upsampling factor of yields the best performance, it reduces the simplicity bias with minimum change to the rest of the data distribution.
Upsampling factor.
Separating epoch. Fig 15a shows that the early epoch that can best separate the examples is when the change training error starts to shrink. At this time, examples that are learned can be separated from the rest by clustering model outputs analyzed in theoretical results in Section 3.3 Figure 15b demonstrates the performance of USEFUL when separating examples at different epochs in early training. Too early too late epochs do not cluster examples well, i.c., some examples with fast- learnable features fall into the slow-learnable clusters and vice versa. This ablation study shows that upsampling correct examples and with enough amount is important for method to achieve its best. Note that there is no universal separating epoch. This is reasonable because each data has a different data distribution, i.e. slow-learnable and fast- leatnale features, thus, a different theoretical time T
in Theorems 3.2 and 3.3
37

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
Our claims match the theoretical and experimental results
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We do not have theoretical results for non-CNN architectures.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
38

---

Justification:
We provide full assumptions and proofs in Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide detailed experimental settings in the Experiment section and in Appendix.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
39

---

Answer: [Yes]
Justification: We use open datasets, which have been cited and described in the paper. We also provide our code for reproducing the experimental results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results. See the NeurlPS code and data submission guidelines (https
ssiinlolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide detailed experimental settings in the Experiment section and in Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification:
We report the error bar for averaging multiple seeds in our experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the eтror bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
40

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπхо rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide the information on our computational resources in Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: Our paper conform to every aspect in the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification: This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of work, none which we feel must be specifically highlighted here.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
41

---

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
There are no potential harms to our models.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We cite all data, code, and models used properly.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
42

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do not release assets.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
43


Reviews:
[{"review_id": "s4pdozXMf4", "reviewer": "Reviewer_Y69v", "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.", "questions": "See Weakness.", "limitations": "N.A.", "rating": 4, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written and easy to follow.\n- The paper has a theoretical analysis to analyze the learning progress and derive the proposed method.\n- 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.\n- 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.\n- 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. \nSome small questions:\n- What's the exact formulation of the Data distribution?\n- 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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedback and acknowledging our theoretical results and comprehensive experiments. We discuss the questions below.\n\n1. **USEFUL vs resampling methods for long-tail data & example difficulty.** \n- As discussed in the [general comment](https://openreview.net/forum?id=yySpldUsU2&noteId=nTUcV9AnO0), 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.\n- We conducted new experiments on long-tail CIFAR10 with an imbalance ratio of 10. Table 3 in the [PDF](https://openreview.net/forum?id=yySpldUsU2&noteId=nTUcV9AnO0) 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](https://openreview.net/forum?id=yySpldUsU2&noteId=nTUcV9AnO0) 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](https://openreview.net/forum?id=yySpldUsU2&noteId=3xcoRIZ69p)).\n    \n2. USEFUL’s solution is not as flat as SAM.\n- **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.**\n- 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.\n3. Formulation of the **Data distribution**.\n- 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].\n4. **\"patch\" meaning in Definition 3.1**.\n- It is similar to the patch in ViT. Each patch is a small region in the image and each image consists of P patches.\n5. More experiments on popular models and big datasets, e.g., **Transformer** **ImageNet-1k**\n- 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].", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Qmi8KMhqJg", "reviewer": "Reviewer_eKCz", "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\n- 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.", "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?\n- 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.", "limitations": "Please see strengths and weaknesses.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Simple easy-to-implement method that uses SAM and upsampling to improve in-distribution generalization\n- 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:\n    - https://arxiv.org/abs/2105.05612\n    - https://arxiv.org/abs/2301.13293\n    - https://arxiv.org/abs/2107.09044 (does not focus on simplicity bias explicitly, but similar to method proposed in paper)\n    - 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?\n- 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.\n- Lacking discussion about when this method would fail. I can imagine two scenarios where the method would not work:\n    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? \n    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? \n- “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.\n- 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.\n- 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..)\n- Writing is repetitive at times, especially the theory section (3.3)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your feedback! \n1. Comparison with simplicity bias baselines\n- 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](https://openreview.net/forum?id=yySpldUsU2&noteId=nTUcV9AnO0) for detailed discussion).\n- 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**).\n2. Example difficulty metrics vs USEFUL clusters & training on the k examples with highest difficulty scores per class\n- 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.\n3. Novelty w.r.t [64]\n- As discussed in our [general comment](https://openreview.net/forum?id=yySpldUsU2&noteId=nTUcV9AnO0), **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**.\n- 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**. \n4. Would USEFUL fail if: \n- Most training examples have one or more slow-learnable features\n  - 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.\n  - 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.\n- 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. \n5. Comparing {SGD, SAM, SAM+USEFUL} w/ and w/o TA\n- 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. \n6. Larger datasets\n- Please see Q5 of [Reviewer Y69v](https://openreview.net/forum?id=yySpldUsU2&noteId=xu392aAv28). We do not believe there is anything specific to ImageNet that does not work with our method.\n\n**Questions**:\n1. If #slow-learnable features >> #number of fast-learnable features\n- 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.\n2. Why is max alignment between weight vector and ground-truth feature v_e the right way to evaluate the feature’s contribution?\n- 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.\n3. 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?\n- 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nlj9GM1obg", "reviewer": "Reviewer_EFb1", "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).", "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?\n2. 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,\n\nPlease refer to the Weakness section for remaining questions.", "limitations": "Yes, the authors have addressed limitations.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "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.\n2. 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.\n3. 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,\n    - 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. \n    - 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.\n    - 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.\n2. 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.\n3. 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. \n    \n [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.\n\n[2] Shah, Harshay, et al. \"The pitfalls of simplicity bias in neural networks.\" Advances in Neural Information Processing Systems 33 (2020): 9573-9585.\n    \n [3] Geirhos, Robert, et al. \"Shortcut learning in deep neural networks.\" Nature Machine Intelligence 2.11 (2020): 665-673.\n    \n [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).\n    \n [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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback and recognizing the originality and novelty of our work, and our comprehensive experiments.\n\n1. Our work (ID) vs [1-5] (OOD).\n- As the reviewers correctly mentioned and we discussed in our [general comment](https://openreview.net/forum?id=yySpldUsU2&noteId=nTUcV9AnO0), 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).\n2. How is the separating epoch chosen?\n- 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.\n- Is there a separating epoch number that works across various datasets?\n    - 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.\n- The average gain on most of the datasets with USEFUL is less than 1% with additional cost for training.\n    - 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.\n\n**Questions**:\n\n1. **Separating epoch** used for all datasets\n- 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.\n2. What is the takeaway from Figure 1? What kind of examples are usually clustered in fast learnable cluster vs slow learnable cluster?\n- 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gO04npIXWy", "reviewer": "Reviewer_F4yx", "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.", "questions": "1) Do you think your method will also generalize on Long-Tail dataset while keeping an upscaling factor constant? \n\n2)) Any ideas or heuristics about how to find the optimal epoch/layer to perform the clustering without running an expensive ablation?", "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.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedback, and acknowledging our well-motivated work and our comprehensive experiments and ablations.\n\n1. When and why one should choose the last output activation vector to define the clustering instead of intermediate activation vector?\n- 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.\n2. How to find the separating epoch?\n- 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.\n3. How robust is the scaling factor of 2 on long-tail distribution?\n- 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](https://openreview.net/forum?id=yySpldUsU2&noteId=nTUcV9AnO0), 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.\n  \n  **Table 3**: Test error (avg of 3 runs) on long-tailed CIFAR10.\n\n  | Ratio | SGD | SGD + USEFUL | SAM | SAM + USEFUL |\n  | --- | --- | --- | --- | --- |\n  | 1:10 | 10.01 ± 0.21 | 9.53 ± 0.13 | 8.85 ± 0.08 | 8.22 ± 0.04 |\n  | Balancing | 9.77 ± 0.17 | 9.25 ± 0.11 | 8.31 ± 0.11 | 7.93 ± 0.02 |\n\n4. The authors claim that their method is generalizing to OOD tasks while providing experiments on only the WaterBird dataset.\n- 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their valuable feedback and recognizing the originality of our work. We’d like to first briefly emphasize the scope and contribution of our work:

- Our work shows, for the first time, that reducing the simplicity bias benefits **in-distribution (ID)**. Previously, the benefits of reducing simplicity bias have been only shown to **out-of-distribution (OOD)**, in particular to mitigate spurious correlations.
- **Prior methods: OOD setting** [long-tail data and spurious correlations]. Here, there is a shift between the training and test data distributions:
    - **Long-tail data** [Long-tailed CIFAR10, Inaturalist, ImageNetLT]: (sub)classes are highly imbalanced in training but are (more) balanced in test data (distribution shift). Upsampling the minority (sub)classes in training data improves the test performance on corresponding groups that are (much) larger in test data. Here, **the benefit comes from balancing the training data and matching the training and test data distributions**.
    - **Spurious features** [Waterbirds, CelebA, CMNIST, CIFAR-MNIST]: spurious features are simple features with a high correlation with a label at training time but not at test time (distribution shift). Due to simplicity bias, spurious features are learned *instead* of the more complex predictive features. This yields a poor (worst-group) accuracy on examples without the spurious feature at test time. Reducing the simplicity bias (or training with SAM) suppresses learning spurious features. Here, **the benefit comes from mitigating the spurious feature to learn more/other features**.
- **Our method: ID setting** [CIFAR10-100, TinyImageNet, CINIC10]**.** Here, training and test data have the same distribution, there is no spurious feature in the training data, and all features are predictive on the test set. Our work shows for the first time that reducing the simplicity bias benefits ID. Our contribution is orthogonal to the above work: **we proved that reducing the simplicity bias allows learning the *predictive* features *at a more uniform speed*, and showed that this benefits ID.** Unlike the OOD setting, *the benefit of reducing the simplicity bias in the ID setting is not attributed to balancing the training data distribution or suppressing the spurious features to learn other features*. It is attributed to the **(more uniform) speed/time of learning the same set of predictive features**. Our findings are more unexpected than its counterpart in the OOD setting and opens up a new direction for future research.
- **Our method: OOD setting** [Long-tailed CIFAR10, Waterbirds]. While our theory and main contribution is to show the benefits of reduced simplicity bias to ID, our new experiments show that our method can also benefit the long-tail setting, by reducing the simplicity bias and learning features of a class at a more uniform speed. Benefits of reducing the simplicity bias to the spurious setting has been studied by several prior work. While our experiments on Waterbirds shows that our method is also applicable to this setting, we note that **this is not our main contribution**.

We provide new experimental results on example difficulty, long-tail data, noisy labels, and spurious baselines in the attached **PDF**.

Meta-review: This paper proposes a theoretically inspired method for modifying the distribution of the training data to improve generalization performance. Rather than building on generalization bounds, which are often too loose in practice to be useful, the paper utilizes SAM, a recent training scheme that has been empirically shown to be very effective for improving generalization, to drive its theoretician findings. Specifically, by analyzing the training dynamics of SAM and comparing that against SGD, the paper proves that SAM tends to learn hard and easy features at a more even pace, compared to SGD which heavily relies on easy features for training. Based on this theory, the authors hypothesize that the success of SAM may stem from its even-paced learning of both types of features. Inspired by this, they propose an algorithm for reweighting samples during training. The idea is to first train a model for a few epochs to get a rough initial model, and then use that model to cluster the data from each class into two categories, where one of them contains only samples that do not have any easy features. Then, by upsampling data that do not have easy features, the rest of the training balances the pace of learning easy and hard features. The proposed method is empirically verified on a few small but standard benchmark datasets, and the results confirm that the proposed reweighting method can indeed improve generalization performance when compared against SGD.

The paper has received a borderline score. Reviewers find the paper well-written and the theoretical analysis of SAM/SGD training dynamics, which demonstrates their connection to the speed at which they learn easy/hard features, to be original and intriguing. The majority of reviewers are also satisfied with the level of empirical verification of the algorithm, though they would appreciate it if the empirical analysis had covered larger datasets.

Common concerns among reviewers were the lack of novelty of the reweighting algorithm itself, as well as missing literature associated with data reweighting. The authors clarified that most of this literature relates to out-of-distribution (OOD) scenarios, while the paper focuses on in-distribution generalization. Other questions/concerns were about how the proposed algorithm works in the presence of label noise, and the effect of the proposed algorithm on the sharpness of the solution. The authors provided a detailed response to these concerns.

Unfortunately, two reviewers who gave the lowest initial rating (borderline reject) to this submission did not engage in the rebuttal, which makes it difficult to determine the fate of the paper. However, after carefully reading the authors' responses to these reviewers, I believe the provided responses can clarify some of the raised concerns. Furthermore, in agreement with the reviewers, I find the theoretical connection between SAM/SGD dynamics and the speed of easy/hard feature learning very intriguing, and I also find the proposed algorithm based on this analysis very simple to implement. Considering all the reviews and responses, I believe the contributions made by this work and the clarity of its presentation warrant an accept.
Decision: Accept (poster)