Phase Transitions for Feature Learning in Neural Networks

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 xi,yi), where the covariate vectors boldsymbol xi∈ℝd are isotropic, and responses yi only depend on boldsymbol xi through a k-dimensional projection boldsymbol Theta_*^sf Tboldsymbol xi. 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→∞, n/d→δ, 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 δ> δalg, for δalg a threshold depending on the data distribution, and is impossible (within a certain class of algorithms) below δalg. Here we derive an analogous threshold δNN for two-layer networks. Our characterization of δNN opens the way to study the dependence of learning dynamics on the network architecture and training algorithm. The threshold δ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 δNN corresponds to a phase transition in the spectrum of the Hessian in this second phase.