Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Phase Transitions for Feature Learning in Neural Networks

Andrea Montanari, Zihao Wang

TL;DR

The paper analyzes how two-layer neural networks learn latent low-dimensional structure in multi-index models under proportional growth, showing a Hessian-driven phase transition that governs feature learning. It introduces a Hessian-based threshold $\delta_{\mathrm NN}$, separating easy directions learned in $O(1)$ steps from hard directions that are learned once negative eigenvalues aligned with the hard subspace emerge. Using discrete-time dynamical mean-field theory (DMFT), the authors characterize the bulk Hessian spectrum and derive a finite-dimensional outlier condition whose solutions determine when feature learning begins, both at finite time and in the long-time limit. The results explain phenomena like grokking and quantify how network architecture and training details affect the feature-learning threshold, providing a principled spectral framework for understanding learning dynamics in high dimensions.

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 Θ}_*^{\sf T}{\boldsymbol x}_i$. Feature learning amounts to learning the latent space spanned by ${\boldsymbol Θ}_*$. 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δ$, 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 $δ> δ_{\text{alg}}$, for $δ_{\text{alg}}$ a threshold depending on the data distribution, and is impossible (within a certain class of algorithms) below $δ_{\text{alg}}$. Here we derive an analogous threshold $δ_{\text{NN}}$ for two-layer networks. Our characterization of $δ_{\text{NN}}$ opens the way to study the dependence of learning dynamics on the network architecture and training algorithm. The threshold $δ_{\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 $δ_{\text{NN}}$ corresponds to a phase transition in the spectrum of the Hessian in this second phase.

Phase Transitions for Feature Learning in Neural Networks

TL;DR

The paper analyzes how two-layer neural networks learn latent low-dimensional structure in multi-index models under proportional growth, showing a Hessian-driven phase transition that governs feature learning. It introduces a Hessian-based threshold , separating easy directions learned in steps from hard directions that are learned once negative eigenvalues aligned with the hard subspace emerge. Using discrete-time dynamical mean-field theory (DMFT), the authors characterize the bulk Hessian spectrum and derive a finite-dimensional outlier condition whose solutions determine when feature learning begins, both at finite time and in the long-time limit. The results explain phenomena like grokking and quantify how network architecture and training details affect the feature-learning threshold, providing a principled spectral framework for understanding learning dynamics in high dimensions.

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 i.i.d. pairs , where the covariate vectors are isotropic, and responses only depend on through a -dimensional projection . Feature learning amounts to learning the latent space spanned by . In this context, we study the gradient descent dynamics of two-layer neural networks under the proportional asymptotics , , while the dimension of the latent space and the number of hidden neurons are kept fixed. Earlier work establishes that feature learning via polynomial-time algorithms is possible if , for a threshold depending on the data distribution, and is impossible (within a certain class of algorithms) below . Here we derive an analogous threshold for two-layer networks. Our characterization of opens the way to study the dependence of learning dynamics on the network architecture and training algorithm. The threshold 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 corresponds to a phase transition in the spectrum of the Hessian in this second phase.
Paper Structure (91 sections, 37 theorems, 335 equations, 11 figures)

This paper contains 91 sections, 37 theorems, 335 equations, 11 figures.

Table of Contents

  1. Introduction
  2. Main results: An informal overview
  3. Background
  4. Hard directions are not learned within $O(1)$ steps
  5. Hessian phase transitions after gradient descent
  6. Outline of the paper.
  7. Related literature
  8. Optimal feature learning.
  9. Feature learning in neural networks.
  10. Spiked Models.
  11. Summary of notations
  12. Results
  13. Problem setup and assumptions
  14. Preliminaries: Discrete time DMFT
  15. Result I: Hard directions cannot be learned in $O(1)$ time
  16. ...and 76 more sections

Key Result

Lemma 2.1

Let $\phi$ be a test function which is locally Lipschitz continuous and has at most quadratic growth. Under proportional scaling $n/d:=\delta_n \rightarrow \delta \in (0,+\infty)$ and Assumption assumption_regularity, we have where $(V(0),\dots,V(t),V_*,\varepsilon)$ is the discrete-time DMFT process. Furthermore, the joint distribution satisfies for every measurable $\psi(\cdot)$ such that the

Figures (11)

  • Figure 1: Empirical success probabilities and predicted thresholds for learning a single-index model with link function $h(z,{\varepsilon})=z^2$ (noiseless phase retrieval). The model is a single-neuron network with ${\sf GeLU}$ activation, trained via gradient descent; we declare success when the correlation between the learned feature and the target direction exceeds $1/2$. Orange curves correspond to optimal spectral initialization, blue curves to random initialization; within each color family, lighter to darker shades indicate increasing dimensions $d=1000, 2000, 3000, 4000$. Dots are success rates over $60$ independent trials (error bars show the standard deviation), and solid lines are sigmoid fits.
  • Figure 2: Grokking phenomenon for a ${\sf GeLU}$ neuron learning the noiseless phase retrieval problem ($d=5000$, $\delta=n/d=17.5$, single run). Here $\rho(t) := |\langle{\boldsymbol \theta}(t),{\boldsymbol \theta}_*\rangle|/\|{\boldsymbol \theta}(t)\|$ denotes the cosine similarity between the learned and target parameters. Light-colored markers show data points; solid lines connect them for visual clarity. Initially, training loss decreases while test loss stays near its initial value and misalignment remains high (overfitting phase). Subsequently, the misalignment drops to zero---indicating successful learning of the signal direction---while both losses plateau at around $0.6$ and the generalization error drops to zero (grokking phase).
  • Figure 3: Cartoon of feature learning via GD on the empirical risk landscape. Within $O(1)$ steps, the network learns easy directions (blue arrow) to a saddle point. Escaping this point requires moving along hard directions (green arrow), governed by the Hessian's negative eigenvalues.
  • Figure 4: Phase transitions for ${\sf GeLU}$ with single-neuron learner. Left: Success rate with sigmoid fits; each data point shows mean $\pm$ standard deviation across 60 independent runs. Right: Final correlation with median and 30th/70th percentile error bars (smooth splines fitted). Both metrics exhibit sharp phase transitions matching the predicted threshold $\delta^*\approx 6.0$ (dash-dot line).
  • Figure 5: Phase transitions for $\mathsf{Quad}$ with single-neuron learner. Left: Success rate with sigmoid fits; each data point shows mean $\pm$ standard deviation across 60 independent runs. Right: Final correlation with median and 30th/70th percentile error bars (smooth splines fitted). Both metrics exhibit sharp phase transitions matching the predicted threshold $\delta^*\approx 3.6$ (dash-dot line).
  • ...and 6 more figures

Theorems & Definitions (60)

  • Definition 1.1
  • Remark 2.1
  • Lemma 2.1
  • Lemma 2.2
  • Lemma 2.3
  • Lemma 2.4
  • Theorem 1
  • Lemma 2.5
  • Theorem 2
  • Remark 2.2
  • ...and 50 more