Noise Stability Optimization for Finding Flat Minima: A Hessian-based Regularization Approach
Authors: Hongyang R. Zhang
Dongyue Li
Haotian Ju
Venue: Accepted by TMLR ()

Abstract: The training of over-parameterized neural networks has received much study in recent literature. An important consideration is the regularization of over-parameterized networks due to their highly nonconvex and nonlinear geometry. In this paper, we study noise injection algorithms, which can regularize the Hessian of the loss, leading to regions with flat loss surfaces. Specifically, by injecting isotropic Gaussian noise into the weight matrices of a neural network, we can obtain an approximately unbiased estimate of the trace of the Hessian. However, naively implementing the noise injection via adding noise to the weight matrices before backpropagation presents limited empirical improvements. To address this limitation, we design a two-point estimate of the Hessian penalty, which injects noise into the weight matrices along both positive and negative directions of the random noise. In particular, this two-point estimate eliminates the variance of the first-order Taylor's expansion term on the Hessian. We show a PAC-Bayes generalization bound that depends on the trace of the Hessian (and the radius of the weight space), which can be measured from data.

We conduct a detailed experimental study to validate our approach and show that it can effectively regularize the Hessian and improve generalization. First, our algorithm can outperform prior approaches on sharpness-reduced training, delivering up to a 2.4% test accuracy increase for fine-tuning ResNets on six image classification datasets. Moreover, the trace of the Hessian reduces by 15.8%, and the largest eigenvalue is reduced by 9.7% with our approach. We also find that the regularization of the Hessian can be combined with alternative regularization methods, such as weight decay and data augmentation, leading to stronger regularization. Second, our approach remains highly effective for improving generalization in pretraining multimodal CLIP models and chain-of-thought fine-tuning.

Published in Transactions on Machine Learning Research (09/2024)
Noise Stability Optimization for Finding Flat Minima: A Hessian-based Regularization Approach
Hongyang R. Zhang
ho.zhang@hhorthesster...l
Northeastern University, Boston
Dongyue Li
l.. H.donggnu@horrteesstee...
Northeastern University, Boston
Haotian Ju
ju.h@hortheastemm.co edu
Northeastern University, Boston
Reviewed on OpenReview:
Abstract
The training of over-parameterized neural networks has received much study in recent lit- erature. An important consideration is the regularization over- wer-parameterized networks due to their highly nonconvex and nonlinear geometry. In this paper, we study noise injec- vion algorithms, which can regularize the Hessian of the loss, leading to regions with flat loss surfaces. Specifically, by injecting isotropic Gaussian noise into the weight matrices of neural network, we obtain approximately unbiased estimate of the trace of the Hessian. However, naively implementing the noise injection via adding noise to the weight matrices before backpropagation presents limited empirical improvemenss. To address this limitation, we design a two-point estimate of the Hessian penalty, which injects noise into the weight matrices along both positive and negative directions of the random noise. In particular, this two-point estimate eliminates the variance of the first-order Taylor's expan- sion term the Hessian. We show PAC-Bayes generalization bound that depends the trace of the Hessian (and the radius of the weight space), which can be measured from data.
We conduct a detailed experimental study to validate our approach and show that it can effectively regularize the Hessian and improve generalization. First, our algorithm can out- perform prior approaches
sharpness-roducet training, delivering up to 2.4% test curacy increase for fine-tuning ResNets six image classification dalasets. Moreover, the trace of the Hessian reduces by 15.8%, and the largest eigenvalue is reduced by 9.7% with our approach. We also find that the regularization of the Hessian can be combined with alternative regularization methods, such as weight decay and data augmentation, leading to stronger regularization. Second, our approach remains highly effective for improving generalization in pretraining multimodal CLIP models and chain-of-thoughtt fine-tuning.
1 Introduction
The loss landscape and its geometry properties are recurring theme in the study of neural networks (Keskar et al., 2017; Hochreiter & Schmidhuber, 1997). Recently, the design of training methods such sharpness-aware minimization and stochastic weight averaging has led to empirical advances in wide variety of setlings (Izmailov et al., 2018; Foret et al., 2021; Wortsman et al., 2022). The Uheoretical study of these training methods has also been explored (Andriushchenko & Flammarion, 2022). For instance, recent work shows that sharpness-aware minimization (Foret al., 2021) has an implicit bias flat surface regions by penalizing the largest eigenvalue of the loss Hessian matrix (Wen et al., 2023; Bartlett et al., 2023). In this paper, we study methods that can provide explicit regularization of the trace of the Hessian, and we will

---

