How Hessian structure explains mysteries in sharpness regularization

We thank the reviewer for their very detailed review. The comments have greatly aided us in improving our work.

Concerns with abstract and introduction:

Indeed, as suggested by the reviewer, our paper is not focused on second order methods which use Hessian information for optimization. We agree that many readers may believe we are referring to these methods and that was not our intention. We are editing the text to better reflect which sort of method we are attempting to analyze.

Concerns with section 2

The submission says that the NME "is related to exploring the effects of switching to different linear regions." Firstly, this statement can only possibly make sense in the context of ReLU networks or other piecewise linear activation functions like hardtanh. Secondly, even for ReLU networks, I don't understand what this assertion means precisely.

We agree that this statement only makes sense for piecewise linear networks (ReLU, leaky ReLU, hard tanh, etc.). The idea here is that the second derivative of these networks is given by a delta function - where the location of non-zero part of the delta function is exactly the boundaries of the piecewise linear regions. This is what the NME captures in these types of models. We will improve discussion of this point in the text.

Concerns with section 4

I think the claims of this section are not convincingly proven… at large rho, the gradient-norm-penalty algorithm on ReLU networks does not find "truly" flat regions, because even though it finds regions where the Hessian is small, the Hessian is a bad local model for the objective, and the objective turns upwards a bit farther away. However, the failure of gradient norm regularization in Figure 2 (a,b) does not look like this hypothetical. We see in this figure that the test accuracy of gradient norm penalty with large rho is smaller than that of unregularized training. This does not look like "insufficient regularization", it looks like some other kind of failure. Could the authors please spell out the precise claimed logical linkage between "quadratic approximation on ReLU nets gives bad local model" and the failure we see in Figure 2 (a,b)?

We are afraid there is a misunderstanding. For any SAM-like algorithm (SAM, gradient penalty, etc.), large values of $\rho$ eventually leads to poor training for all architectures (worse than $\rho = 0$), since practical SAM algorithms approximate the ideal SAM objective (minimax formulation at the start of 4.1). The question at hand is: why are regular SAM and penalty SAM so similar for GELU but very different for ReLU? Our hypothesis is that the structure of the NME, particularly the part relating to the second derivative of the activations, plays a role. For ReLU, the second derivative is theoretically a delta function but in practical implementations is $0$, and this implies that Hessian-vector products do a poor job of capturing local structure.

There really are a lot of potential confounders here. For large $\rho$, the dynamics of SAM and gradient norm penalty can conceivably be very different, especially at large learning rates. For example, perhaps the issue in Figure 2 (a,b) is that ReLU units die, while GeLU units do not die? Could you experiment with leaky-ReLU (which I believe should suffer less from the dead relu problem) to rule out this possibility? If leaky ReLU (with a decently large leak size) breaks down at large rho, that would support your hypothesis; if it doesn't, that would support my hypothesis.

This is an interesting hypothesis and experiment; we would like to point out that for $\beta$-gelu we also get $0$s due to machine precision, especially for larger $\beta$ (our Figure 4).

We chose to do a slightly different experiment inspired by your hypothesis. We trained with ReLU, but with a custom second derivative - a Gaussian with small width centered around $0$. This gives us some approximation of the true second derivative (dirac delta function). The results are detailed in the comment to all reviewers above, but to summarize we found that adding in this approximation to the second derivative brings the accuracy of penalty SAM above the baseline for a considerable range of $\rho$. This to us provides direct evidence that the second derivative (or lack thereof) is an important (but not the only!) factor in the failure of penalty SAM with ReLU. We also conducted an experiment where we removed the second derivative of $\beta$-gelu, which caused penalty SAM to fall below the baseline.

(response cont. below)

We thank the reviewer for their helpful comments, and address each of them below.

The article's findings rely on empirical observations rather than theoretical proofs or extensive experiments… besides, the numerical comparison is only conducted for computer vision tasks. I suggest the authors to add more experiments on various network architecture to make their conclusion more convincing.

