On Linear Stability of SGD and Input-Smoothness of Neural Networks
Authors: Chao Ma
Lexing Ying
Venue: NeurIPS 2021 Spotlight (2021)
Keywords: Stochastic Gradient Descent
Implicit regularization
Linear stability
Sobolev Seminorm

Abstract: The multiplicative structure of parameters and input data in the first layer of neural networks is explored to build connection between the landscape of the loss function with respect to parameters and the landscape of the model function with respect to input data. By this connection, it is shown that flat minima regularize the gradient of the model function, which explains the good generalization performance of flat minima. Then, we go beyond the flatness and consider high-order moments of the gradient noise, and show that Stochastic Gradient Dascent (SGD) tends to impose constraints on these moments by a linear stability analysis of SGD around global minima. Together with the multiplicative structure, we identify the Sobolev regularization effect of SGD, i.e. SGD regularizes the Sobolev seminorms of the model function with respect to the input data. Finally, bounds for generalization error and adversarial robustness are provided for solutions found by SGD under assumptions of the data distribution. 

On Linear Stability of SGD and Input-Smoothness of Neural Networks
Chao Ma
Lexing Ying
Department of Mathematics
Department of Mathematics Stanford University Stanford, CA 94305 Lexing@stanford.cem
Stanford University Stanford, CA 94305
chaoma@stanford.co
Abstract
The multiplicative structure of parameters and input data in the first layer of neural networks is explored to build connection between the landscape of the loss function with respect parameters and the landscape of the model function with respect input data. By this connection, it is shown that flat minima regularize the gradient of the model function, which explains the good generalization performance flat minima. Then, go beyond the flatness and consider high-order moments the gradient noise, and show that Stochastic Gradient Descent (SGD) tends impose constraints these moments by linear stability analysis of SGD around global minima. Together with the multiplicative structure, we identify the Sobolev regularization effect of SGD, i.e. SGD regularizes the Sobolev seminorms of the model function with respect to the input data. Finally, bounds for generalization error and adversarial robustness provided for solutions found by SGD under assumptions of the data distribution.
Introduction
Stochastic gradient descent (SGD) is the most widely used optimization algorithm to train neural networks [25, 4]. By taking mini-batches of training data instead of all the data in each iteration, it was firstly designed as a substitute of the gradient descent (GD) algorithm reduce its computational cost. Extensive researches are conducted on the convergence of SGD, both convex and non-convex objective functions [21, 22, 5, 3]. these studies, convergence usually proven in the cases where the learning rate is sufficiently small, hence the gradient noise is small. In practice, however, SGD is preferred over GD not only for the low computational cost, but also for the implicit regularization effect that produces solutions with good generalization performance [7, 16). Since trajectory of SGD tends to that of GD when the learning rate goes to O, this implicit regularization effect must come from the gradient noise induced by
and moderate learning rate.
When studying the gradient noise of SGD, majority of work treat SGD as an SDE, and studies the noise of the SDE [15, 20, 17, 18, 32). However, SGD is close to SDE only when the learning rate small, and it is unclear that practical setting whether the Gaussian noise of SDE can fully characterize the gradient noise of SGD. Some work resort to heavy-tailed noise like the Levy process [27, 37]. Another perspective to study the behavior SGD is by its linear stability [30, 12]. This is relevant when the learning rate is not very small. The linear stability theory can explain the fast escape of SGD from sharp minima. The escape time derived from this theory depends on the logarithm of the barrier height, while the escape time derived from the diffusion theory based SDE depends exponentially on the barrier height [32]. The former is more consistent with empirical observations [30].
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

An important observation that connects the generalization performance of the solution with the landscape of the loss function is that flat minima tend to generalize better [13, 31]. SGD is shown to pick flat minima, especially when the learning rate is big and the batch size is small [16, 15, 30]. Algorithms that prefer flat minima are designed to improve generalization [6, 14, 19, 35]. On the other hand, though, the reason why flat minima generalize better is still unclear. Intuitive explanations from description length or Bayesian perspective are provided in [13] and [23]. In [10], the authors show sharp minimum can also generalize by rescaling the parameters at flat minimum. Hence, flatness is not a necessary condition of good generalization performance, but is still possible to be sufficient condition. In the study of linear stability in [30], besides the sharpness (a quantity inversely proportional to flatness), another quantity named non- con-uniformity is proposed which roughly characterizes the second order moment of the gradient noise. It is shown that SGD selects solutions with both low sharpness and low non-uniformity.
In this paper, we build a complete theoretical pipeline to analyze the implicit regularization effect and generalization performance of the solution found by SGD. Our starting points are the following two questions: (1) Why SGD finds flat minima? (2) Why flat minima generalize better? Our answers to these two questions go beyond the flatness and cover the non-uniformity and higher-order moments of the gradient noise. This distinguishes SGD from GD, and is out of the scope that can be explained by SDE. For the first question, we extend the linear stability theory of SGD from the second-orde moments of the iterator of the linearized dynamics to the high-order moments. At the interpolation solutions found by SGD, by the linear stability theory, we derive a set of accurate upper bounds of the gradients' moment. For the second question, using the multiplicative structure of the imput layer of neural networks, we show that the upper bounds obtained in the first step regularize the Soboley seminorms of the model function. Finally, bridging the two components, generalization error under some assumptions of the data distribution. The bound works well when the distribution is supported on a low-dimensional manifold (or a union of low-dimensional manifolds). An informal statement of our main result is
main result is bound of
MMain result) Around an interpolation solution of the neural network model, assume that (1)the k-th order moment of SGD's 'ss iterator of the linearized dynamics is stable, (2)with probability at least 1 - the testing data is close to a training data with distance smaller than 6, and (3) both the model function and the target function are upper bounded by a constant M. Then, at this interpolation solution have
generalization error S 82 + eM2,
where n is the number of data. The bound depends on the learning rate and the batch size of SGD as constant factors.
The formal description stated in Theorem 6 of Section 5. Our analysis also provide bounds for adversarial robustness. As byproduct, we theoretically show that flatness of the minimum controls the gradient norm of the model function at the training data. Therefore, searching for flat minima has the effect of Lipschitz regularization, which is shown to be able to improve generalization [24].
Lying at the center of our analysis the multiplicative structure of neural networks, i.e. in each layer the output features from the previous layer is multiplied with a parameter matrix. This structure is rich source of implicit regularization. In this paper, we focus on the input layer, and build connection between the model function's 's derivative with respect to the parameters and its derivatives with respect to the data. Concretely, let f(x, W) be a neural network model, with x being the input data and W being the parameters. We split the parameters by W = (W1, W2), where W1 is the parameter matrix of the first layer, and W₂ denotes all other parameters. Then, the neural network can be represented by the form fx, W)- = f(W, xx, W2) due to the multiplicative structure of W1 and x. Accordingly Ww, /(Wix, W2) and VW],W, x, W2) take the form
Vw, F(W W2) 0f xW,,,,,, W2) xT. a(
Vxf(Wix, W2) W Of(Wix, W2)
(1)
a(Wix)
where @f(Wix. Wa) is usually a long expression produced by back propagation from the output to the dWix first layer. By (1) we have
/Vz W2)11a < |W ll VV j(Wix, W2)1222
||x||2
Hence, (WI W2) is upper bounded by Vw, /Wix, x, W2), given that W, is not too big and is not too small. By the bound (2) we can derive the regularization effect of flatness at interpolation

