Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Learning time-scales in two-layers neural networks

Raphaël Berthier, Andrea Montanari, Kangjie Zhou

TL;DR

This work analyzes gradient-flow dynamics of wide two-layer neural networks under a single-index data model, revealing a structured, multiscale learning process. By reducing the high-dimensional dynamics to a mean-field, $d$-independent flow and applying singular perturbation theory, the authors show that learning proceeds incrementally along Hermite polynomial components of the target function, with distinct time scales controlled by the learning-rate parameter $\varepsilon$. They establish a connection to mean-field PDEs, demonstrate the existence of a zero-risk global infimum under generic conditions, and provide finite-sample SGD guarantees that align with the canonical learning order. The results offer a principled picture of plateaus, waterfalls, and incremental generalization in deep learning, with implications for understanding implicit biases and designing training protocols that exploit feature learning in early phases.

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.

Learning time-scales in two-layers neural networks

TL;DR

This work analyzes gradient-flow dynamics of wide two-layer neural networks under a single-index data model, revealing a structured, multiscale learning process. By reducing the high-dimensional dynamics to a mean-field, -independent flow and applying singular perturbation theory, the authors show that learning proceeds incrementally along Hermite polynomial components of the target function, with distinct time scales controlled by the learning-rate parameter . They establish a connection to mean-field PDEs, demonstrate the existence of a zero-risk global infimum under generic conditions, and provide finite-sample SGD guarantees that align with the canonical learning order. The results offer a principled picture of plateaus, waterfalls, and incremental generalization in deep learning, with implications for understanding implicit biases and designing training protocols that exploit feature learning in early phases.

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.
Paper Structure (55 sections, 18 theorems, 292 equations, 4 figures)

This paper contains 55 sections, 18 theorems, 292 equations, 4 figures.

Table of Contents

  1. Introduction
  2. Theory $\#1$: Dynamics near singular points.
  3. Theory $\#2$: Linear networks.
  4. Theory $\#3$: Kernel regime.
  5. Notations.
  6. Setting and canonical learning order
  7. Further related work
  8. The large-network, high-dimensional limit
  9. Reduction to $d$-independent flow
  10. Elimination of the products $\langle u_i, u_j \rangle$
  11. Connection with mean field theory
  12. A general formulation
  13. Numerical solution
  14. Timescales hierarchy in the gradient flow dynamics
  15. Notations.
  16. ...and 40 more sections

Key Result

Proposition 1

Assume that $\sigma$ is Lipschitz continuous and generic in the following sense: the decomposition of $\sigma$ into Hermite polynomials does not have any coefficient equal to $0$. For any Lipschitz function $\varphi:{\mathbb R}\to{\mathbb R}$, $\|u_*\|_2=1$, and $x \sim {\mathcal{N}}(0, I_d)$ such t

Figures (4)

  • Figure 1: Cartoon illustration of the evolution of the population risk within the canonical learning order of Definition \ref{['def:StandardScenario']}.
  • Figure 2: Simulation of the simplified neuron dynamics of Eqs. \ref{['eq:coupled_MF']}, with the target function of Eq. \ref{['eq:ExampleSimulations']} and ReLU activations. We use learning rate ratios $\varepsilon = 10^{-3}$ (left) and $\varepsilon = 10^{-6}$ (right) and we use $m=10$ neurons. First two rows: evolution of the risk $\mathscrsfs{R}_{\hbox{\tiny\rm mf}}$ of Eq. \ref{['eq:RmfDef']}, in linear and log-scales. Third row: evolution of the first three terms of the sum of \ref{['eq:RiskHermite']}.
  • Figure 3: Same simulation as in Figure \ref{['fig:simulations']} (b). In these plots, we show the evolution of the $a_i$ and the $s_i$ for $i\in\{1, \dots, m\}$ following a discretization of Eqs. \ref{['eq:coupled_MF']}.
  • Figure 4: Comparison between the simplified neuron dynamics \ref{['eq:coupled_MF']} (MF) and projected gradient descent for the two-layer neural network \ref{['eq:First-NNET']} (NN), with the same target function and activation as the simulations in Figure \ref{['fig:simulations']}. We use four different combinations of (learning rare ratio, network width): $(\varepsilon, m) = (1, 10)$ (first row), $(\varepsilon, m) = (1, 50)$ (second row), $(\varepsilon, m) = (10^{-3}, 10)$ (third row), and $(\varepsilon, m) = (10^{-3}, 50)$ (fourth row). Left panel: evolution of the risk $\mathscrsfs{R}_{\hbox{\tiny\rm mf}}$ for NN and MF on a logarithmic scale. Right panel: evolution of the first three components of $\mathscrsfs{R}_{\hbox{\tiny\rm mf}}$ (constant, linear, and quadratic) for NN and MF.

Theorems & Definitions (34)

  • Remark 2.1
  • Definition 1
  • Proposition 1
  • Remark 2.2
  • Remark 4.1
  • Proposition 2: Reduction to $d$-independent flow
  • Corollary 1
  • Proposition 3: Reduction to flow in $\mathbb{R}^{2m}$
  • Proposition 4
  • Remark 4.2
  • ...and 24 more