Sharpness-Aware Minimization Enhances Feature Quality via Balanced Learning
Authors: Jacob Mitchell Springer
Vaishnavh Nagarajan
Aditi Raghunathan
Venue: ICLR 2024 poster (2024)
Keywords: sharpness-aware minimization
representation learning
spurious correlations
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Abstract: Sharpness-Aware Minimization (SAM) has emerged as a promising alternative optimizer to stochastic gradient descent (SGD). The originally-proposed motivation behind SAM was to bias neural networks towards flatter minima that are believed to generalize better. However, recent studies have shown conflicting evidence on the relationship between flatness and generalization, suggesting that flatness does fully explain SAM's success. Sidestepping this debate, we identify an orthogonal effect of SAM that is beneficial out-of-distribution: we argue that SAM implicitly balances the quality of diverse features. SAM achieves this effect by adaptively suppressing well-learned features which gives remaining features opportunity to be learned. We show that this mechanism is beneficial in datasets that contain redundant or spurious features where SGD falls for the simplicity bias and would not otherwise learn all available features. Our insights are supported by experiments on real data: we demonstrate that SAM improves the quality of features in datasets containing redundant or spurious features, including CelebA, Waterbirds, CIFAR-MNIST, and DomainBed.

Published as a conference paper at ICLR 2024
SHARPNESS-AWARY MINIMIZATION ENHANCES FEATURE QUALITY VIA BALANCED LEARNING
Jacob Mitchell Springer', Vaishnavh Nagarajan² & Aditi Raghunathan' Carnegie Mellon University GGoogle Research
{jspringer, radit ( (eccu.eee vaishnavie@google.c com?
ABSTRACT
Sharpness-Awan Minimization (SAM) has emerged as promising alternative op- timizer to stochastic gradient descent (SGD). The tiginalll-proposse motivation behind SAM was to bias neural networks towards flatter minima that are believed to generalize better. However, recent studies have shown conflicting evidence on the relationship between flatness and generalization, suggesting that flatness does fully explain SAM's success. Sidestepping this debate, we identify an orthogonal effect of SAM that is beneficial -ittiiiiiiiiiiioonn argue that SAM implicitly balances the quality of diverse features. SAM achieves this effect by adaptively suppressing well-learned features which gives remaining features opportunity be learned. We show that this mechanism is beneficial in datasets that contain redundant or spurious features where SGD falls for the simplicity bias and would not otherwise learn all available features. Our insights are supported by experi- ments on real data: we demonstrate that SAM improves the quality of features datasets containing redundant or spurious features, including CelebA, Waterbirds, CIFAR-MNIST, and DomainBed.
INTRODUCTION
Cemmpness-Aware Minimization (SAM) has emerged as compelling alternative to stochastic gra- dient descent (SGD) as an optimizer for neural networks (Bahri et al. 2021]. The core motivation of SAM is to not just to minimize the original training objective, but to also find minimum with small Hessian norm (Wen et al.. 2022). This motivation has been inspired by conventional wisdom that "flat" minima may correspond to better in in-distribution generalization comparison to sharper minima (Keskar et al. 2016 Jiang al. 2019) (Foret et al.) 2020). However, there has been a long- standing debate on whether this flatness is even a reliable predictor of generalization (Dinh et al.) 2017 Kaddour et al. 2022 Andriushchenko al.) 2023b Wen et al. 2023 Kaur al.) (023)). To this end, alternative theories about SAM have been put forth to explain its performance gains, typ- ically by studying specific inductive biases of SAM in simple theoretical settings (Andriushchenko & Flammarion 2022, Andriushchenko et al., 2023a Behdin & Mazumder, (2023).
In this paper, we provide complementary perspective the discussion relating flatness to SAM's in-distribution gains. We argue that SAM can lead to certain owl-of-distribution gains via a feature diversifying effect. Specifically, we consider datasets with multiple redundant features (e.g., an image of a cat may be recognized by its fur or eyes or tail, or of truck by its wheels or its front or side doors). In such datasets, SGD is known to suffer from a "simplicity bias" towards learning only the simplest features that suffice for the task, over range of datasets (Shah et al. 2020). However, we show that SAM overcomes this bias to learn multiple highly correlated features. Remarkably, this happens even though the SAM objective was not explicitly designed with this anti-simplicity, diversity bias. This bias makes SAM favorable in downstream tasks with distributed shifts where a diverse set of features become relevant for prediction.
Why does SAM lead this diversif-iiiriifiinngg effect, even when learning only one of the features would suffice to minimize the objective? We construct minimal toy setting with a simple and a complex feature to address this question. In the toy setting, SGD fails to capture an informative representation of the complex feature due to the simplicity bias. By contrast, SAM is able to learn a high-quality representation of both the simple and the complex feature. We argue that the key difference between the two algorithms can be seen directly in their respective gradient update steps. By decomposing the SAM gradient update step into two terms, we are able to identify two separate effects of SAM that lead to the feature-diversifyinn effect. The first effect is a balancing of the

---

Published as a conference paper at ICLR 2024
weight of the training examples during each update step, leading to more uniform update across the dataset. This effect is similar to other methods that explicitly re-weight examples to promote complex-feature learning such as GroupDRO (Sagawa et al. 2019) and Just Train Twice (Liu al. 2021]. Unlike these methods, SAM re-weights points implicitly through its update step. The second effect is a balancing of the effective step size taken to learn each feature: the effective learning rates of well-learned features are suppressed. These effects lead to balanced learning of all features, and thus more diverse representation. SGD, by contrast, does not exhibit these balancing effects.
The feature diversity induced by SAM is particularly important in the setting in which the features which are redundant in the training distribution are not all necessarily predictive of the label in downstream distribution. For example, the Waterbirds dataset (Sagawa et al. 2019) consists of images of water birds and land birds on backgrounds of water and land. In the traming distribution, the background and bird type are highly correlated, but downstream, the background is not predictive of the bird type and vice versa. Ideally, we learn a representation amenable to adaptation to predict bird type, or background. We find that SAM indeed improves the performance of transfer learning with last-layer retraining on the downstream Waterbirds tasks, and on other similar datasets, include CelebA, CIFAR-MNIST. and DomainBed.
Our perspective is complementary to the large body of existing literature that discusses the in- distribution improvements of SAM and variants (Foret et al. 20200 Na et al.
2022; Zhang et al.)
2023 Zhao et al. 2022; Rangwani et al.) [022; Wang et al. 2022, Meng et al. 2023). Importantly,
this finding allows us to contextualize the small set of prior work that has discovered that SAM helps -ut-of-distribution performance (Wang et al. 2023) Cha et al., 2021 Bahri et al.] 2021). We hope that this work will inspire future work to further understand the mechanisms of SAM and will mo- tivate new algorithms for improving downstream performance on out-of-distribution data in similar contexts.
In summary, our contributions are as follows:
1. We identify a feature diversifying effect of SAM in settings in which data has multiple redundant features, and verify that this effect improves feature representation quality.
2. We construct and analyze a minimal toy setting consisting of multiple redundant features where SGD fails to capture an informative representation of all features, but SAM succeeds.
3. We demonstrate that the feature diversifying effect of SAM arises from two interpretable bal- ancing effects: a balancing effect that reweights the training examples be more uniform and balancing effect that suppresses the effective learning rate of well-learned features.
4. We verify that the improved feature representations can be used to improve performance on out- of-distribution downstream tasks in realistic settings.
2 RELATED WORK
Understanding Sharpness-Award Minimization. The success of SAM is commonly explained by the connection between the flatness of the landscape and improved generalization (see Section 5 of Foret et al. (202))) The connection continues to inspire flatness latness-optimizing objectives for wide range of tasks (Abbas et [022 Na al. [022]] Wang et al. al.] [2023] Zhang et al. 2023 Zhao al.) 2022 Rangwani et al. 2022 Wang et al. 2022,] Meng et al. 2023).
However, recent work has questioned whether sharpness and generalization are as linked as pre- viously thought. On the one hand, it appears that sharper minima can generalize well after all, and on the other, even flat minima may sometimes generalize poorly (Dinh et al. 2017 Kaddour let al. 2022; Andriushchenko al. 20236] Wen al. 2023 Kaur al. 2023). In this backdrop, some works on understanding SAM have departed from flatness. These explicitly characterize the implicit bias of SAM in simple theoretical settings show that it can improve in-distribution gen-
eralization (Andriushchenko & Flammarion 2022 Behdin & Mazumder 2023). Our work takes a different approach departing from both flatness and deriving the implicit regularization by looking at how the SAM perturbation changes each gradient step along the trajectory. Further, while these works focus on in-distribution generalization, we focus on the fivture-iierrsffying effect of SAM in and transfer learning settings.
Benefits of SAM beyond in-distribution generalization. SAM has been shown to improve per- formance in across various other settings such as meta learning (Abbas et al. 2022). domain gen- eralization (Wang et al. 2023 Cha et al. 2021), label noise (Foret al. 2020), transfer learning
2

---

