Title: Neglected Hessian component explains mysteries in sharpness regularization
Venue: NeurIPS 2024
Accepted: yes
Keywords: sharpness

flatness

regularization
Abstract: Recent work has shown that methods that regularize second order information like SAM can improve generalization in deep learning. Seemingly similar methods like weight noise and gradient penalties often fail to provide such benefits. We investigate this inconsistency and reveal its connection to the the structure of the Hessian of the loss. Specifically, its decomposition into the positive semi-definite Gauss-Newton matrix and an indefinite matrix, which we call the Nonlinear Modeling Error (NME) matrix. Previous studies have largely overlooked the significance of the NME in their analysis for various reasons. However, we provide empirical and theoretical evidence that the NME is important to the performance of gradient penalties and explains their sensitivity to activation functions. We also provide evidence that the difference in regularization performance between gradient penalties and weight noise can be explained by the NME. Our findings emphasize the necessity of considering the NME in both experimental design and theoretical analysis for sharpness regularization.

Neglected Hessian component explains mysteries in sharpness regularization

Abstract

Recent work has shown that methods that regularize second order information like SAM can improve generalization in deep learning. Seemingly similar methods like weight noise and gradient penalties often fail to provide such benefits. We investigate this inconsistency and reveal its connection to the the structure of the Hessian of the loss. Specifically, its decomposition into the positive semidefinite Gauss-Newton matrix and an indefinite matrix, which we call the Nonlinear Modeling Error (NME) matrix. Previous studies have largely overlooked the significance of the NME in their analysis for various reasons. However, we provide empirical and theoretical evidence that the NME is important to the performance of gradient penalties and explains their sensitivity to activation functions. We also provide evidence that the difference in regularization performance between gradient penalties and weight noise can be explained by the NME. Our findings emphasize the necessity of considering the NME in both experimental design and theoretical analysis for sharpness regularization.

1 Introduction

There is a long history in machine learning of trying to use information about loss landscape geometry to improve gradient-based learning. This has ranged from attempts to use the Fisher information matrix to improve optimization [1], to trying to regularize the Hessian to improve generalization [2]. More recently, first order methods which implicitly use or penalize second order quantities have been used successfully, including the sharpness aware minimization ( SAM ) algorithm [3]. On the other hand, there are many approaches to use second order information which once seemed promising but have had limited success [4]. These include methods like weight noise [5] and gradient norm penalties [6, 7, 8, 9, 10], which have shown mixed success.

Part of the difficulty of using second order information is the difficulty of working with the Hessian of the loss. With the large number of parameters in deep learning architectures, as well as the large number of datapoints, many algorithms use stochastic methods to approximate statistics of the Hessian [1, 11]. However, there is a conceptual difficulty as well which arises from the complicated structure of the Hessian itself. Methods often involve approximating the Hessian via the Gauss-Newton (GN) matrix - which is PSD for convex losses. This is beneficial for conditioners which try to maintain monotonicity of gradient flow via a PSD transformation of the gradient. Thus indefinite part of the Hessian is often neglected due to its complexity.

In this work we show that it is important to consider both parts of the Hessian to understand certain methods that use second order information for regularization. We study the non-PSD part of the Hessian, which we call the Nonlinear Modeling Error (NME). In contrast to commonly held assumptions, this work reveals the NME is key to understanding two previously unexplained phenomena in sharpness regularization:

Training with Gradient Penalties. We show that the performance of gradient penalties is sensitive to the choice of activation functions. Our theoretical analysis reveals that the NME is particularly sensitive to the second derivative of the activation function, and we show that gradient penalties fail to improve performance when these derivatives have poor numerical properties. We also design an intervention that can mitigate this issue. To the best of our knowledge, this work is the first to show that methods using second order information are more sensitive to the choice of activation function.

Training with Hessian penalties. Conventional analysis of weight noise casts it as a penalty on the GN part of the Hessian, but in reality it also penalizes the NME. Our experimental ablations show that the NME exerts a significant influence on generalization performance, and minimizing it is generally bad for training.

We conclude with a discussion about how these insights might be used to design activation functions not with an eye towards forward or backwards passes [12, 13], but for compatibility with methods that use second order information.

2 Understanding the structure of the Hessian

In this section, we lay the theoretical ground work for our experiments by explaining the structure of the Hessian. Given a model z ( θ , x ) defined on parameters θ and input x , and a loss function L ( z , y ) on the model outputs and labels y , we can write the gradient of the training loss with respect to θ as

where the Jacobian J ≡ ∇ θ z . The Hessian ∇ 2 θ L can be decomposed as:

where H z ≡ ∇ 2 z L . The first term, often referred to as the Gauss-Newton (GN) part of the Hessian, is a generalization of the classical Gauss-Newton matrix [14, 15]. If the loss function is convex with respect to the model outputs/logits (such as for MSE and CE losses), then the GN matrix is positive semi-definite. This term often contributes large eigenvalues. The second term has previously been studied theoretically where it is called the functional Hessian [16, 17]; in order to avoid confusion with the overall Hessian we call it the Nonlinear Modeling Error matrix (NME). It is in general indefinite and vanishes to zero at an interpolating minimum θ ∗ where the model 'fits'the data ( ∇ z L ( θ ∗ ) = 0 ), as can happen in overparameterized settings. Due to this, it is quite common for studies to drop this term entirely when dealing with the Hessian. For example, many second order optimizers approximate the Hessian ∇ 2 θ L with only the Gauss-Newton term [11, 18]. It is also common to drop this term in theoretical analysis of the Hessian ∇ 2 θ L [19, 20]. However, we will show why this term should not be ignored.

While the NME term can become small late in training, it encodes significant information during training. More precisely, it is the only part of Hessian that contains second order information from the model features ∇ 2 θ z . The GN matrix only contains second order information about the loss w.r.t. the logits with the term H z . All the information about the model function in the GN matrix is first-order. In fact, the GN matrix can be seen as the Hessian of an approximation of the loss where a first-order approximation of the model z ( θ ′ , x ) ≈ z ( θ , x ) + J δ ( δ = θ ′ -θ ) is used [18]

Thus we can see the GN matrix as the result of a linearization of the model and the NME as the part that takes into account the non-linear part of the model. The GN matrix exactly determines the linearized (NTK) dynamics of training, and therefore controls learning over small parameter changes when the features can be approximated as fixed (see Appendix A.1). In contrast, the NME encodes information about the changes in the NTK [21]. For additional intuition in the ReLU setting, see Appendix A.3.