---

solutions. To see this, let {(x,, 9:) H=z be the training data, and L(W) 2n
E W) yi)² :=1
be the empirical risk given by the square loss. Let W (Wi, W32) be an interpolation solution, i.e. [x; )* y/ for any i = 1,2, n, and L(W*) 0. Define the flatness of the minimum as the sum of the eigenvalues of v²i((W*) i.e. flatness(W**) Tr(V*((W*)). W being the
interpolation solution implies
flatness(W 1 n
Vw fx,, WT]]
Wx),13
Hence, from (2) we can obtain the following equation that gives bounds for the gradients of the model function with respect to the input data at the training data,
| min; x
n
| 12 min; |x
Vxf(xx, W>)
flatness( (W*).
The left hand side of (4) is usually used to regularize the Lipschitz constant of the model function, and such regularization can improve the generalization performance and adversarial robustness of the model. Hence, (4) reveals the regularization effect of flat minima. Later in the paper, we extend the analysis to higher-order moments of the gradient and combine the results with the linear stability theory of SGD to explain the implicit regularization effect of SGD. We also extend the bound on fxf(x, W*) from training data to all x in a neighborhood of the training data, which implies the regularization of flatness is actually stronger than the left-hand-side of (4).
The strong correlation between W fxx, W) and fx W) is justified by numerical experiments in practical settings. Specifically, in experiments we compare the following two quantities:
1/2
1 n2
1 12
gw :=
Ox =
(5)
Figure 1 shows the results for a fully-connected network trained on FashionMNIST dataset and VGG-11 network trained on CIFAR10 dataset. In each plot, 9w and 9x of different solutions found by SGD are shown by scatter plots. The plots show strong correlations between the two quantities. The colors of the points show that SGD with big learning rate and small batch size tends to converge to solutions with small fw and Ix, which is consistent with our theoretical results on the Sobolev regularization effect of SGD. To summarize, the main contributions of this paper are:
We extend the linear stability analysis SGD to high-order moments the iterators. At the solutions selected by SGD, we find a class of conditions satisfied by the gradients of different training data. These conditions cover the flatness and non- aonunfformmty, and also include higher order moments. They characterize the regularization effect of SGD beyond GD and SDE.
2. By exploring the multiplicative structure of the neural networks' input layer, we build relations between the model function's derivatives with respect the parameters and with respect the inputs. By these relations we turn the conditions obtained for SGD into bounds of different Sobolev (semi norms of the model function. In particular, we show that flatness of the minimum regularizes the L norm of the gradient of the model function. This explains how the flatness (as well as other stability conditions for SGD) benefits generalization and adversarial robustness.
3. Still using the multiplicative structure, the bounds for Sobolev seminorms can be extended from the training data a neighborhood around the training data, based on certain smoothness assumption of the model function (with respect to parameters). Then, bounds for generalization error and adversarial robustness are provided under reasonable assumptions of the data distribu- tion. The bounds work well when the data are distributed effectively in a set that consists of low dimensional manifolds.
2 Preliminaries
Basic notations For any vector x = (fi,,.., xT,"² RD and any p 2 1, ||x]| is the conventional
p-norm of x. If g(x): IRd + IR" is a function with vector output, (1(x))l means the p-norm of
the vector g(x), instead of a function norm. For any matrix A E _g ||A||p is the operator norm
3

---

(a)
(b)
Figure I: (a) A fully-connected neural network trained on FashionMNIST. (b) A VGG-11 network trained on CIFARIO. For both (a) and (b), the left panel shows gw and 9x of the solutions found by SGD with different learning rate, while batch size fixed at 20. The right panel shows solutions found by SGD with different batch size, with learning rate fixed at 0.1.
induced by the p-norm on Rh and R. ||||| denotes the Frobenius norm. Let g RR -> R function, $2 C Ie be set in which is defined, N and p 2 1. The Sobolev seminorm Ila.000 defined as Ig = (2/o== DD "xx) where a = (a11 an) is a multi-index with positive integers, |a| = Σ Li=1 Qx, and Ddg is defined Dog 21 extended to fractions, but in this paper we always consider integers. When Ω = (xi,... x, is finite set, we define the Sobolev seminorm 19l....2
gi-ly g 25 The index can be
(x)
For symmetric matrix A E R we let A (AA) be the set of eigenvalues of and Amax (A) be the maximum eigenvalue of A. For two matrices A and B, A @ B denotes the Kronecker product of and B. For k E No, let A A
@ A @ @ A where is multiplied for times. Therefore, if A E R= then A E Im xn However, for a vector x E IR", with an abuse of notation sometimes we also use x@k denote the rank-one tensor product. In this case, x** -
Problem settings In this paper, we consider the learning of the parameterized model f(xx, W), with x being the input data and W being the parameters. Let d be the dimension of the input data and be the number of parameters in W, i.e. has scalar output. We consider the models with multiplicative structure on input data and part of the parameters. With abuse of notation let W = (W1, W2), where W1 mmd is the part of the parameters multiplied with x and W2 denotes other parameters. Then, the model can be written as
IRd and W E R. We assume that the model
f(x, W) = j(Wix, W2).
(6)
We remark that most neural networks have form (6), including convolutional networks, residual networks, recurrent networks, and even transformers. Note that since the x in (6) can also be understood as fixed features calculated using input data, (6) also includes random feature models.
We consider a supervised learning setting. Let j be data distribution supported within TRd,, and f" Rd IR. be target function. A set of training data {(x,, 34:)) }/)) obtained by sampling x..-.., X, i..dd. from 14, and letting y: f (x:) As mentioned the introduction, the model is learned by minimizing the empirical loss function i(W) 3. Σ (xx W) y1)² in the space of parameters. The population loss, or the generalization error, is defined as L(W)- = fx W) - (x))
For any == 1, ..., n, let L:W = of the SGD with learning rate 7 and batch size B
[VV==((((>> ()'x,, W) - y1)² be the loss at x,, Then, the iteration scheme
7 2 VL.( B
W+1 W.
where şt 52.. $2.....(h) a B-dimensional random variable uniformly distributed on the B-tuples in (1.2..., n and independent with W,. In the paper, we study interpolation solutions W found by SGD, which satisfies f fx, W*) = 3/ for any 1, ..., n. Obtaining interpolation solutions is possible in the over-parameterized setting [34], and is widely studied in existing work [36, 33).
Some assumptions will be made when deriving the generalization error bounds. Firstly, we assume the model function with respect to the parameter W to be smooth around W
Definition 1. Let 6, C be positive numbers, and be positive integer. We say the model f(x, W) satisfies the (C, 8, k) -ocaa smoothness condition at data x and parameter W. if for any W' such

---

that ||W 2 S 6, there is
[ 12))2 S C(IVV [xx, 122 + 1).
(8)
Remark 1. In the definition above, we consider the 2k-norm of the gradients for the convenience of later analysis. When & = I, the condition (8) becomes |I w fxx W S C(VV f(x, 12 1). This is weaker than local approximation by Taylor expansion, which usually yields results like f(x, W) - Vw fxx, W S Cá. Here, we only require that the gradient with respect to the parameters does not get exceedingly large when W is close to W
Next, we need the data distribution μ to support roughly on low dimensional manifold (or union of low dimensional manifolds), thus the neighborhoods of training data can well cover all test data.
Definition 2. Let y be probability distribution supported in RRd and {x:} [Xi,-..X,]] be set of points in R4. For positive constants 5 and we say {xi} (6, ccovers if
P Px
min I8i5h
|x - x.12 > =6 E,
(9)
Lee with high probability a point sampled from μ lies close point in {x,)-
Definition 3. Let {x,} fxi,..,,Xx,, } be set of points in RR. For positive constants 5 and K, we say {x+} is (6, K) )sscttetee if the function x(xx) := Ei= 18) (x) satisfies x(x) S K. Here, (x,, 6) is the ball in IRd centered at x, with radius ó, and 1x is the indicator function of set A. Remark 2. The uniform upper bound s(x) S in the above definition can be weakened as an integral upper bound 1 Vx Jx x(x)dx S K, where X = - B(x;, 6) and Vx is the Lebesgue
volume of X.
Definition 4. Let j be a probability distribution supported in IRd. For positive integer n and positive constants 6, E1, E2, say μ satisfies the (nz 6, E e2) -covered condition, if with probability at least 1- $ over the choice of i.i.d. sampled data from μ., {xI,.... x,J. we have
min 15:5n
||X
x1I12> > 8 82.
(10)
On the other hand, for positive integer n and positive constants 8, e and K, we say μ satisfies the (n, 6, 5 -)-scattered condition, if with probability at least 1 E1 over the choice of n i.i.d.
sampled data {x,} from μr, {x,) is (6. scatterdd
Later in the analysis of the generalization error (Section 5), we will assume the model and the data distribution satisfy the conditions in Definition and 4 with appropriately chosen constants.
3 Linear stability theory of SGD
Compared with full-batch GD, SGD adds in each iteration random "noise" to the gradient of the loss function. Hence, it is harder for SGD to be stable around minima than GD. This is shown by the linear stability theory of SGD,
Recall that the iteration scheme of SGD with learning rate 7 and batch size B is given by (7). Let W be an interpolation solution for the learning problem. When {W;} is close to W*, the behavior of (7), including its stability around the minimum, can be characterized by the linearized dynamics W*:
B Vwf(xe: (ww
W441 We
WT) W*).
(1I)
B
The lincarization made by considering the quadratic approximation of Ly near W For ease of notation, let a, = Vw f fx;, W*), 11, = a.aa". In the linearized dynamics (11), without loss of generality we can set n 0 by replacing W, by W. W* Then (11) becomes
B W,
B
W41 W.
a
ag
I
7 H B 1=1
Wr
(12)
B
Next, we define the stability of the above dynamics.
5

---

Definition 5. For any k E N we say W is k-th order linearly stable for SGD with learning rate 77 and batch size B, if there exists a constant (which may depend on k) that satisfies
& |
C
for any Wi, t E NN,, given by the dynamics (12) from any initialization distribution of Wo.
(W,) is the trajectory of the linearized dynamics. It characterizes the performance of SGD around the minimum W*, but does not equal the original trajectory of SGD (W.). practice, quadratic approximation of the loss function works in a neighborhood around the minimum. Hence, linear stability characterizes the local behavior of the iterators around the minimum. The Wo in Definition the starting point of the linearized dynamics. can be a point in the trajectory of the real dynamics, but should not understood as the initialization of the real dynamies (W)). Since W, follows a linear dynamics, the Wo here does not need to be close to W Without linear stability of (W:) (W,) } can never converge to the minimum. It is possible that (W,) oscillates around the minimum at a distance out of the reach of local quadratic approximation. However, empirical observations in [11] show that SGD usually converges and oscillates in region with good quadratic approximation. Hence, linear stability plays an important role practice.
In [30], the following condition on the linear stability for k = 2 is provided.
Proposition 1. (Theorem in [30]) The global minimum W* is 2nd -order linearly stable for SGD with learning rate n and batch size B if
Amax { (U- nH1)² + r - B3
B(n 1) 1
(13)
where H = 1 Σ H. and Σ = Σ H - H
1=1
Shown by the proposition, the spectra of H and E influence the linear stability of SGD. Thus, in [30] the biggest eigenvalue of H and Σ are named sharpness and non-uniformity, respectively. The condition (13) is then be relaxed into a condition sharpness and non-uniformity. Similar analysis was also conducted in [9] for linear least squares problems.
Remark 3. The definition of stability in [30] concerns E| Well? which slightly different from the definition above when E| W/I / - ||W | < Wl,?, the two ways to define linear stability are equivalent.
2. However, since E||W. and |EW@22|| || are equivalent, i.e.
In the current work, we analyze the dynamics of EW which is a closed linear dynamics. As a comparison, the dynamics of E is not closed, i.e. EIJWW+. WWW+| is not totally determined by 2 By (12) we have
B
@k
S ++1
7
H
E
W
(14)
B
1=1
Let Tp be the set of all subsets of (1,2.... n)
with cardinality B, and J = (i,, ia} E IB be batch. Note that Ç is independent with W, Taking expectation for (14) gives
B
86
66
EW 0k ++1
E
Σ
HE:
EW
1
17
Hi,
EW
(15)
B
JEI
j=1
Denote
@k
B
The B =
7 H
(16)
j=1
for each batch J. and T = 11 Eg
To.B be the expectation of T over the choice of batches.
Then we have SW ++1 T EW

---

On the other hand, if EW is understood as tensor in R=xxx-Xx then it is always symmetric tensor (any permutation of the index does not change the entry value). Hence, it has the decomposition EW = Di i= , with X E R and 2 E R [8]. Let M, be the set of symmetric tensors in I 1"x=x- and M = {A M: A= Ei= dee and X 2 O for any i = 1,...,] 7 be the set of "positive semidefinite" symmetric tensors. We provide the following conditions for the linear stability of SGD.
Theorem 1. For any l E N", the global minimum W is t order linearly stable for SGD with learning rate 13 and batch size B, if and only if
1 B TTI:BAlF P S ||A||F
(17)
holds for any AE MÆ if is an even number, or for any A My is is an odd number.
All proofs are provided in the appendix. As a corollary, when k = 2, we can write down another sufficient condition for linear stability than Proposition 1. See Corollary 3 in the appendix.
Recall that we let a; Vw fx,, W*) and H; = a.a, The linear stability conditions in Theorem 1 imply the following bound on the moments of a, Specifically, we have
Theorem 2. If a global minimun W* is t order linearly stable for SGD with learning rate n and batch size B, then, for any j E 21,,,...P), we have
n
2 B
(18)
7
n
where Dis means the j entry of a.,
By Theorem if a global minimum is stable for SGD with some order, then the gradient moment of this order is controlled. Summing from 1 to w, we have / f W S 2wn For k and 2, this gives control to the flatness and non-uniformity of the minimum, respectively. For general k, applying the Hölder inequality on (18), we have the following corollary which bounds the mean 2k-norm of the gradients at the training data.
Corollary 1. If a global minimum W is h order linearly stable for SGD with learning rate n7 and butch size B, then
n
V fx;, W2)
1
n a.
10
M
2B
(19)
7
n
B
11
4 The Sobolev regularization effect
In this section, we build connection between W f(x W) and Vxf(x, W) using the multiplication of Wi and x. For general parameterized model, hard to build connection between the two gradients. However, this is possible for neural network models, in which the input variable is multiplied with a set of parameters (the first-layer parameter) before any non-linear operation. By this connection, at an interpolation solution that is stable for SGD, we can turn the moments bounds derived in the previous section into the bounds on Vx/(x, W). For different k; the moment bound on Vw/(x, W) controls the Sobolev seminorms of f(, W) at the training data with different index p. When = 1, this becomes a bound for the sum of gradient square at the training data, which is used in the literature
regularization term to regularize the Lipschitz constant of the model function. This explains the favorable generalization performance of flat minima.
Recall f(x, W) = f(W1x, Wz), and we can express Vxf and Vw, f as (1). Then, we easily have the following proposition.
Proposition 2. Consider the model J(Wix, W2). For any k E N", any W (W,, W2) and &
W
|Vx/(Wix W2) 28
[Vw W
(20)
2
As corollary, we have the following bound for the Sobolev seminorm using Vwf.
7

---

Xn) Corollary 2. Let f(x, W) = (Wix, W2), and {xi} = Xx.. Let [f(: W (x,)), be the Sobolev seminorm of f(,, W) on {xi} as defined at the beginning of Section 2. Then,
[ff, W W,
Vw w
(21)
min muinisiss Sinn |x*||2
Combining Corollary 2 with the stability condition in Theorem 2, we have the following control for the Sobolev seminorm of the model function at interpolation solutions that are stable for SGD.
Theorem 3. If global minimum W* which interpolates the training data t order linearly stable for SGD with learning rate yj and batch size B, and W = (WI, W2). Then,
I((W)) 12k
2B
|s W 11,2,,11
(22)
millsiinn
|x,12x
By Corollary and Theorem 3, as long as the data are not very small (in practical problems, the input data usually normalized), and the input-layer parameters are very large, the Sobolev seminorms of the model function (evaluated at the training data) is regularized by the linear stability of SGD. The regularization effect on Sobolev seminorm gets stronger for bigger learning rate and smaller batch size. When & is big, the dependence of the bound with p is negligible.
If the solution W* satisfies the local smoothness condition defined in Definition 1, the control on fx x, W can be extended to a neighborhood around the training data. Then, we will be able to control the "population" Sobolev seminorms. First, around certain training data, have the following estimate.
Proposition 3. Assume the model f(x, W) satisfies (C, Supprers k)-local smoothness condition for some C, Supprax >> and le E Nº, at W and x.. Then, for any x that satisfies ||x Sappnae x. 2/ (| 2l.. we have
112 S
|V W [Vx((xWW')l2z CII(W)) 12e
(IIV W + 1).
(23)
||x, 1|2k
Proposition 3 turns the local smoothness condition in the parameter space into local smoothness condition in the data space. This made possible by the multiplicative structure of the network's input layer. Specifically, since in the first layer W1 and x are multiplied together, a perturbation of x can be turned to a perturbation of W1 without changing the value of their product. See the proof of the proposition in Appendix B.2 for more detail.
By the results above, we can obtain the following theorem which estimates the Sobolev seminorm of the model function on a neighborhood of the training data.
Theorem 4. Let W W = (Wi, W;) be an interpolation solution that is 1 order linearly stable for SGD with learning rate 1l and batch size B. Let {xi,..,,X,) } be the training data. Suppose the model f(x, W) satisfies (C, 8mpprox; -lool smoothness condition at W* and for any i = 1, 1,2, 2.2,,.,, ...y n. Con-
sider the set X/= := -''' (x ||x x, < 8) with d < uupprax || / | I| l12- Assume {x1,..., Xn
satisfies (8, K))-satttreed condition defined in Definition 3. Then, we have
If 11., < (KVx;) 20 (C1))
w
*
2B
+1)
(24)
min; |x11128
B
V
where Vx, is the Lebesgue volume of Xg.
By Theorem 4, the model function found by SGD is smooth in a neighborhood of the training data, given that the landscape of the model in the parameter space is smooth. When k = 1, the results implies that flatness regularizes [x W*) 1,22x, which is stronger than the gradient square at training data.
Discussion on the tightness the bounds The bounds this section depend on the norm of Wi. This norm may be big, especially when homogeneous activation functions, such as ReLU, is used. In
8

---

iterations
0
100
1000
10000
20000
50000
[W1 l2
1 1.032
1.034
1.218
2.011
2.257
2.270
Table 1:
The average L2 norm of the first-layer parameter matrix of a fully-connected neural network trained on FashionMNIST dataset.
the ReLU case, for example, we can simply scale the first layer parameters be arbitrarily large and second layer small while not changing the model function.
However, in practice SGD does not find these solutions with unbalanced layers because the homo- geneity of ReLU induces an invariant in parameters. Specifically, let W; and Wi-1 be the parameter matrices in the 1th and (I 1)th layers, then W7-,W-= WoW," does not change significantly throughout the training (it is actually fixed in the zero learning rate limit). This invariant keeps layers balanced and prevents the first layer parameters from getting big alone.
On the other hand, we conduct numerical experiments to check the norm of W1 during training. Table 1 shows the average 12-norm of W1 for the same experiments in part (a) of Figure The model a fully-connected neural network, and the dataset FashionMNIST. Results show that the norm of W1 increases at the beginning, but becomes stable afterwards. It will not keep increasing during the training.
Generalization error and adversarial robustness
Regularizing the smoothness of the model function can improve generalization performance and adversarial robustness. This is confirmed in many practical studies [24, 29, 26). Intuitively, function with small gradient changes slowly as the input data changes, hence when the test data close to one of the training data, the prediction error on the test data will be small. In this section, going along this intuition, we provide theoretical amalysis of the generalization performance and adversarial robustness based on the Sobolev regularization effect derived in the previous sections. Still consider the model [xx W) j[W,x, xx, W2). and an interpolation solution W (WI, W found by SGD, First, based on Proposition and Theorem 2, we show the following theorem, which is similar to Theorem 4 but estimates the maximum value of Vx W*]l in X..
Theorem 5. Let W* (W1, W2) be an interpolation solution that is kdhh order linearly stable for SGD with learning rate 1) and batch size B. Let {x,,..., x,} be the training data. Suppose the model f(x, W) satisfies (C, Sampmas kl-local smoothness condition at W and x, for any 1 1,2, .... ",,,,,],,]]
Recall the definition of Xs in Theorem Then, for any x E Xs, we have
CI((W)) min: ||x||22
( (nnww
22
2B V
Vxf(x, S
+1)
(25)
B
With the result above, we can bound the generalization error if most test data are close to the training data, i.c. the data distribution I satisfics a covered condition. This happens for machine learning problems with sufficient training data, especially for those where the training data lie approximately on a union of some -in-iimnsioonll surfaces. This is common in practical problems [1].
Theorem 6. Suppose parameterized model f(x, W) = f(W1x, W2) is used to solve super- vised learning problem with data distribution p. Let f be the target function. Assume i sat- isfies (n, 5,, -cooeerd condition. Assume with probability no less than learning rate y and batch size B can find an interpolation solution W which is ktth order lin- early stable, and (C, 8appmr, k)-local smoothness condition is satisfied at W and X1,.., Xn with Sappnex 2 &||WTlla min; ([x]]12). Further assume that both [s(,, W*) and (+) are upper bounded by Mr, and xx (-)ll2 is upper bounded by M2, for some constants M, and Mz. Then, with probability 22, over the sampling of training data In we have
E1 SGD with
2
Exxsy [xx, W (x)]11 2dC (
2nww" B
2B V
+1
82 82+2M38 +
min; |x,
(26)

---

The bound (26) tends to O as n - oc as long as 5 decays faster than n sh. From geometric perspective, this happens when the dimension the support of μi is less than k. And the lower the dimension, the faster the decay. It is possible to get rid of the n dependency in the bound when A is sufficiently large. We may use the estimate of Sobolev functions with scattered zeros, e.g. Theorem 4.1 in [2]. We leave the analysis to future work.
Adversarial robustness Neural network models suffer from adversarial examples [28], because the model function changes very fast in some directions so the function value becomes very different after a small perturbation. However, the results in Theorem 5 directly imply the adversarial robustness of the model at the training data. Hence, flatness, well high-order linear stability conditions, also imply adversarial robustness. Specifically, we have the following theorem.
Theorem 7. Let W* (WT, W2) be an interpolation solution that is 1 order linearly stable for SGD with learning rate di and batch size B. Let {X1, xn} be the training data. Sup- pose the model f(x, W) satisfies (C, Ooppros k)-local smoothness condition at W* and x; for any 1 1,,2,... 122,...... Then, for any x that satisfies ||x x:l12 S ô for some {l,...) and
5 S Sapprex min; [kll2/|[Will we have
CVaI(W)))12z min, ||x:|128
(
(nnw
#
2B
1f(x f (x,, W (x)
V
+
6.
(27)
6 Conclusion
In this paper, we connect the linear stability theory of SGD with the generalization performance and adversarial robustness of the neural network model. As a corollary, we provide thcoretical insights of why flat minimum generalizes better. To achieve the goal, we explore the multiplicative structure of the neural network's 's input layer, and build connection between the model's gradient with respect to the parameters and the gradient with respect to the input data. We show that as long as the landscape on the parameter space is mild, the landscape of the model function with respect to the input data is also mild, hence the flatness (as well as higher order linear stability conditions) has the effect of Sobolev regularization. Our study reveals the significance of the multiplication structure between data (or features in intermediate layers) and parameters. It is important source of implicit regularization of neural networks and deserves further exploration.
References
[1] Alessio Ansuini, Alessandro Laio, Jakob H Macke, and Davide Zoccolan. Intrinsic dimension of data representations in deep neural networks. arXiv preprint 1905 2019.
[2] Rémi Arcangéli, María Cruz López de Silanes, and Juan José Torrens. An extension of a bound for functions in sobolev spaces, with applications to (m, s)-spline interpolation and smoothing. Numerische Mathematik, 107(2):181-2)
2007.
[3] Raef Bassily, Mikhail Belkin, and Siyuan Ma. On exponential convergence of sgd in non-convex over-parametrized learning. arXiv preprint arXiv: 11fi::1811.02560 2018.
[4] Léon Bottou. Stochastic gradient learning in neural networks. Proceedings of Neuro-Names, 91(:):12, 1991.
[5] Léon Bottou, Frank E Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. Siam Review, 60(2):223-311, 2018.
[6] Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Chris- tian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-sgd: Biasing gradient descent into wide valleys. Journal of Statistical Mechanics: Theory and Experiment,
2019(12): 2019.
[7] Pratik Chaudhari and Stefano Soatto. Stochastic gradient descent performs variational inference, converges to limit cycles for deep networks. In 2018 Information Theory and Applications Workshop (ITA), pages 1-10. IEEE, 2018
[8] Pierre Comon, Gene Golub, Lek-Heng Lim, and Bemard Mourrain. Symmetric tensors and symmetric tensor rank. SIAM Journal on Matrix Analysis and Applications, 1044 2008.
10

---

[9] Alexandre Défossez and Francis Bach. Averaged weast-mean-squee Bias-variance trade-offs and optimal sampling distributions. In Artificial Intelligence and Statistics, pages 205-213. PMLR, 2015.
[10] Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In International Conference on Machine Learning, pages 1019-1028. PMLR, 2017.
[11] Yu Feng and Yuhai Tu. The inverse variance-flatness relation in stochastic gradient descent is critical for finding flat minima. Proceedings of the National Academy of Sciences, 118(9). 2021.
[12] Niv Giladi, Mor Shpigel Nacson, Elad Hoffer, and Daniel Soudry. At stability's edge: How to adjust hyperparameters to preserve minima selection in asynchronous training of neural networks? arXiv preprint arXiv:1909. 12340, 2019.
[13] Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation, 9(1):1-42, 1997.
[14] Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. arXiv preprint arXiv:1803.05447,, 2018.
[15] Stanislaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in sgd. arXiv preprint arXiv:1711.04623 2017.
[16] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batte batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv:1609.045:
2016.
[17] Qianxiao Li, Cheng Tai, and E Weinan. Stochastic modified equations and adaptive stochastic gradient algorithms. In International Conference on Machine Learning, pages 2101-2110 10.
PMLR, 2017.
[18] Qianxiao Li, Cheng Tai, and E Weinan. Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations. J. Mach. Learn. Res., 20:40-1, 2019.
[19] Tao Lin, Sebastian U Stich, Kumar Kshitij Patel, and Martin Jaggi. Don't use large mini-batches, 2018.
use local sgd. arXiv preprint /
[20] Stephan Mandt, Matthew D Hoffman, and David M Blei. Stochastic gradient descent as approximate bayesian inference. arXiv preprint arXiv: / 704.04289) 2017.
[21] Eric Moulines and Francis Bach. Non-asymptotic analysis of stochastic approximation algo- rithms for machine learning. Advances neural information processing systems, 24:451-459, 2011.
[22] Deanna Needell, Rachel Ward, and Nati Srebro. Stochastic gradient descent, weighted sampling, and the randomized kaczmarz algorithm. Advances in neural information processing systems, 27:1017-1025, 2014.
[23] Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nathan Srebro. Exploring generalization in deep learning, arXiv preprint arXiv: 1w..110047.9777 2017.
[24] Adam M Oberman and Jeff Calder. Lipschitz regularized deep neural networks converge and generalize. arXiv preprint arXiv:8008. 095400 2018
[25] Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, pages 400-407, 1951.
[26] Andrew Ross and Finale Doshi Velez. Improving the adversarial robustness and interpretability of deep neural networks by regularizing their input gradients. In Proceedings of the AAAI Conference Artificial Intelligence, volume 32, 2018.
[27] Umut Simsekli, Levent Sagun, and Mert Gurbuzbalaban. A tail-index analysis of stochastic gradient noise in deep neural networks. In International Conference on Machine Learning, pages 5827-5837. PMLR, 2019.
[28] Christian Szegedy, Wojciech Zaremba, Ilya Sutskever, Joan Bruna, Dumitru Erhan, Ian Goodfel- low, and Rob Fergus. Intriguing properties of neural networks, arXiv preprint arXiv: 131,66199 2013.
11

---

[29] Dániel Varga, Adrián Csiszárik, and Zsolt Zombori. Gradient regularization improves accuracy of discriminative models. arXiv preprint arXiv:1712.09933 2017.
[30] Lei Wu, Chao Ma, and Weinan E. How sgd selects the global minima in over-parameterized learning: A dynamical stability perspective. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, pages 5289-8298, 2018.
[31] Lei Wu, Zhanxing Zhu, et al. Towards understanding generalization deep learning: Perspec- tive of loss landscapes. arXiv preprint arXiv:1706 10239, 2017.
[32] Zeke Xie, Issei Sato, and Masashi Sugiyama. A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. arXiv e-prints, pages arXiv-2002, 2020.
[33] Zitong Yang, Yu Bai, and Song Mei. Exact gap between generalization error and uniform convergence random feature models. arXiv preprint wwXiv:2103.045555 2021.
[34] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. arXiv preprint arXiv:/61J.00 1.035000 2016.
[35] Yaowei Zheng, Richong Zhang, and Yongyi Mao. Regularizing neural networks via adversarial model perturbation. arXiv preprint arXiv:2010.04925, 2020.
[36] Lijia Zhou, Danica Sutherland, and Nathan Srebro. On uniform convergence and low-norm interpolation learning. arXiv preprint arXiv:2006. 05942, 2020.
[37] Pan Zhou, Jiashi Feng. Chao Ma, Caiming Xiong, Steven HOI, et al. Towards theoreti- cally understanding why sgd generalizes better than adam in deep learning. arXiv preprint arXiv:2010.056627 2020.
Checklist
1. For all authors...
(a) Do the main claims made the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] Our generalization bounds work on
effectively low dimensional problems.
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms them?
[Yes]
2. you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] In both Section and the statement of theorems.
(b) Did you include complete proofs of all theoretical results? [Yes]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed reproduce the main experimental results (either in the supplemental material or as a URL)? [Yes] See Section and the URL therein.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? (Yes] Section F
(c) Did you report eтror bars (c.g., with respect to the random seed after running experiments multiple times)? [N/A]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, cloud provider)? [Yes] Section F
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [N/A]
(b) Did you mention the license the assets? [N/A]
(c) Did you include any new assets either in the supplemental material URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you't 're
using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information offensive content? [N/A]
5. If you used crowdsoureing conducted research with human subjects...
12

---

(a) Did you include the full text of instructions given to participants and screenshots, if applica-
ble? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board
(IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent
on participant compensation? [N/A]
13


Reviews:
[{"review_id": "pvf7m9rXrN", "reviewer": "Reviewer_gdzy", "summary": "Prior work on minima stability for SGD considered linear stability in expectation using second order moment. This paper considers linear stability in higher moments, while providing necessary and sufficient conditions for it. Using these conditions and the multiplicative structure of neural networks, the authors show that the Sobolev semi-norm of SGD solutions (w.r.t. the input) is bounded from above. Then, under certain smoothness assumptions on the learned function and the target function, the authors prove generalization bounds and adversarial robustness.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The article is well written, and the proofs are rigorous. I carefully checked the proofs of Theorems 1-3, and they look fine. On the other proofs I went over more generally, and they also seem good. \n\nHere are my comments:\n1.\tWhy should we consider higher moments for linear stability? E.g., why second order is not sufficient to say that a minimum is stable? If a minimum is not stable for some order but is stable in a lower order, what this tells us? Can a minimum be stable for any order? if so, does it have special properties?\n2.\tThe bounds proved by the multiplicative property, namely Theorem 4-7, can be loose. For example, consider the simple case of a two-layer univariate network with ReLU activation (i.e. scalar input/output with a wide hidden layer). Since ReLU is positive homogenous we can arbitrarily increase the $ \\ell_p $-norm of the first layer while preserving the function implemented by the network, and thus making the bounds trivial. This is true for any network with ReLU activation (or its variants).\n3.\tFrom 2., it is not true to say that SGD implicitly regulate the Sobolev semi-norm of the function implemented by the network, since there is no (proven) mechanism in SGD that keeps the the norm of the first layer small.\n4.\tThe authors claim that Corollary 3 is more accurate than previous known result, which is written in Proposition 1 (see row 235), however I did not see any explanation nor proof for this claim (maybe I have missed it...).\n5.\tMost of the results hold only for the quadratic loss.\n6.\tI think that Lemma 1 is used to prove necessity in the proof of Theorem 1, and not sufficiency.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "See above", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "10 Hours "}, "extra_text": {}}, {"review_id": "fpw5V00q7mJ", "reviewer": "Reviewer_Ggg9", "summary": "This paper shows that due to the multiplicative structure of input data and first layer parameters in a MLP neural network the flat minima regularizes the norm of the gradient of the model and then explains its generalizability. It also shows that SGD actually regularizes higher-order moments of the gradient of the noise when it is close to an interpolating minima. Utilizing this property of SGD and multiplicative structure shows that SGD minimizes the Sobolev seminorms of the model. The presented analysis is based on quadratic estimation of the objective around the minima and uses linear stability analysis.       ", "questions": "", "limitations": "", "rating": 7, "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": "Understanding the implicit regularization imposed by SGD has great importance in understanding the generalization performance in deep neural networks. This paper utilizes the product structure between the first layer parameter and the input data and discovers the relation between flat minima and the gradient norm or the Lipschitz constant of the model that improves the stability of the model. \n\nHowever, the major theoretical contribution of the paper is local and based on a quadratic approximation of the model around an interpolating local minima. The major part of the analysis is based on the existing results in the literature for quadratic models. \nBesides, it makes a bunch of assumptions and it does not explain which of them are true in a multi-layer perceptron model. \n\nIn terms of the write-up, there are a bunch of typos here and there that need to be fixed such as stocahstic instead of stochastic in the title. \n\nMy question\n\n\n1- It is not clear how restrictive the assumptions the paper makes. It would be clarifying to explain if these assumptions are true for deep neural networks. \n\n2- The product structure has been shown that holds for MLP. Does it hold for other architectures such as CNN or transformers?\n\n3- In practice we usually normalize the data, therefore some x_i could be very close to zero. Based on this, how tight do you think the upper bound in 4 could be?\n\n4- You assume that SGD with fixed step size would converge to an interpolating solution. It would be helpful to explain under which condition this can happen and why you think MLP satisfies these requirements.   \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "YkGpNsXOCiR", "reviewer": "Reviewer_iuXC", "summary": "This work presents a novel framework of \"Sobolev regularization\" effect of stochastic gradient descents for quantifying and establishing wideness of a global interpolating optimum in SGD. The key idea is that because the neural network input $x$ is multiplids to weights $W_1$ and then provided into the neural network, regularity properties of the weights near the interpolating global optimum imply regularity properties (wideness) of the minimum.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, I feel that the viewpoint presented in this work is novel and worthy of publication. The key insight is that since the input $x$ and the neural network weights are first multiplied and then entered into the network, so regularity (stability) properties on $W$ can be transferred onto regularity (wideness/small Sobolev norm) properties with respect to $x$ is a very nice one.\n\nI have one not-so-minor comment. Definition 5 seems to be incomplete; the relation between $W^*$ and $\\tilde{W}^{\\otimes k}_t$ is not completely specified. (All that is said is that $\\{W_t\\}$ is close to $W^*$.) I think the issue is that the authors have the mental model of an interpolating global minimum as an isolated minimum. However, there is some recent evidence that a connected set of global minima is a more accurate model.\nhttps://proceedings.neurips.cc/paper/2018/hash/be3087e74e9100d4bc4c6268cdbe8456-Abstract.html\nhttp://proceedings.mlr.press/v119/shevchenko20a.html\nCan the authors comment on how having a connected set of interpolating global minima may or may not affect the analysis?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "15J4M1m340H", "reviewer": "Reviewer_PF9o", "summary": "The paper extends previous stability analysis ([30]) to high-order moments of the iterator. They characterize the implicit regularization of SGD by showing that a stable global minimum must satisfy conditions regarding its flatness and non-uniformity. In addition, the authors build relations between the model function’s derivative w.r.t. the parameters and w.r.t. to the inputs. Using these relations, they provide upper bounds for the Sobolev seminorm of the model function at stable SGD interpolating solutions. In addition, under additional smoothness conditions, the authors were able to extend that last result to a neighborhood of the training data. Finally, the authors built upon the Sobolev regularization and use it to provide bounds on the generalization performance and adversarial robustness.", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper examines theoretically how does SGD selects the global minima it converges to. Specifically, the authors analyze the implicit regularization effect and generalization performance of the solution found by SGD using a linear stability analysis. \n\nThe paper main contributions are:\n-\tThe authors prove upper bounds on the gradients’ moment at the solution found by SGD, assuming this solution is stable. These conditions also include control on the flatness and non-uniformity. This result provides a better characterization of the implicit bias of SGD.\n-\tThe authors provide upper bounds for the Sobolev seminorm of the model function at stable SGD interpolating solutions. This implies that the Sobolev seminorms of the model on the training data are regularized by the linear stability. In addition, under some smoothness assumption, this result is extended to a neighborhood of the training data. These results suggest a possible explanation for the good generalization of flat minima.\n-\tUnder additional assumptions on the data distribution, the authors provide bounds on the generalization error and adversarial robustness.\n\nClarity/writing: \n\nI enjoyed reading the paper. The paper is well written, the assumptions and results are clearly stated, and many explanations were provided throughout the paper. \n\nMy main question:\n-\tMany of the upper bounds suggested in the paper include the norm of the first layer weights. Can the author extend on the assumption that the input layer weights are not very large?\n\nMinor comments:\n-\tThere is a typo in the title (Stocahstic --> Stochastic) \n-\tI think that adding “sharpness” and “non-uniformity” as formal definitions will make the paper a bit clearer (these terms were explained verbally but it is easy to miss this)\n\n---\n\nUpdate: Thank you for the detailed rebuttal and for elaborating on the upper limit with experimental results. This helped to elevate my concern.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors thoroughly discussed the limitations of their results.", "review_text": "", "extra_scores": {"ethical_concerns": "None.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "11"}, "extra_text": {}}]

Author rebuttal:

Meta-review: All reviewers have praised the clarity of the paper as well as the soundness and novelty of the results. Some important questions regarding completeness of some definitions and applicability of the current assumptions have been raised in the reviews, which the authors have clarified with a strong rebuttal. I encourage the authors to incorporate these remarks in the camera ready version.
Decision: Accept (Spotlight)