Mean-field Analysis of Piecewise Linear Solutions for Wide ReLU Networks

Abstract

Understanding the properties of neural networks trained via stochastic gradient descent (SGD) is at the heart of the theory of deep learning. In this work, we take a mean-field view, and consider a two-layer ReLU network trained via SGD for a univariate regularized regression problem. Our main result is that SGD is biased towards a simple solution: at convergence, the ReLU network implements a piecewise linear map of the inputs, and the number of "knot" points -i.e., points where the tangent of the ReLU network estimator changes -between two consecutive training inputs is at most three. In particular, as the number of neurons of the network grows, the SGD dynamics is captured by the solution of a gradient flow and, at convergence, the distribution of the weights approaches the unique minimizer of a related free energy, which has a Gibbs form. Our key technical contribution consists in the analysis of the estimator resulting from this minimizer: we show that its second derivative vanishes everywhere, except at some specific locations which represent the "knot" points. We also provide empirical evidence that knots at locations distinct from the data points might occur, as predicted by our theory.

Intrinsic dimensionality and generalization properties of the $\mathcal{R}$-norm inductive bias

“…These norms are also connected to the implicit biases of optimization methods for training neural networks. In the context of univariate functions, the dynamics of gradient descent was shown to be biased towards (adaptive) linear or cubic spline depending on the optimization regime [24,40,45], and these results have been partially extended to the multivariate case [18]. For classification problems, the implicit bias of gradient descent was connected to a variational problem related to R-norm with margin constraints on the data [2], which we explore in our empirical study.…”

Stochastic Langevin Differential Inclusions with Applications to Machine Learning

“…Mean-field analyses seek to explain the uncanny generalization abilities of DNNs by considering the distributional dynamics of idealized networks with infinitely many hidden layer neurons by considering both the stochastic equations corresponding to said dynamics, as well as Fokker-Planck-type PDEs modeling the flow of the distribution of the weights in the network. Analyses along these lines in the contemporary literature include [31,39]. Although this line of works does use limiting arguments involving stochastic and distributional dynamics, they do not directly consider the Langevin differential inclusion and the potential PDE solution for the distribution, as we do here.…”