Original: P ARAMETER A VERAGING FOR SGD S TABILIZES THE I MPLICIT B IAS TOWARDS F LAT R EGIONS
Revised: P ARAMETER A VERAGING FOR SGD S TABILIZES THE I MPLICIT B IAS TOWARDS F LAT R EGIONS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Stochastic gradient descent is a workhorse for training deep neural networks due to its excellent generalization performance.
Revised: Stochastic gradient descent is a workhorse for training deep neural networks due to its excellent generalization performance.

Original: 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.
Revised: 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.

Original: Recently, Izmailov et al.
Revised: Recently, Izmailov et al.

Original: (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.
Revised: (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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: In experiments, we verify our theory and  this learning scheme significantly improves performance.
Revised: In experiments, we verify our theory and show this learning scheme significantly improves performance.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Stochastic gradient descent (SGD) (Robbins & Monro, 1951) is a powerful learning method for training modern deep neural networks.
Revised: Stochastic gradient descent (SGD) (Robbins & Monro, 1951) is a powerful learning method for training modern deep neural networks.

Original: In order to further improve the performance, a great deal of SGD variants such as adaptive gradient methods has been developed.
Revised: In order to further improve the performance, a great deal of SGD variants such as adaptive gradient methods has been developed.

Original: 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 et al., 2017; Luo et al., 2019).
Revised: 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 et al., 2017; Luo et al., 2019).

Original: Therefore, the study of the implicit bias of SGD, explaining why it works so better, is nowadays an active research subject.
Revised: Therefore, the study of the implicit bias of SGD, explaining why it works so better, is nowadays an active research subject.

Original: 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 to have a bias towards a flat minimum.
Revised: 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 to have a bias towards a flat minimum.

Original: Hochreiter & Schmidhuber (1997);
Revised: Hochreiter & Schmidhuber (1997);

Original: Keskar et al.
Revised: Keskar et al.

Original: (2017) suggested the correlation between flatness (sharpness) and generalization, that is, flat minima generalizes well compared to sharp minima, and Neyshabur et al.
Revised: (2017) suggested the correlation between flatness (sharpness) and generalization, that is, flat minima generalizes well compared to sharp minima, and Neyshabur et al.

Original: (2017) rigorously supported this correlation under ℓ 2 -regularization by using the PAC-Bayesian framework (McAllester, 1998; 1999).
Revised: (2017) rigorously supported this correlation under ℓ 2 -regularization by using the PAC-Bayesian framework (McAllester, 1998; 1999).

Original: Furthermore, by the large scale experiments, Jiang et al.
Revised: Furthermore, by the large scale experiments, Jiang et al.

Original: (2020) verified that the flatness measures reliably capture the generalization performance and are the most relevant among 40 complexity measures.
Revised: (2020) verified that the flatness measures reliably capture the generalization performance and are the most relevant among 40 complexity measures.

Original: In parallel, Keskar et al.
Revised: In parallel, Keskar et al.

Original: (2017) empirically demonstrated that SGD prefers a flat minimum due to its own stochastic gradient noise and subsequent studies (Kleinberg et al., 2018; Zhou et al., 2020) proved this implicit bias based on the smoothing effect due to the noise and stochastic differential equation, respectively.
Revised: (2017) empirically demonstrated that SGD prefers a flat minimum due to its own stochastic gradient noise and subsequent studies (Kleinberg et al., 2018; Zhou et al., 2020) proved this implicit bias based on the smoothing effect due to the noise and stochastic differential equation, respectively.

Original: Along this line of research, there are endeavors to enhance the bias aiming to improve performance.
Revised: Along this line of research, there are endeavors to enhance the bias aiming to improve performance.

Original: Especially, stochastic weight averaging (SWA) (Izmailov et al., 2018) and sharpness aware minimization (SAM) (Foret et al., 2020) achieved significant improvement in generalization performance over
Revised: Especially, stochastic weight averaging (SWA) (Izmailov et al., 2018) and sharpness aware minimization (SAM) (Foret et al., 2020) achieved significant improvement in generalization performance over

Original: SGD.
Revised: SGD.

Original: SWA is a cyclic averaging scheme for SGD, which includes the averaged SGD (Ruppert, 1988;
Revised: SWA is a cyclic averaging scheme for SGD, which includes the averaged SGD (Ruppert, 1988;

Original: Polyak & Juditsky, 1992) as a special case.
Revised: Polyak & Juditsky, 1992) as a special case.

Original: Averaged SGD with an appropriately small step size or diminishing step size to zero is well known to be an efficient method that achieves statistically optimal convergence rates for the convex optimization problems (Bach & Moulines, 2011; Lacoste-Julien et al., 2012; Rakhlin et al., 2012).
Revised: Averaged SGD with an appropriately small step size or diminishing step size to zero is well known to be an efficient method that achieves statistically optimal convergence rates for the convex optimization problems (Bach & Moulines, 2011; Lacoste-Julien et al., 2012; Rakhlin et al., 2012).

Original: However, such a small step size strategy does not seem useful for
Revised: However, such a small step size strategy does not seem useful for

Original: [Figure]
Revised: [Figure]

Original: training deep neural networks, and Izmailov et al.
Revised: training deep neural networks, and Izmailov et al.

Original: (2018) found the averaged SGD with not small but large step size works quite well.
Revised: (2018) found the averaged SGD with not small but large step size works quite well.

Original: The success of using a large step size can be attributed to the strong implicit bias as discussed in
Revised: The success of using a large step size can be attributed to the strong implicit bias as discussed in

Original: Izmailov et al.
Revised: Izmailov et al.

Original: 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.
Revised: 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.

Original: After a long run, SGD will finally oscillate according to an invariant distribution covering a flat region.
Revised: After a long run, SGD will finally oscillate according to an invariant distribution covering a flat region.

Original: Then, by taking the average, we can get the mean of this distribution, which is located inside a flat region.
Revised: Then, by taking the average, we can get the mean of this distribution, which is located inside a flat region.

Original: Although this provides a good insight into how the averaged SGD with a large step size behaves, the theoretical understanding remains elusive.
Revised: Although this provides a good insight into how the averaged SGD with a large step size behaves, the theoretical understanding remains elusive.

Original: Hence, the research problem we aim to address is
Revised: Hence, the research problem we aim to address is

Original: Why does the averaged SGD with a large step size converge to a flat region more stably than SGD?
Revised: Why does the averaged SGD with a large step size converge to a flat region more stably than SGD?

Original: In our work, we address this question via the convergence analysis of both SGD and averaged SGD.
Revised: In our work, we address this question via the convergence analysis of both SGD and averaged SGD.

Original: 1. C ONTRIBUTIONS
Revised: 1. C ONTRIBUTIONS

Original: We first explain the idea behind our study.
Revised: We first explain the idea behind our study.

Original: 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 1).
Revised: 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 1).

Original: Since as 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 to a flatter region.
Revised: Since as 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 to a flatter region.

Original: At the same time, the step size 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.
Revised: At the same time, the step size 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.

Original: These observations indicate the bias-optimization tradeoff coming from the stochastic gradient noise and controlled by the step size:
Revised: These observations indicate the bias-optimization tradeoff coming from the stochastic gradient noise and controlled by the step size:

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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 η converges to a distance O ( √ η ) from the solution (Theorem 1), whereas the averaged SGD using the same step size converges to a distance O ( η ) (Theorem 2).
Revised: 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 η converges to a distance O ( √ η ) from the solution (Theorem 1), whereas the averaged SGD using the same step size converges to a distance O ( η ) (Theorem 2).

Original: 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.
Revised: 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.

Original: Clearly, a larger step size regardless of the condition of the objective will diverge, thus it should be appropriately small to achieve sufficient optimization.
Revised: Clearly, a larger step size regardless of the condition of the objective will diverge, thus it should be appropriately small to achieve sufficient optimization.

Original: The better dependence of O ( η ) than O ( √ η ) means that the averaged SGD can work well with a wider range of step sizes than SGD.
Revised: The better dependence of O ( η ) than O ( √ η ) means that the averaged SGD can work well with a wider range of step sizes than SGD.

Revised: Although, a too small step-size does not always bias the solution because the deviation of the solution is O ( η 2 ) , 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.

Original: As a result, we can expect the improvement by the averaged SGD fordifficult datasets such that the stronger implicit bias with the  larger step size is useful.
Revised: 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.

Original: The separation between SGD and averaged SGD regarding the bias can occur even in the simple setup as seen in Figure 1 which depicts obtained parameters by running SGD and averaged SGD 500 times in two cases.
Revised: The separation between SGD and averaged SGD regarding the bias can occur even in the simple setup as seen in Figure 1 which depicts obtained parameters by running SGD and averaged SGD 500 times in two cases.

Original: 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
Revised: 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

Original: SGD with large step size can converge stably to a near biased solution which minimizes the smoothed objective.
Revised: SGD with large step size can converge stably to a near biased solution which minimizes the smoothed objective.

Original: The behavior of the averaged SGD in an asymmetric valley (the top of
Revised: The behavior of the averaged SGD in an asymmetric valley (the top of

Original: Figure 1), 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.
Revised: Figure 1), 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.

Original: ; He et al.
Revised: ; He et al.

Original: We note that this phenomenon is certainly captured by our theory.
Revised: We note that this phenomenon is certainly captured by our theory.

Original: Indeed, Figure 2 shows the convergent point of the averaged SGD is almost the minimizer of smoothed objective for each step size.
Revised: Indeed, Figure 2 shows the convergent point of the averaged SGD is almost the minimizer of smoothed objective for each step size.

Original: Our findings are summarized below:
Revised: Our findings are summarized below:

Original: • SGD and averaged SGD implicitly optimize the smoothed objective, whose strength depends on the step size, up to O ( √ η ) and O ( η ) errors in Euclidean distance from the solution.
Revised: • SGD and averaged SGD implicitly optimize the smoothed objective, whose strength depends on the step size, up to O ( √ η ) and O ( η ) errors in Euclidean distance from the solution.

Original: This explains why these methods reach a flat region, since smoothing eliminates sharp minima.
Revised: This explains why these methods reach a flat region with an approprie step-size, since smoothing eliminates sharp minima.

Original: • 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 a flat region.
Revised: • 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 a flat region.

Original: In other words, averaged SGD better controls the bias-optimization tradeoff than SGD.
Revised: In other words, averaged SGD better controls the bias-optimization tradeoff than SGD.

