Gradient flow dynamics of shallow ReLU networks for square loss and orthogonal inputs

Etienne Boursier (1) , Loucas Pillaud-Vivien (2, 1) , Nicolas Flammarion (3, 1)

1 EPFL - Ecole Polytechnique Fédérale de Lausanne
2 ENPC - École des Ponts ParisTech
3 LIENS - Laboratoire d'informatique de l'école normale supérieure

The training of neural networks by gradient descent methods is a cornerstone of the deep learning revolution. Yet, despite some recent progress, a complete theory explaining its success is still missing. This article presents, for orthogonal input vectors, a precise description of the gradient flow dynamics of training one-hidden layer ReLU neural networks for the mean squared error at small initialisation. In this setting, despite non-convexity, we show that the gradient flow converges to zero loss and characterise its implicit bias towards minimum variation norm. Furthermore, some interesting phenomena are highlighted: a quantitative description of the initial alignment phenomenon and a proof that the process follows a specific saddle to saddle dynamics.

Domains: Machine Learning [cs.LG]; Machine Learning [stat.ML]

Etienne Boursier, Loucas Pillaud-Vivien, Nicolas Flammarion. Gradient flow dynamics of shallow ReLU networks for square loss and orthogonal inputs. NeurIPS 2022 - 36th International Conference on Neural Information Processing Systems, Nov 2022, New Orleans, United States. ⟨hal-04105187⟩