Neglected Hessian component explains mysteries in sharpness regularization
Authors: Yann Dauphin
Atish Agarwala
Hossein Mobahi
Venue: NeurIPS 2024 spotlight (2024)
Keywords: sharpness
flatness
regularization
Primary area: other
TL;DR: Understanding the neglected indefinite part of the Hessian explains important phenomena 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 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
Yann N. Dauphin Google DeepMind ynd@google.. com
Atish Agarwala Google DeepMind thetish@google.c com
Hossein Mobahi Google DeepMind mmobahi@google com
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 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 thcoretical evidence that the NME is important the performance 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.
Introduction
There a long history in machine learning of trying to use information about loss landscape geometry to improve gradient-basee learning. This has ranged from attempts matrix to improve optimization [1], to trying to regularize the Hessian to improve generalization [2]. More recently, first order methods which implicitly use penalize second order quantities have been used successfully, including the sharpness aware minimization (SAM) algorithm [3]. On the other hand, there are many approaches 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, 10], which have shown mixed success.
the Fisher information
Part of the difficulty 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, well as the large number of datapoints, many algorithms use stochastic methods to approximate statistics of the Hessian [1. 11]. However, there is conceptual difficulty well which arises from the complicated structure of the Hessian itself. Methods often involve approximating the Hessian via the Causs-Newton (GN) matrix which is PSD for convex losses. This is beneficial for conditioners which try maintain monotonicity of gradient flow via PSD transformation of the gradient. Thus indefinite part of the Hessian often neglected due its complexity.
In this work we show that 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 understanding two previously unexplained phenomena in sharpness regularization:
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

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 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 eye towards forward or backwards passes [12, 13], but for compatibility with methods that use second order information.
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 z00, x) defined on parameters and input x, and a loss function L(z, y) on the model outputs and labels y, write the gradient of the training loss with respect to 0 as
VeL Jo[=J"(V<<)
(1)
where the Jacobian J Voz. The Hessian V2L can be decomposed as:
VgC= = J VVL Viz
(2)
GN
NME
where = The first term, often referred to as the Gewssn (GN) part of the Hessian, is generalization of the classical Gauss-Newte 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 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 0' where the model "fits" data (V,(())) = 0). can happen in overparameterized settings. Due 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 V2g with only the common to drop this term in theoretical analysis of the Hessian V2L [19, 20]. However, will show why this term should not be ignored.
term [11, 18]. It is also
While the NME term can become small late in training, it encodes significant information during training. More precisely, is the only part of Hessian that contains second order information from the model features features V2z. The GN matrix only contains second order information about the loss w.r.t. the logits with the term Hg. All the information about the model function the GN matrix is first-order. In fact, the GN matrix can be seen the Hessian of an approximation of the loss where a first-order approximation the model (0'', x) RJ x) + J8 (6- 0' 0) is used [18]
(3) 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. In contrast, the NME encodes information about the changes in the NTK [21]. For additional intuition in the ReLU setting, see Appendix A.3.
(zz((( Jo, =e J
=
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 Nawssn tr(G) - E. A, which gives information about the average eigenvalue X, 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- likeliiiiiood of an exponential family distribution, this measure can be expressed as a gradient penalty [23]:
tr G) Ey-p(x) [II (o 5)]
Here, the expectation is taken over labels y sampled from p(z) that is, the probability distribution induced by the model logits z(0) ) via the log- oogliikeihoodd. VeL(e, ý) 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:
(5) where Ny is NME the loss defined with the sampled labels. In this case the NME set to 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.
Votr (G) = Fggppzz) Ny Vo
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 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 WideRes- net 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 10 dataset [25]. All experiments use the WideResnet 28-10 architecture with the same setup and hyperparameters [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 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 mystery: why do gradient penalties recover the benefits of SAM for GELU and not for ReLU? We resolve the mystery with theoretical analysis of the NME, combined with series of ablation studies and design 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 minmax problem,
min max L(e + c) s.t. ||c|| S p.
(6)
For computational tractability, [3] approximates the inner optimization by linearizing w.r.t. around the origin. Plugging the optimal e into the objective function yields
min C (( + [
VoL(O)
(7)
0
3

---

The SAM algorithm upproximately minimizes the objective with the following learning rule:
0 レ - 2- VoL (0 pg): g= VoC(O)/\Vo((O))
(8)
for some step-size parameter 17 > 0.
For small p, there is alternative the SAM rule. We may approximate in (7) by its first order Taylor expansion around the point p = 0 as below.
(0) - p( + pl C (0 VoL(O)
[ [)(( p=0
+ +0(p²)
(9)
VoC(O)
= [()) + VoL(O),
(oC(O) +0
(10)
L(0) + V + +(P)).
(11)
Dropping terms of O(p²) we arrive gradient penalty regularizer. general we can define gradient penalties as additive regularizers of the form
ppnn.p pl V l
(12)
for a base loss Cg. Gradient penalties have recently gained popularity as regularizers [6, 7, 8, 9, 10]. We will focus on the p = case in the remainder the section which we will refer to as Penalty SAM (or PSAM for short).
4.2 Penalty SAM Original SAM
natural question arises: when do SAM and PSAM have similar effects? We find, surprisingly, that the answer is highly dependent the activation function of the architecture. On Resnet50 trained Imagenet, networks trained with GELU activation show similar performance for SAM training and PSAM training (Figure la); in contrast, with ReLU activation PSAM performs significantly worse than SAM as p is increased indeed, becoming worse than the baseline for the best values of for SAM (Figure lb).
In order to understand this difference, we must first look at the update rule for PSAM. Given the base, pre-regularized loss Lu, the update the parameters O, under SGD can be written as:
0,+1 0. -1 VoLo +
|VoColl HV&Co
1
(13)
where H = Vo Therefore, PSAM captures curvature information via an explicit Hessian-grrdien product. This is in contrast SAM, which captures curvature information by evaluating the gradient the alternative point in the 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 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 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 é, feedforward network with L layers on an input X defined iteratively by
h = Wix,, Xt+1 <(hr)
(14)
The gradient of the model output x with respect to a weight matrix W, is given by
dXL J J1041) @'(n)) x1, Jr1 II o' d'(mm) W m. Φ'(x) dx
dø
(a)
aw.

---

Accuracy on Imagenet
70
Accuracy on Imaigonet
79
78
72-
10
70
&
Original BAM
S
Origina BAM
Parocty 5AM
Pareiyy SAM
P
(a) Imagenet with GELU
(b) Imagenet with ReLU
Figure 1: Test Accuracy as p increases across different datasets and activation functions averaged over seeds. PSAM with GELU networks more closely follows the behavior of SAM. For ReL.U networks and large p. there is significant difference between PSAM and SAM
where is the Hadamard (elementwise) product.
The second derivative can be written as:
(uJ+)))) d'()))+ JI((+)) aw m x
Oe' (hi)
(15)
DWOWW,
aw... m
where without loss of generality 2 I. The full analysis of this derivative can be found in Appendix A.4. The key feature is that the majority of the terms have factor of the form
(h,) = ""(ho) ow m =
Oh
d2
("(x)
(I)
(16)
OW.. m
dx
via the product rule dependence on the second derivative ø" of the activation function. On the diagonal m = 1, all the terms depend on ø² We note that similar analysis can be found in Section 8. 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 sensitive the second derivatives of the activation. GELU has a numerically stable second derivative function:
2x GELU(x)
=
-e- 02212 [2 x²
(17)
22
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 sequence of Gaussians centered at zero, the standard deviation a + It integrable to 1 despite not attaining non-zero value anywhere. Therefore ReLU does not have numerically well-posed second derivative, which affects computations involving the NME.
4.4 Ablating activation NME explains the gap
30m
Imagenet
Accuracy on FashionMNIS
CIFAR-10
GeLL
RELU
Gabtnateinng acharin NuT
85.0
whlering acttrren
06.22
08.4
72-
85.
Accuracy
RolU
06.13
98.0
GeLL
GcLU 2btaing ontivatin NME
04.8
85. 25
(a) Imagenet
(b) CIFAR-10
(c) Fashion MNIST
Figure 2: Test accuracy P increases for penalty SAM with parts of the NME ablated (average of seeds). The removal of information from the NME controls the effectiveness of the gradient penalty.
5

---

In practical settings ReLU suffers from a "missing curvature" phenomenology: the second derivative set to 0 in most autodifferentiating frameworks. This means that much of the information in the NME is inaccessible through the Hessian-vece 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:
dr.2 d² GELU = 0
(18)
and leaving the forward and backward propagation of the activation intact (see Appendix C.2 for implementation). This removes terms related to the activation the same way ReI.U does, so we denote GeLU ablating activation NME for brevity. We see that across all 3 datasets in Figure removing this portion of the NME degrades the performance of the GeLU with p 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 approximate ReLU with GELU, The activation function GELU
converges uniformly to ReLU as / + oc, while having well-posee derivatives for any finite 3. We show in Appendix B.3 that gradient penalty performance degrades for large 3 > 10 which we tie to high input sensitivity and increased sparsity of the second derivative. This suggests that even differentiable ReI.U-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 a different way than ablating in the same manner as ReLU.
4.5 Adding synthetic activation NME improves results
Accuracy Imagenet
78-
76-
72-
70-
REUU
OCLU
ReLL with wwtheii vnisss MEE
Figure 3: Test accuracy p increases for penalty SAM with parts of the NME replaced synthetically (average of seeds). The addition of information to the NME improves the performance as increases.
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 delta function which can't 't be implemented numerically. However, we know that the delta function itself can be approximated by a Gaussian distribution. Therefore replace the second derivative of ReLU with such Gaussian
& ReLÚ(x) = 22T
B
820222
(19)
d
where / the kernel width. We see that this prevents the drastic degradation at م = 0.1 of the original ReLU (Figure 3, 8 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 detrimental, and explain why NME minimization is 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 e ~~ - (0, o²) noise with strength
a to the parameters can be approximated to second order by
Ee[C(0 c)] (()) + -..VVL Eele" Eele"He] He)
(20)
[()) + (H)
where the second term has zero expectation since e 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, 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
tr H (G) Ex Eeg (o)
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 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
Hessian Trace Penalty
Eee [[(0+
((0) 2r (G) to (N +0()e13)))
Weight Noise
Gauss. Newton Trace Pemalty
where G = JJ, is the Gauss-Newton, N = VzC Viz the NME.
Hessian Trace penalty This ablation allows to 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 's trace estimator (Ir(H) 1e= Hc))).
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) = Ey--Catz) [|] (o [", to compute the penalty. This is the norm of the gradients, with respect to labels ý 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 ý. 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) = Be-ANoo,1 Ge] to control for the effect of the estimator.
Our experiments show that the methods perform quite differently for similar values of o² (Figure 4) confirming the influence of the NME. We can see that the generalization improvement of the Newss- Trace penalty is consistently greater than either weight noise or Hessian Trace penalty. Its improvement on Imagenet is a significant 1. 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 better trace estimate (Figure 5).
7

---

Accuracy Imagenet
40
Accuracy CIFAR-10
Accuracy on Fashion MNIST
78-
66
28
Gauss Newor Trase Penaty
Gouss Nowton Trace Ponarty
Bauso-Nowter Trrce Penaty Hovalan Trace Panily Weigt Nose
Herean Tracs Paraly
Herean Trace Paralyy
WeightNos No5z
Woight
7/
(a) Imagenet
(b) CIFAR-10
(c) Fashion MNIST
Figure 4: Test Accuracy as o² increases across different datasets and activation functions averaged over 5 seeds. Large 6² reveals stark contrast between the Gauss- trace penalty, which excludes NME, and methods incorporating it, highlighting the NME's influence.
Accuracy CIFAR-10
56
25
Goaso Nowcon ponalty using Hutchmsonn
Hossian pena ty us Hubtinsoon
Figure 5: Test Accuracy 2² increases with both penalties estimated using Hutchinson's estimator with 5 samples (curves are averaged over 2 seeds). Despite the larger number of samples, Hessian trace penalty remains unstable, while Gauss-Newton trace penalty is stable. This suggests the instability is not due the estimator.
In fact, both the weight noise and Hessian trace penalties show severe performance degradation for larger o² with the Hessian trace penalty particular showing degradation as low as 10 least two orders of magnitude lower than the optimal GN trace regularizer values > 10-4 This may be related the fact that the full Hessian trace has no lower bound as it includes the indefinite NME, while the GN is PSD matrix whose trace is bounded by O. Measurements of the trace during training shows that indeed the trace grows large and negative superlinearly the number of training iterations for the Hessian trace penalty (Figure 6, o² = 10-3).
Our experiments suggest that the regularizer than the full Hessian trace penalty. As computed in Section 2.1, this regularizer is gradient penalty whose gradient relies crucially on the information in the NME. These results indicate that minimizing the NME in the loss is counter-pree produer-roduetive. Similarly, Section 4.4 shows that zeroing out portions the NME is also detrimental. This suggests that the NME important understanding the effects of curvature regularization.
trace is better (and indeed, qualitatively different)
Trace CIFAR-10 10
-10000)
-300000
Training Iterations
Figure 6: Trace penalty over training iterations for different methods averaged over 5 seeds. The trace of the Hessian can be negative due to the NME, while the trace the Gauss-Newton cannot. Minimizing the Hessian trace causes it to become very negative, which is detrimental to training stability.

---

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 't generate large eigenvalues, and 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 Vgz- the gradient of the Jacobian. Another intriguing observation was that the gradient of the trace of the Gauss Newton matrix Volr(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 poorly defined pointwise second derivative, and the resulting regularizer harms training. In contrast, GELU has 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 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 promising research direction; they are non-negative, easy 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
[1] James Martens and Roger Grosse. Optimizing Neural Networks with Kronecker-factore Approximate Curvature. In Proceedings of the 32nd International Conference Machine Learning, pages 2408-24044 2417. PMLR, June 2015.
[2] Adepu Ravi Sankar, Yash Khasbage, Rahul Vigneswaran, and Vineeth N Balasubramanian. deeper look 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, 9488, 2021.
[3] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware mini- mization for efficiently improving generalization. arXiv preprint arXiv:2010.01412, 2020.
[4] Jeffrey Dean, Greg Corrado. Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Mare' aurelio Ranzato, Andrew Senior, Paul Tucker, Ke Yang, et al. Large scale distributed deep networks. Advances in neural information processing systems, 25, 2012.
[5] Guozhong An. The effects of adding noise during backpropagation training on a generalization performance. Neural computation, 8(3):643-6) 6744, 1996.
[6] 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
[7] Samuel L. Smith, Benoit Dherin, David G. T. Barrett, and Soham De. On the origin of implicit regularization in stochastic gradient descent, 2021.
[8] 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/forumm watprrmmmmddxxxxteeeee xxx66
[9] Yang Zhao, Hao Zhang, and Xiuyuan Hu. Penalizing gradient norm for efficiently improving generalization in deep learning, 2022.
[10] Patrik Reizinger and Ferene Huszár. SAMBA: Regularized autoencoders perform sharpness- aware minimization. In Fifth Symposium on Advances in Approximate Bayesian Inference, 2023. URL httttp:://openrevieee
http//formmiddddde A UNz.
[11] Hong Liu, Zhiyuan Li, David Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training, 2023.
[12] 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.
[13] James Martens, Andy Ballard, Guillaume Desjardins, Grzegorz Swirszez, 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. arX0v:2110.01765 [cs), October 2021.
[14] Nicol N Schraudolph. Fast curvature matrix- vectox products for second-order gradient descent. Neural computation, 14(7):173-11738, 2002.
[15] James Martens. New Insights and Perspectives on the Natural Gradient Method. Journal of Machine Learning Research, 21(146):1-76, 2020. ISSN 1533-7928.
[16] Sidak Pal Singh, Gregor Bachmann, and Thomas Hofmann. Analytic Insights into Structure and Rank of Neural Network Hessian Maps, July 2021.
[17] Sidak Pal Singh, Thomas Hofmann, and Bernhard Schölkopf. 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.
10

---

[18] James Martens and Ilya Sutskever. Learning recurrent neural networks with hessian- free optimization. In Proceedings of the 28th international conference on machine learning (ICML- 11), pages 1033-1040, 2011.
[19] Chris M Bishop. Training with noise is equivalent tikhonov regularization. Neural computa- tion, 7(1):108-115 116, 1995.
[20] 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: 04454, 2017.
[21] Atish Agarwala, Fabian Pedregosa, and Jeffrey Pennington. Second-order regression models exhibit progressive sharpening to the edge of stability, October 2022.
[22] Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How does sharpness-aware minimization minimize sharpness? arXiv preprint arXiv:2211.00 1.05729, 2022.
[23] 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.
[24] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:1708.077431 2017.
[25] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[26] Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. arXiv preprint 1605 0..5..77111 146, 2016.
[27] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248-255. leee, 2009.
[28] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in I hour, 2018.
[29] Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608, 03983, 2016.
[30] Kevin Ho and John Sum. Note on weight noise injection during training a mlp. Proc. TAAI'2009 2009.
[31] Michael Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in 1076, 1989.
and Computation, 18(3):0059
[32] Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv: 1308.3338.3 2013.
[33] 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 3571-8580. Curran Associates, Inc., 2018
[34] Jachoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide Neural Networks of Any Depth Evolve Linear Models Under Gradient Descent. In Advances in Neural Information Processing Systems 32, pages 8570-8581 Curran Associates, Inc., 2019.
[35] Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On Lazy Training Differentiable Programming. In Advances in Neural Information Processing Systems 32, pages 2937-2947, Curran Associates, Inc., 2019.
[36] Atish Agarwala, Jeffrey Pennington, Yann Dauphin, and Sam Schoenholz. Temperature check: Theory and practice for training models with softmax--rros-entropp losses, October 2020.
11

---

[37] Razvan Pascanu and Yoshua Bengio. Revisiting natural gradient for deep networks. arXiv preprint arXiv: 1301.3340 2013.
[38] Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware minimization. In International Conference on Machine Learning, pages 339-668. PMLR, 2022.
[39] Atish Agarwala and Yann Dauphin. SAM operates far from home: Eigenvalue regularization a dynamical phenomenon. In Proceedings of the 40th International Conference on Machine Learning, pages 152- 168. PMLR, July 2023.
[40] 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://gi //gtthb//thhbtb.. com/ jax.
12

---

A Hessian structure
A.1 Gauss-Newton and NTK learning
In the large width limit (vidth/channels/patches increasing while dataset is fixed), the learning dynamics neural networks are well described by the neural tangent kernel, or NTK [33, 34]. Consider dataset size D, with outputs z(0, X) over the inputs X with parameters O. The (empirical) NTK θ is the D x D matrix given by
G = D J
JJJ
JJJ
Oz
(21)
a0
For wide enough networks, the learning dynamics be written in terms of the model output and the NTK é alone. For small learning rates we can study the gradient flow dynamics. The gradient flow dynamics on the parameters 0 with loss function L (averaged over the dataset) is given by
VoL VoL= - -"Ve
Ô = -
(22)
We can use the chain rule to write down the dynamics of z:
Oz i 100==11-'- D
1
-
(23)
00
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 to small changes in 0. and the dynamics 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 Equation 23 at small times, or around minima, we can linearize with respect to We have:
Oh a Ōz VeL ÔH,
(24)
dz
where H,
8c In the limit of large width, the NTK is constant and the first term vanishes. The Dz8Z
local dynamics depends on the spectrum of θH,, From the cyelic property of the trace, the non-zero part of the spectrum is equal to the non-zero spectrum of JJ H 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. local linear structure of the model.
this sense that the Souss-Newton encodes information about exploiting the
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 slightly different loss [37]]
F
Ey-p. [2),(
(25)
where the only difference is that the labels ÿ 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 we have
Eg [V = Fyg [J'HJ]- + VeL(z. 8) Vi2
JJ EgoWer yJ V.2
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 V2z.
13

---

A.3 Intuitions about NME with ReLU
In piecewise multilinear model like 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 different multilinear region where different neurons are active. 'hee 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 differentiable approximation of ReLU. We define the 3-GELU by
B-GELU(Z) = (Bax)
(26)
where + is the standard Gaussian CDF. We can recover GELU by setting , = 1, ,3-GEUU converges
uniformly to ReLU in the limit / - oc. The second derivative is given by
GELU(x) dx
2288
22 (x/3-1)²
(27)
2
For large B, this function is exponentially small when x >> 8-1, and ( when | 0(8-)) As 3 increases the non-zero region becomes smaller while the non-zen value becomes larger such that the integral is always
The choice of B determines how much information the NME can convey in a practical setting. This second derivative is large only when the input the activation is within distance 1/3 of O. In 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 8, the second derivative serves as an "edge detector"' (more generally. hyperplane detector) of the corresponding multilinear boundaries for the equivalent ReL.U network (denoted by the blue lines in the right panel of Figure 7). Therefore, the NME can be used to probe the usefulness crossing these edges.
NME norm, (8-gelu, f 30)
10a
320
2
10
A PAA7'M'
C
280
260
2
240
-4
-2
0
01
0
-4 -2
01
Figure 7: Loss (left) and Nonlinear Modeling Error matrix (NME) norm (right) as a function of parameters in the same hidden layer of an MLP (MSE loss, one datapoint). For ReLU activation model is piecewise multilinear, and piecewise linear for parameters same layer. Loss is piecewise quadratic for parameters in same layer (left). There is little NME information accessible pointwise and the main features are the boundaries of the piecewise linear regions (blue, right). For 3-GELU, NME magnitude is high only within distance 1/3 of those boundaries. Therefore the NME encodes information about the utility of switching between piecewise multilinear regions.
In fact, the negative of the second order derivative of GELU is closely related to the Laplacian of Gaussian, which is well-known edge-detector detector in image processing and computer vision.
14

---

A.4 Nonlinear Modeling Error and second derivatives of FCNs
We can explicitly compute Jacobian and second derivative of the model for fully connected network. We write feedforward network follows:
h Wix,, Xt41 <(hi)
(28)
The gradient of x with respect to W can be written as:
U dx Oh;
(29)
OW.
Th oW,
which can be written in coordinate-frue free notation
dXL
dXL @ xi Oh:
(30)
oW,
we define the partial Jacobian Jp1 =
8xy 8xi
>
dxL oW;
JL4+) o (b(i) 8
(31)
Here denotes the Hadamard product, in this case equivalent matrix multiplication by (imm) The Jacobian can be explicitly written
JM II (hmm W... m
(32)
Therefore, we can write:
0xL. OW
L-1 II m==++1
o (bmm) W
o (b((hi) s xt
(33)
The second derivative is complicated. Consider
82
[Jaan @'ii)) @ xi]
(34)
dWoaw.. m
aw
for weight matrices W, and W Without loss of generality, assume m 2
We first consider the case where m > In this case, we have
dø' (h:) O, OW OW
Ox 0
(35)
since W., comes after h, If we write down the derivative of J(++))) there are two types of terms. The first comes from the direct differentiation W mi the others come from differentation o' (h,) for 2 m. We have:
0J La+1) O) (hm) ow..
@w
Oe' (ha) (m- (1(1) + llo+1) OW.. W W.J (oJ(0-1)0-1) (36)
ow
The W., derivative projected into a direction B can be written as:
OJ L(++1)
B 1) (hm)BBr BJ 1)(-1)(11)
aW.
(37)
(o(ho) W Ox
J((o+1)
B W (oJ(o-1)(41)
OW...
o=m
15

---

From our previous analysis, we have:
0J L(++1)
B JL(++1) (hm BBJ (m-1)((+1)
OW..
Oe' do'(ho)
J 1) (bo( o(m++) o o' (hm+1) 1) o Bxm))
WoJ (a- Wo ()(+))
aw m
(38)
In total, the second derivative projected into the (A, B) direction for m > 7 is given by:
(A (ABB)
JL(+11)0 [m) (m (m- 1)(+1)++
WWWW
L-1
Oe' (h)))) OW, m
Σ J 1(o+1) (o" (WoJomm4) o' (hm+1) Bxm])
W. (o- (o-1)))))
o' Axx)
(39)
Now consider the case m = Here there no direct differentiation with respect to W but there derivative with respect to o' (hm). The derivative written as:
(A 69 B) = J 1(m+1) o
10 (hm) Bxi Axm +
OW WWW m
L-1
Od (do))) W (o- 1)(m+1)
Σ JL(ot)) (b"( [WJo(m--))o o ('hhm)) Bxm])
ow,
ø Axm
(40)
There two key points: first, all but one of the terms in the off-diagonal second derivative depend on only first derivatives the activation; for deep network, the majority of the terms depend on ø" Secondly, the diagonal, all terms depend on d". Therefore if d" 0, the diagonal of the model second derivative 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 eurvature in deep learning [3]. In this appendix review some additional facts about SAM and gradient penalties.
B.1
USAM
related learning algorithm is unnormalized SAM (USAM) with update rule [38]
+ 77 VoL (0 pg) g VeC(O)
(41)
USAM has similar performance to SAM and is easier to analyze [39]. The unnormalized gradient penalty equivalent PUSAM is
(()) + / + 0
(42)
which corresponds the p 2 case 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 p 7/4. In this section we use linear model highlight some key differences between PUSAM and the discrete effects from finite stepsize SGD.
16

---

Consider a quadratic loss (()) H0 for some parameters O and PSD Hessian H. It is
illustrative to consider gradient descent (GD) with learning rate r and (unnormalized) penalty SAM with radius p.
The gradient descent update rule is
O++1 0, = -- + pH²)
(43)
The "effective Hessian" is given by H+ + H²pH²² (see [39] for more analysis). Solving the linear equation gives us
0, = (1 - 7y H pH3))
(44)
This dynamics is well described by the eigenvalues of the effective Hessian A p^², where À are the eigenvalues of H. The effect of the regularizer is therefore to introduce eigenvalue-deeendent modifications into the Hessian.
There is special setting of p which can be derived from the calculations in [7]. Consider p = n/2. and consider the dynamics after 2t steps. We have:
21
221 = 1
(H))
0.
(45)
2
which can be re-written as
1
" do
02:
27/H 17 H +
4 1"
(46)
To leading order in nHH, this is the same as the dynamics for learning rate 2n, p after t steps:
0: (1 27) 0,
(47)
We note that these two are similar only if Hj <> 1. Under this condition, noH2 = , < H,
and the gradient penalty only has a small effect on the overall dynamics. In many practical learning scenarios, including those involving SAM, πa 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 different learning rate.
In addition, when p set arbitrarily, the dynamics 1j and 2n will no longer match to second order in HH 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 3-GELU experiments
We can understand the failure of NME information accessibility in ReLU via another method constructing sequence of ever closer approximations of ReLU which are still differentiable. We return to the 3-GELU defined in Equation 26:
B-GELU(z) = (Bxr
We recall that limg- 8- GELU(Z)) = RcLU(z)- in fact, with uniform convergence. Note that this uniform convergence does not extend to second derivatives; B-GELU for large 3 is different from standard ReLU implementations at the origin, since it has second derivative 3, and not 0, at the origin.
We can then ask: does B-GELUU with large 3 behave similarly to ReL. 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 3-GELU become more similar ReLU large B.
In contrast, if we train with PSAM with == 0.1 (the best setting we found for regular GELU with 8 = 1), find that performance degrades with B, decreasing rapidly on log scale for 3 > 10 (Figure 8, orange). This once again suggests that choice of activation function matters, and that ReLU-likee activations combine poorly with gradient penalties.
This is interesting because for any finite 3, 9-GELU is in fact infinitely differentiable, Why are there issues with large 3 then? We can return to the form of V2z for the answer. Recall from Equation 15
17

---

80
97
Accuracy Imagener
Accuracy on CIFAR-10
78
96
74
95
72
94
70
68
SGD
93
SGD
66
SGD with Gradient Penally
SGD with Gradient Penally
92
10
10
16
10P
10
B
3
(a) Imagenet
(b) CIFAR-I0
Figure& 8: Accuracy vs 3 for SGD and SGD with gradient penalty (p 0.1) using 3-GELU activations (average seeds). We observe that accuracy decreases with larger 3 with the gradient penalty but not without it. As our theory suggests that the sparsity of the NME increases with 3, this is evidence that it has significant impact on gradient penaltics.
that much of the NME depends the second derivative of the activation function. From Equation 27, the second derivatives are exponentially small outside of a narrow band of width 3-1, Within this band, the derivatives are ((3).
Therefore, for large 8. the elements the NME become "high frequency" functions small change in parameter values can lead to massive change in the function value. Indeed, the statistics are characterized by sparse, large entries. For the Imagenet and CIFARIO 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 8 (blue curve); this trend becomes much stronger after training, particularly in the Imagenet case (orange curve). This gives for the failure at large 9; even though the second derivatives exist, the NME becomes sparse, high frequency estimate of local curvature information and does not provide value during training.
possible explanation
Fraction of nonzeroes in (2)())
Fraction of nonzeroes Vin(())
10a
Before Training
10-1
After Training
10-3
1.
10-
Before Training
10-5
After Training
10a
10
10a
10a
3
Figure 9: Fraction of non-zero activation second derivatives for B-GELU trained on Imagenet (left) and CIFAR1O (right). At initialization, fraction of non-zeros decreases somewhat with (blue); after training. fraction of non-zeros depends strongly on 3 (orange).
Note that we are not claiming that the choice of the activation function a sufficient condition for gradient penalties to work with larger p. There are many architectural changes that can affect the NME matrix and we have shown that the statistics of the activation function is significant one.
C NME ablation details
In this section we go into more detail about how to implement various ablations of the NME shown in the experiments.
18

---

C.1 Gradient penalty, Gauss-Newton only
Accurscy on Imagenet
78
76-
72- 70-
GaLl abbeiing full NME
Figure 10:
Test Accuracy as p increases ablating full NME from the update. We can see ablating the full NME is detrimental to performance.
For cross-entropy loss, can completely remove the NME from the update rule in Equation 13 using a modified version of the penalty. Given model logits z0)), let p(z) be the probability induced by z (via the softmax). Given the true labels y, consider the quantity:
g y, t) Vo((O,y) VoL(0.t)
(48)
for probability vector t.. Differentiating with respect to 0, we have:
Veg t) VgC(O,y)VoC((.t) + V8C(0,t\VoC((..))
(49)
This is simply the sum of two Hcssian vector products.
Consider the special case where = p(z). In this case, V VoC p(z)) 0 since the choice of logits z minimizes the loss with respect true distribution p(z). The first term vanishes. This leaves us with the second term. Expanding the Hessian (g t) into the GN and NME, we have:
VOC(O,t) HO, ()J + VzC(z(0), t) Viz
(50)
For cross-entropy loss, H2 is independent of the labels, and the first term is simply the GN matrix G. For = p(z), we have VzC(z(0), p(z)) = 0 once again. Combining, we have:
Vog 0 y. p(z)) GV
(51)
Note that only differentiating with respect 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 LGNpen,p to be the version of Epen.p with NME set to 0 in the update rule, we have:
1
LeN (0) 2 (00,,
(52)
where differentiate with respect the first coordinate only (easily implemented in any AD framework). For arbitrary gradient penalty, recall that:
VoCpen.pi ( = Vol Cpen, ( (0)p/2
2 200,200)/222 Cpen,
(53)
Therefore, we can implement LGN ( as
LaNpen. (0, (0 =
P plo - 1 p(z))
(54)
where the derivative is taken with respect to the first coordiante, and Ô is set to the value 0 during evaluation. Alternatively, we can apply the stop-gradient operation (from the autodifferentiaaioo framework) to the copy of raised to the p/2 power, and maintain a function with a single input only.
We note that this method works for any loss such that Hg is independent of the labels.
An alternative approach is to use vjp and stop gradient operations. Consider the vector valued function g:
g(0,,02) = (
(55)
19

---

This can be implemented by taking vjp of the model at 02 with cotangent vector v=L(z(01).
We then have:
LGN (, 0)
/ ( 0)
(56)
Once again, derivatives are taken with respect to the first 6, and set 0 on evaluation (or alternatively, use the stop gradient).
third, orthogonal approach 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 GN-vector product. This is analogous to how for simple optimizers, e² regularization of parameters can be added either to the loss 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 activa- tion functions. The code below shows how to implement such a method in JAX [40], using the custom vjp capability.
from jax import grad, custom_vjj _jj import numpy as jup
def cet__cusom_are act_ fn (base_ fn, second_deriv_fn)) """Return version of base_fn replacing second derivative with second_deriv_
deriv_fn.
dbase_fn_ce jup. vectorize vectorine(ggra((bseetttt) (base_fn)) derivative function
# Define custom derivative function, whose oun derivative is second d @custom_ vjp
def new new_derivxx):
return (base_fn_dxx (x)
def new vev_deriv_fwd(x))
Returns prümal output and residuals to be used in backward pass. return @base_fn_dx(x
ficconddderivfiiie
def new_deriv_bud g):
return (res g, new_deriv.e wewiii.defvj(neeeddeee bereeedddrv__eeeee new www_derrv_bblld
# Define new version of fase with custom derivatives
@custom_ _eetto__jj__oo
def mod_ oad_base_ff(x):
return base_fn(x)
def mod_baas_ffwd
return base (x))_fn(x), new new_deriv(x)
def mod_bbase_bwd( (res, g):
return (res 8,
mod_base, fn. defv (mod_ (wweemaase__aa__eeeeeeeaaaeeeeee fwd, mod_ bbas__bdd
return fa
The GELU ablation was obtained by calling get htt__ccsto__ct__ffeee lambda x: 0.), and
the ReLU custom derivative example called get_ wustto_act__ffttereee gauss_pdf_b) where
import jax.numpy numpy jnp
beta 100. beta value, can be adjusted
def gauss_pdff [xff_b(x): return (beta/ jnp sqrt (2+jnp- pi)) -jnp- exp (beta((b))))+)))))
used the setting 3 = 100 in our experiments. This framework can used generally to define custom second derivatives for activations.
20

---

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 perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long 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 not discussed in the paper.
The authors are encouraged to create 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 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might 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 address problems of privacy and faimess.
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 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant 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 complete (and correct) proof?
Answer: [Yes]
21

---

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 cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide 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 supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibiliyy
Does the paper fully disclose all the information needed to reproduce the main ex- perimental 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)?
Question:
Answer: [Yes]
Justification: We use standard setups that are widely available and show the values for the additional hyper-pane parametrassss.ss..
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, 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 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 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 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 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 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 submis- sions 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 reproduce that algorithm. (b) If the contribution is primarily
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 way to access this model for reproducing the results or 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 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., 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 instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
22

---

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.cm
commiisssoooooooo CamPollicyy for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for 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:
/ttpa///giideeeeeeeuuutiiiii / //nips. wttpccccppppccccceeeccccccccccct hetuuli/guuude....////e// /Codeeuooiiisaeeeeeee wetsssumiiiiioo.o ierriissiinnooooo oeeinniyy for more details.
The authors should provide instructions on data access and preparation, including how 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 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 possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, 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 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, confi- dence intervals, statistical significance tests, at least for the experiments that support 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, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
23

---

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 96% CI, if the hypothesis of Normality of enrors 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 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
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
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 well 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 into the paper).
9. Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: /ttp///eerp///e//ee//ee httpppic/tttisccuiiddeo.ooo
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 cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations 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 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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that 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 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.
there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition 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
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: No model is released.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators 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?
Question:
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 dataset.
The authors should state which version of the asset is used and, if possible, include 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 (c.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, paperswithcom thcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's is creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
25

---

Answer:
[NA]
Justification: No new assets.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/coddemoodll as part 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 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 the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included 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 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 expect authors to adhere 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.
26


Reviews:
[{"review_id": "mlVTvhogDG", "reviewer": "Reviewer_NVNp", "summary": "This 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.", "questions": "- 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.", "limitations": "the work analyzes the effect of the NME in various settings and is clear about the particular settings which it studies.", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- 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.", "weaknesses": "overall, I feel the work studies the questions asked in detail. See some questions / writing suggestions below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thorough review. We address the reviewer's questions and suggestions below.\n\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\nThe suggestion to identify specific examples where the NME is large is valuable. We are actively investigating this, but due to time constraints in preparing the rebuttal we don’t have definite results yet.\n\n_small thing: perhaps when you mention the 3 datasets, reverse the order to have them in increasing order of difficulty_\n\nWe will change the order as you suggested.\n\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\nFurther evidence that the size of rho isn’t the sole factor can be found in Figure 3, which shows the gap for ReLU can be narrowed even at larger rho values by using a synthetic activation NME.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GuV1aYLSeA", "reviewer": "Reviewer_ZfRd", "summary": "This 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.", "questions": "See Weakness.", "limitations": "I 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.", "rating": 7, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "**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.", "weaknesses": "**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\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their insightful review. We address the reviewer's questions and suggestions below.\n\n_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\nOur experimental results for SAM are based on the SAM algorithm proposed by Foret et al. Specifically, the SAM update does not utilize any Hessian (but does evaluate an extra gradient). The reason the Hessian is absent from the SAM update is due to the approximation employed by Foret et al (using stop_gradient on the adversarial perturbation). Therefore, our implementation of SAM, which aligns with that of Foret et al, indeed neglects the Hessian. We have not investigated the implications of retaining the Hessian in the SAM formulation because this was already explored by Foret et al. They conducted this experiment and observed that by keeping the Hessian, SAM's performance slightly decreases compared to when the Hessian is neglected (which is consistent with our story). This is discussed in Figure 4 of their paper.\n\nTo derive PSAM from SAM Equation 6, we employ a first-order approximation as shown in Equations 9-11. This differs from the approximation used in the original SAM algorithm [Foret et al., 2021]. While exploring a second-order approximation could be a promising avenue, it falls outside the scope of this current work. \n\n_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\nThe experiments in Section 4.4 show that the off-diagonal elements of the NME by themselves are not sufficient to obtain good results with PSAM. In these experiments we artificially zero out the diagonal elements of the NME of GELU.  It would be interesting to see how well PSAM with ONLY the diagonal elements performs, but we did not have time in the rebuttal period to design and run this experiment.\n\n_I recommend that the theoretical analysis be expressed in a more formal and clear style._\n\nThank you for your suggestion. We will improve the style of the theoretical analysis and address the typos you identified.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LxtNgdteyw", "reviewer": "Reviewer_PSYG", "summary": "The 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.", "questions": "- 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?", "limitations": "NA", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The 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.", "weaknesses": "In 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments, and will fix the errors noted. We address some selected concerns below.\n\n_In terms of the use of second derivatives the manuscript sends a bit of a mixed message._\n\nThe overall point about second derivatives is a bit subtle; our work suggests that second derivatives in the _update rule_ are helpful (get full NME information in HVP). In contrast, the full Hessian trace penalty has second derivative information _in the regularizer_. Minimizing this information during optimization reduces the effect of second derivatives in the _update rule_. It also introduces _third derivatives_ into the update rule; these higher order derivatives in the updates seem to hurt training. It remains an open question whether or not other forms of higher order derivatives are useful.\n\n_143: Should this be SGD (there are not batches to be seen) or a normal gradient descent?_\n\nWe wrote the rule for a single batch; in practice this update rule would be combined with batching/SGD.\n\n_226: \"Minimizing the NME is detrimental\" Shouldn't this state that \"ignoring\" the NME is a problem?_\n\nThis should read “Minimizing the NME trace” is detrimental, as the GN trace penalty (no NME in the regularizer) works well.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wP1szVDxPK", "reviewer": "Reviewer_L6WX", "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.", "questions": "Please see the weaknesses.", "limitations": "This work did not discuss the limitations.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "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.\n\n- The theoretical analysis seems clear and reasonable.\n\n- 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.\n\n-Training with weight noises only implicitly penalize the Hessian. This work often abuse Hessian penalties and weight noises. \n\n-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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their questions, and answer some here.\n\n_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._\n\nWe agree that studying more architectures and datasets would be interesting; for this initial work we focused on very detailed experiments in our chosen settings rather than shallow experiments across a broad range of settings.\n\n_Training with weight noises only implicitly penalize the Hessian. This work often abuse Hessian penalties and weight noises._\n\nIn Sections 5.1 and 5.2, we review the links made _in the literature_ between weight noise and Hessian penalties in order to frame our work. However, in all our experiments we make the procedure clear, and mainly focus on explicit regularization of the Hessian/Gauss Newton trace.\n\n_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._\n\nWe agree that additional research is needed to further improve second order methods/regularizers; however our work showed that the Gauss Newton trace is another useful regularizer. In addition, our synthetic second derivative approach in Section 4.5 shows how we can overcome the poor second derivatives of ReLU in a practical, efficient way.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper investigates the inconsistencies observed in the performance of sharpness-aware minimization compared to other regularization techniques like weight noise and gradient penalties in deep learning. The authors reveal that these inconsistencies are connected to the structure of the Hessian of the loss function, specifically its decomposition into the Gauss-Newton matrix and an indefinite matrix termed the Nonlinear Modeling Error (NME) matrix.

Following the fact that all referees were positive regarding this work, I recommended acceptance to NeurIPS 2025. Nevertheless, I recommend following the recommendation of the referees for the camera-ready version, especially :

- identify specific examples where the NME is large is valuable

- rewrite the parts where reviewers saw inconsistencies

- investigate the role of off-diagonal elements.
Decision: Accept (spotlight)