Original: • Hence, the parameter averaging yields an improvement for difficult datasets such that the stronger implicit bias with the larger step size is useful.
Revised: • Hence, the parameter averaging yields an improvement for difficult datasets such that the stronger implicit bias with the larger step size is useful.

Original: This suggests the use of larger step size for such datasets so that averaged SGD stably converges but SGD itself is unstable to effectively bring out the implicit bias.
Revised: This suggests the use of larger step size for such datasets so that averaged SGD stably converges but SGD itself is unstable to effectively bring out the implicit bias.

Original: Technical difference from Kleinberg et al.
Revised: Technical difference from Kleinberg et al.

Original: The proof idea of Proposition 1 relies on the alternative view of SGD (Kleinberg et al., 2018) which shows the existence of an associated SGD for the smoothed objective.
Revised: The proof idea of Proposition 1 relies on the alternative view of SGD (Kleinberg et al., 2018) which shows the existence of an associated SGD for the smoothed objective.

Original: However, since its stochastic gradient is a biased estimator, they showed the convergence not to the solution but to a point at which a sort of one-point strong convexity holds, and avoid the treatment of a biased estimator.
Revised: However, since its stochastic gradient is a biased estimator, they showed the convergence not to the solution but to a point at which a sort of one-point strong convexity holds, and avoid the treatment of a biased estimator.

Original: Hence, the optimization of the smoothed objective is not guaranteed in their theory.
Revised: Hence, the optimization of the smoothed objective is not guaranteed in their theory.

Original: On the other hand, optimization accuracy is the key in our theory, thus we need nontrivial refinement of the proof under a normal one-point strong convexity at the solution.
Revised: On the other hand, optimization accuracy is the key in our theory, thus we need nontrivial refinement of the proof under a normal one-point strong convexity at the solution.

Original:  P RELIMINARY – S TOCHASTIC G RADIENT D ESCENT
Revised:  P RELIMINARY – S TOCHASTIC G RADIENT D ESCENT

Original: 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.
Revised: 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.

Original: Let f : R d → R be a smooth nonconvex objective function to be minimized.
Revised: Let f : R d → R be a smooth nonconvex objective function to be minimized.

Original: For simplicity, we assume f is nonnegative.
Revised: For simplicity, we assume f is nonnegative.

Original: A stochastic gradient descent, randomly initialized at w 0 , for optimizing f is described as follows: for t = 0 , 1 , 2 , . . .
Revised: A stochastic gradient descent, randomly initialized at w 0 , for optimizing f is described as follows: for t = 0 , 1 , 2 , . . .

Original: [Equation]
Revised: [Equation]

Original: where η > 0 is the step-size and ϵ t +1 :
Revised: where η > 0 is the step-size and ϵ t +1 :

Original: R d → R d is a random field corresponding to the stochastic gradient noise i.e., for any w ∈ R d , { ϵ t +1 ( w ) } ∞ t =0 is a sequence of zero-mean random variables taking values in R d .
Revised: R d → R d is a random field corresponding to the stochastic gradient noise i.e., for any w ∈ R d , { ϵ t +1 ( w ) } ∞ t =0 is a sequence of zero-mean random variables taking values in R d .

Original: A typical setup of the above is an empirical/expected risk minimization in machine learning.
Revised: A typical setup of the above is an empirical/expected risk minimization in machine learning.

Original: Example 1 (Risk Minimization) .
Revised: Example 1 (Risk Minimization) .

Original: Let ℓ ( w, z ) be a loss function consisting of the hypothesis function parameterized by w ∈ R d and the data z ∈ R p .
Revised: Let ℓ ( w, z ) be a loss function consisting of the hypothesis function parameterized by w ∈ R d and the data z ∈ R p .

Original: Let µ be an empirical/true data distribution over the data space and Z be a random variable following µ .
Revised: Let µ be an empirical/true data distribution over the data space and Z be a random variable following µ .

Original: Then, the objective function is defined byf ( w ) =
Revised: Then, the objective function is defined by

Original: E Z ∼ µ [ ℓ ( w, Z )] .
Revised: [Equation]

Original: Given i.i.d.
Revised: Given i.i.d.

Original: random variables { Z t +1 } ∞ t =0 with the same distribution as Z , the standard stochastic gradient at t -th iterate w t is defined as ∇ w ℓ ( w t , Z t +1 ) .
Revised: random variables { Z t +1 } ∞ t =0 with the same distribution as Z , the standard stochastic gradient at t -th iterate w t is defined as ∇ w ℓ ( w t , Z t +1 ) .

Original: In this setting, the stochastic noise ϵ t +1 can be ϵ t +1 ( w ) =
Revised: In this setting, the stochastic noise ϵ t +1 can be ϵ t +1 ( w ) =

Original: ∇ w ℓ ( w, Z t +1 )
Revised: ∇ w ℓ ( w, Z t +1 )

Original: − ∇ f ( w ) .
Revised: − ∇ f ( w ) .

Original: Note that we can further include the ℓ 2 -regularization in the objective f and the perturbation by the data augmentation in the distribution µ .
Revised: Note that we can further include the ℓ 2 -regularization in the objective f and the perturbation by the data augmentation in the distribution µ .

Original: As this example satisfies, we suppose { ϵ t +1 } ∞ t =0 are independent copies each other.
Revised: As this example satisfies, we suppose { ϵ t +1 } ∞ t =0 are independent copies each other.

Original: That is, there is a measurable map from a probability space:
Revised: That is, there is a measurable map from a probability space:

Original: Ω ∋
Revised: Ω ∋

Original: z
Revised: z

Original: (cid:55)→ ϵ ( w, z ) ∈ R d , and then ϵ t +1 can be written as a measurable map from a product probability space:
Revised: (cid:55)→ ϵ ( w, z ) ∈ R d , and then ϵ t +1 can be written as a measurable map from a product probability space:

Original: Ω Z ≥ 0 ∋ { z s +1 } ∞ s =

Revised: Ω Z ≥ 0 ∋ { z s +1 } ∞ s =

Original: (cid:55)→ ϵ ( w, z t +1 ) ∈ R d when explicitly representing them as measurable maps.
Revised: (cid:55)→ ϵ ( w, z t +1 ) ∈ R d when explicitly representing them as measurable maps.

Original: Moreover, we make the following assumptions on the objective function and stochastic gradient noise.
Revised: Moreover, we make the following assumptions on the objective function and stochastic gradient noise.

Original: Assumption 1.
Revised: Assumption 1.

Original: (A1) f : R d → R is nonnegative, twice continuously differentiable, and its Hessian is bounded, 
Revised: (A1) f : R d → R is nonnegative, twice continuously differentiable, and its Hessian is bounded, i.e.,

Original: i e., there is a constant L > 0 such that for any w ∈ R d , − LI ⪯ ∇ 2 f ( w ) ⪯ LI .
Revised: there is a constant L > 0 such that for any w ∈ R d , − LI ⪯ ∇ 2 f ( w ) ⪯ LI .

Original: (A2) Random fields { ϵ t +1 } ∞ t =0 are independent copies each other and each ϵ t +1 ( w ) is differentiable in w .
Revised: (A2) Random fields { ϵ t +1 } ∞ t =0 are independent copies each other and each ϵ t +1 ( w ) is differentiable in w .

Original: Moreover, for any w ∈ R d E [ ϵ t +1 ( w )]
Revised: Moreover, for any w ∈ R d E [ ϵ t +1 ( w )]

Original: = 

Revised: =

Original: and there are σ 1 , σ 2 > 0 such that for any w ∈ R d , E [ ∥ ϵ t +1 ( w ) ∥ 2 ] ≤ σ 21 and E [ 
Revised: and there are σ 1 , σ 2 > 0 such that for any w ∈ R d , E [ ∥ ϵ t +1 ( w ) ∥ 2 ] ≤ σ 21 and E [ ∥ J ⊤ ϵ t

Original: ∥ J ⊤ ϵ t +1 ( w ) ∥ 2 ] ≤ σ 2 . , where J ϵ t +1 is Jacobian of ϵ t +1 .
Revised: +1 ( w ) ∥ 2 ] ≤ σ 2 . , where J ϵ t +1 is Jacobian of ϵ t +1 .

Original: Remark.
Revised: Remark.

Original: The smoothness and boundedness conditions (A1) on the objective function and the zero-mean and the bounded variance conditions (A2) on stochastic gradient noise are commonly assumed in the convergence analysis for the stochastic optimization methods.
Revised: The smoothness and boundedness conditions (A1) on the objective function and the zero-mean and the bounded variance conditions (A2) on stochastic gradient noise are commonly assumed in the convergence analysis for the stochastic optimization methods.

Original: Moreover, if Hessian matrix satisfies − LI ⪯ ∇ 2 w ℓ ( w, z ) ⪯ LI in Example 1, then the last condition on J ϵ t +1 also holds with at least σ 2 = 2 L because J ϵ t +1 ( w )
Revised: Moreover, if Hessian matrix satisfies − LI ⪯ ∇ 2 w ℓ ( w, z ) ⪯ LI in Example 1, then the last condition on J ϵ t +1 also holds with at least σ 2 = 2 L because J ϵ t +1 ( w )

Original: = ∇ 2 w ℓ ( w, Z t +1 )
Revised: = ∇ 2 w ℓ ( w, Z t +1 )

Original: − ∇ 2 f ( w ) .
Revised: − ∇ 2 f ( w ) .

Original:  A LTERNATIVE V IEW OF S TOCHASTIC G RADIENT D ESCENT
Revised:  A LTERNATIVE V IEW OF S TOCHASTIC G RADIENT D ESCENT

Original: An alternative view (Kleinberg et al., 2018) of SGD is the key in our analysis relating to an implicit bias towards a flat minimum.
Revised: An alternative view (Kleinberg et al., 2018) of SGD is the key in our analysis relating to an implicit bias towards a flat minimum.

Original: We introduce this view with a refined convergence analysis and see the bias-optimization tradeoff caused by the stochastic gradient noise with a step size.
Revised: We introduce this view with a refined convergence analysis and see the bias-optimization tradeoff caused by the stochastic gradient noise with a step size.

Original: An alternative view of SGD considers an associated iterations { v t } ∞ t =0 with { w t } ∞ t =0 , which approximately minimizes a smoothed objective function obtained by the stochastic gradient noise.
Revised: An alternative view of SGD considers an associated iterations { v t } ∞ t =0 with { w t } ∞ t =0 , which approximately minimizes a smoothed objective function obtained by the stochastic gradient noise.