2.1 Sharpness regularization and the NME: Case of the Gauss-Newton trace

Sharpness regularization often relies on geometric quantities of the loss landscape such as the largest eigenvalue [3] or combinations of the eigenvalues [22]. To illustrate the impact of the NME, let

us consider the sharpness measure given by the trace of the Gauss-Newton tr ( G ) = ∑ i λ i , which gives information about the average eigenvalue λ i of the Gauss-Newton. This measure, studied in Section 5, also shares some similarities to Section 4's gradient penalties but is simpler to analyze. Surprisingly, even though the quantity purposely ignores the NME, we will see that its gradient still crucially relies on it.

For GN matrix G , if the loss can be written as the log-likelihood of an exponential family distribution, this measure can be expressed as a gradient penalty [23]:

Here, the expectation is taken over labels ˆ y sampled from p ( z ) - that is, the probability distribution induced by the model logits z ( θ ) via the log-likelihood. ∇ θ L ( θ , ˆ y ) is the gradient of this loss defined with the sampled labels; derivatives are taken after the sampling process. For MSE loss, the gradient of tr( G ) can be written as:

where N ˆ y is NME of the loss defined with the sampled labels. In this case if the NME is set to 0 , gradient descent could not effectively minimize this sharpness measure. The rest of the work will explore the relationship between the structure of the Hessian and various curvature regularization techniques beyond this case through experimental and theoretical work.

3 Experimental Setup

Our analysis of the Hessian begs an immediate question: when does the NME affect learning algorithms? We conducted experimental studies to answer this question in the context of curvature regularization algorithms which seek to promote convergence to flat areas of the loss landscape. We use the following setups for the remainder of the paper:

Fashion MNIST We also include results on Fashion MNIST [24]. All experiments use the WideResnet 28-10 architecture with the same setup and hyperparameters as those used in our CIFAR-10 experiments.

CIFAR-10 We provide results on the CIFAR-10 dataset [25]. All experiments use the WideResnet 28-10 architecture with the same setup and hyperparameters as [26], except for the use of cosine learning rate decay. Batch size is 128 . Models are trained on 8 Nvidia Volta GPUs.

Imagenet We conduct experiments on the popular Imagenet dataset [27]. All experiments use the Resnet-50 architecture with the same setup and hyperparameters as [28], except that we use cosine learning rate decay [29] over 350 epochs. Batch size is set to 1024 . Models are trained on using on TPU V3 chips.

4 Explaining the pitfalls of gradient penalties

In this section we explore the sensitivity of gradient penalty regularizers to the activation function via the NME. We first establish gradient penalty regularizers as an approximation of the Sharpness Aware Minimization ( SAM ) learning rule, and then present a mystery: why do gradient penalties recover the benefits of SAM for GELU and not for ReLU? We resolve the mystery with a theoretical analysis of the NME, combined with a series of ablation studies and design a simple intervention which alleviates the issue.

4.1 SAM and gradient penalties

The SAM algorithm originates from seeking a minimum with a uniformly low loss in its neighborhood (hence flat). This is formulated in Foret et al. [3] as a minmax problem,

For computational tractability, [3] approximates the inner optimization by linearizing L w.r.t. ϵ around the origin. Plugging the optimal ϵ into the objective function yields

The SAM algorithm approximately minimizes the objective with the following learning rule:

for some step-size parameter η > 0 .

For small ρ , there is an alternative to the SAM rule. We may approximate L in (7) by its first order Taylor expansion around the point ρ = 0 as below.

Dropping terms of O ( ρ 2 ) we arrive at a gradient penalty regularizer . In general we can define gradient penalties as additive regularizers of the form

for a base loss L 0 . Gradient penalties have recently gained popularity as regularizers [6, 7, 8, 9, 10]. We will focus on the p = 1 case in the remainder of the section which we will refer to as Penalty SAM (or PSAM for short).

4.2 Penalty SAM vs Original SAM

A natural question arises: when do SAM and PSAM have similar effects? We find, surprisingly, that the answer is highly dependent on the activation function of the architecture. On Resnet50 trained on Imagenet, networks trained with GELU activation show similar performance for SAM training and PSAM training (Figure 1a); in contrast, with ReLU activation PSAM performs significantly worse than SAM as ρ is increased - indeed, becoming worse than the baseline for the best values of ρ for SAM (Figure 1b).

In order to understand this difference, we must first look at the update rule for PSAM . Given the base, pre-regularized loss L 0 , the update to the parameters θ t under SGD can be written as:

where H ≡ ∇ 2 θ L 0 . Therefore, PSAM captures curvature information via an explicit Hessian-gradient product. This is in contrast to SAM , which captures curvature information by evaluating the gradient at the alternative point in the ˜ g direction (Equation 8). This lets SAM effectively 'integrate' over in this direction and benefit from higher order information - in contrast to PSAM which only has access to the Hessian. The natural followup question is: how does the activation function influence the Hessian? In the remainder of this section, we provide evidence that it is in fact the NME component that is sensitive to this choice - and in general, it is the most important term to the overall performance of PSAM .

4.3 Effect of Activation functions on the NME

One important feature of the NME is that it depends on the second derivative of the activation function. We can demonstrate this most easily on a fully-connected network, but the general principle applies to most common architectures. Given an activation function ϕ , a feedforward network with L layers on an input x 0 defined iteratively by

The gradient of the model output x L with respect to a weight matrix W l is given by

10

3

10

2

10

1

ρ

66

68

70

72

74

76

78

80

Accuracy on Imagenet

Original SAM

Penalty SAM

(a) Imagenet with GELU

10

3

10

2

10

1

ρ

66

68

70

72

74

76

78

80

Accuracy on Imagenet

Original SAM

Penalty SAM

(b) Imagenet with ReLU

where ◦ is the Hadamard (elementwise) product. The second derivative can be written as:

where without loss of generality m ≥ l . The full analysis of this derivative can be found in Appendix A.4. The key feature is that the majority of the terms have a factor of the form

via the product rule - a dependence on the second derivative ϕ ′′ of the activation function. On the diagonal m = l , all the terms depend on ϕ ′′ . We note that a similar analysis can be found in Section 8.1.2 of [15].

We immediately see how the role of the activation function differs in the NME compared to the gradient or even the GN: only the NME is sensitive to the second derivatives of the activation. GELU has a numerically stable second derivative function:

Therefore, there are no issues computing the NME for GELU networks.

