Sharpness-Aware Minimization Enhances Feature Quality via Balanced Learning

Thank you for your thoughtful review, we appreciate your feedback!

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.

Thanks for pointing this out. We will address all issues of clarity in the revised manuscript. For reference here, would like to clarify the meaning of the legend of Figure 5:

As a preliminary, Figure 5 aims to show that points that are well fit by the easy feature are up-weighted, and points that are poorly fit by the complex feature are down-weighted, on average. As a measure of how well each point is fit by the easy feature, we define the “signed easy feature component” as $y v_{\text{easy}} \Phi_{\text{easy}}$. This term represents the magnitude of the output of the easy feature component of the network signed by whether the prediction is correct. When this component is large (well fit by easy feature), we observe that points are generally up-weighted. When the corresponding component for the hard feature is small (poorly fit by hard feature), we observe that points are generally up-weighted.

In the figure caption, we will emphasize that we are plotting, on the x-axis, how well each feature is fit, and on the y-axis, how much the corresponding point is weighted (on average).

Q1. Is it possible to verify the effect of the found mechanisms?

Q2. Can the authors try disentangle the two mechanisms and verify which is more important for the phenomenon?

We agree with this concern, and we present this as a limitation in our conclusion section. We want to clarify the challenges associated with verifying and quantifying the effects in practice. In particular, we want to explain the challenges in establishing a causal relationship between each effect and the improvement in feature diversity. Further, we argue that this challenge is precisely why we must rely on a toy model for understanding.

For the toy example, we are able to determine that both effects cause improvement in feature quality through the experiment presented in Figure 6. The exact experiment is explained precisely in Appendix C, but we summarize here. To isolate the learning rate effect, we train with SGD, but where the feature gradients term—associated with the learning rate effect—is computed with artificially balanced features.

Ideally, we would like to perform the same experiment on realistic data. However, in order to artificially balance features, we would have to understand precisely how each feature is represented throughout the neural network, in order to make the appropriate perturbation. This is challenging because the community does not currently understand how features are represented throughout the network.

We can approximate this effect by taking advantage of our observation that a single SAM step will balance features. We can compute the feature gradients term under a SAM perturbation, which approximates the term if features were balanced. However, when training with SAM normally, the benefits of both the importance weighting effect and the feature learning effect are compounded. By approximating only one of these effects at every step means that we cannot take advantage of the cumulative balancing step of both terms, and thus such an experiment will underestimate the effect of each term individually.

Due to the challenge of directly verifying the mechanisms on real data, we need an alternative way to verify our hypothesis. This is precisely the reason that establishing an interpretable toy setup that models the behavior of SAM is an important tool for understanding.

Thank you for your thoughtful review, we appreciate your feedback!

We appreciate the suggestions to improve clarity and we hope to address them below. In our revised version, we will address each point explicitly in the paper.

Although the paper introduces several concepts including feature-probing error or phantom parameter, but it is still unclear why SAM can improve feature diversity.

We’d like to ask for further clarification on what you mean that “it is still unclear why SAM can improve feature diversity”. If you are concerned with our definition, we define feature diversity to measure how well each algorithm can achieve small linear probing error for all available features. More formally, we define a measure of feature diversity as the worst case (maximal) ProErr (see Equation 1) over the features. The lower the error, the more diverse the features. We will certainly clarify this precisely in the revisions. If you are concerned with the methodology of the paper, please let us know if there is any particular issue with our argument: in general, we construct a toy setup where we show that when multiple features are present, SAM can promote the learning of features that SGD performs poorly on. We confirm that SAM improves feature diversity on real-world datasets, and we validate our conclusions from the toy setup.

What is the meaning of the ratio $v_{\text{hard}} / v_{\text{easy}}$?

The quantities $v_{\text{hard}}$ and $v_{\text{easy}}$ refer to the last-layer parameters of the architecture (specified in Equation 4 and following text). They affect the gradient—and therefore update step—of the model by placing weight on the hard and easy features, respectively, as described in Equation 6. Thus, the ratio $v_{\text{hard}} / v_{\text{easy}}$ is important because it determines the relative weight placed on the hard feature compared to the easy feature during the update step. For SGD, the easy feature tends to dominate the update step, and thus by placing more weight on the hard feature will “balance” the feature contributions. This corresponds with increasing the ratio $v_{\text{hard}} / v_{\text{easy}}$.

Why do we need to perturb the last layer only?