We point out that we did have some theoretical analysis in Section 2, which gave us the structural breakdown of the NME. This was useful to develop and explain our experiments.

We have conducted additional experiments on the vision datasets; see comment to all reviewers above.

When the loss function is not smooth, we can not get its second derivative. In that case, the NME matrix may not be well-defined, how can we analyze the sharpness regularization by using Hessian structures?

The key point of our theory is that, in the case of ReLU, the second derivative is 0 pointwise almost everywhere, and is undefined at 0 (dirac delta function). In practice this second derivative is given the value 0 at the point 0 in common implementations such as in Jax. This means that for numerical methods, the NME matrix is well-defined pointwise but fails to capture the discontinuity at 0. The gradient penalty method runs but is not effective (Figure 2a/b). In contrast, for Gelu the NME correctly captures the second order information and gradient penalty method works (Figure 2 c/d).

Can the findings on the NME matrix be used to inspire the development of the second-order methods(not the penalty regularization methods)? And how to use it?

We hope to address this in future work, by co-designing activation functions, architectures, and learning algorithms.

What is Cat(z) in eq. (24)?

This is the categorical distribution with logits z. That is, we take a draw from the distribution given by softmax(z).

There are some typos in Appendix B, like ''[cite our ICML]'', ''[cite ICML 2023 paper]''.

Thanks for catching these mistakes; they have been corrected.

How the conclusion are obtained ''These results are evidence that the NME term of the Hessian should not be dropped when applying the analysis of (Bishop, 1995) to weight noise for modern networks.''? As mentioned at the beginning of section 5.2, the Hessian trace penalty consider the NME term but not perform well for both tasks. However, Gauss-Newton penalty do not consider the NME but perform better? Therefore, it seems that NME plays negative effect on this task. Could you please explain this for me?

We apologize for the confusion. The first point is that the NME is not regularized in gradient penalties - it’s involved in hessian vector products in the update rule, but the NME is not differentiated (explicitly or implicitly). See the comment to all reviewers for further discussion on this point.

With regards to Bishop: the issue is that they dropped the term in theoretical analysis but still regularized the term in implementation. In contrast, the GN penalty will drop the term both analytically and in implementation. This suggests that regularizing the trace of the NME hurts training. We will edit this in the main text.

Please let us know if we can clear up any other points for you.

We thank the reviewer for their response, and attempt to address some of their concerns below.

The experimental settings appear to lack sufficient robustness for the results, necessitating further justification to ensure the reliability of the findings. Notably, the experiments rely on only two seeds, and the approximation of expectations in Gauss-Newton and Hessian-trace penalty methods relies on a single sample.

We are confident that the results will hold up over more seeds given the magnitude of the effects compared to the error; we plan to run these experiments before the camera ready. With regards to the stochastic approximation of the GN and Hessian trace: we agree that this will lead to estimation error. We did not exhaustively explore estimation schemes because that is beyond the scope of this work. We used a single sample because that is common in the use of weight noise and often important in terms of efficiency. Please also see the "comment to all reviewers" for additional experiments.

The argument regarding the relationship between feature learning and the NME term is unclear and lacks concrete support. Including intuitive examples to illustrate this connection would be helpful.

We appreciate this concern. The intuition is that the change in the Jacobian matrix $J$ is controlled by the second derivative $\nabla_{\theta}^{2} z$; this term is present in the NME (but not in GN). On the other hand, The Jacobian $J$ is related to the features in the sense of the neural tangent kernel (NTK) which can be computed as $J J^T$. We will add this point to the main text.

The activation Hessians do not reach an exact zero when $beta$-GELU is used. What criteria are applied to consider the activation function Hessian as zero in Figure 4?

Thanks for catching this subtle point; because we use finite precision, we reach floating point 0 even though the derivative is not actually 0. In particular $\beta$-gelu goes as $\exp(-(\beta x)^2)$ for large $x$, so numerical $0$ is reached quite easily.

