Learning time-scales in two-layers neural networks

Abstract

Gradient-based learning in multi-layer neural networks displays a number of striking features. In particular, the decrease rate of empirical risk is non-monotone even after averaging over large batches. Long plateaus in which one observes barely any progress alternate with intervals of rapid decrease. These successive phases of learning often take place on very different time scales. Finally, models learnt in an early phase are typically `simpler' or `easier to learn' although in a way that is difficult to formalize. Although theoretical explanations of these phenomena have been put forward, each of them captures at best certain specific regimes. In this paper, we study the gradient flow dynamics of a wide two-layer neural network in high-dimension, when data are distributed according to a single-index model (i.e., the target function depends on a one-dimensional projection of the covariates). Based on a mixture of new rigorous results, non-rigorous mathematical derivations, and numerical simulations, we propose a scenario for the learning dynamics in this setting. In particular, the proposed evolution exhibits separation of timescales and intermittency. These behaviors arise naturally because the population gradient flow can be recast as a singularly perturbed dynamical system.

64 pages, 15 figures, Found Comput Math (2024)

References in corpus (10)

• A Closer Look at Memorization in Deep Networks
• Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit
• The Landscape of Empirical Risk for Non-convex Losses
• SGD on Neural Networks Learns Functions of Increasing Complexity
• High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation
• Learning Single-Index Models with Shallow Neural Networks
• Neural Networks can Learn Representations with Gradient Descent
• The merged-staircase property: a necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networks
• The Three Stages of Learning Dynamics in High-Dimensional Kernel Methods
• Model, sample, and epoch-wise descents: exact solution of gradient flow in the random feature model

Cited by in corpus (1)

• Phase transitions in the mini-batch size for sparse and dense two-layer neural networks