On Linear Stability of SGD and Input-Smoothness of Neural Networks

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 Descent ( 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 . 1 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 to reduce its computational cost . Extensive researches are conducted on the convergence of SGD , both on convex and non-convex objective functions [ 21 , 22 , 5 , 3 ] . In these studies , convergence is 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 a trajectory of SGD tends to that of GD when the learning rate goes to 0 , this implicit regularization effect must come from the gradient noise induced by mini-batch and a moderate learning rate . When studying the gradient noise of SGD , a 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 is small , and it is unclear that in 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 of 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 on SDE depends exponentially on the barrier height [ 32 ] . The former is more consistent with empirical observations [ 30 ] . 35th Conference on Neural Information Processing Systems ( NeurIPS 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 a flat minimum . Hence , flatness is not a necessary condition of good generalization performance , but it is still possible to be a sufficient condition . In the study of linear stability in [ 30 ] , besides the sharpness ( a quantity inversely proportional to flatness ) , another quantity named non-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-order 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 input layer of neural networks , we show that the upper bounds obtained in the first step regularize the Sobolev seminorms of the model function . Finally , bridging the two components , our main result is a bound of 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 ) Around an interpolation solution of the neural network model , assume that ( 1 ) the k-th order moment of SGD ’ s 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 , and ( 3 ) both the model function and the target function are upper bounded by a constant M . Then , at this interpolation solution we have generalization error . n 1k 2 + `` M2 , 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 is stated in Theorem 6 of Section 5 . Our analysis also provide bounds for adversarial robustness . As a 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 is 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 a rich source of implicit regularization . In this paper , we focus on the input layer , and build connection between the model function ’ 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 W2 denotes all other parameters . Then , the neural network can be represented by the form f ( x , W ) = f̃ ( W1x , W2 ) due to the multiplicative structure of W1 and x . Accordingly rW1 f̃ ( W1x , W2 ) and rxf̃ ( W1x , W2 ) take the form rW1 f̃ ( W1x , W2 ) = @ f ( W1x , W2 ) @ ( W1x ) xT , rxf̃ ( W1x , W2 ) = WT1 @ f ( W1x , W2 ) @ ( W1x ) , ( 1 ) where @ f ( W1x , W2 ) @ W1x is usually a long expression produced by back propagation from the output to the first layer . By ( 1 ) we have krxf̃ ( W1x , W2 ) k2 kW1k2 kxk2 krW1 f̃ ( W1x , W2 ) k2 . ( 2 ) Hence , rxf̃ ( W1x , W2 ) is upper bounded by rW1 f̃ ( W1x , W2 ) , given that W1 is not too big and x is not too small . By the bound ( 2 ) we can derive the regularization effect of flatness at interpolation solutions . To see this , let { ( xi , yi ) } ni=1 be the training data , and L̂ ( W ) : = 12n nP i=1 ( f ( xi , W ) yi ) 2 be the empirical risk given by the square loss . Let W ⇤ = ( W ⇤1 , W ⇤2 ) be an interpolation solution , i.e . f ( xi , W ⇤ ) = yi for any i = 1 , 2 , ... , n , and L̂ ( W ⇤ ) = 0 . Define the flatness of the minimum as the sum of the eigenvalues of r2L̂ ( W ⇤ ) , i.e . flatness ( W ⇤ ) = Tr ( r2L̂ ( W ⇤ ) ) . W ⇤ being the interpolation solution implies flatness ( W ⇤ ) = 1 n nX i=1 krW f ( xi , W ⇤ ) k2 . ( 3 ) 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 , 1 n nX i=1 krxf ( xi , W ⇤ ) k2 kW ⇤1 k22 mini kxik22 1 n nX i=1 krW1f ( xi , W ⇤ ) k22 kW ⇤1 k22 mini kxik2 flatness ( W ⇤ ) . ( 4 ) 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 rxf ( 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 rW f ( x , W ) and rxf ( x , W ) is justified by numerical experiments in practical settings . Specifically , in experiments we compare the following two quantities : gW : = 1 n nX i=1 krW f ( xi , W ) k22 ! 1 2 , gx : = 1 n nX i=1 krxf ( xi , W ) k22 ! 1 2 . ( 5 ) Figure 1 shows the results for a fully-connected network trained on FashionMNIST dataset and a VGG-11 network trained on CIFAR10 dataset . In each plot , gW and gx 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 gW and gx , which is consistent with our theoretical results on the Sobolev regularization effect of SGD . To summarize , the main contributions of this paper are : 1 . We extend the linear stability analysis of SGD to high-order moments of 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-uniformity , 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 to the parameters and with respect to 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 L2 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 to 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 distribution . The bounds work well when the data are distributed effectively in a set that consists of low dimensional manifolds .

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.