Parameter Averaging for SGD Stabilizes the Implicit Bias towards Flat Regions
Authors: Atsushi Nitanda
Ryuhei Kikuchi
Shugo Maeda
Venue: Submitted to ICLR 2023 (2023)
Keywords: stochastic gradient descent
large step size
parameter averaging
implicit bias
flat minima

Abstract: Stochastic gradient descent is a workhorse for training deep neural networks due to its excellent generalization performance. Several studies demonstrated this success is attributed to the implicit bias of the method that prefers a flat minimum and developed new methods based on this perspective. Recently, Izmailov et al. (2018) empirically observed that an averaged stochastic gradient descent with a large step size can bring out the implicit bias more effectively and can converge more stably to a flat minimum than the vanilla stochastic gradient descent. In our work, we theoretically justify this observation by showing that the averaging scheme improves the bias-optimization tradeoff coming from the stochastic gradient noise: a large step size amplifies the bias but makes convergence unstable, and vice versa. Specifically, we show that the averaged stochastic gradient descent can get closer to a solution of a penalized objective on the sharpness than the vanilla stochastic gradient descent using the same step size under certain conditions. In experiments, we verify our theory and demonstrate this learning scheme significantly improves performance. 

Under review as a conference paper at ICLR 2023
PARAMETER AVERAGING FOR SGD STABILIZES THE IMPLICIT BIAS TOWARDS FLAT REGIONS
Anonymous authors Paper under double-blind review
ABSTRACT
Stochastic gradient descent is workhorse for training deep neural networks due to its excellent generalization performance. Several studies demonstrated this success is attributed to the implicit bias of the method that prefers flat minimum and developed new methods based on this perspective. Recently, [zzmailov et al. (018) empirically observed that averaged stochastic gradient descent with large step size can bring out the implicit bias more effectively and can converge more stably to flat minimum than the vanilla stochastic gradient descent. In our work, we theoretically justify this observation by showing that the averaging scheme improves the chas-optimization tradeoff coming from the stochastic gradient noise: a large step size amplifies the bias but makes convergence unstable, and vice versa. Specifically, we show that the averaged stochastic gradient descent get closer to a solution of a penalized objective on the sharpness than the vanilla stochastic gradient descent using the same step size under certain conditions. In experiments, we verify our theory and show this learning scheme significantly improves performance.
INTRODUCTION
Stochastic gradient descent (SGD) Robbins & Monro 1951 training modern deep neural networks. In order to further improve the performance, a great deal of SGD variants such adaptive gradient methods has been developed. However, SGD is still the workhorse because SGD often generalizes better than these variants even when they achieve much faster convergence regarding the training loss (Keskar Socher 2017: Wilson al.] 2017 Luo et al. 2019). Therefore, the study of the implicit bias of SGD, explaining why it works so better, is nowadays an active research subject.
powerful learning method for
Among such studies, flat minima (Hochreiter & Schmidhuber] 1997) has been recognized as an important notion relevant to the generalization performance of deep neural networks, and SGD has been considered have bias towards flat minimum. Hochreiter & Schmidhuber (197)); Keskar et al. (2017) suggested the correlation between flatness (sharpness) and generalization, that is, flat minima generalizes well compared to sharp minima, and Neyshabur al. (2017) rigorously supported this correlation under fz-regularization by using the PAC-Bayesian framework (McAllester) 1998. 1999). Furthermore, by the large scale experiments, Jiang al. (2020) verified that the flatness measures reliably capture the generalization performance and are the most relevant among 40 complexity measures. In parallel, Keskar al. (2017) empirically demonstrated that SGD prefers a flat minimum due to its own stochastic gradient noise and subsequent studies (Kleinberg et al. 2018 Zhou al., 2020) proved this implicit bias based on the smoothing effect due the noise and stochastic differential equation, respectively.
Along this line of research, there are endeavors to enhance the bias aiming to improve performance. Especially, stochastic weight averaging (SWA) (Izmailov et al.) 2018) and sharpness aware minimiza- tion (SAM) (Foret et al. 2020) achieved significant improvement in generalization performance over SGD, SWA is cyclic averaging scheme for SGD, which includes the averaged SGD (Rupperi) 1988) Polyak & Juditsky 1992) as a special case. Averaged SGD with an appropriately small step size or diminishing step size to is well known to be efficient method that achieves statistically optimal convergence rates for the convex optimization problems (Bach & Moulines 201T Lacoste- et al. 2012 Rakhlin al. 2012). However, such a small step size strategy does not seem useful for

---

Under review as a conference paper at ICLR 2023
stop sizew0:00 stop size-0.25 step sicemb.30
step 5,
stop size=0.50
stop siz0-0.05
stop size-0.35
34
35-
stop stzc=0.30
stop 0i20-0.15
StOp sizo-0.11
stop size-0.03
step 5izo-0.01
stop siza-0.00
step sizem0.15
(a) (smoothed) objective
(b) SGD
(c) averaged SGD
Figure 1: We run SGD and averaged SGD 500 times with the uniform stochastic gradient noise for two objective functions (top and bottom). Figure (a) depicts the objective function f (green, 1) = 0) and smoothed objectives F (red and blue, 77 > 0). Figures (b) and (c) plot convergent points by SGD and averaged SGD with histograms, respectively.
training deep neural networks, and Izmailov et al. (0188) found the averaged SGD with not small but large step size works quite well.
The success of using large step size can be attributed to the strong implicit bias as discussed in Izmailov et al. (2018). SGD with a large step size cannot stay in sharp regions because of the amplified stochastic gradient noise, and thus it moves to another region. After a long run, SGD will finally oscillate according to an invariant distribution covering a flat region. Then, by taking the average, we can get the mean of this distribution, which located inside flat region. Although this provides good insight into how the averaged SGD with large step size behaves, the theoretical understanding remains clusive. Hence, the research problem we aim to address is
Why does the averaged SGD with a large step size converge to a flat region more stably than SGD?
In our work, we address this question via the convergence analysis of both SGD and averaged SGD.
1.1 CONTRIBUTIONS
We first explain the idea behind our study. Our analysis builds upon the alternative view of SGD (Kleinberg et al., 2018) which suggested that SGD implicitly optimizes the smoothed objective function obtained by the convolution with the stochastic gradient noise (see the left of Figure I Since pointed out later the smoothed objective is essentially a penalized objective on the sharpness whose strength depends on the step size, the more precise optimization of the smoothed objective with a large step size implies the convergence is known to control the optimization accuracy of SGD, that is, we need to take a small step size at the final phase of training to converge. These observations indicate the bias bias-optimizatiin tradeoff coming from the stochastic gradient noise and controlled by the step size:
flatter region. At the same time, the step size
A large step size amplifies the bias towards a flat region but makes the optimization for the smoothed objective inaccurate, whereas a small step size weakens the bias but makes the optimization accurate.
In our work, we prove that the averaged SGD can improve the above tradeoff. that is, it can optimize the smoothed objective more precisely than SGD under the same step size. Specifically, we prove as long as the smoothed objective satisfies one-point strong convexity at the solution and some regularity conditions, SGD using the step size 7) converges to a distance ((vT)) from the solution (Theorem whereas the averaged SGD using the same step size converges
distance O (Theorem 2
We remark that large step size in our study means the step size with which SGD oscillates and poorly performs but the averaged SGD works. Clearly, larger step size regardless of the condition of
2

---

Under review as a conference paper at ICLR 2023
the objective will diverge. thus it should be appropriately small to achieve sufficient optimization. The better dependence of (()) than O(vī) means that the averaged SGD can work well with wider range of step sizes than SGD. Although, too small step-size does not always bias the solution because the deviation of the solution is (17)), the above dfference of the order can make the separation between SGD and averaged SGD with an appropriately chosen step-size depending on the problem. As a result, we can expect the improvement by the averaged SGD for datasets such that the stronger implicit bias with the appropriately larger step size is useful.
The separation between SGD and averaged SGD regarding the bias can occur even the simple setup as seen in Fig- ure which depicts obtained parameters by running SGD and averaged SGD 500 times in two cases. We observe (a) both methods with the small step size can get stuck at sharp valleys or an edge of a flat region because of weak bias and accurate optimization, (b) SGD with the large step size amplifies the bias and reaches a flat region but is unstable, and (c) averaged SGD with large step size converge stably near biased solution which minimizes the smoothed objective. The behav- ior of the averaged SGD in asymmetric valley (the top of Figure D. that the parameter is biased toward a flat side from an edge of the region, is also known to be preferable property in generalization as well as flat minima (see Izmailov et al. 20188: He et al. (2019). ). We note that this phenomenon is certainly captured by our theory. Indeed, Figure shows the
.0.0
size- 0.10
step strp= 30
stop 5.20
size-0,70
step s.5e
35
The figure plots the origi- nal objective (green), smoothed objec- tives (blue, darker is smoother), and con- vergent points obtained by the averaged SGD which is run for the asymmetric valley objective 500 times for each step size E (0, 1, 0.3, 0.5, 0.7, 0.9).
Figure 2:
convergent point of the averaged SGD is almost the minimizer
of smoothed objective for each step size.
Our findings are summarized below:
SGD and averaged SGD implicitly optimize the smoothed objective, whose strength depends on the step size, up (Vv and (n) errors in Euclidean distance from the solution. This explains why these methods reach a flat region with an approprie step-size, since smoothing eliminates sharp minima.
This means that averaged SGD can optimize the smoothed objective more precisely than SGD under the same step size as long as required conditions uniformly hold regarding the step size, resulting in a stronger bias towards flat region. In other words, averaged SGD better controls the bias-optimization tradeoff than SGD.
Hence, the parameter averaging yields an improvement for difficult datasets such that the stronger implicit bias with the larger step size is useful. This suggests the use of larger step size for such datasets so that averaged SGD stably converges but SGD itself unstable to effectively bring out the implicit bias.
Technical difference from Kleinberg et al. (2018). The proof idea of Proposition Trelies on the alternative view of SGD (Kleinberg al.) 2018)) which shows the existence of an associated SGD for the smoothed objective. However, since its stochastic gradient a biased estimator, they showed the convergence not to the solution but to point at which sort of one-point strong convexity holds, and avoid the treatment of biased estimator. Hence, the optimization of the smoothed objective is not guaranteed in their theory. On the other hand, optimization accuracy is the key in our theory, thus need nontrivial refinement of the proof under a normal one-point strong convexity at the solution.
2 PRELIMINARY - STOCHASTIC GRADIENT DESCENT
In this section, we introduce the problem setup and stochastic gradient descent (SGD) in the general form including the standard SGD for the risk minimization problems appearing in machine learning.
Let f Fdd assume f is nonnegative. A stochastic gradient descent, randomly initialized at Wo. for optimizing f described as follows: for t 0, 1,,,...
be smooth nonconvex objective function to be minimized. For simplicity, we
W++1 w/ 17 (VI(w)) <etr(w:))
3

