Title: Changing the Training Data Distribution to Reduce Simplicity Bias Improves In-distribution Generalization
Venue: NeurIPS 2024
Accepted: yes
Keywords: In-distribution generalization

Simplicity bias

Data modification

Sharpness-aware minimization
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

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 containing 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.

1 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 [47, 56, 82] or 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 can 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, or harmful examples [23, 45, 66]. Nevertheless, it remains an open question if one can change the distribution of a clean training data to further improve the in-distribution generalization performance of models trained on it.

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 [29]. Under this assumption, minimizing the training loss generalizes well on the test data [7]. Nevertheless, for overparameterized neural networks with more parameters than training data, there are many zero 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 to drive the optimization algorithms towards finding more generalizable solutions on the original data distribution.

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 sharpnessaware-minimization (SAM) [22] over (stochastic) gradient descent (GD). SAM finds flatter local minima by simultaneously minimizing the loss value and loss sharpness. In doing so, it outperforms (S)GD and obtains state-of-the-art 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 it with that of GD. We rigorously prove that SAM learns different features in a more uniform speed compared to GD, particularly early in training. In 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-norm 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 a 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 generalization 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 indistribution 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 CIFAR10, STL10, CINIC10, Tiny-ImageNet; ResNet34 on CIFAR100; and VGG19 and DenseNet121 on CIFAR10. We also empirically confirm the benefits of USEFUL to out-of-distribution performance, but we emphasize that this is not the focus of our work.

2 Related Works

Sharpness-aware-minimization (SAM). Motivated by the generalization advantages of flat local minima, sharpness-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 state-of-the-art 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, 72].

