Learning Dynamics of Wide Neural Networks: A Mesoscopic Description Reveals Feature Specialization Mechanisms

Introduction: Understanding Learning Dynamics in Wide Networks

Understanding how stochastic gradient descent (SGD) learns internal representations in wide neural networks is a central problem in machine learning. A recent paper, Mesoscopic learning dynamics of extensive-width neural networks, explores this in depth. The study focuses on two-layer neural networks learning a target function represented by a network of the same architecture, where network width and input dimension grow proportionally. In this regime, conventional methods for describing learning dynamics require tracking a large number of overlaps—quantities measuring the correlation between learned features and target features—and the number of these overlaps diverges, making it difficult to describe with a small number of order parameters.

Wide neural networks are a class of idealized models commonly used in theoretical analysis, referring to networks whose hidden layer width is much larger than the number of training samples or the input dimension. In the limit (width tending to infinity), network behavior is often proven equivalent to some kernel method, such as the neural tangent kernel (NTK), in which case the network undergoes almost no feature learning during training and internal representations remain essentially unchanged. However, finite but "wide" networks occupy a more realistic intermediate regime: they are overparameterized enough to bring optimization benefits, yet retain the ability for features to evolve during training. The extensive-width setting considered in this paper further requires network width and input dimension to grow at the same rate, so that neither width nor dimension alone tends to infinity but together determine the scaling behavior of learning dynamics. This regime differs both from finite-dimensional analyses of fixed-width networks and from the strict infinite-width limit, and thus requires new theoretical tools to characterize its representation learning process.

Overlaps are order parameters used in statistically physics-inspired learning theory to quantify the degree of alignment between "learned features" and "target features." Specifically, if the inner product (after normalization) between some direction of a first-layer weight vector and some hidden feature direction of the target network is large, the network has "discovered" and specialized to that feature. In traditional analyses, when width and dimension are fixed and small, only a few overlaps need to be tracked to describe the dynamics. But in the extensive-width case, the number of overlap pairs that must be tracked diverges as width grows, so closure equations based on a few order parameters no longer hold. This is precisely the direct motivation for the mesoscopic description proposed in this paper: abandon tracking individual overlaps and instead characterize their overall distribution.

Mesoscopic Description: From Diverging Overlaps to Distributional Characterization

To address this challenge, the researchers propose a mesoscopic description that understands the learning process by characterizing the distribution of overlaps. This approach avoids tracking the divergence of individual overlaps and instead focuses on their statistical properties. In the early stage of learning, they find that a Gaussian description of overlap statistics simplifies the learning dynamics into a few collective variables that track the learning of the linear and quadratic components of the target function.

Key insight: Gaussian closure holds in the early stage, enabling dimensionality reduction of high-dimensional learning dynamics.

This finding implies that in extensive-width networks, the learning process can initially be effectively described by a few macroscopic variables, greatly simplifying the analysis of learning dynamics.

Specialization Time Scale and Quadratic Shielding Mechanism

The researchers further estimate the time scale of specialization—the moment when Gaussian closure breaks down and individual learned features begin to strongly correlate with specific target features. They find that in extensive-rank networks, specialization is accelerated by a quadratic shielding mechanism: learning the linear and quadratic components of the target function filters gradient fluctuations, thereby accelerating specialization.

This mechanism explains why networks can rapidly transition from a collective learning phase to a feature discovery phase. Quadratic shielding means that once the network captures the low-order statistics of the target, noise in the gradient is suppressed, focusing learning more on directions relevant to target features.

Specialization refers to the process by which a network transitions from "all hidden neurons having only weak, diffuse correlations with target features" to "some neurons strongly aligning with specific target features." In the early collective stage of learning, neurons behave approximately symmetrically, and overlap statistics can be described by a Gaussian distribution; as training proceeds, this symmetry breaks and individual neurons begin to specialize to specific directions of the target function. The specialization time scale characterizes the critical moment when Gaussian closure fails, and it depends on scaling parameters such as the rank of the target function, the activation function, and the ratio of width to dimension. Understanding this time scale is crucial for determining when simplified statistical descriptions can be safely used and when individual feature discovery must be taken into account.

Activation Functions and Specialization Patterns

Finally, the researchers describe how activation functions determine specialization patterns and characterize the subsequent refinement of learned features. Different nonlinear activation functions lead to different specialization patterns, which affect the quality of features the network ultimately learns.

The choice of activation function not only affects the network's expressive power but also regulates the feature specialization pathways in learning dynamics.

This finding extends the role of activation functions from merely introducing nonlinearity to regulating learning dynamics.

Significance and Outlook

This study shows how collective learning gives way to individual feature discovery and clarifies when large neural networks can adopt simple statistical descriptions during learning. The mesoscopic description provides a new theoretical tool for understanding the learning dynamics of wide networks and may have a profound impact on the future design and analysis of neural networks.

Although the study focuses on two-layer networks, its methodology may generalize to deeper networks. Future work can explore the applicability of the mesoscopic description to more complex architectures and how these insights can be used to improve training algorithms.

One-line Summary

Mesoscopic description reveals the transition mechanism from collective Gaussian phase to feature specialization in wide network learning.

Paragraph Summary

This paper studies the dynamics of two-layer neural networks learning a same-architecture target function when width and input dimension grow proportionally. Traditional methods require tracking a diverging number of overlap order parameters, making closure difficult. The authors propose a mesoscopic description that characterizes the distribution of overlaps rather than individual values. They find that early-stage overlap statistics satisfy Gaussian closure, reducing dynamics to a few collective variables tracking the linear and quadratic components of the target; they then estimate the specialization time scale and reveal that a quadratic shielding mechanism in extensive-rank networks accelerates specialization by filtering gradient fluctuations. Finally, they clarify how activation functions determine specialization patterns and subsequent feature refinement pathways, providing new theoretical tools for understanding representation learning in wide networks.