Stability analysis of SGD through the normalized loss function

1 INTRODUCTION . In the last few years , deep learning has succeeded in establishing state-of-the-art performances in a wide variety of tasks in fields like computer vision , natural language processing and bioinformatics ( LeCun et al. , 2015 ) . Understanding when and how these networks generalize better is important to keep improving their performance . Many works starting mainly from Neyshabur et al . ( 2015 ) , Zhang et al . ( 2017 ) and Keskar et al . ( 2017 ) hint a rich interplay between regularization and the optimization process of learning the weights of the network . The idea is that a form of inductive bias can be realized implicitly by the optimization algorithm . The most popular algorithm to train neural networks is stochastic gradient descent ( SGD ) . It is therefore of great interest to study the generalization properties of this algorithm . An approach that is particularly well suited to investigate learning algorithms directly is the framework of stability ( Bousquet & Elisseeff , 2002 ) , ( Elisseeff et al. , 2005 ) . It is argued in Nagarajan & Kolter ( 2019 ) that generalization bounds based on uniform convergence might be condemned to be essentially vacuous for deep networks . Stability bounds offer a possible alternative by trying to bound directly the generalization error of the output of the algorithm . The seminal work of Hardt et al . ( 2016 ) exploits this framework to study SGD for both the convex and non-convex cases . The main intuitive idea is to look at how much changing one example in the training set can generate a different trajectory when running SGD . If the two trajectories must remain close to each other then the algorithm has better stability . This raises the question of how to best measure the distance between two classifiers . Our work investigates a measure of distance respecting invariances in homogeneous neural networks ( and linear classifiers ) instead of the usual euclidean distance . The measure of distance we consider is directly related to analyzing stability with respect to the normalized loss function instead of the standard loss function used for training . In the convex case , we prove an upper bound on uniform stability with respect to the normalized loss function , which can then be used to prove a high probability bound on the test error of the output of SGD . In the non-convex case , we propose an analysis directly targeted toward homogeneous neural networks . We prove an upper bound on the on-average stability with respect to the normalized loss function , which can then be used to give a generalization bound on the test error . One nice advantage coming with our approach is that we do not need to assume that the loss function is bounded . Indeed , even if the loss function used for training is unbounded , the normalized loss is necessarily bounded . Our main results for neural networks involve a data-dependent quantity that we estimate during training in our numerical experiments . The quantity is the sum over each layer of the ratio between the norm of the gradient for this layer and the norm of the parameters for the layer . We observe that larger learning rates lead to trajectories in parameter space keeping this quantity smaller during training . There are two ways to get our data-dependent quantity smaller during training . The first is by facilitating convergence ( having smaller norms for the gradients ) . The second is by increasing the weights of the network . If the weights are larger , the same magnitude for an update in weight space results in a smaller change in angle ( see Figure 1 ) . In our experiments , larger learning rates are seen to be more favorable in both regards . Our main contributions are summarized as follows : 2 RELATED WORK . Normalized loss functions have been considered before ( Poggio et al. , 2019 ) , ( Liao et al. , 2018 ) . In Liao et al . ( 2018 ) , test error is seen to be well correlated with the normalized loss . This observation is one motivation for our study . We might expect generalization bounds on the test error to be better by using the normalized surrogate loss in the analysis . ( Poggio et al. , 2019 ) writes down a generalization bound based on Rademacher complexity but motivated by the possible limitations of uniform convergence for deep learning ( Nagarajan & Kolter , 2019 ) we take the stability approach instead . Generalization of SGD has been investigated before in a large body of literature . Soudry et al . ( 2018 ) showed that gradient descent converges to the max-margin solution for logistic regression and Lyu & Li ( 2019 ) provides an extension to deep non-linear homogeneous networks . Nacson et al . ( 2019 ) gives similar results for stochastic gradient descent . From the point of view of stability , starting from Hardt et al . ( 2016 ) without being exhaustive , a few representative examples are Bassily et al . ( 2020 ) , Yuan et al . ( 2019 ) , Kuzborskij & Lampert ( 2018 ) , Liu et al . ( 2017 ) , London ( 2017 ) . Since the work of Zhang et al . ( 2017 ) showing that currently used deep neural networks are so overparameterized that they can easily fit random labels , taking properties of the data distribution into account seems necessary to understand generalization of deep networks . In the context of stability , this means moving from uniform stability to on-average stability . This is the main concern of the work of Kuzborskij & Lampert ( 2018 ) . They develop data-dependent stability bounds for SGD by extending over the work of Hardt et al . ( 2016 ) . Their results have a dependence on the risk of the initialization point and the curvature of the initialization . They have to assume a bound on the noise of the stochastic gradient . We do not make this assumption in our work . Furthermore , we maintain in our bounds for neural networks the properties after the “ burn-in ” period and therefore closer to the final output since we are interested in the effect of the learning rate on the trajectory . This is motivated by the empirical work of Jastrzebski et al . ( 2020 ) arguing that in the early phase of training , the learning rate and batch size determine the properties of the trajectory after a “ break-even point ” . Another work interested in on-average stability is Zhou et al . ( 2021 ) . Differently from our work , their approach makes the extra assumptions that the variance of the stochastic gradients is bounded and also that the loss is bounded . Furthermore , our analysis directly exploits the structure of neural networks and the properties following from using homogeneous non-linearities . It has been observed in the early work of Keskar et al . ( 2017 ) that training with larger batch sizes can lead to a deterioration in test accuracy . The simplest strategy to reduce ( at least partially ) the gap with small batch training is to increase the learning rate ( He et al. , 2019 ) , ( Smith & Le , 2018 ) , ( Hoffer et al. , 2017 ) , ( Goyal et al. , 2017 ) . We choose this scenario to investigate empirically the relevance of our stability bound for SGD on neural networks . Note that the results in Hardt et al . ( 2016 ) are more favorable to smaller learning rates . It seems therefore important in order to get theory closer to practice to understand better in what sense larger learning rates can improve stability . 3 PRELIMINARIES . Let l ( w , z ) be a non-negative loss function . Furthermore , let A be a randomized algorithm and denote by A ( S ) the output of A when trained on training set S = { z1 , · · · , zn } ∼ Dn . The true risk for a classifier w is given as LD ( w ) : = Ez∼Dl ( w , z ) and the empirical risk is given by LS ( w ) : = 1 n ∑n i=1 l ( w , zi ) . When considering the 0− 1 loss of a classifier w , we will write L0−1D ( w ) . Furthermore , we will add a superscript α when the normalized losses lα are under consideration ( these will be defined more clearly in the subsequent sections respectively for the convex case and the non-convex case ) . Our main interest is to ensure small test error and so we want to bound L0−1D ( w ) . The usual approach is to minimize a surrogate loss upper bounding the 0− 1 loss . In this paper , we consider stochastic gradient descent with different batch sizes to minimize the empirical surrogate loss . The update rule of this algorithm for learning rates λt and a subset Bt ⊂ S of size B is given by wt+1 = wt − λt 1 B ∑ zj∈Bt ∇l ( wt , zj ) . ( 1 ) We assume sampling uniformly with replacement in order to form each batch of training examples . In order to investigate generalization of this algorithm , we consider the framework of stability ( Bousquet & Elisseeff , 2002 ) . We now give the definitions for uniform stability and on-average stability ( random pointwise hypothesis stability in Elisseeff et al . ( 2005 ) ) for randomized algorithms ( see also Hardt et al . ( 2016 ) and Kuzborskij & Lampert ( 2018 ) ) . The definitions can be formulated with respect to any loss function but since we will study stability with respect to the lα losses , we write the definitions in the context of this special case . Definition 1 The algorithm A is said to be αuni-uniformly stable if for all i ∈ { 1 , . . . , n } sup S , z′i , z E [ |lα ( A ( S ) , z ) − lα ( A ( S ( i ) ) , z ) | ] ≤ αuni . ( 2 ) Here , the expectation is taken over the randomness of A . The notation S ( i ) means that we replace the ith example of S with z′i . Definition 2 The algorithm A is said to be αav-on-average stable if for all i ∈ { 1 , . . . , n } E [ |lα ( A ( S ) , z ) − lα ( A ( S ( i ) ) , z ) | ] ≤ αav . ( 3 ) Here , the expectation is taken over S ∼ Dn , z ∼ D and the randomness of A . The notation S ( i ) means that we replace the ith example of S with z . Throughout the paper , ||·|| will denote the euclidean norm for vectors and the Frobenius norm for matrices . The proofs are given in Appendix A for the convex case and in Appendix B for the non-convex case . 4 CONVEX CASE : A FIRST STEP TOWARD THE NON-CONVEX CASE . Since the convex case is easier to handle , it can be seen as a good preparation for the non-convex case . Consider a linear classifier parameterized by either a vector of weights ( binary case ) or a matrix of weights ( multi-class case ) that we denote by w in both cases . The normalized losses are defined by lα ( w , z ) : = l ( α w ||w|| , z ) , ( 4 ) for α > 0 . In order to state the main result of this section , we need two common assumptions : L-Lipschitzness of l as a function of w and β-smoothness . Definition 3 The function l ( w , z ) is L−Lipschitz for all z in the domain ( with respect to w ) if for all w , w′ , z , |l ( w , z ) − l ( w′ , z ) |≤ L||w − w′|| . ( 5 ) Definition 4 The function l ( w , z ) is β−smooth if for all w , w′ , z , ||∇l ( w , z ) −∇l ( w′ , z ) ||≤ β||w − w′|| . ( 6 ) We are now ready to state the main result of this section . Theorem 1 Assume that l ( w , z ) is convex , β−smooth and L−Lipschitz for all z . Furthermore , assume that the initial point w0 satisfies ||w0||≥ K for some K such that K̂ = K −L ∑T−1 i=0 λi > 0 for a sequence of learning rates λi ≤ 2/β . SGD is then run with batch size B on loss function l ( w , z ) for T steps with the learning rates λt starting from w0 . Denote by αuni the uniform stability of this algorithm with respect to lα . Then , αuni ≤ α 2L2B nK̂ T−1∑ i=0 λi . ( 7 ) What is the main difference between our bound and the bound in Hardt et al . ( 2016 ) ( see theorem 7 in Appendix A ) ? Our bound takes into account the norm of the initialization . The meaning of the bound is that it is not enough to use small learning rates and a small number of epochs to guarantee good stability ( with respect to the normalized loss ) . We also need to take into account the norm of the parameters ( here the norm of the initialization ) to make sure that the “ effective ” learning rates are small . Note that all classifiers are contained in any ball around the origin even if the radius of the ball is arbitrarily small . Therefore , all control over stability is lost very close to the origin where even a small step ( in Euclidean distance ) can lead to a drastic change in the classifier . The norm of the initialization must therefore be large enough to ensure that the trajectory can not get too close to the origin ( in worst case , since uniform stability is considered ) . An alternative if the conditions of the theorem are too strong in some practical scenarios is to use on-average stability ( l = 1 layer in the results of section 5 ) . As a side note , we also incorporated the batch size into the bound which is not present in Hardt et al . ( 2016 ) ( only B = 1 is considered ) . From this result , it is now possible to obtain a high probability bound for the test error . The bound is over draws of training sets S but not over the randomness of A . 1 So , we actually have the expected 1It is possible to obtain a bound holding over the randomness of A by exploiting the framework of Elisseeff et al . ( 2005 ) . However , the term involving ρ in their theorem 15 does not converge to 0 when the size of the training set grows to infinity . test error over the randomness of A in the bound . This is reminiscent of PAC-Bayes bounds where here the posterior distribution would be induced from the randomness of the algorithm A. Theorem 2 Fix α > 0 . Let Mα : = sup { l ( w , z ) s.t . ||w||≤ α , ||x||≤ R } . Then , for any n > 1 and δ ∈ ( 0 , 1 ) , the following hold with probability greater or equal to 1− δ over draws of training sets S : ( 8 ) EAL0−1D ( A ( S ) ) ≤ EAL α S ( A ( S ) ) + α uni + ( 2n α uni +Mα ) √ ln ( 1/δ ) 2n . Proof : The proof is an application of McDiarmid ’ s concentration bound . Note that we do not need the training loss to be bounded since we consider the normalized loss which is bounded . The proof follows the same line as theorem 12 in Bousquet & Elisseeff ( 2002 ) and we do not replicate it here . Note that we need to use that uniform stability implies generalization in expectation which is proven for example in theorem 2.2 from Hardt et al . ( 2016 ) . Furthermore , a bound holding uniformly over all α ’ s can be obtained using standard techniques . Theorem 3 Let C > 0 . Assume that lα ( w , z ) is a convex function of α for all w , z and that αuni is a non-decreasing function of α . Then , for any n > 1 and δ ∈ ( 0 , 1 ) , the following hold with probability greater or equal to 1− δ over draws of training sets S : EAL0−1D ( A ( S ) ) ≤ inf α∈ ( 0 , C ] { EA max ( Lα/2S ( A ( S ) ) , L α S ( A ( S ) ) ) + α uni + ( 2n α uni + Mα ) √ 2 ln ( √ 2 ( 2 + log2 C − log2 α ) ) + ln ( 1/δ ) 2n } . In the next section , we investigate the non-convex case . We exploit on-average stability to obtain a data-dependent quantity in the bound . Note that it is also argued in Kuzborskij & Lampert ( 2018 ) that the worst case analysis of uniform stability might not be appropriate for deep learning .

Summary: This paper considers the generalization bound for stochastic gradient descent. The authors leverage normalized loss function to analyze the stability of SGD algorithms which further yields the generalization bound. They provide the on-average stability result for non-convex optimization under the ReLU neural network setting. The theoretical results deepen our understanding of the performance of the SGD algorithm and an experiment is provided to illustrate theoretical findings.