There have been recent efforts to understand the generalization benefits of SAM. The most popular explanation 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 overfitting under much weaker signal strength compared to (S)GD [12]. More recently, SAM is shown to also benefit out-of-distribution (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 a 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 [34] 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, 63]. 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, 75]. 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 a 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 D.7, we empirically confirm that existing simplicity bias mitigation methods also improve the in-distribution performance, but to a 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-dependent (non-generalizable) features or strong spurious correlations in the training data, or shift between training and test distribution. We focus on in-distribution 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. Long-tailed 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 [15, 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 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 [1, 44, 59], 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. 3.1. 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 a vector v , we use ∥ v ∥ 2 to denote its Euclidean norm. Given two sequence { x n } and { y n } , we denote x n = O ( y n ) if | x n | ≤ C 1 | y n | for some absolute positive constant C 1 , x n = Ω( y n ) if | x n | ≥ C 2 | y n | for some absolute positive constant C 2 , and x n = Θ( y n ) if C 3 | y n | ≤ | x n | ≤ C 4 | y n | for some absolute constant C 3 , C 4 > 0 . Besides, we use ˜ O ( · ) , ˜ Ω( · ) , and ˜ Θ( · ) to hide logarithmic factors in these notations. Furthermore, we denote x n = poly ( y n ) if x n = O ( y D n ) for some positive constant D, and x n = polylog ( y n ) if x n = poly (log( y n )) .

Data distribution. We use a popular data distribution used in recent works on feature learning [2, 8, 11, 12, 18, 32, 40] to represent data as a combination of two features and noise patches. Additionally, we introduce a probability α to control the frequency of fast-learnable features in the data distribution.

Definition 3.1 (Data distribution) . A data point ( x, y ) ∈ ( R d ) P × {± 1 } is generated from the distribution D ( β e , β d , α ) as follows. We uniformly generate the label y ∈ {± 1 } . We generate x as a collection of P patches: x = ( x (1) , x (2) , . . . , x ( P ) ) ∈ ( R d ) P , where

Slow-learnable Feature. One and only one patch is given by β d · y · v d with ∥ v d ∥ 2 = 1 , ⟨ v e , v d ⟩ = 0 , and 0 ≤ β d < β e ∈ R .

Fast-learnable feature. One and only one patch is given by β e · y · v e with ∥ v e ∥ 2 = 1 with a probability α ≤ 1 . With a probability of 1 -α , this patch is masked, i.e. 0.

Random noise. The rest of P -2 patches are Gaussian noise ξ that are independently drawn from N (0 , ( σ 2 p /d ) · I d ) with σ p as an absolute constant.

For simplicity, we assume P = 3 , and the noisy patch together with two features form an orthogonal set. Coefficients β e and β d 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 state-of-the-art architectures, we analyze a two-layer nonlinear CNN which is also used in [8, 11, 18, 32, 40]. Unlike linear models, CNN can handle a data distribution that does not require a fixed position of patches as defined above. Formally,

where w j ∈ R d is the weight vector of the j -th filter, J is the number of filters (neurons) of the network, and σ ( z ) = z 3 is the activation function, i.e., the main source of non-linearity. W = [ w 1 , . . . , w J ] ∈ R d × J is the weight matrix of the CNN. Following [8, 18, 32], we assume a mild overparameterization with J = polylog ( d ) . We initialize W (0) ∼ N (0 , σ 2 0 ) , where σ 2 0 = polylog ( d ) /d .

3.2 Empirical Risk Minimization: GD vs SAM

Consider a N -sample training dataset D = { ( x i , y i ) } N i =1 in which each data point is generated from the data distribution in Definition 3.1. The empirical loss function of a model f ( x ; W ) reads

where l is the logistic loss defined as l ( z ) = log(1 + exp( -z )) . The solution W ⋆ of the empirical risk minimization (ERM) minimizes the above loss, i.e., W ⋆ := arg min W L ( W ) .

GD. Typically, ERM is solved using gradient descent (GD). The update rule at iteration t of GD with learning rate η > 0 reads

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 at iteration t reads

where ρ ( t ) = ρ > 0 is the inner step size that is usually normalized by gradient norm, i.e., ρ ( t ) = ρ/ ∥ ∥ ∇L ( W ( t ) ) ∥ ∥ F .

3.3 Comparing Learning Between fast-learnable & slow-learnable Features for GD & SAM

Next, we present our theoretical results on 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., ⟨ w ( t ) j , v e ⟩ and ⟨ w ( t ) j , v d ⟩ . We first prove that early in training, both GD and SAM only learn fast-learnable feature. Then, we show SAM learns slow-learnable and fast-learnable features at a more uniform speed.

Theorem 3.2 ( GD Feature Learning ) . Consider training a two-layer nonlinear CNN model initialized with W (0) ∼ N (0 , σ 2 0 ) on the training dataset D = { ( x i , y i ) } N i =1 with distribution D ( β e , β d , α ) with α 1 / 3 β e > β d . For a small-enough learning rate η , after training for T GD iterations, w.h.p., the model: (1) learns the fast-learnable feature v e : max j ∈ [ J ] ⟨ w ( T GD ) j , v e ⟩≥ ˜ Ω(1 /β e ) ;(2) does not learn the slow-learnable feature v d : max j ∈ [ J ] ⟨ w ( T GD ) j , v d ⟩ = ˜ O ( σ 0 ) .

Theorem 3.3 ( SAM Feature Learning ) . Consider training a two-layer nonlinear CNN model initialized with W (0) ∼ N (0 , σ 2 0 ) on the training dataset D = { ( x i , y i ) } N i =1 with distribution D ( β e , β d , α ) with α 1 / 3 β e > β d . For small-enough learning rate η and perturbation radius ρ , after training for T SAM > T GD iterations, w.h.p., the model: (1) learns the fast-learnable feature v e : max j ∈ [ J ] ⟨ w ( T SAM ) j , v e ⟩ ≥ ˜ Ω(1 /β e ) ; (2) does not learn the slow-learnable feature v d : max j ∈ [ J ] ⟨ w ( T SAM ) j , v d ⟩ = ˜ O ( σ 0 ) .

The detailed proof of Theorems 3.2 and 3.3 are deferred to Appendices A.1 and A.2.

Discussion. Note that a larger value of ⟨ w ( t ) j , v ⟩ for v ∈ { v e , v d } indicates better learning of the feature vector v by neuron w j at iteration t . From the above two theorems, the growth rate of the fast-learnable feature is significantly faster than that of the slow-learnable feature. As a small portion (1 -α ) of the dataset does not have the fast-learnable feature, the model needs to learn the slow-learnable feature to improve the performance.

Next, we show that SAM learns fast-learnable and slow-learnable features more evenly. We denote by G ( t ) e = max j ∈ [ J ] ⟨ w ( t ) j , v e ⟩ and G ( t ) d = max j ∈ [ J ] ⟨ w ( t ) j , v d ⟩ the alignment of model weights with fast-learnable and slow-learnable features, when training with GD. Similarly, we denote by S ( t ) e and S ( t ) d the alignment of model weights with fast-learnable and slow-learnable , 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 η and the perturbation radius ρ are sufficiently small. Starting from the same initialization, the growth of fast-learnable and slowlearnable features in SAM is more balanced than that in SGD, i.e., for every iteration t ∈ [1 , T 0 ] :

We prove Theorem 3.4 by induction in Appendix A.2 and back it by toy experiments in Section 5.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 SAMthan it does in SGD. Thus, the slow-learnable feature benefits SAM, by reducing its overreliance on the fast-learnable features. We note that as neural networks are nonlinear, a 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-learnable feature exists in all examples, i.e., β e = β d = α = 1 , the gap in Eq. 5 is significant as we confirm in Figure 8 in Appendix D.

Remark. The network often overfits slow-learnable features that are learned late during the training and do not learn them in a 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 β d . 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 t during early training, we have the following results:

The slow-learnable feature has a larger contribution to the normalized gradient of the 1-step SAM update, compared to that of GD.

Amplifying the strength of the slow-learnable feature increases its contribution to the normalized gradients of GD and SAM.

There exists an upsampling factor k s.t. the normalized gradient of the 1-step GD update on D ( β e , kβ d , α ) recovers the normalized gradient of the 1-step SAM update on D ( β e , β d , α ) .

Algorithm 1 UpSample Early For Uniform Learning (USEFUL)

Input: Original dataset D , Model f ( · , W (0) ) , Separating epoch t , Total epochs T .

Train the model f ( · , W (0) ) on D for t epochs.

for every class c ∈ D do

D = D ∪ C 2 , where C 2 is the cluster with higher average loss

end for

Train f ( · , W (0) ) on D for T epochs

Output: Model f ( · , W ( T ) )

Slow-learnable

airplane

automobile

bird

cat

deer

dog

frog

horse

ship

truck

Fast-learnable

Discussion. Proof of Theorem 3.5 is given in Appendix A.4. We see that we can learn features at a more uniform speed by training on a new dataset D ( β e , β ′ d , α ) with a larger strength β ′ d > β d . But, the value of coefficient β ′ d varies by the model weights and gradient at each iteration t .

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 (S)GD 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 a 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 at a more uniform speed, it still suffers from simplicity bias and learns fast-learnable features earlier (although less so than GD) as 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 slow-learnable features in the training data. This drive the network to learn fast-learnable and slow-learnable features at 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 of a toy example and CIFAR-10 in Fig. 2. Motivated by our theory, we seek to find a cluster of examples with similar model outputs early in training. To do so, we 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 y j = c , we find:

0.0

0.2

0.4

0.6

0.8

PC1

0.00

0.25

0.50

0.75

1.00

PC2

y=1

slow-learnble

fast-learnable

4

3

2

1

0

1

2

PC1

4

2

0

2

PC2

y=1

slow-learnble

fast-learnable

0

100

200

300

400

500

600

Iteration

0.0

0.5

1.0

1.5

2.0

G

(

t

)

e

G

(

t

)

d

S

(

t

)

e

S

(

t

)

d

(a)

β

d

= 0

.

2

0

100

200

300

400

500

600

Iteration

0.0

0.5

1.0

1.5

2.0

G

(

t

)

e

G

(

t

)

d

S

(

t

)

e

S

(

t

)

d

(b)

β

d

= 0

.

4

0

100

200

300

400

500

600

Iteration

0.0

0.2

0.4

0.6

G

(

t

)

d

:

d

=0.2, Acc = 95.21%

S

(

t

)

d

:

d

=0.2, Acc = 96.90%

G

(

t

)

d

:

d

=0.4, Acc = 96.90%

S

(

t

)

d

:

d

=0.4, Acc = 99.49%

(c)

β

d

= 0

.

2

vs.

β

d

= 0

.

4

where µ i is the center of cluster S i . The cluster with lower average loss will contain the majority of examples containing fast-learnable features, whereas the remaining examples contain slow-learnable features in the training data. Examples of images in fast-learnable and slow-learnable clusters of CIFAR-10 found by USEFUL are illustrated in Fig. 1.

The choice of clustering. Our choice of clustering is 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 a 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-learnable 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 hyperparameter tuning and performs well for separating examples, as we confirm in our ablation studies.

Step 2: One-shot upsampling of slow-learnable features. Next, we upsample examples that are not in the cluster of points containing fast-learnable features. This speeds up learning slowlearnable features and encourages the model to learn different features at a more uniform speed. Thus, 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 1-shot algorithm that finds fast-learnable examples at an early training iteration and upsample the remaining examples by a factor of k = 2 effectively improves the performance. Notably, in contrast to dynamic sampling or reweighting, USEFUL upsamples examples only once and restart training on the modified but fix distribution.

When to separate the examples. It 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 t to separate examples is when the change in training error starts to shrink as visualized in Figure 15a. More details can be found in Appendices C.2 and D.8.

The pseudocode of USEFUL is illustrated in Alg. 1 and the workflow is shown in Appendix Fig. 7.

1x

1.3x

2x

2.6x

2x

2.6x

Time to Train w.r.t. SGD

3.0

3.5

4.0

4.5

5.0

5.5

Test Error (%)

5.07

4.79

4.23

4.04

4.06

3.49

CIFAR10

1x

1.5x

2x

3x

2x

3x

Time to Train w.r.t. SGD

18

20

22

24

24.06

22.58

21.25

20.4

18.9

17.72

CIFAR100

SGD

SGD+USEFUL

SAM

SAM+USEFUL

SAM+USEFUL+TA

1x

1.5x

2x

3x

2x

3x

Time to Train w.r.t. SGD

12

14

16

18

20

22

24

Test Error (%)

21.45

18.41

19.63

17.72

14.03

11.81

STL10

1x

1.45x

2x

2.9x

2x

2.9x

Time to Train w.r.t. SGD

9.0

9.5

10.0

10.5

11.0

11.5

12.0

11.05

10.89

9.82

9.2

11.03

9.57

CINIC10

1x

1.6x

2x

3.2x

2x

3.2x

Time to Train w.r.t. SGD

29

30

31

32

33

34

35

34.04

33.73

32.45

31.83

29.93

29.21

Tiny-ImageNet

5 Experiments

Outline. In Sec. 5.1, we empirically validate our theoretical results on 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. 5.4. Additional experimental results are deferred to Appendix D where we 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. Wefurther 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 D.8.

Settings. We used common datasets for image classification including CIFAR10, CIFAR100 [41], STL10 [13], CINIC10 [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 CIFAR100 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 set weight decay to 0.0005. We also fixed ρ = 0 . 1 for SAM in all experiments unless explicitly stated. We used a linear learning rate schedule starting at 0.1 and decay by a factor of 10 once at epoch 100 and again at epoch 150. More details are given in Appendix C.

5.1 Toy Datasets

Datasets. Following [18], our toy dataset consists of training and test sets, each containing 10K examples generated from the data distribution defined in 3.1 with dimension d =50 and P =3 . We set β e =1 , β d =0 . 2 , α =0 . 9 , and σ p / √ d =0 . 125 . We also consider a scenario with larger β d =0 . 4 . We shuffle the order of patches randomly to confirm that our theory holds with arbitrary order of patches.

Training. We used the two-layer nonlinear CNN in Section 3.1 with J = 40 filters. For GD, we set the learning rate to η = 0 . 1 and did not use momentum. For SAM, we used the same base GD optimizer and chose a smaller value of the inner step, ρ = 0 . 02 , than other experiments to 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 learn the fast-learnable feature. In particular, the blue dotted line ( G e ) accelerates quickly at around epoch 250 while the orange dotted line ( S e ) increases drastically much later at around epoch 450. That is: (1) GD learns the fast-

1x

1.3x

2x

2.6x

2x

2.6x

Time to Train w.r.t. SGD

5.0

5.5

6.0

6.5

7.0

Test Error (%)

6.68

6.34

5.8

5.54

5.74

5.11

VGG19

1x

1.5x

2x

3x

2x

3x

Time to Train w.r.t. SGD

8.0

8.5

9.0

9.5

10.0

10.5

11.0

10.43

10.0

9.4

8.89

9.49

8.46

DensetNet121

1x

1.45x

2x

2.9x

2x

2.9x

Time to Train w.r.t. SGD

17

18

19

20

21

20.27

19.57

20.11

18.99

20.37

17.89

ViT-S

SGD

SGD+USEFUL

SAM

SAM+USEFUL

SAM+USEFUL+TA

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 ( t ) e -S ( t ) d ) is much smaller than that of GD ( G ( t ) e -G ( t ) d ) . 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-learnable feature strength β d from 0.2 to 0.4 in Figure 3b, the same conclusion for the growth speed of fast-learnable and slowlearnable features holds. Notably, there is a clear increase in the classification accuracy of the model trained with either GD or SAM by increasing β d , as can be seen in Figure 3c. 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 of SAM. The percentages of examples found for upsampling by USEFUL for CIFAR10, CIFAR100, STL10, CINIC10, and Tiny-ImageNet are roughly 30%, 50%, 50%, 45%, 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. TrivialAugment [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) slow-learnable 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, we confirm the versatility of our method, by applying it to different model architectures including 3-layer MLP, CNNs (ResNet18, VGG19, DenseNet121), 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 a lower level than that of SAM alone. Detailed results for 3-layer MLP is given in Appendix D.2.

Settings: 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.1, 0.2, 0.4. In Appendix D.4, we confirm that USEFUL applied to ASAM [42]-a SAM variant which uses a scale-invariant sharpness measure-further reduces the test error.

USEFUL vs Random Upsampling. Fig. 6 shows that USEFUL considerably outperforms SGD and SAM on randomly upsampled CIFAR10 & CIFAR100. This confirms that the main benefit of USEFUL is due to the modified distribution and not longer training time. In Appendix D.8, we also confirm that upsampling outperforms upweighting for SAM & SGD.

SGD

SAM

3.8

4.0

4.2

4.4

4.6

4.8

5.0

5.2

Test Error (%)

CIFAR10

SGD

SAM

20

21

22

23

24

25

CIFAR100

Orig

Rand

USEFUL

5.4 USEFUL's Solution has Similar Properties to SAM

SAM & USEFUL Find Sparser Solutions than SGD. [3] showed that SAM's solution has a better sparsity-inducing property indicated by the L1 norm than the standard ERM. Fig. 10 shows the L1 norm of ResNet18 trained on CIFAR10 and ResNet34 trained on CIFAR100 at the end of training. We see that USEFUL drives both SGD and SAM to find solutions with smaller L1 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 ( λ max ) and the bulk of the spectrum ( λ max /λ 5 ) [30], which are commonly used metrics for sharpness [10, 31, 37, 74]. Table 1 illustrates that SGD+USEFUL on CIFAR10 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 is 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.

SAM & USEFUL Reduce Forgetting Scores. Forgetting scores [69] 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 in Appendix D.3 that both SAM and USEFUL successfully reduce the forgetting scores, thus learn slow-learnable 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.

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 to 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%, respectively. We also show the applicability of USEFUL to fine-tuning a ResNet50 pre-trained on ImageNet.

6 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-aware minimization (SAM), and attributed its superior performance over GD to mitigating the simplicity bias, and learning features at a more speed. Inspired by SAM, we upsampled the examples that contain slow-learnable features to alleviate the simplicity bias. This allows learning features more uniformly, thus improving the performance. Our method boosts the performance of image classifiers trained with SGD or SAM and easily 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 a UCLA Hellman Fellowship.

References

Amro Abbas, Kushal Tirumala, Dániel Simig, Surya Ganguli, and Ari S Morcos. Semdedup: Data-efficient learning at web-scale through semantic deduplication. arXiv preprint arXiv:2303.09540 , 2023.

Zeyuan Allen-Zhu and Yuanzhi Li. Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. arXiv preprint arXiv:2012.09816 , 2020.

Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware minimization. In International Conference on Machine Learning , pages 639-668. PMLR, 2022.

Maksym Andriushchenko, Dara Bahri, Hossein Mobahi, and Nicolas Flammarion. Sharpnessaware minimization leads to low-rank features. arXiv preprint arXiv:2305.16292 , 2023.

Robert Baldock, Hartmut Maennel, and Behnam Neyshabur. Deep learning through the lens of example difficulty. Advances in Neural Information Processing Systems , 34:10876-10889, 2021.

Peter L Bartlett, Philip M Long, and Olivier Bousquet. The dynamics of sharpness-aware minimization: Bouncing across ravines and drifting towards wide minima. Journal of Machine Learning Research , 24(316):1-36, 2023.

Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconciling modern machinelearning practice and the classical bias-variance trade-off. Proceedings of the National Academy of Sciences , 116(32):15849-15854, 2019.

Yuan Cao, Zixiang Chen, Misha Belkin, and Quanquan Gu. Benign overfitting in two-layer convolutional neural networks. Advances in neural information processing systems , 35:2523725250, 2022.

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.

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.

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-23062, 2022.

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.

Adam Coates, Andrew Ng, and Honglak Lee. An analysis of single-layer networks in unsupervised feature learning. In Proceedings of the fourteenth international conference on artificial intelligence and statistics , pages 215-223. JMLR Workshop and Conference Proceedings, 2011.

Elliot Creager, Jörn-Henrik Jacobsen, and Richard Zemel. Environment inference for invariant learning. In International Conference on Machine Learning , pages 2189-2200. PMLR, 2021.

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.

Luke N Darlow, Elliot J Crowley, Antreas Antoniou, and Amos J Storkey. Cinic-10 is not imagenet or cifar-10. arXiv preprint arXiv:1810.03505 , 2018.

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A largescale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition , pages 248-255. Ieee, 2009.

Yihe Deng, Yu Yang, Baharan Mirzasoleiman, and Quanquan Gu. Robust learning with progressive data expansion against spurious correlation. arXiv preprint arXiv:2306.04949 , 2023.

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.

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 arXiv:2110.03141 , 2021.

Ayoub El Hanchi, David Stephens, and Chris Maddison. Stochastic reweighted gradient descent. In International Conference on Machine Learning , pages 8359-8374. PMLR, 2022.

Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. arXiv preprint arXiv:2010.01412 , 2020.

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.14108 , 2023.

Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An investigation into neural net optimization via hessian eigenvalue density. In International Conference on Machine Learning , pages 2232-2241. PMLR, 2019.

Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. Advances in neural information processing systems , 30, 2017.

Katherine Hermann and Andrew Lampinen. What shapes feature representations? exploring datasets, architectures, and training. Advances in Neural Information Processing Systems , 33: 9995-10006, 2020.

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.15556 , 2022.

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-17128, 2020.

Gareth James, Daniela Witten, Trevor Hastie, Robert Tibshirani, et al. An introduction to statistical learning , volume 112. Springer, 2013.

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.

Stanisław Jastrzkebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in sgd. arXiv preprint arXiv:1711.04623 , 2017.

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.

Tyler B Johnson and Carlos Guestrin. Training deep models faster with robust, approximate importance sampling. Advances in Neural Information Processing Systems , 31, 2018.

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 in neural information processing systems , 32, 2019.

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.

Simran Kaur, Jeremy Cohen, and Zachary Chase Lipton. On the maximum hessian eigenvalue and generalization. In Proceedings on , pages 51-65. PMLR, 2023.

Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv:1609.04836 , 2016.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Polina Kirichenko, Pavel Izmailov, and Andrew Gordon Wilson. Last layer re-training is sufficient for robustness to spurious correlations. arXiv preprint arXiv:2204.02937 , 2022.

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.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.

Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. Asam: Adaptive sharpnessaware minimization for scale-invariant learning of deep neural networks. In International Conference on Machine Learning , pages 5905-5914. PMLR, 2021.

Ya Le and Xuan Yang. Tiny imagenet visual recognition challenge. CS 231N , 7(7):3, 2015.

Katherine Lee, Daphne Ippolito, Andrew Nystrom, Chiyuan Zhang, Douglas Eck, Chris Callison-Burch, and Nicholas Carlini. Deduplicating training data makes language models better. arXiv preprint arXiv:2107.06499 , 2021.

Junnan Li, Richard Socher, and Steven CH Hoi. Dividemix: Learning with noisy labels as semi-supervised learning. arXiv preprint arXiv:2002.07394 , 2020.

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.

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.

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.

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.

Samuel G Müller and Frank Hutter. Trivialaugment: Tuning-free yet state-of-the-art data augmentation. In Proceedings of the IEEE/CVF international conference on computer vision , pages 774-782, 2021.

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: Theory and Experiment , 2021(12):124003, 2021.

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.6614 , 2014.

Vardan Papyan. The full spectrum of deepnet hessians at scale: Dynamics with sgd training and sample size. arXiv preprint arXiv:1811.07062 , 2018.

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, high-performance deep learning library. Advances in neural information processing systems , 32, 2019.

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 , 34:1256-1272, 2021.

Hieu Pham, Melody Guan, Barret Zoph, Quoc Le, and Jeff Dean. Efficient neural architecture search via parameters sharing. In International conference on machine learning , pages 40954104. PMLR, 2018.

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:71874-71910, 2023.

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.

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-67, 2020.

Nicolas Roux, Mark Schmidt, and Francis Bach. A stochastic gradient method with an exponential convergence _rate for finite training sets. Advances in neural information processing systems , 25, 2012.

Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. arXiv preprint arXiv:1911.08731 , 2019.

Mark Schmidt, Nicolas Le Roux, and Francis Bach. Minimizing finite sums with the stochastic average gradient. Mathematical Programming , 162:83-112, 2017.

Harshay Shah, Kaustav Tamuly, Aditi Raghunathan, Prateek Jain, and Praneeth Netrapalli. The pitfalls of simplicity bias in neural networks. Advances in Neural Information Processing Systems , 33:9573-9585, 2020.

Vatsal Shah, Anastasios Kyrillidis, and Sujay Sanghavi. Minimum norm solutions do not always generalize well for over-parameterized problems. stat , 1050:16, 2018.

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.

Jacob Steinhardt, Pang Wei W Koh, and Percy S Liang. Certified defenses for data poisoning attacks. Advances in neural information processing systems , 30, 2017.

Damien Teney, Ehsan Abbasnejad, Simon Lucey, and Anton Van den Hengel. Evading the simplicity bias: Training a diverse set of models discovers solutions with superior ood generalization. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 16761-16772, 2022.

Rishabh Tiwari and Pradeep Shenoy. Overcoming simplicity bias in deep networks using a feature sieve. In International Conference on Machine Learning , pages 34330-34343. PMLR, 2023.

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.05159 , 2018.

Guillermo Valle-Perez, Chico Q Camargo, and Ard A Louis. Deep learning generalizes because the parameter-function map is biased towards simple functions. In International Conference on Learning Representations , 2018.

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-8778, 2018.

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.

Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How does sharpness-aware minimization minimize sharpness? arXiv preprint arXiv:2211.05729 , 2022.

Yeming Wen, Kevin Luk, Maxime Gazeau, Guodong Zhang, Harris Chan, and Jimmy Ba. An empirical study of large-batch stochastic gradient descent with structured covariance noise. arXiv preprint arXiv:1902.08234 , 2019.

Yu Yang, Eric Gan, Gintare Karolina Dziugaite, and Baharan Mirzasoleiman. Identifying spurious biases early in training through the lens of simplicity bias. arXiv preprint arXiv:2305.18761 , 2023.

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 10665-10673, 2021.

Li Yuan, Yunpeng Chen, Tao Wang, Weihao Yu, Yujun Shi, 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.

Runtian Zhai, Chen Dan, Zico Kolter, and Pradeep Ravikumar. Understanding why generalized reweighting does not improve over erm. arXiv preprint arXiv:2201.12293 , 2022.

Hongyi Zhang. mixup: Beyond empirical risk minimization. arXiv preprint arXiv:1710.09412 , 2017.

Peilin Zhao and Tong Zhang. Stochastic optimization with importance sampling for regularized loss minimization. In international conference on machine learning , pages 1-9. PMLR, 2015.

Yaowei Zheng, Richong Zhang, and Yongyi Mao. Regularizing neural networks via adversarial model perturbation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 8156-8165, 2021.

Barret Zoph and Quoc V Le. Neural architecture search with reinforcement learning. arXiv preprint arXiv:1611.01578 , 2016.

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 a vector v , we use ∥ v ∥ 2 to denote its Euclidean norm. Given two sequence { x n } and { y n } , we denote x n = O ( y n ) if | x n | ≤ C 1 | y n | for some absolute positive constant C 1 , x n = Ω( y n ) if | x n | ≥ C 2 | y n | for some absolute positive constant C 2 , and x n = Θ( y n ) if C 3 | y n | ≤ | x n | ≤ C 4 | y n | for some absolute constant C 3 , C 4 > 0 . In addition, we use ˜ O ( · ) , ˜ Ω( · ) , and ˜ Θ( · ) to hide logarithmic factors in these notations. Furthermore, we denote x n = poly ( y n ) if x n = O ( y D n ) for some positive constant D, and x n = polylog ( y n ) if x n = poly (log( y n )) .

First, we have the following assumption for the model weight initialization.

Assumption A.1 (Weight initialization) . Assume that we initialize W (0) ∼ N (0 , σ 2 0 ) such that for all j ∈ [ J ] , ⟨ w (0) j , v e ⟩ , ⟨ w (0) j , v d ⟩ ≥ ρ > 0 .

The above assumption is reasonable because we later show that both sequences ⟨ w ( t ) j , v e ⟩ and ⟨ w ( t ) j , v d ⟩ 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 αN is an integer and the first αN data examples have the fast-learnable feature while the rest do not. Before going into the analysis, we denote the derivative of a data example i at iteration t to be

Lemma A.2 (Gradient) . Let the loss function L be as defined in Equation 2. For t ≥ 0 and j ∈ [ J ] , the gradient of the loss L ( W ( t ) ) with regard to neuron w ( t ) j is

Proof. We have the following gradient

With the above formula of gradient, we have the following equations:

Fast-learnable feature gradient. The projection of the gradient on v e is

Slow-learnable feature gradient. The projection of the gradient on v d is

Noise gradient. The projection of the gradient on ξ i is

̸

Derivative of data example i . For 1 ≤ i ≤ αN , l ( t ) i can be rewritten as

while for αN +1 ≤ i ≤ N , l ( t ) i can be rewritten as

Note that 0 < l ( t ) i < 1 due to the property of the sigmoid function. Furthermore, we similarly consider that the sum of the sigmoid terms for all time steps is bounded up to a logarithmic dependence [11]. The sigmoid term is considered small for a κ such that

which implies κ ≥ ˜ Ω(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 (Fast-learnable feature update.) . For all t ≥ 0 and j ∈ [ J ] , the fast-learnable feature update is

where g 1 ( t ) = sigmoid ( ∑ J j =1 -β 3 d ⟨ w ( t ) j , v d ⟩ 3 -β 3 e ⟨ w ( t ) j , v e ⟩ 3 ) .

Proof. Plugging the update rule of GD, we have

where the last equality holds due to Lemma B.5.

Similarly, we obtain the following update rule for slow-learnable features.

Lemma A.4 (Slow-learnable feature update.) . For all t ≥ 0 and j ∈ [ J ] , the fast-learnable feature update is

which gives

Proof. Plugging the update rule of GD, we have

From Lemma B.5, we have for 1 ≤ i ≤ αN,l ( t ) i = Θ(1) g 1 ( t ) and for αN +1 ≤ i ≤ N, Θ(1) g 1 ( t ) ≤ l ( t ) i ≤ 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 T 0 > 0 be such that max j ∈ [ J ] ⟨ w ( T 0 ) j , v e ⟩ ≥ ˜ Ω(1 /β e ) . For t ∈ [0 , T 0 ] , the fast-learnable feature update has the following rule

Proof. Let T 0 > 0 be such that either max j ∈ [ J ] ⟨ w ( T 0 ) j , v e ⟩ ≥ ˜ Ω(1 /β e ) or max j ∈ [ J ] ⟨ w ( T 0 ) j , v d ⟩ ≥ ˜ Ω(1 /β d ) . We will show later that the first condition will be met and we have max j ∈ [ J ] ⟨ w ( T 0 ) j , v d ⟩ ≤ ˜ Ω(1 /β d ) for all j ∈ [ J ] and t ∈ [0 , T 0 ] .

Recall that g 1 ( t ) = sigmoid ( ∑ J j =1 -β 3 d ⟨ w ( t ) j , v d ⟩ 3 -β 3 e ⟨ w ( t ) j , v e ⟩ 3 ) . Then, for t ∈ [0 , T 0 ] , we have

where the first inequality holds due to ⟨ w ( t ) j , v e ⟩ ≤ κ/ ( J 1 / 3 β e ) and ⟨ w ( t ) j , v d ⟩ ≤ κ/ ( J 1 / 3 β d ) for t ∈ [0 , T 0 ] [18][Lemma E.3]. Therefore, similar to [18, 32], we have g 1 ( t ) = Θ(1) in the early iterations. This implies the result in Lemma A.3 as

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 T 0 > 0 be such that max j ∈ [ J ] ⟨ w ( T 0 ) j , v e ⟩ ≥ ˜ Ω(1 /β e ) . For t ∈ [0 , T 0 ] , the fast-learnable feature update has the following rule

Wenext show that GD will learn the fast-learnable feature quicker than learning the slow-learnable feature.

Lemma A.7. Assume η = ˜ o ( β d σ 0 ) . Let T 0 be the iteration number that max j ∈ [ J ] ⟨ w ( T 0 ) j , v e ⟩ reaches ˜ Ω(1 /β e ) = ˜ Θ(1) . Then, we have for all t ≤ T 0 , it holds that max j ∈ [ J ] ⟨ w ( T 0 ) j , v d ⟩ = ˜ O ( σ 0 ) .

Proof. Among all the possible indices j ∈ [ J ] , we focus on the index j ⋆ = arg max j ∈ [ J ] ⟨ w (0) j , v e ⟩ . Therefore, for C t = αβ 3 e = Θ(1) , we apply Lemma B.2 with two positive sequences ⟨ w ( t ) j ⋆ , v e ⟩ and ⟨ w ( t ) j , v d ⟩ defined in Lemmas A.5 and A.6 and get

for all j ∈ [ J ] .

Theorem A.8 (Restatement of Theorem 3.2) . We consider training a two-layer nonlinear CNN model initialized with W (0) ∼ N (0 , σ 2 0 ) on the training dataset D = { ( x i , y i ) } N i =1 that follows the data distribution D ( β e , β d , α ) with α 1 / 3 β e > β d . After training with GD in Equation 3 for T GD iterations where

for all j ∈ [ J ] and t ∈ [0 , T GD ) , we have

After training for T GD iterations, with high probability, the learned weight has the following properties: (1) it learns the fast-learnable feature v e : max j ∈ [ J ] ⟨ w ( T GD ) j , v e ⟩ ≥ ˜ Ω(1 /β e ) ; (2) it does not learn the slow-learnable feature v d : max j ∈ [ J ] ⟨ w ( T GD ) j , v d ⟩ = ˜ O ( σ 0 ) .

Proof. From the results of Lemmas A.5- A.7, it remains to calculate the time T GD. Plugging v = ˜ Ω(1 /β e ) , m = M = ˜ Θ( η ) αβ 3 e , z 0 = ˜ O ( σ 0 ) into Lemma B.3, we have T GD as

A.2 Proof of Theorem 3.3

Before going into the analysis, we denote the derivative of a data example i at iteration t to be

where ϵ ( t ) = ρ ( t ) ∇L ( W ( t ) ) is the weighted ascent direction at the current parameter W ( t ) . We denote the weight vector of the j -th filter after being perturbed by SAM as

where ρ ( t ) = ρ/ ∥ ∥ ∇L ( W ( t ) ) ∥ ∥ F .

First, we have the following inequalities regarding the gradient norm:

Thus,

Lemma A.9 (Gradient) . Let the loss function L be as defined in Equation 2. For t ≥ 0 and j ∈ [ J ] , the gradient of the loss L ( W ( t ) + ϵ ( t ) ) with regard to neuron w ( t ) j,ϵ is

Proof. We have the following gradient

With the above formula of gradient, we have the projection of perturbed weight on v e is

From Equations 32 and 34, we have

Similarly, the projection of perturbed weight on v d is

Fast-learnable feature gradient. The projection of the gradient on v e is

Slow-learnable feature gradient. The projection of the gradient on v d is

Noise gradient. The projection of the gradient on ξ i is

̸

Derivative of data example i . For 1 ≤ i ≤ αN , l ( t ) i,ϵ can be rewritten as

while for αN +1 ≤ i ≤ N , l ( t ) i can be rewritten as

We present the detailed proofs that build up to Theorem A.15. We begin by considering the update for the fast-learnable and slow-learnable features.

Lemma A.10 (Fast-learnable feature update.) . For all t ≥ 0 and j ∈ [ J ] , the fast-learnable feature update is

Proof. Plugging the update rule of SAM, we have

where the last equality holds due to Lemma B.10. Combining with Equation 35, we obtain the desired inequalities.

Similarly, we obtain the following update rule for slow-learnable features.

Lemma A.11 (Slow-learnable feature update.) . For all t ≥ 0 and j ∈ [ J ] , the slow-learnable feature update is

where g 2 ( t ) = sigmoid ( ∑ J j =1 -β 3 d ⟨ w ( t ) j,ϵ , v d ⟩ 3 -β 3 e ⟨ w ( t ) j,ϵ , v e ⟩ 3 ) .

Proof. Plugging the update rule of GD, we have

From Lemma B.5, we have for 1 ≤ i ≤ αN,l ( t ) i = Θ(1) g 1 ( t ) and for αN +1 ≤ i ≤ N, Θ(1) g 1 ( t ) ≤ l ( t ) i ≤ 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 T 0 > 0 be such that max j ∈ [ J ] ⟨ w ( T 0 ) j , v e ⟩ ≥ ˜ Ω(1 /β e ) . For t ∈ [0 , T 0 ] , the fast-learnable feature update has the following rule

Proof. Let T 0 > 0 be such that either max j ∈ [ J ] ⟨ w ( T 0 ) j , v e ⟩ ≥ ˜ Ω(1 /β e ) or max j ∈ [ J ] ⟨ w ( T 0 ) j , v d ⟩ ≥ ˜ Ω(1 /β d ) . We will show later that the first condition will be met and we have max j ∈ [ J ] ⟨ w ( T 0 ) j , v d ⟩ ≤ ˜ Ω(1 /β d ) for all j ∈ [ J ] and t ∈ [0 , T 0 ] .

Recall that g 2 ( t ) = sigmoid ( ∑ J j =1 -β 3 d ⟨ w ( t ) j,ϵ , v d ⟩ 3 -β 3 e ⟨ w ( t ) j,ϵ , v e ⟩ 3 ) . Then, for t ∈ [0 , T 0 ] , we have

where the first inequality holds due to ⟨ w ( t ) j,ϵ , v e ⟩ ≤ ⟨ w ( t ) j , v e ⟩ ≤ κ/ ( J 1 / 3 β e ) and ⟨ w ( t ) j,ϵ , v d ⟩ ≤ ⟨ w ( t ) j , v d ⟩ ≤ κ/ ( J 1 / 3 β d ) for t ∈ [0 , T 0 ] . Therefore, we have g 2 ( t ) = Θ(1) in the early iterations. Replacing g 2 ( t ) = Θ(1) into the results of Lemma A.10, we obtain the desired results.

Similarly, we obtain the following simplified update rule for slow-learnable features in the early iterations.

Lemma A.13 (Slow-learnable feature update in early iterations) . Let T 0 > 0 be such that max j ∈ [ J ] ⟨ w ( T 0 ) j , v e ⟩ ≥ ˜ Ω(1 /β e ) . For t ∈ [0 , T 0 ] , the fast-learnable feature update has the following rule

We next show that SAM will learn the fast-learnable feature quicker than the slow-learnable one.

Lemma A.14. Assume η = ˜ o ( β d σ 0 ) . Let T 0 be the iteration number that max j ∈ [ J ] ⟨ w ( T 0 ) j , v e ⟩ reaches ˜ Ω(1 /β e ) = ˜ Θ(1) . Then, we have for all t ≤ T 0 , it holds that max j ∈ [ J ] ⟨ w ( T 0 ) j , v d ⟩ = ˜ O ( σ 0 ) .

Proof. Among all the possible indices j ∈ [ J ] , we focus on the index j ⋆ = arg max j ∈ [ J ] ⟨ w (0) j , v e ⟩ . Therefore, for C t = αβ 3 e = Θ(1) , we apply Lemma B.2 with two positive sequences ⟨ w ( t ) j ⋆ , v e ⟩ and ⟨ w ( t ) j , v d ⟩ defined in Lemmas A.12 and A.13 and get

for all j ∈ [ J ] .

Theorem A.15 (Restatement of Theorem 3.3) . We consider training a two-layer nonlinear CNN model initialized with W (0) ∼ N (0 , σ 2 0 ) on the training dataset D = { ( x i , y i ) } N i =1 that follows the data distribution D ( β e , β d , α ) with α 1 / 3 β e > β d . After training with SAM in Equation 3 for T SAM iterations where

for all j ∈ [ J ] and t ∈ [0 , T SAM ) , we have

After training for T SAM iterations, with high probability, the learned weight has the following properties: (1) it learns the fast-learnable feature v e : max j ∈ [ J ] ⟨ w ( T SAM ) j , v e ⟩ ≥ ˜ Ω(1 /β e ) ; (2) it does not learn the slow-learnable feature v d : max j ∈ [ J ] ⟨ w ( T SAM ) j , v d ⟩ = ˜ O ( σ 0 ) .

Proof. With the results of Lemmas A.12- A.14, it remains to calculate the time T SAM. Plugging v = ˜ Ω(1 /β e ) , m = M = ˜ Θ( η ) αβ 3 e , z 0 = ˜ O ( σ 0 ) into Lemma B.3, we have T SAM as

Comparing Eq. 22 and 48, we can see that SAM learns the fast-learnable features later than GD. Particularly, if we remove the approximate notations, we have the following inequality

which holds due to Assumption A.1 about weight initialization in Appendix A.1, i.e., ( σ 0 ≥ ρ ≥ 0 ).

A.3 Proof of Theorem 3.4

In this section, we show that SAM learns fast-learnable and slow-learnable features at a more uniform speed. To ease the notation, we denote G ( t ) e = max j ∈ [ J ] ⟨ w ( t ) j , v e ⟩ and G ( t ) d = max j ∈ [ J ] ⟨ w ( t ) j , v d ⟩ for model weights trained with GD. Similarly, we denote S ( t ) e and S ( t ) d for model weights trained with SAM. We use ˆ S ( t ) e and ˆ S ( t ) d to denote the inner products with perturbed weights. We simplify Equation 34 and 36 for early iterations t ≤ T 0 as

Before introducing the theorem, we assume that the model is initialized in favor of the fastlearnable feature, i.e. G (0) e -G (0) d ≥ ρ . 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.1).

Theorem A.16 (Restatement of Theorem 3.4) . Consider the training dataset D = { ( x i , y i ) } N i =1 that follows the data distribution D ( β e , β d , α ) using the two-layer nonlinear CNN model initialized with W (0) ∼ N (0 , σ 2 0 ) . Assume that the fast-learnable feature strength is significantly larger α 1 / 3 β e > β d . Training the same model initialization, we have that for every iteration t ≤ T 0

Proof. We prove this by induction. For t = 0 , the above hypotheses immediately hold because we use train two methods from the same initialization. Particularly, we have 0 < S (0) d = G (0) d < G (0) e = S (0) e and ˆ S (0) e -ˆ S (0) d ≥ S (0) e -ρ -S (0) d ≥ ( G (0) e -G (0) d ) -ρ ≥ 0 .

Assume that the induction hypotheses hold for t , i.e.

S

(

t

d

G

(

t

d

(63)

We need to prove that they also hold for t + 1 . From Lemma A.15 and the first two induction hypotheses,

Then, ˆ S ( t +1) e -ˆ S ( t +1) d ≥ S ( t +1) e -ρ -S ( t +1) d ≥ 0 . From Equation 35 and Lemma A.12,

Similarly, we have S ( t +1) d ≤ G ( t +1) d . From Equations 53 and 54,

Combining with Equations 35 and 37, we have

Therefore, the induction hypotheses hold for t +1 .

)

)

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 = { ( x i , y i ) } N i =1 that follows the data distribution D ( β e , β d , α ) using the two-layer nonlinear CNN model initialized with W (0) ∼ N (0 , σ 2 0 ) . Assume that the noise is sufficiently small (ref. Lemmas B.4 and B.9) and the fast-learnable feature strength is significantly larger α 1 / 3 β e > β d . From any iteration t during early training, the normalized gradient of the one-step SAM update has a larger weight on the slow-learnable feature compared to that of GD.

Proof. First, recall the gradients of GD and SAM are as follows.

Because the noise is sufficiently small, the above equations can be simplified as

Note that in early training, we have an approximation for the logit terms as l ( t ) i = l ( t ) i,ϵ = Θ(1) , we can further simplify the gradients as

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 feature compared to GD, it is sufficient to show the ratio of coefficients in SAM is larger than GD. In other words, we need to verify that

The last inequality holds due to αβ 3 e > β 3 d and ⟨ w ( t ) , v e ⟩ ≥ ⟨ w ( t ) , v d ⟩ from Theorem A.16.

From Equation 81 in the above proof, it can be seen clearly that amplifying the slow-learnable feature strength in either GD or SAM, i.e., increasing β d , places a larger weight on the slow-learnable feature. Thus, we have the next theorem.

Theorem A.18 (Restatement of Theorem 3.5) . Consider the training dataset D = { ( x i , y i ) } N i =1 that follows the data distribution D ( β e , β d , α ) using the two-layer nonlinear CNN model initialized with W (0) ∼ N (0 , σ 2 0 ) . Assume that the noise is sufficiently small (ref. Lemmas B.4 and B.9) and the fast-learnable feature strength is significantly larger α 1 / 3 β e > β d . We have the following results for one-step upsampling, i.e. increasing β d , from any iteration t during early training

The normalized gradient of the one-step SAM update has a larger weight on the slowlearnable feature compared to that of GD.

Amplifying the slow-learnable feature strength puts a larger weight on the slow-learnable feature in the normalized gradients of GD and SAM.

There exists an upsampling factor k such that the normalized gradient of the one-step GD update on D ( β e , kβ d , α ) recovers the normalized gradient of the one-step SAM update on D ( β e , β d , α ) .

Proof. The first result has already been proved in Lemma A.17. Now, consider increasing the slowlearnable feature strength from β d to β ′ d . Similar to the proof of Corollary A.17, to verify that the normalized of the new normalized gradient of GD favors the slow-learnable feature, it is sufficient to show

which is trivial because β ′ d > β d . Similarly, we can verify the result for SAM. Now, let's find the new coefficient β ′ d = kβ d ( k > 1) such that training one-step GD on D ( β e , β ′ d , α ) can recover the normalized gradient of the one-step SAM update on the original data distribution D ( β e , β d , α ) . Using Equation 81, we have

Therefore, with β ′ d = ( 1 -3 ρ ( t ) β 3 d ⟨ w ( t ) ,v d ⟩ 1 -3 ρ ( t ) αβ 3 e ⟨ w ( t ) ,v e ⟩ ) 2 / 3 β d , the normalized gradient of the one-step GD update on D ( β e , β ′ d , α ) is similar to that of the one-step SAM update on D ( β e , β d , α ) .

B Auxiliary Lemmas

Lemma B.1 (Claim D.20, [2]) . Considering an increasing sequence x t ≥ 0 defined as x t +1 = x t + ηC t ( x t -ρ ) 2 for some C t = Θ(1) , 0 ≤ ρ ≤ x 0 , then we have for every A > x 0 , every δ ∈ (0 , 1) , and every η ∈ (0 , 1] :

Proof. For every g = 0 , 1 , . . . , let T g be the first iteration such that x t ≥ (1 + δ ) g x 0 . Let b be the smallest integer such that (1 + δ ) b x 0 ≥ A . Suppose for notation simplicity that we replace x t with exactly A whenever x t ≥ A . By the definition of T g , we have

(96)

Recall b is the smallest integer such that (1 + δ ) b x 0 ≥ A , so we can calculate

Thus, the two desired inequalities are proved.

Lemma B.2 (Lemma D.19, [2].) . Let { x t , y t } t =1 ,... be two positive sequences that satisfy

for some C t = Θ(1) . Suppose x 0 ≥ y 0 S 1+2 G 1 -3 G where S ∈ (0 , 1) , G ∈ (0 , 1 / 3) and 0 < η ≤ min { G 2 x 0 log( A/x 0 ) , G 2 y 0 log(1 /G ) } , 0 ≤ ρ < O ( x 0 ) , and for all A ∈ ( x 0 , O (1)] , let T x be the first iteration such that x t ≥ A . Then, we have y T x ≤ O ( G -1 y 0 ) .

Proof. Let T x be the first iteration t in which x t ≥ A . Apply Lemma B.1 for the x t sequence with C t = C t and threshold A , we have

These imply that

Let T y be the first iteration t in which y t ≥ A . Apply Lemma B.1 for the y t sequence with η = Sη, C t = C t , ρ = 0 and threshold A ′ = G -1 y 0 , we have

Compare Equation 107 and 110. Choosing δ = G and η ≤ min { G 2 x 0 log( A/x 0 ) , G 2 y 0 log(1 /G ) } , together with x 0 ≥ y 0 S 1+2 G 1 -3 G we have T x ≤ T y .

Lemma B.3 (Lemma K.15, [32].) . Let { z t } T t =0 be a positive sequence defined by the following recursions

where z 0 > ρ ≥ 0 is the initialization and m,M > 0 are some constants. Let v > z 0 , then the time T v such that z T v ≥ v for all t ≥ T v is

Proof. Let n ∈ N ⋆ . Let T n be the first time that z t ≥ 2 n z 0 . We want to find an upper bound of T n . We start with the case n = 1 . By summing the recursion, we have:

Now, we want to bound z T 1 -z 0 . Using again the recursion and z T 1 -1 ≤ 2 z 0 , we have

Combining Equation 113 and 114, we get a bound on T 1 as

Now, let's find a bound for T n . Starting from the recursion and using the fact that z t ≥ 2 n -1 z 0 for t ≥ T n -1 we have

On the other hand, by using z T n -1 ≤ 2 n z 0 we upper bound z T n as

Besides, we know that z T n -1 ≥ 2 n -1 z 0 . Therefore, we upper bound z T n -z T n -1 as

Because z t ≥ z 0 , we obtain

Combining Equations 116 and 119 yields

Summing Equation 120 for n = 2 , . . . , n we have

Lastly, we know that 2 n z 0 ≥ v this implies that we can set n = ⌈ log( v/z 0 ) log(2) ⌉ in Equation 122.

We make the following assumptions for every t ≤ T as the same in [32].

Lemma B.4 (Induction hypothesis D.1, [32]) . Throughout the training process using GD for t ≤ T , we maintain that, for every i and j ∈ [ J ] ,

Lemma B.5 (Lemma G.4, [18]) . For every i , we have l ( t ) i = Θ(1) g 1 ( t ) , where

Lemma B.6 (Lemma K.5, [32]) . Let X ∈ R d be a Gaussian random vector, X ∼ N (0 , σ 2 I d ) . Then with probability at least 1 -o (1) , we have ∥ X ∥ 2 2 = Θ( σ 2 √ d ) .

Lemma B.7 (Lemma K.7, [32]) . Let X and Y be independent Gaussian random vectors on R d and X ∼ N (0 , σ 2 I d ) , Y ∼ N (0 , σ 2 0 I d ) . Assume that σσ 0 ≤ 1 d . Then, with probability at least 1 -δ , we have

Lemma B.8 (Bound on noise inner products) . Let N = O ( poly ( d )) . The following hold with probability at least 1 -o (1) :

Proof. For the first inequality, Lemma B.7 implies that with probability at least 1 -1 dN ,

Taking a union bound over n = 1 , . . . , N gives the result.

The second statement is proved similarly.

Lemma B.9 (Bound on the noise component for SAM) . Assume that ρ = o ( σ 0 ) and ω (1) ≤ N ≤ O ( poly ( d )) . Throughout the training process using SAM for t ≤ T , we maintain that, for every i and j ∈ [ J ] ,

Proof. Let χ t = max {|⟨ w ( t ) j,ϵ , ξ i ⟩|} , α = max i { 1 n ∑ n k =1 |⟨ ξ k , ξ i ⟩|} . Combined with l ( t ) k,ϵ ≤ 1 , the noise gradient update rule can be bounded as

Suppose that a ( t ) satisfies the differential equation

Observe that a ( t ) is increasing so, by the Mean Value Theorem there exists τ ∈ ( t, t +1) such that

So an fast-learnable induction shows that a ( t ) ≥ χ t .

Now solving for a ( t ) ,

We conclude that

Using the high probability tail bounds B.8,

where σ = σ p √ d . Substituting these bounds gives

Now Theorem A.15 and ρ = o ( σ 0 ) implies that ηT 0 = ˜ Θ( 1 σ 0 β 3 ϵ ) . Combined with the assumption that N = Ω(1) , the second term in the denominator is of lower order than the first term, so

Lemma B.10. For every i , we have l ( t ) i = Θ(1) g 1 ( t ) and l ( t ) i,ϵ = Θ(1) g 2 ( t ) , where

The proof is the same as [18][Lemma G.4].

C Additional Experimental Settings

C.1 Datasets and Training Details

Datasets. The CIFAR10 dataset [41] consists of 60,000 32 × 32 color images in 10 classes, with 6000 images per class. The CIFAR100 dataset [41] is just like the CIFAR10, except it has 100 classes

Original Dataset

...

Model

USEFUL

Class-Wise

Selected Samples

...

Model

containing 600 images each. For both of these datasets, the training set has 50,000 images (5,000 per class for CIFAR10 and 500 per class for CIFAR100) 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 CIFAR10. The dataset is created by merging CIFAR10 with images extracted from the ImageNet database, specifically selecting and downsampling images from the same 10 classes present in CIFAR10. Tiny-ImageNet [43] comprises 100,000 images distributed across 200 classes of ImageNet [17], with each class containing 500 images. These images have been resized to 64×64 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 CIFAR10 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, 4], we trained Pre-Activation ResNet18 on all datasets except for CIFAR100 which was trained with ResNet34. We trained our models for 200 epochs with a batch size of 128 and used basic data augmentations such as 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 ρ = 0 . 1 for SAM unless further specified. For all datasets, we used a learning rate schedule where we set the initial learning rate to 0.1. 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.

