Spectral phase transitions and trainability in neural network learning dynamics

Technical Summary: Spectral Phase Transitions and Trainability in Neural Network Learning Dynamics

Problem Statement
The emergence of low-dimensional structures within the spectra of neural network weight matrices is a well-documented empirical feature of trained models. However, the dynamical origin of this phenomenon during the learning process remains an open problem. While empirical studies have identified distinct dynamical regimes in neural network trainability—characterized by sharp performance changes relative to hyperparameters like step size (learning rate) and initial weight variance—a unified dynamical theory linking these regimes to spectral structure is absent. Specifically, the relationship between stochastic optimization, the formation of low-dimensional spectral structures, and the trainability of the network requires a theoretical framework that explains how informative representations emerge from random initializations.

Methodology
The authors propose formulating neural network training as the stochastic evolution of an initially random matrix ensemble, driven by Stochastic Gradient Descent (SGD) updates. The core theoretical framework relies on Random Matrix Theory (RMT), specifically the Baik-Ben Arous-Péché (BBP) transition, which describes the detachment of isolated eigenvalues from a random spectral bulk.

The methodology proceeds through several analytical and numerical stages:

1. Solvable Linear Model: The authors first analyze a linear teacher-student model with a rank-1 teacher matrix. In this setting, the gradient update is treated as a finite-rank perturbation to an initial random weight matrix. The weight covariance matrix is updated analytically, allowing for the derivation of time-dependent effective variance and signal strength.
2. Phase Diagram Construction: By varying the step size ($\epsilon$) and initial weight variance ($\sigma_W^2$), the authors construct a trainability phase diagram. This diagram identifies critical boundaries: an optimal step size ($\epsilon_o$) separating monotonic convergence from oscillatory/unstable dynamics, and a critical step size ($\epsilon_c$) marking the onset of the BBP transition.
3. Extensions to Nonlinearity and Depth: The formalism is extended to nonlinear activation functions (Hyperbolic Tangent, ReLU) and multilayer networks. The authors employ the Gaussian equivalence principle and Wick-Isserlis's theorem to compute effective rescaling factors for the bulk variance and signal strength, demonstrating that the qualitative phase structure persists.
4. Stochastic and Finite-Data Effects: The analysis incorporates finite dataset sizes and stochastic gradients. The authors introduce a load parameter $\alpha = P/D$ (dataset size over input dimension) to study how finite-sample fluctuations deform the spectral bulk and modify the BBP transition boundaries. In the stationary limit of SGD, a fluctuation-dissipation relation is derived, linking the weight covariance to the Hessian of the loss landscape and gradient noise.
5. Empirical Validation: The theoretical predictions are tested on realistic datasets (UTKFace) using deep neural networks. The authors measure test loss and a "self-overlap" order parameter (the alignment of top singular vectors across different random seeds) to map the empirical phase diagram.

Key Contributions and Results

• Dynamical BBP Transition: The paper establishes that the formation of informative representations in neural networks corresponds to a dynamical BBP transition. During training, the competition between the random spectral bulk (inherited from initialization) and the finite-rank gradient signal (driven by SGD) causes isolated eigenvalues to detach from the bulk. The corresponding eigenvectors acquire nonzero alignment with the learning signal.
• Trainability Phase Diagram: The authors define a phase diagram spanned by step size and initial variance, identifying three distinct regimes:
  • Ferromagnetic Phase: Occurs between the critical step size ($\epsilon_c$) and the optimal step size ($\epsilon_o$). Here, a stable BBP transition occurs, leading to the emergence of isolated eigenvalues aligned with the signal. This regime corresponds to successful trainability and low test loss.
  • Disordered Phase: Occurs when the step size is too small ($\epsilon < \epsilon_c$) or in specific large-step regions. The signal remains submerged within the random bulk, preventing spectral alignment and slowing information retrieval.
  • Paramagnetic Phase: Occurs when the step size is too large ($\epsilon > \epsilon_o$). Optimization becomes unstable or oscillatory. While isolated eigenvalues may appear transiently due to strong perturbations, they are not stable and do not align with the signal; the network effectively performs a Principal Component Analysis (PCA) of the input data rather than learning task-specific representations.
• Transient Behavior and Finite Data: The study reveals that the BBP transition can be transient. In the proportional regime (finite $\alpha$), isolated eigenvalues may detach from the bulk for a finite range of step sizes before merging back as the step size increases. Finite-sample fluctuations deform the bulk spectrum, controlling the relative scaling between signal and noise.
• Stationary Spectral Structure: In the stationary limit of SGD, the spectral density of the weight matrix is shown to satisfy a fluctuation-dissipation relation determined by the Hessian and gradient fluctuations. This provides a mechanism for the emergence of heavy-tailed spectra observed in deep networks, linking them to the soft directions of the loss landscape.
• Empirical Confirmation: Numerical simulations on the UTKFace dataset confirm the theoretical phase diagram. The "ferromagnetic" region correlates with low test loss and high self-overlap across different training seeds, indicating the emergence of a shared, low-dimensional signal subspace. In contrast, the "paramagnetic" regime shows high overlap with the principal components of the input data (PCA) but fails to learn task-relevant nonlinear features.

Significance
The paper claims to provide a unified dynamical framework for understanding neural network learning. By interpreting training as a spectral phase transition, it links optimization hyperparameters (step size, initialization) directly to the emergence of low-dimensional spectral structures and trainability. The authors suggest that spectral analysis offers a perspective that connects the stability of optimization, the geometry of the loss landscape, and the formation of representations.

The work posits that the BBP transition is a common mechanism underlying representation learning, where the system undergoes an abrupt change in dynamical behavior as it crosses critical thresholds in parameter space. While the exact analytical treatment is limited to simplified settings (linear models, first-order steps), the authors argue that the qualitative mechanism and the resulting phase diagram persist in nonlinear, multilayer, and realistic settings. The paper concludes that this spectral perspective may serve as an organizing principle for understanding the interplay between stochastic learning dynamics, spectral phase transitions, and the collective behavior of complex learning systems.