Strength of Minibatch Noise in SGD
Authors: Liu Ziyin
Kangqiao Liu
Takashi Mori
Masahito Ueda
Venue: ICLR 2022 Spotlight (2022)
Keywords: stochastic gradient descent
minibatch noise
discrete-time SGD
noise and fluctuation
exact solvable models

Abstract: The noise in stochastic gradient descent (SGD), caused by minibatch sampling, is poorly understood despite its practical importance in deep learning. This work presents the first systematic study of the SGD noise and fluctuations close to a local minimum. We first analyze the SGD noise in linear regression in detail and then derive a general formula for approximating SGD noise in different types of minima. For application, our results (1) provide insight into the stability of training a neural network, (2) suggest that a large learning rate can help generalization by introducing an implicit regularization, (3) explain why the linear learning rate-batchsize scaling law fails at a large learning rate or at a small batchsize and (4) can provide an understanding of how discrete-time nature of SGD affects the recently discovered power-law phenomenon of SGD.

Published as a conference paper at ICLR 2022
STRENGTH OF MINIBATCH NOISE IN SGD
Liu Ziyin', Kangqiao Liu", Takashi Mori, & Masahito Ueda
The University of Tokyo
ABSTRACT
The noise in stochastic gradient descent (SGD), caused by minibatch sampling, is poorly understood despite its practical importance in deep learning. This work presents the first systematic study of the SGD noise and fluctuations close to a local minimum. We first analyze the SGD noise in linear regression in detail and then derive a general formula for approximating SGD noise in different types of minima. For application, our results (1) provide insight into the stability of training neural network, (2) suggest that large learning rate can help generalization by introducing an implicit regularization, (3) explain why the linear learning rate- batchsize scaling law fails at large learning rate at small batchsize and (4) can provide an understanding of how discrete-time nature of SGD affects the recently discovered power-law phenomenon of SGD.
INTRODUCTION
Stochastic gradient descent (SGD) is the simple and efficient optimization algorithm behind the success of deep learning (Allen-Zhu et al., 2019; Xing et al., 2018; Zhang et al., 2018; Wang et al., 2020; He and Tao, 2020; Liu et al., 2021; Simsekli et al., 2019; Wu et al., 2020). Minibatch noise, also known as the SGD noise, is the primary type of noise in the learning dynamics of neural networks. Practically, minibatch noise is unavoidable because modern computer's memory is limited while the size of the datasets we use is large; this demands the dataset to be split into "minibatches" for training. At the same time, using minibatch is also recommended practice because using smaller batch size often leads to better generalization performance (Hoffer et al., 2017). Therefore, understanding minibatch noise in SGD has been one of the primary topics in deep learning theory. Dominantly many theoretical studies take two approximations:
(1) the continuous-ting time approximation, which takes the infinitesimal step-size limit; (2) the Hessian approximation, which assumes that the covariance matrix of the SGD noise is equal to the Hessian H. While these approximations have been shown to provide some qualitative understanding, the limitation of these approximations is not well understood. For example, it is still unsure when such approximations are valid, which hinders our capability to assess the correctness of the results obtained by approximations.
In this work, we fill this gap by deriving analytical formulac for discrete-time SGD with arbitrary learning rates and exact minibatch noise covariance. summary, the main contributions are: (I) derive the strength and the shape of the minibatch SGD noise in the cases where the noise for discrete-time SGD is analytically solvable; (2) we show that the SGD noise takes a different form in different kinds of minima and propose general and more accurate approximations. This work is organized as follows: Sec. 2 introduces the background. Sec. 3 discusses the related works. Sec. outlines our theoretical results. Sec. derives new approximation formulae for SGD noises. In Sec. 6, we show how our results can provide practical and theoretical insights to problems relevant to contemporary machine learning research. For reference, the relationship of this work to the previous works is shown in Table 1.
2 BACKGROUND
In this section, we introduce the minibatch SGD algorithm. Let (a,, 9: J=i be training set. We can define the gradient descent (GD) algorithm for a differentiable loss function L as w, = W;-1 AV L(w, ))), where À is the learning rate and w RRP is the weights of the model. We consider an additive loss function for applying the minibatch SGD.
Definition 1. A loss function L((x,, J.NN w) is additive if L(1xi NN
w)
N EN (11 Vi, w) for some differentiable, non-negative function E(.).

---

Published as a conference paper at ICLR 2022
Table I: Summary of related works on noise and stationary distribution of SGD. This work fills the gap of
the lack of theoretical results for the actual SGD dynamics, which is discrete-time and with minibatch noise.
Setting
Artificial Neise
Hessian Approsination Noise
Minthotok Noie
Comiinous time
Sato and Nakapawa (2015) Welling and Tih (20TT)
Jastrzebaki al, (2018; Zhu al. (2019)
Biane 12000) Moxi al. (202))
Mandi et (201 77, Meng (2020)
(2030); (2.21)
Discrete-time
Yada (2019); Cimma es al. (2019)
Liu et al. (2023)
This Work
Liu al. (202))
This definition is quite general. Most commonly studied and used loss functions are additive, e.g., the mean--quare error (MSE) and cross-entropy loss. For an additive loss, the minibatch SGD with momentum algorithm can be defined.
Definition 2. The minibatch SGD with momentum algorithm by sampling with replacement computes the update to the parameter with the following set of equations:
g+ EieB, V( Di, Wa-1);
m. = jum,11 + g:;
(1)
W. W Am,
where μ [0. is the momentum hyperparameter, S |Bal is the minibatch size, and the set B, = (is) are S i.i.d. random integers sampled uniformly from [1,N].
One can decompose the gradient into a deterministic plus a stochastic term. Note that Eg[gg] = VL is equal to the gradient for the GD algorithm. We use Eg()) to denote the expectation over batches, and use Ew (-) to denote the expectation over the stationary distribution of the model parameters. Therefore, we can write g. [g[[g]] Th, where T := 1 Eien, Vi(2i, Vi, Wr-1) EB [g:] is the noise term; the noise covariance is C( w.) := cov(n), nt). Of central importance to is the averaged asymptotic noise covariance C := lim fluctuation Σ := lim,-se cov(W,, w.). Σ gives the strength and shape of the fluctuation of w around a local minimum and is another quantity of central importance to this work. Throughout this work, C called the "noise" and ) the "fluctuation".
[m,[[C(()), Also, we consider the asymptotic model
3 RELATED WORKS
Noise and Fluctuation in SGD. Deep learning models are trained with SGD and its variants. To understand the parameter distribution in deep learning, one needs to understand the stationary distribution of SGD (Mandt et al., 2017). Sato and Nakagawa (2014) describes the stationary distribution of stochastic gradient Langevin dynamies using discrete-time Fokker-Planck equation. Yaida (2019) connects the covariance of parameter Σ to that of the noise C' through the fluctuation- dissipation theorem. When Σ is known, one may obtain by Laplace approximation the stationary distribution of the model parameter around a local minimum w² as N(w", Σ). Therefore, knowing E can be of great practical use. For example, it has been used to estimate the local minimum escape efficiency (Zhu et al., 2019; Liu et al., 2021) and argue that SGD prefers flatter minimum; it can also be used to assess parameter uncertainty and prediction uncertainty when a Bayesian prior is specified (Mandt et al., 2017; Gal and Ghahramani, 2016; Pearce et al., 2020). Empirically, both the fluctuation and the noise are known to crucially affect the generalization of a deep neural network. Wu et al. (2020) shows that the strength and shape of the Σ due to the minibatch noise lead to better generalization of neural networks in comparison to an artificially constructed noise.
Hessian Approximation of the Minibatch Noise. However, it is not yet known what form C' and E actually take for SGD in a realistic learning setting. Early attempts assume isotropic noise in the continuous-time limit (Sato and Nakagawa, 2014; Mandt al., 2017). In this setting, the noise is an isotropic Gaussian with C ~ ID and Σ is known to be proportional to the inverse Hessian H-1 More recently, the importance of noise structure was realized (Hoffer et al., 2017: Jastrzebski et al., 2018; Zhu et al., 2019; HaoChen et al.. 2020). "Hessian approximation" which assumee C' N HoHl for some unknown constant Co. has often been adopted for understanding SGD (see Table 1); this assumption is often motivated by the fact that C' = J,,. H, where Ju, is the Fisher information matrix (FIM) (Zhu et al., 2019); the fluctuation can be solved to be isotropic: Σ Lp. However, it is not known under what conditions the Hessian approximation is valid, while previous works have argued that it can be very inaccurate (Martens, 2014; Liu et al., 2021: Thomas et al., 2020; Kunstner et al., 2019). However, Martens (2014) and Kunstner et al. (2019) only focuses on the natural gradient descent (NGD) setting; Thomas et al. (2020) is closest to ours, but it does not apply to the case with momentum, matrix learning rate, or regularization.
2

---

Published as a conference paper at ICLR 2022
Discrete-time SGD with Large Learning Rate. Recently, it has been realized that networks trained at a large learning rate have dramatically better performance than networks trained with a vanishing learning rate (lazy training) (Chizat and Bach, 2018). Lewkowycz et al. (2020) shows that there is qualitative difference between the lazy training regime and the large learning rate regime; the performance features two plateaus in testing accuracy in the two regimes, with the large learning rate regime performing much better. However, the theory regarding discrete-time SGD at a large learning rate is almost non-existent, and it is also not known what Σ may be when the learning rate is non-vanishing. Our work also sheds light on the behavior of SGD at a large learning rate. Some other works also consider discrete-time SGD in similar setting (Fontaine al., 2021; Dieuleveut et al., 2020; Toulis et al., 2017), but the focus is not on deriving analytical formulae or does not deal with the stationary distribution.
SGD NOISE AND FLUCTUATION IN LINEAR REGRESSION
This section derives the shape and strength of SGD noise and fluctuation for linear regression; concurrent to our work, Kunin et al. (2021) also studies the same problem but with continuous-tima approximation; our result thus more general. To emphasize the message, we discuss the label noise case in more detail. The other situations also deserve detailed analysis; we delay such discussion to the appendix due to space constraints. Notation: denotes the minibatch size. w e RRD is the model parameter viewed in a vectorized form; À E I. denotes scalar learning rate; when the learning rate takes the form of a preconditioning matrix, we use A RR DDD E DDD denotes the covariance matrix of the input data. When matrix X is positive semi-definite, we write X 2 O; throughout, we require A 5 O. ? E R denotes the weight decay hyperparameter; when the weight decay hyperparameter is matrix, For two matrices X, Y, the commutator is defined as [X, Y] XY YX. Other notations are introduced in the context.' The results of this section are numerically verified in Appendix A.
write I' E DDD Ja is the momentum hyperparameter in SGD.
4.1 KEY PREVIOUS RESULTS
When N >> S, the following proposition is well-known and gives the exact noise due to minibatch sampling. See Appendix E.1 for derivation.
Proposition 1. The noise covariance of SGD as defined in Definition is
Σ (w) - VL(W)V((W)) C(w) = SN G-1
(2)
where the notations &(w):: 1x1; Vs, w) and L(w) (1 w) are used.
This gradient covariance matrix C is crucial to understand the minibatch noise. The standard literature often assumes C(w)+ H (w)); however, the following well-known proposition shows that this approximation can easily break down.
Proposition 2. Let w. be the solution such that I(ww)= = O, then C(w.) = 0.
Proof. Because & is non-negative for all i, I(w.) 0 implies that f,(w.) 0. The differentiability in turn implies that each (w.,w))= = O; therefore,
0.
This proposition implies that there is noise if our model can achieve zero training loss (which is achievable for an overparametrized model). This already suggests that the Hessian approximation C' ~ H is wrong since the Hessian is unlikely to vanish in any minimum. The fact that the noise strength vanishes at L = suggests that the SGD noise might at least be proportional to L(w), which will show to be true for many cases. The following theorem relates C' and E of the discrete-tima SGD algorithm with momentum for a matrix learning rate.
Theorem 1. (Liu et al., 2021) Consider rumning SGD on quadratic loss function with Hessian H, learning rate matrix A, momentum μ. Assuming ergodicity, then
++ 112
(1-1)(AHE+ ΣHA)
AHEHA+ 1 1-
(AHAHE+ EHAHA) = ACA.
(3)
Propostion and Theorem allow one solve C. and Σ. Equation (3) can be seen as a general form of the I Lyapunov equation (Lyapunov, 1992) and is hard to solve in general (Hammarling, 1982; Ye et al., 1998; Simoncíni, 2016). Solving this analytical equation in settings of machine learning relevance is one of the main technical contributions of this work.
We use the word global minimum to refer to the glohal minimum of the loss function, i.c., where I. 0 and local minimum refers to minimum that has non-negative loss, i.c., T. O.
3

