Phase Transitions for Feature Learning in Neural Networks 

Abstract: 

According to a popular viewpoint, neural networks learn from data by first identifying effective low-dimensional representations, and subsequently fitting the best model in this low-dimensional space. A sequence of recent works provides a rigorous formalization of this phenomenon when learning multi-index models.

In this setting, we are given n i.i.d. pairs of covariates vectors and responses, where the d-dimensional inputs are isotropic, and responses only depend on the inputs through a low-dimensional projection. Feature learning amounts to learning the latent projection space.

In this context, we study the gradient descent dynamics of two-layer neural networks under the proportional asymptotics. Earlier work establishes that feature learning via polynomial-time algorithms is possible if the sample-to-dimension ratio is larger than a threshold depending on the data distribution, and is impossible below this threshold. Here we derive an analogous threshold for a class of two-layer networks, thus allowing us to study the dependence of learning dynamics on the target function, loss, activation, width, and initialization. The threshold corresponds to a phase transition in the spectrum of the Hessian after the first phase of training, and provides a quantitative explanation for empirical phenomena such as grokking.

Bio: 

Zihao Wang is a second-year PhD student in Stanford Mathematics, advised by Prof. Andrea Montanari. He previously earned his B.S. in Mathematics from Peking University in 2024, where he was advised by Prof. Lei Wu. He was named a Weiming Scholar at Peking University in 2024. In 2023, he visited Princeton University, hosted by Prof. Jason D. Lee. His research aims to understand (and optimize) the mechanisms behind modern machine learning with tools from probability, statistics, and optimization. His current focus includes feature learning and topics in post-training such as reasoning.