Paper: How Hessian structure explains mysteries in sharpness regularization (ICLR-2024)

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Theory):
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.

Author response:
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.

Reviewer comment (Evaluation):
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?

Author response:
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.

Reviewer comment (Experiments):
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?

Author response:
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.

Reviewer comment (Experiments):
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?

Author response:
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.

Reviewer comment (Theory):
**Feature exploitation and exploration analogy is interesting but also crude:** Sure, the second term in the Hessian tells more directly about feature exploration than the GN part. But the GN term during training is not fixed, but evolves. And its evolution or change is naturally dependent on the second term in the Hessian. Generally speaking, the linearization argument isn't fully correct either, as that inherently considers a fixed Hessian (like linearize around initialization or about some other point). Consequently, the whole analogy is rather crude and rife with ambiguities.

Author response:
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.

Reviewer comment (Theory):
**Explanation of gradient penalty via the activation function:** I do agree that this could be a valuable perspective in general, and indeed the various experimental results do offer a partial hint towards what may be the core issues that impede the efficacy of using gradient norm penalty for networks with ReLU. Yet, I don't think the argument is convincing enough to fully attribute this to the nature of the activation function. (a) 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, and in what manner it influences the overall Hessian. (b) The results in Figure 3 seem somewhat misleading for they do not quite match the value corresponding to ReLU ($\beta \rightarrow \infty$) shown in Figure 2. For instance, with $\rho=0.1$, the accuracy in ImageNet and CIFAR10 for PSAM with ReLU seems to be about 69% and 94% from naked eye. But that in Figure 3 seems to be much smaller. This weakens the presented argument.

Author response:
`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.

Reviewer comment (Evaluation):
**Weight noise does not work:** This section as said elsewhere is poorly written, and perhaps due to that or in general, comes across as a bit disconnected, and is also not properly fleshed out. I do get more or less the whole pretext around this, but the execution is lacking. The "roughly matches" the results of the Hessian trace penalty is a bit too rough. Besides, how are the results in the case of ReLU? Right now, with the fairly decent results, I would also not claim that weight noise does not work.

Author response:
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.

Reviewer comment (Novelty):
**Wrong citations and a general sloppiness in discussing related work:**  
- Moosavi-Dezfooli et al. 2019 citation about regularizing the Hessian to improve generalization is an incorrect in this context. They are talking about Hessian with respect to the inputs, not parameters as is the focus of this paper. It kind of betrays the rigour (or lack of it) when such a thing happens on the second citation of the paper.  
- The more accurate citation for USAM is Andriushchenko & Flammarion (2022), not what it currently comes across to be as Agarwala & Dauphin 2023.  
- "The second term to our knowledge does not have a name so we call it the Nonlinear Modeling Error matrix": It has been called functional Hessian in the works of Singh et al., 2021, 2023 analyzing Hessian rank for instance, — where this matrix has also been analyzed in theoretical detail.  
- Many second-order optimizers approximate Hessian only with Gauss-Newton: It would make the discussion balanced by adding that the reason for doing so is to also ensure guaranteed decrease in loss at each step (through a PSD precondition).

Author response:
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).