Training with different architectures. We used the same training procedures for Pre-Activation ResNet18, VGG19, and DenseNet121. 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 set ρ = 0 . 01 . For ViT-S [77], we adopted a Pytorch implementation at https://github.com/lucidrains/vit-pytorch. In particular, the hidden size, the depth, the number of attention heads, and the MLP size are set to 768, 8, 8, and 2304, respectively. We adjusted the patch size to 4 to fit the resolution of CIFAR10 and set both the initial learning rate and ρ to 0.01.

Computational resources. Each model is trained on 1 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 , 6 , 7 , 8 , 10 } for CIFAR10 and { 12 , 14 , 16 , 18 , 20 , 22 } for CIFAR100. Particularly, we separated examples of CIFAR10 at around epoch 8 while that of CIFAR100 is near epoch 20. Near this point, the gain in training error diminishes significantly as shown in Figure 15a, which shows a sign that the model successfully learns fast-learnable features. In addition, we reported the results for different separating epochs in Appendix D.8. 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.

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-learnable example which is repeated in the upsampled dataset.

Hessian spectra. We approximated the density of the Hessian spectrum using the Lanczos algorithm [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 ( λ max ) and the bulk of spectra ( λ max /λ 5 ) [30].

0

100

200

300

400

500

600

Iteration

0.0

0.1

0.2

0.3

0.4

0.5

G

(

t

)

e

G

(

t

)

d

S

(

t

)

e

S

(

t

)

d

1x

1.6x

2x

3.2x

2x

3.2x

Time to Train w.r.t. SGD

38

39

40

41

42

43

39.82

39.0

39.65

38.76

42.59

39.8

MLP

SGD

SGD+USEFUL

SAM

SAM+USEFUL

SAM+USEFUL+TA

35000

36000

37000

38000

39000

L1 Norm

CIFAR10

98000

100000

102000

104000

106000

108000

CIFAR100

SGD

SGD+USEFUL

SAM

D Additional Results

D.1 An Extreme Case of Toy Datasets

We consider an extreme setting when two features have identical strength and no missing fastlearnable features, i.e., β e = β d = α = 1 , the gap between LHS and RHS in Equation 5 is not small as shown in the figure 8. 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 MLP, USEFUL successfully reduces the test error of both SGD and SAM by nearly 1%. Additionally, SGD+USEFUL yields better performance than SAM alone.

36000

38000

40000

42000

Number of Examples

38883

39675

41392

Forget Score [0-19]

3000

4000

5000

6000

6117

5325

3606

Forget Score [20-39]

SGD

SGD+USEFUL

SAM

D.3 SAM & USEFUL Reduce Forgetting Scores

We used forgetting scores [69] to partition examples in CIFAR10 into different groups. Forgetting scores count the number of times an example is misclassified after being correctly classified during training and is an indicator of the learning-speed of examples. Figure 11 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 2 and 3 illustrate 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, the prediction is always correct) of examples in fast-learnable and slow-learnable clusters. Iteration learned is highly correlated with prediction depth [5], which is another notion of data difficulty. It can be seen clearly that examples in fast-learnable clusters have a lower difficulty score for every metric, indicating that USEFUL successfully identifies fast-learnable examples early in training.

