Stability analysis of SGD through the normalized loss function

Abstract: We prove new generalization bounds for stochastic gradient descent for both the convex and non-convex case.  Our analysis is based on the stability framework. We analyze stability with respect to the normalized version of the loss function used for training. This leads to investigating a form of angle-wise stability instead of euclidean stability in weights. For neural networks, the measure of distance we consider is invariant to rescaling the weights of each layer.  Furthermore, we exploit the notion of on-average stability in order to obtain a data-dependent quantity in the bound. This data dependent quantity is seen to be more favorable when training with larger learning rates in our numerical experiments.This might help to shed some light on why larger learning rates can lead to better generalization in some practical scenarios.

Reviews and author responses:

This paper discusses the uniformly statbility and on average stability of SGD on convex and nonconvex optimization algorithms. With proper choice of step size, the stabilities can be theoretically bounded. The convex case extends from Hardt et al. For the nonconvex case, the result is concrete and can be applied to neural nets.

A few questions/suggestions about writing. If they are addressed I'm glad to raise my rating.

1. I'd like to see more discussions about the bound for non-convex theorem. The quantity in Thm4 is a bit hard to interpret. Perhaps some examples can be given in appendix, such as matrix factorization, neural networks, and compute $\epsilon_{av}^\alpha$ with respect to data and parameters.  
2. I'd like to see more comparison with related literatures, specifically carefully compare to Hardt et al. Thm5 is kind of similar to Rademacher complexity which is widely used in other generalization type of papers, such as Du and Allen-Zhu's papers, so I'd like to see them as well.
3. If the length is not limited for appendix, it's more clear to put Hardt et al proof in it. Even if the exact same procedures are used, the sequence of lemmas and critical ideas can be listed so readers know which thm is quoted/replaced for the scope of this paper. 
4. If possible I want to see an end-to-end result in appendix. Say we take an nn and train it with some algorithm, what is the time complexity, training error and generalization error, and what is the overall algorthm for this ml task and how to tradeoff, etc. 
5. Defs need more clarification. With $l(w,z)$ and $l^\alpha(A(S),z)$, $S = \{z_1,...,z_n\}$, I feel that $z$ is the data. So $w$ and $A(S)$ are also data? When saying "$A(S)$ is the output of $A$", does it mean the algorithm recovers the parameter, and uses the model with this parameter to label the training set, and you observe the difference with true training labels? In Def. 2, is $z$ a single random sample or the set of samples $z_1,...,z_n$? Since $S= \{z_1,...,z_n\}$ also shows up, I guess that $z$ is another single random sample to avoid redundancy. In $l^\alpha(A(S^i),z)$, $z$ replaces some data, and what about $l^\alpha(A(S),z)$, why does $z$ appear here? And what does the superscript $\alpha$ mean?

===================== Update =====================

Raised to 6. A theoretical calculation on 4. would be great, if it can be done before publication. NN can be simplified, say linear, or one hidden layer linear. 

We list below some answers to your concerns:

1.  We refer to the answer to question 1 from AnonReviewer4.

2. For a very large number of iterations, we must be careful to choose small learning rates or have a large initialization to enforce good uniform stability (with respect to the normalized loss). The problem is that a uniform stability analysis is a worst case analysis and if the quantity \hat{K} can be too close to $0$ (in a potential worst case), the trajectory could get very close to the origin. Remark 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 can lead to a drastic change in the classifier. 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$).
    
    Stochasticity is actually crucial for stability. Full batch gradient descent is not stable. The fact that the bound is better for small batch size is aligned with empirical studies. The bound does indeed get bad very fast if we grow the batch size. We think that the current framework would not be good if the goal is to obtain strong generalization guarantees for large batch training. However, the bound can be good for guarantees on small batch training.

3. This is true. First, to be technically more accurate we should probably consider subgradients instead of gradients. Indeed, the loss function is differentiable only almost everywhere. Second, the layerwise smoothness could be replaced by a condition on the Hessian that would hold only almost everywhere. However, we do not believe that these extra technicalities would offer new interesting insights.

4. We agree that it is not clear to what extent the growing norm assumption will be satisfied in practice. This assumption is responsible for the added $\epsilon$ term which is a constant term not decaying when $n$ grows. This is indeed not desirable and future work could try to remove or replace this assumption. Note that in the overparameterized regime, once the training set is fit perfectly, the norm of the parameters will be growing to keep reducing the training loss. We conjecture that it might be possible to exploit such properties in the analysis.

5. If the norm of the parameters is too small (relative to the norm of the gradient), the exploration of the hypothesis space is more important leading to a worse stability (see also figure 1, for this intuition).
    

We provide below comments and answers on your 5 questions:

1. Interpretation of bound in theorem 4: First, to evaluate the bound, we need to find the best $t_0$. There is a tradeoff here between two quantities. A small $t_0$ is better for the term $M_{\alpha}(\frac{Bt_0}{n}+\epsilon)$ but is worst for the remaining term. This establishes the best \say{burn in} period. The amount of \say{exploration} before $t_0$ is not effecting the generalization bound. The amount of exploration measured by $\zeta_t$ (and through the learning rate via the value of $c$ in the bound) is affecting the bound only after iteration $t_0$. The bound will be better if we can reach a region in parameter space such that the classifier is then effectively not changing much. This is measured through the norm of the gradient but also takes into account the norm of the parameters. The bound also involves the batch size $B$. In order to have a good guarantee of generalization, a small batch size is required.

2. In the convex case, we have a discussion under theorem $1$ comparing to the result in Hardt et al.(2016). We added the statement of this result in the appendix for ease of comparison. Notice that there is no on-average result in Hardt et al.(2016). We understand the question as asking for more comparison with other generalization bounds in the litterature (outside of stability). It is indeed yet to be seen how the kind of bounds presented in this paper compares to PAC-Bayes bounds or margin based bounds currently in the litterature. The paper (Jiang et al.,2019) presents an extensive empirical study to this effect. We might try to incorporate our results within their study in the future.
    
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, Samy Bengio,''
    Fantastic Generalization Measures and Where to Find Them''

3. For the scope of the paper,  we refer to the answers given to question 1 and 3 from Anonreviewer 4 above.

4. We provide a new experiment in section $6.2$ where we evaluate the bound in theorem $5$ on MNIST.

5. $A(S)$ denotes the output of the algorithm $A$ (SGD with some choice of batch size and learning rate in our case) when we run it on training set $S$. The ouput will be represented with some weights $w$. It is the neural network or the linear classifier obtained at the end of training. The quantity $l(w,z)$ is the loss on example $z=(x,y)$ (data) for the classifier represented by $w$ (weights).
    
    In definition 2, $z$ does not belong to the original training set. It is indeed a single random variable sampled from $\mathcal{D}$. \say{what about $l^{\alpha}(A(S),z)$
, why does $z$ appear here?}  We use the same $z$ to evaluate the loss for both the classifier resulting from $S$ and the classifier resulting from $S^{(i)}$. This $z$ is also used to replace the $i^{th}$ example in $S$ to form $S^{(i)}$. 
\say{what does the superscript $\alpha$ mean?} It means that we use the $l^{\alpha}$ loss (normalized loss). This is defined in equations $4$ and $12$.


We provide further explanations for your 3 comments below:

1. A first advantage of considering the normalized loss is that since it is bounded, we can then exploit generalization bounds requiring this assumption. Note that the standard cross-entropy loss is unbounded (if we do not constrain the weight space). We can therefore obtain the results in theorem $2$ and $5$ bounding the true risk of the output of unconstrained SGD when minimizing the cross-entropy training loss.
    
    As mentioned in the related work, previous empirical study (Liao et al., 2018) is observing that test error is well correlated with the empirical normalized loss. In order to develop generalization bounds better correlated with test error in practice, it therefore makes sense to consider the normalized loss in the analysis. 
    
    When looking at the quantity $\zeta_t$ in theorem $4$, we can see the norm of the parameters (at the different layers) in the denominator. This means that when we reach larger norms of parameters, stability (with respect to the normalized loss) is less negatively affected. The intuitive reason is the following: the same magnitude of step  results in a smaller change in the classifier if the parameters are larger (see the new figure 1). In (Hoffer et al.,2017), it is observed that small batch training and larger learning rates (finding solutions generalizing better) are reaching larger norms of parameters (see also our figure $4$). Using standard euclidean distance in the analysis of stability would lead us to believe that this behavior is highly undesirable. Our analysis shows that this behavior can actually be favorable to the on-average stability with respect to the normalized loss.



2. In Kuzborskij and Lampert (2018), the data dependent quantity is measured at the initialization and is not affected by the learning rate. The behavior shown in our experiments about learning rates (section $6.1$) cannot be captured by their results. Furthermore, Kuzborskij and Lampert (2018) do not provide probabilistic bounds (we could provide them but we would need to assume that the loss function is bounded). At the end of section $6.1$, we say \say{Remark also that considering only the norm of the gradients without the norm of the parameters would lead to a less favorable quantity compared to considering both the norm of the gradients and the norm of the parameters.} The standard approach (like the one used in Kuzborskij and Lampert (2018) for example) would not benefit form the norm of the parameters.



3. Lemma $1$ to $5$ are not in previous analysis of stability. Lemma $1$ is crucial and allows for a layerwise analysis for neural networks. The layerwise smoothness is a new property that we introduce. The fundamental structure of the proofs (for our theorems $1$ and $4$) remains similar to Hardt et al. (2016). 
    

We want to thank all the reviewers for their precious time and their comments on our paper. We will try to give the clarification needed.

This paper develops new stability bounds for SGD. The main difference from the existing studies is that they consider stability bounds for normalized loss functions where the parameters are normalized to have a norm of $1$. This paper considers both convex and nonconvex cases. For the convex case, the authors develop uniform stability bounds and high-probability bounds. For the nonconvex case, the authors develop on-average stability bounds for neural networks. Experimental results are also given.

Comments:
1. The motivation should be further clarified. In particular, it is not quite clear to me the advantage of considering normalized loss as compared to the standard loss. At least I can not see the advantages from the theoretical results.

2. In Thm 1, the authors assume that the norm of the initial weight is larger than the summation of step sizes. I think this is a very strong assumption and can be violated if we run sufficient number of iterations. Furthermore, it seems that the argument is incremental. The analysis heavily uses the existing stability bounds in Hardt 2016. It seems that the authors do not introduce new techniques. The stability bounds grows as a linear function of the batch size. This dependency on the batch size is not desirable since it means that the generalization would decrease by a large amount if we increase the batch size.

3. In the non-convex case, the authors consider RELU neural networks and assume the smoothness of loss functions. However, RELU networks are not smooth. Therefore, the smoothness assumption in the theoretical results is violated.

4. In Thm 4, the authors assume a growing norm assumption: the norm of layer-wise weights are non-decreasing after some iterations. This is a very strong assumption and can be easily violated in practice. Although the authors give some examples to show this can happen, I still believe it is far from true in general.

5. The stability bound in Thm 4 depends on $1/\|w\|$. If we encounter weights with very small norms, the bound would be very large. Furthermore, both the stability bound and the generalization bound involves $\epsilon$, which is not decreasing w.r.t. either $T$ or $n$. This shows that the result is not converging which is a basic requirement in statistical learning theory.

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.  

Reasons for score: 

Overall, this paper is well-written, and the idea of using normalized loss function to assist the theoretical analysis is interesting. My major concerns are listed below. Hopefully, the authors can address them in the rebuttal period. 

Here are some comments:

1) The idea of the normalized loss function should be discussed in detail. It is not very clear what is the motivation behind it and how this assists the theoretical analysis.

2) In the experiment, the authors compare their analysis with the one in Hardt et al. (2016). However, I think they should compare it with the one in Kuzborskij & Lampert (2018). Throughout this paper, this comparison should also be provided in their theoretical results. 

3) Their main theoretical results seem to be direct corollaries following from Hardt et al. (2016) and Elisseeff et al. (2005). The proof is trivial and it is preferred to explain their technical innovation.


Meta-review: The paper considers generalization analysis of SGD using stability analysis. The authors argued the use of normalized version of the loss function, and angle-wise stability. However, the reviewers pointed out that both motivation and novelty of the current work are not strong enough for it to be accepted by ICLR.

Decision: withdrawn-rejected-submissions