On the global convergence of gradient flow for wide shallow models beyond homogeneous nonlinearities

Abstract

A surprising phenomenon in the training of neural networks is the ability of gradient descent to find global minimizers of the training loss despite its non-convexity. Following earlier work, we investigate this behavior for wide shallow models. Existing global convergence results primarily concern models with positively one-homogeneous nonlinearities, such as ReLU activations, and models with scalar output weights and bounded nonlinearities, such as sigmoid activations. We study a broader class of models, including multi-head attention layers and two-layer networks with bounded or asymptotically positively one-homogeneous activations and vector output weights. Building upon [Chizat and Bach, 2018], we prove that, in the limit of many hidden neurons or attention heads, non-global minimizers of the training loss are unstable under mean-field gradient flow dynamics by constructing "escape regions" in the parameter space. Our global convergence statements are conditional in the following sense: if the mean-field gradient flow converges in W2, then its limit must be a global minimizer. We revisit the bounded nonlinearity, scalar-output setting of [CB18], giving an escape region construction adapted to unbounded nonlinear parameter domains. We also propose new constructions for nonlinearities with at most linear growth under a non-degeneracy assumption and for asymptotically positively one-homogeneous nonlinearities. Finally, we show the well-posedness and stability estimates for the mean-field training dynamics under sub-Gaussian initializations.