Published as a conference paper at ICLR 2024
in language models (Bahri et al. 2021). Our work augments these findings by offering one unify- ing factor that can explain SAM's gains beyond in-distribution generalization SAM learns higher quality representations of hard-to-leamm features.
Feature learning in the presence of multiple predictive features. It has been shown that neural networks learn some features more easily than others, tendency that is dubbed as a "simplicity bias" (Shah et al., 2020, Kalimeris et al., [2019] Morwani et al. 2023; Rahaman et al.] 2019). Besides different learning speeds, the presence of one feature also impacts the quality (i.e. probing error) of another feature in the learned representation as shown in (Pezeshki et al. (02T). This has inspired various fixes such as regularizers (Pezeshki et al. 202T). selecting freezing of parameters (Ye al.) 2023) and data augmentation (Plumb et al. (022T) This problem has also been addressed in a line of work aiming to improve worst-group error in the presence of spurious correlations, notably Sagawa et al. (2019), but we defer detailed discussion to Appendix [A] In contrast to these algorithms that are explicitly designed to address the simplicity bias, we show that SAM does so without being explicitly incentivized do so.
We elaborate on further related work including the broader literature on shortcut learning and spuri- ous correlations, feature diversity and finetuning, and variants of the SAM algorithm in Appendix
3 SETUP AND PRELIMINARIES
Task. We consider the setting of classification where we map input X to output y y Given training data (21, !n ), (I2, U2), (Xn, yn) where Ti, y/ ~ Diminn our goal is to learn a neural network classifier : X + V. We are interested in the diversity of the representations used by classifier. Hence it is convenient to parameterize the neural network classifiers f as linear classifiers on top of feature representations, i.e. f(u, 0: x) = arg max(u" (z))), where v E R and bo Rd + R* Note that the feature map D is itself neural network. We use w (v, to denote the weights of the network when we do not need to discuss the features explicitly.
Multiple predictive features. In our setting, there several predictive features. As example, inspired by Sagawa et al. (2019), we consider the CelebA inputs as having two redundant features: gender feature { malee female' and hair color {blond, dark), both, in part, predictive of the label. Formally, each input æ has a set of features a1 (z), a2 (2), am (x) where each a; (z) E A, denotes the discrete (ground truth) value associated with the 7 feature. We are particularly interested in the setting where several features are correlated with the label y at varying strengths.
3.1 EVALUATING FEATURE QUALITY
In this work, we seek to compare the feature quality of the representations learned by different meth- ods. To measure feature quality, our idea is to train linear probe on the representation to discrim- inate the corresponding feature, and measure its performance (Rosenfeld et al. 2022 Kirichenko et al. 2022). Importantly, the training must be done on dataset only the desired feature is corre- lated with the label, so that probe only picks up only this feature from the representation.
Formally, we construct the linear probe dataset follows. Consider balanced distribution Doal where there is an equal number of points from each configuration of the features, i.c., there are [A1] |A2l x... I... subpopulations that have equal probability masses in D Then, for any feature i, we define the feature-probing error for this representation in terms the error of a linear probe in predicting the true feature value ax(x)) when trained on D Formally,
ProErr,(0) (8) def EewDD.loo (o(z), (I())) for u = arg min Eg- (e("" a.) (1)
where Pan corresponds to the empirical distribution over training data from Daal used to train the linear probe and l is some suitable classification loss such as the logistic loss.
3.2 TRAINING ALGORITHMS
Empirical risk minimization via Stochastic Gradient Descent (SGD). Stochastic gradient de- scent (SGD) is the de facto approach to minimizing the empirical risk over training data. Given a batch of training samples, {221, vI),... (IB. UR)). loss function l, and model (V, 0), we
B define the empirical batch loss ((w) def (L/B) Σ (f(w; (w;ri)) v.). The SGD update is:
w n V. L(w),
(2)
3

---

Published as a conference paper at ICLR 2024
Table 1: Comparing the in-distribution testing error and balanced distribution probing error for each feature for SAM and SGD, as defined in Section 3 For CelebA, the hard feature is gender, for Waterbirds, it is background, for CIFAR-MNIST, FMNIST.
CIFAR, and for FMNIST-MNIST, it is
Test
Easy ft. probe
Hard ft. probe
Test
Easy ft. probe
Hard ft. probe
CelebA
Waterbirds
SGD
4.7 + 0.07
10.9 + 1.15
20.9 + 0.93
5.2 0.13
11.7 + 2.28
21.8 + 01.99
SAM
4.3 1: 0.10
8.7 + 0.51
15.1 + 1.10
4.6 - 0.12
7.8 + 0.50
19.7 + 2.00
CIFAR-MNIST
FMNIST-MNIST
SGD
0.1 0.02
0.02
12.7 + 1.05
0.5 0.52
5.8 + 4.15
11.6 1.93
SAM
0.0 +: 0.01
0.1 = 0.03
10.2 +: 0.51
0.0 0.02
0.3 + 0.10
10.5 0.36
where 7) 0 is the learning rate. Unless otherwise mentioned, we use the cross-entropy loss for multi-class classification and logistic loss for binary classification.
Sharpness-Aware Minimization (SAM). In recent years, SAM has been proposed as an alterna- tive SGD, For model weights w, the SAM update
B w w w Σ V (ww) where w (wtt/ettw//V.cCtt)
def
(3)
for perturbation radius p O. We refer the reader to Foret et al. (2020) for derivation.
The SAM phantom parameter. Comparing SGD (Equation 2 and SAM (Equation Bl notice that SAM takes descent step in the direction of the gradient computed at different point: tú which we call the phantom parameter. The phantom parameter is computed by taking ascent step of fixed norm in the direction of the gradient at the current parameter w. This is an important characterization which we exploit in the rest of this paper. By studying how the phantom parameter ú relates the real parameter w, we can study how SAM changes the learning trajectory directly.
SAM IMPROVES FEATURE DIVERSITY
In this section, we demonstrate that SAM can empirically improve the quality of multiple redundant features, even when SGD would fail result of the simplicity bias (Shah et al. 2020). This quality particularly important when the features relevant to downstream performance are unknown at training time, and thus necessary to learn a high quality representation of all available features. We evaluate SAM on datasets in which multiple features have been labeled, where SGD tends to learn higher quality representation of the easier feature. With multiple labeled features, we can evaluate the quality of the representation of each feature individually. We first describe our experimental setup, and then present our results which we compare the feature quality of SAM and SGD.
4.1 DATASETS AND MODELS
Datasets. We use four datasets in our experiments each annotated by two features: CelebA (Lu al., 2015), Waterbirds (Sagawa et al.] 2019). CIFAR-MNIST (binary) (Shah et al.) 2020). and
FMNIST-MNIST (5-class) (Kirichenko et al. 2022]. The simplicity bias of SGD suggests that the casier of the two features will learned better. Thus, we refer the feature which attains the lower probing error for SGD trained models as the "easy" feature and the other has the "hard" feature. We describe the datasets in detail in Appendix B
Training setup. Following prior work (Kirichenko et al. 2022 Sagawa et al. 2019). we train an ImageNet-pretrainee ResNet-18 on CelebA and Waterbirds that has been initialized with weights pretrained on ImageNet. For CIFAR-MNIST we train a three-layer MLP from scratch. We apply standard data augmentation and weight decay (detailed in Appendix B and tune both algorithms over different learning rates, and for SAM, we tune over the parameter. We select the optimal hy- perparameters based on the validation set. For CIFAR-MNISS and CelebA, train for 20 epochs, and for Waterbirds we train for 100 epochs. We repeat each experiment four times using different random seeds, and report the mean and standard deviation of the errors.
4

---

Published as a conference paper at ICLR 2024
4.2 COMPARING SAM AND SGD
Our goal is to quantify the quality of the features learned by SGD-trained models vs. SAM-trained models. To this end, in all our datasets, we evaluate both the easy and hard features each model in terms of their feature probing error. As described in Section E] we compute the probing error by training a last-layer probe on a small dataset in which the feature of interest is uncorrelated with the other feature. Then, report the accuracy of this probe on holdout test dataset in which the feature of interest is again uncorrelated with the other feature. In accordance with our main hypothesis, find that SAM consistently achieves lower probing error for both features all our datasets, in comparison to SGD (Table D. We have thus verified that SAM implicitly improves the quality of multiple redundant features.
This suggests that even though SAM tries to optimize for the loss (and its sharpness) on the training distribution-and it does so successfully-
the hood, it also somehow improves feature quality on multiple different, balanced distributions in which the correlation between features is broken and examples are labeled by a feature of interest. We are interested in understanding how this unexpected "under the hood" improvement of the feature quality arises in SAM.
5 UNDERSTANDING SAM FEATURE LEARNING IN A TOY SETUP
In this section, design minimal, controlled setup that helps us isolate two core mechanisms within SAM. In our setup, data has an easier and a harder feature that are both predictive of the label. SGD exhibits the simplicity bias and learns a poor quality representation of the harder feature. In contrast, we show that SAM learns good representation of both features.
S
H
Wayya commlee
5
SGD
LSAM
5.1 A MINIMAL SETUP
Figure 1: Illustration of the toy ample. (i) The toy data distribu- tion. We vary the complexity of the spiral component the data by tightening the spiral. boundary of classifiers trained with SGD and with LSAM along a 2D slice where the other feature.
Feature distribution. We consider 4D setup where the first two coordinates correspond to a linear feature (the "easier" feature) and the second two coordinates corre- spond to a spiral feature (the "harder" feature). We denote the individual components x = [Taasy, Thard] The associated at- tributes (see Section 3 for each feature 01 (1r), 02 (a) are such that 01 (1()) = 1 if x resides on the right branch of the linear feature, and zero otherwise; similarly, (2(x)) 1 if resides on the right spiral, and zero otherwise. We assume binary label y E (-1,1} such that each feature independently is fully predictive of it.. y aj (2:) and y (2)). We can vary the complexity of the spiral feature by varying the number of rotations of the spiral. Here, the complexity refers to the amount of rotation of the spiral, in degrees.
R4
(ii) Decision
Disentangled architecture. We are interested in a setting where it is possible to precisely measure the quality of representations. To this end, we create a network with an explicitly disentangled representation space (rather than hoping that a trained linear probe would precisely disentangle it). Specifically, we define an architecture whose last-layer representation takes the form
def [(easy(1)) Paard def [Bo[(Feasy 0)), Po (0ardl))) )]],
D (I)
(4)
where bo: R- to IR is a three-layer neural network. This is followed by last layer weights [Veasy, Uhard] such that the final classification takes the form, sign(fw((z)) = thard (arra())))) where w = (v, 0) are the parameters of the neural network. Then, we can think of our feature probe for either features as simply & Peasy and h respectively.
T Peasy (1r) +
Probing error. Since we have a disentangled representation space, measuring the probing error is easy-there is no need to train a linear probe. since each feature is represented separately in the last layer. We can simply measure the error fo-1 (6.140aae (4e (x), y) or lo (#) y) to compute the probing error for the casy and hard feature, respectively.
LSAM: A simplification of SAM. For this controlled study, we consider a minimal version of SAM where we only perturb the last layer V; we call this LSAM. While the final descent of L.SA still updates all parameters; it is only the phantom ascent that is restricted to the last layer.
5

---

Published as a conference paper at ICLR 2024
Specifically, following the Section form:
notation, the phantom parameters take the following simple
def a -V (0ll..0//vvc 2)122. I def
(5)
This is nearly identical to the SAM phantom parameter described by Equation 3 except that only v is perturbed, and not O. As we will shortly see (Section 6). the above simplified version of SAM is sufficient (and in fact, helpful) in cleanly identifying our two fundamental mechanisms that can be linked to improved feature quality.
5.2 LSAM LEARNS THE HARDER SPIRAL FEATURE, WHILE SGD DOES NOT
We note that this dataset, both SAM and SGD learn the label y per- fectly. However, the behaviors are starkly different at the feature-level. In Figure [ we plot the contour of Peasy and Phard which correspond to how the models have learned either features. We observe, visually, that SGD has no trouble learning the lin- ear feature but has poor representa- tion of the hard spiral feature. How- ever, LSAM manages learn good representation of both features. We quantify this observation by measur- ing the spiral feature's 's probing error for varying complexities of the spiral feature (i.e., varying tightness levels), and for varying values of the SAM perturbation radius p (Figure A)) We find that while the learned spiral feature's 's quality is poorest when م 0 (i.e., SGD), progressively improves as we increase p. In effect, this toy example isolates core property of SAM that sets it apart from SGD: SAM learns higher quality representations of hard-to-learn features compared to SGD, even in the presence of a fully predictive easy feature. This is analogous to the feature-diversifying effect we observe in real-world datasets in Section 4
Ft. weight ratio (Tues) Roar
Algarihm SCD LS4M
Complecity 400
Algotthm
Complexity
SGD
Probing error (hard
0.4
552
15MM
600
100
0.2
0.0
0.20
3.50
0.25
1.00
Phantom parameter (9)
Phantom parameter (p)
represeent
Figure 2: (A) Test-set probing error of harder feature function of the phantom parameter p. for multiple complex- ities for the hard feature. SGD corresponds with p 0. (B) Phantom weight ratio as function of p. We plot uhard Ressy (solid lines) and thard/ Heasy (dashed lines) for different val- ues of when running LSAM. Note that SGD corresponds to when p = 0. As increases, the ratio Uhard Teasy in- creases.
WHY DOES SAM LEARN THE HARD-TO-LEARN FEATURE?
We have so far established that LSAM achieves significantly lower probing error on a hard-to-learn spiral feature in comparison to SGD (See Figure A)). In this section. we will describe the two mechanisms by which LSAM offers this Ceature-diversifying benefit without being explicitly trained to do so-we will argue that LSAM re-weights the training examples to increase weight on points that may be helpful for learning the hard feature, and that LSAM increases the effective step size of updates that fit the hard feature.
6.1 SAM RE- RE-BALANCES PHANTOM PARAMETERS
The first step understanding this mechanism is to understand how the phantom parameters Heasy and Gnard behave as we vary the perturbation radius p. We plot the ratio Unard /eeas as a function of p in Figure B), which shows that the ratio increases as we increase p. This implies that as is increased, LSAM places more relative weight on the hard feature during the update step, as described by Equation 3 Note that the re-balancing effect is not the result of the real weight thard and leass changing, but rather the result of the perturbation applied by LSAM when computing the phantom parameters. The ratio of the real parameters Chard/ Ueasy remains nearly constant as a function of the perturbation radius p (Figure 2 B), dashed line).
6.2 THE FEATURE- DIVERSIFYIE EFFECTS OF SAM
With knowledge that LSAM reweights the ratio Shard/Weasy our main idea is to decompose the gradient update into two terms in which we can interpret the effect of this reweighting. For binary

---

Published as a conference paper at ICLR 2024
Toy
CelebA
Waterbirds
SGD
LSAM
SGD
"SAM
SGD
SAM
0.5
0.5
0.5
0.0
0.0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
1.0
0.0
1.0
0.0
0.5
0.0
0.5
1.0
Fraction training examples
Fraction training examples
Fraction training examples
SGD (1)5AM
Real parameter
Phantom parameter
Figure 3: Lorenz curves for the real and phantom importance weight X; and AG. The dotted diagonal line represents the Lorenz curve for uniform distribution. The closer this curve is to this diagonal, the more equally the importance weights are spread. In blue, we plot the Lorenz curves for an SGD) checkpoint. In orange, we plot the Lorenz curves for LSAM checkpoint. The update step gradient is computed at real parameter for SGD, and the phantom parameter for SAM. We include the curves for the toy (left), CelebA (center), and Waterbirds (right).
classification, logistic loss on sample (Ii, Mı), we can write the gradient update to parameters as:
VoL = o (- Vi-aaVeaaae (Veasy Peasy| (z;) + Uhard Phard (ard (x;))
(eass Vo 4 easy (2%) + Thard V Phard (x;))
Importance weighting A:
Sum feature gradients sh.
(6)
where is the sigmoid function. If replace thard and Veasy with the phantom parameters Whard: Peasy. we get analogous terms X; and 0; for ISAM. The scalar importance weighting term X; represents the contribution of each individual example to the final loss gradient. The vector sum of feature gradients term gi is a weighted combination of the individual feature gradients. The effect of LSAM the importance weighting term is example-specific: the degree to which each point is up up-weighted by LSAM X, AA, depends on the example. However, the effect of LSAM the sum of feature gradients term is constant across all examples in a batch: the degree to which the effective learning rate of each feature is up-weighted by LSAM, Uhard Uhard and /easy/ Veasy, is the same for all examples within a batch.
In the remainder of the section, explore effects that arise when uhard Teasy Uhard/ Teasy
that enable LSAM to learn the harder feature.
Effect one: the importance weighting effect. We con- sider the effect of increasing Unard Veasy on the A term in Equation 6) which corresponds to how heavily each point is weighted the gradient. At a high level, we argue that this
Easy ft.
Hard ft.
Density
effect leads to two key beneficial behaviors that can be tested experimentally; first, examples are weighted more uniformly, and second, examples that are poorly fit by the hard feature are upweighted, even when they are well-fit by the easy fea- ture. Both of these effects promote the learning of the harder feature by encouraging additional points to contribute to the learning of the harder feature that would otherwise have con- tributed negligibly. In particular, the second effect leads to the up-weighting of points that incorrectly fit by the harder fea- ture but have a low loss because they are fit well by the easy feature. The effect of LSAM similar to the intervention im- posed by algorithms that explicitly up-weight examples that belong to the minority class in imbalanced datasets (Liu et al., 2021 Sagawa et al.) 2019), but it is remarkable that LSAM induces this importance weighting effect implicitly without being designed to do so.
1.0
Learning re-weighting factor (0//e))
Figure 4: The learning rate weighting factor for the easy fea- ture Ueasy/ /Veasy and the hard fea- ture uaard /hard/Uhard. Uhard- plotted as a distri- bution over batches in the dataset (see Equation
To verify that points are more uniformly weighted in the update step, visualize the distribution of their importance weights as a Lorenz curve, which plots the cumulative fraction of the total importance contributed by the top k points function of & (Figure This means that the closer the Lorenz curve is to a diagonal line, the more uniformly the importance weights are distributed across all points. We use the Lorenz curves compare the behavior of an LSAM step SGD)
7

---

Published as a conference paper at ICLR 2024
step given fixed model parameters (Figure left, comparing dashed to solid). We also compare the Lorenz curves of the importance weights computed using the real parameters and computed using the phantom parameters for a model that has been trained with SGD and one that has been trained with LSAM (Figure left, comparing solid line in LSAM plot to dashed line in SGD plot). We find that the importance weights are more uniformly weighted under LSAM in both setrings. We plot additional Lorenz curves corresponding to different points during training in Appendix G
To verify that points that are well fit by the easy feature but poorly fit by the hard feature are upweighted. We need a metric to measure the the degree to which the easy and hard features are fit by the model. In our toy setting, we can compute the contribution of each feature explic- itly. We can thus measure the degree to which the easy and hard features are fit by the model as their contribu- tion signed by the label: y (1) and (ar). When either term is positive, the point is classified correctly by the corresponding feature, with the magnitude of the term indicating the margin. In order to summarize how points weighted as function of the signed contributions, we partition the points into bins based on yyaasy(( (z)) and Praaaa (z)). Thus, we plot the median importance weight for each bin, which gives us the relationship between the importance weight of each example and their signed fea- ture contribution. We see that typical examples that are poorly-fit by the easy feature (yyPeas((a) (1)) is small) are as-
Fasy
Hard
Toy
25-
Easy contribution
Hard contribution
Lasy
Hard
1.0
CalebA
0.5
Easy contribution
Hard contribution
Importerce mo Aoadaee
Lasy
Hard
Waterbirds
signed less importance by SAM, and examples that are poorly-fit by the hard feature (arna (a) is small) are as- signed more importance by SAM (Figure This verifies our claim that LSAM up- up-weights points that are well-fit by the easy feature but poorly- fit by the hard feature in the toy setting. We plot additional importance weight plots corresponding to different points during training in Ap- pendix
Eryy RLdFF
Red paureii
Plonsirr paaeiii
Figure 5: Median importance weight- ing as a function of the contribution of each feature. We partition the data into bins based on the contribution of the easy and hard features (yWeasy Peasy and (hardd) as defined in Section 6.2 For each of these bins, we plot the me- dian importance weight term A; for the points in the bin. We include the corre- sponding plots for the toy (top). CelebA (center), and Waterbirds (bottom).
Verifying the importance weighting effect in real- world setup. We have thus far shown that LSAM in- duces the importance weighting effect in the toy setting. However, what happens when we relax our assumptions and move to a more realistic setting? We aim to test the importance weighting effect when we run the usual SAM algorithm without disentangled repre- sentation space on real-world datasets. We can evaluate importance weight uniformity directly. We plot the Lorenz curve for the importance weights evaluated at both the real parameter and the phan- tom parameter for CelebA and Waterbirds in Figure 3 We see that the points are more uniformly weighted under a SAM perturbation, suggesting that the importance weighting effect is present in the real-world datasets as well.
To verify that LSAM up- up-weights points that are well-fit by the easy feature but poorly-fit by the hard feature, we used the disentangled representation space to measure the degree to which the easy and hard features are fit by the model. However, in the real setting, we cannot directly measure the contribution of each feature. Instead, we rely on extracting the feature contributions from the linear probes Ceasy and thard that best classifies the respective feature (refer to Section
B With
these probes, we can define the signed contribution of the easy and hard feature as yoeasy ((x) and yhhaad (xx), respectively. Thus, we plot the median importance weight for both features as a function of yu yueasy 2 (z) and y for CelebA and Waterbirds. Analogous to the toy experiment, we see that examples that are poorly-fit by the easy feature are assigned less importance by SAM, and examples that are poorly-fit by the hard feature are assigned more importance by SAM (Figure 5)
Effect two: the learning rate scaling effect. We now turn our attention to the second term, gi. from Equation 6 The 9; term is weighted combination of the feature gradients with weights Teasy and Unard corresponding to the easy and hard feature respectively. These weights can be viewed as the "learning rates" as we take gradients steps in the direction of learning different fea-
8

---

Published as a conference paper at ICLR 2024
tures. Notably, these learning rates constant across points within batch. When LSAM causes Weard/Weasy > Uhard/ Usasy. 9: corresponds with a larger effective step size associated with the hard feature in comparison to 9:. As result the effective learning rate for the hard feature is higher and leads to more progress on the hard feature. We plot the ratio between the learning rates of the phan- tom parameters and the real parameters hhard/ Dhard thard and / /Vasy for each feature in Figure We see that LSAM increases the learning rate for hard features and decreases the learning rate for each feature. We call this the learning rate scaling effect. Unlike the importance weighting effect which acts on each data point separately, the learning rate scaling effect is more global across the entire training batch.
Unfortunately, we cannot directly verify the learning rate scaling effect without a disentangled rep- resentation space. We discuss this further in Appendix C.1
Contribution of each effect. To re- cap, we show that LSAM induces (i) an importance weighting effect and
Learnnng rats effect
Importance weighting offect
Combincd offect
(ii) a learning rate scaling effect, both
of which in turn arise from by re- balancing the phantom parameters. We have verified that each of these ef- fects are present in the toy example, and we would intuitively expect each effect to have a diverrifiiirii effect. But how do these effects com- pare in improvement in feature qual- ity, and which effect if any is more important? Are these effects simply correlations caused by SAM's improved features, or do these effects causally improve SAM?
Ft. (fawt/oo))
Ft. ratla
Figure 6: Simulations of the learning rate, importance weighting, and combined effects. Both effects offer gains in hard feature probing accuracy and together account for all the benefits offered by LSAM in the toy setting, See Ap- pendix C for detailed description of the simulations.
We answer these questions in the toy setting, where it is possible to explicitly add or remove each effect. To isolate each effect, we train with algorithm that applies only one effect at a time: either the importance weighting effect or the learning rate effect.
To isolate the importance weighting effect, we train with an algorithm in which we fix the ratio Uhard/ Veasy- = /hard/ / Peasy to a constant value in the importance weighting,
A = a (- Ve( (Vaaa (Pess((z)) + hhard naard (z:)
Ii is as usual for SGD
where Yeasy = and hhadd is hyperparameter.
We similarly isolate the learning rate effect and test the combination of both effects by fixing Uhard/Veasy analogously. The details are described precisely in Appendix C
In order to compare these algorithms with LSAM, we compute the probing error for the hard feature against the ratio between the hard feature contribution and the easy feature contribution. For our algorithms, we have manually set this value "hard/ Peasy (see Equations 7.. For LSAM, we consider this value to be the mean value of Uhara/Ueasy over training (From Figure BB). This leaves both our new algorithms and LSAM to be directly comparable as a function of this ratio, which we plot (Figure 6). We see that both individual interventions (Figure 6 left and center) improve the probing error for the hard feature. This suggests that both effects are causal and also significant (offer comparable gains) in explaining LSAM's improved spiral probing error in the toy setting. Further, the combined effect (Figure o right) matches the performance of LSAM and thus appears to provide complete picture of understanding LSAM in our toy.
7 CONCLUSION
In this work, we offer insight into the dynamics of SAM that lead to the improvement of feature diversity, in contrast with the usual sharpness-based arguments. We have shown how SAM can pro- mote balanced feature learning in the presence of multiple redundant features, and that this can lead to improved performance on out-of-distribution tasks. We have also demonstrated that SAM leads to improvements in feature quality for transfer learning on real data, including the Waterbirds dataset, CelebA, MIIST, and DomainBed. We hope that our insights provide a new perspective on the dynamics of SAM without relying on the flatness- based arguments, and that this work will inspire future work to further understand the mechanisms of SAM and will foster new algorithms for improving downstream performance on -fitt--iistiiuttio data in similar contexts.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
We would like to thank Christina Baek, Jeremy Cohen, and Suhas Kotha for their feedhack.
This material is based upon work supported by the National Science Foundation Graduate Research Fellowship under Grant No. DGE2140739. Any opinion, findings, and conclusions or recommen- dations expressed in this material are those of the authors(s) and do necessarily reflect the views of the National Science Foundation. This work was supported in part by the A12050 program at Schmidt Sciences (Grant #G2264481). We gratefully acknowledge the support of Apple.
REFERENCES
Momin Abbas, Quan Xiao, Lisha Chen, Pin-Yu Yu Chen, and Tianyi Chen. Sharp-maml: Sharpness- aware model-agnostic meta learning. In International Conference on Machine Learning, pp- 10-32. PMLR, 2022.
Samira Abnar, Mostafa Dehghani, Behnam Neyshabur, and Hanie Sedghi. Exploring the limits of large scale pre-training. arXiv preprint arXiv:21 2021.
Michael A Alcorn, Qi Li, Zhitao Gong, Chengfei Wang. Long Mai, Wei-Shinn Ku, and Anh Nguyen. Strike (with) a pose: Neural networks are easily fooled by strange poses of familiar objects. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 4845- 4854, 2019.
Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware mini- mization. In International Conference on Machine Learning, pp. 639-668. PMLR, 2022.
Maksym Andriushchenko, Dara Bahri, Hossein Mobahi, and Nicolas Flammarion. Sharpness-awann minimization leads to low-rank features. arXiv preprint arXiv:2305. 2023a.
Maksym Andriushchenko, Francesco Croce, Maximilian Müller, Matthias Hein, and Nicolas Flam- marion. A modern look at the relationship between sharpness and generalization. arXiv preprint arXiv:2302.07011. 2023b.
Saeid Asgari, Aliasghar Khani, Fereshte Khani, Ali Gholami, Linh Tran, Ali Mahdavi Amiri, and Ghassan Hamarneh. Masktune: Mitigating spurious correlations by forcing to explore. Advances in Neural Information Processing Systems, 35:23284-23296, 2022.
Dara Bahri, Hossein Mobahi, and Yi Tay. Sharpness- aware minimization improves language model generalization. arXiv preprint arXi::21 2021.
Kayhan Behdin and Rahul Mazumder. Sharpness-aware minimization: An implicit regularization perspective. arXiv preprint arXiv:2302. 1836, 2023.
Wieland Brendel and Matthias Bethge. Approximating cnns with bag-of-local-fenet features models works surprisingly well imagenet. arXiv preprint arXhv:9904.00660, 2019.
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 Infor- mation Processing Systems, 34:22405-22418, 2021.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In International Conference on Machine Learning, 1019-1028. PMLR, 2017.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 11929, 2020.
Jiawei Du, Hanshu Yan, Jiashi Feng, Joey Tianyi Zhou, Liangi Zhen, Rick Siow Mong Goh, and Vincent YF Tan. Efficient sharpness-aware minimization for improved training of neural net- works. arXiv preprint 1700v:2110.03111, 2021.
Jiawei Du, Daquan Zhou, Jiashi Feng. Vincent Tan, and Joey Tianyi Zhou. Sharpness-awar aware training for free. Advances in Neural Information Processing Systems, 35:23439-23451, 2022. 2022.
10

---

Published as a conference paper at ICLR 2024
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-awart minimiza- tion for efficiently improving generalization. arXiv preprint arXiv:2010.01412, 2020.
Robert Geirhos, Patricia Rubisch, Claudio Michaelis, Matthias Bethge, Felix A Wichmann, and Wieland Brendel. Imagenet-trained cnns are biased towards texture; increasing shape bias im- proves accuracy and robustness. arXiv preprint 1. 2018.
Robert Geirhos, Jörn-Henrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix A Wichmann. Shortcut learning in deep neural networks. Nature Machine Intelligence, 2(11):665-0) 1711666,,,,,, 2020.
Ishaan Guirajani and David Lopez-Paz. In search of lost domain generalization, 2020.
Suchin Gururangan, Swabha Swayamdipta, Omer Levy, Roy Schwartz, Samuel R Bowman, and Noah A Smith. Annotation artifacts natural language inference data. arXiv preprint arXiv:1803.02324, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Katherine Hermann and Andrew Lampinen. What shapes feature representations? exploring datasets, architectures, and training. Advances in Neural Information Processing Systems, 33: 9995-10006, 2020.
Katherine Hermann, Ting Chen, and Simon Kornblith. The origins and prevalence of texture bias in convolutional neural networks, Advances in Neural Information Processing Systems, 33:19000- 19015, 2020.
Minyoung Huh, Pulkit Agrawal, and Alexei A Efros. What makes imagenet good for transfer learn- ing? arXiv preprint
2016.
Pavel Izmailov, Polina Kirichenko, Nate Gruver, and Andrew G Wilson. On feature learning in the presence of spurious correlations. Advances in Neural Information Processing Systems, 35: 38516-38532, 2022.
Samyak Jain, Sravanti Addepalli, Pawan Sahu, Priyam Dey, and R Venkatesh Babu. Dart: Diversiff-aggrrgaate-rpaa training improves generalization of neural networks. arXiv preprint arXiv:2302.14688 2023.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. arXiv preprint 1912.0211 2019.
Jean Kaddour, Linqing Liu, Ricardo Silva, and Matt Kusner. When do flat minima optimizers work? Advances in Neural Information Processing Systems, 35:16577-16595, 2022.
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.
Simran Kaur, Jeremy Cohen, and Zachary Chase Lipton. On the maximum hessian eigenvalue and generalization. In Proceedings on, pp. 51-65. PMLR, 2023.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On large-batel training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv:1609.0488 04836, 2016.
Polina Kirichenko, Pavel Izmailov, and Andrew Gordon Wilson. Last layer re- re-training is sufficient for robustness spurious correlations. arXiv preprint arXiv:2204.02937, 2022.
Alexander Kolesnikov and Christoph H Lampert. Improving weakly-supervised object localization by micro-annotation. arxiv preprint arXiv:/605.05538 2016.

---

Published as a conference paper at ICLR 2024
Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Joan Puigcerver, Jessica Yung, Sylvain Gelly, and Neil Houlsby. Big transfer (bit): General visual representation learning. In Computer Vision- ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28 2020, Proceedings, Part V 16, pp. 491-507. -9riner.Spr 2020.
Simon Kornblith, Jonathon Shlens, and Quoc V Le. Do better imagenet models transfer better? In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pp. 2661- 2019.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Ananya Kumar, Aditi Raghunathan, Robbie Jones, Tengyu Ma, and Percy Liang. Fine- tuning can distort pretrained features and underperform out-of-distribution. arXiv preprint 10005 2022.
Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. Asam: Adaptive sharpness- aware minimization for scale-invariant learning of deep neural networks. In International Con- ference on Machine Leurning, pp. 5905-5914. PMLR, 2021.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 36(11):2278-3224 1998.
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, pp. 6781-6792. PMLR, 2021.
Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, and Yang You. Towards efficient and scalable sharpness-awe aware aware minimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12360-12377 2022.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
Dhruv Mahajan, Ross Girshick, Vignesh Ramanathan, Kaiming He, Manohar Paluri, Yixuan Li, Ashwin Bharambe, and Laurens Van Der Maaten. Exploring the limits of weakly supervised pretraining. In Proceedings of the European conference on computer vision (ECCV), pp. 181- 196, 2018.
Thomas McCoy, Ellie Pavlick, and Tal Linzen. Right for the wrong reasons: Diagnosing syntactic heuristics in natural language inference. arXiv preprint arXiv: 2019.
Xuran Meng, Yuan Cao, and Difan Zou. Per-example gradient regularization improves learning signals from noisy data. arXiv preprint 0xXiv:2303.17940 2023.
Mazda Moayeri, Phillip Pope. Yogesh Balaji, and Soheil Feizi A comprehensive study of image classification model sensitivity to foregrounds, backgrounds, and visual attributes. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 19087-19097, 2022.
Depen Morwani, Jatin Batra, Prateek Jain, and Praneeth Netrapalli. Simplicity bias in 1-hidden layer neural networks. arXiv preprint arXiv:2302.00457 2023.
Clara Na, Sanket Vaibhav Mehta, and Emma Strubell. Train flat, then compress: Sharpness-ane aware minimization learns more compressible models. arXiv preprint 1rXXi:2205.1269, 2022.
Behnam Neyshabur, Hanie Sedghi, and Chiyuan Zhang. What being transferred in transfer learn ing? Advances in neural information processing systems, 33:512-523, 2020.
Matteo Pagliardini, Martin Jaggi, François Fleuret, and Sai Praneeth Karimireddy, Diversity through disagreement for better transferability. In NeurIPS 2022 Workshop on Distribution Shifts: Con- necting Methody and Applications, 2022.
12

---

Published as a conference paper at ICLR 2024
Sinno Jialin Pan and Qiang Yang. A survey transfer learning ieee transactions knowledge and data engineering. 22 (10), 1345, 2009.
Mohammad Pezeshki, Oumar Kaba, Yoshua Bengio, Aaron Courville, Doina Precup, and Guil- laume Lajoie. Gradient starvation: A learning proclivity in neural networks. Advances in Neural Information Processing Systems, 34: 1256-1222. 2021.
Gregory Plumb, Marco Tulio Ribeiro, and Ameet Talwalkar. Finding and fixing spurious patterns with explanations. arXiv preprint arXiv:2106.00 1211:2200000112, 2021.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp- 9748-8763. PMLR, 2021.
Nasim Rahaman, Aristide Baratin, Devansh Arpit, Felix Draxler, Min Lin, Fred Hamprecht, Yoshua Bengio, and Aaron Courville. On the spectral bias of neural networks. In International Confer- on Machine Learning, pp. 5301-5310. PMLR, 2019.
Harsh Rangwani, Sumukh K Aithal, Mayank Mishra, et al. Escaping saddle points for effective generalization on class-imbalanced data. Advances Neural Information Processing Systems, 35:22791-22805, 2022.
Amir Rosenfeld, Richard Zemel, and John K Tsotsos. The elephant in the room. arXiv preprint arXiv:1808.00 03300, 2018.
Elan Rosenfeld, Pradeep Ravikumar, and Andrej Risteski. Domain-adjusted regression Erm may already learn features sufficient for out-of-distribution generalization. arXiv preprint crXiv:2202.06683 2022.
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 generaliza- tion. arXiv preprint arXiv: 1911::/9110007/11 2019.
Luca Scimeca, Seong Joon Oh, Sanghyuk Chun, Michael Poli, and Sangdoo Yun. Which short- cut cues will dnns choose? study from the parameter-space perspective. arXiv preprint vxXiv:2110.03090 2021.
Harshay Shah, Kaustav Tamuly, Aditi Raghunathan, Prateek Jain, and Praneeth Netrapali.. The pitfalls of simplicity bias in neural networks. Advances in Neural Information Processing Systems, 33:9573-9585, 2020.
Ali Sharif Razavian, Hossein Azizpour, Josephine Sullivan, and Stefan Carlsson. Cnn features off- the-shelf: an astounding baseline for recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition workshops, pp. 806-813, 2014.
Rakshith Shetty, Bernt Schiele, and Mario Fritz. Not using the car to see the sidewalk quantifying and controlling the effects of context in classification and segmentation. In Proceedings of the JEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 8218-8226, 2019.
Sahil Singla and Soheil Feizi. Salient imagenet: How to discover spurious features deep learning? 2021.
arXiv preprint
Chen Sun, Abhinav Shrivastava, Saurabh Singh, and Abhinav Gupta. Revisiting unreasonable ef- fectiveness of data in deep learning era. In Proceedings of the IEEE international conference on computer vision, pp. 843 852, 2017.
Hao Sun, Li Shen, Qihuang Zhong, Liang Ding, Shixiang Chen, Jingwei Sun, Jing Li, Guangzhong Sun, and Dacheng Tao. Adasam: Boosting sharpness-aware minimization with adaptive learning rate and momentum for training deep neural networks. arXiv preprint arXiv:2303. 00565, 2023.
Damien Teney, Ehsan Abbasnejad, Simon Lucey, and Anton Van den Hengel. Evading the simplic- ity bias: Training diverse set of models discovers solutions with superior ood generalization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 16761-16772, 2022.
13

---

Published as a conference paper at ICLR 2024
Dong Wang. Yicheng Liu, Liangji Fang, Fanhua Shang, Yuanyuan Liu, and Hongying Liu. Balanced gradient penalty improves deep long-tailed learning. In Proceedings of the 30th ACM Interna- tional Conference on Multimedia, pp. 5093-50010 2022.
Pengfei Wang, Zhaoxiang Zhang, Zhen Lei, and Lei Zhang. Sharpness-aware gradient matching for domain generalization. arXiv preprint arXiv:2303. 10353, 2023.
Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How does sharpness-aware minimization minimize sharpness? arXiv preprint (at::2221000722 2022.
Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. Sharpness minimization algorithms do not only minimize sharpness to achieve better generalization. arXiv preprint arXlv:2307.11007, 2023.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Pashion-mnise a novel image dataset for benchmark- ing machine learning algorithms. arXiv preprint arXiv: (00.::/77 2017.
Kai Xiao, Logan Engstrom, Andrew Ilyas, and Aleksander Madry. Noise or signal: The role of image backgrounds in object recognition. arXiv preprint (rXiv:2006.09994 2020.
Yao-Yuan Yang. Chi-Ning Chou, and Kamalika Chaudhuri. Understanding rare spurious correla- tions in neural networks. arXiv preprint arXiv:2202.05//8, 2022.
Haotian Ye, James Zou, and Linjun Zhang. Freeze then train: Towards provable representation learning under spurious correlations and feature noise. In International Conference on Artificial Intelligence and Statistics, pp. 1968-8990 PMLR, 2023.
Xiaohua Zhai, Joan Puigcerver, Alexander Kolesnikov, Pierre Ruyssen, Carlos Riquelme, Mario Lucic, Josip Djolonga, Andre Susano Pinto, Maxim Neumann, Alexey Dosovitskiy, et al. A large-scale study of representation learning with the visual task adaptation benchmark. arXiv preprint arXiv:1910.04867, 2019.
Xingxuan Zhang, Renzhe Xu, Han Yu, Hao Zou, and Peng Cui. Gradient norm aware minimization seeks first-order flatness and improves generalization. arXiv preprint arXiv:2303003/08, 2023.
Yang Zhao, Hao Zhang, and Xiuyuan Hu. Penalizing gradient norm for efficiently improving gener- alization in deep learning. In International Conference on Machine Learning, pp. 26982-26992. PMLR, 2022.
Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha Dvomek, Sekhar Tatikonda, James Duncan, and Ting Liu. Surrogate gap minimization improves sharpness-aware training. arXiv preprint arXiv:2203.0080 08065, 2022.
A ADDITIONAL RELATED WORK
Shortcut learning and spurious correlations. Many works have demonstrated shortcut learn-
ing in neural networks image and natural language classification tasks (Geirhos et al., 2018 2020)
Scimeca et al. 2021| Yang et al. 2022.
Hermann et al. [2200] Hermann & Lampinen 2020 Brendel
Bethge| 2019) Alcorn et al. 2019;
Shetty et al. 2019 Singla & Feizi, [021] Rosenfeld et al.
2018 Kolesnikov & Lampert 2016; Moayeri et al. 2022; Gururangan et al. 2018 McCoy et al.
2019
Xiao et al., 2020 Sagawa et al. 2019). In contrast, we focus purely on feature probing error which looks at how well a feature is captured in the representation irrespective of how much the final classifier relies on different features. We also care about probing error with respect to all predictive
features and do not delineate specific features spurious. Recent works (Kirichenko et al. 2022
Izmailov et al. 2022 Rosenfeld et al. (022) have shown that just last layer retraining on a suited dataset can significantly overcome reliance spurious correlations suggesting that the non-spurious features are reasonably well-learnt by the representation.
Feature diversity and finetuning. Multiple methods have been proposed to improve feature di- versity to evade the simplicity bias (Teney al.) (022),, ignore spurious features (Asgari et al. 2022 improve transfer learning (Pagliardini et al.) (2022), and improve generalization (Jain et al. 2023 rine-tuning and other retraining has been popular method to adopt learned features to tasks other than the pre-training objective (Pan & Yang (009). Retraining has been adopted to recover from
failure modes (Rosenfeld al.) 2022 Kirichenko et al. 2022 Kumar al. 2022) and improve
14

---

Published as a conference paper at ICLR 2024
generalization novel tasks (Radford et al. 2021: Sharif Razavian et al.
2014 Huh et al. 2016
Sun et al., 2017 Mahajan et al. 2018; Kolesmikov et al. 2020 Zhai et al. 2019 Dosovitskiy et al.
2020
Neyshabur al. 2020:
Abnar et al. 2021 Korbblith al. 2019).
Alternative SAM variants.
Multiple variants have been proposed to generalization notions of 2021 Zhuang et al., 2022). to improve the efficiency of SAM (Du et al. 2022
flatness Kwon et al.
2021
Liu et al. 2022) and generalize SAM to different optimizers (Sun et al. 2023).
B OMITTED EXPERIMENTAL DETAILS
Description of datasets. CelebA is large-scale face attributes dataset with 40 binary attributes. Following Kirichenko al. 2022). we train classifier to predict the "hair color" attribute, with values [blond, dark-hair). However, the "gender" attribute, with values {female, male} is highly correlated with "hair color" due to the imbalance in the dataset, and thus additionally be useful to predict the label, for the training distribution. Waterbirds is synthetic dataset consisting of images of birds superimposed onto different backgrounds. The data is annotated by "bird type". with values (water bird, land bird}. which refers to whether the pictured bird primarily lives in water or on land. The "background" attribute, with values fwater, land), highly correlated with "bird type". CIFAR-MNIST is dataset consisting images from MNIST (LeCun et al. 1998) and CIFAR-10 (Krizhevsky et al. 2009) that have been concantenated together, where the label of the CIFAR image and the label of the MNIST image are perfectly correlated for all training examples. Following Shah et al. (2020), we restrict the dataset to the binary setting. The "CIFAR" attribute can attain values |arplanee automobile] and refers the label of the CIFAR component of the image. Similarly, the "MNIST" attribute can attain values {0, 1}. For all datasets, we use the standard vrain/validation/test split, and when validation set is not provided, the training set. For computational efficiency, we down-scale the CelebA images to 64 x 64 pixels, and the Waterbirds images 96 x 96 pixels. For FashionMNIST-MNIS dataset (Xiao al. 2017 we restrict to the first five classes, associated with the digits 0-5 of MNIST.
random 90/10 split of
Reporting test error. Since the validation and testing datasets of Waterbirds and CelebA differ in distribution from the training set, be consistent with Kirichenko al. (2022) when reporting the testing error, we weight the testing error of each group by its corresponding frequency in the training dataset.
Architectures For the experiments involving CIFAR-MNIST and variants, train on a simple convolutional architecture including three convolutional layers, followed by ReLUs, with a final linear decoding layer. The architecture is defined by the following geudo-PyTorech
torch nn. Sequential
torch mn Conv2d 32, kernel-size=5, stride=2, padding=2).
torch mn Mm..clU(inplace-Tnn torch. .m. ....nn.......... (32, 64, kernel size=3, stride=2, padding=1] torch .nn (m.RLlU(inplace-Trre)) torch mn (64, 128, ternel.size=3, stride=2, padding=1). torch an ReLU(inplace-Truu torch an Flatten
torch mn Linear features.. num.classes)
For the experiments involving CelebA and Waterbirds, we use an ImageNeet ResNet-18 specified by He et al. (2016).
For the toy experiments, the representation neural network d had the architecture specified by the following Pseudo-PyTorch:
torch nn Sequential
torch nn. Linear featueeess 100),
torch nn LayerNorm (1000)) torch mn ..eRRU))
torch mn. Linear (100, 100). torch mn LayerNorm (100). torch. nn Rel.U)
15

---

Published as a conference paper at ICLR 2024
torch. mn Linear (100, n_classes)
Parameters and sweeps. For the toy experiments, we choose a constant learning rate of 0.01, batch size of 5, 300 training points, momentum, and no weight decay.
For the CIFAR-MNISH and FMNIST-MNIST experiments, we sweep over the learning rates (0.01, 0.05, 0.1} and the phantom hyperparameter p over {0.0, 0.01, 0.03, 0.05, 0.07, 0.1, 0.2]. We use batch size of 100, cosine learning rate schedule, momentum parameter of 0.9, and no weight decay. We normalize the images by the mean pixel value. Otherwise, do not use data augmentation.
For the CelebA and Waterbirds experiments, we sweep over the learning rates (0.0005, 0.001, 0.005, 0.01} and the p parameter {0.0, 0.01, 0.02, 0.05, 0.07}. We use batch size of 128, cosine learning rate schedule, momentum parameter of 0.9, and weight decay of 10-1, We use data augmentation described by the following pseudo-PyTorche
WATERBIRDS.RAANSFOOSSAAAAAAA transforms RandomResizeddrugg
transforms. Compose ([
scale (0.7. 1.0).
ratio (0.75 4./3.)
haappolltioonnmnteolllaaioooonno BILINEAR))
transforms RandomHoriizontllFiip (),
transforms ToTensor ()
transforms Normalize (WATERBIRD WATERBIRDS.STD)
1)
T transforms, Compose ([
transforms. Resize ((64, 64)),
transforms. ToTensor () transforms. Normalize (CELEBA CELEBA
1)
C ADDITIONAL DETAILS FOR THE TOY SETTING
Algorithms for the interventions To simulate only the importance weighting effect, we re-weight points in the importance weighting term X; to the weight of each point computed as if we had fixed the feature ratio. Similarly, to simulate the learning-rate effect, we re-weight each feature in the sum of feature gradients term by the weight that would assigned if we had fixed the feature ratio. In order to verify that each effect is important and to quantify the relative strenght of each effect, we intervene individually on each effect, separately, by fixing the ratio Teasy/ Uhard, but only when computing the ferm corresponding with that effect. This means that the loss gradient for the importance weighting intervention is computed as,
/
B
9.w. (0)
Σ
Tye (Veasy Peasy + Uhard Phard (2;)) (Veasy Vo Voleeasy(2)) + Uhard V Phard (2;)).
(8)
and the loss gradient for feature gradient effective learning rate intervention is computed as,
13
Ih.r (0)
(Feasy & Peasy (x:) + "aard (Veasy V Peaaa (I;) + hhadd V (Prrddd())) ys
For both equations, veasy and "hard are fixed and do not change during training. We train with SGD, but we compute the gradient as described in Equations 8 and To simulate the effects together, we combine the two interventions by fixing the ratio Ueasy/ Uhard in both terms. This means that the loss gradient for the combined intervention is computed as,
B
(0)
Σ
Tv. (Veasy Peasy (xt) + whard (ardd (Vaasy Ve (z+) + vhard V (eeddrr
(10)
16

---

Published as a conference paper at ICLR 2024
C.I THE LEARNING RATE EFFECT IS DIFFICULT VERIFY IN PRACTICE.
Unfortunately, in the absence of explicitly disentangled architecture (like in our toy setting), it is difficult estimate the learning rate scaling effect. The key challenge in achieving a similar obser- vation of the learning-rai effect on real data is understanding precisely how features are represented. Since, SAM perturbs the weights every layer (unlike LSAM). we would need to understand how each feature is represented at every layer of the neural network. Based on the general principle that SAM perturbs weights as function of how much they affect the output, we suspect, intuitively, that well-learned features will be inhibited by SAM at every layer. While we cannot verify the effect explicitly for a realistic setup, we believe it contributes to SAM's gains in feature probing error.
D THEORETICAL INTUITION AND ANALYSIS
In this section, we present a brief discussion of the theoretical intuition behind our results. In partic- ular, we will aim to understand the importance weighting effect.
Preliminaries. For simplicity, we will aim to understand the dynamics when training with a batch a single example x. We will assume, as is typical, that we have a neural network,
(0) = w" TuToo (x(
(11)
parameterized by 0, and where (x) is the representation of x under the neural network. We assume that the neural network is differentiable and well-approximated by a first-order Taylor expansion. We assume that the loss function L())= exp(-u5(0) is exponential, though note that cross-entropy is almost identical to exponential, for correctly classified points. Finally, recall from the main body of the paper that we can separate the loss gradient into two terms,
VoL(O) = -
exp(-uf(0))
Vaf
(12)
imporiance weighting term À feature gradient tern
We define the phantom importance weighting parameter Ā as the importance weighting term evalu- ated at the phantom parameter 0.
Phantom parameter. Recall that the SAM perturbation with exponential loss is defined as,
VoC(O)
g + P
(13)
&oC(O)II
VO (0)
+ P
(14)
Vo cx)(
P y eexp( (f
(15)
(o)))= (of
PM (0)]
(16)
Importance weighting, We first aim to understand the ratio between the phantom and real impor- tance weighting term X/X. We have,
X/A exp ( (1x=((00)eepppwf((((
(17)
= exp - (-y(((()-
(18)
which arises from plugging in the definition of the importance weighting term with exponential loss. For convenience, we will consider the log of the ratio.
log (5/2) = f(O)).
-
(19)
The effect based upon a first-order Taylor expansion of the logit function f(0).
f(0) f(0) - -100)ppp//fi (0) Vol(0)
+ +0 +0(
(20)
[o)l
z (()) V
(21)
17

---

Published as a conference paper at ICLR 2024
We can plug this into the expression for the ratio of importance weighting terms, and get,
log (2/2) ~ -y - py pyy -
f(0))
(22)
- [Va((0)
(23)
(24)
since y² = 1. We can see that the ratio of importance weighting terms is approximately proportional the magnitude of the feature gradient [[)ll-
For better sense of the dynamics, we compute [ for a few different architectures.
Example 1: LSAM. We consider the case of LSAM, in which only the last layer is perturbed to construct the phantom parameter. Since this is equivalent to considering the representation function é to be a constant, we only need to compute the feature gradient norms with respect to the last layer, w. We have,
V (()]l = (x)
(25)
This implies that the log importance re-weighting factor λ/A is proportional to the norm of the representation do (z).
Example 2: Two-layer linear network. We consider the case of a two-layer linear network, with single hidden layer. Let f() = v T Wr where (e, W). For convenience, we will compute the squared norm of the feature gradient,
Vef(0)]1
of f(O)ll W x|| ||
(26)
(27)
(28)
Example 3: Multi-layer linear network. We consider the case of a multi-layer linear network.
Let (0) = W1 WL-1 where (v, Wi.... WL-1). For convenience, will compute the squared norm of the feature gradient,
Σ
Ww. 1())
(29)
= W ... WT- 1
W1 ... W. 1 ... W2-10/12
(30)
,ll,, .....W W
(31)
where a, |[W ... w 1 constant with respect to z. This means that the squared norm of the feature gradient is proportional to weighted sum of the squared norms of the representations at each layer.
Example 4: Multi-layer perceptrons with ReLU. We can also compute the gradient of non- linear architectures, such as multi-layer perceptrons with ReLU activations. f() = v WW2 oW be defined by sequence of layers with the ReLU acti- vation function (xx) = max(0, x). For a particular input, we can compute the output of particular layer as,
particular, let
(x) W.W..WW......WW...... (WL
(32)
observe that for given input if we define the matrix A, = diag(L(o,(xx) > 0}), then observe that the output of the network can be written as,
f(0) T
AL-WWL-
(33)
18

---

Published as a conference paper at ICLR 2024
For convenience, we will compute the squared norn of the feature gradient,
I-1
[vf(0)]== 1
I [w)j1
(34)
= WA, AL- A W
(35)
W,WW2 A;- 1. 1 ... A 1
agll ay ... A..
(36)
Σ a,)lo,(0)11²
(37)
(38)
where aj = || W, A Unlike the multi-layer linear case, this constant depends on through Ay. This means that the squared norm of the feature gradient is proportional weighted sum of the squared norms of the representations at each layer, but where the weights depend on the input itself. Due this, general interpretation of the importance weighting effect is more difficult to obtain.
General multi-layer neural networks. In general, mutli-layer neural networks can introduce non- linearities that can make this gradient difficult to analyze. However, we suspect that the computation for mutli-layer linear networks will provide sensible approximation for the feature gradient norm for multi-layer neural networks.
E DATA DISTRIBUTION
In this section, we describe the data distribution for the toy setup.
The data the toy setup a concatenation two features that are independently drawn conditional on a label y. Before precisely defining the entire distribution, we will specify the feature distributions individually.
Easy feature distribution. We sample the easy feature, conditioned on the label y/ E (-1, 1). from a latent variable 2. The latent variable = is sampled from the uniform distribution over the interval [0, 1] if y = and from a uniform distribution over the interval [-1,0] if y = -1. The easy feature is then defined as the 2-dimensional vector [z, =]. and multiplied by a fixed scale parameter aeasy- Together,
((0,1)) 1) (u(-1,0)
yfy=1
Zeasy
(39)
yf
Teasy
Ocasy
[Zeasy Zeasyl
(40)
Hard feature distribution. Similarly, we sample the hard feature, conditioned on the label y E 1,,, from a latent variable 2. The latent variable z² is sampled from the uniform dis- tribution over the interval [0, (2xx/360)² if y = 1 and from uniform distribution over the interval [-(2xx (2mx/360)2 0) if y = 1, where is fixed scalar representing the complexity. The hard feature is then defined as the 2-dimensional vector [-2 cos 2, 2 sin 2), multiplied by fixed scale parameter Whard and added to some 2-dimensional uniform noise 7) (( 0.5) [0, 0.5).. Together,
(u( (2Tx/360)) { (U(-(2mxx360)) 0) if y -1
if y = =
(41)
hhard
Thard
Ohard
Shard Shard Zhard sin 2hard] +
(42)
Noise. For the variants that we present in the appendix, we add Gaussian noise E ~ (O0 o²), feature label noise u R(y), and feature dropout noise B(g), where R(p) is the Rademacher
19

---

Published as a conference paper at ICLR 2024
distribution with probability p of -1 and probability 1-p 1, and (Qq) is the Bernoulli distribution with probability q. The noise is added to the easy feature after it is scaled by a.
I feasy hard) =
(seasy,
[Feeasy hard) 2 "(easy
(43)
In general, we choose Geasy
2 and Ghard 0.25.
F DOMAINBED
Table 2: Comparing the domain transfer performance of SAM and SGD on the DomainBed datasets.
OfficeHome
PACS
Domain
Domain
A
C
P
R
A
C
P
S
SGD
0.687
0.716
0.853
0.777
SGD SAM
0.910
0.897
0.955
0.893
SAM
0.709
0.745
0.858
0.784
0.914
0.921
0.964
0.908
VLCS
Domain
C
L
S
SGD
0.993
0.736
0.748
0.812
SAM
0.996
0.748
0.776
0.830
In this section, present comparison of SAM and SGD DomainBed (Gulrajani & Lopez- Paz, 2020). DomainBed is a benchmark consisting of datasets with realistic domain shifts. Thus, the results on DomainBed
indicative of SAM's performance under real-world domain shifts in which the features that are useful for in-distribution classification may differ from the features that are useful for citto--distribution classification. We compare SAM and SGD on three of the DomainBed datasets: OfficeHome, PACS, and VLCS.
F.I DATASETS
OfficeHome is a dataset consisting of images from four domains: Art (A), Clipart (C), Product (P), and Real-World (R).
PACS is a dataset consisting of images from four domains: Art (A), Cartoon (C), Photo (P), and
Sketch (S).
VLCS is dataset consisting of images from four domains: Caltech101 (C), LabelMe (L), SUN09 (S), and VOC2007 (V).
We use the standard test/train splits of each of the datasets.
F.2 EXPERIMENTAL SETUP
To evaluate domain transfer performance for a given dataset and domain, we train a classifier on the training set of all domains except the target domain. Following the setup of the main paper, extract the quality of the representation by training a linear probe on the training dataset of the target domain.
More precisely, given the representation & of the classifier that has been pre-trained on the other domains, we aim minimize the following loss on the training set of the target domain:
((u) D [M(u (()).v)
(44)
where Durget refers to the training distribution of the target domain, is the linear probe we are optimizing, E is cross-entropy loss, and x and y are the input and label the training example. We evaluate by measuring the accuracy of the linear probe on the test set of the target domain.
We perform a hyperparameter sweep over P E {0, 0.03, 0.05, 0.13 and the learning rate 7 E (0.005, 0.01, 0.02). We validate small held-out validation set (10% of the training dataset) selected from the training dataset split. We report the results the best hyperparameter setting on the test set of the target domain.
20

---

Published as a conference paper at ICLR 2024
F.3 RESULTS
We present the results of SAM and SGD on the three datasets in Table 2 We find that SAM out- performs SGD all three datasets, confirming that SAM improves the representation quality of the neural network in variety of domain transfer settings.
G EXTENDED FIGURES
We provide extended versions of Figures and in Figures and 8 respectively, These figures show the Lorenz curves for the real and phantom importance weights A and \; and the median importance weighting as a function of the contribution of each feature, respectively. We plot cach result over multiple checkpoints over the duration of training,
Results. We observe that the trends discussed in Section 6.2 are consistent across the duration of training. The phantom importance weights X are more uniformly distributed than the real im- portance weights λ₁. Further, the median importance weighting is higher for points with higher contribution of the easy feature than the hard feature.
Toy
CelebA
Waterbirds
Epoch
Epoch
Epoch
1.0
SGD
LSAM
SGD
SAM
1.0 SGD
SAM
0.5
0.0
0.0
0.0
0.0
0.5
1.0
0.0
1.0
Fraction training
Fraction training examples
Fraction
examples
Epoch 10
Epoch
Epoch
SGD
LSAM
SGD
SAM
SGD
SAM
0.5
0.5
0.0
0.0
0.0
0.5
0.0
0.5
0.5
training examples
training examples
training examples
Epoch
Epoch 13
Epoch
SGD
LSAM
SGD
SAM
SGD
SAM
0.5
0.5
0.5
0.0
0.0
0.0
0.5
1.0
0.5
1.0
1.0
0.0
0.5
0.0
0.5
0.0
0.5
1.0
Fraction training examples
Fraction of training examples
Fraction of training examples
SGD
(L)SAM
Real parameter
Phantom parameter
Figure 7: Extended version of Figure 3 Lorenz curves for the real and phantom importance weight
A; and Åς The dotted diagonal line represents the Lorenz curve for uniform distribution. The closer this curve is to this diagonal, the more equally the importance weights are spread. In blue, we plot the Lorenz curves for each SGD checkpoint. In orange, we plot the Lorenz curves for each LSAM checkpoint. The update step gradient is computed at real parameter for SGD, and the phantom parameter for SAM. We include the curves for the toy (left), CelebA (center), and Waterbirds (right). We observe that within the SAM checkpoints, the weights of points evaluated the phantom parameter are closer to uniform than when evaluate at the real parameter. Further, in comparison an SGD checkpoint, the phantom parameter of SAM weights points more uniformly. We plot each result over multiple checkpoints over the duration of training (epoch in bold).
H VARIANTS OF THE TOY SETUP
21

---

Published as a conference paper at ICLR 2024
Toy
CelebA
Waterbirds
Epach
X10
Easy
Lary
Hard
25-
6.0
10
0.5
0.0
contribution
Hand buttin
ft.
Easy
sontribttion
Epoch 10
Epoch
Epoch
Hibe
Easy
Easy
0.3
10
Fany contribution
Hand ft contriutiio
Easy ft. bution
Hard intrinbution
contribation
x10
Epoch 71
Easy
Lasy
Hard
Easy
25
100
contribution
Easy Iution
contribatiion
contribation
SGD
(LJSAM
Roal parameter
Phantom parameter
Figure 8: Extended version of Figure Median importance weighting function of the tribution of each feature. We partition the data into bins based on the contribution of the easy and hard features (yUeasy Peasy and y'eeasy Phard), as defined in Section 6.2 For each of these bins, plot the median importance weight term X; for the points in the bin. We include the corresponding plots for the toy (top), CelebA (center), and Waterbirds (bottom). We plot each result over multiple checkpoints over the duration of training (epoch in bold).
22

---

Published as a conference paper at ICLR 2024
A1
B1
C1
mo.n
dropout
Posits enc (hed
0.3
408
Phartion
Phurtom
A2
B2
C2
Comples
Photing end (bont R.)
Pridig une (hed
Complox
6.3
n.2
Complec
01
018
00
100
Prantor
(1]
Phontem
(x]
Phanzom perarizz (x)
Figure 9: Variants the toy setup to include noise in the feature distributions (compare with Fig- ure 2). We include three types of noise distributions: (A) label noise, (B) Gaussian noise, and (C) feature dropout. For each type noise, we apply the noise to: (1) both the simple and complex fea- tures, and (2) only the complex feature. We plot the probing error of the hard feature as a function of the LSAM phantom parameter Note that = 0 corresponds to the baseline SGD model.
23

---

Published as a conference paper at ICLR 2024
Peciigg orer (t.d
Batch size n.
200
200
Phantom paratettor (0)
Figure 10: Variant of the toy setup which we vary the batch size (compare with Figure 2 We
plot the probing error of the hard feature as function of the LSAM phantom parameter p. Note that
p = 0 corresponds to the baseline SGD model.
24

---

Published as a conference paper at ICLR 2024
In this section, present variants of the toy setup include noise in the feature distributions and to vary batch size. In general, our results are consistent with the results from the main paper. We observe that LSAM is robust wide range of noise distributions and batch sizes.
H.1 NOISE DISTRIBUTIONS
In order to explore how LSAM behaves when features include noise, we train classifiers using the toy setup with three types of noise distributions.
Label noise. We add label noise by randomly flipping the labels of fraction of the features when generating the training, as defined precisely in Section Note that this is different from how label noise is typically defined:
not flipping the labels of the training examples, but rather the labels of the features. We vary the fraction label noise from 0% to 10%.
Gaussian noise. We add Gaussian noise to the features, defined precisely in Section E the standard deviation of the Gaussian noise from 0 to 1.0.
We vary
Feature dropout. We add feature dropout to the features, in which, with some probability, replace features with O, as defined precisely in Section E We vary the probability of dropout from 0 to 0.4.
Hard-features fentud-eturr--nyy noise. For each type of noise, we also consider a variant in which the noise is only added to the hard feature.
In general, we find that LSAM improves the probing error of the hard feature in comparison to SGD for all types of noise (Figure 9.. This suggests that the mechanisms of LSAM that improve feature representation quality are robust to noise in the feature distributions.
H.2 BATCH SIZE
We also vary the batch size in the toy setup. We plot the probing error of the hard feature as function of the LSAM phantom parameter p in Figure I0 Note that p 0 corresponds to the baseline SGD model. We observe that LSAM is generally robust to wide range of batch sizes, and that the optimal value of p is similar across batch sizes. However, as batch size increases be very large, the improvement of L.AAM degrades, consistent with [CITE]. In addition, we observe that for large p. the probing error of the hard feature is more sensitive batch size.
25


Reviews:
[{"review_id": "rSG4PfG68u", "reviewer": "Reviewer_LFgC", "summary": "Several learning theories, e.g., PAC Bayes, explain that flatter minima generalize better.\nThus, SAM is formulted to find flatter minima, and is known to improve in-distribution\ngeneralization.\nHowever, the beyond in-distribution performance is not well understood.\nThis paper focuses on  a feature-diversifying effect of SAM that is relevant for downstream generalization.\nAs seen in simplicity bias or gradient starvation,SGD exhibits a bias toward learning simple/easy-to-fit feature, even if the training data has multiple latent features that are all useful for prediction.\nIn contrust, this paper explans that SAM also learn the hard-to-fit features.\nThe authors shows this property of SAM by constructing a toy setting consisting of an easy-to-fit and a hard-to-fit features.", "questions": "Q1.What does it take in the future to extend the theory from toy setting to reality image identification? Please explain how difficult it would be. \nQ2: As a question related to Q1, in the toy setting, easy-to-fit and hard-to-fit are separated as latent features, but this may not necessarily be the case in real-world data. For example, is it possible that at the time of the feature extractor, the SGD is trained in such a way that the signals of hard-to-fit features do not reach the output layer in the first place? Is there any way to know that easy-to-fit and hard-to-fit features are disentangled in latent features in a real-world problem?\nQ3: How much is the influence of the terms after the second order in Eq.(20), which can be ignored since the hessian trace is small for the purpose of finding the fat minima? In this case, what is the effect during the learning process?\nQ4. The word \"diversity\" is used in this paper; however, it actually refers only to two characteristics, easy-to-fit and hard-to-fit.\nIn terms of beyond-in-distribution setting,  transferabiliity is also important. For example, robust/non-robust features in the analysis of adversarial examples, etc. could be considered. Please explain the relationship with other features.\nQ5. There are oftern noisy features as well as easy-to-fit and hard-to-fit features in real images. Is it possible to analyze the effect of SAM on noisy features that are typically hard-to-fit.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "This paper seems to be the first to show the relationship between easy-to-fit and hard-to-fit features and learning by SAM. \nThis perspective is useful in that it opens up the possibility that SAMs can solve problems that SGD has, such as gradient starvation.", "weaknesses": "While the analysis of the toy setting is interesting, it is unclear to what extent the analysis is applicable to learning with real images such as the winter bird and CelebA datasets.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vJcB47ouQ3", "reviewer": "Reviewer_NKN8", "summary": "This paper studies why Sharpness-Aware Minimization (SAM) can generalize well. They provide an explanation based on the weight that SAM/SGD put on different features, explaining SAM can up-weight features that are harder to learn.", "questions": "-", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "I think the idea of balancing features is interesting, and it makes sense to study SAM from this perspective. I agree that sharpness of the loss cannot alone explain the differences of SAM/SGD, and I appreciate this work's attempt to go beyond.", "weaknesses": "The paper is inherently more on the empirical side, which I think is fine. However, I think some simplified theory could increase the contribution of the work. Based on this, I'd rate the paper as borderline accept.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "23Auk2tF1M", "reviewer": "Reviewer_BgWc", "summary": "This paper empirically studies the mechanism of Sharpness-Aware Minimization (SAM), which improves generalization beyond in-distribution data (e.g., on downstream or out-of-distribution tasks). The authors demonstrate that SAM has a significant impact on improving representation quality through feature diversity, i.e., promoting a neural network to learn both easy and hard features. Through toy experiments, the authors isolate two core mechanisms within SAM that lead to better feature quality: 1) importance sampling, and 2) learning rate scaling. The former one is further verified on real-world datasets.", "questions": "Please see weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- Explaining the effectiveness of SAM from the perspective of feature quality is novel.\n- The authors identify that SAM’s feature improvement cannot be solely explained by in-distribution improvements.\n- The authors carefully design a toy experiment, based on which two reasonable mechanisms are isolated to explain why SAM can learn the hard-to-learn features.", "weaknesses": "- The authors argue that due to the ability to learn diverse features (both easy and hard), SAM can improve the generalizability beyond in-distribution. However, this raises a question: are all these features beneficial to the out-of-distribution generalization or downstream tasks? For example, in Table 1, the hard-to-learn feature in CelebA is the \"gender feature\" while the task is to predict the \"hair color\". In this case, the easy features are correctly related to the labels, and there exists spurious-correlation between the hard \"gender feature\" and the labels. The SAM indeed decreases both the probe errors for easy and hard features, but the latter is harmful for ood generalization. Thus, it seems contradictive between \"SAM learns diversity features\" and \"SAM improved generalizability beyond in-distribution\". The authors should provide more explanations.\n- While the authors identify two underlying effects to explain why SAM learns diversity features with better quality, the learning rate effect is hard to verify in real-world datasets.\n- Minor: In Sec3.1, Eq 1, replace the $\\phi_\\theta$ to $\\Phi_\\theta$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4GHqkzpeL0", "reviewer": "Reviewer_KbQq", "summary": "Sharpness-aware minimization (SAM) is an alternative method to the traditional stochastic gradient descent (SGD) for training neural networks. The core idea behind SAM is to guide models towards \"flatter\" minima, believed to enhance generalization. However, recent works revealed that flatness is not the only reason for good generalization. This study proposes an additional advantage of SAM: it improves the quality of dataset features, especially in datasets with redundant or unnecessary features. SAM achieves this by suppressing already well-learned features, allowing lesser-known features to be recognized. This process leads to a more balanced model feature update. Consequently, SAM enhances representation learning by promoting feature diversity. This sheds light on SAM's effectiveness in out-of-distribution performance, offering a new perspective on SAM's benefits.", "questions": "- In experiments such as CIFAR-MINIST, do all the data have a strong feature and a weak feature, or only some of the data have a strong feature while others do not? In equation 6, the authors divide the gradient into the direction of the easy feature and the hard feature, so it seems that each data has both strong and weak features.\n\n- A recent theoretical paper [1] shows that SAM can prevent noise memorization by deactivating the neurons that align with noise and, therefore, can get better in-distribution performance. The authors provide a novel view of the strong-weak features but haven't considered the noise in the model. So, I am wondering how the noises (which are independent of the label ) influence the SAM and SGD's out-of-distribution performance. \n\n-  The batch size $m$ of SAM updates plays an important role [2, 3, 4]. It is observed and proved that m-SAM (SAM with batch size $m$) can have different implicit biases and, therefore, get different performances. It would be better if the author could include the batch size in the methodology and have some discussions. \n\n- I checked the appendix in detail and found that the authors calculate the ratio of importance weighting terms for different neural network architectures. The authors only calculate the gradient for deep linear networks and leave the general networks blank. I agree that the general network's gradient is complex and hard to compute. However, adding some discussion for the two-layer neural networks can help the authors better understand the intuition of this paper,  which is still simple and easy to calculate [1, 4]. \n\n[1] Chen et al. \"Why Does Sharpness-Aware Minimization Generalize Better Than SGD?\" arXiv preprint arXiv:2310.07269, 2023.\n\n[2] Foret et al. \"Sharpness-aware minimization for efficiently improving generalization.\" In ICLR, 2021.\n\n[3]  Andriushchenko et al. \"Towards understanding sharpness-aware minimization.\" In International Conference on Machine Learning, pp. 639–668. PMLR, 2022.\n\n[4] Wen et al. \"Sharpness minimization algorithms do not only minimize sharpness to achieve better generalization.\" arXiv preprint arXiv:2307.11007, 2023.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Generally, I think this paper provides an interesting view of SAM. \n\n- This paper provides a new point of view that the SAM algorithm can boost feature diversity. In detail, SGD may be biased towards learning the strongest features from the training data, while SAM promotes the learning of a diverse set of features. Such an explanation is different from the view of the sharpness-based or Hessian-based analysis. \n\n- The new view indicates that the SAM can not only benefit in-distribution tasks but also downstream or out-of-distribution tasks. The empirical study supports their results. The findings may inspire the community to develop new algorithms to improve feature learning.\n\n- This paper is well-organized and presented with comprehensive experiments and interesting illustration figures.", "weaknesses": "- Incomplete Understanding of Mechanisms: The discovery that SAM can result in a variety of features is interesting. However, the underlying mechanisms of this finding are not entirely comprehended, particularly in the theoretical aspects.\n\n- Clarity in Data Model Description: The study explores SAM from a feature learning perspective. Nevertheless, the data model described in the paper lacks clarity. It would enhance the paper's comprehensibility if the authors could explicitly present the formula of the toy model (perhaps in the Appendix).\n\n- Elaboration on Feature Measurement: The paper introduces a method to measure strong and weak features. I think it would be an independent significant contribution. An in-depth elaboration on this measure would further strengthen the paper's impact.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FXhdPz19u6", "reviewer": "Reviewer_qpzp", "summary": "The paper focuses on the effectiveness of Sharpness-Aware Minimization (SAM).\nThe authors claims that SAM can enhance downstream generalization through feature diversity.\nEmpirical analysis of SAM trues to demonstrate that SAM offers feature-diversifying benefit.", "questions": "See Weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "- The paper provides clear motivation.\n- The paper provides several experimental results from various perspectives.", "weaknesses": "- Although the paper introduces several concepts including feature-probing error or phantom parameter, but it is still unclear why SAM can improve feature diversity.\n- What is the meaning of the ratio $v_{\\text{hard}}/v_{\\text{easy}}$? Why we can say the high ratio means re-balancing effect?\n- Why do we need to perturb the last layer only? It would be nice to describe in detail.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gF8i8Z7Vxf", "reviewer": "Reviewer_Jhr2", "summary": "This paper shows that SAM can enhance the quality of features in datasets containing redundant or spurious features and this effect can't be explained by in-distribution alone. They show two possible mechanisms for this effect, namely (1) SAM will upweight examples that can't predict well using so-called 'hard features', and (2) the modified gradient in SAM optimizer will attribute more to 'hard features' instead of well-learned 'easy features'. They support the claims with a simple controlled toy experiment and also experiment on real datasets.", "questions": "* **Q1.** Is it possible to verify the effect of the found mechanisms? (see weaknesses)\n\n* **Q2.** Can the authors try disentangle the two mechanisms and verify which is more important for the phenomenon?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* **Originality.** To the reviewer's knowledge, this paper is very novel. The phenomenon that SAM may learn more diverse features than SGD even when both reach nearly perfect generalization hasn't been reported before.\n\n* **Quality.** The experiments conducted are complete and thorough.\n\n* **Significance.** The reviewer believes that the finding will deepen the understanding of the mechanism of SAM.", "weaknesses": "* **Clarity.** The reviewer feels that clarity requires improvement in the paper. For example, the legend of Figure 5 is hard to follow. Also, the equality (1) seems to contain a typo. \n\n* **Verification of mechanism.** A possible way to verify whether the effect pointed out in the paper really contributes to the phenomenon is to try designing different variants using the discovered mechanism. The paper now only verifies that the importance weighting mechanism exists in real-world experiments but hasn't connected it with the diversity of features.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)