In contrast, the second derivative of ReLU is 0 everywhere except the origin, where it is undefined. In theoretical settings, the ReLU second derivative is defined as the Dirac delta 'function' - more formally, the measure attained by a sequence of Gaussians centered at zero, as the standard deviation σ → 0 . It is integrable to 1 despite not attaining non-zero value anywhere. Therefore ReLU does not have a numerically well-posed second derivative, which affects computations involving the NME.

4.4 Ablating activation NME explains the gap

10

2

10

1

ρ

66

68

70

72

74

76

78

80

Accuracy on Imagenet

ReLU

GeLU

GeLU ablating activation NME

(a) Imagenet

10

2

10

1

ρ

95.25

95.50

95.75

96.00

96.25

96.50

96.75

Accuracy on CIFAR-10

ReLU

GeLU

GeLU ablating activation NME

(b) CIFAR-10

10

9

10

7

10

5

10

3

10

ρ

94.8

95.0

95.2

95.4

95.6

95.8

Accuracy on FashionMNIST

ReLU

GeLU

GeLU ablating activation NME

(c) Fashion MNIST

1

In practical settings ReLU suffers from a 'missing curvature' phenomenology: the second derivative is set to 0 in most autodifferentiating frameworks. This means that much of the information in the NMEis inaccessible through the Hessian-vector product that at the heart of the update rule in Equation 13. Therefore PSAM suffers from a mismatch between the true NME information and the information available to the implemented algorithm. In contrast, SAM does not suffer from this issue; any curvature information is gained via differences in first derivatives - which by the fundamental theorem of calculus, are equivalent to integrals of second derivatives (and therefore of NME information).

We can confirm that the missing second derivative information in Equation 16 for ReLU networks is key for gradient penalties by removing it in GELU networks, and performing training with PSAM . This can be done by taking the second derivative of the GeLU to be zero:

and leaving the forward and backward propagation of the activation intact (see Appendix C.2 for implementation). This removes terms related to the activation in the same way ReLU does, so we denote as GeLU ablating activation NME for brevity. We see that across all 3 datasets in Figure 2 removing this portion of the NME degrades the performance of the GeLU with ρ similar to that of the ReLU, even though only the NME has been affected.

We can see similar results when we use a different method to approximate ReLU with GELU. The activation function GELU( βx ) /β converges uniformly to ReLU as β →∞ , while having well-posed derivatives for any finite β . We show in Appendix B.3 that gradient penalty performance degrades for large β > 10 3 - which we tie to high input sensitivity and increased sparsity of the second derivative. This suggests that even differentiable ReLU-like activation functions can fail to have compatibility with gradient penalties. We also confirm in Appendix C that ablating the full NME is also detrimental to generalization, though in a different way than ablating in the same manner as ReLU.

4.5 Adding synthetic activation NME improves results

10

2

10

1

ρ

66

68

70

72

74

76

78

80

Accuracy on Imagenet

ReLU

GeLU

ReLU with synthetic activation NME

Using this insight, we can modify the ReLU activation function to improve performance as follows. The issue with the ReLU is that the second derivative is a delta function - which can't be implemented numerically. However, we know that the delta function itself can be approximated by a Gaussian distribution. Therefore we replace the second derivative of ReLU with such a Gaussian

where β is the kernel width. We see that this prevents the drastic degradation at ρ = 0 . 1 of the original ReLU (Figure 3, β = 100 ). This suggests that it may be possible to design interventions to approximately access NME information in cases where second derivatives have poor numerical properties.

5 Explaining the pitfalls of Hessian penalties

In this section, we explore the impact of the NME on the effectiveness of Hessian penalties like weight noise. We first review the previously claimed link between gradient penalties and weight noise

[30], and then present a mystery: if this link is correct, why is there a difference in regularization boost between gradient penalties and weight noise? In contrast to the previous section where the NME solely featured in the update rule, weight noise implicitly minimizes the NME. We will show through ablations that minimizing the NME is detrimental, and explain why NME minimization is a poor strategy.

5.1 Weight noise equivalence assumes zero NME

We first review the analysis of training with noise established by [19]. Though the paper considers input noise, the same analysis can be applied to weight noise. Adding Gaussian ϵ ∼ N (0 , σ 2 ) noise with strength hyper-parameter σ to the parameters can be approximated to second order by

where the second term has zero expectation since ϵ is mean 0 , and the third term is a variation of the Hutchison trace estimator [31]. (We note that though the second term vanishes in expectation, it still can have large effects on the training dynamics.) [19] argues that we can simplify the term related to the Hessian by dropping the NME in Equation 2 for the purposes of minimization, which in combination with 4 yields

The argument is that for the purposes of training neural networks, the NME can be dropped because it is zero at the global minimum we are trying to reach. However, the hypothesis that the NME has negligible impact in this setting has not been experimentally verified and we address this gap in the next section.

5.2 Minimizing the NME is detrimental

In order to study the impact of the NME in this setting, we evaluate ablations of weight noise to determine the impact of the different components. Recalling Equation 20, the methods we will consider are given by

where G ≡ J T H z J is the Gauss-Newton, N ≡ ∇ z L· ∇ 2 θ z is the NME.

Hessian Trace penalty This ablation allows to us to single out the second order effect of weight noise, as it's possible the higher order terms from weight noise affect generalization. We implement this penalty with Hutchinson's trace estimator (tr ( H ) = E ϵ ∼N (0 , 1) [ ϵ T H ϵ ] ).

Gauss-Newton Trace penalty This ablation removes the NME's contribution, enabling us to isolate and measure its specific influence on the model. We use the estimator from Equation 4, tr ( G ) = E ˆ y ∼ Cat ( z ) [ ∥∇ θ L ( θ , ˆ y ) ∥ 2 ] , to compute the penalty. This is the norm of the gradients, with respect to labels ˆ y sampled via the distribution Cat ( z ) induced by the model logits via the softmax [23]. This is an alternative gradient penalty, a point we will return to later. We do not pass gradients through the sampling of the labels ˆ y , but find similar results if we pass gradients using the straight-through estimator [32]. We also run experiments with the Hutchison's estimator tr ( G ) = E ϵ ∼N (0 , 1) [ ϵ T G ϵ ] to control for the effect of the estimator.

Our experiments show that the methods perform quite differently for similar values of σ 2 (Figure 4) - confirming the influence of the NME. We can see that the generalization improvement of the Gauss-Newton Trace penalty is consistently greater than either weight noise or Hessian Trace penalty. Its improvement on Imagenet is a significant 1 . 6% . In contrast, the other methods provide little accuracy improvement. While these experiments use different estimators with a single sample, results are consistent even when compared with the same estimator and with 5 samples to get a better trace estimate (Figure 5).