---

Published as a conference paper at ICLR 2022
4.2 RANDOM NOISE IN THE LABEL
We first consider the case when the labels contain noise. The loss function takes the form
L(w)
w Iz - --))
(4)
2N
1
where Z; E IR. are drawn from a zero-mean Gaussian distribution with feature covariance := Eh [xxx]]] , and y u²x; + G,, for some fixed and G IR is drawn from distribution with zero mean and finite second momentum o². We redefine w u w and let N - co with D held fixed.
The following lemma finds C' function of E.
Lemma 1. (Covariance matrix for SGD noise in the label) Let N + and the model be updated according to Eq. (1) with loss function in Eq. (4). Then,
C 5 (AEAA T(
(5)
The model fluctuation can be obtained using this lemma.
Theorem 2. (Fluctuation of model parameters with random noise in the label) Let the assumptions be the same as in Lemma I and [A, A] = 0. Then,
o² (1 Kyg AG E
5155)
)
5
where K2 1- [TAG,, with G = -101-DDDD (t# )
THAGG,"1
Remark. This result is numerically validated in Appendix A. The subscript refers to momentum.
To obtuin results for vanilla SGD, one can set μ = 0, which has the effect of reducing Gy. +
21 (1+ }) AA. From now on, we focus on the case when μ 0 for notational simplicity, but we note that the results for momentum can be likewise studied. The assumption [A, A] is not too strong because this condition holds for a scalar learning rate and common second-order methods such as Newton's 's method.
If 02 = O, then 0. This means that when there no label noise, the model parameter has a vanishing stationary fluctuation, which corroborates Proposition 2. When scalar learning rate À << 1 and << S, we have
Aa Σ
ID.
2S
which is the result one would expect from the continuous- time theory with the Hessian approximation (Liu et al., 2021; Xie et al., 2021; Zhu et al., 2019), except for correction factor of σ². Therefore, Hessian approximation fails to account for the randomness in the data of strength o². We provide a systematic and detailed comparison with the Hessian approximation in Table 2 of Appendix B.
Moreover, it is worth comparing the exact result Theorem 2 with Eq. (7) in the regime of non- vanishing learning rate and small batch size. One notices two differences: enhancement, appearing in the matrix G, and taking the form -X(1 1 L/S) compared with the result in Liu et al. (2021), this term is due to the compound effect of using large learning rate and a small batchsize: (2) an isotropic enhancement term K, which causes the overall magnitude of fluctuations to increase; this term does not appear the previous works that are based on the Hessian approximation and is due to the minibatch sampling process alone. As the numerical example Appendix A shows, at large batch size, the discrete-time nature of SGD is the leading source of fluctuation; small batch size, the isotropic enhancement becomes the dominant source of fluctuation. Therefore, the minibatch sampling process causes two different kinds of enhancement to the fluctuation, potentially increasing the exploration power of SGD at initialization but reducing the convergence speed.
(I) an anisotropic
Now, combining Theorem 2 and Lemma 1, one obtain explicit form of the noise covariance. Theorem 3. The noise covariance matrix of minibatch SGD with random noise in the label is
C=
A+
1
Ki 5
(AAG,) + THAAG A.
52
4

---

Published as a conference paper at ICLR 2022
By definition, C is the FIM. The Hessian approximation, in sharp contrast, can only account for the term in orange. A significant modification containing both anisotropic and isotropic (up to Hessian) is required to fully understand SGD noise, even in this simple example. Additionally, comparing this result with the training loss (127), one can find that the noise covariance contains one term that is proportional to the training loss. In fact, will derive in Sec. 5 that containing a term proportional to training loss is a general feature of the SGD noise. We also study the case when the input is contaminated with noise. Interestingly, the result is the same with the label noise case with o² replaced by a more complicated term of the form Tr[AK-) BU]. We thus omit this part from the main text. detailed discussion can be found in Appendix E.3.1. In the next section, we study the effect of regularization on SGD noise and fluctuation.
4.3 LEARNING WITH REGULARIZATION
Now, we show that regularization also causes unique SGD noise. The loss function for I' L2 regularized linear regression is
N
[w J² + w Tw -w--u)TAw u) wwwT T... 2
(ww = 2.V 1-1
(9)
2
where F is symmetric matrix; conventionally, one set I = =TID with scalar y O. For conciseness, we assume that there is no noise in the label, namely yi u"ar with a constant vector u. One important quantity in this case will be uu := U. The noise for this form of regularization can be calculated but takes a complicated form.
Proposition 3. (Noise covariance matrix learning with L2 regularization) Let the algorithu be updated according to Eq. (1) on loss function (9) with N + co and [A,1] =0. Then,
1 (AEA+ + THAAJJ++ TH ATUJA+
C = 5
(10)
where A'
K-'T with K := A.
Notice that the last term PAUTT in C is unique to the regularization- noise: it is rank-1 because U rank-1. This term is due the mismatch between the regularization and the minimum of the original loss. Also, note that the term Tr[AE] is proportional to the training loss. Define the test loss to be Tuest == lim F, [ (w: T A(w; u)]. we can prove the following theorem. We will show that one intriguing feature of discrete-time SGD is that the weight decay can be negative.
Theorem 4. (Test loss and model fluctuation for la regularization) Let the assumptions be the same as in Proposition Then
X (Tr[AK +)) 1 TT
Liest
(1I)
25
where R:= TrAK 6 1
r:= THA K 9 -0]
T-3TI(ARK-1 G-1j with G 2In > > 3 A2). Moreover
let [T. U1 1- T-2K
then
An G (A ( Σ = T= T (1+
Ar _A)
(12)
This result is numerically validated in Appendix A. The test loss (11) has interesting consequence. One can show that there exist situations where the optimal is negative. When discussing the test loss, we make the convention that if w. diverges, then Lrest co.
Corollary 1. Le arg min, Ltest There exist a, X and S such that 7" <.
The proof shows that when the learning rate is sufficiently large, only negative weight decay is allowed. This agrees with the argument in Liu al. (2021) that discrete-time SGD introduces an implicit L2 regularization that favors small norm solutions. A too-large learning rate requires negative weight decay because a large learning rate already over-regularizes the model and one needs
2Some readers might argue that discussing test loss is meaningless when N + so; however, this criticism does not apply because the size of the training set is not the only factor that affects generalization. In fact, this section's 's crucial message is that using large learning rate affects the generalization by implicitly regularizing the model and, if one over-regularizes, one needs to offset this effect.
5

---

Published as a conference paper at ICLR 2022
to introduce an explicit negative weight decay offset this over-regularization effect of SGD. This is a piece of direct evidence that using a large learning rate can help regularize the models. has been hypothesized that the dynamics of SGD implicitly regularizes neural networks such that the training favors simpler solutions (Kalimeris et al., 2019). Our result suggests one new mechanism for such a regularization.
NOISE STRUCTURE FOR GENERIC SETTINGS
The results in the previous sections suggest that (I) the SGD noises differ for different kinds of situations, and (2) SGD noise contains term proportional to the training loss in general. These two facts motivate us to derive the noise covariance differently for different kinds of minima. Let f(w, 2) denote the output of the model for a given input I E RDD Here, we consider a more general case; f(w,x) x) may be any differentiable function, e.g., non-linear deep neural network. The number of parameters in the model is denoted by P. and hence w E RRP For training dataset (Zi, Mi) 1,,,,.... the loss function with a L2 regularization given by
1 Lr(w) Lo()) + w Tww, 2
(13)
where [o(w)= N E. ((S([, v.). the loss function without regularization, and Hn is the Hessian of Ly- We focus on the MSE loss {(f(w, y:) [[(w. .z;)) - =x12/2. Our result crucially relies on the following two assumptions, which relate to the conditions of different kinds of local minima.
Assumption 1. (Fluctuation decays with batch size) Σ is proportional to S- ,i.e. E= 0(5
This is justified by the results in all the related works (Liu et al., 2021; Xie al., 2021; Meng et al., 2020; Mori et al., 2021), where E is found to be O(S")).
Assumption 2. (Weak homogeneity) LL Es] is small; in particular, it is of order o L).
This assumption amounts assuming that the current training loss L reflects the actual level of approximation for each data point well. In fact, since L 2 O, one can easily show that IL. L O(L). Here, we require slightly stronger condition for a more clean expression, when 6 = can still get similar expression but with condition can be an important and interesting future work. The above two conditions allow us to state general theorem formally.
constant that hinders the clarity. Relaxing this
Theorem 5. Let the training loss be Lr Lo + the neighborhood of a local minimum w". Then,
+ W/ I'w and the models be optimized with SGD in
(Lo(w) C(w)= 5
Holw) VL - o(Lo).
(14)
The noise takes different forms for different kinds of local minima.
Corollary 2. Omitting the terms of order o( ([n)), when T # 0,
C= 2Lo(w")
1 S Ho(w")
www"w""T=(("" +(( w²)))
(15)
S
When r= =
and Lo ) #0.
2Lo(w")
C C= (lolw')+((57)) +(( w²
(16)
s
When -= =
and Lo Lo(w") (w*)= =0.
== (I= H Ho( (2)
(17)
Remark. Assumption 2 can be replaced by weaker but more technical assumption called the "decoupling assumption", which has been used in recent works to derive the contimuous time distri- bution of SGD (Mori et al., 2021; Wojtowytsch, 2021). The Hessian approximation was invoked in most of the literature without considering the conditions of its applicability (Jastrzebski et al., 2018; Zhu el al,, 2019; Liu et al., 2021; Wu al, 2020; Xie al., 2021). Our result does provide such

---

Published as a conference paper at ICLR 2022
conditions for applicability. As indicated by the two assumptions, this theorem is applicable when the batch size not too small and when the local minimum has a loss close to O. The reason for the failure of the Hessian approximation is that, while the FIM is equal to the expected Hessian J = R, there is no reason to expect the expected Hessian to be close to the actual Hessian of the minimum.
The proof is given in Appendix C. Two crucial messages this corollary delivers are (1) the SGD noise is different in strength and shape in different kinds of local minima and that they need to be analyzed differently; (2) the SGD noise contains a term that is proportional to the training loss Lo in general. Recently, it has been experimentally demonstrated that the SGD noise is indeed proportional to the training loss in realistic deep neural network settings, both when the loss function is MSE and cross-entropy (Mori et al., 2021); our result offers a theoretical justification. The previous works all treat all the minima as if the noise is similar (Jastrzebski et al., 2018; Zhu et al., 2019; Liu et al., 2021; Wu et al., 2020; Xie et al., 2021), which can lead inaccurate or even incorrect understanding. For example, Theorem 3.2 in Xie et al. (2021) predicts high escape probability from a sharp local or global minimum. However, this is incorrect because a model at a global minimum has zero probability of escaping due to a vanishing gradient. In contrast, the escape rate results derived in Mori et al. (2021) correctly differentiate the local and global minima. We also note that these general formulae are consistent with the exact solutions we obtained in the previous section than the Hessian approximation. For example, the dependence of the noise strength on the training loss in Theorem 2, and the rank-1 noise of regularization are all reflected in these formulae. In contrast, the simple Hessian approximation misses these crucial distinctions. Lastly, combining Theorem 5 with Theorem 1, one can also find the fluctuation.
Corollary 3. Let the noise be as in Theorem 5, and omit the terms of order O(S-2) and O(Iw w (²). Then, when I O and when A, Ho(w") and I commute with each other, P,
S I- (2 Iw"w w" (H) + I)* 2I
A I+p (H, + r)] When r and Lo Lo(w")
P.E = 5(1=a) o
2hm P.A (2l
T4p Ho) When I'= 0 and Lo (w") O, P.E = 0. Here the superscript pseudo inverse, P. diag(L, ...,, I, .....
, 0) is the projection operator with non-zero entries, S D is the rank of the Hessian Ho, and r' s D is the rank of Hu r. For the null space Ho, E can be arbitrary.
is the
APPLICATIONS
One major advantage of analytical solutions that they can be applied in simple "plug-in" manner by the practitioners or theorists to analyze new problems they encounter. In this section, we briefly outline a few examples where the proposed theories can be relevant.
6.1 HIGH-DDMENSIONAL. REGRESSION
We first apply our result to the high-dimensional regression problem and show how over-and- underparametrization might play a role in determining the minibatch noise. Here, we take N, D + oo
with the ratio a NNDD held fixed. The loss function is L(w)= a EN (wTx:; 1; vi)² As in the standard literature (Hastie et al., 2019), we assume the existence of label noise: y, = uu 2, + e,, with Var[e,] 0². A key difference between our setting and the standard high- high-dimensional setting is that, in the standard setting (Hastie et al., 2019), one uses the GD algorithm with vanishing learning rate À instead of the minibatch SGD algorithm with non-vanishing learning rate. Tackling the high-dimensional regression problem with non-vanishing À and a minibatch noise is another main technical contribution of this work. In this setting, we can obtain the following result on the noise covariance matrix.
Proposition 4. Let A= EN zız, and suppose assumptions 1 and 2 hold. With fixed S, then C= (TH[Æ] ÂΣ) À max
(0. (0. 2 (1 )) Å.
We note that this proposition follows from Theorem 5, showing an important theoretical application of our general theory. An interesting observation is that one E-independent term proportional to o² emerges in the underparametrized regime (a > 1). However, for the overparametriized regime, the noise is completely dependent Σ, which is sign that the stationary solution has fluctuation. This shows that the degree of underparametrization also plays a distinctive role in the fluctuation. In fact, one prove the following theorem, which is verified in Appendix A.2.
7

---

Published as a conference paper at ICLR 2022
10
0.5
0.01 A n.il
= 0.5/d
VorL1
varti
10-1
10-1
10-2
weggh strengtt
Figure I: Realistic learning settings with neural networks and logistic regression. Left: Variancc of training loss of neural network with width d and tanh activation on the MNIST dataset. We see that the variance explodes after 2 200. In contrast, rescaling the learning rate by 1/d results in a constant noise level in training. This suggests that the stability condition we derived for high-dimension regression is also useful for understanding deep learning. Middle: Stability of Adam with the same setting. Adam also experiences similar stability problem when the model width increases. Right: Logistic regression MNIST trained with SGD; with À 1.5, 32. We see that the optimal performance is also achieved at negative weight decay strength 7 suggesting that a large learning rate can indeed introduce effective regularization.
Theorem 6. When a stationary solution exists for w, we have TH(AE]- = max (o, Aa² 9 (1-2) k).
Tv(G-1 A] where R ITG TA] with G 2ID - X(1 - })) Å.
6.2 IMPLICATION FOR NEURAL NETWORK TRAINING
It is commonly believed that the high-dimensional linear regression problem can be minimal model for deep learning. Taking this stance, Theorem suggests a technique for training neural networks. For SGD to converge, positive semi-definite Σ must exist; however, Σ 2 0 if and only if & 0. From R O, have Σe 2/Aa:-1+1/5 S, where a; are the eigenvalues of Á. This means that each summand should have the order of D/S. Thus the upper bound of should have the order of 25/aD, where a is the typical value of asss s. One implication of the dependence the dimension is that the stability of a neural network trained with SGD may strongly depend on its width d, and one may rescale the learning rate according to the width stabilize neural network training. See Figure 1-Left and Middle. We train a two-layer tanh neural network on MNIST and plot the variance of its training loss in the first epoch with fixed X 0.5. We see that, when d 2 200, the training starts destabilize, and the training loss begins to fluctuate dramatically. When rescaling the learning rate by 1/d, we see that the variance of the training loss is successfully kept roughly constant across all d. This suggests a training technique worth being explored by practitioners in the field. In Figure 1 1-Middle, we also use Adam for training the same network and find a similar stabilizing trick work for Adam.
6.3 NATURAL LEARNING EXAMPLE WITH NEGATIVE WEIGHT DECAY
Sec. 4.3 shows that loo- learning rate introduces an effective L2 regularization that can be corrected by setting the weight decay to be negative. This effect can be observed in more realistic learning settings. We train logistic regressor on the MNIST dataset with large learning rate (of order ((1). Figure 1- R-Right confirms that, at a large learning rate, the optimal weight decay can indeed be negative. This agrees with our argument that using a large lcarning rate can effectively regularize the training.
6.4 SECOND-ORDER METHODS
Understanding stochastic second-order methods (including the adaptive gradient methods) is also important for deep learning (Agarwal et al., 2017; Zhang and Liu, 2021; Martens, 2014; Kunstner et al., 2019). In this section, Newton's 's method (DNM) and natural gradient descent (NGD). We provide more accurate results than those derived in Liu et al. (2021). The derivations are given in Appendix D.2. For DNM, the preconditioning learning rate matrix is defined as A := XA-1 The model fluctuation is shown to - +b)) X.
apply theory to two standard second-order methods: damped
be proportional to the inverse of the Hessian: Σ Ag² A-1, where g 2(1 2(1- - -) [I=H
gS-XD
The main difference with the previous results is that the fluctuation now depends explicitly on the dimension D, and implies a stability condition: S 2 >D/g, corroborating the stability condition we derived above. For NGD, the preconditioning matrix is defined by the inverse of the Fisher information that A == -(w)-1 3 c We show that Σ à solution when a = 0, which also contains a correction related to compared to the result Liu et al. (2021) which is Σ 2/2 Itp + I-p 1
2 I-D 1+p 1- 3/ A- is one
83 A-1, A consequence is that J ~ E-! The surprising
8

---

Published as a conference paper at ICLR 2022
fact is that the stability of both NGD and DNM now crucially depends on D; combining with the results in Sec. 6.1, this suggests that the dimension of the problem may crucially affect the stability and performance of the minibatch-bra based algorithms. This result also implies that some features derived are shared across many algorithms that depend on minibatch noise and that our results may be relevant broad class of optimization algorithms other than SGD.
6.5 FAILURE OF THE X 5 SCALING LAW
One well-known technique in deep learning training is that one can scale linearly as one increases the batch size S to achieve high high-efficiency training without hindering the generalization performance; however, is known that this scaling law fails when the learning rate is too large, or the batch size too small (Goyal et al., 2017). In Hoffer et al. (2017), this scaling law is established on the ground that X/S. However, our result in Theorem suggests the reason for the failure even for the simple setting of linear regression. Recall that the exact E takes the form:
Ao² E= S (1+ " )"
for a scalar A. One notices that the leading term is indeed proportional to X/S. However, the discrete-time tireteetimeeeee SGD results in second-order correction in S, and the term proportional to 1/5² does not contain a corresponding A; this explains the failure of the scaling law in small S, where the second-order contribution of S becomes significant. To understand the failure at large X, we need to look at the term G₁
G 2.1=2(1-))) - -2 A )
1
One notices that the second term contains a part that only depends on but not on S. This part is negligible compared to the first term when À small: however, becomes significant the second term approaches the first term. Therefore, increasing À changes this part of the fluctuation, and the scaling law no more holds if X large.
6.6 POWER LAW TAIL IN DISCRETE -IME TIME SGD
exp. proposeed theary Mcri(2021) Mcng(2020)
It has recently been discovered that the SGD noise causes a heavy- tail distribution (Simsekli et al., 2019; 2020), with a tail decaying like power law with tail index |8 (Hodgkinson and Mahoney, 2020). In continuous-time
the stationary distribution has been found to obey Student's 's t-like distribution, p(w) L-(18))22 ~
(o² + aa=)) a (Meng et al., 2020; Mori et al., 2021; Wo- jtowytsch, 2021). However, this result is only established for continuous-time approximations to SGD and one does not know what affects the exponent for discrete-time SGD. Our result Theorem 2 can serve as tool to find the discrete-time correction to the tail index of the stationary distribution. In Appendix D.3, we show that the tail index of discrete-time SGD in 1d can be estimated 25 - S. clear discrete-time contribution is -(S + 1) as (), = at which depends only on the batch size, while 25 1 is the tail index in the continuous citiiuuuus-tmme time limit (Mori et al., 2021). See Figure 2; the proposed at agrees with the experiment. Knowing the tail index 3 is important for understanding the SGD dynamics because 3 is equal to the smallest moment of w that diverges. For example, when 3 S 4, then the kurtosis of w diverges, and one expects to see outliers of w very often during training: when 3 < 2, then the second moment of w diverges, and one does not expect w to converge in the minimum under consideration. Our result suggests that the discrete-time dynamics always leads to heavier tail than the continuous-time theory expects, and therefore is more unstable.
0
50
200
150
Figure 2: Comparison of the pro- posed theory with the continuous- time theory on the SGD station- ary distribution for aA 1. The proposed theory agrees with the experiment exactly.
7 OUTLOOK
In this work, we have presented systematic analysis with focus exactly solvable results to promote our fundamental understanding of SGD. One major limitation is that we have only focused on studying the asymptotic behavior of SGD in local minimum. For example, Ziyin et al. (2022) showed that SGD can converge to a local maximum when the learning rate is large. One important future step is thus understand the SGD noise beyond strongly convex landscape.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGEMENT
Liu Ziyin thanks Jie Zhang, Junxia Wang, and Shoki Sugimoto. Ziyin is supported by the GSS Scholarship of The University of Tokyo. Kangqiao Liu was supported by the GSGC program of the University of Tokyo. This work was supported by KAKENHI Grant Numbers JP18H01 145 and JP21 from the Japan Society for the Promotion of Science.
REFERENCES
Agarwal, N., Bullins, B., and Hazan, E. (2017). Second-order stochastic optimization for machine learning in linear time. The Journal of Machine Learning Research, 18(1):4148-415)
Allen-Zhu, Z., Li, Y., and Song, Z. (2019). A convergence theory for deep learning via over- parameterization. In International Conference Machine Learning, pages 242-252. PMLR.
Amari, S.-I. (1998). Natural gradient works efficiently learning. Neural Comput., 10(2):251-276.
Blanc, G., Gupta, N., Valiant, G., and Valiant, P. (2020). Implicit regularization for deep neural networks driven by an munsten--uhlenbed like process. Conference on learning theory, pages 483-513. PMLR.
Chizat, L. and Bach, (2018). note on lazy training in supervised differentiable programming. arXiv preprint arXiv: 1812.07956, 8.
Clauset, A., Shalizi, C. R., and Newman, M. E. (2009). Power-law distributions in empirical data. SIAM review, 51 514))6611-733
Dieuleveut, A., Durmus, A., Bach, F. et al. (2020). Bridging the gap between constant step size stochastic gradient descent and markov chains. Annals of Statistics, 48(3):1348 1348-1382.
Fontaine, X., Bortoli, V. D., and Durmus, A. (2021). Convergence rates and approximation results for sgd and its continuous- counterpart.
Gal, Y. and Ghahramani, Z. (2016). Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learning, pages 1050-1059. PMLR.
Gitman, Lang, H., Zhang, P., and Xiao, L. (2019). Understanding the role of momentum in stochastic gradient methods. In Advances in Neural Information Processing Systems, pages 9633-9644
Goyal, P., Dollár, P., Girshick, R., Noordhuis, P., Wesolowski, L., Kyrola, A., Tulloch, A., Jia, Y., and He, K. (2017). Accurate, large minibatch sgd: Training imagenet in I hour. arXiv preprint arXiv:1706.02677
Hammarling, S. (1982). Numerical solution of the stable, non-negative definite lyapunov equation Yyapunov equation. IMA Journal of Numerical Analysis, 2(3):303-323,
HaoChen, J. Z., Wei, C., Lee, D., and Ma, T. (2020). Shape matters: Understanding the implicit bias of the noise covariance. arXiv preprint arXiv:2006.086800
Hastie, T., Montanari, A., Rosset, S., and Tibshirani, R. J. (2019). Surprises in high-dimensional ridgeless least squares interpolation. arXiv preprint 1903.08500.
He, F. and Tao, D. (2020). Recent advances in deep learning theory. arXiv preprint arXiv:2012.1093:
Hodgkinson, L. and Mahoney, M. W. (2020). Multiplicative noise and heavy tails in stochastic optimization. arXiv preprint arXiv:2006.0629.
Hoffer, E., Hubara, I., and Soudry, D. (2017). Train longer, generalize better: closing the generaliza- tion gap in large batch training of neural networks. In Advances in Neural Information Processing Systems, pages 1731-1741.
10

---

Published as a conference paper at ICLR 2022
Janssen, P. H. M. and Stoica, P. (1988). On the expectation of the product of four matrix-valuee gaussian random variables. IEEE Transactions on Automatic Control, 33(9):867-870
Jastrzebski, S., Kenton, Z. Arpit, D., Ballas, N., Fischer, A., Storkey, A., and Bengio, Y. (2018).
Three factors influencing minima SGD.
Kalimeris, D., Kaplun, G., Nakkiran, P., Edelman, B., Yang, T., Barak, B., and Zhang, H. (2019). Sgd on neural networks learns functions of increasing complexity. Advances in Neural Information Processing Systems, 32:3496-3506.
Kunin, D., Sagastuy-Brena J., Gillespie, L., Margalit, E.,, Tanaka, H., Ganguli, S., and Yamins, D. L. (2021). Rethinking the limiting dynamics of sgd: modified loss, phase space oscillations, and anomalous diffusion. arXiv preprint 1XXv:2107.09133.
Kunstner, F., Balles, L., and Hennig, P. (2019). Limitations of the empirical fisher approximation for natural gradient descent. arXiv preprint arXiv: /905. /2558.
Levy, M. and Solomon, S. (1996). Power laws are logarithmic boltzmann laws. Internationaa Journal of Modern Physics C, 7(04))595-601
Lewkowycz, A., Bahri, Y., Dyer, E., Sohl-Dickstein, J., and Gur-Ari, G. (2020). The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.022/8.
Liu, K., Ziyin, L., and Ueda, M. (2021). Noise and fluctuation of finite learning rate stochastic gradient descent. arXiv preprint arXiv:2012.03666.
Lyapumov, A. M. (1992). The general problem of the stability of motion. International journal of control, 55(3):531-534
Mandt, S., Hoffmmn, M. D., and Blei, D. M. (2017). Stochastic gradient descent as approximate bayesian inference. J. Mach. Learn. Res., 18(1):88734907.
Martens, J. (2014). New insights and perspectives on the natural gradient method. cite 141i 199Comment: New title and abstract. Added multiple sections, including proper introdaction/outline and one convergence speed. Many other revisions throughout.
Meng, Q., Gong, S., Chen, W., Ma, Z.-M., and Liu, T.-Y. (2020). Dynamic of stochastic gradient descent with state-dependent noise. arXiv preprint arXiv: 13719.
Mori, T., Ziyin, L., Liu, K., and Ueda, M. (2021). Logarithmic landscape and power-law escape rate of sgd. arXiv preprint arX1v:2105.09557
Pearce, T., Leibfried, F,, and Brintrup, A. (2020). Uncertainty neural networks: Approximately bayesian ensembling. In International conference artificial intelligence and statistics, pages 234-244. PMLR.
Sato, I. and Nakagawa, H. (2014). Approximation analysis of stochastic gradient langevin dynamics by using fokker-planck equation and ito process. International Conference Machine Learning, pages 982-990. PMLR.
Simoncini, V. (2016). Computational methods for linear matrix equations. SIAM Review, 58(3):377~ 441.
Simsekli, U., Sagun, I.., and Gurbuzbalaban, M. (2019). A tail-index analysis of stochastic gradient noise in deep neural networks. In International Conference on Machine Learning, pages 5827- 5837. PMLR.
Simsekli, U., Sener, O., Deligiannidis, G., and Erdogdu, M. A. (2020). Hausdorff dimension, heavy tails, and generalization in neural networks. Advances in Neural Information Processing Systems, 33.
Thomas, V., Pedregosa. F., Merriënboer, B., Manzagol, P.-A., Bengio, Y., and Le Roux, N. (2020). On the interplay between noise and curvature and its effect on optimization and generalization. In International Conference Artificial Intelligence and Statisties, pages 3503-3513. PMLR.
11

---

Published as a conference paper at ICLR 2022
Toulis, P., Airoldi, E. M., et al. (2017). Asymptotic and finite-sample properties of estimators based on stochastic gradients. Annals of Statistics, 45(44:1694-1727.
Wang, X., Zhao, Y., and Pourpanah, F. (2020). Recent advances in deep learning. International Journal of Machine Learning and Cybernetics, 1(44)7777770.
Welling, M. and Teh, Y. W. (2011). Bayesian learning via stochastic gradient langevin dynamics, In Proceedings of the 28th international conference on machine learning (ICML-II), pages 681-688. Citeseer.
Wojtowytsch, S. (2021). Stochastic gradient descent with noise of machine learning type. part ii: Continuous time analysis. arXiv preprint arXiv:2/6600288.
Wu, J., Hu, W., Xiong, H., Huan, J., Braverman, V., and Zhu, Z. (2020). On the noisy gradient descent that generalizes as sgd. In International Conference on Machine Learning, pages 10367-10376. PMLR.
Xie, Z., Sato, I., and Sugiyama, M. (2021). diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. In International Conference Learning Representations.
Xing. C., Arpit, D., Tsirigotis, C., and Bengio, Y. (2018). A walk with sgd. arXiv preprint arXiv:1802.08770.
Yaida, S. (2019). Fluctuation-dissipation relations for stochastic gradient descent. In International Conference on Learning Representations.
Ye, H., Michel, A. N., and Hou, L. (1998). Stability theory for hybrid dynamical systems. IEEE transactions on automatic control, 43(4):461-474
Zhang, C., Liao, Q. Rakhlin, A., Miranda, B., Golowich, N., and Poggio. T. (2018). Theory of deep learning iib: Optimization properties of sgd. arXiv preprint arXiv: arXiv:1801.000 1801.02254.
Zhang, Z. and Liu, Z. (2021). On the distributional properties of adaptive gradients. In Uncertainty in Artificial Intelligence, pages 419-429. PMLR.
Zhu, Z., Wu, J., Yu, B., Wu, L., and Ma, J. (2019). The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. In International Conference Machine Leurning, pages 7654-7663. PMLR.
Ziyin, L., Li, B., Simon, J. B., and Ueda, M. (2022). SGD converge to local maxima. In International Conference on Learning Representations.
12

---

Published as a conference paper at ICLR 2022
0.40
experiment
1.0
0.35
expertment
Exacl solution
exact solution
cts-time hessian approx.
0.8
cts-time hessian spprox.
Varte]
0.25
hessian approx.
tost loss
0.6
time hessian approx.
0.20
0.15
0.4
0.10
0.05
0.9
5
20
40
-1.0
-0.5
0.0
0.5
1.0
1.5
2.0
2.5
3.0
Figure Left: experiments with label noise. The parameters are set to be 1.5 and X Right: Experiments with Ly regularization with weight decay strength The parameters are set to be a = 0.5, 1. This the standard case with a vanishing optimal yy The vertical lines show where our theory predicts divergence.
experiment proposed theury
cto time h-apprxx
10-
Vedoll
d-time IL 2021)
0.5
0.2
-1.
CmC REDPOS c-dime h-sppro IU) 20211 prososed theary
-1.
0.00
0.25
0.50
0.75
1.00
1.50 1755
2.00
-25
(a) a= 1, 10
(b)
(c) 10
Figure 4: Comparison between theoretical predictions and experiments. (a) 1d experiment. We plot Σ as increasing function of X. We that the continuous- approximation fails to predict the divergence at a learning rate and the prediction Liu et al. (2021) severely underestimates the model fluctuation. In contrast, experiments. The straight line shows where the proposed theory predicts divergence in the variance, which agrees with experiment exactly. The Hessian has eigenvalues and 0.5, and À 1.5. For large batch size, the discrete-time Hessian approximation is quite accurate; for small S, the Hessian approximation underestimates the overall strength of the fluctuation. In contrast, the continuous-time result is both inaccurate in shape and in strength.
result is accurate throughout the entire range of learning rates. (b)-(c) 2d
A EXPERIMENTS
A.1 LABEL NOISE AND REGULARIZATION
Theorem 2 can be verified empirically. We run 1d experiment in Figure 4(a) and high dimensional experiments in Figures 4(b)-(c), where choose D = 2 for visualization. We that the continuous Hessian approximation fails badly for both large and small batch sizes. When the batch size is large, both the discrete-time Hessian approximation and our solution give accurate estimate of the shape and the spread of the distribution. This suggests that when the batch size is large, discreteness the determining factor of the fluctuation. When the batch size is small, the discrete Hessian approximation severely underestimates the strength of the noise. This reflects the fact that the isotropic noise enhancement is dominant at a small batch size.
In Figure 3-Leftt we run Id experiment with X = 1, N 10000 and 0² 0.25. Comparing the predicted Σ, we see that the proposed theory agrees with the experiment across all ranges of The continuous theory with the Hessian approximation fails almost everywhere, while the recently proposed discrete theory with the Hessian approximation underestimates the fluctuation when is small. In Figure 3-gight, we plot standard case where the optimal regularization strength is vanishing.
Now, validate the existence of the optimal negative weight decay predicted by formula. For illustration, we plot in Figure 5 the test loss (11) for a 1d example while varying either S X. The orange vertical lines show the place where the theory predicts divergence in the test loss. We also
13

---

Published as a conference paper at ICLR 2022
201
experiment RE
precosed sheery
bost 1os
102
dtscime h-spprox
u-ame hasspror. U.u 2021)
304
Figure 5: 1d experiments with L2 regularization with weight decay strength y. The parameters are set to be a = 4, À 1, S 64. This shows case where the optimal is negative. The vertical lines show where our theory predicts divergence.
0.15
TAAII
0.10
exp.
0.05
proposed thecry
apprax.
0.00
Figure 6: High-dimensional linear regression. We see that the predicted fluctuation coefficient agrees with the experiment well. The slight deviation is due a finite training time and finite N and D. the other hand, naive Hessian approximation results in qualitatively wrong result.
plot a standard case where optimal is close to in Appendix A. Also, we note that the proposed theory agrees better with the experiment.
A.2 HIGH- NIGH-DIMENSIONAL REGRESSION
See Figure 6-Left. We vary N with D = 1000 held fixed. We set À = 0.01 and S 32. We see that the agreement between the theory and experiment is good, even for this modest dimension number D. The vertical line shows where the over-to-une sarerouunddrprraeetizttion transition takes place. As expected, there is no fluctuation when 1, and the fluctuation gradually increases as a oo. On the other hand, the Hessian approximation gives wrong picture, predicting fluctuation to rise when there is no fluctuation and predicting a constant fluctuation just when the fluctuation starts to rise.
14

---

Published as a conference paper at ICLR 2022
Table Comparison with previous results. For notational conciseness, compare the case when all the relevant matrices commute. The model fluctuation Σ, the expected training loss Ltrain and the expected test loss Lust calculated by continuous- and discrete- time theories with Hessian approximation C' s H are presented. Exact solutions to these quantities obtained in the present work are shown the rightmost column.
Hessian Approximation
Exact Solation
Cts-time Approximation
D-time Solution
This Work
Σ
Y }(2lp - AA)
>>
Label Noise
lo
22 A(1+ 4]
Input Noise
II
$(210 - AK)
AR
Au 1200 - (1+ }]
lg Regularization
lD
3(210 -K)
Eq. (12)
Larnin
Luzin
Larsin
Label Noise
23
[M[A] + 30²²
Eq. (20)
(1 4))
Input Noise
[T[]+ [IAA BU] T[[] [TAAK'T
Eq. (28)
TT(AK BBU (x)
Lg Regularization
Eq. (37)
Eq. (151)
Lawe
Label Noise
4s TH[
Tr [ 1]
Input Noise
TH[A++ + TT AB'U)
Eq. (29)
A THAKK TTTB"T Eq. (11)
L2 Regularization
T+|A++ TTA 2r
Eq. (38)
B COMPARISON WITH CONVENTIONAL HESSIAN APPROXIMATION
We compare our results for the three cases with the results obtained with the conventional Hessian approximation of the noise covariance, i.e., C H, where H is the Hessian of the loss function. We summarize the analytical results for special case in Table 2.
B.1
LABEL NOISE
We first consider discrete-ting time dynamics with the Hessian approximation. The matrix equation is
A
ΣA + AΣ XAEA
A.
(18)
S
Compared with the exact result (3), is a large-S limit up to the constant 2². This constant factor is ignored during the approximation that J(w)= := Egl [ Eg[vv" == H (w), which is exact only when 1({2;), w) is a negative log likelihood function of w. Solving the matrix equation yields
Σ= S
(2 -
(19)
The training loss and the test loss are
Lerain X Tr[A(2Ip XA)-
2
(20)
25
X
Luest
Tr[AA(2ID - XA)
(21)
25
On the other hand, by taking the large-S limit directly from the exact equation (3), the factor o² is present:
ΣA+ AE
-AEA= of
(22)
5
For the continuous-time limit with the Hessian approximation, the matrix equation is
X
ΣA 1Σ
S A,
(23)
which is the small-A limit up to the factor 0². The variance is
X
Σ= =
Ip.
(24)
28
15

---

Published as a conference paper at ICLR 2022
The training and the test error
1
Lirain
[A]
=
(25)
45
2
Loss
Ir[A].
(26)
45
Again, taking the small-A limit directly from the exact result (3) shows the presence of the factor o² on the right hand side of the matrix equation.
B.2 INPUT NOISE
The case with the input noise similar to the label noise. This can be understood if we replace A by K and o² by Tr[AK BU]. The model parameter variance resulting from the discrete-time dynamics under the Hessian approximation is
X Σ=
(21D
(27)
S
The training and the test error
X
1 TM[AK- BU], 1 TT AB'U].
Lerain
(28)
25
Ltest
(T(A(21D-AK]])
(29)
25
The large-S limit from the exact matrix equation (144) results in prefactor Tr[AK- BU] in the fluctuation:
== -U
(30)
For the continuous-time limit, we take A + 0. The Hessian approximation gives
ID.
(31)
28
Ltrain
Tr[K] [[[AK
(32)
4S
2
Ltest
4S
Tr[A]+ TT[BTABBU)]
2
AB'U].
(33)
The large-S limit again produces prefactor Tr[AK BUJ.
B.3 La REGULARIZATION
For learning with regularization, there a more difference between the Hessian approximation and the limit taken directly from the exact theory. We first adopt the Hessian approximation for the discrete-time dynamics. The matrix equation is
ΣK KE YKEK X
K,
(34)
S
which is similar the previous subsection. However, is different from the large-S limit of the exact matrix equation (154):
ΣK + K K (TH(A-
(35)
S
This significant difference suggests that the conventional Fisher-to-Hessian approximation J H fails badly. The fluctuation, the training loss, and the test loss with the Hessian approximation are
Σ
-2T
(36)
S
TT[K(2TD (K) 1 TT[AK-TUUJ,
Ltrain
(37)
25
Liest
TT[A(2ID 1 TT[[K-2r'up
(38)
25
16

---

Published as a conference paper at ICLR 2022
while the large-S limit of the exact theory yields
3 5 (PUJA
-
(39)
TArain
25
25
+ $ TTT[AK]TU]]
(40)
Liest
25
28
2 TTITAAK
(41)
The continuous-time results are obtained by taking the small-A limit on Eqs. (36)-(38) for the Hessian approximation and on Eqs. (39)-(41) for the limiting cases of the exact theory. Specifically, for the Hessian approximation, we have
== JD.
(42)
28
A
Ltrain
[r[K]+ TT[AAK
(43)
48
1
Luest
Tr[A]
(44)
4S
The small-A limit of the exact theory yields
A ==
A
A
(45)
28
25
1
Lirain
TMAAK T
Tr[
U + [HAAA--UU)
(46)
45
45
2
Liest
31
(47)
4S
45
17

---

Published as a conference paper at ICLR 2022
C PROOF OF THE GENERAL FORMULA
C.I PROOF OF THEOREMS 5 AND COROLLARY
We restate the theorem.
Theorem 7. Let the training loss be Lr Lo 1 Tw and the models be optimized with SGD in the neighborhood of a local minimum w". When T + 0. the noise covariance is given by
C= 2Lo(w")
Ho(w") 3 J. T+ + +WWWWwW11 w w3))
5
(48)
S
When 1 0 and Lo(w) (w")
(**)+² * #0.
2Lo +0(5+)++++++(())))
C=
(49)
s
When D = 0 and Lo (w")=0. =0
C= == (T+(HH((w"'" --Ho" Ho( EZ) + O([w- - -w²)).
(50)
Proof. We will use the following shorthand notations: 4== ((S([(WW)))))).. vi), Eg of 0 := 8%
37
The Hessian of the loss function without regularization Ho(w)= = VV Lo (w) is given by
Ho(w)= 1 (V
1 N VVvVs(w,
(51)
The last term of Eq. (51) can be ignored when Lo 1. since
1 N
1/2
N
1/2
Lvv T f(w, (w)
(VVTf(w,x:)
N Σ (e))
VV T f(w, f(w,
$
N
103/1/22
VV
xz) 12
(2l
Σ v [V f(w,,;) x:) 1B
where || ||r stands for the Frobenius norm', and we have defined the variable (12) := N E
Since E = 1 for the mean-square error, we obtain
Σ Ho(w)= = N
O +
(52)
near a minimum. The Hessian with regularization Hr(w): "VTI.r(w) is just given by Ho(w) + r.
On the other hand, the SGD noise covariance C(w) is given by Eq. (2). By assumption 2, the SGD noise covariance is directly related to the Hessian:
C(w)
v W/(w,x:)V/(wwa) -
)
Σ
N
1 (L((W)
SN
2 N SN
(l: Lo) (ww:))///..::://....)"
2 w)
1 vLr +
Ho (w)
(53)
S
This finishes the proof.
IIn the linear regression problem, the last term of Eq. (51) does not exist since VV f(w, =0
18

---

Published as a conference paper at ICLR 2022
Now prove Corollary 2.
Proof. Near a minimum w' the full loss Lr(w), we have
vLr(w) Ho (w w') Iw" O([W- - w P
(54)
within the approximation (r(w) = LI (L(www)) Equations (14) and (54) give the SGD noise covariance near minimum of Lr(w).
(w-w") )² y + O( w w 1²)).
Now is worth discussing two different cases separately: (1) with regularization and (2) without regularization. We first discuss the not small enough, and the SGD noise covariance is not proportional to the Hessian. Near a local or global minimum w we obtain
when regularization is present. In this case, the regularization
the first term of the right-hand side of Eq. (54) is negligible, and hence
Ew[C(w)] 2Lo(w") Ho(w")- "w'' T 5
s
Ew Howw")(w- (om--w)(e w')(ww w' (w Ho O(w w*l))
[s
(ww(w") Holw")- Tw" w T 0 + O(Ww-O('')) w
(55)
s
S
where we have used the fact that =w". The SGD noise does vanish even global minimum of Lr(w). Note that this also agrees with the exact result derived in Sec. 4.3: together with an anisotropic noise that proportional to the Hessian, a rank-1 noise proportional to the strength of the regularization appears. This rank-1 noise is signature of regularization.
w"
On the other hand, as we will see below, the SGD noise covariance proportional to the Hessian near minimum when there is no regularization, i.e., P O. We have
C(w)= 2Lo(w w) Ho(w)- Ho Ho(w")(ww w (w """)"Ho('"" O(I w*l?). S
1
(56)
For this case, we need to differentiate between a local minimum and a global minimum. When Lo (w") ) is not small enough (e.g. at local but not global minimum).
C(w)= 2Lo(w*)
Ho(w) (w - ) Ho (w (w w Ho(w)
S
S Hoo(")((w (ww")))www>)) -o:-ww') "Ho'"" + O(I - (²)
2L (w")
Ho(w)+0(C) + O(w - w²13)
S
2o((w")
Ho ( O(ww w²1²),
(57)
and so, leading order,
C= 2Lo(w") Ho(w"),
(58)
S
which is proportional to the Hessian but also proportional the achievable approximation error. On the other hand, when Lo(w") is vanishingly small (e.g. at a global minimum), we have 2L (w) (w- w*) Ho (w") (w w"), and thus obtain
C(w)= [ - (ww-") T (w w² "oo""" Ho(w") (o(w")(w-w")(mm) w²) Ho(
O(S + O(lw-O(www w?
(59)
i.e.
E [ = (Tr[HeE]4 HoΣ) Hy 0(S-)) + +0([-07]])) (3))
(60)
This completes the proof.
19

---

Remark. It should be noted that the second term the right-hand side of Eq. (59) would rypically be much smaller than the first term for large D. For example. when Ho( w alp with a > 0, the first and the second terms are respectively given by (a² [S) ww and The Frobenius norm of the former is given by Da [S][ wwwwwwll I while that of the latter is given Published as a conference paper at ICLR 2022
(w-w"" w (w-w")
by (a²|S)[ w w" w=I,, which indicates that in Eq. (59), the first term dominant over the second term for large D. Therefore the second term of Eq. (59) can be dropped for large D, and Eq. (59) is simplified as
C(w) (w-w') Holw" (w-w') Ho(w"); [[C]
(61)
Ho.
Again, the SGD noise covariance is proportional the Hessian.
In conclusion, long as the regularization is small enough, that the SGD noise covariance near a minimum proportional to the Hessian a good approximation. This implies that the noise is multiplicative, which is known to lead to heavy tail distribution (Clauset et al., 2009; Levy and Solomon, 1996). Thus, we have studied the nature of the minibatch SGD noise three different situations. As example, we have demonstrated the power of this general formulation by applying it to the high-dimensional linear regression problem in Sec. 6.1.
C.2 PROOF OF COROLLARY 3
Proof. We prove the case where I and L(w") 0 as an example. Substituting Theorem into Theorem yyells
1 AHo AHOΣ 5(1-
2Lo
2ID
Ho.
(62)
1+
where we have assumed necessary commutation relations. Suppose that the Hessian Ho is of rank-r with s D. The singular-value decomposition and its Moore-Penrose pseudo inverse are given by Ho USV' and Hi VS+ respectively, where U and V are unitary, S is a rank-r diagonal matrix with elements being singular values of Hy. and S* is obtained by inverting every non-zero entry of S. Multiplying Hn to both sides of the above equation, we have
2Lg P.E = 5(1-w)
P.A 2ID
Hu
(63)
1+ P
where P. = diag(1, 1,..,,,,1111. 1.0.... ................................ 0) is the projection operator with r non-zero entries. When the Hessian full-rank, i.e., r D, the Moore-Penreee pseudo inverse is nothing but the usual inverse. The other cases can be calculated similarly.
20

---

Published as a conference paper at ICLR 2022
D APPLICATIONS
D.1 INFINITE-DIMNNIINALL LIMIT OF THE LINEAR REGRESSION PROBLEM
Now we apply the general theory in Sec. to linear regressions in the high-dimensional limit, namely
N, D oo with a == N/D held fixed.
D.1.1 PROOF OF PROPOSITION
The loss function
I(w)=
2N Σ w T
xr yi )²
(64)
with yt uf te be written
L(w) (w- (w-u- (W---AA - At -
(65)
2N F=1
47
where ;= 1
N Eis1 TęI; is an empirical covariance for the training data and v N Σ J36. The symbol (.)* denotes the Moore-Penrose pseudoinverse. We also introduce the the averaged traing
loss: Lirain := Ew [W[]]
The minimum of the loss function is given by
w' u Á'v IIr,
(66)
where RDD an arbitrary vector and is the projection onto the null space of Å. Since 1 IT= A Å, w" is also expressed as
w A (Au+ v) Ir..
(67)
In an underparameterizeed regime a > 11 = almost surely holds as long as the minimum eigenvalue of A (not A) is positive (Hastie et al., 2019). In this case, Å ' and we obtain
11
w u A- for > 1.
(68)
On the other hand, in an overparameterized regime a 1, II # 0 and there are infinitely many global minima. the ridgeless regression, we consider the global minimum that has the minimum norm [[w]], which corresponds to
w' AAu + v) (1- II-II)+++0 Â'v for ridgeless regression with a < <1
(69)
In both cases, the loss function is expressed
L(w) 2 (w-w") A(w -ww v Áv
(70)
2N
Asymptotically, w, converges to stationary point w with fluctuation Σ obeying the following equation (Theorem I:
XĀΣ PEA - - X²c.
(71)
The SGD noise covariance C is given by Eq. (14). In the present case, the Hessian is given by H = and we also have
À
2 N
Σ (w" x4 - 21)²
Σ 2L(w).
(72)
N
On the other hand, A VL(W)VL(w)) A( w (w- w )T Å, and hence Ew [VVL(W)VL(W))) ÁΣA. Therefore we obtain
2Lerain A AΣA.
C C=E [[W[W[
(73)
S 5
21

---

Published as a conference paper at ICLR 2022
Now, find Lerain First, define X RNND as Xsx (2t)s, and غ RR as E cz. Then w IJu Åtv (1 II)u+ (XIX)+XTX X
With this notation, we have A= X and the loss function expressed as
L(w) ( (w- (ww))'(w-w'') w")
x X)
1
x(x
(74)
2N
2N 1-1
We therefore obtain
1 [=[
o² 2
Lirain マ
El
x)
x c]
(75)
2N
Here,
Bi We can prove that the following identity is almost surely satisfied (Hastie et al., 2019) as long as the smallest eigenvalue of A (not A) is positive:
x)=]] Tr(1 IT).
(76)
Tr(1 - 11) mi,
(77)
We therefore obtain
1
-
Tre AS.
0² for a >
1 T[[A
o²
a²
2
+
2
Ltrain
min D, .N
(78)
2N
2
Tr[AE]
Tr
for a s
By substituting Eq. (78) into Eq. (73), we obtain the following SGD noise covariance:
ΛΣ) A a² (1 2)² for a 1,
Tr [T[AE]]
s
C=
(79)
1
[Tr[AE]- ΛΣ) A S
for a
This finishes the proof. D
D.1.2 PROOF THEOREM 6
Proof. We have to solve this equation:
ÅΣ ΣÅ =
(80)
where C' is given in Proposition 4. Using the similar trick of multiplying by G 2lp as in Appendix E.2.2, one obtains
A
TM(AT]
2g² (1- 2) h
for a > 1; for s 1,
(81)
0
where & 1-
T-66-1 A]
with = 2In X(1 )
Substituting the above trace into the matrix equation, we have
5 (1 21+ G-1 for a > 1; for a $ 1.
(82)
D.2 SECOND-ORDER METHODS
Proposition 5. Suppose that run DNM with A= XA with random noise in the label. The model
fluctuation is
2a2
E=
g$ AD A
(83)
where g 2(1 - -1) - I- Itp
3)) A.
22

---

Published as a conference paper at ICLR 2022
Proof. Substituting A XA- into Eqs. (3) and (5) yields
gD (Tr[A>] + A
(84)
where g === 2(1-
1- + 1 X. Multiplying and taking trace on both sides, we have I+p
ADo²
Tr[AE]= 9S - AD
(85)
Therefore, the model fluctuation is
Σ Ao
A-1
(86)
g$ AD
Proposition Suppose that we run NGD with [[ label. The model fluctuation is
c-1 with random noise in the
A
0²
X²g² + -4
2
1 02
0² 2
4
21+D
1 1++1++/
1
+
1 D
A
(87)
where g == 1+D 1+m +
11
I-#
Proof. Similarly to the previous case, the matrix equation satisfied by E is
1+ > C-l AEAC" P "I (C AC AE+ EAC EAC =
A cor
(1-
AE+ AE+ AE+EAC") 1-
(88)
Although it is obvious how directly solve this equation, it is possible to guess solution according to the hope that Σ be proportional to _/ in turn, A (Amari, 1998; Liu et al.. 2021).
We assume that > -A-1 and substitute into the above equation to solve for r. This yields one solution without claiming its uniqueness. By simple algebra, this x is solved to be
1
82g+42 I+DI+H
a2
2
x 7
I+D +
(89)
21+
1+
Let 0. We obtain the result in Sec. 6.4.
D.3 ESTIMATION OF TAIL INDEX
In Mori et al. (2021); Meng et al. (2020), it is shown that the (Id) discrete-time SGD results in distribution that is similar Student's 's t-distribution:
p(w) (a² aw?
(90)
where o² is the degree of noise in the label, and is the local curvature the minimum. For large w, this distribution is power-law distribution with tail index:
p(w)) | (1+/6)
(91)
and it is not hard to check that 9 also equal to the smallest moment of w that diverges: E[w"1 co. Therefore, estimating 3 can of great use both empirically and theoretically.
25 In continuous-time, it is found that Bats 2 (Mori et al., 2021). For discrete-time SGD, hypothesize that the discrete-time nature causes change in the tail index 3 = Bota e, and we are interested in finding e. We propose "semi-continuous" approximation give the formula to estimate the tail index. Notice that Theorem 2 gives the variance of the discrete-tima SGD, while Eq. (90) can be integrated to give another value of the variance, and the two expressions must be equal for consistency. This gives an equation that B must satisfy:
f (w (w El 2 Var[w],
(92)
23

---

Published as a conference paper at ICLR 2022
discrete-time correction
exp.
1750
Mori (2021)/ Mang (2020)
discrete correction
1500
300
150
Mon /20211 Meng (2020)
1030
100
750
100
250
discrete-time serrete-tiimeeee correction
Mon /2021) Meng (2020)
0
0
130
50
100
150
50
100
150
5
$
Figure 7: Tail index 3 the stationary distribution of SGD in a 1d linear regression problem. Left to Right: aA 0.2, 1.0, 1.8.
this procedure gives the following formula:
25 aA
B(,
Bets +,
(93)
and one immediately recognizes that + +55+ 1) is the discrete-time contribution to the tail index. See Figure for additional experiments. We see that the proposed formula agrees with the experimentally measured value of the tail index for all ranges of the learning rate, while the result of Mori et al. (2021) only correct when 0*. Hodgkinson and Mahoney (2020) also studies the tail exponent of discrete-time SGD; however, their conclusion is only that the "index decreases with the learning rate and increases with the batch size". In contrast, our result give the functional form of the tail index directly. In fact, this is the first work that gives any functional form for the tail index of discrete-time SGD fluctuation to the best of our knowledge.
The following proposition gives the intermediate steps in the calculation.
Proposition 7. (Tail index estimation for discrete-time SGD) Let the parameter distribution be 1+3
p(w)- (o² aw²) and Var[w]] be given by Theorem 2. Then
25 S.
B(A
(94)
aA
Proof. The normalization factor for the distribution exists if > 0:
1+8 2
N
(95)
π
r
If [ > 2, the variance exists and the value is
Var[w]=
(96)
a(B
By equating Eq. (96) with the exact variance (6), we are able to solve for an expression of the tail
index
25 S.
3(A,S) aA
(97)
24

---

Published as a conference paper at ICLR 2022
E PROOFS AND ADDITIONAL THEORETICAL CONSIDERATIONS
E PROOF OF PROPOSITION
The with-replacement sampling is defined in Definition 2. Let us here define the without--rellaeement sampling.
Definition 3. A minibatch SGD without replacement computes the update to the parameter with the following set of equations:
gi 3 Eis, Vi Vi, Wc-1);
(98)
Wr W Agr,
where 5: |B,] S N is the minibatch size, and the set B, is element uniformly-randommy drawn from the set of all S-size subsets of NY.
From the definition of the update rule for sampling with or without replacement, the covariance matrix of the SGD noise can be exactly derived.
Proposition 8. The covariance matrices of noise in SGD due to minibatch sampling as defined in Definitions 2 and 3 with arbitrary N are
I { VLVE - (L()) C(w) -
(with replacement) (without replacement)
S(N-)) & EN V4V4 [[(W)V((w)"]
(99)
where the shorthand notation ( (w) == 1( Vi, w) is used.
In the limit of = N S. two cases coincide. In the N >>>>> S limit, both methods of sampling have the same noise covariance as stated in Proposition 1:
1 N Σ VAVE (V(()) SN 1-1
C(w)=
-
(100)
Remark. We also note that different way of defining minibatch noise exists in Hoffer et al. (2017). The difference that our definition requires the size of each minibatch to be exactly S, while Hoffer et al. (2017) treats the batch size also as random variable and is only expected to be S. In comparison, our definition agrees better with the common practice.
Now prove Proposition 8.
Proof. We derive the noise covariance matrices for sampling with and without replacement. We first derive the case with replacement. According to the definition, the stochastic gradient for sampling with replacement can be rewritten as
N
gn 8n
(101)
where gn == Vln and
Sn 1, if / multiple n's are sampled in S, with s
(102)
The probability of sn assuming value / is given by the multinomial distribution
P(s, = - (5)
S (5) N
(
1 1 N
(103)
Therefore, the expectation value of s, is given by
8 EB[Sa] [P(sn =1)
S N
(104)
1=0
which gives
Ep [g[g] = g=
Σ = VL(w).
(105)
N
n=1
25

---

Published as a conference paper at ICLR 2022
For the covariance, first calculate the covariance between Sn and Syy, Due to the properties of the covariance multinomial distribution, we have for n # n
Eg[sn = cov( sns Sw + [[sn]
S
=
N
N2
S(S 1).
(106)
N2
and for n'
[n[ Var[s, sn +[ssp]
S
S
N N N
SN + S(S-1)
(107)
N
Substituting these results into the definition of the noise covariance yields
C(w) EB ág"
]
EB[
N
N
Σ
Σ
gn.n
T [B[S
gg"
52
n'=1
N
T S(S-1)
SN + S(S-1) 11)
5(S-1)
Σ
Sn
Σ
gn
T
gg
S2
N2
52
N2
N²
N
T
NS
gngn
588
5
1
VLVE - -L())(WW
(108)
$
Then, derive the noise covariance for sampling without replacement. Similarly, according the definition, the stochastic gradient for sampling without replacement can be rewritten as
g
8n Sn n=1
(109)
where
8n
(110)
1, n S.
The probability of that is sampled in S from N is given by
(N-1) S
P(sn 1) 8-1
(111)
N
5
The expectation value of is then given by
S
EB[ [s[] P(nn =
1)
(112)
N
which gives
g Σ gn VL().. n=1
(113)
For the covariance, we first calculate the covariance between S, and sw By definition, we have for # n
Ep Sw] P(s, 1, 1) 1|s/, 1)P(5, 1)
(N-2) 8-1
(N-1) S(S-
(114)
(N-1)
(N)
N(N-1)
26

---

Published as a conference paper at ICLR 2022
and for n'
= P(s, =
(115)
N
Substituting these results into the definition the noise covariance yields
C(w) = Eslgg"
1
N
gagn T BB [81
T
5
BB
N
gn T S(S - 1)
1 N
S
S(S 1) N(N-1 1)
Σ
gm
T
gg
5
N(N -)
52
N
N S
N S
NS N - Σ gngn
5(N -1 gg
n=l
N-S
Σ
VAVE VI(W)VL(())
(116)
S(N - 1)
N
1-1
E.2
PROOFS IN SEC. 4.2 PROOF OF LEMMA
E.2.1
Proof. From the definition of noise covariance (2), the covariance matrix for the noise the label is
C(w)= NS Σ (L( [L((w=)V
N [=
N E(w" as - -.) (w" xw
Σ w
Tu
ci)
5
N
S N [w"" 2:0i;; 8; w Ex
1
N
w
|
aa ww)
(a xF
(117)
= - 5 (Aww A+ Tr[[Aw o2+
(118)
where we have invoked the law of large numbers and the expectation value of the product of four Gaussian random variables in the third line evaluated as follows.
Because N is large, we invoke the law of large numbers to obtain the kj, component of the matrix as
N E(wTr,
D Σ
D Σ
lim No000
T n, w) = Es[w xrw T x w] jkk EB T,
Earwy
(119)
N
Because the average is taken with respect to x and each is Gaussian random variable, we apply the expression for the product of four Gaussian random variables [[523227)) E[ar [ E 2| B[
x224 [xx121 [ [[x22a]= 2 ]E[ [=2]] [Erra] (Janssen and Stoica, 1988) to obtain
EB
Σ warzzaa Σ arwy
D
EB
Σ magoraay
EB
Σ xep wy wy
+ EB
Σ
EB
Σ way
D
+ En
Σ Wyar
Σ
Jo acpreeer Wae
[aryane]
= =((Aww) T A) [TTAAW]A
(120)
Writing = Ew [ww]w=FJ we obtain
[[[[[w]]=:'ss (AEA- o2A).
(121)
27

---

Published as a conference paper at ICLR 2022
This method has been utilized repeatedly in this work.
E.2.2 PROOF OF THEOREM
Proof. We substitute Eq. (5) into Eq. (3) which a general solution obtained in a recent work (Liu al., 2021):
+ n 1 1- p²
(1=))AAE++AA EAA)
/ +
1- (AAME) ACA.
(122)
To solve it, we assume the commutation relation that [A, = AA- AA 0. Therefore, the above equation can be alternatively rewritten as
(1- 1DD
4(1-M 1 2 1
}AA
AA EAA+ AAE (1 WD -1- -
1(1 1- 4 1 AA 2 1 S
1
- T[AEJAA= S AA
(123)
To solve this equation, we first need to solve for Tr(AE].
Multiplying Eq. (123) by G ==
22 - DD 1-1 + ( AA AA] and taking trace, we obtain
Tr(AE]-
1 THAAE]TH[MGGT)
TTHAA,,,
(124)
S
S
which solves to give
d² TH[AG,,]]]
S 1 [TAGG == o²
Tr[AE]
5 Kgp
(125)
Therefore, Σ is
σ² S
A 2(1 - (i- 5 AA
(126)
5
E.2.3 TRAINING ERROR AND TEST ERROR FOR LABEL NOISE
In the following theorem, we calculate the expected training and test loss random noise in the label.
Theorem 8. (Approximation error and test loss for SGD noise
label) The expected approxi- mation error, or the training loss, is defined as Lurain := E L (ww ]; the expected test loss is defined as Lres 2 E EB ((wTx) x:){] For SGD with noise in the label given by Eq. (5), the expected approximation error and test loss are
02 ((1 1+
Ax s
Ltrain
(127)
2g
Ltest
25 k.
(128)
Remark. Notably, the training loss decomposes into two additive terms. The term that is proportional the bias, caused by insufficient model expressivity to perfectly fit all the data points, while the second term that is proportional to AK/S is the variance in the model parameter, induced by the randomness of minibatch noise.
Remark. When the learning rate X is vanishingly small, the expected test loss diminishes whereas the training error remains finite as long as label noise exists.
Proof. We first calculate the approximation error. By definition.
Lurain Ew [w][[w)]
TT[AA]]
0
1
1 2
2
2
Am
1
(129)
2
s
28

---

Published as a conference paper at ICLR 2022
The test loss is
7 Bw [w' "w]Aww= 1
T[[A]]
Liest
Ao
(130)
25
E.3 MINIBATCH NOISE FOR RANDOM NOISE IN THE INPUT
E.3.1 NOISE STRUCTURE
Similar to label noise, noise in the input data also cause fluctuation. We assume that the training data points Σ; 2 77; can be decomposed into signal part and random part. As before, we assume Gaussian distributions. 24 NO, A) and πι (O, B). The problem remains analytically solvable if we replace the Gaussian assumption by the weaker assumption that the fourth-order moment exists and takes some matrix form. For coneiseness, we assume that there no noise in the label, namely y/ u 2; with a constant vector One important quantity in this case will be uu := U. Notice that the trick w u = w no more works, and so we write the difference explicitly here. The loss function then takes the form
L(w) = 2N
- u)) w T m = (w - A) u) wWTTBo Bw. [[(w-u)
(131)
1=1
The gradient VL= = (A+ (w -ww-U u) + Bu vanishes at w. == B)++ which is the minimum of the loss function and the expectation of the parameter convergence. It can be seen that, even at the minimum w., the loss function remains finite unless = O, which reflects the fact that in the presence of input noise, the network not expressive enough to memorize all the information of the data. The SGD noise covariance for this type of noise is calculated in the following proposition.
Proposition 9. (Covariance matrix for SGD noise in the impul) Let the algorithm be updated according to Eq. (1) or (98) with random noise in the input while the limit N + oo is taken with D held fixed. Then the noise covariance
C C=1 (KEK Tr[KE]K- TT[AK-
(132)
where K = A+ B, and E:= Ew [(w ((w-w.)(s - w.
Remark. It can be seen that the form of the covariance (132) of input noise is similar to that of label noise (5) with replacing A by K and 0² by Tr[AK-1BU suggesting that these two rypes of noise share a similar nature.
Defining the test loss as Lrest := J Nw EB [(wT -- )² Proposition can then be used to calculate the test loss and the model fluctuation.
Theorem 9. (Training error, test loss and model fluctuation for noise in the input) The expected training loss is defined as Lerain == Ew [w]](w)], and the expected test loss is defined as Lrest == EE [wTT= --rox uFx2) 2) For SGD with noise in the input given in Proposition (9), the expected approximation error and test loss are
Ltrain TTAAK BU] (0 1+ X )x)
(133)
X Tr[AK BUJK' + ! T AB'U],
Luest
(134)
25
where x'
I-x K-GG-T with G' := 21 1₂ K, and B' = K-1 B. Moreover, let [K U] 0.
Tv Tr[Kd-11
Then the covariance matrix of model parameters is
Σ= XTr[AK- BU]
1+ -
An' [ | 1 )
(135)
S
Remark. Note that [K is necessary only for an analytical expression of Σ. It can be obtained by solving Eg. (144) even without invoking [K, [K O. In general, the condition that [K, U. = =0 does not hold. Therefore, only the training and test error can be calculated exactly.
Remark. The test loss is always smaller than or equal to the training loss because all matrices involved here positive semidefiite.
29

---

Published as a conference paper at ICLR 2022
E.3.2 PROOF OF PROPOSITION 9
Proof. We define E
::=Ew Ew [(w (w--w. (.))(.-..))) Then,
Bw [w[ww - B)+ AUA(A AUA B) =+ AUA" EA.
(136)
Ew (w- - B'UB" == EB.
(137)
where we use the shorthand notations A' (A+ A, B' := + (A+ and ΣA E+ AT Eg 2 + BUU We remark that the covariance matrix 2 here still satisfies the matrix equation (3) with the Hessian being K == A+ B.
The noise covariance
1 C(w)
8. TTT, 2, (w w. TT 1 VL(W)VL(w))
S E (w
5
S [A(w u)(ww (ww A Bww BB A u)w T B+ Bw (w-wu)
Tr[A(w- uw ()- K Tr[Bww"]K]
(138)
In Eq. (138), there are four terms without trace and two terms with trace. We first calculate the traceless terms. For the latter two terms, we have
E [(w - u)w "] = E A'UBT Ew [w(w - u)
(139)
= - -
AT
(140)
Because A' + B' Ip. after simple algebra the four traceless terms result in 2(A+ B)Σ B).
The two traceful terms add to Tr[AEB BEA]K. With the relation AB' BA', what inside the trace is
AΣ BEA = KE AK BU.
(141)
Therefore, the asymptotic noise is
C' = [w]]
[KEK+ Tr[AE BEA]K}
(142)
3 (KEK Tr[KEJK BUJK
(143)
E.3.3 PROOF OF THEOREM
Proof. The matrix equation satisfied by
ΣK 3K+KE KE x KEK
1
(144)
By using similar technique in Appendix E.2.2, the trace Tr[KE] can be calculated to give
Tr[K]] ATT[AK BB]
(145)
S
where Ke' TH TIK 1
with G = 2Ip (11
5 K.
TYK 1j
With Eq. (145), the training error and the test error can be calculated. The approximation error is
Lirain = [=]][[]]]= 'T[AA BEA] 2
TT[AK- BBU] (1 >
)
(146)
30

---

Published as a conference paper at ICLR 2022
The test loss takes the form of bias-variance tradeoff:
Fww [(w - (w- A(w n)] ! T.(AEn]
Tetst 2 T
FB [w
x u²xx) x²²))]
2
TMAAK BU] (1+ An'
TT[AG" 1 - TT[BBT AB'U]
25
TMAKK BU]R++ Tr[BTAB'U]] AB'U].
(147)
28
Let [KKU]=0, Then Σ can be explicitly solved because it is function of and U. Specifically,
= ATr[AK-
BB]
Ardd)
An'
> x1++ 5 K
1+ S
2ID
(148)
E.4 PROOFS IN SEC. 4.3
E.4.1 PROOF OF PROPOSITION
Proof. The covariance matrix of the noise
C(w)= 1 N SN E[(w u) u))TTaa x. + Tw] [1 uw w wT
(
(A(w - uw - u) Ar]Aw u)(w 5
(149)
Using similar trick as in Appendix E.3.2, the asymptotic noise is
1 C= 5 [AEA+ [AEA+TAE]4++ T< Tr[I" ATU]A+TAUAT)
(150)
F,4.2 PROOF OF THEOREM 4
Besides the test loss and the model fluctuation, we derive the approximation error here as well.
Theorem. (Training error, loss and model fluctuation for learning with L2 regularization) The expected training loss is defined Tarain := w [.( w ], and the expected test loss is defined as Liest = Eg [wWTr ; u ==) For noise induced by L2 regularization given in Proposition 3, let [A, O. Then the expected approximation eтror and test loss are
Ltrain À (TAKITATTTTAAA (1+ Ar )
A
T
Ar s (151)
(TAAG
25
25
TUK
Liest
X
(Tr
UR r) TT[AK-
(152)
25
where K Tr/4 +
-- r= MAKK 9 U0
I- with G 2Ip >K - Moreover,
if A, I and U commute with each other, the model fluctuation is
Am シ= S B= ( A
K A
Ar
A
A
(153)
5
Remark. Because I may not be positive semidefinite, the test loss can be larger than the training loss, which is different from the imput noise case.
31

---

Published as a conference paper at ICLR 2022
Proof. The matrix equation obeyed by E is
A
X (TH[A- AK
ΣK KE- - AKEK
AEA
Tr[AE]A+
g
5
(154)
where we use the shorthand notation := T. Let [A, r] 0. Using the trick in Appendix E.2.2, the trace term Tr[AE] is calculated
TMEE]]
(155)
TH42K=1G-1 114K- and G 2ID A >(K A²))
where x
1-3 1- -TT
The training error is
Ltrain Tr[AEr rEAl
TMKE]] TT[AK TU]
X
IHAK T (1+ An A T
Ar
Ar
28
5 25
5
TT].
(156)
The test loss is
Leest
En EB (w T x- - u 2
Faw (w - -1) A(w (w- u)] Tr[ Tr[ 2
A
(Tr[AK 2T r) + _ TT[AK
(157)
25
Let A, and U commute with each other. Then,
X == -T[AK-2-u7(1
dr ) s
A A ( aAK +
dr
5 x)K-
(158)
s
E.4.3 PROOF OF COROLLARY 1
For a 1d example, the training loss and the test loss have simple form. We use lowercase letters for 1d cases.
Corollary 4. For 1d SGD with La regularization, the training loss and the test loss are
-a++))-I(aat)
Ga(u 7)
Lirain a7
2
(159)
2(a+7))
2(a+7) - 7)2 2a²
an2
2 >(a+7)
Luest 2(a+ 2(a+
(160)
(a+ 5
Proof. The training error and the test loss for 1d cases be easily obtained from Theorem E.4.2.
Now we prove Corollary 1.
Proof. The condition for convergence is 1- - 0. Specifically,
2
2
<0.
(161)
S
32

---

Published as a conference paper at ICLR 2022
For given y, the learning rate needs to satisfy
2(a 4(++7) 7) (a+ aa
A
(162)
For given X, y needs to satisfy
1 ax
1 ax
_a²sz
<9<
(163)
which indicates a constraint on A:
a7
(164)
2
If y is allowed to be non-negative, the optimal value can only be due the convergence condition. Therefore, negative optimal y requires an upper bound on it being negative, namely
1 - ax
2
u2 X2
3
0.
(165)
Solving it, we have
aX
(166)
1
2 3
By combining with Eq. (164), necessary condition for the existence of negative optimal y is
2
S
2 +
25--))>
(167)
Hence, a negative optimal 7 exists, if and only if
2
S and S
aa V
(168)
1+
2
For higher dimension with r TTp it is possible calculate the optimal y for minimizing the test loss (11) as well. Specifically, the condition is given by
1 d f(7) = 0,
Laest dy
(169)
2d:9(1)
where
S(7) = TT AK-2 In 2
(170)
9(>) = 1- A 1- 1T[AA2
(171)
Although it is impossible to solve the equation analytically, it can be solved numerically.
33


Reviews:
[{"review_id": "x0kSNc1phLI", "reviewer": "Reviewer_pu11", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n- The paper is very well-written and easy to read.\n- It provides clear explanations on the theorems and formulae with breakdown of the expression.\n- The results are novel and close a few important gaps.\n- Implications discussed in Section 6 are particularly interesting and match the empirical observations with the new theory.\n- To the extent I have checked the theorems are sound and derivations are accurate.\n\nMinor comments:\n- Pg. 2: The following should be stated as an assumption that might not hold in practice. “One can decompose the gradient into a deterministic plus a stochastic term.”\n- Pg. 3: The following sentence would not hold for loss functions without a global minimum at zero: “This proposition implies that there is no noise if our model can achieve zero training loss (which is achievable for an overparametrized model)”.\n- Pg. 5: “By definition, C = J is the FIM”. J is used without proper definition.\n- Pg. 8: Is there a typo in this sentence: “Figure 1-Right confirms that, at a large learning rate, the optimal weight decay can indeed be optimal.”\n- Could your analysis be extended to handle learning rate schedules such as step decay? For example, it has been observed that the gradient variance can increase for a while after a learning rate drop in common neural network training [1]. That would not exactly match Theorem 5 that predicts shrinking covariance by the loss. Any thoughts?\n\n[1] Faghri, F., Duvenaud, D., Fleet, D. J., & Ba, J. (2020). A Study of Gradient Variance in Deep Learning. arXiv preprint arXiv:2007.04532.\n", "summary_of_the_paper": "This paper studies the properties of the gradient noise in mini-batch SGD using discrete-time analysis for a fixed learning rate. Their analysis is more general than prior work and considers SGD with momentum, a learning rate matrix (that subsumes preconditioning methods), and regularization. They provide closed form expressions for the noise covariance for various machine learning problems including linear regression with label noise (Theorem 3), with regularization (Proposition 3), and more general models under reasonable assumptions (Theorem 5). Their theory matches the empirical observations about the effect of mini-batch size and learning rate on the gradient noise for small batch sizes and large learning rates where there was previously a gap in theory. Section 6 illustrates the significance of the results as it is able to explain various observations in machine learning theory and deep learning practice as well as providing ideas for practitioners.", "summary_of_the_review": "This work makes a clear contribution and is very well written.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "wb7AwDCwf_n", "reviewer": "Reviewer_4YZA", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n- Clarity: The paper is well written, and the main technical results are well presented.\n\n- Novelty: Understanding the generalization capability of SGD is an active research topic. The paper proposes a slightly different approach than previous works, finding an analytical solution to the noise shape of SGD in simple scenario rather than relying on continuous-time or Hessian-based approximation. This approach allows to highlight the importance of some factor, such as the loss level, that were neglected in previous works. I think this finding would be of interest to the community.\n\n- Significance: The analysis also provides some insight on more applicative issue. It highlights why the linear scaling of the learning rate might not work for small batches or large learning rate. It also provides justification for the use of negative weight decay or insight about stability in second order approaches.\n \nWeakness:\n- The analysis of the noise structure in the generic setting relies on two assumptions and it is unclear if those assumptions are realistic in practice, in particular assumption 2 about loss homogeneity.\n", "summary_of_the_paper": "This paper investigates the importance of noise in mini-batch SGD. The main contribution of this paper is to derive an analytic solution to the shape and strength of SGD minibatch noise for linear regression with random noise in the label, linear regression with additional L2 regularization and non-linear regression given some assumption about the model fluctuation and loss homogeneity.\n\nThose analysis reveal that 1) the SGD noise is proportional to the loss level 2) the shape of the noise differs for different loss minima. This is in contrast with previous work looking at this problem using various approximation which neglects those effect.\n\nAuthors finally highlight various applications for those findings. In particular, the analytical formulation they derive for the noise explains why the linear scaling rule of the learning rate does not work for small batch or high learning rate.\n", "summary_of_the_review": "\nThis paper investigates the role minibatch SGD noise from a different perspective and provide novel insight about the strength and the shape of the noise. Paper finding could be of interest to the community, I therefore recommend acceptance. \n\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "hC66eWOXif", "reviewer": "Reviewer_oBNf", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Strengths:**\n- Notation is clear and the background nicely sets up the distinction between the noise $C$ and the fluctuation $\\Sigma$ as the two central objects of study in this work\n- I like where you discuss the \"crucial messages\" a theoretical claim delivers.   \n- I think your thought experiments considering interpolating regimes as a means of demonstrating a more carful understanding of noise (beyond crude approximations) is needed are very good.\n\n**Weaknesses:**\n- The major weakness of this work is while you clearly explain how your theory differs from previous works and introduces fewer \"approximations\", I keep coming back to a line in your introduction, \"the limitation of these approximations is not well understood\".   So, I would expect to not just understand how your theory differs, but be clearly explained and demonstrated (empirically) where the limitations in the previous approximations are and how your theory provides explanation or avoids these limitations.  I feel that many of the experiments in section A of your appendix are very valuable to demonstrate this and I would pull these up to the main and discuss directly.\n- Section 6 applications discuss a broad range of claims and connections to your theory, but most of them are not supported empirically (such as the connection to second order methods, $\\lambda - S$ scaling law, high dimensional regression) and its left to the reader to assume that expressions you give explain the limitations in the previous approximations.  Again if the limitations are not well understood how should we evaluate the more complex expressions you derive in this work.  For example, in section 6.5 Failure of $\\lambda - S$ scaling law, you write \"it is known that this scaling law fails when the learning rate is too large, or the batch size is too small\", but you do not provide empirical evidence or citation of this fact.  You then state that \"our result in Theorem 2 suggests the reason for the failure\" and demonstrate that only \"the leading term is indeed proportional to $\\lambda/S$\" but you don't demonstrate that including the higher order terms corroborates the evidence demonstrating how the scaling law fails.  I think spending section 6 to focus deeply on a few applications highlighting your theory would be more valuable than trying to address a broad range of applications shallowly.\n- In section 4.3, it is not clear to me how \"[L2] regularization also causes a unique SGD noise\".  How could regularization introduce gradient noise given that the gradient for the regularization term is independent of the batch or label and is a deterministic function of the parameters?   Indeed considering the expression for $\\eta_t = \\frac{1}{S} \\Sigma_{i \\in B_t} \\nabla \\ell (x_i,y_i,w_{t-1}) - \\mathbb{R}_B[\\hat{g}_t]$ given in background it seems clear that this term would be independent of $\\gamma$. I can understand how regularization could effect the fluctuation $\\Sigma$, but not how it could change the noise $C$.  Please explain the sentence \"this term is due to the mismatch between the regularization and the minimum of the original loss\" in more detail.\n- Throughout your work you borrow heavily upon the theory developed in Liu et al. 2021 (i.e. it seems that almost all your theory builds upon Theorem 1 which is from their work), but you don't address this work directly in related work and scatter throughout the body references to how your expressions differs from theirs.  I would devote a whole paragraph either in related work or discussion to discuss how your work builds up or overlaps with their work.\n\n**Minor comments:**\n- You write in section 4.2 \"a Hessian approximation fails to account for the randomness in the data of strength $\\sigma^2$, but why wouldn't the  $\\sigma^2$ be absorbed into the scalar constant $c_0$ in the Hessian approximation?  Its not clear to me that this is a setting that is a limitation for the Hessian approximation.\n- Consider discussing the recent empirical work \"Stochastic Training is Not Necessary for Generalization\"", "summary_of_the_paper": "This work studies the noise (covariance of gradients) and fluctuation (covariance of parameter distribution) in the limiting setting of linear regression and for a general loss trained with discrete SGD.", "summary_of_the_review": "In summary, I think this is a good paper and leaning towards accept.  I think if the authors address some of the weakness discussed above then I will raise my score. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "VUK0_zZS8iu", "reviewer": "Reviewer_5twY", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros: \nThe structure of the paper is well-orgainzed. The problem is well-motived with proper assumptions. The proofs look clear.  The mentioned applications in the paper are practical and novel, it is interesting to see this work could promote the understanding of many empirical findings.\nCons: \nI have some concerns as following list though.\n1. in Figure 4, the proposed theory saw a drop around $\\lambda~1.7$, can you explain this phenomenon and what does it mean when $\\lambda>1.7$. same question for Figure 3 right and Figure 5 where exact solution has a sudden drop/rise.\n2. in equation 53, the first equality replies on approximating $(l_i')^2$ by average. this could be an issue if there exist some relationship between $l_i'$ and $\\nabla f$, for instance $f$ is quadratic function.\n3. in equation 54, the sign of $\\Gamma w^\\star$ should be positive? Is this a typo?\n4. in equation 55, how do we get $\\mathcal{O}(S^{-2})$ term for the first equality? same for equation 57\n5. in equation 117, why there's $t-1$ subscript for $w$ after first equality?\n6. line spacing needs fix: equation 131", "summary_of_the_paper": "This paper studies the minibatch noise for discrete time SGD and discusses its approximation on different applications. The novelty of the paper stands on the derivation of minibatch noise covatiance of discrete time SGD. For special cases with label noise and L2 regularization, this work derives the exact solution of covariance. For a more general setting, the author gives the covariance replying on two assumptions. In application section, a few experiments are presented to show the strength of the theory.", "summary_of_the_review": "This paper finds the novelty of discrete sgd noise. Compared with previous with Hessian noise approximation, this paper found the shape of sgd noise on a generic setting. Experiments shows much better consistency with this work. Throughout the paper, the theoretical derivations are solid and numerical experiments are well-design. Suggest acception.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: All the reviewers think that the work is significant and new. Therefore, they support the paper to be published at ICLR 2022. Given the strong results and the “accept” consensus from the reviewers, I accept the paper as “spotlight”. The authors should implement all the reviewers’ suggestions into the final version.
Decision: Accept (Spotlight)