1x

1.15x

2x

2.3x

2x

2.3x

Time to Train w.r.t. SGD

22

24

26

28

30

Test Error (%)

29.92

24.12

23.71

22.45

28.03

24.87

Train ResNet18 from scratch

1x

1.05x

2x

2.1x

2x

2.1x

Time to Train w.r.t. SGD

6

7

8

9

10

11

12

13

12.05

10.41

9.49

8.46

8.41

6.95

Finetune pre-trained ResNet50

SGD

SGD+USEFUL

SAM

SAM+USEFUL

SAM+USEFUL+TA

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 [19]. Following the recommended settings in ASAM, we trained it with a perturbation radius ρ = 1 . 0 , which is 10 times that of SAM. Other settings are identical to the standard settings in Appendix C. Table 4 demonstrates the results for training ResNet18 on CIFAR10. Both SAM and ASAM can be combined with USEFULto improve the test classification error. When using TA, ASAM shows a slightly better performance than SAM.

D.5 USEFUL shows promising results for the OOD settings

While our main contribution is providing a novel and effective method to improve the in-distribution generalization performance, we conducted new experiments confirming the benefits of our method to improving out-of-distribution (OOD) generalization performance. As a few very recent works [9, 72] showed the benefit of SAM in improving OOD performance, it is expected that USEFUL also extend to this setting.