1

10

9

10

7

10

5

10

3

10

1

σ

2

70

72

74

76

78

80

Accuracy on Imagenet

Gauss-Newton Trace Penalty

Hessian Trace Penalty

Weight Noise

(a) Imagenet

10

9

10

7

10

5

10

3

10

1

σ

2

92

93

94

95

96

97

Accuracy on CIFAR-10

Gauss-Newton Trace Penalty

Hessian Trace Penalty

Weight Noise

(b) CIFAR-10

10

9

10

7

10

5

10

3

10

σ

2

92

93

94

95

96

97

Accuracy on Fashion MNIST

Gauss-Newton Trace Penalty

Hessian Trace Penalty

Weight Noise

(c) Fashion MNIST

10

9

10

7

10

5

10

3

10

1

σ

2

92

93

94

95

96

97

Accuracy on CIFAR-10

Gauss-Newton penalty using Hutchinson

Hessian penalty using Hutchinson

In fact, both the weight noise and Hessian trace penalties show severe performance degradation for larger σ 2 - with the Hessian trace penalty in particular showing degradation as low as 10 -6 , at least two orders of magnitude lower than the optimal GN trace regularizer values ≥ 10 -4 . This may be related to the fact that the full Hessian trace has no lower bound as it includes the indefinite NME, while the GN is a PSD matrix whose trace is bounded by 0 . Measurements of the trace during training shows that indeed the trace grows large and negative superlinearly in the number of training iterations for the Hessian trace penalty (Figure 6, σ 2 = 10 -3 ).

Our experiments suggest that the Gauss-Newton trace is a better (and indeed, qualitatively different) regularizer than the full Hessian trace penalty. As computed in Section 2.1, this regularizer is a gradient penalty whose gradient relies crucially on the information in the NME. These results indicate that minimizing the NME in the loss is counter-productive. Similarly, Section 4.4 shows that zeroing out portions of the NME is also detrimental. This suggests that the NME is important in understanding the effects of curvature regularization.

0

50

100

150

200

Training Iterations

400000

300000

200000

100000

0

Trace on CIFAR-10

Gauss-Newton Trace Penalty

(

σ

2 =0

.

001)

Hessian Trace Penalty

(

σ

2 =0

.

001)

6 Discussion

Our theoretical analysis gives some understanding of the structure of the Hessian - in particular, the Nonlinear Modeling Error matrix. This piece of the Hessian is often neglected as it is generally indefinite and doesn't generate large eigenvalues, and is 0 at an interpolating minimum. However, the NME is the only part of the Hessian that encodes important second order information about the features, as it depends on ∇ 2 θ z - the gradient of the Jacobian. Another intriguing observation was that the gradient of the trace of the Gauss Newton matrix ∇ θ tr( G ) can be written in terms of an NME-vector product. We also saw that the NME, particularly the diagonal, is sensitive to the second derivative of the activation function.

Our experiments suggest that these second derivative properties can be quite important when training with gradient penalty regularizers. ReLU has a poorly defined pointwise second derivative, and the resulting regularizer harms training. In contrast, GELU has a well defined one and gains benefits from modest values of the regularizer. Our ablation experiments showed that removing the second derivatives prevented gradient penalties from usefully using the NME information. We also found an alternative approximation for ReLU second derivatives which added NME information and improved training.

These results suggest that some second order methods may benefit from tuning the NME. This is especially true for methods which result in Hessian-vector products in update rules (like the gradient and Hessian trace penalties studied here). Another interesting avenue for research is to replace explicit second order methods with implicit second order methods which use first order information at discrete intervals - analogous to how SAM avoided sensitivity to bad second derivatives by 'integrating' over a direction via differences, which accessed averaged quantities and higher order information.

Our experiments on Hessian trace penalties confirmed that the NME is important to understanding the successes and failures of those methods. It is intriguing that the variant which performed best, the GN trace penalty, can itself be written as an alternative gradient penalty. Exploring other gradient penalties is a promising research direction; they are non-negative, easy to compute, and generally contain NME information in their update rules.

7 Conclusion

Our work sheds light on the complexities of using second order information in deep learning. We have identified clear cases where it is important to consider the effects of both the Gauss-Newton and Nonlinear Modeling Error terms, and design algorithms and architectures with that in mind. This insight may unlock new classes of second order algorithms which use loss landscape geometry in qualitatively different ways.

References

James Martens and Roger Grosse. Optimizing Neural Networks with Kronecker-factored Approximate Curvature. In Proceedings of the 32nd International Conference on Machine Learning , pages 2408-2417. PMLR, June 2015.

Adepu Ravi Sankar, Yash Khasbage, Rahul Vigneswaran, and Vineeth N Balasubramanian. A deeper look at the hessian eigenspectrum of deep neural networks and its applications to regularization. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 35, pages 9481-9488, 2021.

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

Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Marc'aurelio Ranzato, Andrew Senior, Paul Tucker, Ke Yang, et al. Large scale distributed deep networks. Advances in neural information processing systems , 25, 2012.

Guozhong An. The effects of adding noise during backpropagation training on a generalization performance. Neural computation , 8(3):643-674, 1996.

David G. T. Barrett and Benoit Dherin. Implicit gradient regularization. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021 . OpenReview.net, 2021. URL https://openreview.net/forum?id=3q5IqUrkcF .

Samuel L. Smith, Benoit Dherin, David G. T. Barrett, and Soham De. On the origin of implicit regularization in stochastic gradient descent, 2021.

Jiawei Du, Zhou Daquan, Jiashi Feng, Vincent Tan, and Joey Tianyi Zhou. Sharpness-aware training for free. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022. URL https: //openreview.net/forum?id=xK6wRfL2mv7 .

Yang Zhao, Hao Zhang, and Xiuyuan Hu. Penalizing gradient norm for efficiently improving generalization in deep learning, 2022.

Patrik Reizinger and Ferenc Husz´ ar. SAMBA: Regularized autoencoders perform sharpnessaware minimization. In Fifth Symposium on Advances in Approximate Bayesian Inference , 2023. URL https://openreview.net/forum?id=gk3PAmy_UNz .

Hong Liu, Zhiyuan Li, David Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training, 2023.

Jeffrey Pennington, Samuel Schoenholz, and Surya Ganguli. Resurrecting the sigmoid in deep learning through dynamical isometry: Theory and practice. In Advances in Neural Information Processing Systems , volume 30. Curran Associates, Inc., 2017.