The toy setup (LSAM) uses a last-layer-only perturbation for simplicity and for interpretability. Despite its simplicity, LSAM exhibits many of the core behaviors of SAM, such as improving feature diversity via the balancing effect. In addition to presenting an understanding of LSAM in the toy setup, we verify that the conclusions we make from the toy setup generalize to SAM---where we perturb every layer---on real data. Please refer to Section 6.3, “Verification on Real-world datasets”.

Thank you for your thoughtful review, we appreciate your feedback!

It would enhance the paper's comprehensibility if the authors could explicitly present the formula of the toy model.

Of course! We have added this to the revised paper. Please see Appendix E in the revised copy.

An in-depth elaboration on [the linear probing error] measure would further strengthen the paper's impact.

We will add discussion of this measure in detail in the appendix in the form of extended related work, in which we will discuss the following. The linear probing measure—in which we measure feature quality by measuring the error of the linear classifier on top of the representation layer that minimizes error of the desired feature—has been applied previously. Previous work has shown that the last layer representations often encode features that are useful for out-of-distribution generalization even when the neural network performs poorly OOD (without linear probing) [1]. Surprisingly, linear probing error is often close to the error as if the neural network was finetuned (with all parameters) on the downstream data [2], which makes it a strong measure of feature quality. In addition, linear probing is often itself used for downstream generalization when there is only limited downstream data available, or in compute constrained environments, notably for [3].

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 all cases, every data point has both features: for example, in CIFAR-MNIST, the “hard” CIFAR feature was {truck, airplane}, and an image of a truck or an image of an airplane would be present, and similarly for MNIST, {0, 1}. This is also true of every other dataset that we consider.

We have added experiments to the toy setting but where each feature is not necessarily present in every training example. Consistent with our results, we find that in this setting, SAM improves performance. Please see our “General Response” for specific details.

So, I am wondering how the noises (which are independent of the label ) influence the SAM and SGD's out-of-distribution performance.

It would be better if the author could include the batch size in the methodology and have some discussions.

We have added experiments to test various types of noise and varied batch size in our toy setting. Consistent with our results, we find that in both of these settings, SAM improves performance. When we vary batch size, consistent with prior work, small batch size tends to improve performance with SAM. Please see our “General Response” for specific details.

Adding some discussion for the two-layer neural networks can help the authors better understand the intuition of this paper.

We have added the analysis of the importance weighting ratio for a multi-layer ReLU network (MLP) to the revised paper (Appendix D. Example 4). Similar to the multi-layer linear case, the importance weighting ratio is a weighted sum of the intermediate layer outputs. However, unlike the multi-layer linear case, these weights can depend on the inputs.

Please let us know if you have any additional questions or want further clarification!

[1] https://openreview.net/forum?id=Zb6c8A-Fghk#

[2] https://arxiv.org/abs/2202.06856

[3] https://arxiv.org/abs/2212.07143

I have read the author's rebuttal and appreciate their efforts in addressing most of my concerns about the presentation. In light of their responses, I am inclined to revise my evaluation, raising the score to an 8. However, my revised score is a conservative estimate, oscillating between 6 and 8 due to the lack of theoretical understanding. Additionally, for greater transparency and clarity, I strongly encourage the authors to include their code's release in the paper's final version and polish their appendix.

Thank you for your thoughtful review, we appreciate your feedback!

Are all these features beneficial to the out-of-distribution generalization or downstream tasks?

[...] It seems contradictive between "SAM learns diversity features" and "SAM improved generalizability beyond in-distribution".

We apologize for any confusion and clarify that indeed feature diversity is beneficial to OOD/downstream tasks. There is no contradiction here. In fact, feature diversity is key for better OOD/downstream performance. We explain why below.

Prior works that study Waterbirds and CelebA typically treat these settings as having a “core/good” feature (e.g., hair) and a “spurious/bad” feature (e.g., gender) which is merely correlated with the label. These works only care about distribution shifts where the spurious feature in particular is no longer predictive of the label. However, our work takes a more general view where we want to tackle distributions where either one of the features (hair or gender) could possibly become unpredictive of the label. Thus, we argue that ideally we want our learned representation to have both the “hair” and “gender” feature: this is because a downstream task may require using any one of these features (either hair or gender). If both features are learned, then a linear probe can successfully pick up whichever feature is relevant downstream. We clarify that this does not contradict the “spurious correlation” literature, where the nomenclature can be somewhat misleading: there, the practice is to call the easier feature the “spurious feature” since their interest is specifically in learning the feature that SGD cannot learn (the non-spurious one). However, in the real-world, where any feature (either the easy or the hard one) can be “spurious” in that its correlation with the label may disappear in the downstream task.