Published in Transactions on Machine Learning Research (09/2024)
(Wl + U
W
U;
UUi
Vf(
UD
Figure 1: An illustration of one update step in algorithm. At each iteration i, we sample random variable U, from vero-mean distribution P (e.g., an isotropic Gaussian with variance o²), where is a hyper-parameter that controls the strength of the noise injection (hence the regularization). We query the gradient of f, at f(W, + U,), and f(W, U,), and take their average. This results a two-point noise injection scheme, whose computation cost is the same as aharpness-aware minimization (Foret et al., 2021),
and twice the cost of running SGD. Notice that in practice, we can also implement an extension of this algorithm, which samples multiple Us. For details, see Algorithm 1.
show provable generalization guarantees of our methods. More formally, given input function Rd that represents the empirical risk of neural network and d-dimensional distribution P with mean zero, we consider minimizing the noise-perturbed function
R R
F(W)
E [f(W U]]. UNP
(1)
Minimizing this perturbed function can improve the resilience of the neural network to noise injection, leading to flatter loss surfaces and improved regularization (Nagarajan Kolter, 2020). By analyzing the perturbed loss of a fine-tuned model, one can identify measure of the sharpness of loss surfaces based on the trace of the Hessian (Ju et al., 2022; 2023). We remark that the minimization problem of the form (1) traces back to carlier works randomized smoothing (Duchi et al., 2012), which have provided a detailed study of convergence rates for nonsmooth stochastic optimization. Our work differs from this line of lileralure in that we focus on evaluating the regularization effect of penalizing the Hession trace upon neural network training.
Although noise injection algorithms can be theoretically motivated as improving generalization (and sta- bility), its practical implication is not evident (Hinton & Van Camp, 1993; An, 1996; Graves, 2011). To motivate our study, we begin by running several empirical studies to compare the performance of (standard) SGD and weight-perturber SGD (WP-SGD), which first injects random noise into the weight matrices of neural network before computing its gradient in SGD. As mentioned above, this would provide a randomized smoothing effect to the loss surface (Duchi et al., 2012). We will fine-tune (pretrained) ResNets on three image classification tasks for this empirical study. To ensure the robustness of the analysis, we also vary the distribution of P and the variance of U. Our overall finding is that WP-SGD (or randomized smoothing) does not offer clear benefits over SGD, which also consistent with recent studies of weight noise injection (Orvicto et al., 2023; Dauphin et al., 2024) (soc Section 2.2, Table 2 for the complete results). However, we hypothesize that these results may be due to the randomness of the noise injection (upon the Hessian penalty term) rather than the ineffectveneess of regularizing the Hessian trace.
Our approach to mitigate the randomness of the noise injection on the Hessian penalty involves two parts. First, we retrieve the gradient at W - U to cancel out the first-order expansion term of W + U (recall that U is a random sample from P). Meanwhile, the second-order expansion term remains the same after this cancellation. We term this modification a two-point noise injection scheme, which is reminiscent of two-point gradient estimates in zeroth-order optimization (Duchi et al., 2015). The difference in our setting is that this two-point averaging cancels out the first-order gradient term, thereby eliminating its variance the Hessian
2

---

Published in Transactions on Machine Learning Research (09/2024)
Table 1: Comparison between our approach (NSO) and SAM (Foret et al., 2021), based on inductive bias, gencralization guarantee, and convergence rate. In particular, the inductive bias of SAM is based on the results of Wen et al. (2023). The list of notations used in the table is explained as follows. 4 refers to the Hessian matrix of the loss function f. {1 and Tr[:] refer to the largest eigenvalue and the trace of an input matrix. refers to the trace norm, taken over the maximum of the entire hypothesis space and data distribution (including unseen test data). is the radius of the fine-tuning region measured in & distance. n is the number of samples in the training dataset. T is the total number of iterations run by our algorithm.
Approach
Inductive Bias Generalization Guarantee
Convergence Rate
Sharpness-Awart Aware Minimization (SAM) Noise Stability Optimization (NSO)
T
(Theorem 2.1)
(Section 4)
penalty. Second, we sample multiple perturbations U1, U,,U.....U U3 at each epoch and take their averaged two-point (noisc-injected) gradients. See Figure 1 for an illustration of one step.
A primary advantage of our approach compared to prior sharpness minimization algorithms is that our approach can provide an approximately unbiased estimate of the Hessian trace. We empirically validate this claim across three real-world scttings (sec Figure 2, Section 2.2 for an illustration). By utilizing this property, we show a PAC-Bayes bound that depends on the trace of the Hessian and the radius of the weight hypothesis space. We briefly describe this result, leaving a formal statement to Theorem 2.1. Tet a be an upper bound on the trace of the Hessian measured within the hypothesis space and the data distribution (in practice, one may take this as the union of training and testing data). Let r be the radius of the hypothesis space, measured in l distance. Suppose there are n. empirical samples from an unknown distribution. We show generalization bound that scales as O 0o2 Our proof utilizes linear PAC-Bayes bound (Catoni, 2007: McAllester, 2013; Alquier, 2021), and we optimize the variance of the prior and posterior distributions to derive the result. A detailed proof sketch is presented in Section 2.3.
Next, we validate our approach with a detailed empirical study. First, we compare our approach with four prior approaches for the setting of fine- finettuiigg pretrained ResNets, including sharpness-awan minimization (Foret et al., 2021), tested on six image classification datasets. We show that our algorithm can reduce the trace and the largest cigenvalue of the loss Hessian matrix by 15.8% and 9.7%, respectively. Our approach also improves test accuracy by 2.4%. Second, we show that by combining our approach with regularization methods such as data augmentation and distance-based regularization (Gouk et al., 2022), we can further regularize the Hessien, lcading to 13.6% lower trace values and 16.3% lower test loss values (averaged over six (asks). Third, we extend our approach to pretraining and chain-of-though fine-tuning. The details can be found in Section 3.2. Overall, our algorithm can consistently provide better regularization of Hessian and improved test accuracy across these different settings and datasets. Some of these empirical results are not completely explained by our theory, and we discuss the limitations in Section 7.
Finally, we analyze the convergence of our algorithm using techniques from the stochastic optimization literature (Ghadimi & Lan, 2013; Lan, 2020; Carmon et al., 2020; Drori & Shamir, 2020; Zhang, 2023), leading to matching upper and lower bounds. We also present a case study of Hessian regularization in over- parametrized matrix sensing and show that it is equivalent to nuclear norm regularization for this setting. Our work raises several new questions that may be worth revisiting: can acceleratod gradient descent methods be applied to design flat-minima optimizers? Can recent advances in zeroth-order optimization be leveraged to better regularize the training of transformer neural networks?
In summary, the contributions of this paper are thrco-fold. First, we present an algorithm that can explicitly regularize the Hessian frace and show a PAC-Bayes generalization bound that could be measured from data. Second, we conduct experiments on multiple settings to validate our approach by comparing downstream performance and Hcssian statistics with prior sharpness minimization algorithms and alternative regulariza- tion methods. Third, we analyze the convergence of our algorithm using stochastic optimization techniques. In Table 1, we highlight the key aspects of our approach compared prior approaches.
3

---

Published in Transactions on Machine Learning Research (09/2024)
Organization: The rest of this paper is organized as follows. In Section 2, will present our approach. We will start by presenting the motivating experiments. Then, we describe our algorithm and PAC- bound that depends on the Hessian. In Section 3, we present our experiments for validating the proposed approach. Section 4 presents an analysis of the convergence rate. Section provides a case study of the regularization effect of the Hessian trace in the over-paramelerrveed matrix sensing problem. In Section 6, we discuss the related works. Finally, in Section 7, we state the conclusion and the limitations of this work. We provide complete proof of our theoretical results in Appendix A-B, We provide additional experimental details in Appendix C.
2 Our Approach
In this section, we present our approach. First, to set up the stage, we will study the straightforward implementation of noise injection by directly adding noise to the weight matrices of the neural network before computing the gradients in backpropagation. We term this procedure weight-perturbed SGD (or WP-SGD in short), also known randomized smoothing (Duchi et al., 2012). We will compare the empirical performance of these two approaches to evaluate the effect of noise injection. Then, we describe our algorithm and provide empirical measurements of the trace of the Hessian, along with the actual perturbation gaps observed in practice. Finally, we will show a PAC-Bayes generalization bound that depends on the trace of the Hessian, which can be measured from data to compare different methods.
2.1 Motivating Experiments
We compare the results from running WP-SGD standard SGD. We choose the setting of fine-tuning pretrained foundation models, as overfitting is a common problem for this setting (Wortsman al., 2022), and strong regularization needed (Li & Zhang, 2021; Ju et al., 2022). We will fine-tune a pretrained ResNet-34 several image classification datasets, including aircraft recognition (Aircraft) (Maji et al., 2013), indoor scene recognition (Caltech-256) (Griffin et al., 2007), and medical image classification (retina images for diabetic retinopathy classification) (Pachade et al., 2021). To implement WP-SGD, we sample a random vector from P and add it to the model weights at each ileration before computing the gradient. We set P as the isotropic Gaussian and adjust its standard deviation between 0.008, 0.01, and 0.012 via cross-validation.
We report our results in Table 2, which indicate that the performance gap is less than 0.5%, R 0.75 standard deviations based on five independent runs. Furthermore, varying P does not change the results. In particular, we test four types of P, including Gaussian, Laplace, uniform, and Binomial. We adjust standard deviations between 0.008, 0.01, and 0.012 via cross-validation. We find that using Laplace and uniform distributions achieves a performance comparable to that of Gaussian. However, using Binomial results in worse results. These experiments suggest that the straightforward implementation of noise injection does not offer clear benefits over SGD.
2.2 Description of Our Algorithm
In our approach, we make two modifications to WP-SGD. First, we add the perturbation from both the positive and negative directions during the noise injection, as shown in line 5. Second, we average over multiple noise injections to reduce the variance from noise injection, as described in line 7. As for the first modification, recall that P is symmetric distribution. We use Taylor's expansion on both f(W + U) and f(W U) as follows:
(U) = - + 0(1
(2)
- -C. - +01E)
J(W - -)=
(3)
4

---

Published in Transactions on Machine Learning Research (09/2024)
Table 2: Comparing the outcome of running WP-SGD to standard SGD across four different P, measured over three image classification datasets. Recall that WP-SGD refers to normal weight perturbation (without the paired perturbation). To be concise, we have included the results of running our approach (i.e., NSO). All the results and their standard deviations are based five independent runs.
Aircraft
Indoor
Retina Discase
p
Train Acc.
Test Acc.
Train Acc.
Test Acc.
Train Acc.
Test Acc.
SGD
None
100.0% = 0.0
59.8% + 0.7
100.0% + 0.0
76.0% = 0.4
100.0% = 0.0
61.7% + 0.8
WP-SGD
Gaussian
98.4% + 0.2
60.1% + 0.1
99.0% + 0.3
76.3% + 0.0
100.0% + 0.0
62.3% + 0.5
WP-SGD
Laplace
98.3% + 0.1
60.3% 0.3
98.9% +: 0.1
76.4% 0.3
100.0% 0.0
62.0% +: 0.1
WP-SGD
Uniform
98.6% +:2 0.3
60.3%
+ 0.5
98.6% + 0.3
76.6% + 0.1
100.0% + 0.0
62.3%
+ 0.0
WP-SGD
Binomial
19.6% + 0.1
0.1 11.3%
18.2% + 0.9
10.7% = 0.1
58.1% = 0.1
57.1% + 0.0
NSO
Gaussian
95.8% +:: 0.4
62.3% + 0.3
95.7% 1 0.2
77.4% 1 0.3
100.0% + 0.0
66.6% 1 0.7
NSO
Laplace
96.5% + 0.3
61.9% + 0.3
96.1% + 0.3
77.1% = 0.1
100.0% 0.0
65.9% 1 0.1
NSO
Uniform
96.4% + 0.4
61.9% 0.5
96.4%
+ 0.2
76.8% + 0.2
100.0% + 0.0
65.7% 1 0.1
NSO
Binomial
20.1% + 0.1
14.3% + 0.3
22.8%
+ 0.1
17.9% = 0.2
59.2% 0.1
57.8% + 0.1
We have that E [0) = 0, and E [UUT]] = Σ. Thus, by taking the average of equations (2) and (3), we get
! ((((+ + f(W - -D) = F(W) = UND
f(W) - (-W)) + (1I (
E
(4)
We can see that the two-point estimate eliminates the first-order gradient term, potentially reducing its vari- ance in estimating the Hessian term. The second modification reduces the variance of the stochastic gradient, using the fact that each perturbation is independent of the others. The entire procedure is summarized Algorithm 1. As remark, two-point gradient estimators are commonly used in zeroth-order optimization (Duchi et al., 2015). However, their use in designing flat minima optimizers has not been explored much.
Algorithm 1 Noise stability optimization (NSO) for regularizing the Hessian of neural networks
Input: Initialization Wo E Rd, a function ʃ Rdd + R Require: An estimator R - Rd that for any W, returns g(W) s.t. E Parameters: # perturbations k, # epochs T, step sizes U/T-11
1: for = 0......, T do
2: / Compute the two-point averaged stochastic gradient for each independent noise injection
3: for j == 0, 0,1,,..,4 --1 do
0,0) + sampled independently from P
c(i) + g(W + U(1) + g(W, - (((i)) 6: end for
5:
W+1 +- W, 7i 2k -
end for
Mcasurements of the Hessian trace and the perturbation gap: Next, we provide several examples to measure the approximation quality of equation (4). Following the experimental setup of Section 2.1, we will fine-tune foundation model a downstream task. After training, we will set W as the model weight at the last epoch for all the measurements.
To measure equation (4), we then add U to W, where U is sampled from an isotropic Gaussian. We will measure f(W + ((W), averaged over 100 independent samples of U, and we measure this and 2 (W) by taking the average over the training dataset.
The results are shown in Figure We can see that f provides an accurate approximation F(W) (W) a
for various values of a. In particular, the approximation error of equation (4) using the Hessian trace is less
5

---

Published in Transactions on Machine Learning Research (09/2024)
than 3%. As remark, the range of o² differs across architectures because of the differing scales of their weights. More details about the neural network architectures can be found in Appendix C.
x10
MLP
x10
BERT
x10
GNN
Gap
Gap
Gap
Trace x
Trace x
Trace x
2-
0.020
0.025
0.030
0.0070
0.0075
0.0080
0.040
0.045
0.050
a
0
Figure 2: Illustration of the approximation quality of equation (4), We report all measurements based on the network weight at the last epoch of fine-tuning. We can see that the perturbation gap (i.e., F(W) f(W) in equation (4)) and _- [(W)] are at the same order. Recall that a refers to the standard deviation of the Gaussian noise injected into the weight matricos. More specifically, will decide the strength of noise injection the strength of regularization on the Hessian trace.
2.3 Generalization Guarantee and Proof Sketch
Next, we present. PAC-Bayes generalization bound that depends on the trace of the Hessian. Our bound be related to the notion of trace norm, which has been used in earlier works for quantifying sample complexity in the context of matrix recovery (Srebro & Shraibman, 2005).
Concretely, suppose we have a pretrained model in the fine-tuning setting. This can be viewed as our prior belief the target hypothesis in PAC-Bayes analysis. Once we have learned a model (though fine-tuning), we can view this as the posterior in PAC-Bayes analysis. Let D S X x be an unknown data distribution, sup- ported on the feature space and the label space y. Given random samples (I1, f1), (I2, (22,92).... 92)..... ..... (Ia Un) drawn from D, the empirical loss (measured by loss function f) applied to model fw (with W E TRP) is:
L(W) Σ ((J
((w((a)) (ax).
The population loss is L(W)- B(z.)) [e(sw (1) is sufficient to think that the empirical loss is less than the population loss, and the goal is to bound the gap between i(() and L(W) from above (Shalev- Shwartz & Ben-David, 2014).
Let W be any learned hypothesis within the hypothesis space, denoted as H. Our generalization bound will apply uniformly W within the hypothesis space. We state our result, including the required assumptions, as follows.
Thcorem 2.1. Assume that the loss function & is bounded between O and C for fired constont C > on the dala distribution D. Suppose ((Jw(-),-) wnice-differentialle W and the Hessian matria (-)((Jw(-),)) is Lipschitz continuous within the hypothesis space. Suppose for any W in H. the trace norm of the Hessian is less than
a: = max max Tr [V W WEH (2.D)-
(5)
and the Ez-norm of W is at most r for any W E H. Then, for any W in H. with probability at least for any 8 0, the following must hold, for any close to zero:
8
Car² S (W))((+0)))) (1+)) -1+
0 1og(6-1))
(6)
n

---

Published in Transactions on Machine Learning Research (09/2024)
Proof Sketch: We provide a high-level illustration of the proof of Theorem 2.1. Let Q denote the posterion distribution. Specifically, we consider Q as being centered at the learned hypothesis W (which could be anywhere within the hypothesis space), given by a Gaussian distribution NNW,, o² Idp), where ldp denotes the p by p identity matrix. Given a sample U ~N o² Idp), let the perturbed loss be given by
(o((w()), - BB [e([wwo(<):)))
Then, let (W) be the averaged value of ():) taken over n empirical samples from the training datasct. Likewise, let (W) be the population average of (e(fw()) (),,), in expectation over an unseon data sample from the underlying data distribution.
Having introduced the notations, start with the linear PAC-Bayes bound (Catoni, 2007; McAllester, 2013; Alquier, 2021) (see Theorem A.1 for referenco), stated as follows, which holds with probability 1 1- 5 for any 8 E (0, 1) and B & (0, 1): Le(W) - 5
1 LWW) + C(KI(Q) P) + l)g(o=)))
(8)
2 -8)n
where P refers to the prior distribution, C refers to the upper bound the loss value l. For analyzing fine-tuning, we view P centered at the pretrained model, with covariance matrix o² ldp. By Taylor's expansion of lg (see Lemma A.4 for the precise statement), we show that:
σ²
Lo(W) L(W)+ +
E (M.g)>D
Tr ([]l
(9)
2
0² Tr (f.)) +0
Lo(W)= L(W)+ +
(10)
2n
Since the Hessian operator is Lipschitz contunuous by the assumption of Theorem 2.1, we can bound the gap between the above two quantities with e-covering arguments (see Lemma A.5 for the precise statement). By plugging in these results back to the PAC-Bayes bound of equation (8), after some calculation, can get:
1 [(W)+
(1
Cr2/20²
+0 -0(0²+
0² V' +0 Vn
log(d
L(W) S
(11)
23
2 B)m
n
In particular, the above uses the fact that the fn-nrrm of W is less than for any W E H (the KI. divergence is discussed in Proposition A.2). By choosing 0² and 3 to minimize equation (11), we will obtain equation (6). This summarizes the high-level proof idea. The complete proof can be found in Appendix A.1.
Remark 2.2. We highlight two key aspects of our results. The first that our PAC-Bayes bound is non- vacuous, meaning that it matches the scale of empirically observed gaps when measured in practice; this is based the trace measurements in Figures 2 and 3. The second is that this non-vacuous bound has practical implications, meoning that we can ulilize this bound to design oplimization algorilhmss that improve generalization.
These are non-trivial to achieve. To give some context, prior work has provided a PAC-Bayes margin bound for multi-layer neural networks, which depends on the product of the spectral norm of the network layers (Neyshabur et al., 2018). While this paper provides important insights regarding the generalization deep nelworks, the bound is vocuous when measured in practice. Arora el al. (2018) provide another data-dependent PAC-Bayes bound based on compression techniques. Their work started with an experiment in which they injected noise into the network layers and showed that deep nets can absorb the noise after retraining. However, their bound remains orders of magnitude higher than the actual generalization errors observed in practice.
In contrast, our bound matches the scale of empirically observed gaps. To achieve this, we start from the line of work on data-dependent PAC- MCC-Bayes bounds. We build on the line of work on distance from the initialization (Nagarajan & Kolter, 2020), which is ideal for underslanding fine-tuning (Tá &s Thang, 2021).
7

---

Published in Transactions on Machine Learning Research (09/2024)
Our key breakthrough is to connect noise stability in PAC-Bayes bound with the loss Hession matrix (cf. equations (9) and (10)). Then, we can measure the Hession of loss landscapes from dala.
We additionally note that few existing works have considered using PAC- C-Bayes bounds to design algorithms.. The reason is that for new algorithm designs, we need to connect the BACeB bound with data in a non- vacuous way. The work of Dziugaite & Roy (2017) has provided compulational framework to achieve non- vacuous generalization bounds. Instead, our result provides an analytical expression that can be leveroged in algorithm design. To operationalize the design, we utilize the explicit dependence of our result the Hessian to design the regularization scheme.
3 Experiments
We now turn empirical validations of our algorithm. First, we apply our approach to fine-tune pretrained ResNets on various image classification datasets. We find that NSO can more significantly regularize the Hessian of the loss surface, resulting in reductions in the trace and the largest eigenvalue by 15.8% and 9.7%, respectively. After controlling computation costs, it can outperform four sharpness-reducing meth- ods by up to 2.4%. In addition, we justify our algorithm design through detailed ablation analysis. We also show that our approach is compatible with alternative regularization techniques, including distance- based regularization and data augmentation, and combining these methods with our approach leads to more significant regularization and test performance. Second, we show similar results for pretraining and chain- of thought fine-tuning The experiment code for reproducing our empirical findings can be found online at: httm:///rttcooeeeaaaa//lllliooio aarieeesssarrrrneeeeeeeeeeeeeeee Nod.sssstabbii1tt--ppmitatttt
https://githu/: httechu...om/vitttteeeeee
3.1 Comparison with Sharpness Minimization Methods
We compare Algorithm 1 with five sharpness-reduucing training methods, including sharpness-aware minimization (SAM) (Foret et al., 2021), unnormalized SAM (USAM) (Agarwala & Dauphin, 2023), adaptive variants of SAM (ASAM), and random SAM (RSAM) (Liu et al., 2022). During the comparison, we control for the same amount of computation (for Algorithm 1, we will set the number of sampled injections & as 1). Thus, all the methods under consideration will use twice the computation of SGD. For NSO, we sample perturbation from an isotropic Gaussian distribution and adjust o between 0.008, 0.01, and 0.012. For SAM, we adjust the & norm of the perturbation between 0.01, 0.02, and 0.05. For each method, we run it with both momentum and weight decay. We ensure that all the training methods carefully adjusted. See Appendix C for the details.
3.1.1 Empirical Findings
In Table 3, we report the comparison between NSO, SGD, SAM, unnormalized SAM (USAM), and adaptive SAM (ASAM). We find that our approach reduces the trace of Hessian by 15.8% on average. The largest eigenvalue of the Hessian is also reduced by 9.7%. This finding is intriguing since SAM has been motivated by a min-max problem. As for test accuracy, our approach can provide up to 2.4% lift, with an average improvement 1.2%. Additional comparisons are deferred to Table Appendix C.
Figure illustrates the measurements between SGD, WP-SGD, and NSO. Curiously, we find that the trace of the Hessian also decreases for SGD, possibly due to implicit norm control of SGD. While both WP-SGD and NSO reduce the trace of the Hessian, our approach penalizes the Hessian more. Besides, the generalization gap and the test loss are consistently lower during NSO training.
As a remark, the regularization effect of noise injection should be orthogonal to training methods such as momentum, weight decay, learning rate scheduling, etc. To this end, we performed comparisons without using either momentum or weight decay. Our approach can again reduce the trace of the Hessian by 17.7% compared to the five tharpness-reducing methods on average, with up to 1.8% higher test accuracy.
8

---

Published in Transactions on Machine Learning Research (09/2024)
Table 3: Comparison between our approach (NSO) with SGD, sharpness-aware minimization (SAM), un- normalized SAM (USAM), and adaptive SAM (ASAM). We fine-tune the ResNet-34 network on six image classification datasets and report the test accuracy and the trace of Hessian using the model in the last epoch of training. The results are averaged over five random seeds.
CIFAR-10
CIFAR-100
Aircrafts Caltech-256
Indoor
Retina
# Training
45,000
45,000
3,334
7,680
4,824
1,396
# Validation
5,000
5,000
3,333
5,120
536
Basic
Statistics
# Test.
10,000
10,000
3,333
5,120
1,340
250
# Classes
10
100
100
256
67
5
SGD
4128 + 83
13188 + 221
5471 + 65
3674 +
3629
28607 226
Trace (4)
SAM
2429 + 87
9227 + 286
4499 + 70
3285 + 95
3159 75
15444 = 173
USAM
2352 + 61
7382 + 222
4298 + 94
3174 = 52
3072 51
12068 + 246
ASAM
2445 1 63
9960 + 313
4475 + 69
3339 +
3014 53
14155 136
NSO
1728 =
5244
3678
2958 77
2737 + 90
10970 146
SGD
96.1% + 0.1
82.8% + 0.1
60.5% + 0.7
80.0% + 0.1
76.7% = 0.4
62.2% + 0.8
Test Acc. (t)
SAM
97.0% 0.2
84.0% 0.4
62.3%
0.3
77.0%
77.2% = 0.3
65.0% 0.3
USAM
96.9% + 0.2
83.7% + 0.2
61.9% + 0.3
76.9% + 0.2
76.7% = 0.3
64.7% + 0.1
ASAM
97.1% + 0.1
84.2% + 0.3
62.4% + 0.5
77.3% + 0.2
77.2% = 0.2
65.2% 1 0.3
NSO
97.6% = 0.4
84.9% + 0.3
63.2% + 0.3
78.1% + 0.5
78.2% 0.3
67.0% = 0.4
ReaNet 34
104
ResNet 34
RosNet 34
0.8
2.0
Generalizatoon Gap
0.6
WP-SGD
Test Loss
0.7
NSO
Trace
1.5
0.4
0.6
1.0
0.2
0.5
0.0
0
10
20
30
10
20
10
20
30
f
BERT Base
10
BERT Base
BERT Base
1.4
Generalization Gap
0.9
SGD
WP-SGD
1.2
NSO
Trace
0.7
0.6
Tost
1.0
0.4
0.3
0.8
0.1
0.0
0
4
0
4
2
t.
Figure Comparison between SGD, WP-SGD, and NSO for fine-tuning ResNet-34 and BERT-Base re- spectively, on an image and text classification dataset. We evaluate the lest loss, the trace of the Hessian, and the generalization gap for the trained model at each epoch. For WP-SGD and NSO, we sample noise from isotropic Gaussian with standard deviation a 0.01 in both settings.
3.1.2 Ablation Analysis
Next, conduct ablation studies of modifications in our approach: the use of negative perturbations and the sampling of multiple perturbations.
Comparing using negative cancellation or not after controlling computation costs: our algorithm uses negative perturbations to zero out the first-order term in Taylor's expansion of F(W). We validate this by comparing the performance between using or not using the negative perturbation. We control for the same amount of computation costs to ensure fair comparison. In particular, we sample two independent perturbations and take their averaged stochastic gradient. We find that using the nega-
Recall that
9

---

Published in Transactions on Machine Learning Research (09/2024)
tive perturbation achieves 3.6% improvement in test accuracy (on average) over not using the negative perturbation, Le.,, randomized smoothing.
As discussed in Section 2.2, our intuition on why NSO can be expected to generalize better than randomized smoothing is that can better regularize the Hessian. In particular, even though, in theory, the expectation of J(W++ + +U) and ( + + -U) over are both equal to F(W). However, the two-point scheme cancels out the gradient expansion term compared randomized smoothing at every epoch. More precisely, believe that the improved regularization from our approach stems from its better estimate of the Hessian penalty term. As illustrated in Figure 3, NSO consistently reduces the trace of the Hessian and achieves lower generalization errors compared to randomized smoothing throughout model training. At the end of the training, NSO yields 10.6% smaller trace of the Hessian on average than randomized smoothing.
Increasing the number of noise injections Recall that increasing the number perturbations k can reduce the variance of the estimated gradient. Thus, we consider increasing & in NSO and compare that with specific implementation of WP-SGD that uses the same amount of computation. Using k = 2 perturbations improves the test accuracy by 1.2% average compared to k = 1.
Varying the learning rate and the number of epochs. We provide a detailed comparison between NSO and WP-SGD when varying the learning rate and the number of epochs. The learning rate is varied be- tween 0.0001, 0.0002, 0. 0.0005, 0.001, 0.002, and 0.005. The number of epochs is varied between 10, 20, 30, 40, 50, and 60. We report the test loss from running an image classification task in Figure 4.
WP-SGD
1.8
WP-SGD
Test loss
NSO
Test loss
NSO
1.6
1.4
1.2
16-4 20-4 4 55-4 1e-3 20-3
5
10
20
30
40
50
60
Learning rate
Number of epochs
Figure 4: Results of varying the learning rate and the number of epochs for running our approach and WP-SGD, We report the test loss from the last epoch and average the results over five random seeds.
Remark 3.1 (Noise variance scheduling as k increases). A natural question whether one can gradually increase or decrease the regulorization strength by during training, similar to learning rate scheduling. To facilitate this discussion, we lest two scheduuess for adjusting a. The first. schedule is to increase to specified value at a linear rate. The second schedule exponentially increases to reach specified value. Our preliminary experiments show that neither schedule offers significant performance improvements over using a constant noise voriance. One might also consider other scheduling schemes; we leave this to future work.
3.1.3 Detailed Comparison with Sharpness-Award Minimization (SAM)
Varying the radius of SAM: We provide detailed comparison to SAM by varying the perturbation radius of SAM (denoted as p). To illustrate this comparison, we vary P between 0.001, 0.002, 0.005, 0.01, 0.02, and 0.05. We report both the validation accuracy and the trace of the Hessian for SAM and unnormalized SAM on an image classification dataset. We present the results in Table 4. We observe that using smaller p (i.e., less than 0.01) results in worse results. Thus, choose p between 0.01, 0.02, and 0.05 in our experiments.
Varying the batch size of SAM: size. In particular, we vary the batch size between 8, 16, 32, and 64 for fine-tuning ResNet-34 on two image classification datasets. The results are shown in the leftmast two panels of Figure 5. We use the same number of epochs for each batch size configuration to ensure a fair comparison. On the indoor dataset, our approach is less sensitive to different batch sizes than SAM. Across all the batch sizes and datascts, our approach consistently provides more robust regularization of the Hessian compared SAM. The best results are achieved when the batch size is 32. Thus, this particular setting in our experiments.
Next, measure the sensitivity of our approach concerning the batch
10

---

Published in Transactions on Machine Learning Research (09/2024)
Table 4: Results varying the perturbation radius of SAM (denoted as p) and unnormalized SAM. We report both the test accuracy and the trace of the Hessian based on the model trained at the last epoch. We report the averaged results and their standard deviations across five random seeds.
P
0.001
0.002
0.005
0.01
0.02
0.05
Trace (1)
SAM
4920 + 158
4347 + 166
4016 + 80
3918 + 94
3159 + 75
3028 + 78
Unnormalized SAM
4352 + 169
3990 70
3723 + 87
3427 + 57
3072 + 51
3048 22
Test Accuracy
SAM
73.6 + 0.2
74.4 + 0.4
74.8 + 0.6
75.2
0.3
76.6
0.5
73.8 + 0.7
(1)
Unnormalized SAM
74.1 0.1
74.1 0.7
74.7 = 0.5
74.6 = 0.3
76.3 = 0.3
73.1 0.6
1.4
SAM
2.2
SAM
1.6
dfst zgg
1.6
w/o data aug.
NSO
N80
dist, reg.
w/ data sug.
Test loss
1.2
Test loss
1.8
Test Loss
Test Loss
1.2
1.2
1.0
1.4
0.8
1.0
0.8
0.8
32
32
20
30
20
30
Batch size
Batch size
(a)
Loss on Indoor dataset
(b) Loss
Aircraft dataset
(c) Test loss, w/ dist. reg.
(d) Test loss, w/ data aug.
10P
x10#
x10
x10*
SAM
SAM
w/o dist.
w data aug.
N8O
ASO
cist. reg.
w/ aug.
Thoce
Traon
Thaoe
Thooe
3-
4
3
16
32
64
16
32
64
0
10
20
30
10
20
30
Batch size
Batch size
(e) Hessian, Indoor dataset
(I) Hessian, Aircraft dataset
(g) Hessian, w/ dist. reg.
(h) Hessian, w/ data aug.
Figure 5: Results of varying the batch size of our approach and SAM
two image classification datasets
(indoor scene recognition and Aireraft detection). We report the test loss and the trace of Hessian using the model from the last epoch of training. The results are averaged over five random sceds. The regularization provided by noise injection can be combined with distance-based regularization and data augmentation to reduce the test loss and the Hessian trace.
3.1.4 Combining Algorithm 1 with Alternative Regularization Methods
In this section, we show that the regularization of the Hossian serve complement to existing, alterna- live regularization methods. To validate this, combine our training approach with data augmentation and distance-based regularization (Gouk et al., 2022). In particular, the latter approach has been used to regular- ize fine-tuning algorithms. We use a popular scheme for data augmentation that applies random horizontal flipping and random cropping sequentially to each training image. As for distance-based regularization, we penalize the f distance between the fine-tuned model and the pretrained initialization.
The results are shown in Figure within the two rightmost panels. Combining our approach with each regularization method further reduces the trace of the loss Hessian matrix by 13.6% (on average). This further leads to 16.3% lower test loss of the fine-tuned network, suggesting that our approach can be used on top of these preexisting regularization methods.
3.2 Applying Algorithm to Pretraining and Fine-tuning
We apply our approach to pretraining randomly initialized models by replacing SGD to train contrastive language-image (CLIP) models dataset of image-caption pairs. In particular, use the Conceptual
11

---

Published in Transactions on Machine Learning Research (09/2024)
Caption dataset, which contains 3.3 million image caption pairs. Each caption briefly describes the corre- sponding image, with ten tokens on average. We use a 12-layer Vision Transformer as the image encoder and a 12-layer GPT-2 transformer as the text encoder. We train the encoders jointly to maximize the cosine similarity between the embodding of image caption pairs following the protocol of Radford et al. (2021).
Table presents the results. For each algorithm, we evaluate the trace of the loss Hessian and recall scores (of the top-10 scored images in retrieving images from texts) on the development set. The results show that our approach can reduce the trace of the Hessian by 17% compared to both SAM and SGD. In addition, our approach achieves 1.4% higher recall scores in image retrieval.
Lastly, we apply our algorithm to fine-tuning pretrained language models on thain-of-thought reasoning datasets. The lask is to generate the reasoning process, i.e., a chain of thoughts and the answer for a given commonsense reasoning question. We fine-tune pretrained GPT-2 models on two question-answering datascts: Commonsense QA and Stratogy QA. Table shows that our approach can yield 25% lower trace values than SAM and SGD. In addition, we can obtain 5.3% higher test accuracy.
Table 5: Results for pretraining CLIP the Conceptual Caption and chain-of-thought fine-tuning on Com- monsense/Stratee QA. We report the recall score of image retrieval/test accuracy and trace/Ar using the model at the last epoch. We report the averaged results and standard deviations over five random seeds.
Conceptual Caption
Trace (L)
X1 (1)
Recall@10 (t)
SGD
220 + 24
41 + 2.8
36.1% + 0.3
SAM
144 + 20
30 + 1.1
36.9% + 0.4
NSO
119 + 34
22 + 1.2
37.5% + 0.3
Commonsenso QA
Trace (L)
A1 (4)
Test Accuracy (t)
SGD
372 + 34
19 + 0.8
27.7% = 1.8
SAM
288 + 15
15 + 0.3
32.7% +\2 1.4
NSO
208 + 31
13 + 0.6
39.2% + 1.4
Strategy QA
Trace (L)
X1 (4)
Test Accuracy (t)
SGD
294 + 13
44 + 1.5
68.9% +\2 1.0
SAM
249 + 33
42 + 2.6
71.1% + 1.2
NSO
193 + 31
33 + 1.8
75.2% + 1.2
4 Convergence Rates
We now study the convergence of Algorithm 1. Recall that our algorithm minimizes f(W) plus regulariza- tion term the Hessian trace. As is typical with regularization, the penalty is usually small relative to the loss value. Thus, we aim to find stationary point of F(W) instead of f(W) because otherwise, we would not have the desired regularization. We state the convergence to an approximate stationary point such that [[V((W)] is small, building on the following gradient oracle assumption (see, e-g., Ghadimi & Lan (2013); Duchi et al. (2015)).
Assumption 4.1. Given a random seed 2, let 9= : Rdd Rdd be a continuous function that gives an unbiased estimate of the gradient: Ez [9 (W)) Vf(W), for any W E Rd. Additionally, the variance is bounded in the sense that Ez [[] (f))- ((W)) S o²
To help understand the above assumption, suppose there is a dataset of size n. Then, in SGD, the stochastic gradient would be unbiased estimate of the gradient of the entire dataset. As for the variance of the gradient estimator, note that as long the l norm of the gradient remains bounded, which will always hold in practice, then the last equation of the above assumption will hold. We now state an upper bound on the convergence rate of Algorithm 1.
12

---

Published in Transactions on Machine Learning Research (09/2024)
Proposition 4.2. Suppose Assumption 4.1 holds. Let P be distribution that is symmetric at zero and let HP)) = |[ Let C and D) be fized, posilive constants. Let Wo E Rd denole an arbilrary inilialization. Suppose F(Wo) min F(W) S D², and Vf is C-Lipschitz continuous. There erists a fixed learning rate 1) < C such that if we run Algorithm with l ll for all for T steps, the algorithm returns W (where is a random integer between 1,2,..., T), such that in expectation over the randommess of W,:
E [V(())) & 2 (P)(P))
2CD T
(12)
KT
As remark, existing sharpness-reduccing methods such as SAM seem to suffer from oscillation around the local basin (Bartlett et al., 2023). Thus, the convergence behavior of SAM seems challenging to analyze for nonconvex functions. By contrast, our algorithm is amenable to stochastic optimization techniques. Our proof slightly extends the proof of Theorem 2.1, Ghadimi & Lan (2013), to tackle noise injection and other variations. For details, see Appendix B.1.
Lower bounds: Next, we construct an example to match the rate of equation (12), essentially showing that this is tight under the same set of assumptions. We use an example from the work of Drori & Shamir (2020). The difference is that we have deal with the perturbations added the objective. For = 1.... ......... let e Rd be the basis vector in dimension d, whose i-th coordinate is while the remaining coordinates are all zero. Let f: Rd R be defined as
T-1
f(W) 2G
(W, eo)2 2
Σ
(WW,)))))),
(13)
where h is a piece-wise quadratic function parameterized by as, defined as follow:
Co I
|x| S aii
c(=)-=a))
+ Caf
S | < 3 0s:
h₂(x)
c (12)-2a))
3 a S |x| S 20
20, < |z].
One can verify that for each piece above, Vhg is C-Tipschitz. As a result, provided that G s C-1, Vf is C-Lipschitz, based on the definition of f in equation (13).
The stochastic function F requires setting the perturbation distribution P. We set P by truncating an isotropic Gaussian N(0, o² Ida) that the i-th coordinate is at most 2 2- for i 1,...,7 T. Additionally, set the initialization Wo to satisfy (Wo, = 0 for any i 2 while (Wo, 6o) + 0. Finally, gradient oracle to satisly that the i-th step's gradient noise & = (Ei, e141)e++1, which means that & is along the direction of the basis vector e++1. In particular, this implies only coordinate i i 1 is updated in step i, as long as (Ei, 1--1) I 2- With this construction, we state the lower bound below.
choose the
Theorem 4.3. Let the learning rates .o...., 11 be at most C-1, Let D O be fixed value. When either Er- 14 S VKT, T 11 < C-1 for any epoch i, then for the above construction, the following must hold
min IE [IVF(W))]]] 2 DD 32k T
C
(14)
1
We remark that the above construction requires T < d. Notice that this purely for technical reasons. We briefly illustrate the key steps of the proof. At step i, the gradient noise plus the perturbation noise is less than a + 2- = a; at coordinate + (by triangle inequality). Thus, d,(,W,, e+1)) = 0, which holds for all prior update steps. This implics
VS(W,) G (We, 6o).
13