James Martens, Andy Ballard, Guillaume Desjardins, Grzegorz Swirszcz, Valentin Dalibard, Jascha Sohl-Dickstein, and Samuel S. Schoenholz. Rapid training of deep neural networks without skip connections or normalization layers using Deep Kernel Shaping. arXiv:2110.01765 [cs] , October 2021.

Nicol N Schraudolph. Fast curvature matrix-vector products for second-order gradient descent. Neural computation , 14(7):1723-1738, 2002.

James Martens. New Insights and Perspectives on the Natural Gradient Method. Journal of Machine Learning Research , 21(146):1-76, 2020. ISSN 1533-7928.

Sidak Pal Singh, Gregor Bachmann, and Thomas Hofmann. Analytic Insights into Structure and Rank of Neural Network Hessian Maps, July 2021.

Sidak Pal Singh, Thomas Hofmann, and Bernhard Sch¨ olkopf. The Hessian perspective into the Nature of Convolutional Neural Networks. In Proceedings of the 40th International Conference on Machine Learning , pages 31930-31968. PMLR, July 2023.

James Martens and Ilya Sutskever. Learning recurrent neural networks with hessian-free optimization. In Proceedings of the 28th international conference on machine learning (ICML11) , pages 1033-1040, 2011.

Chris M Bishop. Training with noise is equivalent to tikhonov regularization. Neural computation , 7(1):108-116, 1995.

Levent Sagun, Utku Evci, V Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of the hessian of over-parametrized neural networks. arXiv preprint arXiv:1706.04454 , 2017.

Atish Agarwala, Fabian Pedregosa, and Jeffrey Pennington. Second-order regression models exhibit progressive sharpening to the edge of stability, October 2022.

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

Colin Wei, Sham Kakade, and Tengyu Ma. The implicit and explicit regularization effects of dropout. In International conference on machine learning , pages 10181-10192. PMLR, 2020.

Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:1708.07747 , 2017.

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

Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. arXiv preprint arXiv:1605.07146 , 2016.

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

Priya Goyal, Piotr Doll´ ar, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in 1 hour, 2018.

Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608.03983 , 2016.

Kevin Ho and John Sum. Note on weight noise injection during training a mlp. Proc. TAAI'2009 , 2009.

Michael F Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in Statistics-Simulation and Computation , 18(3):10591076, 1989.

Yoshua Bengio, Nicholas L´ eonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432 , 2013.

Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural Tangent Kernel: Convergence and Generalization in Neural Networks. In Advances in Neural Information Processing Systems 31 , pages 8571-8580. Curran Associates, Inc., 2018.

Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha SohlDickstein, and Jeffrey Pennington. Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent. In Advances in Neural Information Processing Systems 32 , pages 8570-8581. Curran Associates, Inc., 2019.

L´ ena¨ ıc Chizat, Edouard Oyallon, and Francis Bach. On Lazy Training in Differentiable Programming. In Advances in Neural Information Processing Systems 32 , pages 2937-2947. Curran Associates, Inc., 2019.

Atish Agarwala, Jeffrey Pennington, Yann Dauphin, and Sam Schoenholz. Temperature check: Theory and practice for training models with softmax-cross-entropy losses, October 2020.

Razvan Pascanu and Yoshua Bengio. Revisiting natural gradient for deep networks. arXiv preprint arXiv:1301.3584 , 2013.

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

Atish Agarwala and Yann Dauphin. SAM operates far from home: Eigenvalue regularization as a dynamical phenomenon. In Proceedings of the 40th International Conference on Machine Learning , pages 152-168. PMLR, July 2023.

James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax .

A Hessian structure

A.1 Gauss-Newton and NTK learning

In the large width limit (width/channels/patches increasing while dataset is fixed), the learning dynamics of neural networks are well described by the neural tangent kernel , or NTK [33, 34]. Consider a dataset size D , with outputs z ( θ , X ) over the inputs X with parameters θ . The (empirical) NTK ˆ Θ is the D × D matrix given by

For wide enough networks, the learning dynamics can be written in terms of the model output z and the NTK ˆ Θ alone. For small learning rates we can study the gradient flow dynamics. The gradient flow dynamics on the parameters θ with loss function L (averaged over the dataset) is given by

We can use the chain rule to write down the dynamics of z :

In the limit of infinite width, the overall changes in individual parameters become small, and the ˆ Θ is fixed during training. This corresponds to the linearized or lazy regime [35, 36]. The NTK encodes the linear response of z to small changes in θ , and the dynamics is closed in terms of z . For finite width networks, this can well-approximate the dynamics for a number of steps related to the network width amongst other properties [34].

In order to understand the dynamics of Equation 23 at small times, or around minima, we can linearize with respect to z . We have:

where H z = ∂ 2 L ∂ z ∂ z ′ . In the limit of large width, the NTK is constant and the first term vanishes. The local dynamics depends on the spectrum of ˆ ΘH z . From the cyclic property of the trace, the non-zero part of the spectrum is equal to the non-zero spectrum of 1 D J T H z J - which is the Gauss-Newton matrix.

Therefore the eigenvalues of the Gauss-Newton matrix control the short term, linearized dynamics of z , for fixed NTK. It is in this sense that the Gauss-Newton encodes information about exploiting the local linear structure of the model.

A.2 GN and second order information

The GN part of the Hessian may seem like it must contain second order information about the model due to its equivalence to the Fisher information matrix for losses that can be written as negative log-likelihoods, like MSE and cross-entropy. For these, the Fisher information itself can be written as the Hessian of a slightly different loss [37]:

where the only difference is that the labels ˆ y are sampled from the model instead of the true labels. However, the NME is 0 for this loss. For example, in the case of MSE using Equation 2 we have

The second term vanishes because we are at the global minimum for this loss. Therefore, as expected, the Fisher information (and therefore, the Gauss Newton matrix) carry no information about the model second derivatives ∇ 2 θ z .

A.3 Intuitions about NME with ReLU

In a piecewise multilinear model like a ReLU network, we can think of the GN part of the Hessian as exploiting the linear (NTK) structure, while the NME gives information on exploration - namely, the benefits of switching to a different multilinear region where different neurons are active. The NME is made out of the sum of Dirac delta functions, whose 'spikes' are given by the boundaries between the multilinear regions corresponding to ReLU reaching the saturated regime. The NME is small outside those regions.

We can gain additional intuition by constructing a differentiable approximation of ReLU. We define the β -GELU by

