Phase Transitions for Feature Learning in Neural Networks

TL;DR

This paper investigates the phase transition thresholds in neural network feature learning, analyzing how the network's ability to learn low-dimensional representations depends on data and architecture parameters.

Contribution

It derives a new threshold $oldsymbol{ ext{delta}}_{ ext{NN}}$ for two-layer neural networks, extending previous theoretical results to understand learning dynamics.

Findings

Identifies a phase transition in the Hessian spectrum during training.

Derives a threshold $oldsymbol{ ext{delta}}_{ ext{NN}}$ for neural network feature learning.

Provides insights into how network architecture influences learning dynamics.

Abstract

According to a popular viewpoint, neural networks learn from data by first identifying low-dimensional representations, and subsequently fitting the best model in this space. Recent works provide a formalization of this phenomenon when learning multi-index models. In this setting, we are given $n$ i.i.d. pairs $({\boldsymbol x}_i,y_i)$, where the covariate vectors ${\boldsymbol x}_i\in\mathbb{R}^d$ are isotropic, and responses $y_i$ only depend on ${\boldsymbol x}_i$ through a $k$-dimensional projection ${\boldsymbol \Theta}_*^{{\sf T}}{\boldsymbol x}_i$. Feature learning amounts to learning the latent space spanned by ${\boldsymbol \Theta}_*$. In this context, we study the gradient descent dynamics of two-layer neural networks under the proportional asymptotics $n,d\to\infty$, $n/d\to\delta$, while the dimension of the latent space $k$ and the number of hidden neurons $m$ are kept fixed. Earlier work establishes that feature learning via polynomial-time algorithms is possible if $\delta> \delta_{\text{alg}}$, for $\delta_{\text{alg}}$ a threshold depending on the data distribution, and is impossible (within a certain class of algorithms) below $\delta_{\text{alg}}$. Here we derive an analogous threshold $\delta_{\text{NN}}$ for two-layer networks. Our characterization of $\delta_{\text{NN}}$ opens the way to study the dependence of learning dynamics on the network architecture and training algorithm. The threshold $\delta_{\text{NN}}$ is determined by the following scenario. Training first visits points for which the gradient of the empirical risk is large and learns the directions spanned by these gradients. Then the gradient becomes smaller and the dynamics becomes dominated by negative directions of the Hessian. The threshold $\delta_{\text{NN}}$ corresponds to a phase transition in the spectrum of the Hessian in this second phase.