Mesoscopic learning dynamics of extensive-width neural networks
Abstract
Understanding how stochastic gradient descent learns internal representations in wide neural networks trained on rich target functions is a central problem in machine learning. We study this question in two-layer neural networks learning target functions represented by networks of the same architecture, with arbitrary nonlinearity, and with network width and input dimension growing proportionally. In this regime, standard approaches for describing the learning dynamics require tracking a diverging number of overlaps, which measure correlations between learned and target features, leaving no obvious description in terms of a few order parameters. Here, we develop a mesoscopic description of learning by characterizing the distribution of overlaps. In the early stages of learning, we find that a Gaussian description of overlap statistics reduces the learning dynamics to a few collective variables that track learning of the target function's linear and quadratic components. We then estimate the timescale of specialization, when Gaussian closure breaks down and individual learned features begin to strongly correlate with specific target features. We find that specialization in extensive-rank networks is accelerated via a quadratic shielding mechanism, whereby learning the linear and quadratic components of the target function filters gradient fluctuations. We finally describe how the activation function determines the pattern of specialization and characterize the subsequent refinement of learned features. Together, these results show how collective learning gives way to individual feature discovery and clarify when learning in large neural networks admits a simple statistical description.
The following analysis is generated by AI and may contain errors. It is separate from the authors’ abstract and is not peer review.
Research summary
The paper develops a mesoscopic theory for learning dynamics in extensive-width neural networks, deriving closed-form equations for the pre-specialization regime and predicting the timescale of feature specialization via a quadratic shielding mechanism. By bridging the gap between narrow-network order parameter analyses and infinite-width mean-field theories, it provides a rigorous framework for understanding how collective learning transitions to individual feature discovery in high-dimensional settings.
Methodology
The paper introduces a "mesoscopic" description of learning dynamics for extensive-width (proportional width and input dimension) two-layer neural networks. It bridges the gap between the Saad-Solla approach (which fails in the extensive limit due to diverging order parameters) and mean-field/NTK theories (which often miss feature learning details or require infinite width). The core method involves tracking the distribution of overlaps between teacher and student neurons via their empirical moments. The authors derive a "Gaussian closure" regime where the dynamics close on a few collective variables (linear and quadratic components), and then analyze the breakdown of this closure to predict the onset of "specialization" (individual feature alignment). The derivation of the "quadratic shielding" mechanism, which explains how learning low-order components filters gradient noise to accelerate specialization, is a particularly insightful theoretical contribution.
Results & evaluation
The paper provides strong numerical validation of the theoretical predictions. It compares the derived ODEs for the order parameters ($Q_1, Q_2$) against simulations of SGD for various activation functions (ReLU, tanh, ELU, etc.). The agreement between the theoretical loss plateaus and the specialization timescales is precise. The experiments confirm the existence of the distinct learning phases (collective alignment, specialization, refinement) and validate the scaling laws for specialization time with respect to network width and activation function properties.
Reproducibility
The paper is highly reproducible from a theoretical standpoint, as it provides detailed derivations in the appendices. The experimental setup is standard (teacher-student setup with Gaussian inputs), and the specific activation functions and network sizes are clearly defined. While code is not explicitly linked in the text provided, the mathematical framework is sufficient for replication.
Limitations
The analysis is restricted to two-layer networks with fixed second-layer weights (committee machines). Extending this to deeper networks or trainable readouts is left for future work. The "Gaussian closure" is an approximation that breaks down during specialization, and the paper provides estimates for the transition rather than a fully closed dynamic description of the specialization phase itself. The assumption of online SGD (one-pass) limits direct applicability to multi-epoch training regimes common in practice, though the authors argue relevance to LLM pre-training.
Broader impact
This work contributes to the fundamental understanding of how neural networks learn representations in high-dimensional regimes. It provides a rigorous framework for analyzing feature learning that goes beyond the NTK and mean-field limits. The insights into "quadratic shielding" and the role of activation functions in specialization patterns could inform the design of better initialization schemes or learning rate schedules for wide networks. It connects statistical physics methods with modern deep learning theory, offering tools to predict training dynamics and generalization errors in the extensive-width limit. The paper develops a mesoscopic theory for learning dynamics in extensive-width neural networks, deriving closed-form equations for the pre-specialization regime and predicting the timescale of feature specialization via a quadratic shielding mechanism. By bridging the gap between narrow-network order parameter analyses and infinite-width mean-field theories, it provides a rigorous framework for understanding how collective learning transitions to individual feature discovery in high-dimensional settings.
Author of this paper? Send a correction.