Original: We here define v t as a parameter obtained by the exact gradient descent from w t , that is, v t = w t − η ∇ f ( w t ) and we analyze the update of v t instead of w t .
Revised: We here define v t as a parameter obtained by the exact gradient descent from w t , that is, v t = w t − η ∇ f ( w t ) and we analyze the update of v t instead of w t .

Original: Since w t +1 = v t − ηϵ t +1 ( w t ) , we get v t +

Revised: Since w t +1 = v t − ηϵ t +1 ( w t ) , we get v t +

Original: =
Revised: =

Original: v t − ηϵ t +1 ( w t ) − η ∇ f ( v t − ηϵ t +1 ( w t )) .
Revised: v t − ηϵ t +1 ( w t ) − η ∇ f ( v t − ηϵ t +1 ( w t )) .

Original: As shown in Appendix A. 1, under a specific setting given later, w (cid:55)→ v = w − η ∇ f ( w ) becomes a smooth invertible injection and its inverse is differentiable, thus, we identify ϵ ′ t
Revised: As shown in Appendix A. 1, under a specific setting given later, w (cid:55)→ v = w − η ∇ f ( w ) becomes a smooth invertible injection and its inverse is differentiable, thus, we identify ϵ ′ t

Original: +1 ( v ) with ϵ t +1 ( w ) through the map w (cid:55)→ v .
Revised: +1 ( v ) with ϵ t +1 ( w ) through the map w (cid:55)→ v .

Original: Then, we get an update rule of v t :
Revised: Then, we get an update rule of v t :

Original: [Equation]
Revised: [Equation]

Original: For convenience, we refer to the rule (2) as an implicit stochastic gradient descent in this paper.
Revised: For convenience, we refer to the rule (2) as an implicit stochastic gradient descent in this paper.

Original: Since, the conditional expectation of ϵ ′ t +1 ( v t ) at v t is zero, we expect that the implicit SGD (2) minimizes the following smoothed objective function:
Revised: Since, the conditional expectation of ϵ ′ t +1 ( v t ) at v t is zero, we expect that the implicit SGD (2) minimizes the following smoothed objective function:

Original: [Equation]
Revised: [Equation]

Original: where ϵ ′ is an independent copy of ϵ ′ 1 , ϵ ′ 2 , . .
Revised: where ϵ ′ is an independent copy of ϵ ′ 1 , ϵ ′ 2 , . .

Original: However, we note that this implicit SGD is not a standard SGD because ∇ f ( v t − ηϵ ′ t +1 ( v t )) is a biased estimate of ∇ F ( v ) (i.e., ∇ F ( v )
Revised: However, we note that this implicit SGD is not a standard SGD because ∇ f ( v t − ηϵ ′ t +1 ( v t )) is a biased estimate of ∇ F ( v ) (i.e., ∇ F ( v )

Original: ̸
Revised: ̸

Original: = E [ ∇ f ( v − ηϵ ′ ( v ))] ) in general 1 , and thus we need a detailed convergence analysis.
Revised: = E [ ∇ f ( v − ηϵ ′ ( v ))] ) in general 1 , and thus we need a detailed convergence analysis.

Original: which The function (3) is actually a smoothed function of f by the convolution using the stochastic gradient noise ηϵ ′ and the level of smoothness is controlled by the step-size η as seen in the left of Figure 1 which depicts the original objective f corresponding to η = 0 and smoothed objectives F .
Revised: The function (3) is actually a smoothed function of f by the convolution using the stochastic gradient noise ηϵ ′ and the level of smoothness is controlled by the step-size η as seen in the left of Figure 1 which depicts the original objective f corresponding to η = 0 and smoothed objectives F .

Original: In this figure, we can observe how a nonconvex function is smoothened and its sharp local minima are eliminated by an 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).
Revised: In this figure, we can observe how a nonconvex function is smoothened and its sharp local minima are eliminated by an 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).

Original: Hence, we expect that stochastic gradient descent can avoid sharp minima and converges to a flat region.
Revised: Hence, we expect that stochastic gradient descent can avoid sharp minima and converges to a flat region.

Original: Indeed, by taking Taylor expansion of f , we see that F ( v ) is an approximation of the function f ( v ) plus the penalization on the high (positive) curvature of f along the noise direction in expectation:
Revised: Indeed, by taking Taylor expansion of f , we see that F ( v ) is an approximation of the function f ( v ) plus the penalization on the high (positive) curvature of f along the noise direction in expectation:

Original: [Equation]
Revised: [Equation]

Original: 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 ) .
Revised: 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 ) .

Original: In this paper, we make the following one-point strong convexity at the solution v ∗ to show the convergence of
Revised: In this paper, we make the following one-point strong convexity at the solution v ∗ to show the convergence of

Original: F ( v t ) .
Revised: F ( v t ) .

Original: Let v ∗ =
Revised: Let v ∗ =

Original: arg min v ∈ R d F ( v ) .
Revised: arg min v ∈ R d F ( v ) .

Original: We note that F and v ∗ depend on the value of η , but we do not explicitly denote this dependency for simplicity.
Revised: We note that F and v ∗ depend on the value of η , but we do not explicitly denote this dependency for simplicity.

Original: Assumption 2.
Revised: Assumption 2.

Original: (A3)
Revised: (A3)

Original: There is c > 0 such that for any v ∈ R d , ∇ F ( v ) ⊤ ( v − v ∗ ) ≥ c ∥ v − v ∗ ∥ 2 .
Revised: There is c > 0 such that for any v ∈ R d , ∇ F ( v ) ⊤ ( v − v ∗ ) ≥ c ∥ v − v ∗ ∥ 2 .

Original: For instance, this assumption holds for the function in the bottom-left in Figure 1 with sufficiently large η and for the function in the top-left figure with any η in a certain interval (0 , η 0 ] .
Revised: For instance, this assumption holds for the function in the bottom-left in Figure 1 with sufficiently large η and for the function in the top-left figure with any η in a certain interval (0 , η 0 ] .

Original: Assumption (A3) is a normal one-point strong convexity, whereas Kleinberg et al.
Revised: Assumption (A3) is a normal one-point strong convexity, whereas Kleinberg et al.

Original: (2018) assumed a different condition: E [ ∇ f ( v − ηϵ ′ ( v ))]
Revised: (2018) assumed a different condition: E [ ∇ f ( v − ηϵ ′ ( v ))]

Original: ⊤ ( v − v ◦ ) ≥ c
Revised: ⊤ ( v − v ◦ ) ≥ c

Original: ∥ v − v ◦ ∥ 2 at some parameter v ◦ and showed the convergence to v ◦ .
Revised: ∥ v − v ◦ ∥ 2 at some parameter v ◦ and showed the convergence to v ◦ .

Original: If ∇ F ( w ) =
Revised: If ∇ F ( v ) =

Original: E [ ∇ f ( v − ηϵ ′ ( v ))] , then v ◦ should be v ∗ and both assumptions coincide.
Revised: E [ ∇ f ( v − ηϵ ′ ( v ))] , then v ◦ should be v ∗ and both assumptions coincide.

Original: However, as noted above ∇ F ( w ) ̸
Revised: However, as noted above ∇ F ( v ) ̸

Original: = E [ ∇ f ( v − ηϵ ′ ( v ))] in general, and hence v ◦ is not necessarily v ∗ .
Revised: = E [ ∇ f ( v − ηϵ ′ ( v ))] in general, and hence v ◦ is not necessarily v ∗ .

Original: Our aim is to clarify how precisely SGD and averaged SGD can minimize F ( v ) .
Revised: Our aim is to clarify how precisely SGD and averaged SGD can minimize F ( v ) .

Original: That is why we make the normal one-point strong convexity at v ∗ and need a much more detailed analysis.
Revised: That is why we make the normal one-point strong convexity at v ∗ and need a much more detailed analysis.

Original: Moreover, our proof allows for a larger step size than that in Kleinberg et al.
Revised: Moreover, our proof allows for a larger step size than that in Kleinberg et al.

Original: because of the different proof techniques.
Revised: because of the different proof techniques.

Original: Theorem 1.
Revised: Theorem 1.

Original: Under Assumption (A1) , (A2) , and (A3) , run SGD for T -iterations with the step size η ≤ 12 L , then a sequence { v t } ∞ t =0 of the implicit SGD satisfies the following inequality:
Revised: Under Assumption (A1) , (A2) , and (A3) , run SGD for T -iterations with the step size η ≤ 12 L , then a sequence { v t } ∞ t =0 of the implicit SGD satisfies the following inequality:

Original: [Equation]
Revised: [Equation]

Original: Remark.
Revised: Remark.

Original: If σ 1 = 0 , then SGD is nothing but deterministic gradient descent and f = F because of the absence of stochastic gradient noise.
Revised: If σ 1 = 0 , then SGD is nothing but deterministic gradient descent and f = F because of the absence of stochastic gradient noise.

Original: Hence, SGD converges to a minimizer of f according to the classical optimization theory, which is recovered by Theorem 1 with σ 1 = 0 .
Revised: Hence, SGD converges to a minimizer of f according to the classical optimization theory, which is recovered by Theorem 1 with σ 1 = 0 .

Original: This theorem shows the convergence of SGD to the minimum of the smoothed objective F up to distance O ( √ η ) from v ∗ as long as F satisfies required assumptions even if the original objective f has local minima.
Revised: This theorem shows the convergence of SGD to the minimum of the smoothed objective F up to distance O ( √ η ) from v ∗ as long as F satisfies required assumptions even if the original objective f has local minima.

Original: This is also true for w t since ∥ w t − v t ∥
Revised: This is also true for w t since ∥ w t − v t ∥

Original: =
Revised: =

Original: O ( η ) .
Revised: O ( η ) .

Original: Thus, convergence to a flatter region is expected through an explicit expression as a regularized objective (4).
Revised: Thus, convergence to a flatter region is expected through an explicit expression as a regularized objective (4).

Original: Moreover, we can 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 η .
Revised: Moreover, we can 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 η .

Original: On the other hand, a small step size clearly weakens the bias.
Revised: On the other hand, a small step size clearly weakens the bias.

Original: Thus, the step size η controls the bias-optimization tradeoff coming from the stochastic gradient noise.
Revised: Thus, the step size η controls the bias-optimization tradeoff coming from the stochastic gradient noise.

Original:  A VERAGED SGD WITH L ARGE S TEP - SIZE
Revised:  A VERAGED SGD WITH L ARGE S TEP - SIZE