To confirm that the features learned are useful for OOD/downstream tasks, we have added additional experiments. We show that SAM improves domain-transfer performance on DomainBed, a realistic dataset of domain shifts. Please see our “General Response” for more specific details about both additions.

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

We agree with the concern that it is still unclear how the learning rate effect generalizes to a realistic setup. We present this as a limitation in our conclusion section. The key challenge in achieving a similar observation of the learning-rate effect on real data is understanding precisely how features are represented. Since SAM perturbs the weights at every layer, 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 a function of how much they affect the output, we suspect that well-learned features will be inhibited by SAM at every layer.

Nonetheless, due to the challenge of directly verifying the mechanisms on real data, we need an alternative way to verify our hypothesis. This is precisely the reason that establishing an interpretable toy setup where the latent features are explicitly disentangled—here we can understand the behavior of SAM.

Please let us know if you have any additional questions or want further clarification!

Thank you for your support of our contributions. We are excited to present work that moves beyond the notion that sharpness can explain generalization. We hope that our empirical but principled insights can serve as useful guidance for future theoretical characterizations of SAM, and can motivate the development of principled algorithms.

We agree that there is currently a gap between the assumptions of the toy setting and a more realistic setting. We hope to clarify the specific challenges that we face in generalizing our results to more realistic settings. However, we want to emphasize that we do not require all of these assumptions to be held strictly, and we observe improvements in SAM even when these assumptions are broken. All discussed below:

It is unclear to what extent the analysis is applicable to learning with real images What does it take in the future to extend the theory from toy setting to reality image identification? Is it possible that the signals do not reach the output layer in the first place.

We present analysis to verify the learning rate effect in real data in Section 6.3. To summarize, we verify the two main predictions from the toy: (1) the distribution over importance weights becomes more uniform under SAM (Figure 4), and (2) SAM up-weights points that are well-fit by the easy feature and poorly fit by the complex feature (Figure 5).

We agree with the concern that it is still unclear how the learning rate effect—in which we argue that SAM balances the step size in the direction of each feature—generalizes to a realistic setup. The key challenge in achieving a similar observation of the learning-rate effect on real data is understanding precisely how features are represented. Since SAM perturbs the weights at every layer, 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 a function of how much they affect the output, we suspect that well-learned features will be inhibited by SAM at every layer. We expect that if the signal of a particular feature does not reach the output layer, SAM will promote this feature.

This challenge is precisely the reason that establishing an interpretable toy setup is an important tool for understanding.

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

We agree, as you mention, that features may not be entirely disentangled in the latent space. Nonetheless, we observe the benefits of feature learning with SAM and can identify the importance weighting effect with realistic data. This suggests that our analysis holds even when the features are not entirely separated as latent features.

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.

As mentioned above, the community does not have a solid understanding of how features are represented by neural networks. Thus it may be difficult, in general, to determine if features are disentangled. Recent work suggests that features are often at least partially disentangled (see Appendix B.4 of [1]).

However, We suspect, for the reasons enumerated above, that SAM should promote harder features even when they are not entirely disentangled.

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?

Exactly as you mention, since SAM is optimizing for minima with a small spectral norm and since $\rho$ is relatively small (in practice, we find that performance is typically maximized for $\rho$ around 0.05 or so), we suspect that the second order term is unlikely to contribute meaningfully to the dynamics.

The word "diversity" is used in this paper; however, it actually refers only to two characteristics, easy-to-fit and hard-to-fit. [...] For example, robust/non-robust features.

We call a feature “easy-to-fit” or “hard-to-fit” relative to another depending on which feature SGD learns better. In any real-world scenario, such preferential learning is bound to happen in the presence of multiple features. In your robustness example, non-robust features are known to be preferentially learned by SGD over robust features [2], thus our insights would apply there too. Consistent with this, SAM has been shown to improve adversarial robustness [3]. Besides, our results on DomainBed show that SAM indeed helps learn more transferable features in a realistic setting.

Is it possible to analyze the effect of SAM on noisy features that are typically hard-to-fit.

Great point! We have added new experiments based on this suggestion. We find that SAM improves domain-transfer performance on DomainBed, a realistic dataset of domain shifts. In addition, we have added experiments that explore variants of the toy in which we add various types of noise to the features. Consistent with our results, we find that LSAM improves performance in this setting. Please see our “General Response” for more specific details about both additions.

[1] https://openreview.net/pdf?id=Zb6c8A-Fghk

[2] https://arxiv.org/abs/2006.07710

[3] https://arxiv.org/abs/2305.05392