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

TL;DR

This paper develops a new theory showing that gradient descent with large step sizes in univariate ReLU networks converges to smooth, stable minima that generalize well, even in noisy, nonparametric regression settings.

Contribution

It introduces a novel generalization framework based on minima stability for univariate ReLU networks trained with large step sizes, demonstrating near-optimal rates without regularization.

Findings

Gradient descent with fixed large step size converges to smooth minima.

Stable minima have bounded weighted first order total variation.

Achieves near-optimal MSE bounds of rom data.

Abstract

We study the generalization of two-layer ReLU neural networks in a univariate nonparametric regression problem with noisy labels. This is a problem where kernels (\emph{e.g.} NTK) are provably sub-optimal and benign overfitting does not happen, thus disqualifying existing theory for interpolating (0-loss, global optimal) solutions. We present a new theory of generalization for local minima that gradient descent with a constant learning rate can \emph{stably} converge to. We show that gradient descent with a fixed learning rate $\eta$ can only find local minima that represent smooth functions with a certain weighted \emph{first order total variation} bounded by $1/\eta - 1/2 + \widetilde{O}(\sigma + \sqrt{\mathrm{MSE}})$ where $\sigma$ is the label noise level, $\mathrm{MSE}$ is short for mean squared error against the ground truth, and $\widetilde{O}(\cdot)$ hides a logarithmic factor. Under mild assumptions, we also prove a nearly-optimal MSE bound of $\widetilde{O}(n^{-4/5})$ within the strict interior of the support of the $n$ data points. Our theoretical results are validated by extensive simulation that demonstrates large learning rate training induces sparse linear spline fits. To the best of our knowledge, we are the first to obtain generalization bound via minima stability in the non-interpolation case and the first to show ReLU NNs without regularization can achieve near-optimal rates in nonparametric regression.