Original: Izmailov et al.
Revised: Izmailov et al.

Original: (2018) empirically demonstrated that averaged SGD converges to a flat region and achieves better generalization even when SGD oscillates with a relatively large step size. 
Revised: (2018) empirically demonstrated that averaged SGD converges to a flat region and achieves better generalization even when SGD oscillates with a relatively large step size. We ηϵ ′ ( v )))]

Revised: = (cid:82) ∇ ( f ( v − ηϵ ′ ( v, z )))

Revised: d P ( z ) , whereas ∇ f ( v − ηϵ ′ ( v )) in Eq.

Revised: (2) means ∇ f ( w ) | w = v − ηϵ

Revised: ′ ( v ) which does not involve the derivative of ϵ ′ ( v ) .

Original: We theoretically attribute this phenomenon to that the averaged SGD can get closer to v ∗ than SGD using the same step size under certain settings.
Revised: Therefore, ∇ F ( v ) ̸ = E [ ∇ f ( v − ηϵ ′ ( v ))] in general. theoretically attribute this phenomenon to that the averaged SGD can get closer to v ∗ than SGD using the same step size under certain settings.

Original: In other words, parameter averaging can improve the bias-optimization tradeoff and bring out the implicit bias more effectively.
Revised: In other words, parameter averaging can improve the bias-optimization tradeoff and bring out the implicit bias more effectively.

Original: In the averaged SGD, we run normal SGD (1) and take the average as follows:
Revised: In the averaged SGD, we run normal SGD (1) and take the average as follows:

Original: [Equation]
Revised: [Equation]

Original: Our aim is to show lim T →∞ w T can be closer to v ∗ than { w t } ∞ t =0 and { v t } ∞ t =0 by clarifying the dependency of this limit on the step size η .
Revised: Our aim is to show lim T →∞ w T can be closer to v ∗ than { w t } ∞ t =0 and { v t } ∞ t =0 by clarifying the dependency of this limit on the step size η .

Original: Preferably, the implicit SGD (2) is more useful in analyzing the averaged SGD because the average v T = 1 T (cid:80) Tt
Revised: Preferably, the implicit SGD (2) is more useful in analyzing the averaged SGD because the average v T = 1 T (cid:80) Tt

Original: =

Revised: =

Original: v t is consistent with w T as confirmed below.
Revised: v t is consistent with w T as confirmed below.

Original: By the definition, we see
Revised: By the definition, we see

Original: [Equation]
Revised: [Equation]

Original: where the noise term (cid:80) Tt =

Revised: where the noise term (cid:80) Tt =

Original: ϵ t +1 ( w t ) /
Revised: ϵ t +1 ( w t ) /

Original: ( T +1) is zero in expectation and its variance is upper bounded by σ/ ( T +1) under Assumption (A2) .
Revised: ( T +1) is zero in expectation and its variance is upper bounded by σ/ ( T +1) under Assumption (A2) .

Original: Hence, w T +1 − v T converges to zero in probability by Chebyshev’s inequality; for any r > 0 , P [ ∥ w T +1 − v T ∥ > r ] ≤ σ 21 / ( T +1) r 2 → 0 as T → ∞ , and the analysis of lim T →∞ w T reduces to that of lim T →∞ v T .
Revised: Hence, w T +1 − v T converges to zero in probability by Chebyshev’s inequality; for any r > 0 , P [ ∥ w T +1 − v T ∥ > r ] ≤ σ 21 / ( T +1) r 2 → 0 as T → ∞ , and the analysis of lim T →∞ w T reduces to that of lim T →∞ v T .

Original: We further make the additional assumptions on the smoothed objective F : R d → R and give the theorem that shows the convergence of the averaged SGD.
Revised: We further make the additional assumptions on the smoothed objective F : R d → R and give the theorem that shows the convergence of the averaged SGD.

Original: Assumption 3.
Revised: Assumption 3.

Original: (A4)F is three times continuously differentiable and ∇ 3 F is uniformly bounded, i.e., there is
Revised: (A4)

Original:  M > 0 such that for any v, w ∈ R dand k ∈ { 1 , 2 , . . . , } , ∥∇3 F ( v ) ww ⊤ ∥
Revised: There is M > 0 such that for any v ∈ R d , ∥∇ F ( v ) − ∇ 2 F ( v ∗ )( v − v ∗ )

Original: ≤ 2 M ∥ w ∥2 , where ∇ 3 F ( v ) ww ⊤ is a vector whose k -th element is (cid:80) di,j =1 ∂ [Equation] v )
Revised: ∥

Original: [Equation] w i w ′ j .
Revised: ≤ M ∥ v − v ∗ ∥ 2 .

Original: (A5) ∇ 2 F ( v ∗ ) is positive, i.e., there is µ > 

Revised: (A5) ∇ 2 F ( v ∗ ) is positive, i.e., there is µ >

Original: such that ∇ 2 F ( v ∗ ) ⪰ µI .
Revised: such that ∇ 2 F [Equation]

Original: Remark.
Revised: Remark.

