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 conducted a stability analysis of Stochastic Gradient Descent (SGD) for empirical risk minimization induced by the so-called normalized loss function. Here, the normalization is taken with respect to parameters involved in an individual loss; see (4) for the definition of the normalized loss function. The paper should be regarded as a theoretical paper. The main results are stability bounds of SGD for convex and nonconvex ERM schemes.