Spurious correlation (Waterbirds). As can be seen in Figure 12, both SAM and USEFUL effectively improve the performance on the balanced test set though the model is trained on the spurious training

set. When training ResNet18 with SGD from scratch, USEFUL decreases the classification errors by 5.8%. In addition, it can be successfully applied to the pre-trained ResNet50 (on ImageNet), opening up a promising application to the transfer learning setting.

Long-tail distribution (Long-tail CIFAR10). We conducted new experiments on long-tail CIFAR10 [15] with an imbalance ratio of 10. Table 6 shows that USEFUL can also improve the performance of SGD and SAM, by reducing the simplicity bias on the long-tail data. Figure 13 visualizes 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 performance, as we confirmed in the last row.

0.0

2.5

5.0

7.5

Class

0

2000

4000

6000

Count

Long-tailed CIFAR10 dataset

Original

Upsample

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 CIFAR10 with 10-20% random label flip successfully boosts the performance of both SGD and SAM when training on data with corrupt labels.

128

256

512

Batch size

4.8

4.9

5.0

5.1

5.2

5.3

Test Error (%)

0.1

0.2

0.4

Learning rate

5.0

5.5

6.0

6.5

7.0

1

2

3

4

5

Upsampling factor

4.80

4.85

4.90

4.95

5.00