Given the approximations outlined in Section 5.1, the hyperparameter $\rho$ can be aligned with specific weight noise conditions for different methods in Section 5.2. Wouldn’t it then be necessary to compare the results under these matched hyperparameters to draw more precise conclusions regarding the discrepancy between methods?

We cross-validated over parameters because all the methods differ in nontrivial ways - they estimate different quantities, with different errors and biases. The cross-validation ensures we don’t miss a well-performing setting for any method.

In the case of the Hessian-trace penalty, not only $\rho$ but also the noise standard deviation in the trace estimator can be varied to control the regularization strength. Is the performance similar when both hyperparameters are controlled to achieve comparable regularization strength? If not, can the results be explained through the lens of the NME?

The Hessian trace is worse than the weight noise at all matching hyperparameter values. We omitted detailed discussion due to space constraints, but the Hessian trace penalty can rapidly drive large negative eigenvalues in the NME.

Let us know if there are any other questions we can answer.

Please let us know if we can clarify anything else! If we have substantially addressed your criticisms we hope you will consider revising your review score.

We thank the reviewer for their very thorough response; their feedback will help us improve this work tremendously.

With regards to the explore/exploit analogy: we agree that the analogy is not precise. We included it to give some insight into the differences between the GN and NME, at least in the case of ReLU and similar activations. The analogy refers to what information the GN and NME immediately surface; as pointed out by the reviewer, the time evolution of the GN depends on the second derivative of the model with respect to parameters, and therefore is related to the NME.

The off-diagonal terms in the NME will still have a highly non-trivial contribution, and these are present regardless the almost everywhere zero entries on its diagonal. There is not much of a quantification as to the precise significance of the diagonal entries of the NME.

We agree that the off-diagonal term is still present in ReLU, and may in fact contribute to learning. However we can test the effects of the second order terms directly. We defined an “augmented ReLU” function whose second derivative was a Gaussian (approximation of the true second derivative, a 0-width dirac delta). The details of the experiment are in the comment to all reviewers above; the upshot is that adding this second derivative term causes penalty SAM to achieve higher accuracy than baseline for a non-trivial range of $\rho$. We also conducted an experiment where we set the second derivative of gelu to be $0$ (diminished gelu). Removing this second derivative led to penalty SAM performing below the unregularized baseline.

We thank the reviewer for noticing the apparent discrepancy in Figure 3. This is due to the fact that at finite $\beta$ the Beta-Gelu has second derivative $\beta$ for input $0$, while the ReLU implementation in Jax has second derivative $0$ for input $0$. (This is a common choice of second derivative for ReLU.) This leads to some differences in the training dynamics between large $\beta$, large $\rho$, and ReLU + large $\rho$. We will clarify this in the paper.

We will add/fix the suggested references, and improved our overall discussion of the references. We have also improved our overall discussion of second order methods, and have noted that PSD-ness is useful to have decreasing loss for preconditioning methods (at small learning rate).

We appreciate the criticism that Section 5 is confusing. We posit that the lack of success of weight noise is well known in the community, and indeed in our own analysis weight noise does not lead to significant improvement over SGD (0.3% accuracy increase over SGD on Imagenet, with error of 0.1% over 2 samples; 0.1% increase with error 0.1% on CIFAR10). We acknowledge the issues with the presentation of this section. We have provided additional insight into section 5 in our comment to all reviewers, and we are in the process of improving the discussion in the main text.

Minor questions:

Can you add the training accuracy for the plots in Figure 2? We will add this.

How is the gradient norm penalty exactly implemented? I guess using a mini-batch? Of the same size as that for SAM? Yes, exactly - gradient is computed on one minibatch, then autodiff is used to differentiate the gradient norm on the same minibatch.

Figure 4: Is this computed across all neurons? Over the entire dataset? This is computed across all activations over a minibatch of 1024 samples.

The Hessian trace is approximated over how many samples? Hessian trace is computed over the minibatch, which is 1024 samples.

Please let us know if we can answer any more questions; happy to clarify any points we made above.