where Φ is the standard Gaussian CDF. We can recover GELU by setting β = 1 . β -GELU converges uniformly to ReLU in the limit β →∞ . The second derivative is given by

For large β , this function is exponentially small when x ≫ β -1 , and O ( β ) when | x | = O ( β -1 ) . As β increases the non-zero region becomes smaller while the non-zero value becomes larger such that the integral is always 1 .

The choice of β determines how much information the NME can convey in a practical setting. This second derivative is large only when the input to the activation is within distance 1 /β of 0 . In a deep network this corresponds to being near the boundary of the piecewise multilinear regions where the activations switch on and off in an equivalent ReLU newtork.

We can illustrate this using two parameters of an MLP in the same layer, where with ReLU activation the model is in fact piecewise linear with respect to those parameters (Figure 7). For GELU with large β , the second derivative serves as an 'edge detector' 1 (more generally, hyperplane detector) of the corresponding multilinear boundaries for the equivalent ReLU network (denoted by the blue lines in the right panel of Figure 7). Therefore, the NME can be used to probe the usefulness of crossing these edges.

||

L·∇

||∇

θ

1

4

2

0

2

4

θ

2

4

2

0

2

4

L

240

260

280

300

320

4

2

0

2

4

θ

1

4

2

0

2

4

θ

2

0

10

1

10 0

10 1

10 2

z

2

θ

z

NME norm, (

β

-gelu,

β

=30

)

A.4 Nonlinear Modeling Error and second derivatives of FCNs

We can explicitly compute the Jacobian and second derivative of the model for a fully connected network. We write a feedforward network as follows:

The gradient of x L with respect to W l can be written as:

which can be written in coordinate-free notation as

If we define the partial Jacobian J l ′ l ≡ ∂ x l ′ ∂ x l , l ′ > l

Here ◦ denotes the Hadamard product, in this case equivalent to matrix multiplication by diag( ϕ ′ ( h m )) .

The Jacobian can be explicitly written as

The second derivative is more complicated. Consider

for weight matrices W l and W m . Without loss of generality, assume m ≥ l .

We first consider the case where m>l . In this case, we have

since W m comes after h l . If we write down the derivative of J L ( l +1) , there are two types of terms. The first comes from the direct differentiation of W m ; the others come from differentation of ϕ ′ ( h n ) for n ≥ m . We have:

The W m derivative projected into a direction B can be written as:

Therefore, we can write:

From our previous analysis, we have:

In total, the second derivative projected into the ( A , B ) direction for m>l is given by:

Now consider the case m = l . Here there is no direct differentiation with respect to W m , but there is a derivative with respect to ϕ ′ ( h m ) . The derivative is written as:

There are two key points: first, all but one of the terms in the off-diagonal second derivative depend on only first derivatives of the activation; for a deep network, the majority of the terms depend on ϕ ′′ . Secondly, on the diagonal, all terms depend on ϕ ′′ . Therefore if ϕ ′′ ( x ) = 0 , the diagonal of the model second derivative is 0 as well.

B SAM and gradient penalties

The gradient penalties studied in Section 4 are related to the Sharpness Aware Minimization algorithm ( SAM ) developed to combat high curvature in deep learning [3]. In this appendix we review some additional facts about SAM and gradient penalties.

B.1 USAM

A related learning algorithm is unnormalized SAM ( USAM ) with update rule [38]

USAM has similar performance to SAM and is easier to analyze [39]. The unnormalized gradient penalty equivalent PUSAM is

which corresponds to the p = 2 case of the gradient penalty.

B.2 Penalty SAM vs. implicit regularization of SGD

The analysis of [7] suggested that SGD with learning rate η is similar to gradient flow (GF) with PUSAM with ρ = η/ 4 . In this section we use a linear model to highlight some key differences between PUSAM and the discrete effects from finite stepsize SGD.

Consider a quadratic loss L ( θ ) = 1 2 θ T H θ for some parameters θ and PSD Hessian H . It is illustrative to consider gradient descent (GD) with learning rate η and (unnormalized) penalty SAM with radius ρ .

The gradient descent update rule is

The 'effective Hessian' is given by H + ρ H 2 (see [39] for more analysis). Solving the linear equation gives us

This dynamics is well described by the eigenvalues of the effective Hessian λ + ρλ 2 , where λ are the eigenvalues of H . The effect of the regularizer is therefore to introduce eigenvalue-dependent modifications into the Hessian.

There is a special setting of ρ which can be derived from the calculations in [7]. Consider ρ = η/ 2 , and consider the dynamics after 2 t steps. We have:

To leading order in η H , this is the same as the dynamics for learning rate 2 η , ρ = 0 after t steps:

We note that these two are similar only if η H ≪ 1 . Under this condition, ηρ H 2 = 1 2 η 2 H 2 ≪ η H , and the gradient penalty only has a small effect on the overall dynamics. In many practical learning scenarios, including those involving SAM , ηλ can become O (1) for many eigenvalues during training [39]. In these scenarios there will be qualitative differences between using penalty SAM and training with a different learning rate.

In addition, when ρ is set arbitrarily, the dynamics of η and 2 η will no longer match to second order in η H . This provides further theoretical evidence that combining SGD with penalty SAM is qualitatively and quantitatively different from training with a larger learning rate.

B.3 β -GELU experiments

We can understand the failure of NME information accessibility in ReLU via another method constructing a sequence of ever closer approximations of ReLU which are still differentiable. We return to the β -GELU defined in Equation 26:

We recall that lim β →∞ β -GELU ( x ) = ReLU( x ) - in fact, with uniform convergence. Note that this uniform convergence does not extend to second derivatives; β -GELU for large β is different from standard ReLU implementations at the origin, since it has second derivative β , and not 0 , at the origin.

We can then ask: does β -GELU with large β behave similarly to ReLU? Training with SGD, we see that performance is close to invariant across 6 orders of magnitude of β (Figure 8, blue curves average and standard deviation over 5 seeds). This is consistent with the uniform convergence which implies that forward and backwards passes of β -GELU become more similar to ReLU at large β .

In contrast, if we train with PSAM with ρ = 0 . 1 (the best setting we found for regular GELU with β = 1 ), we find that performance degrades with β , decreasing rapidly on a log scale for β ≥ 10 3 (Figure 8, orange). This once again suggests that choice of activation function matters, and that ReLU-like activations combine poorly with gradient penalties.

This is interesting because for any finite β , β -GELU is in fact infinitely differentiable. Why are there issues with large β then? We can return to the form of ∇ 2 θ z for the answer. Recall from Equation 15