5.05

SGD

SGD+USEFUL

0

25

50

75 100125150175200

Epoch

0

20

40

60

80

Train Error

CIFAR10

0

25

50

75 100125150175200

Epoch

0

25

50

75

100

CIFAR100

Epoch 20

(a) Training error trajectories

4

5

6

7

8

9

10

Epoch

4.6

4.8

5.0

5.2

5.4

Test Error (%)

CIFAR10

12

14

16

18

20

22

Epoch

22.0

22.5

23.0

23.5

24.0

24.5

CIFAR100

SGD+USEFUL

SGD

(b) Test classification error

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 CIFAR10. 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 lr ∈ {1e-1, 1e-2, 1e-3, 5e-4, 1e-4}, number of epochs ∈ {1, 2, 4, 8} for training the reference model, and the weight decay ∈ {1e-3, 5e-4, 1e-4} for training GroupDRO. For JTT, we tuned the separating epoch ∈ {4, 5, 6, 7, 8, 10} and upsampling factor ∈ {2, 3, 4, 5}, while lr and weight decay follow the standard training of ResNet18 on CIFAR10. Table 9 shows that all methods successfully reduce the simplicity bias, yielding an improvement over SGD. While EIIL requires tuning 3 hyperparameters with a total of 60 combinations, USEFUL only requires one hyperparameter, which is the separating epoch within a small range (around the time when the slope of the training loss curve diminishes).

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 5, when coupling with either SGD or SAM, USEFUL clearly outperforms upweighting loss on both CIFAR datasets. It 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 is to upsample misclassified examples (MISCLASSIFICATION) 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 at the same epoch t . Table 7 shows that USEFUL selects a 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 is 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 is a standard theoretical assumption [12, 60, 62, 73]. In practice, both for separating fast-learnable vs slowlearnable 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 a larger learning rate does not yield better generalization for the problems we considered in contrast to other settings [46, 58]. As shown in Figure 14 middle, increasing the learning rate has an adverse effect on the model performance. Indeed, for a 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.

