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

Authors: tuan  hai
dang  nguyen
paymon  haddad
eric  gan
baharan  mirzasoleiman
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 23:43:32

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.

Main paper: 
 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. 
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.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. 
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.Notation. We use lowercase letters, lowercase boldface letters, and uppercase boldface letters to denote scalars (a), vectors (v v v), and matrices (W W W ). For a vector v v v, we use ∥v v v∥ 2 to denote its Euclidean norm. Given two sequence {x n } and {y n }, we denotex 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 Õ(•), Ω(•), and Θ(•) to hide logarithmic factors in these notations. Furthermore, we denotex 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 x 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 x x as a collection of P patches: x x x = (x x x (1) , x x x (2) , . . . , x x x (P ) ) ∈ (R d ) P , where • Slow-learnable Feature. One and only one patch is given byβ d • y • v v v d with ∥v v v d ∥ 2 = 1, ⟨v v v e , v v v d ⟩ = 0, and 0 ≤ β d < β e ∈ R.• Fast-learnable feature. One and only one patch is given by β e • y • v v v e with ∥v v v e ∥ 2 = 1 with a probability α ≤ 1. With a probability of 1 -α, this patch is masked, i.e. 0 0 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,EQUATIONwhere w w 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 W = [w w w 1 , . . . , w w 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 W W (0) ∼ N (0, σ 2 0 ), where σ 2 0 = polylog(d)/d.Consider a N -sample training dataset D = {(x x 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 x x; W W W ) readsEQUATIONwhere l is the logistic loss defined as l(z) = log(1 + exp(-z)). The solution W W W ⋆ of the empirical risk minimization (ERM) minimizes the above loss, i.e., W W W ⋆ := arg min W W W L(W W W ).GD. Typically, ERM is solved using gradient descent (GD). The update rule at iteration t of GD with learning rate η > 0 readsEQUATIONSAM. 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 readsEQUATIONwhere ρ (t) = ρ > 0 is the inner step size that is usually normalized by gradient norm, i.e., ρ (t) = ρ/ ∇L(W W W (t) ) F .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 w w (t) j , v v v e ⟩ and ⟨w w w(t) j , v v 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. v v v d : max j∈[J] ⟨w w w (TSAM) j , v v v d ⟩ = Õ(σ 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 w w(t) j , v v v⟩ for v v v ∈ {v v v e , v v v d }indicates better learning of the feature vector v v v by neuron w w 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 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 ]:G (t) e = max j∈[J] ⟨w w w (t) j , v v v e ⟩ and G (t) d = max j∈[J] ⟨w w w (t) j , v v v d ⟩EQUATIONWe 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 SAM than 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:1. The slow-learnable feature has a larger contribution to the normalized gradient of the 1-step SAM update, compared to that of GD.2. Amplifying the strength of the slow-learnable feature increases its contribution to the normalized gradients of GD and SAM. 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.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes]• 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). 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). 
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 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 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. Test Error (%) TrivialAugment (TA) further boosts SAM's performance (except for CINIC10). Remarkably, USE-FUL consistently boosts the performance across all scenarios and achieves (to our knowledge) SOTA performance for ResNet18 and ResNet34 on the selected datasets when combined with SAM and TA. 
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.We further conduct ablation studies on the effect of our data selection strategy for upsampling, training batch size, learning rate, upsampling factor, and separating epoch in Appendix D.8.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.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- 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.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.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. 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. 
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. 
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)? 
 
 Appendix: 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 v v), and matrices (W W W ). For a vector v v v, we use ∥v v v∥ 2 to denote its Euclidean norm. Given two sequence {x n } and {y n }, we denotex 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 Õ(•), Ω(•), and Θ(•) to hide logarithmic factors in these notations. Furthermore, we denotex n = poly(y n ) if x n = O(y D n )for some positive constant D, andx 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 W W (0) ∼ N (0, σ 2 0 ) such that for all j ∈ [J], ⟨w w w (0) j , v v v e ⟩, ⟨w w w (0) j , v v v d ⟩ ≥ ρ > 0.The above assumption is reasonable because we later show that both sequences ⟨w w w (t) j , v v v e ⟩ and ⟨w w w (t) j , v v 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 beEQUATIONLemma 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 W W (t) ) with regard to neuron w w wEQUATIONProof. We have the following gradient∇ w w w (t) j L(W W W (t) ) = - 1 N N i=1 exp(-y i f (x x x i ; W W W (t) )) 1 + exp(-y i f (x x x i ; W W W (t) )) • y i f ′ (x x x i ; W W W (t) ) = - 3 N N i=1 l (t) i y i P p=1 ⟨w w w (t) j , x x x (p) ⟩ 2 • x x x (p) = - 3 N N i=1 l (t) i β 3 d ⟨w w w (t) j , v v v d ⟩ 2 v v v d + β 3 e ⟨w w w (t) j , v v v e ⟩ 2 v v v e + y i ⟨w w w (t) j , ξ ξ ξ i ⟩ 2 ξ ξ ξ i - 3 N N i=αN +1 l (t) i β 3 d ⟨w w w (t) j , v v v d ⟩ 2 v v v d + y i ⟨w w w (t) j , ξ ξ ξ i ⟩ 2 ξ ξ ξ iWith the above formula of gradient, we have the following equations:Fast-learnable feature gradient. The projection of the gradient on v v v e isEQUATIONSlow-learnable feature gradient. The projection of the gradient on v v v d isEQUATIONNoise gradient. The projection of the gradient on ξ ξ ξ i isEQUATIONDerivative of data example i. ForEQUATIONi can be rewritten asEQUATION)while for αNEQUATIONi can be rewritten asEQUATIONNote 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 thatEQUATIONwhich 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 ⟨w w wEQUATIONwhereg 1 (t) = sigmoid J j=1 -β 3 d ⟨w w w (t) j , v v v d ⟩ 3 -β 3 e ⟨w w w (t) j , v v v e ⟩ 3 .Proof. Plugging the update rule of GD, we have⟨w w w (t+1) j , v v v e ⟩ = ⟨w w w (t) j -η∇ w w w (t) j L(W W W (t) ), v v v e ⟩ = ⟨w w w (t) j , v v v e ⟩ + 3ηβ 3 e N αN i=1 l (t) i ⟨w w w (t) j , v v v e ⟩ 2 = ⟨w w w (t) j , v v v e ⟩ + Θ(η)αβ 3 e g 1 (t)⟨w w w (t) j , v v v e ⟩ 2, 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 isEQUATIONwhich givesΘ(η)β 3 d g 1 (t)⟨w w w (t) j , v v v d ⟩ 2 ≤ ⟨w w w (t+1) j , v v v d ⟩ -⟨w w w (t) j , v v v d ⟩ ≤ Θ(η)β 3 d (αg 1 (t) + 1 -α)⟨w w w (t) j , v v v d ⟩ 2 (17)whereg 1 (t) = sigmoid J j=1 -β 3 d ⟨w w w (t) j , v v v d ⟩ 3 -β 3 e ⟨w w w (t) j , v v v e ⟩ 3 .Proof. Plugging the update rule of GD, we have⟨w w w (t+1) j , v v v d ⟩ = ⟨w w w (t) j -η∇ w w w (t) j L(W W W (t) ), v v v d ⟩ = ⟨w w w (t) j , v v v d ⟩ + 3ηβ 3 d N N i=1 l (t) i ⟨w w w (t) j , v v v d ⟩ 2From Lemma B.5, we have forEQUATIONi = Θ(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 w w(T0) j , v v v e ⟩ ≥ Ω(1/β e ).For t ∈ [0, T 0 ], the fast-learnable feature update has the following rule ⟨w w wEQUATIONProof. Let T 0 > 0 be such that either max j∈[J] ⟨w w w(T0) j , v v v e ⟩ ≥ Ω(1/β e ) or max j∈[J] ⟨w w w (T0) j , v v v d ⟩ ≥ Ω(1/β d ).We will show later that the first condition will be met and we have max j∈[J] ⟨w w w(T0) j , v v 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 w w (t) j , v v v d ⟩ 3 -β 3 e ⟨w w w (t) j , v v v e ⟩ 3 .Then, for t ∈ [0, T 0 ], we haveg 1 (t) = 1 1 + exp( J j=1 -β 3 d ⟨w w w (t) j , v v v d ⟩ 3 -β 3 e ⟨w w w (t) j , v v v e ⟩ 3 ) ≥ 1 1 + exp(κ + κ) = 1 1 + exp( Ω(1)),where the first inequality holds due to ⟨w w w(t) j , v v v e ⟩ ≤ κ/(J 1/3 β e )and ⟨w w w(t) j , v v 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 asEQUATIONSimilarly, 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 w w(T0) j , v v v e ⟩ ≥ Ω(1/β e ).For t ∈ [0, T 0 ], the fast-learnable feature update has the following rule ⟨w w wEQUATIONWe next show that GD will learn the fast-learnable feature quicker than learning the slow-learnable feature.Lemma A.7. Assume η = õ(β d σ 0 ). Let T 0 be the iteration number that max j∈[J] ⟨w w w(T0) j , v v v e ⟩ reaches Ω(1/β e ) = Θ(1). Then, we have for all t ≤ T 0 , it holds that max j∈[J] ⟨w w w(T0) j , v v v d ⟩ = Õ(σ 0 ).Proof. Among all the possible indices j ∈ [J] , we focus on the index j ⋆ = arg max j∈[J] ⟨w w w EQUATIONfor all j ∈ [J]. T GD = Θ(1)EQUATIONfor all j ∈ [J] and t ∈ [0, T GD ), we haveEQUATION⟨w w wEQUATIONAfter training for T GD iterations, with high probability, the learned weight has the following properties:(1) it learns the fast-learnable feature v v v e : max j∈[J] ⟨w w w T GD = Θ(1)EQUATIONA.2 Proof of Theorem 3.3Before going into the analysis, we denote the derivative of a data example i at iteration t to beEQUATIONwhere ϵ ϵ ϵ (t) = ρ (t) ∇L(W W W (t) ) is the weighted ascent direction at the current parameter W W W (t) . We denote the weight vector of the j-th filter after being perturbed by SAM asEQUATIONwhereρ (t) = ρ/ ∇L(W W W (t) ) F .First, we have the following inequalities regarding the gradient norm:EQUATIONThus,EQUATIONEQUATIONLemma 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 W W (t) + ϵ ϵ ϵ (t) ) with regard to neuron w w wEQUATIONProof. We have the following gradient∇ w w w (t) j,ϵ ϵ ϵ L(W W W (t) + ϵ ϵ ϵ (t) ) = - 1 N N i=1 exp(-y i f (x x x i ; W W W (t) + ϵ ϵ ϵ (t) )) 1 + exp(-y i f (x x x i ; W W W (t) + ϵ ϵ ϵ (t) )) • y i f ′ (x x x i ; W W W (t) + ϵ ϵ ϵ (t) ) = - 3 N N i=1 l (t) i,ϵ ϵ ϵ y i P p=1 ⟨w w w (t) j,ϵ ϵ ϵ , x x x (p) ⟩ 2 • x x x (p) = - 3 N αN i=1 l (t) i,ϵ ϵ ϵ β 3 d ⟨w w w (t) j,ϵ ϵ ϵ , v v v d ⟩ 2 v v v d + β 3 e ⟨w w w (t) j,ϵ ϵ ϵ , v v v e ⟩ 2 v v v e + y i ⟨w w w (t) j,ϵ ϵ ϵ , ξ ξ ξ i ⟩ 2 ξ ξ ξ i - 3 N N i=αN +1 l (t) i,ϵ ϵ ϵ β 3 d ⟨w w w (t) j,ϵ ϵ ϵ , v v v d ⟩ 2 v v v d + y i ⟨w w w (t) j,ϵ ϵ ϵ , ξ ξ ξ i ⟩ 2 ξ ξ ξ iWith the above formula of gradient, we have the projection of perturbed weight on v v v e is ⟨w w wEQUATIONFrom Equations 32 and 34, we haveEQUATIONSimilarly, the projection of perturbed weight on v v v d is ⟨w w wEQUATIONFast-learnable feature gradient. The projection of the gradient on v v v e is ⟨∇ w w w (t)EQUATIONSlow-learnable feature gradient. The projection of the gradient on v v v d isEQUATIONNoise gradient. The projection of the gradient on ξ ξ ξ i isEQUATIONDerivative of data example i. ForEQUATIONi,ϵ ϵ ϵ can be rewritten asEQUATIONwhile for αNEQUATIONi can be rewritten asEQUATIONWe 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 isEQUATIONwhereg 2 (t) = sigmoid J j=1 -β 3 d ⟨w w w (t) j,ϵ ϵ ϵ , v v v d ⟩ 3 -β 3 e ⟨w w w (t) j,ϵ ϵ ϵ , v v v e ⟩ 3 .Proof. Plugging the update rule of SAM, we have⟨w w w (t+1) j , v v v e ⟩ = ⟨w w w (t) j -η∇ w w w (t) j,ϵ ϵ ϵ L(W W W (t) + ϵ ϵ ϵ (t) ), v v v e ⟩ = ⟨w w w (t) j , v v v e ⟩ + 3ηαβ 3 e N N i=1 l (t) i,ϵ ϵ ϵ ⟨w w w (t) j,ϵ ϵ ϵ , v v v e ⟩ 2 = ⟨w w w (t) j , v v v e ⟩ + Θ(η)αβ 3 e g 2 (t)⟨w w w (t) j,ϵ ϵ ϵ , v v v e ⟩ 2, 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 isEQUATIONwhereg 2 (t) = sigmoid J j=1 -β 3 d ⟨w w w (t) j,ϵ ϵ ϵ , v v v d ⟩ 3 -β 3 e ⟨w w w (t) j,ϵ ϵ ϵ , v v v e ⟩ 3 .Proof. Plugging the update rule of GD, we have⟨w w w (t+1) j , v v v d ⟩ = ⟨w w w (t) j -η∇ w w w (t) j,ϵ ϵ ϵ L(W W W (t) ), v v v d ⟩ = ⟨w w w (t) j , v v v d ⟩ + 3ηβ 3 d N N i=1 l (t) i ⟨w w w (t) j,ϵ ϵ ϵ , v v v d ⟩ 2From Lemma B.5, we have forEQUATIONi = Θ(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 w w (T0) j , v v v e ⟩ ≥ Ω(1/β e ). For t ∈ [0, T 0 ], the fast-learnable feature update has the following ruleEQUATIONProof. Let T 0 > 0 be such that either max j∈[J] ⟨w w w(T0) j , v v v e ⟩ ≥ Ω(1/β e ) or max j∈[J] ⟨w w w (T0) j , v v v d ⟩ ≥ Ω(1/β d ).We will show later that the first condition will be met and we have max j∈[J] ⟨w w w(T0) j , v v 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 w w (t) j,ϵ ϵ ϵ , v v v d ⟩ 3 -β 3 e ⟨w w w (t) j,ϵ ϵ ϵ , v v v e ⟩ 3 .Then, for t ∈ [0, T 0 ], we haveg 2 (t) = 1 1 + exp( J j=1 -β 3 d ⟨w w w (t) j,ϵ ϵ ϵ , v v v d ⟩ 3 -β 3 e ⟨w w w (t) j,ϵ ϵ ϵ , v v v e ⟩ 3 ) ≥ 1 1 + exp(κ + κ) = 1 1 + exp( Ω(1)),where the first inequality holds due to ⟨w w w(t) j,ϵ ϵ ϵ , v v v e ⟩ ≤ ⟨w w w (t) j , v v v e ⟩ ≤ κ/(J 1/3 β e )and ⟨w w w(t) j,ϵ ϵ ϵ , v v v d ⟩ ≤ ⟨w w w (t) j , v v 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 w w (T0) j , v v v e ⟩ ≥ Ω(1/β e ). For t ∈ [0, T 0 ], the fast-learnable feature update has the following ruleEQUATIONWe next show that SAM will learn the fast-learnable feature quicker than the slow-learnable one.Lemma A.14. Assume η = õ(β d σ 0 ). Let T 0 be the iteration number that max j∈[J] ⟨w w w(T0) j , v v v e ⟩reaches Ω(1/β e ) = Θ (1). Then, we have for all t ≤ T 0 , it holds that max j∈[J] ⟨w w w(T0) j , v v v d ⟩ = Õ(σ 0 ).Proof. Among all the possible indices j ∈ [J] , we focus on the index j ⋆ = arg max j∈[J] ⟨w w w EQUATIONfor all j ∈ [J]. EQUATIONfor all j ∈ [J] and t ∈ [0, T SAM ), we haveEQUATION⟨w w wEQUATIONAfter training for T SAM iterations, with high probability, the learned weight has the following properties: (1) it learns the fast-learnable feature v v v e : max j∈[J] ⟨w w w EQUATIONComparing 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 inequalityEQUATIONwhich holds due to Assumption A.1 about weight initialization in Appendix A.1, i.e., (σ 0 ≥ ρ ≥ 0).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 w w (t) j , v v v e ⟩EQUATIONBefore introducing the theorem, we assume that the model is initialized in favor of the fastlearnable feature, i.e. G EQUATION)EQUATION)EQUATIONEQUATIONProof. 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 Ŝ(0) e - Ŝ(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.EQUATIONEQUATIONEQUATIONEQUATIONWe need to prove that they also hold for t + 1. From Lemma A.15 and the first two induction hypotheses, ρ + S (t+1) d = ρ + S (t) d + Θ(η)β 3 d ( Ŝ(t) d ) 2 < S (t) e + Θ(η)αβEQUATIONSimilarly, we haveS (t+1) d ≤ G (t+1) d. From Equations 53 and 54,EQUATION)EQUATIONCombining with Equations 35 and 37, we haveEQUATIONG (t) d -S (t) d + Θ(η)β 3 d ((G (t) d ) 2 -( Ŝ(t) d ) 2 ) < G (t) e -S (t) e + Θ(η)αβ 3 e ((G (t) e ) 2 -( Ŝ(t) e ) 2 ) (72) G (t+1) d -S (t+1) d < G (t+1) e -S (t+1) e (73)S (t+1) e -S (t+1) d < G (t+1) e -G (t+1) d (74)Therefore, the induction hypotheses hold for t + 1.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 x 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 W 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.EQUATION)EQUATIONBecause the noise is sufficiently small, the above equations can be simplified asEQUATION∇ w w w (t)ϵ ϵ ϵ L(W W W (t) + ϵ ϵ ϵ (t) ) = - 3 N N i=1 l (t) i,ϵ ϵ ϵ β 3 d ⟨w w w (t) ϵ ϵ ϵ , v v v d ⟩ 2 v v v d - 3 N αN i=1 l (t) i,ϵ ϵ ϵ β 3 e ⟨w w w (t) ϵ ϵ ϵ , v v v e ⟩ 2 v v v e (78)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 asEQUATION∇ w w w (t)EQUATIONBoth 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 thatEQUATION⟨w w wEQUATION)EQUATION)EQUATIONThe last inequality holds due to αβ 3 e > β 3 d and ⟨w w w (t) , v v v e ⟩ ≥ ⟨w w w (t) , v v v d ⟩ from Theorem A. 16 .From Equation 81in 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 x 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 W 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 1. The normalized gradient of the one-step SAM update has a larger weight on the slowlearnable feature compared to that of GD.2. Amplifying the slow-learnable feature strength puts a larger weight on the slow-learnable feature in the normalized gradients of GD and SAM.3. There exists an upsampling factor k such that the normalized gradient of the one-step GD update on D(β 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 showEQUATIONwhich 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β 3 d ⟨w w w (t) ϵ ϵ ϵ , v v v d ⟩ 2 β 3 e ⟨w w w (t) ϵ ϵ ϵ , v v v e ⟩ 2 = (β ′ d ) 3 ⟨w w w (t) , v v v d ⟩ 2 β 3 e ⟨w w w (t) , v v v e ⟩ 2(87)EQUATION)EQUATIONEQUATIONEQUATIONTherefore, with β ′ d = 1-3ρ (t) β 3 d ⟨w w w (t) ,v v v d ⟩ 1-3ρ (t) αβ 3 e ⟨w w w (t) ,v v ve⟩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]:EQUATIONProof. 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 haveEQUATIONThese imply that t∈[Tg,Tg+1)EQUATION)Recall b is the smallest integer such thatEQUATIONThus, 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 satisfyx t+1 ≥ x t + η • C t (x t -ρ) 2 , y t+1 ≤ y t + Sη • C t y 2 t , for some C t = Θ(1). Suppose x 0 ≥ y 0 S 1+2G 1-3G where S ∈ (0, 1), G ∈ (0, 1/3) and 0 < η ≤ min{ G 2 x0 log(A/x0) , G 2 y0 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 Tx ≤ 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 haveEQUATION)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 haveEQUATION)Compare Equation 107 and 110. Choosing δ = G and η ≤ min{G 2 x0 log(A/x0) , G 2 y0 log(1/G) }, together with x 0 ≥ y 0 S 1+2G 1-3G 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 recursionsz t+1 ≥ z t + m(z t -ρ) 2 , z t+1 ≤ z t + M (z t ) 2 ,where z 0 > ρ ≥ 0 is the initialization and m, M > 0 are some constants. Let v > z 0 , then the timeEQUATIONProof. 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:EQUATIONBecause z t ≥ z 0 , we obtainEQUATIONNow, we want to bound z T1 -z 0 . Using again the recursion and z T1-1 ≤ 2z 0 , we haveEQUATIONCombining Equation 113 and 114, we get a bound on T 1 asEQUATIONNow, 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 haveEQUATIONEQUATIONOn the other hand, by using z Tn-1 ≤ 2 n z 0 we upper bound z Tn asEQUATIONBesides, we know that z Tn-1 ≥ 2 n-1 z 0 . Therefore, we upper bound z Tn -z Tn-1 asEQUATION)Combining Equations 116 and 119 yieldsEQUATIONSumming Equation 120 for n = 2, . . . , n we haveEQUATIONLastly, we know that 2 n z 0 ≥ v this implies that we can set n = log(v/z0) 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],|⟨w w wEQUATIONLemma B.5 (Lemma G.4, [18] ). For every i, we have l(t) i = Θ(1)g 1 (t),whereEQUATIONLemma 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 max |⟨w w w (0) j,ϵ ϵ ϵ , ξ ξ ξ i ⟩| = Õ(σσ 0 √ d) max i 1 n n k=1 |⟨ξ ξ ξ k , ξ ξ ξ i ⟩| = Õ( σ 2 d N + σ 2 √ d)Proof. For the first inequality, Lemma B.7 implies that with probability at least 1 -1 dN ,EQUATIONTaking 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],|⟨w w wEQUATIONProof. Let χ t = max{|⟨w w 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χ t+1 ≤ χ t + 3ηαχ 2 tSuppose that a(t) satisfies the differential equationa ′ = 3αηa 2 a(0) = χ 0Observe that a(t) is increasing so, by the Mean Value Theorem there exists τ ∈ (t, t + 1) such thata(t + 1) -a(t) = a ′ (τ ) = 3αηa(τ ) 2 ≥ 3αηa(t) 2So an fast-learnable induction shows that a(t) ≥ χ t . Now solving for a(t),a(t) = 1 1 a(0) -3αηt , t ≤ 1 3αηa(0).Using the high probability tail bounds B.8,a(0) = χ 0 = Õ(σσ 0 √ d) α = Õ( σ 2 d N + σ 2 √ d)where σ = σp √ d . Substituting these bounds givesa(t) = 1 Ω( 1 σσ0 √ d ) -Õ(ηt( σ 2 d N + σ 2 √ d))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, soa(T 0 ) = Õ(σσ 0 √ d).We conclude that |⟨w w wEQUATIONLemma B.10. For every i, we have lEQUATIONEQUATION)The proof is the same as [18] [Lemma G.4]. 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.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. 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.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.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 Long-tail distribution (Long-tail CIFAR10). We conducted new experiments on long-tail CI-FAR10 [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. 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. 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- 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. Justification: Our claims match the theoretical and experimental resultsGuidelines:• 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.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.38th Conference on Neural Information Processing Systems (NeurIPS 2024). 


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

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

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

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

Reviews: review: paper_summary: paper_summary: This paper proposes an algorithm for changing the distribution of training data to improve the generalization of the model on origin data distribution. The paper is inspired by Sharpness Aware Minimization, which aims at finding a flat minimum meaning that it has a good generalization capability. This paper divides features into two categories: fast-learnable features and slow-learnable features and derives some observations like "SGD and SAM only learn fast-learnable or easy features early in training" and "SAM learns slow-learnable and fast-learnable features at a more uniform speed". The authors propose the method dubbed as USEFUL to train the model on some slow-learnable features repeatedly. The experiments show the effectiveness of USEFUL on CIFAR10 and CIFAR100 datasets.
strength_weakness: strength_weakness: There are some questions based on the presentation of this paper, I will not hesitate to improve my score if the following question are solved.
- Difference between this paper and methods for long-tailed data distribution or measuring the difficulty of learning examples. Algorithms for long-tailed data distribution are usually based on resampling training data or reweighing loss value. The proposed USEFUL is similar to the resampling methods except that USEFUL focuses on the features that are hard/slow to learn. Some references for understanding: [Shi, Jiang-Xin, et al. "How re-sampling helps for long-tail learning?." Advances in Neural Information Processing Systems 36 (2023).](https://arxiv.org/pdf/2310.18236), [Shrivastava, Abhinav, Abhinav Gupta, and Ross Girshick. "Training region-based object detectors with online hard example mining." Proceedings of the IEEE conference on computer vision and pattern recognition. 2016.](https://arxiv.org/pdf/1604.03540v1) and some references based on it, [A Re-Balancing Strategy for Class-Imbalanced Classification Based on Instance Difficulty](https://openaccess.thecvf.com/content/CVPR2022/papers/Yu_A_Re-Balancing_Strategy_for_Class-Imbalanced_Classification_Based_on_Instance_Difficulty_CVPR_2022_paper.pdf), [Active Teacher for Semi-Supervised Object Detection](https://openaccess.thecvf.com/content/CVPR2022/papers/Mi_Active_Teacher_for_Semi-Supervised_Object_Detection_CVPR_2022_paper.pdf), I believe a discussion of these references in paper should be helpful.
- The relation between the proposed USEFUL and SAM? It seems like the motivation of USEFUL is changing the data distribution to get a flat minimum like SAM. But the results in Appendix D.2, *i.e.*, 53.8 for SGD 41.8 for SGD+USEFUL 12.4 for SAM in Table 1($\lambda_{max})$, do not show effectiveness compared with SAM. It could show the effectiveness on SGD but it's far from being comparable to SAM. 
Some small questions:
- What's the exact formulation of the Data distribution?
- What's the "patch" meaning in Definition 3.1? Is that the same as the patch in ViT or the channel of the image? It's a little confusing.
- The experiments mainly focus on traditional architecture, e.g., n-Layer CNN, ResNet. More experiments on popular models and big datasets, e.g., Transformer ImageNet-1k, would be better.
questions: questions: See Weakness.
limitations: limitations: N.A.
paper_summary: paper_summary: - Proves for a 2-layer CNN with fixed second layer weighsts trained on a toy dataset,  SAM learns slow-learnable and fast-learnable features more uniformly in the early epochs compared to SGD
- Based on this analysis, proposes a simple clustering-based upsampling strategy for reducing simplicity bias / excessive reliance on fast-learnable features. The results show that this improves in-distribution generalization of standard small-scale image classification tasks.
strength_weakness: strength_weakness: - No baselines. There are several papers now that try to reduce simplicity bias in order to improve performance:
    - https://arxiv.org/abs/2105.05612
    - https://arxiv.org/abs/2301.13293
    - https://arxiv.org/abs/2107.09044 (does not focus on simplicity bias explicitly, but similar to method proposed in paper)
    - simpler baselines: there are several papers that propose “example difficulty” metrics (https://arxiv.org/abs/2106.09647). How well do this correlate with the clusters found in your method? If you just train on the k examples with the highest difficulty scores (per class), does this fare worse than the proposed method?
- Limited novelty due to findings in [64] (Sharpness-aware minimization enhances feature quality via balanced learning). This paper also shows that SAM improves feature diversity (on real datasets + backed up with analysis on a toy dataset) and improves performance on transfer-learning tasks.
- Lacking discussion about when this method would fail. I can imagine two scenarios where the method would not work:
    1. Most training examples have one or more slow-learnable features. In this case, the clustering approach would “remove” most of the points in the dataset, and train on very points for multiple epochs. This could result in overfitting and performance that is worse than training. There’s an implicit assumption that there is some sort of one-to-one relation between examples and features. In the case where all examples contain an “easy” (e.g. patch) and a “hard” feature (e.g. CIFAR), would this method improve performance over SGD? 
    2. In noisy datasets, low-quality examples or mislabeled examples would require more time to learn, and this method would cluster them and train on them for longer. That is, it would group examples that are “high-quality” and hard-to-learn with “low-quality” points. In this case, would the proposed method improve performance over SGD? 
- “SB of SGD has been long conjectured to be the reason for the superior generalization performance of overparameterized models, by providing capacity control or implicit regularization” This incorrectly cites https://arxiv.org/abs/2006.07710v2, which shows that too much simplicity bias can lead to robustness and in-distribution generalization issues.
- Unfair evaluation. The experiments compare SAM+TA augmentation and SAM+USEFUL+TA to SGD (no TA). I think there should be two plots, complaring {SGD, SAM, SAM+USEFUL} w/ and w/o TA.
- Experiments on larger datasets. The image classification used here are fairly small-scale. I would like to how well this method scales to ImageNet-scale datasets (TinyImageNet is not a good proxy..)
- Writing is repetitive at times, especially the theory section (3.3)
questions: questions: - How specific is the analysis to the toy distribution setting in which there is just 1 slow-learnable and 1 fast-learnable feature? How do things change if the number of slow-learnable features >> number of fast-learnable features?
- Why is max alignment between weight vector and ground-truth feature v_e the right way to evaluate the feature’s contribution to the model? Isn’t it hypothetically possible that SAM solutions rely more on the simpler feature if more weight vectors rely on v_e instead of v_d, even when max-alignment with slow feature is higher for SAM? Some discussion connecting this metric to “feature reliance” vis-a-vis model outputs would be great.
limitations: limitations: Please see strengths and weaknesses.
paper_summary: paper_summary: This work aims to modify the training data distribution to improve in-distribution generalization. First, the authors theoretically analyse a 2-layer CNN and compare the feature learning dynamics (fast learnable and slow-learnable features) of Gradient Descent (GD) and Sharpness-Aware Minimization (SAM). It is then shown that SAM mitigates simplicity bias compared to GD. The authors then propose USEFUL (UpSample Early For Uniform Learning), a method that upsamples the examples in the training set that contains slow-learnable features.  USEFUL first clusters the examples with similar outputs early in the training and then upsamples the slow-learnable clusters. The main idea behind USEFUL is to learn features at a uniform speed (similar to SAM) by changing the training data distribution. USEFUL can be trained with SGD, SAM and SAM + Trivial Augment. Results on CIFAR-10, CIFAR-100, STL10, TinyImageNet indicate that USEFUL is across datasets and architectures. Additonal ablation and analysis show that USEFUL learns similar properties to SAM (for e.g less sharp solutions).
strength_weakness: strength_weakness: 1. The authors explicitly mention that their focus in this paper is only on “in-distribution generalization”. I am a bit confused by this given the motivation of simplicity bias and learning features uniformly. To elaborate more on this point,
    - Springer et al [1] also show that SAM implicitly balances the quality of diverse features (similar to the observations made in Section 3 of this paper. The experimental results in [1] is focused more on datasets with multiple predictive features like CelebA, CIFAR-MNIST. 
    - Past work on simplicity bias and shortcut learning [2, 3, 4, 5] has focused on similar datasets like CelebA, Waterbirds, CIFAR-MNIST, Colored-MNIST to name a few.
    - While the authors have shown encouraging results on Waterbirds dataset in Appendix D5, it would be good to show the complete results on various groups and on other datasets as well.
2. Connection to [1]. Springer et al [1] made a very similar observation as to Section 3 in this paper. It would be great if the authors can clarify the differences with the observations in [1] and this work. Particularly, [1] also shows that SAM mitigates simplicity bias and that  SAM learns higher quality representations of hard-to-learn features. The authors briefly discuss this in Related Works section but a more detailed answer would be helpful.
3. I just wanted to understand the practical usefulness of the proposed method. This method has one additional hyperparameter i.e the separating epoch. The authors have reported the best separating epoch for all the datasets which is epoch 8 for CIFAR-10 and epoch 20 for CIFAR-100 (Appendix C.2). How is this hyperparameter chosen? Is there a separating epoch number that works across various datasets? This is especially relevant given that that the average gain on most of the datasets with USEFUL is less than 1% with additional cost for training. 
    
 [1] Springer, Jacob Mitchell, Vaishnavh Nagarajan, and Aditi Raghunathan. "Sharpness-Aware Minimization Enhances Feature Quality via Balanced Learning." The Twelfth International Conference on Learning Representations.

[2] Shah, Harshay, et al. "The pitfalls of simplicity bias in neural networks." Advances in Neural Information Processing Systems 33 (2020): 9573-9585.
    
 [3] Geirhos, Robert, et al. "Shortcut learning in deep neural networks." Nature Machine Intelligence 2.11 (2020): 665-673.
    
 [4] Kirichenko, Polina, Pavel Izmailov, and Andrew Gordon Wilson. "Last layer re-training is sufficient for robustness to spurious correlations." arXiv preprint arXiv:2204.02937 (2022).
    
 [5] Teney, Damien, et al. "Evading the simplicity bias: Training a diverse set of models discovers solutions with superior ood generalization." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.
questions: questions: 1. I did not find the separating epoch used for most of the datasets (except CIFAR-10 and CIFAR-100). Could you please point me to that?
2. What is the takeaway from Figure 1? Please mention the observations regarding the Figure involving the fast and slow learnable features. What kind of examples are usually clustered in fast learnable cluster vs slow learnable cluster,

Please refer to the Weakness section for remaining questions.
limitations: limitations: Yes, the authors have addressed limitations.
paper_summary: paper_summary: It is known that usually deep neural networks will learn “easy examples" that contain fast-learnable features first while learning more complex examples in a second time. The authors argue that mitigating such simplicity bias is the reason method like SAM are outperforming SGD. Based on such analysis, the authors introduce their methods coined as USEFUL that consists in two setups: 1) Identifying the examples with fast-learnable features using a clustering method based on layer output similarity 2) Upsampling by a constant factor the remaining examples with slow-learning features. By doing so, the authors can significantly increase model performances and training time on different classification tasks using different optimizers. They assess their methods across a wide range of dataset and different hyper-parameters and outperform random clustering baseline.
strength_weakness: strength_weakness: It is not clear when and why one should choose the last output activation vector to define the clustering instead of intermediate activation vector. It is also not clear at which epoch one should decide to do the clustering since for a dataset like CIFAR10 the optimal performances are achieved at epoch 8 while for CIFAR100 it is epoch 20. So, finding the correct hyper-parameters for the clustering might be costly and thus impact how fast convergence can really be (if we consider this needed additional ablation on clustering epoch). In addition, the authors mention that they are using an upscaling factor of 2, but I am wondering how robust this is when using long-tail distribution. For example, I am not sure that on something like ImageNet-LT or Inaturalist, we will get the best performances by using a constant factor.  I would also be a bit more cautious about some of the claims made in the papers. For example, the authors claim that their method is generalizing to OOD tasks while providing experiments on only the WaterBird dataset.  So, it would be better to write about promising preliminary results than claiming generalization on OOD.
questions: questions: 1) Do you think your method will also generalize on Long-Tail dataset while keeping an upscaling factor constant? 

2)) Any ideas or heuristics about how to find the optimal epoch/layer to perform the clustering without running an expensive ablation?
limitations: limitations: The authors did not really discuss any limitations (outside the fact that their theoretical result does not extend to CNN) or societal impact. I think that one limitation that could have been highlighted is the smaller scale of the experiment and the focus on classification tasks. Another limitation is the lack of results on OOD or long-tail benchmarks which would seem to be well suited for this type of work.

Comments: title: Response to Reviewer
Your feedback is greatly appreciated
Your feedback is greatly appreciated
Your feedback is greatly appreciated
Response to Reviewer
Response to Rebuttal
comment: Thank you for reading our rebuttal and keeping a positive view of our paper. We are glad that our rebuttal, especially the new long-tail experiment, has addressed your concerns. We assure that we will modify our language in the revised version.
I would like to thank the authors for addressing my concerns. I really appreciate the new long tail experiment and the fact that the authors will revise some of the paper's language. I am maintaining my score.
We hope our rebuttal has effectively addressed your concerns. As the discussion phase is nearing its conclusion, we are wondering if there's anything further we can clarify or elaborate on. Your feedback is greatly appreciated.
We hope our rebuttal has effectively addressed your concerns. As the discussion phase is nearing its conclusion, we are wondering if there's anything further we can clarify or elaborate on. Your feedback is greatly appreciated.
We hope our rebuttal has effectively addressed your concerns. As the discussion phase is nearing its conclusion, we are wondering if there's anything further we can clarify or elaborate on. Your feedback is greatly appreciated.
Thank you for reading our rebuttal, and we’re glad that you found our findings interesting and striking. We also believe that this is an important result and we are quite excited about it.

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

> reducing the simplicity bias benefits in-distribution (ID).

This is an interesting result and a bold claim.
Does this mean that past methods that focus on mitigating simplicity bias (and show results in the OOD scenarios) also improve the performance in IID setup by similar margins? The additional simplicity bias (JTT and EIIL) baselines discussed in the rebuttal pdf seem to indicate so. I would love to hear the authors' comments on this. In the future, the authors can consider running experiments with a couple of other methods that mitigate simplicity to verify the generality of the claim.

Most of my concerns have been resolved and thus, I have increased my score to 6.