Original: The boundedness of the third-order derivative (A4) is typically required to show the superiority of the averaging scheme as in Dieuleveut et al.
Revised: (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 et al.

Original: (2020) which studies the Markov chain view of SGD for strongly convex optimization problems.

Original: The positivity of Hessian (A5) is only required at v ∗ , which is consistent with nonconvexity.
Revised: The positivity of Hessian (A5) is only required at v ∗ , which is consistent with nonconvexity.

Original: For instance, examples in Figure 1 satisfy (A5) .
Revised: For instance, examples in Figure 1 satisfy (A5) .

Original: Theorem 2.
Revised: Theorem 2.

Original: Under Assumption (A1) – (A5) , run the averaged SGD for T -iterations with the step size η ≤ 12 L , then the average v T satisfies the following inequality:
Revised: Under Assumption (A1) – (A5) , run the averaged SGD for T -iterations with the step size η ≤ 12 L , then the average v T satisfies the following inequality:

Original: [Equation]
Revised: [Equation]

Original:  Theorem 2 says that the averagedSGD can optimize the smoothed objective F with better accuracy of
Revised: The variance of the averaged parameter v T is typically small, hence we evaluate the distance of E [ v T ] to v ∗ . Indeed, this is reasonable because the central limit theorem holds for averaged SGD under the mild conditoin even for nonconvex problems (Yu et al., 2020). Theorem 2 says that the averaged

Original:  O ( η ) than O ( √ η ) achieved by SGD using the same step size as long as the required conditions (onepoint strong convexity at minimizer and regularity for smoothed objectives) are satisfied uniformly for η in a certain interval (0 , η 0 ] ( ∃ η 0 < 1 ).
Revised: SGD can optimize the smoothed objective F with better accuracy of O ( η ) than O ( √ η ) achieved by SGD using the same step size as long as the required conditions (one-point strong convexity at minimizer and regularity for smoothed objectives) are satisfied uniformly for η in a certain interval (0 , η 0 ] ( ∃ η 0 < 1 ).

Original: These uniform requirements hold for valleys like the top-left case of Figure 1 and likely holds in the final phase of training deep neural network because of the observation that the parameter eventually falls in a better-shaped valley (see Figure 4).
Revised: These uniform requirements hold for valleys like the top-left case of Figure 1 and likely holds in the final phase of training deep neural network because of the observation that the parameter eventually falls in a better-shaped valley (see Figure 4).

Original: 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 in SWA (Izmailov et al., 2018), whose benefit is well known even in the convex optimization (Rakhlin et al., 2012; Mücke et al., 2019).
Revised: 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 in SWA (Izmailov et al., 2018), whose benefit is well known even in the convex optimization (Rakhlin et al., 2012; Mücke et al., 2019).

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We evaluate the empirical performance of SGD and averaged SGD on image classification tasks using CIFAR10 and CIFAR100 datasets.
Revised: We evaluate the empirical performance of SGD and averaged SGD on image classification tasks using CIFAR10 and CIFAR100 datasets.

Original: To evaluate the usefulness of the parameter averaging for the other methods, we also compare SAM (Foret et al., 2020) with its averaging variant.
Revised: To evaluate the usefulness of the parameter averaging for the other methods, we also compare SAM (Foret et al., 2020) with its averaging variant.

Original: We employ the tail-averaging scheme where the average is taken over the last phase of training.
Revised: We employ the tail-averaging scheme where the average is taken over the last phase of training.

Original: [Table] [Table]
Revised: [Table] [Table]

Original: We use the CNN architectures: ResNet (He et al., 2016) with 50-layers (ResNet-50), WideResNet (Zagoruyko & Komodakis, 2016) withlayers and width 10 (WRN-28-10), and
Revised: We use the CNN architectures: ResNet (He et al., 2016) with 50-layers (ResNet-50), WideResNet (Zagoruyko & Komodakis, 2016) withlayers and width 10 (WRN-28-10), and

Original: Pyramid Network (Han et al., 2017) with 272 layers and widening factor 200.
Revised: Pyramid Network (Han et al., 2017) with 272 layers and widening factor 200.

Original: In all settings, we use the standard data augmentations: horizontal flip, normalization, padding by four pixels, random crop, and cutout (DeVries & Taylor, 2017), and we employ the weight decay with the coefficient 0 .
Revised: In all settings, we use the standard data augmentations: horizontal flip, normalization, padding by four pixels, random crop, and cutout (DeVries & Taylor, 2017), and we employ the weight decay with the coefficient 0 .

Original: 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.
Revised: 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.

Original: To see the dependence on the step size, we use two decay schedules for the parameter averaging.
Revised: To see the dependence on the step size, we use two decay schedules for the parameter averaging.

Original: Tablesummarizes milestones labeled by the symbols: ‘s’ , ‘m’ , and ‘l’ .
Revised: Tablesummarizes milestones labeled by the symbols: ‘s’ , ‘m’ , and ‘l’ .

Original: The initial step size and a decay factor of the step size are set to 0 .
Revised: The initial step size and a decay factor of the step size are set to 0 .

Original: 1 and 0 . 

Revised: 1 and 0 .

Original: in all cases.
Revised: in all cases.

Original: The averages are taken from 300 epochs for the schedules ‘s’ and ‘l’ , and from 160 epochs for the schedule ‘m’ .
Revised: The averages are taken from 300 epochs for the schedules ‘s’ and ‘l’ , and from 160 epochs for the schedule ‘m’ .

Original: These hyperparameters were tuned based on the validation sets.
Revised: These hyperparameters were tuned based on the validation sets.

Original: For a fair comparison, we run (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.
Revised: For a fair comparison, we run (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.

Original: We evaluate each method 5 times for ResNetand WRN-28-10, and 3 times for Pyramid network.
Revised: We evaluate each method 5 times for ResNetand WRN-28-10, and 3 times for Pyramid network.

Original: The averages of classification accuracies are listed in Table 2 with the standard deviations in brackets.
Revised: The averages of classification accuracies are listed in Table 2 with the standard deviations in brackets.

Original: We observe from the table that the parameter averaging for SGD improves the classification accuracies in all cases, especially on CIFAR100 dataset.
Revised: We observe from the table that the parameter averaging for SGD improves the classification accuracies in all cases, especially on CIFAR100 dataset.

Original: Eventually, the averaged SGD achieves comparable or better performance than SAM.
Revised: Eventually, the averaged SGD achieves comparable or better performance than SAM.

Original: Moreover, we also observe improvement by parameter averaging for SAM in most cases, which is consistent with the observations in Kaddour et al.
Revised: Moreover, we also observe improvement by parameter averaging for SAM in most cases, which is consistent with the observations in Kaddour et al.

Original: with ResNet-50 and WRN-28-10.
Revised: with ResNet-50 and WRN-28-10.

Original: Comparing results on CIFAR100 and CIFAR10, the large step size is better, and the small step size is relatively poor on CIFAR100 dataset, whereas the small step size generally works on CIFAR10 dataset.
Revised: Comparing results on CIFAR100 and CIFAR10, the large step size is better, and the small step size is relatively poor on CIFAR100 dataset, whereas the small step size generally works on CIFAR10 dataset.

Original: If we use the step-size strategy ‘l’ for CIFAR10, then the improvement becomes small (see Appendix B for this result).
Revised: If we use the step-size strategy ‘l’ for CIFAR10, then the improvement becomes small (see Appendix B for this result).

Original: 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 a small step size would be sufficient for simple datasets such that the normal SGD already achieves high accuracies.
Revised: 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 a small step size would be sufficient for simple datasets such that the normal SGD already achieves high accuracies.

Original: Moreover, we note that the averaged SGD on CIFAR quite works well with the large step-size schedule ‘l’ , but SGD itself does not converge and poorly performs under this schedule as seen in Figure 3.
Revised: Moreover, we note that the averaged SGD on CIFAR quite works well with the large step-size schedule ‘l’ , but SGD itself does not converge and poorly performs under this schedule as seen in Figure 3.

Original: The accuracy of SGD temporarily increases at the 300 epochs because of the decay of the step size, it decreases thereafter.
Revised: The accuracy of SGD temporarily increases at the 300 epochs because of the decay of the step size, it decreases thereafter.

Original: However, the average of such parameters achieves significantly high accuracy as expected by our theory.
Revised: However, the average of such parameters achieves significantly high accuracy as expected by our theory.

Original: Finally, we observe in Figure 4 the loss landscape around the convergent point is in better shape and forms an asymmetric valley.
Revised: Finally, we observe in Figure 4 the loss landscape around the convergent point is in better shape and forms an asymmetric valley.

Original: Therefore, we expect that the loss function around the solution uniformly satisfies the required conditions in our theory.
Revised: Therefore, we expect that the loss function around the solution uniformly satisfies the required conditions in our theory.

Original: Specifically, Figure 4 depicts the section of train and test loss functions across parameters obtained by the averaged SGD and SGD.
Revised: Specifically, Figure 4 depicts the section of train and test loss functions across parameters obtained by the averaged SGD and SGD.

Original: The middle figure is the close-up view at the edge and plots each parameter.
Revised: The middle figure is the close-up view at the edge and plots each parameter.

Original: The right figure depicts the smoothed objectives with Gaussian noises in addition to train and test losses in log-scale.
Revised: The right figure depicts the smoothed objectives with Gaussian noises in addition to train and test losses in log-scale.

Original: We observe in Figure 4 the phenomenon that SGD converges to an edge and averaged SGD converges to a flat side.
Revised: We observe in Figure 4 the phenomenon that SGD converges to an edge and averaged SGD converges to a flat side.

Original: This phenomenon can be explained by our theory because the minimizer of the smoothed asymptotic valley is shifted to a flat side as confirmed in a synthetic setting (Figure 2) and deep learning setting (the right of Figure 4).
Revised: This phenomenon can be explained by our theory because the minimizer of the smoothed asymptotic valley is shifted to a flat side as confirmed in a synthetic setting (Figure 2) and deep learning setting (the right of Figure 4).

Original: 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.
Revised: 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.

Original: Finally, we observe that averaged SGD achieves a lower test loss which makes about 2% improvement in the classification error on CIFAR100 dataset.
Revised: Finally, we observe that averaged SGD achieves a lower test loss which makes about 2% improvement in the classification error on CIFAR100 dataset.

Original: These observations are also consistent with the experiments conducted in He et al.
Revised: These observations are also consistent with the experiments conducted in He et al.

Original:  R ELATED L ITERATURE AND D ISCUSSION
Revised:  R ELATED L ITERATURE AND D ISCUSSION

Original: Flat Minimum.
Revised: Flat Minimum.

Original: Keskar et al.
Revised: Keskar et al.

Original: (2017) and Hochreiter & Schmidhuber (1997) showed a flat minimum generalizes well and a sharp minimum generalizes poorly.
Revised: (2017) and Hochreiter & Schmidhuber (1997) showed a flat minimum generalizes well and a sharp minimum generalizes poorly.

Original: However, the flatness solely cannot explain generalization because it can be easily manipulated (Dinh et al., 2017).
Revised: However, the flatness solely cannot explain generalization because it can be easily manipulated (Dinh et al., 2017).

Original: Neyshabur et al.
Revised: Neyshabur et al.

Original: (2017) rigorously proved the sharpness combined with ℓ 2 -norm provides a generalization bound and Jiang et al.
Revised: (2017) rigorously proved the sharpness combined with ℓ 2 -norm provides a generalization bound and Jiang et al.

Original: (2020) verified this correlation through large scale experiments.
Revised: (2020) verified this correlation through large scale experiments.

Original: Keskar et al.
Revised: Keskar et al.

Original: (2017) also argued that SGD converges to a flat minimum and He et al.
Revised: (2017) also argued that SGD converges to a flat minimum and He et al.

Original: (2019) argued the averaged SGD tends to converge to an asymmetric valley.
Revised: (2019) argued the averaged SGD tends to converge to an asymmetric valley.

Original: Several works (Kleinberg et al., 2018;
Revised: Several works (Kleinberg et al., 2018;

Original: Zhou et al., 2020) studied the stochastic gradient noise to theoretically prove the existence of an implicit bias towards a flat region or asymmetric valley.
Revised: Zhou et al., 2020) studied the stochastic gradient noise to theoretically prove the existence of an implicit bias towards a flat region or asymmetric valley.

Original: Moreover, many works Izmailov et al.
Revised: Moreover, many works (Izmailov et al., 2018;

Original: ; Foret et al. ; Damian et al. (2021); Orvieto et al. (2022) studied the techniques to further bring out the bias of SGD.
Revised: Foret et al., 2020; Damian et al., 2021; Orvieto et al., 2022) studied the techniques to further bring out the bias of SGD.

Original: In particular, SAM and SWA achieved a significant improvement in the generalization performance.
Revised: In particular, SAM and SWA achieved a significant improvement in the generalization performance.

Original: 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.
Revised: 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.

Original: Markov Chain Interpretation of SGD.
Revised: Markov Chain Interpretation of SGD.

Original: Dieuleveut et al.
Revised: Dieuleveut et al.

Original: Yu et al.
Revised: Yu et al.

Original: (2020) provided the
Revised: (2020) provided the

Original: Markov chain interpretation of SGD.
Revised: Markov chain interpretation of SGD.

Original: They showed the marginal distribution of the parameter of SGD converges to an invariant distribution for convex and nonconvex optimization problems, respectively.
Revised: They showed the marginal distribution of the parameter of SGD converges to an invariant distribution for convex and nonconvex optimization problems, respectively.

Original: Moreover, Dieuleveut et al.
Revised: Moreover, Dieuleveut et al.

Original: (2020) showed the mean of the invariant distribution, attained by the averaged SGD, is at distance O ( η ) from the minimizer of the objective function, whereas SGD itself oscillates at distance O ( √ η ) in the convex optimization settings.
Revised: (2020) showed the mean of the invariant distribution, attained by the averaged SGD, is at distance O ( η ) from the minimizer of the objective function, whereas SGD itself oscillates at distance O ( √ η ) in the convex optimization settings.

Original: Izmailov et al.
Revised: Izmailov et al.

Original: (2018) also attributed the success of SWA to such a phenomenon.
Revised: (2018) also attributed the success of SWA to such a phenomenon.

Original: That is, Izmailov et al.
Revised: That is, Izmailov et al.

Original: (2018) explained that SGD travels on the hypersphere because of the convergence to Gaussian distribution and the concentration on the sphere under a simplified setting, and thus averaging scheme allows us to go inside of the sphere which may be flat.
Revised: (2018) explained that SGD travels on the hypersphere because of the convergence to Gaussian distribution and the concentration on the sphere under a simplified setting, and thus averaging scheme allows us to go inside of the sphere which may be flat.

Original: We can say our contribution is to theoretically justify this intuition by extending the result obtained by Dieuleveut et al.
Revised: We can say our contribution is to theoretically justify this intuition by extending the result obtained by Dieuleveut et al.

Original: (2020) to a nonconvex optimization setting.
Revised: (2020) to a nonconvex optimization setting.

Original: In the proof, we utilize the alternative view of SGD (Kleinberg et al., 2018) in a non-asymptotic way under some conditions not on the original objective but on the smoothed objective function.
Revised: In the proof, we utilize the alternative view of SGD (Kleinberg et al., 2018) in a non-asymptotic way under some conditions not on the original objective but on the smoothed objective function.

Original: Combination with the Markov chain view for nonconvex objective (Yu et al., 2020) may be helpful in further relaxing required conditions.
Revised: Combination with the Markov chain view for nonconvex objective (Yu et al., 2020) may be helpful in more detailed analyses.

Original: Step size and Minibatch.
Revised: Step size and Minibatch.

Original: SGD with a large step size often suffers from stochastic gradient noise and becomes unstable.
Revised: SGD with a large step size often suffers from stochastic gradient noise and becomes unstable.

Original: This is the reason why we should take a smaller step size so that SGD converges.
Revised: This is the reason why we should take a smaller step size so that SGD converges.

Original: In this sense, the minibatching of stochastic gradients clearly plays the same role as the step size and sometimes brings additional gains.
Revised: In this sense, the minibatching of stochastic gradients clearly plays the same role as the step size and sometimes brings additional gains.

Original: For instance, Smith et al.
Revised: For instance, Smith et al.

Original: (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.
Revised: (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.

Original: We remark that our analysis can incorporate the minibatch by dividing σ 21 and σ 22 in Theorem 1 and 2 by the minibatch size, and we can see certain improvements of optimization accuracy as well.
Revised: We remark that our analysis can incorporate the minibatch by dividing σ 21 and σ 22 in Theorem 1 and 2 by the minibatch size, and we can see certain improvements of optimization accuracy as well.

Original: Then, both SGD and averaged SGD share the same dependency on the minibatch size and thus controlling step size seems  beneficial for parameter averaging.
Revised: 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.

Original: Edge of Stability.
Revised: Edge of Stability.

Original: Recently, Cohen et al.
Revised: Recently, Cohen et al.

Original: (2021) showed the deterministic gradient descent for deep neural networks enters Edge of Stability phase.
Revised: (2021) showed the deterministic gradient descent for deep neural networks enters Edge of Stability phase.

Original: In the traditional optimization theory, the step size is set to be smaller than 1 /L to ensure stable convergence and we also make such a restriction.
Revised: In the traditional optimization theory, the step size is set to be smaller than 1 /L to ensure stable convergence and we also make such a restriction.

Original: On the other hand, the Edge of Stability phase appears when using a higher step size than 2 /L .
Revised: On the other hand, the Edge of Stability phase appears when using a higher step size than 2 /L .

Original: In this phase, the training loss behaves non-monotonically and the sharpness finally stabilizes around 2 /η .
Revised: In this phase, the training loss behaves non-monotonically and the sharpness finally stabilizes around 2 /η .

Original: 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.
Revised: 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.

Original: There are works (Arora et al., 2022;
Revised: There are works (Arora et al., 2022;

Original: Ahn et al., 2022) which attempted to rigorously justify Edge of Stability phase.
Revised: Ahn et al., 2022) which attempted to rigorously justify Edge of Stability phase.

Original: Interestingly, their analyses are based on a similar intuition to ours, but we consider a different regime of step sizes and a different factor (stochastic noise or larger step size than 2 /L ) brings the implicit bias towards flat regions.
Revised: Interestingly, their analyses are based on a similar intuition to ours, but we consider a different regime of step sizes and a different factor (stochastic noise or larger step size than 2 /L ) brings the implicit bias towards flat regions.

Original: We believe establishing a unified theory is interesting future research.
Revised: We believe establishing a unified theory is interesting future research.

Original: Averaged SGD.
Revised: Averaged SGD.

Original: The averaged SGD (Ruppert, 1988; Polyak & Juditsky, 1992) is a popular variant of
Revised: The averaged SGD (Ruppert, 1988; Polyak & Juditsky, 1992) is a popular variant of

Original: SGD, which returns the average of parameters obtained by SGD aiming at stabilizing the convergence.
Revised: SGD, which returns the average of parameters obtained by SGD aiming at stabilizing the convergence.

Original: 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 / √ T ) and
Revised: 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 / √ T ) and

Original: O (1 /T ) for non-strongly convex and strongly convex problems (Nemirovski et al., 2009; Bach &
Revised: O (1 /T ) for non-strongly convex and strongly convex problems (Nemirovski et al., 2009; Bach &

Original: Moulines, 2011; Rakhlin et al., 2012; Lacoste-Julien et al., 2012).
Revised: Moulines, 2011; Rakhlin et al., 2012; Lacoste-Julien et al., 2012).

Original: 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.
Revised: 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.

Original: When applying a non-diminishing step size schedule, the non-zero optimization error basically remains.
Revised: When applying a non-diminishing step size schedule, the non-zero optimization error basically remains.

Original: What we do in this paper is to characterize it as the implicit bias toward a flat region.
Revised: What we do in this paper is to characterize it as the implicit bias toward a flat region.

Original: CONCLUSION
Revised: CONCLUSION

Original: In this paper, we showed that parameter averaging improves the bias-optimization tradeoff caused by the stochastic gradient noise.
Revised: In this paper, we showed that parameter averaging improves the bias-optimization tradeoff caused by the stochastic gradient noise.

Original: Specifically, we proved that averaged SGD optimizes the smoothed objective functions up to O ( η ) -error, whereas SGD itself optimizes it up to O ( √ η ) -error in terms of Euclidean distance from the solution, where η is the step size.
Revised: Specifically, we proved that averaged SGD optimizes the smoothed objective functions up to O ( η ) -error, whereas SGD itself optimizes it up to O ( √ η ) -error in terms of Euclidean distance from the solution, where η is the step size.

Original: Therefore, parameter averaging significantly stabilizes the implicit bias toward a flat region, and we can expect improved performance for difficult datasets such that the stronger bias induced by a larger step size is helpful.
Revised: Therefore, parameter averaging significantly stabilizes the implicit bias toward a flat region, and we can expect improved performance for difficult datasets such that the stronger bias induced by a larger step size is helpful.

Original: Finally, we observed the consistency of our theory with the experiments on image classification tasks.
Revised: Finally, we observed the consistency of our theory with the experiments on image classification tasks.

Original: In addition to the above discussion, another interesting research direction is to investigate what type of noise is strongly related to generalization performance.
Revised: In addition to the above discussion, another interesting research direction is to investigate what type of noise is strongly related to generalization performance.

Original: Appendix A P ROOFS A.I MPLICIT S TOCHASTIC G RADIENT D ESCENT
Revised: Appendix A P ROOFS A.I MPLICIT S TOCHASTIC G RADIENT D ESCENT

Original: Let denote by φ : R d → R d a change of variables from w to v introduced in Section 3, i.e., v = φ ( w ) =
Revised: Let denote by φ : R d → R d a change of variables from w to v introduced in Section 3, i.e., v = φ ( w ) =

Original: w − η ∇ f ( w ) .
Revised: w − η ∇ f ( w ) .

Original: Lemma A. Under Assumption (A1) and η ≤ 12 L , the function φ is injective and invertible, and its inverse φ − 1 defined on on Im φ is differentiable.
Revised: Lemma A. Under Assumption (A1) and η ≤ 12 L , the function φ is injective and invertible, and its inverse φ − 1 defined on on Im φ is differentiable.

Original: Proof.
Revised: Proof.

Original: For w, w ′ ∈ R d , we suppose φ ( w ) =
Revised: For w, w ′ ∈ R d , we suppose φ ( w ) =

Original: φ ( w ′ ) .
Revised: φ ( w ′ ) .

Original: Then, it holds that
Revised: Then, it holds that

Original: [Equation]
Revised: [Equation]

Original: where we used L -Lipschitz continuity of ∇ f due to (A1) .
Revised: where we used L -Lipschitz continuity of ∇ f due to (A1) .

Original: Therefore, we see w = w ′ and φ is an injection.
Revised: Therefore, we see w = w ′ and φ is an injection.

Original: Moreover, since J φ ( w ) = I − η ∇ 2 f ( w ) ⪰ (1 − ηL )
Revised: Moreover, since J φ ( w ) = I − η ∇ 2 f ( w ) ⪰ (1 − ηL )

Original: I ⪰ 12 I .
Revised: I ⪰ 12 I .

Original: Thus, φ is invertible and φ − 1 , which is defined on Im φ , is differentiable because of the injectivity and the inverse map theorem.
Revised: Thus, φ is invertible and φ − 1 , which is defined on Im φ , is differentiable because of the injectivity and the inverse map theorem.

Original: Using φ , we see ϵ ′ ( v ) =
Revised: Using φ , we see ϵ ′ ( v ) =

Original: ϵ ( φ − 1 ( v )) for v ∈
Revised: ϵ ( φ − 1 ( v )) for v ∈

Original: Im φ .
Revised: Im φ .

Original: Let (Ω , F , P ) be a probability space such that ϵ ′ ( v ) can be represented as a measurable map z ∈ Ω
Revised: Let (Ω , F , P ) be a probability space such that ϵ ′ ( v ) can be represented as a measurable map z ∈ Ω

Original: (cid:55)→ ϵ ′ ( v, z ) .
Revised: (cid:55)→ ϵ ′ ( v, z ) .

Original: Note that we use ϵ ′ ( v ) and ϵ ′ ( v, z ) depending on the situation.
Revised: Note that we use ϵ ′ ( v ) and ϵ ′ ( v, z ) depending on the situation.

Original: For a function g : R d → R d , we denote by J g ( w ) Jacobian of g , i.e.,
Revised: For a function g : R d → R d , we denote by J g ( w ) Jacobian of g , i.e.,

Original: J g ( w ) =
Revised: J g ( w ) =

Original: ( ∂g i ( w ) /∂w j ) di,j =1 .
Revised: ( ∂g i ( w ) /∂w j ) di,j =1 .

Original: Lemma B. Under Assumption (A1) and (A2) , we get for any v ∈
Revised: Lemma B. Under Assumption (A1) and (A2) , we get for any v ∈

Original: Im φ ⊂ R d ,
Revised: Im φ ⊂ R d ,

Original: [Equation]
Revised: [Equation]

Original: Moreover, if η ≤ 12 L , then
Revised: Moreover, if η ≤ 12 L , then

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: The first equality of the statement can be confirmed by the direct calculation as follows:
Revised: The first equality of the statement can be confirmed by the direct calculation as follows:

Original: [Equation]
Revised: [Equation]

Original: Next, we evaluate the last term below.
Revised: Next, we evaluate the last term below.

Original: By the chain rule and inverse map theory,
Revised: By the chain rule and inverse map theory,

Original: [Equation]
Revised: [Equation]

Original: Note that from assumption for any w ∈ R d , J φ ( w ) =
Revised: Note that from assumption for any w ∈ R d , J φ ( w ) =

Original: I − η ∇ 2 f ( w ) ⪰ (1 − ηL )
Revised: I − η ∇ 2 f ( w ) ⪰ (1 − ηL )

Original: I ⪰ 12 I .
Revised: I ⪰ 12 I .

Original: Hence,
Revised: Hence,

Original: [Equation]
Revised: [Equation]

Original: Finally, we get
Revised: Finally, we get

Original: [Equation]
Revised: [Equation]

Original: This finishes the proof.
Revised: This finishes the proof.

Original: A.P ROOF OF T HEOREM
Revised: A.P ROOF OF T HEOREM

Original: The following proposition is the restatement of the well-known convergence result to a stationary point using the coordinate v .
Revised: The following proposition is the restatement of the well-known convergence result to a stationary point using the coordinate v .

Original: Proposition A. Under Assumption (A1) , (A2) , and η ≤ 12 L , we get
Revised: Proposition A. Under Assumption (A1) , (A2) , and η ≤ 12 L , we get

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: It is known that (A1 ) derives the following (Nesterov, 2004): for any w, w ′ ∈ R d ,
Revised: It is known that (A1 ) derives the following (Nesterov, 2004): for any w, w ′ ∈ R d ,

Original: [Equation]
Revised: [Equation]

Original: Substituting the update Eq.
Revised: Substituting the update Eq.

Original: (1) into this inequality with w ′ =
Revised: (1) into this inequality with w ′ =

Original: w t +1 and w = w t , and taking the conditional expectation E [ ·|F t ] , we get
Revised: w t +1 and w = w t , and taking the conditional expectation E [ ·|F t ] , we get

Original: [Equation]
Revised: [Equation]

Original: Thus, we have E [ f ( w t +1 )] ≤ E [ f ( w t )] − 3 η 4 E [ ∥∇ f ( w t ) ∥ 2 ]+ η 2 σ 21 L.
Revised: Thus, we have E [ f ( w t +1 )] ≤ E [ f ( w t )] − 3 η 4 E [ ∥∇ f ( w t ) ∥ 2 ]+ η 2 σ 21 L.

Original: By summing up this inequality, we get
Revised: By summing up this inequality, we get

Original: [Equation]
Revised: [Equation]

Original: where we used the nonnegativity of f .
Revised: where we used the nonnegativity of f .

Original: By dropping the term with t = 0 of the sum in the left hand side and using w t +1 = v t − ηϵ ′ t +1 ( v t ) , we finally get
Revised: By dropping the term with t = 0 of the sum in the left hand side and using w t +1 = v t − ηϵ ′ t +1 ( v t ) , we finally get

Original: [Equation]
Revised: [Equation]

Original: Using the above results, we prove Theorem 1, which is restated below.
Revised: Using the above results, we prove Theorem 1, which is restated below.

Original: Theorem A. Under Assumption (A1) , (A2) , and (A3) , run the stochastic gradient descent with
Revised: Theorem A. Under Assumption (A1) , (A2) , and (A3) , run the stochastic gradient descent with

Original: T -iterations with the step size η ≤ 12 L , then the implicit SGD satisfies the following inequality:
Revised: T -iterations with the step size η ≤ 12 L , then the implicit SGD satisfies the following inequality:

Original: [Equation]
Revised: [Equation]

Original: Proof of Theorem A. To evaluate ∥ v t +1 − v ∗ ∥ 2 for the implicit SGD (2)
Revised: Proof of Theorem A. To evaluate ∥ v t +1 − v ∗ ∥ 2 for the implicit SGD (2)

Original: [Equation]
Revised: [Equation]

Original: we first give several bounds as follows.
Revised: we first give several bounds as follows.

Original: By Assumption (A3)
Revised: By Assumption (A3)

Original: [Equation]
Revised: [Equation]

Original: Young’s inequality
Revised: Young’s inequality

Original: [Equation]
Revised: [Equation]

Original: and Lemma B
Revised: and Lemma B

Original: [Equation]
Revised: [Equation]

Original: we get
Revised: we get

Original: [Equation]
Revised: [Equation]

Original: By Assumption (A2) and Young’s inequality again, we get
Revised: By Assumption (A2) and Young’s inequality again, we get

Original: [Equation]
Revised: [Equation]

Original: Combining the above two inequalities
Revised: Combining the above two inequalities

Original: [Equation]
Revised: [Equation]

Original: we get
Revised: we get

Original: [Equation]
Revised: [Equation]

Original: Taking the expectation regarding all histories and summing up over t =
Revised: Taking the expectation regarding all histories and summing up over t =

Original: [Equation]
Revised: [Equation]

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: T
Revised: T

Original: [Equation]
Revised: [Equation]

Original: we get cη
Revised: we get cη

Original: T (cid:88)
Revised: T (cid:88)

Original: [Equation]
Revised: [Equation]

Original: where we used Proposition A.
Revised: where we used Proposition A.

Original: Therefore
Revised: Therefore

Original: [Equation]
Revised: [Equation]

Original: we conclude
Revised: we conclude

Original: [Equation]
Revised: [Equation]

Original: A.P ROOF OF T HEOREM
Revised: A.P ROOF OF T HEOREM

Original: We give several statements used to prove Theorem 2.
Revised: We give several statements used to prove Theorem 2.

Original: Lemma C. Under the same assumptions as in Theorem A, run the stochastic gradient descent with
Revised: Lemma C. Under the same assumptions as in Theorem A, run the stochastic gradient descent with

Original: T -iterations with the step size η ≤ 12 L , then the implicit SGD satisfies the following inequality:
Revised: T -iterations with the step size η ≤ 12 L , then the implicit SGD satisfies the following inequality:

Original: [Equation]
Revised: [Equation]

Original: Proof of Lemma C. By the simple calculation, we get
Revised: Proof of Lemma C. By the simple calculation, we get

Original: [Equation]
Revised: [Equation]

Original: [Table]

Original: [Equation]

Original: [Equation]

Original: We bound Eq. (7) as follows: [Equation] t =
Revised: where we used Theorem A

Original: E (cid:34)(cid:114)

Original: E

Original: [Equation]

Original: (cid:105)(cid:35) +

Original: [Equation]

Original: (cid:88) t = (cid:114)

Original: [Equation]
Revised: [Equation]

Original: (cid:104)(cid:13)(cid:13) [Equation] [Equation] where we used Assumption ( A1 ) and ( A2 )
Revised: Next

Original: [Equation]

Original: We bound Eq.

Original: (8) as follows:

Original: [Equation]
Revised: [Equation]

Original: where we used Theorem A. We bound Eq. (9) as follows:
Revised: we show the second inequality by using Lemma B as follows:

Original: [Equation]
Revised: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]
Revised: [Table]

Original: where we used Assumption ( A2 ).

Original: By combining the above inequalities, we obtain the first result in the statement.

Original: Next we show the second inequality by using Lemma B as follows:

Original: [Equation]
Revised: [Equation]

Original: Proposition B. Under the same assumptions as in Theorem A
Revised: Proposition B. Under the same assumptions as in Theorem A

Original: [Equation]
Revised: [Equation]

Original: run the stochastic gradient descent with T -iterations with the step size η ≤ 12 L , then the implicit SGD satisfies the following inequality:
Revised: run the stochastic gradient descent with T -iterations with the step size η ≤ 12 L , then the implicit SGD satisfies the following inequality:

Original: [Equation]
Revised: [Equation]

Original: Proof of Proposition B. Using Lemma C
Revised: Proof of Proposition B. Using Lemma C

Original: [Equation]
Revised: [Equation]

Original: we get
Revised: we get

Original: [Equation]
Revised: [Equation]

Original: We here prove Theorem 1 which is restated below.
Revised: We here prove Theorem 1 which is restated below.

Original: Theorem B. Under Assumption (A1) – (A5)
Revised: Theorem B. Under Assumption (A1) – (A5)

Original: [Equation]
Revised: [Equation]

Original: run the averaged SGD for T -iterations with the step size η ≤ 12 L , then the average v T satisfies the following inequality:
Revised: run the averaged SGD for T -iterations with the step size η ≤ 12 L , then the average v T satisfies the following inequality:

Original: [Equation]
Revised: [Equation]

Revised: σ 2 1 LM 3 cµ (cid:18) 1 + 2 ησ 2c (cid:19) .

Original: Proof.
Revised: Proof.

Original: We apply Taylor’s theorem to ∇ F around v ∗ and obtain ∇ F ( vt ) = ∇ 2 F ( v ∗ )( vt − v ∗ ) 
Revised: We define R ( v ) = ∇ F ( v ) −∇ 2 F ( v ∗ )( v − v ∗ ) .

Original: + R ( v t ) , where R : R d → R d satisfies ∥ R ( v ) ∥
Revised: Then, by (A4) , we see ∥ R ( v ) ∥

Original: ≤ M ∥ v − v ∗ ∥ 2by (A4) .
Revised: ≤ M ∥ v − v ∗ ∥ 2 .

Original: By taking average  over t ∈ { 0 , 1 , . . . , T } and rearranging terms, we get 
Revised: By taking average of R ( v t ) over t ∈ { 0 , 1 , . . . , T } and rearranging terms, we get ∇ 2 F ( v ∗ )( v T − v ∗ ) =

Original: [Equation]
Revised: T + T (cid:88) t = ∇ F ( v t ) − T + T (cid:88) t =

Original: [Equation]
Revised: R ( v t ) .

Original: Therefore, we get
Revised: Therefore, we get

Original: [Equation]
Revised: [Equation]

Revised: E

Revised: [Equation]

Revised: +

Revised: [Equation]

Original: The latter and former terms can be bounded by Theorem A and
Revised: E [Equation] [Equation] ≤

Revised: [Equation]

Revised: E (cid:34)

Revised: [Equation]

Revised: (cid:88)

Original: Proposition B. Thus, we finally get µ E [ ∥ v T − v ∗ ∥ ] ≤ O (cid:16) T − 12 (cid:17) +4 σ 1 σ 2 η 32 L 12 √
Revised: [Equation] ∇ F [Equation] +

Revised: M

Revised: T

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: E (cid:34) T (cid:88) t = ∥ v t − v ∗ ∥(cid:35) . The latter and former terms can be bounded by Theorem A and Proposition B. Thus, we finally get µ ∥ E [ v T ] − v ∗ ∥ ≤ O (cid:16) T − 12 (cid:17) + 4 σ 1 σ 2 η 32 L 12 √ + 2 ησ 21 M c + 8 η 2 σ 21 LM 3 c (cid:18) 1 + 2 ησc (cid:19) .

Original: B A DDITIONAL E XPERIMENTS
Revised: B A DDITIONAL E XPERIMENTS

Original:  We run SGD and averaged SGD on CIFAR10 dataset with the step size strategy ‘l’ under the same settings as in Section 5.
Revised: [Table] We run SGD and averaged SGD on CIFAR10 dataset with the step size strategy ‘l’ under the same settings as in Section 5.

Original: Table 3 lists the results including this case.
Revised: Table 3 lists the results including this case.

Original: We observe that the large step size ‘l’ does not work so well on CIFAR10 dataset compared to other schedules.
Revised: We observe that the large step size ‘l’ does not work so well on CIFAR10 dataset compared to other schedules.

Original: We hypothesize this is because CIFAR10 is not so difficult dataset and does not require stronger bias induced by a larger step size.
Revised: We hypothesize this is because CIFAR10 is not so difficult dataset and does not require stronger bias induced by a larger step size.

Original: We also validate the cosine annealing strategy for the step size, which is frequently used due to its excellent performance.
Revised: We also validate the cosine annealing strategy for the step size, which is frequently used due to its excellent performance.

Original: We used the symbols ‘s’ , ‘m’ , and ‘l’ for the cosine annealing depending on the last step sizes which are set to 0 , 0 . [Table] and 0 .
Revised: We used the symbols ‘s’ , ‘m’ , and ‘l’ for the cosine annealing depending on the last step sizes which are set to 0 , 0 . 004 , and 0 .

Original: 02 , respectively.
Revised: 02 , respectively.

Revised: The parameter averaging for averaged SGD is taken over the last quarter of the training.

Original: From the table, we observe the usefulness of parameter averaging for cosine annealing schedule as well. 
Revised: From the table, we observe the usefulness of parameter averaging for cosine annealing schedule as well. [Table] Finally [Table]

Revised: we run SGD, SGD with a large step size

Original: [Table]
Revised: [Table]

Revised: and averaged SGD to train the standard convolutional neural network on Fashion MNIST dataset to confirm how efficiently sharpness and classification accuracy can be optimized by each method.

Revised: We note the large step size used for SGD is the same as that for averaged SGD.

Revised: We plot the trace of Hessian ∇ 2 f ( w ) and test loss functions in Figure

Original: [Table]
Revised: [Table]

Revised: From this figure

Original: [Table]
Revised: [Table]

Revised: we observe that the averaged SGD converges to a flatter region and achieves the highest classification accuracy on the test dataset as expected in our theory

Original: [Table]
Revised: [Table]

Revised: C M OTIVATING E XAMPLE C.P ROBLEM S ETUP

Revised: In this section, we present a motivating example that verifies the convergence to a flat minimum and a certain separation between SGD and averaged SGD.

Revised: We consider a one-dimensional objective function f : R → R defined below: for p, δ > 0 ,

Revised: [Equation]

Revised: where g δ : R → R is a scaled mollifier:

Revised: [Equation]

Revised: g δ ( w ) =

Revised: δg 1 ( w/δ ) is a scaling of the well-known mollifier of g 1 which is an infinitely differentiable function with a compact support.

Revised: That is, g δ is a smooth function whose support is [ − δ, δ ] .

Revised: Because of the coefficient p of g δ , the function f ( w ) has a local minimum in [ − δ, δ ] , which can be the global minimum.

Revised: See Figure 6 (right).

Revised: The maximum values of the first and second derivatives of g 1 p are bounded.

Revised: Thus, we define constants

Revised: C 1 , C 2 by

Revised: [Equation]

Revised: 1 ( w ) | .

Revised: Since g ′′ δ ( w ) = 1 δ g ′′ 1 ( w/δ ) , we see the second derivative of g δ is bounded by C 2 pδ − 1 .

Revised: Hence,

Revised: Lipschitz smoothness (boundedness of Hessian) L of f is 1 + C 2 pδ − 1 .

Revised: Next, we consider the uniform noise on the interval [ − r, r ] for r > 0 , i.e., ϵ ∼ U [ − r, r ] and suppose ϵ ( w, z )

Revised: =

Revised: ϵ ( z )(=

Revised: ϵ ′ ( v, z )) where Ω ∋ z

Revised: (cid:55)→ ϵ ( w, z ) is an explicit representation of the random noise.

Revised: In other words, noise distribution does not change in w .

Revised: In this case, we see σ 21 =

Revised: E

Revised: [ ϵ 2 ] ≤ r 2 and σ 2 = 0 .

Revised: The smoothed objective F with the noise ϵ ′ and step-size η is

Revised: [Equation]

Revised: We consider the following problem setup:

Revised: [Equation]

Revised: Note we can choose arbitrarily small δ > 0 and large r which satisfy the above inequalities.

Revised: For appropriate smoothing, we choose the step size η so that

Revised: [Equation]

Revised: Note that a step-size η that satisfies the condition (10) exists and it also satisfies η ≤ 1 /L

Revised: =

Revised: δ/ ( δ + C 2 p ) required in the theory.

Revised: C.C ONVERGENCE OF SGD AND A VERAGED SGD

Revised: Under the above setup (8)–(10), we can estimate constants appearing in the convergence results of

Revised: SGD and averaged SGD as follows (for the detail see the next subsection):

Revised: [Equation]

Revised: Moreover, the minimum of the smoothed objective is v ∗ =

Revised: p , a sharp minimum ( ∼ 0 ) can be eliminated by smoothing.

Revised: Therefore

Revised: [Equation]

Revised: for SGD we obtain by Theorem 1,

Revised: [Equation]

Revised: We see from this inequality, η ∗ = 2 C 1 δ r is the best choice of the step-size, resulting in

Revised: T +

Revised: [Equation]

Revised: where we apply Jensen’s inequality to derive the bound on L 1 -norm.

Revised: This result means SGD avoids a sharp minimum (i

Revised: [Equation]

Revised: e., v ∼ 0 under small δ > 0 ) and converges to a flat minimum v ∗ =

Original: s
Revised: p

Original: [Table]
Revised: [Equation]

Original: [Table]
Revised: and a too large noise will affect the convergence to v ∗ based on our step-size policy. Moreover, for averaged SGD we obtain by Theorem 2, [Equation] Hence, for η ∗ = 2 C 1 δ r we obtain

Original: [Table]
Revised: [Equation]

Revised: This bound means averaged SGD will get closer to v ∗ =

Revised: p as long as SGD approaches a neighborhood of v ∗

Original: AveragedSGD s
Revised: According to the above results, both SGD and averaged SGD converge to a 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 . C.E STIMATION OF C ONSTANTS

Revised: We verify the estimations of constants in (11).

Revised: L , σ 21 , and σ 2 are already obtained, thus, mu, c , and

Revised: M remain.

Revised: Minimum and estimation of µ .

Revised: We first see that under our problem setting, the local minimum around the origin is eliminated and p is the optimal solution of F , i.e., v ∗ = p .

Revised: The smoothed function G δ ( v ) def =

Revised: E [ g δ ( v − ηϵ ′ )] and its derivative G ′ δ ( v ) are calculated as follows:

Original: [Table]
Revised: [Equation]

Revised: By taking into account supp( g δ )

Original: [Table]
Revised: = [ − δ, δ ] , the smoothed objective G δ ( v ) is constant on {| v | ≤ ηr − δ }∪{| v | ≥ ηr + δ } , and thus, G ′ δ is non-zero only on supp( G ′ δ ) = [ − ηr − δ, − ηr + δ ] ∪ [ ηr − δ, ηr + δ ] . See Figure 6 (left).

Revised: Since ηr + δ < p/ 4 < p under (10), v = p is still a local minimum of F .

Revised: We evaluate the bound on G ′ δ on supp( G ′ δ ) below.

Original: [Table]
Revised: for v ∈ [ ηr − δ, ηr + δ ] the support of g ′ δ ( v − ηt ) in t ∈ R is [( v − δ ) /η, ( v + δ ) /η ] , we get [Equation] where we used | g ′ δ ( v ) | = | g ′ 1 ( v/δ ) | ≤ pC 1 .

Revised: A bound on [ − ηr − δ, − ηr + δ ] is also obtained in the same way.

Original: s
Revised: Thus, we see [Equation] If there are additional stationary points of F , they should exist in [ ηr − δ, ηr + δ ] = supp( G ′ δ ) \ [ − ηr − δ, − ηr

Original: [Table]
Revised: + δ ] because of the sign of G ′ δ and supp( G ′ δ ) ⊂ ( −∞ , p/ 4) . However, since ηr + δ ≤ p/and pC 1 δ ηr ≤ p/ 2 under (10), we see [Equation] Hence, v ∗ = p is the unique local minimum (i.e., optimal solution) of F and we can conclude µ = 1 .

Original: [Table] [Table] .07 l [Table]
Revised: Estimation of c .

Original: [Table] [Table] m [Table]
Revised: From the above argument, we get

Original: [Table] SAM [Table] [Table]
Revised: F ′

Original: [Table]
Revised: [Equation]

Original: s
Revised: p

Original: [Table]
Revised: [Equation]

Original: AveragedSAM
Revised: =

Original: [Table]
Revised: [Equation]

Original: s
Revised: − p )

Original: [Table] [Table] [Table]
Revised: [Equation] ( else ) .

Original: 55 l
Revised: Clearly, p/ 2 ≤ 2( p − v ) / 3 for v ≤ ηr + δ ≤ p/ 4 .

Original: [Table] [Table]
Revised: Thus, F ′ ( v )( v − p ) ≥ ( v − p ) 2 / 3 on v ∈ [ ηr − δ, ηr + δ ] and we conclude c = 1 / 3 .

Original: [Table] m
Revised: Estimation of M .

Original: [Table]
Revised: Noting v ∗ = p and F ′′ ( p ) = 1 , we have

Original: [Table]
Revised: [Equation]

Original: Finally, we run SGD, SGD with large step size, and averaged SGD to train the standard convolutional neural network on Fashion MNIST dataset to confirm how efficiently sharpness and classification accuracy can be optimized by each method.{+verify M = 
+}
Revised: Because of the problem setup, it is enough to {+verify M = 
+}

Original: We note the large step size used for SGD is the same as that for averaged SGD. We plot the trace of Hessian ∇ 2 f ( w )  and test loss functions in Figure 5.
Revised: p satisfies | G ′ δ ( v ) | ≤ M | v − p | 2 on v ∈ [ ηr − δ, ηr + δ ] . Since | G ′ δ ( v ) | ≤ p/ 2 and v ≤ p/ 4 for v in this interval, we have [Equation]

Original: From this figure, we observe that the averaged SGD converges to a flatter region and achieves the highest classification accuracy on the test dataset as expected in our theory.
Revised: This concludes M = 89 p .