Upsampling factor. 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 1. As illustrated in Figure 14 right, while all factors from 2 to 5 bring performance improvement, the upsampling factor of 2 yields the best performance, as it reduces the simplicity bias with minimum change to the rest of the data distribution.

Separating epoch. Fig 15a shows that the early epoch that can best separate the examples is when the change in training error starts to shrink. At this time, examples that are learned can be separated from the rest by clustering model outputs as analyzed in our theoretical results in Section 3.3. Figure 15b demonstrates the performance of USEFUL when separating examples at different epochs in early training. Too early or too late epochs do not cluster examples well, i.e., some examples with fastlearnable 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 our 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-learnable features, thus, a different theoretical time T in Theorems 3.2 and 3.3.

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 be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it 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 a 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 on 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 a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not 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 to address problems of privacy and fairness.

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 acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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 a complete (and correct) proof?

Answer: [Yes]

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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a 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 or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental 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)?

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, a 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 a 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 a 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 it 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 to 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 a 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 submissions 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 to reproduce that algorithm.

(b) If the contribution is primarily a new 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 a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe 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., to 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

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 NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) 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 a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) 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. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 of 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, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

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, confidence 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 are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error 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).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It 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 a 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 of range (e.g. negative error 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 in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

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, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as 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 https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: Our paper conform to every aspect in the NeurIPS 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 a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration 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 our work, none which we feel must be specifically highlighted here.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or 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.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are 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 to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: There are no potential harms to our models.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 to 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 a best faith effort.

12. Licenses 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 use 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 a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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.

If this information is not available online, the authors are encouraged to reach out to the asset'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 new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model 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 as 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 crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, 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 nor 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 NeurIPS 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIt 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\nIt 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.\n\nQUESTIONS:\n1) 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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).\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nYes, the authors have addressed limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n- 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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)\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nPlease see strengths and weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-14"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\nThere 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.\n\nQUESTIONS:\nSee Weakness.\n\nLIMITATIONS:\nN.A.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-20"}