---

Published in Transactions on Machine Learning Research (09/2024)
Recall that F(W)) < D2. This condition imposes how large the a,'s can be. In particular, will set 2//0/ /ag//ve /v in the proof. Then, based on the definition of f(Wo),
Caf since (Wo + U,e+1) S as. 4
In Lemma B.3, then argue that the learning rates in this case must satisly Σi= "I S O(VT). When the learning rate is fixed and at least 1/2"1//3,, construct a piece-wise quadratic function (similar to equation (13)), now with a fixed a. This is described in Lemma B.4. In this case, the gradient noise grows by 1- up to T steps. We then carefully set a to lower bound the norm the gradient. Combining these two cases, we conclude the proof of Theorem 4.3. For details, see Appendix B.2. As is typical in lower-bound constructions, our result holds for a specific instance (with a particular learning rate range).
The proof can also be extended adaptive learning rate schedules. Notice that the above construction holds for arbitrary learning rates defined as function of previous iterates. Then, set the width of each function he, Ot, proportional to 0, for any Ne that may depend previous iterates, as long as they satisfy the constraint that Ei=o TJs S O(VT).
Extension:
We can show similar lower bound for the momentum update rule. Recall this is defined as
M,+1 M.M n,G. and Wi+1 W. M4+1,
(15)
for i 0, .,, 1, where G2 is the specific gradient step i. To handle this case, will need fine-grained control on the gradient, result and ils proof to Appendix B.3.
consider quadratic function as f(W) [|]] We leave the
Remark 4.4. The novelty of results lies in analyzing sharpness minimization using techniques from stochastic optimization. This oppears to be new, and hope work can inspire further studies, such designing accelerated sharpness minimization methods.
5 Regularization Effect of Hessian Trace in Over-Parameterized Matrix Sensing
Before proceeding, let give an example of the regularization effect of penalizing the Hessian trace. We consider the matrix sensing problem, whose generalization properties are particularly well-understood in the nonconvex factorization setting (Li et al., 2018). Let there be an unknown, rank-r positive semi-definite ma- trix X = U E dRd The input consists of a list of d by d Gaussian measurement matrix A1, A2,... An The labels are given by y; = (A,, X*), for every = 1,22... 2..... n. The empirical loss is
L(W
";)², 2n (AA,, WWT - where W e dddd
(16)
When the loss reaches near zero (which implies the gradient also reaches near zero), it is known that multiple local minimum solutions exist (Li al., 2018), and the Hessian becomes
JAWB d/il d d
By prior results (Recht et al., 2010), among all X = WWT such that i(W)= = 0, X* has the lowest nuclear norm. Thus, the regularization placed on L(W) is similar to muclear norm regularization after interpolating the training datasct. We formalize this discussion and state the result below.
Proposition 5.1. In the setting above, for any W that satisfies i(W) = 0, the following must hold with high probabilily:
Tr V( < Tr [V[L((W)]] + +0 On
(17)
14

---

Published in Transactions on Machine Learning Research (09/2024)
A similar statement holds if the trace operator is replaced by A in equation (17). To this, we look the quadratic form of the Hessian to find the maximum eigenvalue. Let be a de dimension vector with length equal to one, ||u]| = One can derive that:
A L[W)=(
max TV2L(W)==
Σ (AW,
dn Σ LAWI
max
WERd ||
The last step is by setting d 11 whose length is equal to one. The proof of Proposition 5.1 can be found in Appendix A.2.
Simulation: We conduct a numerical simulation to compare algorithmic behaviors. We generate a low- rank matrix U E Rdxr from the isotropic Gaussian. We set = 100 and r = 5. Then, we test three algorithms: gradient descent (GD), weight-perturbed gradient descent ( (WP-GD), and Algorithm 1 (NSO). In particular, we will implement the full gradient update rather than using the stochastic updates. We use an initialization Ug E Rddd where each matrix entry is sampled independently from standard Gaussian (V
Recall that WP-GD and NSO require setting We choose between 0,001, 0.002, 0.004, 0.008, 0.0016. NSO additionally requires sctting the number of sampled perturbations k. We set A = 1 for faster computing. As for the learning rate, we choose fixed for each run and vary its value between 0.001, 0.0025, 0.005, and 0.01. We find that setting y as either 0.005 or 0.01 would be too large, leading the loss values to explode. Hence, we report the results for setting 7j as 0.0025 or 0.001.
Our findings are illustrated Figure 6.
We see that all three algorithms can reduce the training MSE to near zero, shown in Figure 6a. From the trends, we can see that the training loss has fully converged for all cases.
GD suffers from overfitting to training data, while both WP-GD and NSO can generalize to the validation samples. Moreover, NSO reduces this validation loss further in Figure 6b.
Finally, we can see in Figure 6c that our algorithm can indeed produce a more accurate estimate of the ground truth matrix X*, as measured by the Frobenius norm distance between W.W. and X*.
The results for varying learning rates can be found in the bottom panel from Figure 6d to 6f. The comparative results remain consistent.
6 Discussions and Related Work
Using noise injection during neural network training has appeared in very early studies of machine learning roscarch (Hinton & Van Camp, 1993; An, 1996). Graves (2011) test variety of variational inference ap- proaches with different prior and posterior distributions with recurrent neural networks. Cohen et al. (2019) examine the use of randomized smoothing (with different smoothing distributions) against different fp ad- versaries for certified robustness. Camuto et al. (2020) propose a layer-wise regularization scheme motivated by adaptation patterns of weights through deeper layers. Yang et al. (2020) show how to turn any classifier that classifies well under Gaussian noise into a new classifier robust to adversarial perturbation under the l norm. One of the implications of their work is that smoothing with Gaussian noise naturally confers adversarial robustness in the &a norm. Bisla et al. (2022) conduct an extensive empirical study to explore the connection between sharpness and generalization for training neural networks. Orvieto et al. (2023) analyze Taylor's expansion of the stochastic objective after noise injection, examining the induced regularization in various neural network training settings, and find that layer-wise perturbation can improve generalization.
There is also a line of work on Hessian and sharpness in the edge of stability regime during gradient descent dynamics (Cohen et al., 2021). In particular, the edge of stability refers scenarios where the lcarning rate goes out of bounds beyond the Lipschitz continuity of a function, which inversely proportional the largest eigenvalue of the Hessian matrix. Long & Bartlett (2024) identify the edge of stability regime
15

---

Published in Transactions on Machine Learning Research (09/2024)
0.05
GD
10
10
Normalizzd Distance
Training MSE
NSO (k 1)
Validation MSB
0.04
WP-GD
GD
1011
GD
NSO (A 1]
NSO (A = 1)
0.03
10-2
WP-GD
10-2
WP GD
0.02
10-3
10-1
0.00
5000
1000
Number of epochs
Number of epochs
Number of epochs
(a) Training loss
(b) Validation loss
(c) Normalized distance
10
NSO (A 1)
10
NSO (k = 13
NSO (A: = 1) WP-GD
Normalized Distance
Thining MSE
WP-GD
Validation MSE
10"
WP-GD
102
30-1
10-1
10-1
10
10:2
10-2
10-3
10-
10-1
0
5006
10000
15000 20000
0
3000
10000
15000
20000
5000
10000
15000
20000
Number epochs
Number epochs
Number of epochs
(d) Training loss
(e) Validation loss
(I) Normalized distance
Figure 6: Comparing the training loss, validation loss, and the normalized Frobenius norm distance, i.e.
[.W,T-X 1I between GD, our approach (NSO), and weightt-prturrbed (WP) GD (which computes the full
|
gradient as opposed to the stochastic gradient). For the top panel, the learning rate is fixed at 0.0025 for all the runs. For the bottom panel, the lcarning rate set at 0.0001. set as 0.008 for WP-GD and NSO. Also, we trained sufficiently long until the loss curves fully converged.
for the SAM algorithm, highlighting the differences of these regimes between SAM and gradient descent. Agarwala & Dauphin (2023) present a detailed study of the gradient dynamics of SAM, documenting various respects of this algorithm. They first analyze the full-batch gradient descent with unnormalized SAM in a quadratic regression model. This analysis suggests that at initialization, full-batch SAM presents limited suppression of the largest eigenvalue of the Hessian matrix. They also show that as the batch size decreases, the regularization of SAM becomes stronger. This work underscores the intricate dynamics of SAM due its connection the min-max problem, which computationally intractable (Daskalakis al., 2021). Dauphin al. (2024) provide in-depth comparison between SAM and weight noise by examining the structure of the Hessian during training. Our results in Section 2.1, which show that weight noise remains ineffective (for fine-tuning), are consistent with the findings of this work. Wu et al. (2020) study the structure of the Hessian and conduct experiments on how the Hessian structure changes based on architecture and the training method.
Randomized smoothing has been studied stochastic optimization under various contexts, for instance, es- timating gradients in zeroth-order optimization (Duchi al., 2015), and for nonsmooth convex optimization problems (Duchi et al., 2012). In particular, Duchi et al. (2012) analyze the convergence rates of stochastic optimization algorithms and examine convolution-base based smoothing technique for nonsmooth stochastic op- timization problems by drawing stochastic gradient samples from the smoothed problem with an appropriate choice of smoothing density. They show that with the ability to issue several queries to the stochastic oracle, the original problem can be solved with faster convergence rates than a simple stochastic oracle. Besides, recent rescarch has investigatod the query complexity of finding stationary points of nonconvex functions (Carmon et al., 2020; Arjevani al., 2023). These results provide fine-grained characterization of the complexity of iterative methods under different orders of gradient oracles.
16

---

Published in Transactions on Machine Learning Research (09/2024)
The findings from work suggest several avenues that seem ripe for future work. Can recent advancements in optimization be used to design better noise injection algorithms with faster convergence? Can we better understand the effect of noise injection on the Hessian during training (e.g,, in tensor regression where saddle points are known to exist (Li et al., 2020))? In particular, our work highlights the need for more accurate measurements to undersland the learning mechanisms of complex models.
7 Conclusion and Limitations
This paper examines the regularization and generalization effects of noise-injection methods for training neural networks. The study begins by noting that a straightforward implementation of injecting noise into weight matrices (of a neural network) before computing the gradient SGD does not. perform well practice. Thus, an alternative, two-point noise injection scheme proposed and shown be effective through extensive experiments. In particular, this new algorithm can be used to regularize the Hessian and improve generalization. The results are tested on fine-tuning, pretraining, and instruction tuning. As complement, a PAC-Bayes generalization bound is provided to support the rationale of this approach. Finally, this paper presents a detailed convergence analysis of the proposed algorithm.
Limitations: In Theorem 2.1, we have shown that the generalization error of a training algorithm can be bounded by the trace the Hcssian of the loss matrix, scaled by the distance of the hypothesis space. Notice that this result applies to both Algorithm 1 (NSO) and the naive noise injection algorithm (WP-SGD). As shown in Figure 3, this result can provide descriptive measure to explain different algorithms. Since the Hessian measurements can be used on both algorithms, they can only distinguish one from another after taking the measurements from the data. Thus, generalization theory should be interpreted with this data-dependent lens in mind. We hope future work could work addressing such limitations, along with designing more principled optimization algorithms for training neural networks.
Acknowledgements
H. Z. would like to thank Huy Nguyen, Zhiyuan Li, and Guanghui Lan for the discussions and for pointing out several references during various stages of this work. The authors would also like to thank the anonymous reviewers and the action editor for their constructive feodback. We acknowledge financial support from NSF award IIS-2412008.
References
Atish Agarwala and Yann Dauphin. Sam operates far from home: eigenvalue regularization dynamical phenomenon. In International Conference on Machine Learning, pp. 152-168. PMLR, 2023. 8, 16
Pierre Alquier. User-friendly introduction to pac-bayes bounds. arXiv preprint arXiv:2110.11210 2021. 3, 7
Guozhong An. The effects of adding noise during backpropagation training on generalization performance. Neural computation, 8(3):643-674, 1996. 2, 15
Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware minimization. In ICML, 2022. 1
Yossi Arjevani, Yair Carmon, John C Duchi, Dylan J Foster, Nathan Srebro, and Blake Woodworth. Lower bounds for non-convex stochastic optimization. Mathematical Programming, 199(1):165-25) -2446, 2023. 16
Sanjeev Arora, Rong Ge, Behnam Neyshabur, and Yi Zhang. Stronger generalization bounds for deep nets via a compression approach. In ICML, 2018. 7
Peter Bartlett, Philip M Long, and Olivier Bousquet. The dynamics of sharpness-aware minimization: Bouncing across ravines and drifting towards wide minima. Journal of Machine Learning Research, 24 (316):1-36, 2023. 1, 13
17

---

Published in Transactions on Machine Learning Research (09/2024)
Devansh Bisla, Jing Wang, and Anna Choromanska. Low-pass filtering sgd for recovering flat optima in the deep learning optimization landscape. In International Conference on Arlificial Intelligence and Stalistics. pp. 8299 8339. PMLR, 2022. 15
Alexander Camuto, Matthew Willetts, Umut Simsekli, Stephen Roberts, and Chris Holmes. Explicit regularisation in gaussian noise injections. Advances in Neural Information Processing Systems, 33:16603- 16614, 2020. 15
Yair Carmon, John C Duchi, Oliver Hinder, and Aaron Sidford. Lower bounds for finding stationary points i. Mathematical Programmming, 184(1-2):71-12) -120,, 2020. 3, 16
Olivier Catoni. Pac-bayesian supervised classification: the thermodynamics of statistical learning. arXiv preprint ar arXio:0712.022 0248, 2007. 3,
Jeremy Cohen, Elan Rosenfeld, and Zico Kolter. Certified adversarial robustness via randomized smoothing. In international conference machine learning, pp.
PMLR, 2019. 15
Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. ICLR, 2021. 15
Constantinos Daskalakis, Stratis Skoulakis, and Manolis Zampetakis. The complexity of constrained min-
max optimization. In Symposium on Theory of Compuling, 2021. 16
Yann N Dauphin, Atish Agarwala, and Hossein Mobahi. Neglected hessian component explains mysteries in sharpness regularization. arXiv preprint arXiu:241111000, 2024. 2, 16
Yoel Drori and Ohad Shamir. The complexity of finding stationary points with stochastic gradient descent. In ICML, 2020. 3,
John C Duchi, Peter L. Bartlett, and Martin J Wainwright. Randomized smoothing for stochastic optimiza- tion. SIAM Journal on Optimization, 22(2):674 701, 2012. 2, 4, 16
John C Duchi, Michael I Jordan, Martin J Wainwright, and Andre Wibisono. Optimal rates for zero-order convex optimization: The power of two function evaluations. IEEE Transactions on Information Theory, 2015. 2, 5, 12, 16
Gintare Karolina Dziugaite and Daniel Roy. Computing nonvacuous generalization bounds for deep (stochas- tic) neural networks with many more parameters than training data. UAI, 2017. 8
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. ICLR, 2021. 1, 2, 3,
Saeed Ghadimi and Guanghui Lan. Stochastic first-and zeroth-order methods for nonconvex stochastic programming. STAM Journal Optimization, 23(4) -2668, 2013. 3, 12, 13, 27
Henry Gouk, Timothy M Hospedales, and Massimiliano Pontil. Distance-based regularisation of deep net- works for fine-tuning. In Ninth International Conference on Learning Representations 2021, 2022. 3. 11
Alex Graves. Practical variational inference for neural networks. Advances in neural information processing systems, 24, 2011. 2, 15
Gregory Griffin, Alex Holub, and Pietro Perona. Caltech-256 object category dataset. 2007. 4
Geoffrey E Hinton and Drew Van Camp. Keeping the neural networks simple by minimizing the description length of the weights. In Proceedings of the sixth annual conference on Computational learning theory, pp. 5-13, 1993. 2, 15
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural compulation, 9(1):1 1997.
18

---

Published in Transactions on Machine Learning Research (09/2024)
Pavel Izmailov, Dmitri Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. UAT, 2018.
Haotian Ju, Dongyue Li, and Hongyang R Zhang. Robust. fine-tuning of deep neural networks with hessian- based generalization guarantees. ICML, 2022. 2,
Haotian Ju, Dongyue Li, Aneesh Sharma, and Hongyang R Zhang. Generalization in graph neural networks: Improved pac-bayesian bounds on graph diffusion. AISTATS, 2023.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. ICLR, 2017. 1
Guanghui Lan. First-order and stochastic optimization methods for machine learning, volume 1. Springer, 2020.
Dongyue Li and Hongyang R Zhang. Improved regularization and robustness for fine-tuning in neural networks. Advances in Neural Information Processing Systems, 34:27249-2277 2021. 4, 7
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in over pre-paaraeterrized matrix sensing and neural networks with quadratic activations. In Conference On Learning Theory, pp. 2-47. PMLR, 2018. 14
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Learning wvrrparaaeerrieer two-layer neural networks beyond ntk. In Conference on learning theory, pp. 2613 2682. PMLR, 2020. 17
Yong Liu, Siqi Mai, Minhao Cheng, Xiangning Chen, Cho-Jui Hsieh, and Yang You. Random sharpness- aware minimization, Advances in Neural Information Processing Systems, 2022. 8
Philip M Long and Peter L Bartlett. Sharpness- aware minimization and the edge of stability. Journal of Machine Learning Research, 25 25(179):1-22 2024. 15
Subhransu Maji, Esa Rahtu, Juho Kannala, Matthew Blaschko, and Andrea Vedaldi. Fine-grained visual classification of aircraft. arxiv preprint orXiv:1306.515. 2013. 4
David McAllester. A pac-bayesian tutorial with dropout bound. arXiv preprint arXiv:1307.2118 2013. 3,
7, 21
Vaishnavh Nagarajan and Zico Kolter. Deterministic pac-bayesian generalization bounds for deep networks via generalizing noise-resilience. ICLR, 2020. 2,
Behnam Neyshabur, Srinadh Bhojanapalli, and Nathan Srebro. A pac-bayesian approach to spectrally- normalized margin bounds for neural networks. In International Conference on Learning Representations, 2018.
Antonio Orvieto, Anant Raj, Hans Kersting, and Francis Bach. Explicit regularization in overparametrized models via noise injection. AISTATS, 2023. 2, 15
Samiksha Pachade, Prasanna Porwal, Dhanshree Thulkar, Manesh Kokare, Girish Deshmukh, Vivek Sa- hasrabuddhe, Luca Giancardo, Gwenolé Quellec, and Fabrice Mériaudeau. Retinal fundus multi-discase image dataset (rfmid): A dataset for multi-disease detection research. Data, 6(2):14, 2021. 4
Alec Radford, Jong Wook Kim, Chris Hallacy Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and llya Sutskever. Learning transferable visual modcls from natural language supervision. In ICML, 2021. 12
Benjamin Recht, Maryam Fazel, and Pablo A Parrilo. Guaranteed minimum-rant solutions of lincar matrix equations via nuclear norm minimization. SIAM review, 52(3):471-501, 2010. 14, 25
Shai Shaley-Shwartz and Shai Ben-David. Understanding machine learning: From theory to algorithms. Cambridge university press, 2014. 6
19

---

Published in Transactions on Machine Learning Research (09/2024)
Nathan Srebro and Adi Shraibman. Rank, trace-norm and max-norm. In International conference
computational learning theory, pp. 345-560. Springer, 2005. 6
Martin J Wainwright. digh-dimensional statistics: non-asymptotic vieupoint, volume 48. Cambridge University Press, 2019. 24, 25
Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How does sharpness-aware minimization minimize sharpness? ICLR, 2023. 1, 3
Mitchell Wortsman, Gabriel IIharco, Samir Ya Gadre, Rebecca Roclofs, Raphael Gontijo-Topes, Ari S Mor- cos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, and Ludwig Schmidt. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In International conference on machine learning, pp. 23965-23998, PMLR, 2022. 1, 4
Yikai Wu, Xingyu Zhu, Chenwei Wu, Annie Wang, and Rong Ge. Dissecting hessian: Understanding common structure of hessian in neural networks. arXiv preprint arXiv:2010.04966 2020. 16
Greg Yang, Tony Duan, Edward Hu, Hadi Salman, llya Razenshteyn, and Jerry Li. Randomized smoothing of all shapes and sizes. In International Conference Machine Learning, pp. 15
PMLR, 2020.
Tong Zhang, Mathematical analysis of machine learning algorithms. Cambridge University Press, 2023. 3
20

---

Published in Transactions on Machine Learning Research (09/2024)
A Omitted Proofs from Section 2
We state a few standard notations. Given two matrices X, Y having the same dimension, let (X,Y) Y) Tr[XT y] denote the matrix inner product of X and Y. Let || denote the spectral norm (largest singular value) of X, and let | denote the Frobenius norm of X. We use the big-O notation f(z) (g((z) to indicate that there exists a fixed constant C' independent of æ such that f(x) S C- g(x) for large enough x.
A.1 Proof of the PAC-Bayes Bound
We will use the following PAC-Bayes bound. For reference, see, c.g., Thoorem 2, McAllester (2013). Theorem A.1. Suppose the loss function ((fw(x), y) lies in bounded range [0, CJ given any a E X with label y- For any 8 E 10 and 8 E (0,1), with probability al least 1 - 6, the following holds:
Le(W) 2 8 ( C(KL(2]P)) + log 2 23 -3)
(18)
This result provides flexibiiity in setting B. Our results will set 3 to balance the perturbation error of Q and the KL divergence between P and Q. We will need the KL divergence between the prior P and the posterior Q in the BaC-Bayesinn analysis. This is stated in the following result.
Proposition A.2. Suppose P = N) and Q = N(Y,)) are both Gaussion distributions with mean veclors given by X IR,YY Y E IR, and population covariance malrix E RRxp, The KL divergence between P and Q is equal
KL(2I))= = 1 -(x (x -x Y).
Specificully, 8 o2Idp. Idp, then the above simplifies to
KL(QIP)) | YI
20
We will use Taylor's expansion on the perturbed loss. This is stated precisely as follows.
Claim A.3. Let fw twice- difie-differentiable parameterized by weight vector W Rp. Let Rp be another vector with dimension p. For any W and U, the following identity holds
((fw4u(I),)) v) ((ff(x)) + U V((f(xx).)) + + (V( + (a(((fw(s).v))
where (e(E(fw(z). is a second-order error term in Taylor's 's expansion.
Proof. The proof follows by the fact that go fw is difeeeiiffeeetiiable. Let 7 E RP be a vector with the same dimension W and U from the mean value theorem. There must exist an n between W and U + W such that the following equality holds:
This completes the proof of the claim.
We provide Taylor's expansion of to e based on the above.
Lemma A.4. In the setting of Theorem 2.1, suppose each parameter is perturbed by independent noise drawn from o2). Let lo(f((x) y) be the perturbed loss with noise perturbation injection vector on W. There exist some faoed value C2 that do not grow with n and 1/5 such that
1 002 [V((((((())) SCO
(elsw(x),u) ((w()).)) 2
21

---

Published in Transactions on Machine Learning Research (09/2024)
Proof. We take the expectation over U for both sides of the equation in Claim A.3. The result becomes
E ( v)] B [((fw(z)) + U (e(ff())))) U. ((f u) + Re(((fw()))))) Then, we use the perturbation distribution Q on Ey (Ww+tr/wwrrrre (x), y)l, and gel
Co(fw(z),)) = El((((()))) + U ve((((()) v) + + UU
+ (x).
Since E = 0, the first-order term will be zero expectation. The second-order term becomes equal
B L [[v - o² (+)))
(19)
The expectation of the error term (2(e(fw v)) be -1( [0| (V((((((0)))))
B B[RA(((((()):)))
1 | || - (1l()
E |uil C||UUl|| SCO0..
Thus, the proof is complete.
The last piece we will need is the uniform convergence of the Hessian operator. The result uses the fact that the Hessian matrix is Lipschitz continuous.
Lemma A.5. the setting of Theorem 2.1, there exist some fixed values C2, C3 that do not grow with n and 1/6, such that with probobility least 8 for any 8 > O, over the rondomness of the n training ezamples, have.
C2 CoVlog(Csn/0
E (a.N)~
[V(((((())
(20)
The proof will be deferred to Section A.1.2. With these results ready, we will provide proof of the Hessian- based generalization bound.
A.1.1 Proof of Theorem 2.1
Proof of Theorem 2.1. First, we separate the gap of L(W) and L(W) into three parts:
L(W)- - (ww -(W) Lo(W) Lo(W)
1 Le()) 1 -
3
3
By Lemma A.4, we can bound the difference between L(W) and Lo(W) by the Hessian trace plus error:
L(W)- 3 fi(w)= E
o²
Tr [V(((([((
+ca0 (Io(W) - ( )
(x.))>
2
Tr (1. + Cios).
2
After re-arranging the terms, can get the following:
1 (imms
0²
m
o²
L(W) 1003-11
E (D
Tr
[()((
2
n
E1
1 1 8
1 C + Le(W)
-
(WW)
(21)
3
3
E,
22

---

Published in Transactions on Machine Learning Research (09/2024)
We will examine Ej by separating it into two parts:
o²
o²
E1
Tr (1)
- IF
Tr [V99)(((())
(22)
3
n1 i=1
2
(x.0)r
2
1 -
8
o²
Tr [V(((((())
(23)
B
2
(DD
We can the uniform convergence result of Lemma A.5 to bound equation (22), leading to:
0²
Tr [+)
E
TTr (V
(x).)))
28
(a.N)~
t=1
0²
Vp-
E [=((([(<)))))))
(by Cauchy-Schward
2/
(2,9)~
0² Vp- C2v C2V/og((Can/6)
(24)
2m
In particular, the second slep also uses the fact that the Hessian matrix symmetric p by p matrix. As for equation (23), we recall that
a: max Tr [V(((w(().
(a.n)~D
Combined with equation (24), we have shown that
E x o² Vp- C2//og(C33/0)
- 3-
+
a
(25)
23/n
B
As for Eg, we will use the PAC-Bayes bound of Theorem A.1. In particular, we the prior distribution p the distribution of U and the posterior distribution as the distribution of W + U. Thus,
C( [[I] 2 + log
ES S C(KL(Q)P) + log 1)
2 C(
2 1og
(26)
28(1-3)m
2 3)m
23 8)n
The last step is because [[ S by the assumption of the hypothesis space. Combining equations (21), (25), (26), we claim that with probability at least - 25, the following must be true:
L(W)- 1 L(Wss 0² VD.V22220 C2n/og(C3n/0) + B 2 a
1-30² 1+
C
2 log 4
Ci03 3
(27)
23
B
2 - B)n
Thus, we will now choose and 8 E (0, 1) minimize the term above. In particular, we will set a as
r
C
(28)
By plugging in to equation (27) and re-arranging terms, the gap between L(W) and L(W) becomes:
1 Car2 (i(m) 3V n
C2 (og(Co/)))) o² 1 C,03+ 2m
C
L(W) /
log
2
3)m
8
Let be fixed value close to and independent of N and 8-1, and let (1- (1-8)/3. We get
Car + where
L(W) < (1 (1 ( + (
&= Cz C2/2plog (0ap/
0²+ (1+ 8
C log 1 1) C100 + 2 3)n
2
8
Notice that & is of order O(n $ + n : log(8 11)) 1m s Olog(6 ')n 2). Therefore, we have finished the proof of equation (6).
23

---

Published in Transactions on Machine Learning Research (09/2024)
Remark A.6. When f is strongly convez, the lowest eigenvalue of the Hession is bounded from below. Once the algorithm reaches the global minimizer, our result from Theorem 6 can be used to provide a generalization bound based on the trace of the Hessian. Notice that the noise injection will add some bias to this minimizer, leading to sub-opii oub-optimal empirical loss. To remedy this issue, one can place the regularization of Hessian as a construinl, simalar to how l2 le-regularizatinn can implemented as constraind.
A.1.2 Proof of Lemma A.5
In this section, we provide the proof of Lemma A.5, which shows the uniform convergence of the loss Hessian.
Proof of Lemma A.5. Let C, c > 0, and let S = ( R || S C). There cxists an c-cover of $ with respect to the f--norm at most max ((=0)).1 clements; see, e.g., Example 5.8 (Wainwright, 2019). Let T S S denote the set of this cover. Recall that the Hessian ())([[()))))) is Ci-Lipschitz for all (W + U) E Then we have
(VWWW VW (z), willl S Ci 2|U||2
For parameters d,e > O, let N be the e-ooeer of S with respect to the fz-norm. Define the event
- "
E VW ET,
(2.0)~D
i=1
F
By the matrix Bernstein inequality, we have
no²
Pr[E] 2 1- 11-4-WW] exe
2a
Next, for any W E S, pick some W +U T such that ||0|| S. We have
[V(([w+
(V((([
<G C1e
(2,0)/
(((((wo(()))
Σ (l(((s((3)))) wi)1
ScIUU S Cie.
Therefore, for any W E S, we obtain:
(((()(())))))))11
[V((([w(<).)))]]
2C 6.
n
(2.9)/D
We will also set the value of 8 and First, set 5/(2Ci) that conditional E,
((((w((s)))))111
E
(le((f(
< 26.
(2.0)D
The event E happens with probability of at least:
no²
no²
1 - 4Tpp exp
- 4p ex-
log [T'l
2a
2a
We have log [T] S plog(3//e) = p
we set
4pa log(3/C
n
24

---

Published in Transactions on Machine Learning Research (09/2024)
that log(37CCan/a)) 2 (because 2 3C1
and 1),
then get
60CIVīi
plog(6CC(/0) n2²2 /202/ =p =oo
2p log (3TCCin/a)
3CCIV
=p log
2plog (3TCC1m/a)
aV/plo((+CC/m//a)
Sp log (3r (3rCCIn/a) 2p log (3TCCin/a)
(r 2 1ogg((333C)))/ 2 (3CCnn/a e)
log (3rCCIn/a) S -plog(et).
Therefore, with probability greater than
4WW exp(-nd / 2 1 4p(et)
the following estimate holds:
16pa log(3rCC(n/aa
- E
n 1=1
(0.y)D
Denote 8' = 4p(eT) C2 = 4avD. and C3 = 12pCCt/(ea). With probability greater than 1 - 8', the final result is:
log log(Can/88)
(0))
E (m.y)~
S SC21 F
This completes the proof of Lemma A.5.
A.2 Proof of Proposition 5.1
Proof of Proposition 5.1. We can calculate the gradient as
vi(W) _ Σ ((A, wW W.W.W.
(29)
i=1
For a particular entry Wi of W, for any 1 1 j. S d, the derivative of the gradient with respect to W is
(A,WI,.A,WW + (AA, v.) aw,* a
((A,W)
(30)
n
When Ż(W) is zero, the second term of equation (30) above must be zero, because (A,, WWT) is equal to for any i = I,...,
We use the assumption that A, is a random Gaussian matrix, in which every entry is drawn from a normal distribution with mean zero and variance one. Notice that the expectation ||A,W]] satisfies:
E [ = (r [WT AT A.W]] = Tr [W (d (d Idaxa) = d. Tr W T W| = =d d
Thus, by concentration inequality for x2 random variables (e.g., Wainwright (2019, equation (2.19))), the following holds for any < 1,
JAWI [II11 2 ed [W]]?] 2eep (- nc²
Pr
(31)
n
This implics that c must be smaller than ((n-1/2) with high probability. As a result, the average of ||4,W]| must be d ||W]| plus some deviation error that scales with n 11/2 times the expectation.
By Theorem 3.2, Recht al. (2010), the minimum Frobenius norm ([[W)) solution that satisfies L(W) (for Gaussian random matrices) is precisely U*. Thus, conclude that equation (17) holds.
25

---

Published in Transactions on Machine Learning Research (09/2024)
B Omitted Proofs from Section 4
B.1 Proof of Proposition 4.2
Recall that each iteration involves two sources of randomness stemming from Os and (0, Let us define
Jz=1, respectively,
& 6,
V VF(W, + V.(W. - -VF()))
2k
j=1
(G - TF( -
&,
2k
j=1
for i= "Q...,TT T 1. One can see that both 5, and & have mean zero The former is by the symmetry of P. The latter is because Is unbiased under Assumption 4.1. The following result gives their variance.
Lemma B.1. In the setting of Proposition 4.2, for any i= = ,,...,T, we have
[KS]] k and E [10.11]] < (PH
E
(32)
k
The last step uses smoothness to show that |/VF(W)) keeps reducing.
Proof. Let us bound the variance of de and & for i = 0, 0......,7 - 1. First, we see that
E
|0:1
E
VS(W (/W -
01 U5
Up
U.
2h:
||
E
VF(W,+ + vf(W. -2)---()--- - -VF(2))(
(33)
2
2
2 V +,) Vf(: - -U)) VE(WES [7]
(34)
of
where in the second line use that U. and U are independent when ji * ja, in the last line use fact that U?...., U. are identically distributed. In the second step, we use the fact that for two independent random variables U, and any continuous functions h(U), 9VV), h(U) and g(V) are still independent (recall that f is continuous since If U and V are independent, we have Pr(U A, E B) Pr(U E A] Pr(V E B], for any A, E Borel( Thus, if h and g are continuous functions, we obtain
)ieeeeiifeeenti...l.... We include a short proof of this fact for completeness.
Pr(h(U) E gV) =r (A), E (B)]
= e =- Pr(V E (B)] = Pr(M(U) A] Pr(g(V) e B).
Thus, we have shown that
1 E Ump
VF(W.) +)) -V)(W((.))) |
E
[16.1]]
(35)
lo
Next, we deal with the variance of the two-point stochastic gradient. We will show that
1 (V (V/(W+D) +U) (oWW-u)) -V) - - CHH(P).
(36)
26

---

Published in Transactions on Machine Learning Research (09/2024)
We mainly use the Lipschitz continuity of the gradient of F. The left-hand side of equation (36) is equal to
Vs(W +0) VF(W)) + ( V -U) - VF(W)) |
E
0
1
E
VF(W +U) - VF(W)]
VV((W - U) VF(W)]?
(by Cauchy-Schwarde
0 2
F
V/(W + +0- VF(W)1
(by symmetry of P since has mean zero)
2
+ +
2 1
IV +U) vs( + +
VV(W +0) VS(W
20
1
E
ca |U U'|
1 C²
E ,
|17
+ +140! = HP)
(by equation (38))
2 U.
As for the variance of E., we note that U ......ub! are all independent from each other. Therefore,
+ 19 (W +U) vs(W + +0) + W - U) - -(W- u3)
E
(0t0,400 031*
4k
U.
a² k (W vS( + +]] + 1l - U) - VSW -U)] S
U>|
E U.z
2k
The first step uses the fact that both 9-(-) and f()) are continuous functions The second step above uses Cauchy-Schward inequality. The last step uses the variance bound of 9z(,), Thus, the proof finished.
In the next step, use result from Theorem 2.1, Ghadimi & Lan (2013). Our proof follows from their work, but we deal with some extra technical details related to the noise injection.
Lemma B.2 (Slightly adapted from Theorem 2.1, Ghadimi & Lan (2013)). In the setting of Proposition 4.2, for any yo -1 less than C-1 and a vandom variable according distribution Pr( j] Σ 7 for any j 0, ... T 1, the following holds:
7/
TETO
|18,11? + +E [I&,1]]
2C
[VF(W))))
D²
(37)
T-1
ET.1 778
2i=0 97:
Proof. First, let us show that VF is C-Tipschitz continuous. To see this, we apply the Lipschitz condition of the gradient inside the expectation of F(W).. For any W1, W2 Rd, by definition,
[f[ + U)] V E [s( + Ui] unp
IVF(W))-VF(W))) VF(22)
V E Unp
[V +U)
VV( -U]]
UnP
< E U<p
[I] + +U V/(W, U\Il S C [ Wall
Since VFW is C-Lipschitz continuous, we have the following domination inequality:
IF(W)) F(W)) (VF(W)), Wo - WiY < C I2 - -1,l 2
(38)
27

---

Published in Transactions on Machine Learning Research (09/2024)
Based on the above incquality, we have
F(W+11) SF(W.)+ + (VF(W)), - W.) W.) + C .z
1 +.) +
C's
=F(W.) 7/4 (VF 8, +
VF(W:)) +
||o; + VF(W))]
2
=F(W.)- ( -w Ct²
/F(W3))| - (n: - Cui (VF(W.), G) + Cnh | + =11
Summing up the above inequalities for i = 0,1,..., T 1, we obtain T-1
T-1
Cnh
Σ F(W++1) S
Σ F(W.)
|VF(W))]
T
Σ
(n. Ort)
(++)
+ Σ
T-1
C
16; + E
2
1=0
1=0
which implies that
T-1
(n
Cuf
IIVE(W)) SF(W)) - F(W))-
- Σ
T-1
n1 ( (VF(.))))) +5))
C T-1
Σ
m |1o + Evll
i=0 2
2 1=
1=0
T-1
C T-1
<D²
Σ 1=0
nh (n- 6
+
Σ
n 16, + Ell?
(39)
2 1=0
where in the last step, we use the fact that
F(W)) F(WT) < F(W)) - min F(W)= - D WERd
For any 0, 1,.....777 T 1, notice that as long as 0 < m S t, then 1/1 S 2n: Onl. Hence, we have
T-1
Chii
IVF(W))]
s Σ 1/t
|VF(W)))
2 1=
Th
which implics that
1 T-1
T-1
C T-1 ni 16, + Eall
Σ n IVF(W,)1 < D
- (n- ((n= (mi
+ E.) + (VF(W)) (n))
Σ
(40)
2
Additionally, since U, is drawn from a distribution with mean zero. Hence, by symmetry, we get that
1 E [V U.) Vf(W, U.) 0.
[. E T 20,
(41)
Thus, if we take the expectation over Up. UI.... UT Eo. . -1, then E (VVE(W)), 6; + =+)==00 0. Recall
that is a random variable whose probability mass is specified in Lemma B.2. We can write equation (40) equivalently as (below, we take expectation over all the random variables along the update since W. is a
function of the previous gradient updates, for each 0......,2 T 1, recalling that Pr|t = i] Σ |= 75
ET-1 7 E |VF(W)))
[IVF(W))]
E 1s UU....UU Er.
Σ1- 7
2D2 + C CΣTT
1 "h E [16: + E/ll
Σ1-1 176
2D + CEP'1 1/ E ||
+ E
| [|
2
ΣT- 176
where we the fact that 6; and & are independent for any i. Hence, equation (37) is proved.
7
28

---

Published in Transactions on Machine Learning Research (09/2024)
Based on the above result, we now finish the proof of Proposition 4.2.
Proof of Proposition 4.2. Let the step sizes be fixed for all epochs. Thus, equation (37) becomes
D2 + Cn T-1
2 [F(W))) S Tn
[18:111 + E [Isa1*]):
(42)
By Lemma B.1,
ST. + CH)(P)
E [1o:ll + [IE]IP] (E ==
(43)
For simplicity, let us denote A H(C
The proof is divided into two cases.
Case 1: A is large. More precisely; suppose that A 2 2CD2/T. Then, minimixing over 1) above leads to the following upper bound on the right-hand side of equation (42):
2CD2
(44)
T
which is obtained by setting 17 2CD?, Thus, we conclude that equation (42) must be less than
2D2
One can verify that this step size is less than since A is at least
CAT
A
2CD +CH
(45)
T
KT
Case 2: A is small. In this case, suppose A <2CD2/ /T. Then, the right-hand side of equation (42) must less than
2D
20 D2n
2CD2
(46)
Tn
T
T
Thus, combining equations (45) and (46), have completed the proof of equation (12).
B.2 Proof of Theorem 4.3
Recall our construction from Section 4 as follows. Let be the basis vector for the t-th dimension, for 1=0,1 0, 1, 0,1...., 1. Define f(W) as
S(W) 2G (W, e0)2
Σ he
where a quadratic function parameterized by a:, defined as follow:
Co C(|a--a;)
x S a;
Cof
a < | s g
h₂(x)
C(la 20
8 a < | < 20;
20; < |z]
For technical reasons, we define a truncated perturbation distribution P. Given a sample U from a d- dimensional isotropic Gaussian VO,,,))) we truncate the i-th coordinate U so that U, = min(U,, a.), for some fixed a, > that we will specify below, for all = 0, I,...,d 1. Let P denote the distribution of Ú. The proof of Theorem 4.3 divided into two cases. First, we examine the case when the averaged learning rate is OTT2//2)
29

---

Published in Transactions on Machine Learning Research (09/2024)
Lemma B.3. In the setting of Theorem 4.3, suppose the learning rates satisfy that Σi= S /
DAKT 2
consider the function J(W) constructed in equalion (13), we have
Co²
min E [VF(WW)) 2 DV 32KT
Proof. We start by defining a gradicnt oracle by choosing the noise vectors {Et {Ge) to be independent random variables such that
& = (5x, ec+ ec+ 1)e1+1 and (&,, ex+1) S
(47)
Vk'
where e++1 is basis vector whose (t + 1) entry is one and otherwise is zero. In other words, only the (t + 1)-th coordinate of & is nonzero. Otherwise, the rest of the vector remains zero. We use & to denote the averaged noise variable as
1 k
&
where <( is defined following the condition specified in equation (47). Thus, also conclude that
K&,, ((t)l
Vk
We consider the objective function S(W): Idd + defined above (see also equation (13), Section 4), with
2n for i 0.1....,77
a, Vk"
(48)
We will analyze the dynamics of Algorithm 1 with the objective function f(W) and the starting point Wo DVG eo, where = max {C-1 2 m. For the first iteration, we have
W1 Wo 110 (1 (VS(W) 08 VSWW )) Ev)
= - -1- W noEo-
where is random draw from the truncated distribution P with (U, = min(Pi, ai} for a; 2e-1d Next, from the construction of h1, we get
V f(W, + U) + VS(W- -
= GG -W
(N6(To ei) (U,e1))er M.( e1) (U)e)))))).
Here, using the fact that ag (W)e s Vk and (Eo,e1) s VE obtain
2 from equation (48) above, and the truncation of U, which implies
vk
2nd
21p0 = 0g Vk
In055, e1) + (U,e)) <
os, and similarly (no e1) - (U, (i,ei)l Vk
which implics that
Mo (Ev. e1) + (U, e1)) = (o - -U) 0.
This is the first update. Then, in the next iteration,
(Wi,e0) W2 = Wi 1 (G (W-((W,, + +E1)
-(1 m G G-G---WWW noo5 E
30

---

Published in Transactions on Machine Learning Research (09/2024)
Similarly, we the fact that 0; 2n.e and the fact that |U, c++1) I I De
which renders the gradient At the t-th iteration, suppose have that
vk zero similar the above reasoning. This holds for any i ........ ..,, T VE
1-1
W, = Wo II (1 1 1l-nnn-1)
Σ
nki.
1=0
Then by induction, at the (t 1)-th iteration, we must have
W4+1 We - ne (Wer eo) +
1
1. Σ
Wo II [ n/ G/
7cE..
(49)
Next, from the definition of he above, we have that
F(Wo)=
min ([W)=((M) WeR*
(the minimum be attained at zero)
T-1 C 2m:00 Σ
1 (DVG)++ +
(since (Wo + e(+1) aa))
2G
Vk
1=0
The above must be at most D?; which implies that we should set the learning rates to satisfy (after some calculation)
kD² 2C
T 7:
Σ
(50)
8=0
We note that for [0, 1], 1 2/2 2 exp(log $). Thus, applying this to the right-hand side of equation (49), we obtain that for any t.
II 1=
(1 (1-1-GG [1- G. (1-
(51)
2'
where we recall that G = Max((C_,222' mn. Our calculation so far shows that for all the h; except ho, the algorithm has not moved at all from its initialization at Wo under the above gradient noise. We thus conclude that
min 1SST
/V
= min ISiTT
G- (Wo, .o))
(by the construction of F(-))
G (DVG)²
(by equations (49) and (51))
D
min {c,
(recall the definition of above)
2 -TT-1 7i
22
{c, 2D DDV
v2C²²
Co²
min
DV
(by equation (50))
32kT
In the first step, we use the fact that (Es, ea) = 0, for all 0 = 1....... - 1. Thus, we have proved that equation (14) holds for W. for any i = 1,2...., T. The proof of Lemma B.3 is finished.
Next, let us consider another case of the lower bound.
Lemma B.4. In the selling of Theorem 4.3, suppose the learning rales salisfy that Σr- 2
DATT C
and T/: = for some fixed v S C-11 Then, consider the function from equation (13), we have that
Coa min E [P(W))) 2 DV 3247
31

---

Published in Transactions on Machine Learning Research (09/2024)
Proof. We define the functions 9, parametrized by fixed, positive constants a 1-p"
22no, as follows:
I-p
C Cx²
a
R awb b
g(z)
c (1] x)2
I IA
Sa,
One verify that Vg is C-Lipschitz continuous, but g is not twice-differentiabl like function:
We also consider
f(W) s((W, (W, et+1)²
(52)
From the definition of f, its gradient C-Lipschitz continuous. Similar to equation (47), we define an adversarial gradient oracle by choosing the noise vectors {6))-0) to be independent random variables such that
& = Bl = and e++)) ca,
where e is a fixed constant. We use & to denote the averaged noise variable as
&
Suppose (5()) Jez are iiddd random variables for any t, we have
o²
(&,, e++1 S ca and E [|5.1]
(53)
Next, analyze the dynamics of Algorithm 1 with the objective function f(W) and the starting point
es In this case, by setting 17s = 17 for all 0,1, ,, Recall that 7 C-1, Denote by = Cn, which is strictly less than one.
Wo Ei= Ea ee.
Since hy an even function, derivative M odd. For the first iteration, we have
W1 Wo 7 5 Vf(WW U) Vf(Wo U)) + +o) (1- Cn)Wo neO-
where U is truncate distribution of p NOO,, Ida) with (U, min ao} and an cnj.
Using the fact that a =
1-p" I-p 201o, (U, oo)l < ena and (En, <0) S co, we have
g'(nio, en) + (U, en)) + (1 ea) - (U,)o)) CCn(Co,ep)
Then, in the next iteration,
W₂ = W1 - -n(c) ( (WW--
- CnEo + + -Cno = =11 - (1- -n))nooo - n
Similarly, we use the fact that L-p" as gx(x) -Cr, for any = 1,2, 1.2....17 1. At the t-th iteration, suppose that
2enjor and the fact that (U, enjr, which renders the gradient T-p
W, W (1- Cn) (1-)
E -nn-nnn---
1=0
32

---

Published in Transactions on Machine Learning Research (09/2024)
Then by induction, at the (t 1)-th iteration, have
WW, es) - - 1-1
W41 = W. n((c) Σ
cn) + E1)
1=1
1=0
= -1--Cn) Wo (1- ((1-Cn)(---nfe
(54)
1=0
Next, from the definition of F above, we have that
dC (V
D2)2
C C(2(1- pcno=
F(W)) min F(W) = F(W)) 2 Cd. WER"
+
(since (Wo + U,eo) sa)
4
p)
which must be at most D Thus, we must have (after some calculation)
D² - p)2
20 p²
We conclude that
min E 135S
[F(W)))
= min E
C (W 60)
K3TT
1=1
D2
0²
min K
(d - -)
(V Cd
+ k
p² Σ
p) ()2((-1-))
1=0
CCD pp) + d² P (1 - p)
2 min CDD a²
min K
k 2 p
k p-pl
kD2
Co²
2 DV 16k
(after some calculation)
2T 2 CV
kD2 TTTCC
Thus, we have proved this lemma.
Taking both Lemma B.3 and B.4 together, we thus conclude the proof of Theorem 4.3.
B.3 Proof of momentum lower bound
In this section, prove the following result.
Theorem B.5. There exists a quadratic function I such that for the iterates WI,..., Wr generated by equation (15), we must have: minigest IE [IV 2 O Co2
We will focus on a perturbation distribution P equal to the isotropic Gaussian distribution for this result. In this case, we know that F(W)= = f(W) + d. For the quadratic function f(W) = [[W]| its gradient is C-Lipschitz continuous. We set the initialization Wo E Idd such that
F(Wo)
min F(W)= = D WeRd
This condition can be met when we set Wo as a vector whose Euclidean norm is equal to
T-1 C-
D. 2 maax
Σ
ni
1=0
33

---

Published in Transactions on Machine Learning Research (09/2024)
The case when μ 0. We begin by considering the case when μ 0. In this case, the update reduces to SGD, and the iterate W441 evolves as follows:
We+1 -11- We 1e5r,
(55)
プラ=ı cli) and the noise perturbation U. cancelled out between the plus and minus perturbations. The case when f 0 builds on this simpler case, as we will describe below.
where we denote & as the averaged noise
The key observation is that the gradient noise sequence f1, Ea...., forms martingale sequence:
For any = 1,2..... ,, conditioned on the previous random variables c/ for any if < i and any j = 1,2....,,k the expectation of & is equal to zero.
In addition, the variance of &, is equal to k= since conditional on the previous random variables, the < are all independent from each other.
The martingale property allows us to characterize the SGD path of [[&l as shown in the following result. Lemma B.6. In the setting of Theorem B.5, for any step sizes no...... -/r-1 less than C-1, and any t = 1, 1,...,T T, the expected gradient of We, E [IVF(W.))) is equal to
Co²
2CD2 2
(1
Cnj) 2 +
Cni)² II (1 (1- (1-Cn;)²
k
f=0
1=0
j=++1
Proof. By itcrating over equation (55), we can get
1-1
1-1
1-1
W₁ Wo II (1-C))
Cmj)
Σ
n II (1-Cns)
i=0
j=it1
Meanwhile,
( = IVF( C² ||,
Thus, by squaring the norm of W, and taking the expectation, we can get
[VF(W))) C² || II (1 Cnj) + Σ mi II (1- (1- j=i+1
(56)
1=0
Above, we use martingale property a), which says the expectation of & is equal to zero for all i. In addition. based on property b), equation (56) is equal to
I (n-Cm) 2 E [[3]9]] -=++1
0²02 t-1
1-1
II
(1-CmH)
Th
h
j=it1
To see this, based on the martingale property of E again, the cross terms between &, and E, for different i, are equal to zero in expectation:
EE [G. = 0, for all ssjiss.
Additionally, the second moment of & satisfies:
E
||
for any i I....,T T. k
34

---

Published in Transactions on Machine Learning Research (09/2024)
Lastly, let Wo be vector such that
||oll DV2C F(W))
min F(W) s D?. WE
Setting ||Wll = DV2C-1 in equation (56) leads to
C²0² 1-1
E [VF(W))]]] 2CD2 II
(1- (1p)
+
2
II (1-C1s)
k
==
i=0
j=i+1
Thus, we conclude the proof of this result.
We now present the proof for the case when Ei=o S O(VT). For this result, we will use the following quadratic function:
T-1
f(W) 2k IWIP, where K = max (aax(C-1,2
Σ n
(57)
1=0
Lemma B.7. Consider f given in equation (57) above. For any slep sizes no..... following holds for the stochastic objective F:
1T-1 177-1 less than C-11, the
[IVF(WOI)) KEeT
D
min
2 max mcx m
Proof. The norm of the gradient of F(W) is equal to
[VF(W)]|= |W]]-
(58)
Following the update rule in NSO, similar to equation (55), W, evolves follows:
(1
W+11
17 K
W.
n
(59)
where & has variance equal to o²/k, according to the proof of Lemma B.6. By iterating equation (59) from the initialization, we can get a closed-form form equation for wr( for any 2, T:
t-1
W. Wo II
1
Σ 0=0
EhE
II
(60)
==0
j=k-1
Following equation (58). we can show that
//FF(W)
= - ?|
Thus, in expectation,
|[F(())])] = E [|
= =
1 - 11-
+ R
ThEs II (1 (1-k (1-x-2z)
3=0
1=0
j=i+1
=x 2
1 1- x - m13
+
m
1 -
n/3 E ||
z=0
j=i+1
2DD2x
(-1
1 - x 11 )
+
0²
2
[-1
1 (1- 113}} 2
(61)
k
==
t==
j=it1
35

---

Published in Transactions on Machine Learning Research (09/2024)
where the definition of initialization Wo and the variance of & in the last step. In order to tackle equation (61), we note that for all E [0, }},
2 exp log 2 =).
(62)
2
Hence, applying equation (62) to the right-haue hand side of equation (61), we obtain that for any i ==
1-1,
II
f max SC-x Ei 76
2 exp
log
1
Σ
2
1
2 ==
max (2C))
Σi= m1
/=1
Thus, equation (61) must be least
E [IVF(WO)]]] 2D 0²
(63)
The above result holds for any = 1.2...., Therefore, we conclude that
D²
min E [[F(W)]]] 2 max (Cx n1
Thus, the proof of Lemma. B.7 is finished.
Next, consider the other case, which is when the learning rates fixed.
Lemma B.8. There erists convez quadratic functions f such that for any gradient oracle satisfying Assump- tăon 4.1 and any distribution P with mean zero, if n C-1 for any = 1....,T, T, or if 2i 1/. S VT, then the following must hold:
7-
Co²
min E |IF 2 D
(64)
15ITT
32k T
Proof. By Lemma B.7, there exists a function such that the left-hand side of equation (64) is at least
D2
CD2
D2x
2
(65)
2 max CC 2i- m/ 2 max (1,x-11 2.-
2i=
4VT
which holds E " S v 1 for any fixed x > 0.
On the other hand, if E/= 2 VT and f = 11 for a fixed then > /VT By setting = 17 for all í in Lemma B.6, the left-hand side equation (64) is equal to
(1-2(1-CC))) o² t-1
min ISET
Σ (()t-())
Recall that 7 < C-1, Thus, p On must be less than one. With some calculations, we simplify the above to
2CDD2 pp) 0²p² (1-p)21
min 151ST
k
1 2
min I8tS
p²p
(1 (1- (1-p)² 2) 12CD²
k(2 ))
0²
(66)
k(2 (22-p) p)
Tf 2CD < d2p
x(2-0}: the above is the smallest when t. t= In this case, equation (66) is equal to
o²p²
1
(1
0
k
+ 6²
2CD2
36

---

Published in Transactions on Machine Learning Research (09/2024)
2CD2 d'p k(2-p}>
the above is the smallest when 1 T. In this case, equation (66) is at least
oap o²p 1/2 p)
oCx
(67)
2k
2k
VT
To conclude the proof, set that the right- right-hand side of equations (65) and (67) match each other. This leads to
20²C
x
1:D2
Thus, by combining the conclusions from both equations (65) and (67) with this value of x, we finally conclude that if Ei= Vi S v for all i 0, 1- 1) < C-1, then in both cases, there exists function f such that equation (64) holds. This completes the proof of Lemma B.8.
The case when > 0. In this case, since the update of We also depends on the momentum update, it becomes significantly more involved. One can verify that the update from step to step + 1 is based on
X₂ =
1 On
1 P
(68)
Our analysis examines the eigenvalues of the matrix X,XX and the first entry in the corresponding eigenvec- tors. Particularly, we show that the two entries are bounded away from zero. Then, we apply the Hölder's inequality to reduce the case of μ > 0 to the case μ = O, Lemma B.8 in particular.
Proof. First, consider a quadratic function
f(W) = 2C ||
Clearly, f(W) is C-Lipschitz continuous. Further, F(W) = f(W) + d, for P being the isotropic Gaussian. Let Wo be vector whose Euclidean norm equals DV2C. Thus,
F(Wo) min F(W) = D² WERd
As for the dynamic of momentum SGD, recall that M4+1 MM, TTGr and W.+1 W, M
We consider the case where We = n for all steps t. In this case, we can write the above update into matrix notation as follows:
W41 Me+1 - | -cn
-
1-Cn On # [][WA We Me
& On
On
a
&
Let X,, = [1 Cn. -nn, μ4] denote the 2 by 2 matrix (that depends on ()) above. Similar Lemma B.6, we can apply the above iterative update to obtain the formula for We+1 as:
W441 M4+
Xt
Wo
Σ Cm
E E ]. {
(69)
Mo
By multiplying both sides by the vector é1 = [T, and then taking the Euclidean norm of the vector (notice that this now only evolves that W,1 voctor on the left, and the W, vector on the right), we now obtain that, in expectation over the randomness of the & the following holds:
C2I202 k 1=0
[1 (CD'leTx(ee)"
Σ le x
(70)
37

---

Published in Transactions on Machine Learning Research (09/2024)
Above, similar to Lemma B.6, we have set the length of Wo appropriately so that its size is equal to DV2C-1, which has led to the CD term above. Recall that Mo is equal to zero in the beginning. To get the first term above, we follow this calculation:
Wo
Wo
Wo
]
6
x 1
Mo
= Tr
e
x #
Mo
Mo
X 61
CD2
=Tr
ej eTxt
X
61
0
= (e] (eT x)
We use == [] to denote the vector of ones. Now, we focus on the 2 by matrix X, (recall this is the coefficient matrix on the right side of equation (69)). Let its singular values be denoted as A and X2. In addition, to deal with equation (70), let an and o2 denote the first entry of X, left singular vectors, corresponding to and b, respectively. Thus, we write
(x) aî>? + aZ
(71)
Now, one can verify that Xf and 23 are the roots of the following quadratic equation over a:
x² - ((1 - Cn)² + (Ct))² + 22²²)= + =.
(72)
This can be checked by first laking X, times X then using the definition of the eigenvalues by caleulating the determinant of XT ld 0. Thus, we have that X and X2 equal to:
(1 - Cn) + (Cn)² + 24²² + ((1 Cn) + (Cn)) + 2n²22 4/12
A1 A2
(73)
2
Now, of (and a², respectively) satisfies that:
Cn((1 (1) af = (1 - - -Cn(11 Cn) + +2²
(74)
By enumerating the possible values of Cn between 0 and 1, one can verify that for fixed value of μt, of and az are both bounded below from zero. Therefore, we can claim that from equation (71),
a + a x
(75)
By the Hölder's inequality;
(AP* + ) (1 1)-- 2 \ι + (1 (1- + (Cn)² + 242
(76)
2 (1 Cn) - (Cn)²
(77)
which implics that
X + x > ((1 - (1)Cn) (Cm)))
(78)
2(2:-1)
Now, we consider two cases. If Cn < 1/2, then the above is greater than (1 -1-6 CT)2 which holds for any 0.I,..., 1. By way of reduction, we can follow the proof of Lemma B.8 to complete this proof. If
On > 1/2, then the above is greater than (Cn)2t, Again by following the proof steps in Lemma B.8, we can show that
Co2
min E Will
DV
K
This completes the proof of Theorem B.5.
38

---

Published in Transactions on Machine Learning Research (09/2024)
C Experiment Details
We describe the setup for Figure 2, ran on (1) a two-layer Multi-Layer Perceptron (MLP) trained on the MNIST digit classification dataset, (2) a twelve-layer BERT-Base model trained on the MRPC sentence classification dataset from the GLUE benchmark and (3) a two-layer Graph Convolutional Network (GCN) trained on the COLLAB node classification dataset. We set both MLP and GCN with a hidden dimension of 128 for model architoctures and initialize them randomly. We initialize the BERT model from pretrained BERT-Base- BenTT-Bae-Uncased. We train each model the provided training set for the training process until the training loss close to zero. Specifically, epochs. We use the model of the last epoch to measure the error in the approximation. We do this for 100 times and again measure the perturbed loss &e on the training set. We take the gap between & and l. Our measurements show that the error between the actual gap and the Hessian approximation is within 3%.
train the MLP, BERT, and GCN models for 30, 10, and 100
Table 6 reports additional comparisons between our approach and several baselines, including label smoothing (LS), random SAM (RSAM), and Bayesian SAM (BSAM). We report the test accuracy and the trace of the Hessian for the model weights the last epoch of training six image classification datasets. We observe that NSO also further reduces the trace of the Hessian and improves the test accuracy over the basclines. The largest eigenvalue reduces by 9.7%.
Table 6: Comparison between our approach (NSO), label smoothing (TS), random-SAM (RSAM), and Bayesian SAM (BSAM). Also included is the largest eigenvalue of the Hessian.
CIFAR-10
CIFAR-100
Aircraft
Caltech-256
Indoor
Retina
LS
2690 +: 85
10669 +: 363
5699 72
3482 +:2 85
3650 + 82
17681 193
RSAM
2379 + 89
9762 + 422
4665 = 95
3224 t 97
3425 + 70
16950 + 257
Trace (L)
BSAM
2768 + 54
9787 465
4750 =
3498 t 38
3162 73
16238 286
NSO
1728 + 79
5244 = 89
3678 + 83
2958 + 77
2737 + 90
10970 = 146
LS
96.9% + 0.1
83.8% + 0.1
59.0% + 0.2
76.6% + 0.2
76.5% + 0.3
64.2% + 0.7
Test
RSAM
96.8% + 0.1
84.0% 0.1
60.9%
76.4% + 0.1
76.8% + 0.5
65.9% + 0.3
Acc
BSAM
96.9% + 0.1
83.9%
0.2
61.0% 0.3
76.8% + 0.3
76.4%
+ 0.3
65.4% + 0.2
(t)
NSO
97.6% + 0.4
84.9% + 0.3
63. = 0.3
78.1% = 0.5
78.2% 0.3
67.0% + 0.4
CIFAR-10
CIFAR-100
Aircraft
Caltech-256
Indoor
Retina
SGD
1442 1 63
4639 1 95
1152 = 40
1064 + 44
1087 +: 56
8276 + 91
LS
1311 +
3051 + 95
1144 =
893 79
764 + 75
4296 74
A1
SAM
1326 + 72
2625 + 91
890 + 90
948 + 95
887 + 53
4033 + 52
(L)
USAM
1245 + 43
2299 = 98
592 = 32
782 + 38
755 + 58
3893 55
ASAM
1383 - 73
2638 + 86
615 = 95
795 + 72
697 + 36
3925 - 56
RSAM
1356 + 69
2901 - 121
895 = 74
779 68
988 + 65
4537 58
BSAM
1375 1 86
2788 1 177
972 1 79
843 97
939 1 73
4123 1 87
NSO
1070
74
2059 45
579 1 59
643 1 57
639 1 72
3681 1 66
Finally, we report the hyper-parameters for the experiments in Section 3. These include a learning rate of 0.0002, momentum of 0.99, weight decay of 0.0001, batch size of 32, and training epochs of 60. We reduce the learning rate by 0.1 every epochs. We choose these syper-parameters based on grid search on the validation split. The range in which conduct a grid search is as follows: Tearning rate: 0.005, 0.002, 0.001, 0.0005, 0.0002, and 0.0001; Momentum: 0.9, 0.95, 0.99; Weight decay: 0.01, 0.001, 0.0001; Epochs: 20, 40, and 60; Batch size: 16, 32, and 64.
Fach baseline method may have its own set of hyper-parameters, which are adjusted via a grid search. For label smoothing, we choose the weight of the loss caleulated from the incorrect labels between 0.1, 0.2, and 0.3; For SAM and BSAM, we choose the &2 norm of the perturbation between 0.01, 0.02, and 0.05; For ASAM, we choose the f norm of the perturbation for the weights between 0.5, 1.0, and 2.0; For RSAM, we choose the & norm of the perturbation between 0.01, 0.02, and 0.05 and the standard deviation for sampling perturbation betweeu 0.008, 0.01, and 0.012.
39


Reviews:
[{"review_id": "jCI7ovpaWW", "reviewer": "Reviewer_v3GM", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- This paper offers theoretical guarantees and extensive empirical evidence to support the effectiveness of the proposed algorithm.\n- The insights into noise injection algorithms are promising and could potentially stimulate future research on the geometric properties of training large models, such as transformers, as well as the development of new regularization schemes.\n\nWeaknesses:\n- I noticed that the authors uploaded a revision. The revised paper seems somewhat unpolished. \n- To support the arguments in the paper, the algorithms could be tested on a wider range of applications and datasets.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Recent literature has extensively studied the training of overparameterized neural networks. This paper introduces a Hessian-based regularization approach for training these networks, focusing on noise injection algorithms. These algorithms have a desirable theoretical property that allows for an unbiased estimation of the Hessian. However, the paper finds that naive noise injection does not perform well empirically. To address this issue, the authors propose a new two-point noise injection algorithm that effectively cancels out the first-order gradient term. This algorithm demonstrates strong performance in practice, as evidenced by tests on various image classification datasets. Additionally, the paper derives generalization bounds and convergence rates to theoretically support the proposed algorithm.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- It would be helpful to polish the writing of the paper. \n- To generate broader interest, it would be beneficial to demonstrate the effectiveness of the proposed algorithms in additional application domains.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "1EDuKWKSEb", "reviewer": "Reviewer_Z2QG", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n\n- The paper is well structured, and the idea is clearly presented. \n- The results unequivocally show the superior performance of the proposed algorithm.\n- The authors have done a very thorough numerical analysis to validate many aspects of the proposed work also when combined with other existing methods. I found particularly interesting the ablation study. \n- I find the main idea behind this work to be intuitive, natural, yet very effective.  \n\n\n## Weaknesses\n- I find the writing often lacking clarity. For example, some sentences are particularly convoluted, and repetitions sometimes occur (e.g., I assume there are some problems with cross-referencing in latex as it sometimes appears as \"equation equation X\").\n- I find the structure suboptimal. The **Results** section on page 2 feels a bit out of place and breaks the flow of the paper. Moreover, I find the paper very hard to follow from Example 1.1 onwards. Although I assume the goal of the authors was to provide examples to enhance clarity, this had the exact opposite effect on me (for context, I am not an expert on this specific subfield, so most of the techniques discussed in this work were new to me). \n- For being part of the **introduction** Example 1.1 already gets to technical and I felt lost already because of the heavy notation. \n- I found the **related work** section to be hard to parse. In fact, this is, at times, repetitive with respect to the previous section. Moreover, the authors discuss some evidence from their empirical analysis in the **related work** section, which I found confusing. I'd recommend leaving out any comment on the present work and just focusing on the findings, advantages, and limitations of previous works. \n- From page 5 until page 8, I found the paper extremely hard to follow. Specifically, I find the notation particularly heavy and dense. I'd like to know if this can be restructured in a way that highly technical details and proofs are left out and the focus is on the main contribution, i.e., injecting noise in both positive and negative directions.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a new algorithm to inject noise in the weights during training of a Neural Network. The noise injection is known to make the training more stable and improve the generalization of the model, thus making it more robust with respect to being trapped in local minima. A crucial difference of the proposed method wrt extisting algorithms is to use negative perturbations to the model's weights in order to zero out the first-order in Taylor's expansion of the function of the weights, i.e., $F(W)$ in the paper, and in turn lead to a better estimate of the gradient. \n\nThe theoretical insights presented in the first part of this paper are validated with an extensive set of numerical experiments showcasing the advantages of the proposed method.\n\nAs a side note to the authors and the AE: I've noticed a new version of the manuscript (revised) has already been submitted. Unfortunately, the version I reviewed was the original one. The authors should please keep this in mind when going through my comments. It could be that some of the concerns raised below have already been addressed.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "### Requested changes, recommendations and questions\n\n- I'd recommend running a spellchecker through the manuscript and checking all the cross-refs in LaTeX. \n- I'd recommend editing the structure of the manuscript by removing (or at the very least heavily restructuring) page 2. In particular, I'd recommend focusing on the very general idea in very intuitive and easy terms, anticipating the structure of the paper and clearly listing the main contributions. \n- I'd recommend making a sketch/visual that represents the core idea of this work to put on page 2 instead of (or along with) the example 1.1. This should help the reader to better follow the flow.\n- Make sure to introduce all acronyms. For example, Sharpness-Aware Minimization has never been explicitly defined as SAM, and it took me a while to realize this. The same may apply elsewhere. \n- Algorithm 1: I'd suggest to add a comment on line 2: highlighting that the loop is over different perturbations. \n- I feel that Figure 1 is out of place on page 3. I'd rather, as mentioned earlier, show a visual describing how the algorithm works and leave this for the result section, after the theory part. \n- Currently, the bullet points that summarize the contributions of this work are very verbose and do not go straight to the points. I'd recommend shortening them and stating more clearly what the modified algorithm actually does. \n- I cannot see where the variable $\\epsilon$ in the middle of page 5 has been defined. \n- To which extent are Assumptions 3.2 and 3.3 *standard*? I think the clarity of the paper would benefit if the authors could give a brief explanation of why those assumptions could be considered \"easy to realize\" in practice.\n- Section 3.2 uses a specific example to show that gradient norms bounds are tight. I wonder how general this result is and if it is really necessary for the development of the main story of the paper. While reading the manuscript I perceived this section rather as an add-on instead of a necessary contribution (in light of the results shown in the **Experiments** section). In order to enhance clarity and streamline the reading, I wonder if it would be a suitable option to relegate all these details to the appendix and perhaps spend only a couple of sentences describing this in the main text. It may also be the case that I am overlooking something here. In that case, I'd be grateful if the authors could make me aware of this. \n- If necessary, I'd introduce Example 1.1 at the beginning of Section 4.1. People who are not strictly familiar with the field may have forgotten what Example 1.1 was when they reached this point of the paper. \n- I find Table 2 a bit misplaced. This is not really a *contribution of this work* but rather an ablation study of SGD and WP SGD. While I think this is still interesting to show, I'd add NSO there as well to compare directly with the method introduced in the manuscript. \n\n### Open Question to Authors\n\n- At the bottom of page 11, the authors note that using $k\\geq 3$ perturbations does not bring any obvious improvement. Have the authors made any investigation about the correlation between the number of perturbations and the level of noise injected? In other words, if I think of making parallelism with annealing, I might imagine that having more levels of perturbations, each of which injects less noise at each step, might make the procedure even more robust. Do the authors have an intuition as to whether this speculation is correct and why?", "broader_impact_concerns": "I could not see any concerns regarding the ethical implications of the submitted work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Z29XiWVhIJ", "reviewer": "Reviewer_X1ow", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The overall idea of the method is simple and interesting. It feels like an in-hindsight-obvious modification of weight noise - which is to the authors' credit. The experimental results in the fine-tuning settings seem promising in general. There are however some major concerns with the experiments; these are listed in \"requested changes\".\n\nThe presentation overall is very confusing. It seems to go back and forth between the theory, experiments, and different versions of the algorithms without guidance for the reader. More specific critiques and suggestions can be found in \"requested changes\".\n\nI'm overall confused at the point that sampling multiple perturbations per iteration is essential to the method's success, and that simultaneously the method works with only 1 sample.\n\nThe experimental section would be strengthened if the authors included some experimental results on how NSO operates in (pre)training rather than fine-tuning tasks - an area where SAM tends to do well. I don't think a negative result here would weaken the paper but would instead help practitioners understand more about NSO. There is also a question about the batch sizes used in \"requested changes\" that I ask the authors address.\n\nI believe the method (and the paper) are promising, but further work is needed to clarify the claims and the presentation. I believe the theoretical analysis is sound but I haven't checked some of the details carefully.\n\nEdit: Claims and evidence update to yes following rebuttal period.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors analyze training under weight noise, a scheme which has largely been abandoned for its lack of efficacy despite its potential as an implicit regularizer of statistics of the Hessian. The authors then propose Noise Stability Optimization, which trains with pairs of noised points, where each pair involves perturbations $\\pm U$ respectively. They show that this specific coupling reduces the influence of low order terms from a Taylor expansion of the weight-noise objective, and improves finetuning over SAM and related methods in a variety of settings. They additionally provide theoretical analysis of PAC-Bayes generalization bounds as well as convergence analysis of the new method.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please edit the abstract. The first part reads like a mathematical proof and not an informative summary for a reader. The overall presentation could use some work as well; there are already detailed mathematical equations in the first paragraph for example.\n\nThere are a few concerns about the experiments:\n1. It's mentioned in Appendix C that the experiments were done without momentum and weight decay. What happens when these are included? It seems very odd to present results without these basic features of modern deep learning pipelines.\n2. The radii selected for SAM may be too large; in general the ratio of $\\rho/\\eta$ (SAM radius to learning rate) gives a dimensionless number which gives a sense of the strength of SAM.\n3. In that vein, unnormalized SAM is a better comparison for NSO. See for example [2]. In this setting one can search over equivalent regularization strengths. Generally $\\sigma^{2}$ from NSO should be compared to $\\rho$ from SAM - another point which suggests that perhaps the SAM radius was set too large.\n\nSmall changes:\n\nIn the experiments for Table 1, what are the results training with normal weight perturbation and not the paired perturbation introduced in the paper?\n\nThere are many typos and incomplete sentences, e.g.: \"In the case that f is a strongly convex function, the lowest eigenvalue of the Hessian is above from below\" - what does this mean?\n\nKey definitions like $\\ell_{Q}$ are buried in the middle of proofs, even though they are later used in experiments.\n\nTable 2 should be a plot not a table.\n\nOne of the key practical features of SAM seems to be its sensitivity to batch size. What batch sizes were experimented with for the NSO-SAM comparisons? Does NSO have a different sensitivity to batch size compared to SAM?\n\nMinor points:\n* The introduction discusses the link between edge of stability and SAM, [1] is a key reference here. It also may be relevant for points made about the difficulty of analyzing SAM at convergence - something which may not be relevant for analysis of practical algorithms.\n* For the experiments involving the Hessian trace: what activation function was used in the networks? For ReLU networks the trace mainly captures the Gauss-Newton component of the Hessian, see [2] for example\n* Proposition 2.1 does not require the paired noise construction to work, including it in the proof is confusing - unless one can also make the point that it reduces the variance of the linear term in finite samples. I think it should be modified to include such a statement.\n* Reference [2] has some discussion and experiments about training with weight noise which may also be relevant for the analysis.\n\n[1] https://proceedings.mlr.press/v202/agarwala23a\n[2] https://arxiv.org/abs/2401.10809", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "2Eiydzvxpm", "reviewer": "Reviewer_wuA1", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strength:\n\n1) Their algorithm is simple and intuitive.\n\n2) They provided numerous experiments to prove the superiority of their algorithm. Maybe the practitioners will like this paper.\n\n\n# Weakness:\n\n1) Even though they emphasized the connection with PAC-Bayes bound throughout the paper, their result itself is not much related to PAC-Bayes. They don't provide any PAC-Bayes bound. What they did was simply sweep the hessian chores to PAC-Bayes result on page 6,  **Illustration** paragraph. \n\n2) Their theoretical result itself is relatively weak. Their hessian result is empirical, and even the effect of hessian is not clearly stated (just illustration paragraph). Their theoretical result is only about the scale of $\\|\\nabla F(W)\\|$, and they didn't provide enough theoretical comparison between other works. The objective of optimization is reducing the loss (maybe in this paper generalization error) but they didn't prove this smaller gradient norm leads to smaller generalization error mathematically. They all compared numerically, so as a theoretician, I am not clear about the superiority of their result. \n\n3) (Not sure, I will remove this when it is true) I'm not sure whether this work is the first 'canceling trick of perturbation' in the field of optimization. I will also check other reviewers and literature to see whether their work is really novel.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors proposed a new algorithm called Noise Stability Optimization. The idea itself is very simple - to avoid introducing first-order term $U^\\top \\nabla f(W)$, use $\\frac{1}{2}(g(W+U)+g(W-U))$ to cancel out the first-order perturbation. \nThey give upper and lower bounds on the gradient norms of the converged solution of their proposed algorithm. The bounds can also be extended to momentum updates.\n\nThey also conduct multiple experimental evaluations of their algorithm, compared with a number of “sharpness” reducing training methods. By their experiment, they empirically validate the regularization of the algorithm on the Hessian matrix.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "1) Even though they frequently mentioned the word PAC-Bayes, there is no PAC-Bayes generalization bound result in this main paper. It is very misleading (especially for the readers who are experts in PAC-Bayes but not in optimization). For completeness, it would be great if authors clearly state the effect of smaller hessian in the Appendix or somewhere. That said, I've read the related works and illustration paragraphs already. I want something more mathematical and rigorous statement to see how it works. \n\n2) Also, their main result is mainly focusing on minimizing the gradient norm, but it is not crystal clear why minimizing the norm leads the better generalization. Demonstrating the loss convergence, even in a strongly convex environment, would be better to convince readers I think.\n\n3) Is this the first paper using 'canceling perturbation'? I mean, $g(W-U)+g(W+U)$ type of approach to cancel out the first-order term is quite common in many statistical machine learning studies. \n\n4) Minor: It would be great if authors could cite some studies that use Assumptions 2 and 3, to support the generality of those assumptions.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: After a detailed and productive discussion, the paper received uniformly positive reviews, with reviewers appreciating the simplicity of the approach and strong experimental results. I second these impressions and recommend acceptance to TMLR. However, I noticed a few points that still need to be addressed in the final revision of the paper.

1. The theoretical guarantees in this work hold equally for WP-SGD and NSO. This is a significant limitation, because a main (empirical) claim of the paper is that NSO generalizes better than WP-SGD. The revision needs to state this limitation clearly, in a paragraph or section titled “Limitations.”

2. The above point makes the empirical superiority of NSO to WP-SGD mysterious. Viewing NSO/WP-SGD simply as implementation of Hessian-regularized ERM, the lower variance of the NSO gradient estimator should only speed convergence, not improve generalization. Please discuss this point and try to provide some intuition for why NSO generalizes better.

3. The experiment in Section 5 needs some clarification:
   * Are both “NSO” and “WP-GD” applied on exact empirical gradients (as the name of the latter implies?)
   * In light of points 1 and 2 above, does the generalization gap between NSO and WP-GD persist even when thoroughly tuning the learning rate and number of steps for each method? Demonstrate this by showing the test MSE as function of learning rate and number of steps.

4. The novelty of the theoretical components in this paper is also not sufficiently clear. In Section 2 and Section 4, the authors need to clearly state what is the difference between their result and the closest results in prior work, and what are the new ideas. In particular, Theorem 4.2 seems to be a simple corollary of the results in Ghadimi and Lan (2013), i.e., there are no new ideas. To make this clearer please, the proof of Theorem 4.2 should cite Theorem 2.1 of Ghadimi and Lan in a black box way, rather than re-prove it from scratch.

5. The paper is missing a reference to Duchi et al.’s work on randomized smoothing [1], which to my knowledge was the first to propose what you call WP-SGD. More recently, randomized smoothing also became popular as a means for obtaining provable adversarial robustness guarantees [2,3]. The revision should cite and discuss these works.

6. A number of formal statements need improvements:  
    * Generalization bounds such as Theorem 2.1 are missing an all-important “for all” quantifier. For example, Theorem 2.1 needs to have a “for all $W\in\mathcal{H}$” appear right before equation (2). (I did not verify this indeed holds, but if not then Theorem 2.1 is not a PAC guarantee.)  
    * In Theorem 4.2, the statement “$f$ is Lipschitz-continuous” should be “$\nabla f$ is $C$-Lipschitz continuous”; this error seems to partially propagate to the proof. Note that by using the analysis from [1] one can also obtain a guarantee assuming only that $f$ (and not $\nabla f$) is Lipschitz, but this comes as a cost of introducing dimension-dependence to the rate of convergence. 


[1] Duchi, John C., Peter L. Bartlett, and Martin J. Wainwright. "Randomized smoothing for stochastic optimization." SIAM Journal on Optimization 22.2 (2012): 674-701.

[2] Cohen, Jeremy, Elan Rosenfeld, and Zico Kolter. "Certified adversarial robustness via randomized smoothing." international conference on machine learning. PMLR, 2019.

[3] Yang, Greg, et al. "Randomized smoothing of all shapes and sizes." International Conference on Machine Learning. PMLR, 2020.
Decision: Accept (poster)