which can be re-written as

10 0

10 1

10 2

10 3

10 4

10 5

10 6

β

66

68

70

72

74

76

78

80

Accuracy on Imagenet

SGD

SGD with Gradient Penalty

(a) Imagenet

10 0

10 1

10 2

10 3

10 4

10 5

10 6

β

92

93

94

95

96

97

Accuracy on CIFAR-10

SGD

SGD with Gradient Penalty

(b) CIFAR-10

that much of the NME depends on the second derivative of the activation function. From Equation 27, the second derivatives are exponentially small outside of a narrow band of width β -1 . Within this band, the derivatives are O ( β ) .

Therefore, for large β , the elements of the NME become 'high frequency' functions - a small change in parameter values can lead to a massive change in the function value. Indeed, the statistics are characterized by sparse, large entries. For the Imagenet and CIFAR10 examples we can measure sparsity in the second derivative of the activation directly (Figure 9). We see that at initialization, the fraction of nonzero second derivatives drops with β (blue curve); this trend becomes much stronger after training, particularly in the Imagenet case (orange curve). This gives us a possible explanation for the failure at large β ; even though the second derivatives exist, the NME becomes a sparse, high frequency estimate of local curvature information and does not provide value during training.

10 1

10 3

10 5

β

10

5

10

4

10

3

10

2

10

1

10 0

Fraction of nonzeroes in

∇

2

x σ

(

x

)

Before Training

After Training

10 1

10 3

10 5

β

10

1

10 0

Fraction of nonzeroes in

∇

2

x σ

(

x

)

Before Training

After Training

Note that we are not claiming that the choice of the activation function is a sufficient condition for gradient penalties to work with larger ρ . There are many architectural changes that can affect the NME matrix and we have shown that the statistics of the activation function is a significant one.

C NMEablation details

In this section we go into more detail about how to implement various ablations of the NME shown in the experiments.

C.1 Gradient penalty, Gauss-Newton only

1

10

2

10

ρ

66

68

70

72

74

76

78

80

Accuracy on Imagenet

GeLU

GeLU ablating full NME

For cross-entropy loss, we can completely remove the NME from the update rule in Equation 13 using a modified version of the penalty. Given model logits z ( θ ) , let p ( z ) be the probability induced by z (via the softmax). Given the true labels y , consider the quantity:

for a probability vector t . Differentiating with respect to θ , we have:

This is simply the sum of two Hessian vector products.

Consider the special case where t = p ( z ) . In this case, ∇ θ L ( θ , p ( z )) = 0 - since the choice of logits z minimizes the loss with respect to true distribution p ( z ) . The first term vanishes. This leaves us with the second term. Expanding the Hessian ∇ 2 θ L ( θ , t ) into the GN and NME, we have:

For cross-entropy loss, H z is independent of the labels, and the first term is simply the GN matrix G . For t = p ( z ) , we have ∇ z L ( z ( θ ) , p ( z )) = 0 once again. Combining, we have:

Note that we are only differentiating with respect to the first coordinate of g . This is exactly the GN-vector product, rather than the Hessian-vector product found in Equation 13.

Therefore, if we define L GNpen,p to be the version of L pen,p with NME set to 0 in the update rule, we have:

where we differentiate with respect to the first coordinate only (easily implemented in any AD framework). For arbitrary gradient penalty, recall that:

Therefore, we can implement L GNPen,p ( θ ) as

where the derivative is taken with respect to the first coordiante, and ˜ θ is set to the value θ during evaluation. Alternatively, we can apply the stop gradient operation (from the autodifferentiation framework) to the copy of g raised to the p/ 2 -1 power, and maintain a function with a single input θ only.

We note that this method works for any loss such that H z is independent of the labels.

An alternative approach is to use vjp and stop gradient operations. Consider the vector valued function ˜ g :

This can be implemented by taking a vjp of the model at θ 2 with cotangent vector ∇ z L ( z ( θ 1 ) , y ) .

We then have:

Once again, derivatives are taken with respect to the first θ , and set ˜ θ = θ on evaluation (or alternatively, use the stop gradient ).

Athird, orthogonal approach is to modify the update rule directly. That is, the second term in Equation 13 can be added explicitly to the update rule, but with the Hessian-vector product replaced with a GN-vector product. This is analogous to how for simple optimizers, ℓ 2 regularization of parameters can be added either to the loss or to the update rule as weight decay.

C.2 Implementing activation NME modifications

The experiments in Section 4.4 and 4.5 require us to define custom second derivatives for activation functions. The code below shows how to implement such a method in JAX [40], using the custom vjp capability.

We used the setting β = 100 in our experiments. This framework can be used generally to define custom second derivatives for activations.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: All claims supported by theoretical and experimental work.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

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

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: Claims are qualified.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate 'Limitations' section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: Everything is included either in the main text or the appendix.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

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

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

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

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We use standard setups that are widely available and show the values for the additional hyper-parameters.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: Publicly available datasets are used, but the code is not open source.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

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

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

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: Specified in experimental setup.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

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

Answer: [Yes]

Justification: Errors reported for all results over different random seeds.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer 'Yes' if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

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

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

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

Answer: [Yes]

Justification: See experimental setup.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have reviewed and abide by the code of ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: This is an empirical study.

Guidelines:

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

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to

generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

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

Answer: [NA]

Justification: No model is released.

Guidelines:

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

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We used widely available public datasets.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

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

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: No new assets.

Guidelines:

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

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: No crowdsourcing.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

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

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: No human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: SUMMARY:
It is known that SAM can improve generalization, while weight noises and gradient penalties often fail. This work reported that the structure of the Hessian can explain the inconsistency by identifying the key role of NME. This work first studied gradient penalties and show that methods using second order information are sensitive to activation functions due to NME. Moreover, as NME matters to Hessian penalties, weight noises which minimizes NME often behave poorly.

STRENGTHS:
- This work reports a novel and interesting conclusion that the Hessian structure, particularly the overlooked NME, can explain the poor performance of gradient penalties and Hessian penalties. This is significant for understanding the second order information in optimization and generalization of deep learning.

- The theoretical analysis seems clear and reasonable.

- This work empirically verified the theoretical results and identified the key role of NME.

WEAKNESSES:
-The empirical results only include ResNets. May the empirical conclusion also depend on model architectures? How about the results of very simple models (FCN/LR) and very complex models (Transformers)? Note that they have very different loss landscape.