---

Under review as a conference paper at ICLR 2023
where 1) > is the siep-size and E++1 IR + Rd is random field corresponding the stochastic gradient noise i.c., for any w E IR, {4+1 [w] is a sequence of zero aero-mean random variables taking values in Rd. A typical setup of the above is an empirical/exppeted risk minimization in machine learning.
Example 1 (Risk Minimization). Let l(w, 2) be loss function consisting of the hypothesis function parameterized by w E R and the data & E R². Let μ be an empirical/true data distribution over the data space and Z. be random variable following μr. Then, the objective function is defined by
[ww) Ez-a Z)}.
Given ii.dd random variables {Z++1 10c with the same distribution as Z, the standard stochastic gradient at 1-th iterate W. is defined V Voll(r, 12+1). In this setting, the stochastic noise et+11 can be (++((w) = V w Z++1) - V V((w). Note that we can further include the fz-regularization in the objective and the perturbation by the data augmentation in the distribution μ.
As this example satisfies, we suppose {<++1]= are independent copies each other. That is, there is a measurable map from probability space: 2 e(w, =) R and then 61+1 be written as a measurable map from a product probability space: Ω820 Đ { === + <(w, 2++1) E Rd when explicitly representing them as measurable maps. Moreover, we make the following assumptions on the objective function and stochastic gradient noise. Assumption 1.
(A1) f: IR - TR is nonnegative, twice continuously differentiable, and its Hessian is bounded, i.e., there is constant L, > O such that for any R V2 (w) LI.
(A2) Random fields {c(4.1 1 are independent copies each other and each 6241 (w) is differentiable in w. Moreover, for any E Rd "ed"EErrs 1lex+1(w)
and there are σ1, 02 such that for any w E R2d ([[]le+1(ww]++ S of and B1I cx+1 (w) S 02. where Jer1 is Jacobian of t++1-
The smoothness and boundedness conditions (A1) on the objective function and the zero-mean and the bounded variance conditions (A2) stochastic gradient noise commonly assumed in the convergence analysis for the stochastic optimization methods. Moreover, if Hessian matrix satisfies -LI LT V ((,, 2) < LI in Example then the last condition on Jey1 also holds
Remark.
with at least d2 2L because Je (w) Va ((w, Ze-1) V2f(w)
3 ALTERNATIVE VIEW OF STOCHASTIC GRADIENT DESCENT
An alternative view (Kliineerg et al. 2018) of SGD is the key in our analysis relating to an implicit bias towards a flat minimum. We introduce this view with refined convergence analysis and see the bias-optimizationn tradeoff caused by the stochastic gradient noise with step size.
An alternative view of SGD considers an associated iterations {uz J with {wy) which ap- proximately minimizes a smoothed objective function obtained by the stochastic gradient noise. We here define 14 as parameter obtained by the exact gradient descent from wy, that is, Vt W/ f and analyze the update of U2 instead of wy. Since w1-1= = we get Ut+1 = 2/2 - 17:+1( (wz) - f(e, - (1:++1(w;)) As shown in Appendix[A. 1 under a specific setting given later, w
77 (we).
w f becomes smooth invertible injection and its inverse is differentiable, thus, we identify (1) with (+)1(w) through the map w + v. Then, we get an update rule of vy:
Uc+1 0t 176(411 - - (1(4
For convenience, we refer to the rule (2) as implicit stochastic gradient descent in this paper. Since, the conditional expectation of ( at ie is zero, we expect that the implicit SGD (2) minimizes the following smoothed objective function:
F(v) = Eif(r np (v))),
(3)
where is an independent copy of é₁, 62. ... However, we note that this implicit SGD is not a standard SGD because f 7 (vx) is biased estimate of V FV)) (i.c., VF(o) + V
mp' (+))) in genera[] [ and thus we need a detailed convergence analysis.
By the construction, there is probability space (Ω, F, P) such that e'(v) can be represented as measurable map f 3 e(( z). Then, F(r) Elf(v-ne' (v)) f(o me' (v. P))dP(s) and F(r) Vff(-
4

---

Under review as a conference paper at ICLR 2023
The function (3 is actually smoothed function of f by the convolution using the stochastic gradient noise yje' and the level of smoothness is controlled by the step-size l, as seen in the left Figure D which depicts the original objective corresponding
= 0 and smoothed objectives F. In this figure, we can observe how a nonconvex function is smoothened and its sharp local minima are eliminated by appropriately large step size (the bottom-left figure) and how the solution is biased toward the flat side in an asymmetric valley (the top-left figure). Hence, we expect that stochastic gradient descent can avoid sharp minima and converges to a flat region. Indeed, by taking Taylor expansion of f. we see that F(v) is an approximation of the function f(u) plus the penalization on the high (positive) curvature of f along the noise direction expectation:
' Tr (VIOO)((())) +0
F(w)= f(v)+
(4)
2
The above observation indicates the reasonability of imposing some sort of convexity conditions at the solution of the smoothed objective F(v) rather than the original objective f(w). In this paper, we make the following one-point strong convexity at the solution us to show the convergence of F(e)). Let v. = arg min F(v). We note that F and v, depend on the value of 17, but we do not explicitly denote this dependency for simplicity.
Assumption 2.
(A3) There is c > such that for any E TR.,, (u - v.) 2 c/lu v.l1?.
For instance, this assumption holds for the function in the bottom-left Figure with sufficiently large and for the function in the top-left figure with any 7) in certain interval (0, vol.
Assumption (A3) a normal one-point strong convexity, whereas Kleinberg et al. (2018) assumed a to |12 at some parameter and showed (e)( then vo should be v, and both assumptions coincide. However, noted above VF(v) & EVV f(e ne'(v)) in general, and hence is not necessarily v.. Our aim is to clarify how precisely SGD and averaged SGD can minimize F(v). That is why make the normal one-point strong convexity at v, and need much more detailed analysis. Moreover, our proof allows for a larger step size than that in Kleinberg et al. (2018) because of the different proof techniques.
different condition: E/VS(v 176' (v)) (u vo) 2 cllv v
the convergence to v.. If VF( = EfV
Theorem 1. Under Assumption (A1), (A2), and (A3), run SGD for T-iterations with the step size 7 S 21. then a sequence {v, F of the implicit SGD satisfies the following inequality:
fno
8n
(1+ 21102 11+2003
T+1 Σ Bllee v. - (T
3c
Remark. If = O, then SGD is nothing but deterministic gradient descent and f F because of the absence of stochastic gradient noise. Hence, SGD converges to minimizer of according to the classical optimization theory, which is recovered by Theorem with 0 = 0.
This theorem shows the convergence of SGD to the minimum of the smoothed objective up to distance ( from as long as F satisfies required assumptions even if the original objective has local minima. This also true for w; since ||w; - well = O()). Thus, convergence to a flatter regularized objective 4 Moreover, we can
region is expected through explicit expression see from the theorem the optimization accuracy becomes more accurate by using a smaller step size for the problem where the required conditions uniformly hold regarding 17. On the other hand, small step size clearly weakens the bias. Thus, the step size v) controls the bias-optimization tradeoff' coming from the stochastic gradient noise.
AVERAGED SGD WITH LARGE STEP-SIZE
Izmailov et al. (018) empirically demonstrated that averaged SGD converges to flat region and achieves better generalization even when SGD oscillates with a relatively large step size. We
yje' (c))) = / (v((vej (()==="((ee=)))ddP({)) =)))d((s) whereas
(v)) in Eq. 2 means V [(w))u=v--ee(t) which
does involve the derivative of (()). Therefore, F A E f(x me (v)) in general.
5

---

Under review as a conference paper at ICLR 2023
theoretically attribute this phenomenon to that the averaged SGD get closer to than SGD using the same step size under certain settings. In other words, parameter averaging can improve the bias-optimization tradeoff and bring out the implicit bias more effectively. In the averaged SGD, run normal SGD ① and take the average as follows:
T+1 Σ wt
WT+1
T
Our aim to show lim WT can be closer to u, than {wr)i_n and {or by clarifying the dependency of this limit on the step size 7). Preferably, the implicit SGD (2) is more useful in analyzing the averaged SGD because the average UT Et Us consistent with Wy confirmed below. By the definition, we see
WT+1 = ūT
c1+1 (w,),
0
where the noise term Ei=o (+41((x)) (T+]) is zero in expectation and its variance is upper bounded by <f/ar+1) under Assumption (A2). Hence, WT+1 Uy converges to zero in probability by Chebyshev's inequality; for any 0, P हा || > r] < o/(r+1>>² + 0 as T cc, and the analysis of limy-+00 Wy reduces to that of limy-so UT.
We further make the additional assumptions the smoothed objective F Rd R and give the theorem that shows the convergence of the averaged SGD.
Assumption 3.
(A4) There is M > 0 such that for any v E Rd (VF(()) (F a, S M| 2 1.11?.
(A5) 4 F(.)) positive, i.e., there is μ such that V² 2 pl.
Remark. (A4) is used to show the superiority of the averaging scheme. This condition can be derived by the boundedness of the third-order derivative assumed in Dieuleveut al. (2020). The positivity of Hessian (A5) is only required at v., which is consistent with nonconvexity. For instance, examples in Figure satisfy (A5).
Theorem Under Assumption (A5), run the averaged SGD for T T-iterations with the step size S ZL. then the average Uy satisfies the following inequality:
[B|T]] - r.ll S (T+ 40 101021
2n17
8n 8170
3 (1+
2n10 21+210
+
v3p
cp
3cp
The variance of the averaged parameter is typically small, hence we evaluate the distance of E[vr] to Vs. Indeed, this is reasonable because the central limit theorem holds for averaged SGD under the mild conditoin even for nonconvex problems (Yu et al. [2220). Theorem 2 says that the averaged SGD can optimize the smoothed objective F with better accuracy of O(1) than O(vn) achieved by SGD using the same step size long the required conditions (one-point strong convexity minimizer and regularity for smoothed objectives) are satisfied uniformly for in certain interval (0, no] (3%o 1). These uniform requirements hold for valleys like the top-left case of Figure] | and likely holds in the final phase of training deep neural network because of the observation that the parameter eventually falls in a shater- valley (see Figure 4 Therefore, we expect the averaged SGD to outperform the normal SGD in such cases, and we recommend the use of the tail- averaging scheme for deep learning as adopted SWA (Izmailov et al., 2018), whose benefit is well known even in the convex optimization (Rakhlin et al. 2012 Mücke et 2019).
5 EXPERIMENTS
We evaluate the empirical performance of SGD and averaged SGD on image classification tasks using CIFARIO and CIFAR100 datasets. To evaluate the usefulness the parameter averaging for the other methods, we also compare SAM (Foret et al. (020)) with its averaging variant. We employ the tail-averaging scheme where the average is taken over the last phase of training.

---

Under review as a conference paper at ICLR 2023
Table 2: Comparison of test classification accuracies on CIFARIO0 and CIFARIO datasets.
CIFAR 100
CIFAR10
1) ResNet-50
WRN-28-10 10
Pyramid
1
ResNet-50
WRN-28-10
Pyramid
SGD
80.83 (0.21)
81.81 (0.29)
81.43 (0.32)
8 95.95 (0.11)
96,85 (0.16)
96,41 (0.22)
Averaged
5 82.13 (0.22)
83.13 (0.13)
84.23 (0.03)
96.58 (0.14)
97.24 (0.07)
97.07 (0.08)
SGD
82.87 (0.13) 1
84.23 (0.10)
85.12 (0.20)
m
96.89 (0.05)
97.44 (0.04)
97.28 (0.13)
SAM
82.56 (0.14)
83.80 (0.27)
84,59 (0.24)
8 96.34 (0.12)
97.14 (0:05)
97.34 (0.03)
Averaged SAM
$
82.64 (0.12)
84.09 (0.30)
85,40 (0.12)
8
96.33 (0.10)
97.21 (0.05)
97.34 (0.03)
1
82.73 (0.28)
84.55 (0.17)
86.00 (0.04)
m
96.31 (0.11)
97.20 (0.06)
97.35 (0.06)
We use the CNN architectures: ResNet (He et al., 2016) with 50-layers (Res WideResNet (Zagonuyko & Komodakis 2016) with 28 layers and width 10 (WRN-28-10), and Pyramid Network (Han et al. 2017) with 272 layers and widening factor 200. In all settings, we use the standard data augmentations: horizontal flip, normalization, padding by four pixels, random crop, and cutout (De & Taylor (2017). and we employ the weight decay with the coefficient 0.05. Moreover, we use the multi-step strategy for the step size, which decays the step size by a factor once the number of epochs reaches one of the given milestones. To see the dependence on the step size, we use two decay schedules for the parameter averaging. Table ! summarizes milestones labeled by the symbols: 's', 'm'. and '!' The initial step size and decay factor of the step size are set to 0.1 and 0.2 in all cases. The averages are taken from 300 epochs for
Table 1: Decay schedules for (averaged) SGD.
milestones
(80,160,22) 160,240)
m /
(80,160)
[300)
the schedules 's' and 'I, and from 160 epochs for the schedule 'm'. These hyperparameters were tuned based on the validation sets.
For fair comparison, we (averaged) SGD with 400 epochs and (averaged) SAM with 200 epochs because SAM requires two gradients per iteration, and thus the milestones and starting epoch of taking averages are also halved for (averaged) SAM. We evaluate each method 5 times for ResNet-50 and WRN-28- 10, and 3 times for Pyramid network. The averages of classification accuracies are listed in Table 2 with the standard deviations in brackets. We observe from the table that the parameter averaging for SGD improves the classification accuracies in all cases, especially on CIFAR 100 dataset. Eventually, the averaged SGD achieves comparable or better performance than SAM. Moreover, we also observe improvement by parameter averaging for SAM in most cases, which is consistent with the observations Kaddour et al.) (2022).
Comparing results on CIFAR 100 and CIFAR1O, the large step size is better, and the small step size is rel- atively poor on CIFAR 100 dataset, whereas the small step size generally works on CIFAR10 dataset. If we use the step-size strategy 'I' for CIFAR10. 10, then the improvement becomes small (see Appendix B for this result). We hypothesize that this is because the strong bias with a large step size would be useful for difficult datasets, whereas the weak bias with small step size would be sufficient for simple datasets such that the normal SGD already achieves high accuracies. More- over, we note that the averaged SGD on 100 quite works well with the large step-size schedule '', but SGD itself does not converge and poorly performs under this schedule as seen in Figure 3 The accuracy of SGD temporarily increases at the 300 epochs be- cause of the decay of the step size, it decreases thereafter. However, the average of such parameters achieves significantly high accuracy as expected by theory.
Test Test classification accuracy
SGD(WRN)
SGD(ResNet)
ASGD(ResNet)
ASGD(WRN)
Epochs
Figure 3: Test accuracies achieved by SGD and averaged SGD on CIFARI00 dataset with ResNet-50 and WRN-28-10.
7

---

Under review as a conference paper at ICLR 2023
Train
Tain
SGD
ASGD
10
10.-
0.5
Distance
Distarice
Figure 4: Sections of the train (red) and test (blue) loss landscapes across the parameters obtained by averaged SGD (distance=0) and SGD (distance=1) for ResNet-50 with CIFAR dataset. SGD is run with a small step size after running averaged SGD with a large step size. The middle figure is the close-up view the edge. The triangle and circle markers represent convergent parameters by SGD and averaged SGD, respectively. The right figure plots smoothed train loss functions (green, darker is smoother) with Gaussian noises in addition to train and test losses. The blank circles are the minimizers of smoothed objectives.
Finally, we observe in Figure 4 the loss landscape around the convergent point is in better shape and forms an asymmetric valley. Therefore, we expect that the loss function around the solution uniformly satisfies the required conditions in theory. Specifically, Figure depiets the section of train and test loss functions across parameters obtained by the averaged SGD and SGD. The middle figure is the close-up view the edge and plots each parameter. The right figure depicts the smoothed objectives with Gaussian noises in addition to train and test losses in log-scale. We observe in Figure the phenomenon that SGD converges to an edge and averaged SGD converges a flat side. This phenomenon can be explained by our theory because the minimizer of the smoothed asymptotic valley is shifted flat side as confirmed in a synthetic setting (Figure 2 and deep learning setting (the right of Figure 4). Moreover, the right figure indicates the possibility that the smoothed objective with appropriate stochastic gradient noise well approximates test loss, although we employ artificial noise (Gaussian) to depict graphs for simplicity. Finally, we observe that averaged SGD achieves lower test loss which makes about 2% improvement in the classification error on CIFAR 100 dataset. These observations are also consistent with the experiments conducted in He et al. 2019).
6 RELATED LITERATURE AND DISCUSSION
Flat Minimum. Keskar et al. (017 and Hochreiter & Schmidhuber (1997) showed a flat minimum
generalizes well and sharp minimum generalizes poorly. However, the flatness solely cannot explain generalization because it can be easily manipulated (Dinh et al.) (017). Neyshabur et al. (2017) rigorously proved the sharpness combined with 2o-norm provides generalization bound and Jiang et al. (2020) verified this correlation through large scale experiments. Keskar et al. (2017) also argued that SGD converges to a flat minimum and He et al. (2019) argued the averaged SGD tends to converge to asymmetric valley. Several works (Kleinberg al. 2018 /houu et al., 2020) studied the stochastic gradient noise to theoretically prove the existence of an implicit bias towards a flat region asymmetric valley. Moreover, many works (Izmailov et al.) 2018 Foret et al. 2020 Damian et al. 2021: Orvieto et al.] 2022) studied the techniques to further bring out the bias of SGD. In particular, SAM and SWA achieved significant improvement in the generalization performance. In our paper, we show that parameter averaging stabilizes the convergence to a flat region or asymmetric valley, and suggest the usefulness of the combination with the large step size for the difficult dataset which needs a stronger regularization.
Markoy Chain Interpretation of SGD. Dieuleveut al. (2000) Yu et al. (2020) provided the
Markov chain interpretation of SGD. They showed the marginal distribution of the parameter of SGD converges to an invariant distribution for convex and nonconvex optimization problems, respectively. Moreover, Dieuleveut et al. (2020) showed the mean of the invariant distribution, attained by the averaged SGD, is at distance (n) from the minimizer of the objective function, whereas SGD itself oscillates at distance (v)7) in the convex optimization settings. Izmailov et al. (2018) also attributed the success of SWA to such phenomenon. That is, Izmailov et al. (2018) explained that SGD travels on the hypersphere because of the convergence Gaussian distribution and the concentration on the
8

---

Under review as a conference paper at ICLR 2023
sphere under a simplified setting, and thus averaging scheme allows to go inside of the sphere which may be flat. We can say our contribution is to theoretically justify this intuition by extending the result obtained by Dieuleveut et al. (2020) to nonconvex optimization setting. In the proof, we utilize the alternative view of SGD IKleinberg et al. 2018) in a non-asymptotic way under some conditions not the original objective but the smoothed objective function. Combination with the Markov chain view for nonconvex objective (Yu et al. 2020) may be helpful in more detailed analyses.
Step size and Minibatch. SGD with a large step size often suffers from stochastic gradient noise and becomes unstable. This the reason why we should take smaller step size so that SGD converges. In this sense, the minibatching of stochastic gradients clearly plays the same role as the step size and sometimes brings additional gains. For instance, Smith et al. (2017) empirically demonstrated that the number of parameter updates can be reduced, maintaining the learning curves on both training and test datasets by increasing minibatch size instead of decreasing step size. We remark that our analysis can incorporate the minibatch by dividing of and o in Theorem [ byy the minibatch size, and we can see certain improvements of optimization accuracy as well. Then, both SGD and averaged SGD share the same dependency on the minibatch size and thus controlling step size seems more beneficial for parameter averaging.
Edge of Stability. Recently, Cohen al. (2021) showed the deterministic gradient descent for deep neural networks enters Edge of Stability phase. In the traditional optimization theory, the step size is set to be smaller than I/L L to ensure stable convergence and we also make such restriction. On the other hand, the Edge of Stability phase appears when using a higher step size than 2/ L. In this phase, the training loss behaves non-monotonically and the sharpness finally stabilizes around 2/1,. This can be explained as follows (Lewkowycz et al. 2020); if the sharpness around the current parameter is large compared to the step size, then gradient descent cannot stay in such a region and goes to a flatter region that can accommodate the large step size. There are works (Arora al. 2022, Ahn et al.) 2022) which attempted to rigorously justify Edge of Stability phase. Interestingly. their analyses are based similar intuition to ours, but we consider different regime of step sizes and different factor (stochastic noise or larger step size than 2/ L) brings the implicit bias towards flat regions. We believe establishing unified theory is interesting future research.
Averaged SGD. The averaged SGD (Ruppert 1988 Polyak & Juditsky 1992) is a popular variant of
SGD, which returns the average of parameters obtained by SGD aiming at stabilizing the convergence. Because of the better generalization performance, many works conducted convergence rate analysis in the expected risk minimization setting and derived the asymptotically optimal rates O(1/v (1/VT) and O(1/T) for non-strongly convex and strongly convex problems (Nemirovski et al. 2009 Bach & Moulines 2011 Rakhlin et al. 2012 Lacoste-Julien et al., 2012). However, the schedule of step size is basically designed to optimize the original objective function, and hence the implicit bias coming from the large step size will eventually disappear. When applying a non-diminishing step size schedule, the non-zero optimization error basically remains. What we do in this paper is to characterize it as the implicit bias toward flat region.
CONCLUSION
In this paper, we showed that parameter averaging improves the bias-optimization tradeoff caused by the stochastic gradient noise. Specifically, we proved that averaged SGD optimizes the smoothed objective functions up to -rr)-error, whereas SGD itself optimizes it up to ((vT)--rroo in terms of Euclidean distance from the solution, where 7j is the step size. Therefore, parameter averaging significantly stabilizes the implicit bias toward flat region, and we can expect improved performance for difficult datasets such that the stronger bias induced by a larger step size is helpful. Finally, we observed the consistency of our theory with the experiments image classification tasks. In addition to the above discussion, another interesting research direction is to investigate what type of noise is strongly related to generalization performance.
9

---

Under review as a conference paper at ICLR 2023
REFERENCES
Kwangjun Ahn, Jingzhao Zhang, and Suvrit Sra. Understanding the unstable convergence of gradient descent. In Proceedings of the 39th International Conference on Machine Learning, volume 162, pp. 247-257, 2022.
Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on the edge of stability in deep learning. In Proceedings of the 39th International Conference on Machine Learning, pp. 1024, 2022.
Francis Bach and Eric Moulines. Non-asymptotic analysis of stochastic approximation algorithms for machine learning. In Advances in Neural Information Processing Systems, volume 24, pp. 451-459, 2011.
Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In Proceedings of the 9th International Conference on Learning Representations, 2021.
Alex Damian, Tengyu Ma, and Jason D Lee. Label noise sgd provably prefers flat global minimizers. Advances in Neural Information Processing Systems, 34:27449 224411117761, 2021.
Terrance DeVries and Graham W Taylor. Improved regularization of convolutional neural networks with cutout. arXiv preprint arXiv: porXiv::70000 0100 1700: 2017.
Aymeric Dieuleveut, Alain Durmus, and Francis Bach. Bridging the gap between constant step size stochastic gradient descent and markov chains. Annals of Statistics, 48(3): 1348-13338 06 2020.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In Proceedings of the 34th International Conference on Machine Learning, pp. 1019-1028, 2017.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In Proceedings of the 8th International Conference on Learning Representations, 2020.
Dongyoon Han, Jiwhan Kim, and Junmo Kim. Deep pyramidal residual networks. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 5927-5935, 2017.
Haowei He, Gao Huang, and Yang Yuan. Asymmetric valleys: beyond sharp and flat local minima. In Advances Neural Information Prucessing Systems, volume 32, pp. 2553-2564, 2019.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 770-778, 2016.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation, 9(1):1-42, 1997.
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. In 34th Conference Uncertainty in Artificial Intelligence, pp. 876-885, 2018
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. In Proceedings of the 8th International Conference on Learning Representations, 2020.
Jean Kaddour, Linqing Liu, Ricardo Silva, and Matt Kusner. A fair comparison of two popular flat minima optimizers: Stochastic weight averaging sharpness- aware minimization. arXiv preprint arXiv:2202.00661, 2022.
Nitish Shirish Keskar and Richard Socher. Improving generalization performance by switching from adam sgd. arXiv preprint arXiv: 1712.07688, 2017.
Nitish Shirish Keskar, Jorge Nocedal, Ping Tak Peter Tang, Dheevatsa Mudigere, and Mikhail Smelyanskiy. On large- batch training for deep learning: Generalization gap and sharp minima. In Proceedings of the 5th International Conference Learning Representations, 2017.
10

---

Under review as a conference paper at ICLR 2023
Bobby Kleinberg, Yuanzhi Li, and Yang Yuan. An alternative view: When does sgd escape local minima? In Proceedings of the 35th International Conference on Machine Learning, pp. 2698-
Simon Lacoste-Julien, Mark Schmidt, and Francis Bach. A simpler approach to obtaining an (1/t) convergence rate for the projected stochastic subgradient method. arXiv preprint arXiv: 1212.2000 2012.
Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint urXiv:2003.02216 2020.
Liangchen Luo, Yuanhao Xiong, Yan Liu, and Xu Sun. Adaptive gradient methods with dynamic bound of learning rate. In Proceedings of the 7th International Conference on Learning Represen- tations, 2019.
David A McAllester. Some pac-bayesian theorems. In Proceedings of the eleventh amnual conference on Computational learning theory, pp. 230-234, 1998.
David A McAllester. Pac-bayesian model averaging. In Proceedings of the twelfth annual conference Computational learning theory, pp. 164-170, 1999.
Nicole Mücke, Gergely Neu, and Lorenzo Rosasco. Beating sgd saturation with tail-averaging and 12577, 2019.
minibatching, Advances in Neural Information Processing Systems,
Arkadii S Nemirovski, Anatoli Juditsky, Guanghui Lan, and Alexander Shapiro. Robust stochastic approximation approach to stochastic programming. SIAM Journal on Optimization, 19(4) 1609, 2009.
Yurii Nesterov. Introductory Lectures on Convex Optimization: A Basic Course. Kluwer Academic Publishers. 2004.
Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro. Exploring generaliza- tion in deep learning. Advances neural information processing systems, 30, 2017.
Antonio Orvicto, Anant Raj. Hans Kersting, and Francis Bach. Explicit regularization in over- parametrized models via noise injection. arXiv preprint 23X6::2206.046/3, 2022.
Boris T Polyak and Anatoli B Juditsky. Acceleration of stochastic approximation by averaging. SIAM Journal on Control and Optimization, 30(4):838-855 1992.
Alexander Rakhlin, Ohad Shamir, and Karthik Sridharan. Making gradient descent optimal for strongly convex stochastic optimization. In Proceedings of International Conference on Machine Learning 29, pp. 1571-1578, 2012.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, pp. 400-407, 1951.
David Ruppert. Efficient estimations from a slowly convergent Robbins-Monro process. Technical report, Cornell University Operations Research and Industrial Engineering, 1988.
Samuel L Smith, Picter-Jan Kindermans, Chris Ying, and Quoe Le. Don't decay the learning rate, increase the batch size. arXiv preprint arXiv:1711.00489, 2017.
Ashia C Wilson, Rebecca Roelofs, Mitchell Stem, Nati Srebro, and Benjamin Recht. The marginal value of adaptive gradient methods in machine learning. Advances in neural information processing systems, 30, 2017.
Lu Yu, Krishnakumar Balasubramanian, Stanislav Volgushev, and Murat A Erdogdu. An analysis of constant step size sgd in the non-convex regime: Asymptotic normality and bias. arXiv preprint arXiv:2006.07904 2020.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In British Machine Vision Conference 20/6. British Machine Vision Association, 2016.
11

---

Under review as a conference paper at ICLR 2023
Pan Zhou, Jiashi Feng, Chao Ma, Caiming Xiong, Steven Chu Hong Hoi, et al. Towards theoretically understanding why sgd generalizes better than adam in deep learning. Advances in Neural Information Processing Systems, 13:21285-21296 2020.
12

---

Under review as a conference paper at ICLR 2023
Appendix
A PROOFS
A. IMPLICIT STOCHASTIC GRADIENT DESCENT
Let denote by & Rdd + Rd change of variables from to introduced in Section 3 i.e., ((w)= w f(w).
Lemma A. Under Assumption (A1) and S 2T. the function is injective and invertible, and its inverse defined Im.p is differentialee.
Proof. For w. w' E Rd, we suppose p(w) = ((w'). Then, it holds that
w'l f() Vf(w')) - 7/|lw w'll S 2 w w/
where we used L-Lipschitz continuity of V due to (A1). Therefore, we see w w' and ip is an injection. Moreover, since (u( = I 7/ f(w) (1 nL)I = I. Thus, yp is invertible and yp_1, which is defined on Imi, is differentiable because of the injectivity and the inverse map theorem.
Using ya, we see e'(v) = (v)) for v E Imp. Let (St be a probability space such that e((v) can be represented as measurable map E ll e'(v, z). Note that we use e'(v) and e' ('(,,)) =) depending on the situation. For a function g Id > Rd, we denote by [y(w) Jacobian of g. i.e. (w) = (0g,(w))/w))){
Lemma B. Under Assumption (A1) and (A2), we get for any v E Imip c Rd
(F())= ([()=E((V((rr (o))
Moreover, if
2L' then
-F(c)-B(V((e - ( (opeeenee(m)))) (c)())) S 2n [[V ( (T)(1.)
Proof. The first equality of the statement can be confirmed by the direct calculation as follows:
VF(c) = vEIs(v - (v))
/ v(f(e nH (v, =))))P((z)
(1 ( (t) (r,, 2))))((()
EVVf( (())l)
Next, we evaluate the last term below. By the chain rule and inverse map theory,
(v) = Je(0- (e (o))) [o) (v)) (v)).
Note that from assumption for any w E Rd J. (w) I
2 (1- -nn JT. Hence,
ITTT & (1tto===~~= (v))ll2 202.
13

---

Under review as a conference paper at ICLR 2023
Finally, get
1/ -
s $/ 1 [V (e, (IP
S 2200/ / [Vf(v (e,
s 202 / [| 77
This finishes the proof.
A. PROOF THEOREM
The following proposition is the restatement of the well-known convergence result to a stationary point using the coordinate v.
Proposition A. Under Assumption (A1), (A2), and
2L' get
T
[[f(ww)]]+ + 2 3 Lot +
- -
(5)
3n
1=0
Proof. It is known that (A1) derives the following Nesterov
2004): for any w, w/ E Rd.
I. 'l'' w/
f(w') s f(w) Vf(w)) (w' w) 2
(6)
Substituting the update Eq. into this inequality with w' w141 and w = wr, and taking the
conditional expectation [E|-
we get
B[f(wr+)) < f(we) 7/ rl;Tf( L E
[ fw,)) c (w)))
f(w)) - (1- 7L
112 1
2
37
(w;) || + n²oiL
s f(w))
|| 2| pofL 2- By summing up this inequality,
Thus, we have [[f(w++1) S [[](w]]] we get
T+1 [
Elf [wo(l + L + +
3n
where we used the nonnegativity of f. By dropping the term with t = 0 of the sum in the left hand side and using 1
v, nd (w:), we finally get
T
4 B[f(wo]]] Lof + 2).
Ell Vf -
3n
3
Using the above results, we prove Theorem
I which is restated below.
14

---

Under review as a conference paper at ICLR 2023
Theorem A. Under Assumption (A1). (A2), and (A3), run the stochastic gradient descent with T Titterations with the step size S 2L then the implicit SGD satisfies the following inequality:
1
8
(1 (1+
211
Bl vo11 S
Blll00 111
R B[s(r))))
(T + 1)
Tc(T + +1)
c
2no8
+ 8nj
2no
1+
c
3c
= 0 (T + 2
SnPoRL
2n
+
3c
Proof of Theorem A To evaluate - - for the implicit SGD
2 we first give several bounds
as follows. By Assumption (A3), Young's inequality, and Lemma B
get
-2(wr - vo. " (wve - (v))(
= -2 - u.) V (v:) + 2(v, - v.) (w:) /VV - ( (vt))
S <-2ollug 12 + cl c1|0 1= + 1
VVF (v:) B (00))
S <cll|u v. 11 4n 0
[ 7e (o:))
c
By Assumption (A2) and Young's 's inequality again, we get
Blll(+ (1)) Vf( (
- 2101Il (1x)) s 207 2E V f(v - 7 (wt)))
TEL
Combining the above two inequalities, we get
B[[|e++1 v. 11 = Bl[lvz 7k441 (1x) - S(= - (1441 - 1
11 v2112 - 271(1/ v..) " E(VS
(
+ 1711 (1:) + vf(u, - (
I (1 - --))))) ||u - v. 2| + +21j'0
+ 2n² 1+ 2n0 E V f(u, - npl (.t))1
Taking the expectation regarding all histories and summing up over = 0, 1..... T, we get
Ellly - v. 1111 < 10 - u. 1² Elllrr+1 - -., + 27 + 1)
2012 Σ + 2n² 1
T
E [[] f(v - ( (1:)))1
1=
< Billvo
v-l] - El[lut+1 1=1121+ + 2nj² of 1)
8
2n103
4 3
1+ 2010
+ 3
1
1 +
E[f(wb)] +
of L(T 2).
where we used Proposition A Therefore, we conclude
1
8
(1 I++
2002
T Bllluz v. <
Ellleo v111
B[)(w))] c
CT(T + 1)
30(T +1)
2010 2
8n2
2n02
+
|
o7L
c
3c
c
in
+ 811²
(
2110
11-1)
+
1
LfL.
c
3c
c
15

---

Under review as a conference paper at ICLR 2023
A. PROOF OF THEOREM 2
We give several statements used to prove Theorem 2
Lemma C. Under the same assumptions in Theorem the stochastic gradient descent with T-iterations with the step size S z then the implicit SGD satisfies the following inequality:
T
Vf(v - v v66+2(0))
- 00) + 1 (V++)
Anop
16n
1
270 ))
7
3c
T
E
(P(()) E (V - (o-1
S 2202) + 20 2010222
L 1)(T +
Proof of Lemma
C
By the simple calculation, we get
T
T
f((u - -
E (Vl( - (v(2((1)) + +++(
=
||E [wo
vr+illl
E
[1]e0
*1+1l
E
1100 - 11+1ll?]
1
<
2E
|1v0 v.112 + |00+1 - vo.l?
"7
T+1
2
Σ 1=
E
1ler 11]
$ / 0(1) (T+1)
1
Ano1
161207L
(1+
211o
30
- 1 o0))))) + 1 1 1
[anop +
1612
(1
(1 2012
30
where we used Theorem A
Next, we show the second inequality by using Lemma Blas as follows:
T
E
(VF(w))) - E [V (2.2
E
VF(u.) - E [Vf(v nd (vx))
E
21102
[ - (o.)
< 21102
Vf(v: (
1=0
T E -
<2n02
(T + 1)
1=0
217021
O(T) + LT + 1) +2)
16

---

Under review as a conference paper at ICLR 2023
2 1 + 1)(7 + O 201021 V ([T-=1)(
Proposition B. Under the same assumptions as in Theorem run the stochastic gradient descent with T-iterations with the step size n S 21 then the implicit SGD satisfies the following inequality:
T
E
VF(
v3 102n
T+1
Proof of Proposition
Using Lemma
we get
V(())
-F(t))B[[((CC 7é4 (o:))
T t=0
T+1
E
E VV[( 1
T+1
1=0
E
(F((v) B V
ne' (1:+)(1))) (v)))
7+1
IE
Vf(0i (01))
T+1
12
T
202n3
201
L
1 T++1
+ -
4no7
16
2n12
+
1+
17
T
3c
s O
4
v3
We here prove Theorem 2 which is restated below.
Theorem B. Under Assumption (A1)-(A5), the averaged SGD for T-iterations with the step size S TL then the average UT satisfies the following inequality:
Telrl- s (T- (T
4010212 L
2
81
(1+ 2nj03 21+))
+
+
V3n
cp
3cp
Proof. We define R(v)= VF(V) V2F (u-)) Then, by (A4), see R(v)] M M 1ol1o-v,)1 vu.1?. By taking average of R(U)) over E {0, 1,.,,,,,TT and rearranging terms, we get
T}
T
1
V2F(e,) - v..)
VF(w))
R(ve).
T+1
T+1 t=
t=0
Therefore, we get
1 [|[E[Tr]
- 0 < V22
2 )l
T
T
1
<
E
V
+
E
R(v)
T+1
T 1
t=0
t=0
1
M
<
E
VF())
+ T
E
|/1e v.1l
T+1
1=0
1=0
17

---

Under review as a conference paper at ICLR 2023
The latter and former terms can be bounded by Theorem A and Proposition B Thus, finally get
1 [ - v.ll S so - + 40 1010n18
v3 + 211
8n
(1+ 2n1o
3c
B ADDITIONAL EXPERIMENTS
Table 3: Comparison of test classification accuracies on CIFARIO dataset. All methods adopt the multi-step strategy for the step size schedule.
CIFARIO
7
ResNet-50
WRN-28-10
SGD
95.95 (0.10)
96.85 (0.16)
96.58 (0.14)
97.24 (0.07)
Averaged
96.89 (0.05)
97.44 (0.04)
SGD
96.27 (0.16)
97.05 (0.09)
We run SGD and averaged SGD CIFARIO dataset with the step size strategy 'I' under the same settings as in Section [ Table 3
lists the results including this case. We observe that the large step size 'I' does not work so well on CIFAR10 dataset compared other schedules. We hypothesize this because CIFARIO is not so difficult dataset and does not require stronger bias induced by larger step size.
We also validate the cosine annealing strategy for the step size, which is frequently used due to its excellent performance. We used the symbols 's', 'm', and 'I' for the cosine annealing depending on the last step sizes which are set to O, 0.004, and 0.02, respectively. The parameter averaging for averaged SGD is taken over the last quarter of the training. From the table, observe the usefulness of parameter averaging for cosine annealing schedule as well.
Table 4: Comparison of test classification accuracies on CIFAR100 and CIFARIO datasets. All methods adopt cosine annealing for the step-size schedule.
CIFAR 100
CIFAR10
ResNet-50
WRN-28-10
Pyramid
ResNet-50
WRN-28-10
Pyramid
SGD
82.26
82.68
82,97
96.58
97.00
96.66
Averaged SGD
83.89
84.28
85.14
97.01
97.28
97.07
83.21
84.49
85.47
m
96.86
97.51
97.32
SAM
83.35
84.64
86.24
96.40
96.89
97.61
Averaged SAM
83.18
84.94
86.79
96.56
97.14
97.55
83.58
85.26
86.84
96.51
97.19
97.48
Finally, we run SGD, SGD with large step size, and averaged SGD to train the standard convolutional neural network on Fashion MNIST dataset confirm how efficiently sharpness and classification accuracy can be optimized by each method. We note the large step size used for SGD is the same that for averaged SGD. We plot the trace of Hessian V2 and test loss functions in Figure 5 From this figure, we observe that the averaged SGD converges to flatter region and achieves the highest classification accuracy the test dataset expected in theory.
18

---

Under review as a conference paper at ICLR 2023
1330
SGD
Test dassification securacy
0.50
ASGD
SGD
large step size
0.9
Hessian trace
0.00
6300
0.00
200
2200
1106
303
large
Epochs
Epochs
Figure 5: The figure depicts the curve of the trace of Hessian 2 f(w) and test loss functions achieved by SGD, SGD with large step size, and averaged SGD. Each algorithm is run train the standard convolutional neural network Fashion MNIST dataset.
1.0
-0.02
-0.4
-0.05
0.4
-0.2
W
W
Figure 6: The left figure plots the mollifier (blue) and smoothed mollifier Gs (orange), and the right figure plots the objective f (blue) and smoothed objective J (orange). The constants 0.1,7 r 2.0, and 1.0
C MOTIVATING EXAMPLE
C. PROBLEM SETUP
In this section, we present a motivating example that verifies the convergence to flat minimum and a certain separation between SGD and averaged SGD. We consider a one-dimensional objective function f IR- + R defined below: for p, 8 > O,
f(w) = 2 (w p)2 + gs(),,
where : R + R is a scaled mollifier:
(1- 1-
põ exp
(lwl < 5),
go (5(w)=
(/w 2 5).
go(w) 8gi (w/8) is scaling of the well-known mollifier of which is an infinitely differentiable function with compact support. That is, gs is smooth function whose support [-6, 6). Because of the coefficient of the function f(w) has local minimum in [-8, 6), which can be the global minimum. See Figure 6 (right).
The maximum values of the first and second derivatives of are bounded. Thus, we define constants C1, C2 by
C = max { {1. max lsi (mmii))] O2 = max lo (w))-
19

---

Under review as a conference paper at ICLR 2023
Since (w) (w/6), we the second derivative of gō is bounded by C2pō Hence, Lipschitz smoothness (boundedness of Hessian) L of f is 1 1 C2pō
Next, we consider the uniform noise on the interval [-r, for r > 0, 0.i. e U r] and suppose c(w, 2) = <(z)(=<< c((v, 2)) where &
c(w, 2) is explicit representation of the random noise.
In other words, noise distribution does not change in w. In this case, we see of = E|e"] I r² and 02 0. The smoothed objective F with the noise e' and step-size 1)
F(u) = B n7c')l
- p)² + oo +Egg(( n7E')l
(v- p)2 + Elgs(v -)e))
We consider the following problem setup:
8
(8)
4(1 201)
> 20 (5 + Cap).
(9)
Note that we can choose arbitrarily small 5 0 and large which satisfy the above inequalities.
For appropriate smoothing, we choose the step size 7 so that
2
p 5
min
(10)
4r
r 0
A step-size 17 that satisfies the condition required in the theory.
10
exists and it also satisfies 1, S 1/L 5/(5 + C2p)
C. 2 CONVERGENCE SGD AND AVERAGED SGD
Under the above setup 8
10 we can estimate constants appearing in the convergence results of SGD and averaged SGD as follows (for the detail see the next subsection):
8
L =
of r2, 02
0,
(11)
+ C2p
1,
3 M
8
(12)
9p
Moreover, the minimum of the smoothed objective v. p. sharp minimum (~ 0) can be eliminated by smoothing.
Therefore, for SGD we obtain by Theorem
2no?
8n
2no2
Ellle u. (O- +
3c
81
2 6yr²
(1 6nr²)
8
C2p
We from this inequality, 7/- 2C,8 is the best choice of the step-size, resulting in
T
12Ciór + 32C753
T+1 M E El - v. s o (T- +
8 CCp (1 + 12C)
where apply Jensen's inequality to derive the bound L, norm. This result means SGD avoids sharp minimum (i.e., v ~ 0 under small á >> 0) and converges to a flat minimum v. = p. and a too large noise will affect the convergence to v, based our step-size policy.
Moreover, for averaged SGD we obtain by Theorem 2
40102n8L1
+ 2 +
8nPo9LM
(1+ 2no3
E [rl vall 0 ( (T +
V3p
cp
3cp
=
16
3mp 2
+
(1 + 6mr?))
9p
8 + Czp
20

---

Under review as a conference paper at ICLR 2023
Small Delta
Large Delta
SGD
SGD
10-
Averaged SGD
10
Averaged SGD
06
0.6
04
00
00-
-1.0
-0.5
00
0.5
1.0
1.5
20
-1.0
-0.5
0.5
1.0
15
20
Figure The figures plot the convergent points of SGD and averaged SGD for problems with 0.1 and 8 = 0.5.
Hence, for n/+ = 2016 we obtain
u S (0(TT
32 9p
80383 3C +
8 Cap (1 12C;or))
+
This bound means averaged SGD will get closer to v. as long as SGD approaches neighborhood of v.
According to the above results, both SGD and averaged SGD converge to flat region when is small, and averaged SGD converges even when ō is relatively large.
We empirically observed this phenomenon in Figure 7 in which we run SGD and averaged SGD for problems with small 0.1 and relatively large 0.5.
C. 3 ESTIMATION OF CONSTANTS
We verify the estimations of constants in M remain.
L, of, and 02 are already obtained, thus, ma, c, and
Minimum and estimation of μ. We first see that under our problem setting, the local minimum around the origin eliminated and p the optimal solution of F. i.c., U. = p.
The smoothed function Gz(z) def Elgs(v nsE')] and its derivative (v(v)) are calculated as follows:
Gs(c)= = 1 go(v motnntt)___ m 2r dt.
Gz(()) [ g'(l -1)
By taking into account supp(gs) = 1-8, 5]. the smoothed objective Gs(v) is constant on { I np [ 2 +d), and thus, G¹ is non-zero only on supp(G"s) [-nrr Unr+ npr +ō]. See Figure 6 (left). Since ypr + 5 p/4 < p under (10 p is still a local minimum of F.
We evaluate the bound on G/2 on (U)) below. for u E [19 mp + 8) the support of ge - nt)
in t E is [vv 8)/n, (v + 6)/m), get
0 < (v)
[" d6(l (1()
1/1)
_ dl
1
2r
1
lg'(1 -ntt] dt
2r
vis 1
dt pC18
$ PC1
1-8
2r
rp
21

---

Under review as a conference paper at ICLR 2023
where we used (x(o) (0/0/)) S pC1. A bound on [-np 6, -1p + 8] is also obtained in the same way. Thus, we see
pC1 - (v) 0
(e [- - 6, 1pr + 6)),
0 0 ()))) I pC16
(o [nr 6, mpr + 6)),
(,())= =0
(else).
there are additional stationary points of F, they should exist (nr- nr = supp(G?)) 1-mr- 8, -pp- + 6] because of the sign of G' and supp(G's) C p/4), However, since 1/r + 0 p/4
and pC,6 p/2 under 10 see 7
pC18 T/r
max (,
(v) S (nr + -
4
Hence, C, = p is the unique local minimum (i.e., optimal solution) F and we can conclude μ =
Estimation of c. From the above argument, we get
F((((V p) = (v p)² + (b)((( p)
(v p)2
(V [- 8, -pr + 8]),
(u p)2
pCj ar
(v p) (v p)2 + (v p)
(v pyp 6, mpr 5)) (else).
(u p)²
Clearly, p/2 < 2(p - v)/3 for S 1 + z p/4. Thus, F'(v)(v p) 2 (v - 2)2 on
E [np--1, 6 + 6) and conclude 1/3. 6,
Estimation of M. Noting v = p and F" F"(p) = 1, we have
(v) - F" (v.) - v.)l [( - p) + G3((v) - (v - (v-
Because of the problem setup, it is enough to verify M 9p
8 satisfies (v) S M - pl on
E [pr 6, ypr + 6). Since (x) p/2 and p/4 for in this interval, have
9p²
(x) s P < M
S M(r - 2)
2
16
This concludes M 8
9p
22


Reviews:
[{"review_id": "VsE8OYFLwO", "reviewer": "Reviewer_Q2SC", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper aims to understand how averaging the parameters of SGD (averaged SGD) leads to a more stable convergence to flat regions. \nThe authors take the alternative view of SGD from Kleinberg et al. (2018), which tracks $v = w - \\eta \\nabla f(v)$ instead of $w$ itself. Then SGD wrt $w$ can be viewed as a stochastic semi-gradient method wrt $v$ on a smoothed loss, where the smoothening factor depends on eta. Under one-point convexity at the minimizer of the smoothed loss (and also various regularity conditions), they show that both SGD and averaged SGD oscillate around the minimizer, but the distance to the minimizer from the final parameter of averaged SGD is smaller than that from SGD ($O(\\eta)$ v.s. $O(\\sqrt{\\eta})$).", "summary_of_the_review": "I would like to hear the authors' comments on my major concern because I could be wrong. If $v_*$ can indeed change a lot for different learning rates, I would be rather happy to raise my score. Otherwise, I have to vote for rejection, but I encourage the authors to come up with another argument to motivate their convergence analysis and resubmit.\n\n=================\n\nAfter reading the rebuttal, I found the 1-dim example given by the author helpful. This simple example shows that the theory indeed predicts the averaged SGD converges to a flatter solution than SGD, so I have raised my score from 3 to 5.\n\nOverall, I'm not against accepting this paper because the convergence analysis still looks good to me. But I feel that the paper is at the borderline due to the following concerns:\n1. As I pointed out in the initial review, this paper compares the upper bounds for SGD and averaged SGD (even in the 1-dim example, only the upper bounds are compared). A more solid theoretical result should give a lower bound for SGD to ensure that it doesn't converge to flat minimizes.\n2. Though I appreciate the 1-dim example for its simplicity, it is unclear how general it is. It would be better if the authors could have generalized the example to a broad function class. Alternatively, the authors could also make up the generality issue in theory by measuring the empirical loss landscape to see if it is similar to the 1-dim example in some sense. For this reason, I believe that this paper still has much room for improvement.", "strength_and_weaknesses": "### Strength\n\n1. The phenomenon that averaged SGD leads to a more stable convergence to flat regions is interesting to study, and may deepen the current understanding of the training dynamics in the late training phase.\n2. Clear convergence results for both SGD and averaged SGD.\n3. The theoretical result indeed provides a clear separation between SGD and averaged SGD in terms of the distance upper bounds to the minimizer of the smoothed loss.\n\n\n### Weakness\n\n**Major Concern.** The major concern I have is whether the current result indeed shows that larger LR leads to a flatter minimizer. If not, then people may want to use SGD with a smaller LR so that the distance to the minimizer is smaller, and then SGD may beat averaged SGD. Basically, the authors should check whether the minimizer $v_*(\\eta)$ as a function of $\\eta$ can move a lot when $\\eta$ goes from a big value $\\eta$ to a small value $\\eta'$.\n\nHowever, the distance between $v_*(\\eta)$ and $v_*(\\eta')$ should be bounded by $O(\\eta^2)$. The reason is that we can apply the implicit function theorem for solving the equation $\\nabla F(v) = 0$, which gives us a curve $v_*(\\eta)$ wrt $\\eta$. Note that the derivative wrt $\\eta$ is $-(\\nabla^2 F(v_*))^{-1} \\left(\\frac{\\eta^2}{2} \\nabla \\mathrm{Tr}(\\nabla^2 f(v_*) \\mathbb{E}[\\epsilon' \\epsilon'^\\top]) + O(\\eta^3)\\right)$. According to A4 and A5, this should be of order $O(\\eta^2)$.\n\nNote that the theorem for SGD shows a distance bound $O(\\eta)$. So for SGD with a much smaller LR $\\eta'$, the distance from the final parameter to the minimizer of the smoothed loss defined with $\\eta$ would be $O(\\eta^2) + O(\\sqrt{\\eta'}) = O(\\eta^2)$. But for averaged SGD with LR $\\eta$, it is $O(\\eta)$, which is actually much higher.\n\nThe above argument would suggest that averaged SGD with large LR does not encourage flatness at all.\n\n**Minor concerns.**\n1. I understand that one-point convexity makes the analysis possible, but the authors should give some concrete examples of one-point convex functions, and discuss how general this notion is. It should also be stated clearly whether one-point convexity holds in the experiments or not.\n2. Comparing the distance upper bounds for SGD and averaged SGD is reasonable, but this does not mean SGD cannot get much closer to the minimizer because the bound may not be tight. I would recommend the authors prove a lower bound in the SGD case to make the separation more concrete.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "UZaD9T8lO3", "reviewer": "Reviewer_fQ6K", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The manuscript studies the optimization and generalization trade-off caused by stochastic noise, step size, implicit bias, and sharpness of the minimum; and more importantly, why an averaged SGD with a large step size can improve the trade-offs and converge to a flat region more stably than SGD.", "summary_of_the_review": "In general, the manuscript is well-written, and the main concerns of the reviewer are\n* unfair comparison of theoretical/numerical results\n* unsupported experimental setups.  ", "strength_and_weaknesses": "## Strength\n* The paper is well-structured and well-written (in high quality), with a clear related work section. The reviewer was enjoying reading most part of the manuscript.\n* The studied question is of great importance to the optimization community.\n* Extensive numerical results are included in the manuscript in terms of various neural architectures and learning rate schedules; besides, some visualization results are very insightful. \n\n## Weaknesses\n1. The manuscript argues that SGD using the step size $\\eta$ converges to a distance $O(\\sqrt{\\eta})$ from the solution, whereas the averaged SGD using the same step size converges to a distance to $O(\\eta)$. First of all, the reviewer is unclear on how to conclude the $O(\\sqrt{\\eta})$ and $O(\\eta)$ from Theorem 1 and Theorem 2; and authors are required to elaborate on them. More importantly, compared to Theorem 1, Theorem 2 relies on the additional Assumption 3, and thus could derive a tighter rate; however, it does not make much sense to compare two results under different assumptions, and even highlight it as a key contribution.\n2. If the reviewer understands the manuscript correctly, the current convergence theory in the submission cannot capture the training dynamic caused by using a larger step size as the averaged SGD only runs normal SGD and takes the average in the end. Understanding this would be a key step in interpreting the implicit bias in SGD, however, the current draft cannot fill this gap.\n3. The choice of hyper-parameters used in the manuscript seems a bit random and no justification can be found. For example, weight decay with a coefficient of 0.05 was used but the reviewer believes many SOTA training recipes use 1e-4 or 5e-4; Similar comment can be applied to the choice of decay factor as well as the milestones for three cases.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "O5QBnNXpNk", "reviewer": "Reviewer_vTYE", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper presents the convergence analysis for SGD and averaged SGD with fixed step size, showing that iterates converge to the vicinity of the optimizer of a smoothed loss function. Compared to prior work, the author improved existing results with a less strict assumption on the step size, and provide a convergence result for the averaged SGD, showing a better rate over vanilla SGD.", "summary_of_the_review": "Good paper studies a very relevant topic in the theory of machine learning.", "strength_and_weaknesses": "Strength: This paper is well written, and it is easy to follow. The theoretical results are novel and the discussions/remarks are very helpful. While the results seem incremental over some existing work on the implicit bias of SGD, there is still a significant contribution, showing an improved rate using averaged SGD.\n\nWeakness: The author mentioned before Theorem 1, \"different proof techniques\" than (Kleinber et al.), maybe you can elaborate?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "JRpTFYAylgv", "reviewer": "Reviewer_EP51", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper aims to explain why stochastic weight averaging (SWA, Izmailov et al., 2018) helps generalization. The paper adopts the framework of Kleinberg et al., 2018 and shows that SGD implicitly minimizes a regularized objective, where the regularizer depends on both Hessian of the loss and the covariance matrix of the stochastic gradient, under the assumption that the regularized loss is one-point convex. The authors show that the regularization strength is proportional to $O(\\eta^2)$ where $\\eta$ is the learning rate and SGD can reach a solution which is $O(\\sqrt{\\eta})$-close to the minimizer of the regularizers loss. The authors further explain the benefit of stochastic weight averaging by showing that the averaged iterate better optimizes the regularized loss, that is, SWA can get $O(\\eta)$-close to the minimizer. ", "summary_of_the_review": "This paper gives a novel understanding of the implicit bias of SGD, which explains the efficacy of stochastic weight averaging. However, I am not sure how this understanding is related to practice as the optimization guarantee presented in the main theorems is quite weak, not enough to separate SGD and SWA from ERM even. Thus now I tend to reject this paper, but I am willing to increase the score if the authors can address my concerns satisfactorily. \n(A good example is Damian et al.,2021, where they showed that the minimizer of the original loss doesn't meet their guarantee for label noise SGD on the regularized loss)\n\n- Damian, Alex, Tengyu Ma, and Jason D. Lee. \"Label noise sgd provably prefers flat global minimizers.\" Advances in Neural Information Processing Systems 34 (2021): 27449-27461.\n\n================================================\n\nUpdate after rebuttal: I've read the revision and rebuttal by the authors as well as reviews by other reviewers. My concern that \"whether there exists a suitable learning rate, such that the optimization-bias tradeoff gives non-trivial improvement compared to the empirical minimizer of the original training loss\" is not well-addressed by the authors and I decide to keep my score. \n\nThe new example in appendix C given by the authors in the revision shows when $\\delta$ is sufficiently small, SGD and average SGD will converge to the flat local minimizer (which is also the global minimizer of original loss $f$). This is just an optimization result because any algorithm that finds global minimizer will also find the flat minimizer in this case. To make it a result of the implicit bias of SGD, the author should demonstrate there is a case where the sharp local minimizer is indeed a global minimizer but still gradient noise biases SGD to the flatter area, even if the loss there could be slightly higher.  ", "strength_and_weaknesses": "##Pros:\n\n1. The paper is well-written and easy to understand. The proofs of the main theorems (Theorem 1 and 2) seem correct to me.\n\n2. The framework (Kleinberg et al., 2018) is originally developed for understanding optimization under assumptions weaker than convexity, that is, one-point convexity. The adoption of their framework towards understanding the implicit bias of SGD and SWA is quite novel, as well as the bias-optimization tradeoff view. \n\n3. The authors successfully explain the efficacy of SWA using the bias-optimization tradeoff, that is, though SWA implicitly penalizes the same objective as SGD does, SWA optimizes the objective better than SGD.\n\n\n\n##Cons:\n\n1. Though the theory proposed in this paper looks self-consistent, the authors don't really justify whether their theory is relevant to practice, like either directly testing their proposed regularizer using other optimizers so there is no optimization issue, or developing algorithms with better generalization using their insight. Technically speaking, it is not clear to me whether there exists a suitable learning rate, such that the optimization-bias tradeoff gives non-trivial improvement compared to the empirical minimizer of the original training loss.\n\n\tTo be more specific, under Assumption 2, it is fairly easy to see any stationary point $w^*$ of the original loss $f$ satisfies that $\\eta^2 \\nabla_{w^*} {Tr}(\\nabla^2 f(w^*) \\Sigma(w^*)) \\ge c\\|v^*-w^* \\|$ , where $\\Sigma(w^*)$ is the covariance matrix of noise in Eq. (4). By Assumption (3), $\\nabla^3 f(w^*)$ is of constant magnitude, which implies that $\\|w^*-v^* \\| = O(\\eta^2)$.  Immediately we see the optimization guarantee given by both Theorem 1 and 2 are vacuous in terms of the implicit bias. This is because we know even getting $O(\\eta^2)$-close (like $w^*$) to the minimizer of the regularized loss $v^*$ is not enough ($w^*$ can be understood by the solution found SGD with tiny LR)\n\t\n\tThat being said, I don't think the view presented in this paper is wrong. It is still valuable. It could be analysis are not tight or assumptions are too strong. But I do think the authors should address the above issue before publishing this work, which can be viewed as a sanity check to the theory.\n\t\n2. Regarding the relevance to the deep learning practice, another thing I think the authors should reconcile is why in practice typically the learning rate of SGD needs to be annealed and the generalization is still good. \n\n3. The details of the setting for experiments are missing, e.g. Figure 4. Why does ASGD have higher training loss than SGD in Figure 4? Does the same thing happen for Figure 3 (training loss is not reported there) Intuitively by averaging you should get a smaller training loss, e.g., think about SGD on a quadratic loss with isotropic gaussian noise.\n\n**Typos**: \n\n1.page 17. \"here we prove Theorem 1 which ..\". should be Theorem 2.  \n2.page 17. All $\\nabla R$ should be $R$.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject