Stable Minima Cannot Overfit in Univariate ReLU Networks: Generalization by Large Step Sizes

Summary

This paper studies the implicit regularization of large learning rates in gradient descent. The setting is univariate linear regression with two-layer ReLU neural networks. The authors show that, if GD converges to a local minimum, then the function implemented by the neural network at this local minimum has a bounded first-order total variation. This regularity avoids overfitting, as quantified by a generalization bound. Numerical experiments validate the results.

Strengths

Understanding the implicit regularization of optimization algorithms for deep learning is a key topic. This paper studies a fairly realistic setting, which does not necessitate interpolation. To my knowledge, the results are novel and go beyond what was already known on the impact of large learning rates (edge of stability, minima stability). Most previous works focus on parameter space and not function space. The interpretation in terms of function space proposed in the present paper is interesting. The paper is very clearly-written. While I have not read the proofs, the mathematical presentation of the results in the main text is very precise.

Weaknesses

The paper considers a univariate case, some comments on the extension to the multivariate case are given at the end of Section 1.1.

Questions

Do the authors have any insight on the role of overparameterization in the setting they consider? It does not seem to have a large influence on the generalization bounds, but perhaps it would help the optimization (e.g., via a PL inequality)? Minor remarks: - lines 132: the connection between stepsize and L1/L2 regularization has been thoroughly analysed for diagonal linear networks in [2]. - lines 147-153: a relevant theoretical paper on the Edge of Stability is [1]. [1] Damin, Nichani, Lee, Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability, ICLR 2023. [2] Even, Pesme, Gunasekar, Flammarion. (S)GD over Diagonal Linear Networks: Implicit Regularisation, Large Stepsizes and Edge of Stability. NeurIPS 2023.

Rating

8

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

Limitations are properly addressed.

Summary

This paper studies the generalization properties of two-layer ReLU neural networks in a univariate nonparametric regression problem with noisy labels. It proposes a new theory for local minima to which gradient descent (GD) with a fixed learning rate $\eta$ stably converges. The paper shows that GD with a constant learning rate can only find stable local minima whose weighted TV(1) is bounded by $1/\eta-1/2+\tilde{O}(\sigma+\sqrt{\mathrm{MSE}})$. With this property of local minimas, they prove the generalization bound for univariate nonparametric regression. The theoretical results are validated by extensive simulations demonstrating that large learning rate training induces sparse linear spline fits.

Strengths

This paper gives an end-to-end analysis of the generalization of two layer ReLU networks on learning nonparametric univariate functions. The completeness of the work is significant. The theoretical analysis is solid and the proof is well-organized. The authors also present extensive experiments to support their theoretical findings. In general, it is a good paper.

Weaknesses

1. The analysis in this paper has little focus on optimization. There is no rigorous theoretical evidence that GD will find the solutions that satisfies the assumptions (though I know this is an open problem in the literature). 2. This paper is limited to univariate functions, which is not important in practice. The authors claims that the technique can be generated to multivariate functions but they did not show that.

Questions

How large is the family of target functions satisfying the BV(1) condition?

Rating

8

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

The authors have addressed some limitations of their work, acknowledging that their analysis is only for full batch gradient descent and univariate nonparamatric regression.

Summary

The paper studies uni-variate regression with two layer networks and builds on the following observation: If the basin (derived form a quadratic approximation) around a given minimum is too narrow, gradient descent with a fixed step size $\eta$ will escape it, only sufficiently wide basins can capture the iterates of GD with large step sizes. It is thus reasonable to restrict the function space attainable by GD to the set of functions that can be expressed by a minimum with a wide basin. By extending a previous result to the noisy-label setting, they show that the width of the basin (or the hessian of the loss) relates to a (weighted) measure of the TV norm of the function encoded by the minimum. The function space of interest thus becomes a space functions with low TV norm. By computing covering numbers of this space, uniform generalization bounds are derived for this restricted space of functions.

Strengths

The paper is very well written. The proofs in the appendix section are also well written and easy to follow. The minima studied are non-interpolating, unlike the minima studied by a large portion of prior literature (although this work is in the univariate setting). Moreover, although the fact that large stepsizes bias towards solutions with fewer knots appears in prior work, this paper explores this interesting fact further in the noisy-label setting. The paper provides some avenues to understanding generalization when over-fitting is not benign.

Weaknesses

- *The "optimized" assumption*: This assumption which effectively assumes that there exists wide and low minima is somewhat justified in lines 267-270, but I believe it is reasonable to disagree with the statement that it is mild. The experiments provided as evidence for the assumption are conducted with really large noise levels sigmas that are of the same order as $f$. The authors empirical argument could be strengthened by considering noise levels that do not dwarf the signal. It seems difficult to believe that a smooth $f$ can attain a training loss that is lower than $f_0$ when no assumption on the smoothness of $f_0$ is made. The authors should clarify why a condition on $f_0$ is not necessary for this assumption to be mild. - The refined bounds with underparametrization: The bounds are refined by considering under parametrized settings with the width $k$ being smaller than the number of data points but the authors do not discuss its impact on the "optimized" assumption. Why would an underparametrized network be able to attain train losses smaller than a ground truth on which no assumptions are made ? I believe there are trade-offs between $\eta$, the level of underparametrization and the attained train loss that are not mentioned sufficiently by the authors. - A final minor weakness: The work extends [Mulayoff et al 2021] in limited ways but does not exploit the added generalization to derive new conclusions. That stepsize affects the number of knots was explored before. There is now a sigma appearing in the smoothness bounds and an MSE term, yet there are no comments on what is added by their extension. The interplay between sigma and eta is not discussed, for instance, does having a large eta but high noise means that stable functions need no have a low TV norm or is the upper bound loose? I believe such discussions could strengthen their work as some results, namely Thm 4.3, are straightforward applications of generalization bounds for spaces with known metric entropy.

Questions

The weight function $g$ is inherited from [Mulayoff et al 2021] but would really benefit from some more clarifications, even for completeness. It is clear that the interval $\mathcal{I}$ must be introduced to remove the weighing and allow the authors to use covering numbers for bounded TV norm spaces. Could the authors explain why it arises in theorem 4.1 ?

Rating

6

Confidence

3

Soundness

4

Presentation

4

Contribution

3

Limitations

The limitations are discussed.

Summary

This paper studies the generalizability of stable local minima in univariate regression with shallow ReLU networks. Along the way, the authors provide a bound on a weighted total variation norm of networks corresponding to stable solutions which in turn provide a tighter generalization bound.

Strengths

The paper is well written and provides a novel way to obtain generalization bounds using minima stability theory.

Weaknesses

It is not clear to me why the experiments (e.g. figure 3) corroborate the theoretical results in the paper. To illustrate my point, suppose for a given $\eta > 0$, a solution $\theta^*$ is linearly stable. Now suppose one trains a ReLU NN with a different learning rate $\alpha > 0$ and it converges to $\theta^*$. When considering the generalization gap of the solution $\theta^*$ (e.g. theorem 4.4), only the constant $\eta$ is relevant. The point is that the actual learning rate of GD that was used to find stable solutions is irrelevant to the generalizability of the solution (at least in the context of theorem 4.4) . Hence, I feel that more compelling experiments would study generalization errors of stable local minima that are stable for varying $\eta$. In this vein, I feel that the statement, “Meanwhile, the dependence $\frac 1 {\eta}$ ...” (line 296 -298) to be inexact as it is only supported by the experiments.

Questions

In the context of theorem 4.3 and 4.4, should $n_I$ be the number of data points belonging to the interval $I \subset [-x_{max}, x_{max}|$ instead of the length of the interval $|I|$? Otherwise, the RHS of the bound in theorem 4.4 seems to be non-vanishing as $x_{max} $ is at least on the order of the length of the interval $|I|$. Can the authors also provide some intuition on how the weight function $g$ was chosen?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have adequately addressed the limitations.

Thank you for the rebuttal.

I thank the authors for their rebuttal and keep my score.

Any further questions / comments?

Thanks again for your time in reviewing our paper! We have addressed your technical questions above. Could you kindly let us know if our rebuttal satisfactorily resolved your concerns? We would love an opportunity to address any further questions and comments you may have before the author discussion period expires.

Any follow-up questions before discussion period expires?

Thanks again for your time in reviewing our paper! We have addressed your technical questions above. Could you kindly let us know if our rebuttal satisfactorily resolved your concerns? We would love an opportunity to address any further questions and comments you may have before the author discussion period expires.

Any further questions?

Thanks again for your time in reviewing our paper! We have addressed your technical questions above and shared some perspectives. Could you kindly let us know if our rebuttal satisfactorily resolved your concerns? We would love an opportunity to address any further questions and comments you may have before the author discussion period expires.