-Training with weight noises only implicitly penalize the Hessian. This work often abuse Hessian penalties and weight noises. 

-This work provides insights but failed to show how to improve or design better second-order-based methods. This can further support the conclusion of this work.

QUESTIONS:
Please see the weaknesses.

LIMITATIONS:
This work did not discuss the limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-06-23

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors examine the performance of training methods for neural\nnetworks, which utilize (approximate) second order information. They\nnote that often only the curvature of the loss function is taken into\naccount (rather than that of the loss functions) and demonstrate that\nthis approach leads to drawback in training in particular regrading\nsharpness regularization such as gradient penalties / weight\nnoise. They examine the effect of including / excluding second-order information \nin the form of regularizations on generalizability based on numerical experiments.\n\nSTRENGTHS:\nThe authors demonstrate a problem in learning regularizations that has so far no been addressed, namely the varying performance of the regularizations and the mixed success of for example gradient norm penalties. The observations made have the potential to aid in the development of learning algorithms.\n\nWEAKNESSES:\nIn terms of the use of second derivatives the manuscript sends a bit of a mixed message. The initial experiment in Section 4 seems to show that accurate second-order information of the model is useful when using gradient penalties in learning algorithms, the Hessian penalties in Section 5 seem to pain a different picture, where the entire Hessian trace penalty performs the worst, so neglecting Hessian components may in fact not be such a bad idea in general.\n\nQUESTIONS:\n- 83 - 84: It would be helpful to reformulate the defintion\n  of the trace to make clear that the trace being the sum\n  of eigenvalues is an established result rather than a\n  definition of the trace operator to avoid confusion.\n- 123: The function L changes its meaning halfway through\n  the manuscript. Initially, it depends on z and y, i.e.,\n  on model output and labels. Starting with 123 it becomes\n  a function with the sole input being a vector of weights.\n  I think this is supposed to be L(z(theta, x), y) for fixed\n  training data (x, y). Please introduce another function\n  to make this clearer to the reader.\n- Some abbreviations could be given once in full in order\n  to jog the memory of readers who are not very familiar\n  with the precise context (NTK, ReLU / GELU)\n- 143: Should this be SGD (there are not batches to be seen)\n  or a normal gradient descent?\n- 226: \"Minimizing the NME is detrimental\" Shouldn't this state that \"ignoring\" the NME is a problem?\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the importance of considering second order information, specifically the structure of the Hessian of the loss, in deep learning. It decomposes the Hessian into the Gauss-Newton matrix and the Nonlinear Modeling Error (NME) matrix, with focus on the often-overlooked NME. Through empirical and theoretical evidence, the study demonstrates the significance of the NME in the performance of gradient penalties and their sensitivity to activation functions. The difference in regularization performance between gradient penalties and weight noise is also attributed to the NME. The findings underscore the need to consider the NME in experimental design and theoretical analysis for sharpness regularization, potentially leading to new classes of second order algorithms that utilize the loss landscape geometry differently.\n\nSTRENGTHS:\n**Strengths**\n\n1. The paper is clearly written and easy to follow.\n\n2. The paper focuses on an important topic in the area of sharpness learning, which is worthy of investigation.\n\n3. The paper provides a novel understanding of the connections between SAM and gradient norm penalty through the perspective of NME, offering interesting insights.\n\n4. The paper highlights the pitfalls of using different activations with the Hessian, providing valuable guidance for practical training.\n\nWEAKNESSES:\n**Weakness**\n\n1. How do you solve SAM? Have you also neglected the second-order term in your empirical analysis? What would be the effect if this term were not neglected in your situation?\n\n2. Could you provide some demonstrations regarding the off-diagonal elements in the NME? It is suggested that these elements may also play an important role.\n\n3. I recommend that the theoretical analysis be expressed in a more formal and clear style.\n\n4. There are a couple of typos: Line 133, \"$p = 1$\" -> \"$\\rho = 1$\"; Line 208, \"if this is link\" -> \"if this link\".\n\nQUESTIONS:\nSee Weakness.\n\nLIMITATIONS:\nI have not found any discussions about the limitations and potential negative societal impact. But in my opinion, this may not be a problem, since the work only focuses on understanding the sharpness learning. Still, it is highly encouraged to add corresponding discussions.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the influence of the second-order component of the Hessian in sharpness-aware minimization and optimization methods that involve gradient penalties. First, they show that the Hessian decomposes into the component of the Hessian that people usually consider [GN] and a term that includes the second derivative with respect to the parameters [NME]. From there, they demonstrate that a gap between penalty sharpness-aware minimization and general sharpness-aware minimization exists when activations are ReLU but not when activations are GeLU. They study how the NME explains this gap by showing that ablating the NME component in GeLU recovers the ReLU generalization. Further, they also show that including a term related to the NME in the ReLU case removes this gap. Finally, they ask the question of whether explicitly accounting for the NME in sharpness-aware minimization / gradient-penalty via weight noise can improve generalization of solutions, finding that it cannot in general. In the end, this work, to my understanding, deepens our understanding of the role of the usually-ignored NME term in various settings.\n\nSTRENGTHS:\n- identifies a part of the hessian that is not usually considered in sharpness-aware minimization and isolates its effect in various settings, showing that in some cases it is relevant for understanding the behavior, while in others there are reasons to not include it in algorithms.\n- theoretical analysis as well as experiments on real-scale datasets\n- appropriate ablations to study the effects of terms, mainly on outcome. I would be interested in future work looking at dynamics, as well!\n- I like that the work studies the nuances of the effects of the NME, not just tells a single story and leaves it at that.\n\nWEAKNESSES:\noverall, I feel the work studies the questions asked in detail. See some questions / writing suggestions below.\n\nQUESTIONS:\n- In section 2 when deriving the NME, can you include an explicit example where the NME is large? Either analytically or a plot demonstrating the evolution of, e.g., its Frobenius norm over the course of training.\n- small thing: perhaps when you mention the 3 datasets, reverse the order to have them in increasing order of difficulty\n- one thing you do not discuss but I would want to know more about is — PSAM method only should work when rho is small. And indeed, when rho is small, even for ReLU the methods match. Of course, since the GeLU has the match even at the values of rho where there is a gap for ReLU, it seems that the size of rho is not the determining factor. But could you still comment on how to know whether the issue is just the size of rho?\n- any comments on dynamics of training with and without the additional penalty in section 5.2? Seems like that might also be interesting.\n\nLIMITATIONS:\nthe work analyzes the effect of the NME in various settings and is clear about the particular settings which it studies.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-18"}