Mesoscopic learning dynamics of extensive-width neural networks

Shu Zhou; Jean Barbier; Sebastian Goldt

stat.MLOct 16, 2026

Mesoscopic learning dynamics of extensive-width neural networks

Authors: Shu Zhou, Jean Barbier, Sebastian Goldt

Organizations: International School for Advanced Studies (SISSA) Via Bonomea 265, 34136 Trieste, Italy · The Abdus Salam International Centre for Theoretical Physics (ICTP) Strada Costiera 11, 34151 Trieste, Italy

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.

Figures & tables

Figure 1: The stages of learning in an extensive-width neural network. Schematic view of the learning regimes in the extensive-width teacher–student model. Red and blue arrows denote teacher and student hidden neurons, respectively, and the sphere represents the space of weight directions. The colored line schematically illustrates the decrease in generalization error during training. (a) At initialization, student weight vectors are isotropic and have only small random overlaps with the teacher’s. (b) In the early stages of learning, the student learns a quadratic approximation of the teacher function by learning the teacher weights up to a rotation; individual teacher-student overlaps remain small $d^{-1/2}$ . This “Gaussian-closure” regime is described by the simple closed equations derived in section III.2 and accounts for the first two stages of loss reduction, shown by the green curve. (c) At the specialization transition, a positive fraction of student neurons develop macroscopic correlations with teacher neurons, although the matching may remain imperfect. Gaussian closure predictions fail, and the loss undergoes a third stage of reduction shown by the yellow curve; see section III.3 . The activation function determines the type of specialization; see section III.3 . (d) In the refinement phase, matched student neurons converge to their corresponding teacher directions, see section III.4 .

Figure 2: Sequential learning of order parameters. For the three different activation functions ReLU, ELU, and tanh, we show the test loss (top row) and the evolution of the order parameters $Q_{p}$ and $Q_{W_{\ell}}$ during training. The shaded regions indicate the successive regimes dominated by the dynamics of $Q_{1}$ , $Q_{2}$ , and higher-order modes $Q_{p\geq 3}$ , while the bottom row tracks a permutation-fixed student–teacher overlap. Parameters : Network with width $k=200$ , input dimension $d=50$ .

Figure 3: Empirical distribution of the student–teacher overlap matrix entries at representative stages of training. Histograms of the teacher–student overlaps $\Omega^{01}_{ij}$ at representative stages of training. Early on (first two columns), all entries lie within a Gaussian bulk (blue), whose mean and width are determined by $Q_{1}$ and $Q_{2}$ . As training progresses, student neurons become strongly correlated with individual teacher neurons. For each student neuron, we highlight in red all overlaps whose magnitudes exceed the extremal value of $k$ random teacher–student overlaps, which scales like $\sqrt{2\log k}$ . Their appearance marks the emergence of specialization. The solid black curve is a Gaussian fit to the remaining overlaps. Parameters : Width $k=200$ and input dimension $d=50$ .

Figure 4: Generalization loss dynamics under Gaussian closure. We show the test loss eq. 5 obtained by simulating neural networks together with the prediction obtained under Gaussian closure, eq. 11 . Panel (a) uses the activation $\sigma(x)=\mathrm{He}_{1}(x)+\mathrm{He}_{2}(x)/2$ , for which the predicted plateau loss is zero. In this case, Gaussian closure describes the entire training dynamics; panels (b–d) show ReLU, ELU, and tanh activation functions. The blue dashed line shows the theoretical pre-specialization loss plateau after removing the contributions from the first two Hermite sectors. Parameters: $k=400,d=100$ .

Figure 5: Timescale of specialization. We plot the time of specialization measured experimentally against the theoretical prediction, eq. 16 . Symbols indicate the SGD step where networks with varying $k=d$ specialize, defined as the first time at which $Q_{W_{\ell}}^{01}$ deviates from its Gaussian-closure value by more than $0.1$ , with $\ell=1$ for odd $r$ and $\ell=2$ for even $r$ . The detailed data collection procedure is illustrated in Fig. 10 . We fix the learning rate for all networks at $\eta=0.1$ ; the activation functions for $r=3,4,5,6$ are defined as $\sigma_{r}(z)=\mathrm{He}_{1}(z)+\frac{1}{2}\mathrm{He}_{2}(z)+\frac{\mu_{r}}{r!}\mathrm{He}_{r}(z)$ with $\mu_{r}=0.25,0.4,1.2,0.72$ , respectively.

Figure 6: Three scenarios of specialization. We illustrate the dynamics of specialization for three representative activation functions. For each student neuron, we show the largest alignment it has with a teacher neuron (thin blue curves): at each update $t$ and for each student neuron $j$ , we set $\bar{\Omega}_{j}^{01}(t)\coloneq\Omega_{i^{\prime}j}^{01}(t),i^{\prime}=\arg\max_{i\in[k]}|\Omega_{ij}^{01}(t)|$ . The black curves show the corresponding macroscopic order parameter defined by eq. 6 . Insets schematically illustrate the three specialization patterns discussed in the text: (a) sign-degenerate matching for activations with only even nonlinear Hermite components, such as ReLU $(r=4)$ ; (b) student–student anti-alignment for activations with only odd nonlinear Hermite components, such as tanh $(r=3)$ ; (c) many-to-one matching without a sign symmetry for mixed-parity activations, such as ELU $(r=3)$ . Parameters: $d=50$ , $k=200$ , and $\eta=0.1$ .

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

$\Omega_{ij}\Omega_{jm}^{\ell}$	Index relation	Order	Sum of fluctuations	Multiplicity	Overall ( $d^{-3}$ )
$\ell=0$	$i\neq j$	$O(d^{-1/2})$	$d^{2}$	$d$	$O(d^{-3/2})$
	$i=j$	$O(1)$	$1$	$d^{2}$	$O(d^{-1})$
$\ell=1$	$i\neq j\neq m$	$O(d^{-1})$	$d^{3}$	$1$	$O(d^{-5/2})$
	$i=j\neq m$	$O(d^{-1/2})$	$d^{2}$	$1$	$O(d^{-5/2})$
	$i=m\neq j$	$O(d^{-1})$	$1$	$d^{2}$	$O(d^{-2})$
	$i=j=m$	$O(1)$	$1$	$d$	$O(d^{-2})$

Appendix

Table 1: Scaling of all the combinations of the form

\Omega_{ij}\Omega_{jm}^{\ell}

for

\ell=0,1,2

and

\ell\geq 3

. The “Index relation” column lists the index configurations, “Order” gives the typical size of a single term, “Sum of fluctuations” the accumulated standard deviation over all index combinations, and “Multiplicity” the number of equivalent terms. The “Overall” column includes the normalization

1/(k^{2}d)

, which scales as

d^{-3}

at fixed

k/d

; the learning-rate factor is omitted here. The centered-background estimates in this table do not include coherent contributions from nonzero overlap means during training, which are retained explicitly in the contractions below.

Figure 7: Macroscopic first- and second-order dynamics in a ReLU network. The input dimension is $d=100$ , while the width $k$ is varied. Darker curves correspond to larger width-to-dimension ratios. Solid curves show online-SGD simulations, while dashed curves show the closed dynamical predictions from eqs. 8 and 10 . For each run, the theoretical curve is initialized with the empirical order parameters of the corresponding seed. Thus the comparison tests the deterministic macroscopic flow conditioned on the initial state, rather than an ensemble-averaged trajectory.

Figure 8: Fitted closure for high-order loop contractions. Data are collected from simulations with different activation functions in the unspecialized Gaussian-closure regime. The near-unit $R^{2}$ values show that a dictionary built from $Q_{1}^{ab}$ and $Q_{2}^{ab}$ captures the dominant finite-size loop contractions entering $Q_{3}$ and $Q_{4}$ . This regression is used only to validate the Gaussian closure and is not part of the asymptotic scaling argument.

Figure 9: Numerical check of eq. 35 . The x-axis shows the time-averaged ratio in eq. 35 , and the y-axis indicates the dimension $d\ (k=d)$ . Each point is computed over the interval from the saturation of $Q_{1}$ and $Q_{2}$ to specialization onset. We define onset as the first time $Q_{W_{\ell}}$ increases by $0.2$ above its plateau value, tracking the largest-magnitude student–teacher overlap for each student. The lower bound is estimated by the 5th percentile of the data, while the upper bound is the maximum after excluding the largest $1/d$ fraction of values.

Figure 10: Evolution of $Q_{W_{\ell}}^{01}$ for different widths $k$ , with $d=k$ , illustrating how the timescale data in Fig. 5 are collected. We show $r=3$ and $r=4$ as examples. Since the value of $Q_{W_{\ell}}^{01}$ at the beginning of the plateau depends on $d$ , we measure the first time it exceeds this value by a fixed amount of $0.1$ . Alternatively, one could use a fixed absolute threshold, such as $Q_{W_{\ell}}^{01}=0.6$ , for all dimensions. We expect the same asymptotic scaling, provided the subsequent growth introduces no additional slow timescale. Indeed, the reduced dynamics predict that growth between two fixed positive overlap magnitudes takes $O(k^{2}/(\eta\varepsilon^{2}))$ time, which is asymptotically shorter than the plateau. A higher threshold can nevertheless introduce stronger finite-size effects at small $d$ and requires longer simulations at large $d$ .

Figure 11: Experimental verification of quadratic shielding. We measure the mean-square displacement (MSD) eq. 48 of the teacher–student overlaps $\Omega_{ij}^{01}$ on the Gaussian plateau. We use the effective activation $\sigma(z)=\mathrm{He}_{1}(z)+\tfrac{1}{2}\mathrm{He}_{2}(z)+\tfrac{\mu_{r}}{r!}\mathrm{He}_{r}(z)$ , with $r=3,4$ , $k=d=100$ , $\eta=0.1$ . Here, $\tau_{Q_{2}}=kd/(4\eta\mu_{2}^{2})$ is the relaxation time of the $Q_{2}$ dynamics, as per eq. 10 . Left: The MSD slope as a function of the lag time shows a crossover between two noise scales ( $r=3,4,\mu_{r}=0.25$ ) Right: The two limiting MSD slopes as a function of $\mu_{r}$ for short and long lags compared to $\tau_{Q_{2}}$ .

Figure 12: Matched-initialization experiment and full-trajectory macroscopic prediction. The student is initialized with partial teacher information according to eq. 53 , with $s=0.5$ , ELU activation, $d=50$ , $k=200$ , and learning rate $\eta=0.1$ . (a) Student–teacher overlap distributions at different training times, plotted as in Fig. 3 . The matched entries form a diagonal spike, while the remaining overlaps are compared with a Gaussian fit. (b) Time evolution of the macroscopic order parameters. Solid curves show the simulation results, and dashed curves show the theoretical ODE predictions. (c) Population loss along training.

Neural Network Training Dynamics

Neural network training dynamics studies how parameters and representations change during optimization and how those changes shape learning, stability, and generalization.

See all 499 papers

May 23, 2026

Feature Learning in Wide Neural Networks under $μ$ P: Identifiability and Sparse-Dictionary Decomposition of the Mean-Field Limit

We establish four structural results for feature learning in wide two-layer neural networks under the Maximal Update Parametrization (

μ

P). First, we prove global existence and uniqueness of the mean-field limit of noisy gradient descent under

μ

P, identifying the maximal admissible weight

w^*

on the moment sequence of the initialization as the reciprocal parameter-moment-growth boundary, and hence the largest weighted moment class propagated by the flow. The finite-particle approximation has uniform-in-time squared-Wasserstein rate

O(N^{-1})

. Second, we characterize identifiability of the mean-field limit: two admissible parameter measures induce the same network function in

L^2

exactly when their active components agree modulo the finite-rank realization symmetry of the architecture. The orbit depth

D^*_{\mathrm{orb}}

is separated from the moment-variety depth

D^*_{\mathrm{var}}

. Third, under the Barron-Hermite target condition the active support of the long-time limit measure admits a sparse-dictionary decomposition: it is supported on at most

S^*

atoms modulo finite-rank realization symmetry, with

S^*

bounded by an explicit coefficient-threshold number. Fourth, we derive the total feature-learning-error decomposition into statistical, optimization, propagation-of-chaos, and sparse-residual components, with a target-dependent Hermite/Barron tail replacing any initialization-only residual. The four results are tied together by an architectural identity: the triple

(w^*, D^*_{\mathrm{orb}}, S^*)

-- the maximal admissible weight, the orbit identifiability depth, and the sparse-dictionary depth at which the target is realizable -- is the natural learning cell of the architecture-data pair

(σ, ρ)

. The proofs are self-contained except for standard results from

μ

P and mean-field Langevin theory.

Akmal Xodarev

Independent Researcher · Tashkent, Uzbekistan

Representation Learning Parameter Identifiability Neural Network Optimization Mean-Field Theory

Ask PDF

May 18, 2026

Feature Learning in Linear-Width Two-Layer Networks: Two vs. One Step of Gradient Descent

We study feature learning in two-layer neural networks within the linear-width regime, where the number of hidden neurons, sample size, and input dimension scale proportionally. While recent work has analyzed feature learning via a single step of gradient descent on the first layer weights in this regime, such one-step update schemes are fundamentally limited: the update to the weights is approximately rank-one, captures only a single direction, and requires the target function to have an information exponent of one. In this paper, we go beyond one-step updates to provide a full characterization of the features learned during the second step of gradient descent with step-sizes

η_1\asymp N^{α_1}

and

η_2 \asymp N^{α_2}

for

α_1, α_2 \in [0,0.5)

, where

N

is the number of hidden neurons. We derive a spectral characterization of the updated weights, demonstrating they behave as a spiked random matrix with multiple outliers, each corresponding to a learned direction. We show that the number of the outliers is determined by the parameters

α_1, α_2

through

\lfloor \frac{α_2}{1/2 - α_1} \rfloor

. Furthermore, by analyzing the alignment between the learned directions and the target function, we identify a gap between training with independent versus reused batches. While independent batches restrict learning to directions with an information exponent of one, batch reuse enables the second update to capture directions even when the information exponent exceeds one, provided that

α_1, α_2

are chosen properly. This shows that the benefits of batch reuse, previously observed in narrow-width regimes, persist in the linear-width limit as well. By characterizing these early-phase evolutions, our work proposes a tractable framework for studying optimization and feature learning phenomenology in modern overparameterized networks.

Behrad Moniri, Hamed Hassani

University of Pennsylvania, Philadelphia, PA, USA

Representation Learning Neural Network Training Dynamics Gradient Descent Dynamics

Ask PDF

Jul 26, 2026

A Statistical Difference between Single-Layer Learning and Hierarchical Learning in Wide Neural Networks

Hierarchical neural networks are widely used in artificial intelligence, yet their mathematical properties remain incompletely understood. In the infinite-width limit, two different theoretical frameworks have been proposed. One reduces deep learning to kernel regression with a fixed kernel by assuming that the parameters remain close to their initialization, whereas the other allows the parameters to move away from their initialization, requiring the kernel itself to be optimized. In this paper, we study a three-layer neural network with a finite but large number of hidden units. We show that training the input-to-hidden weights yields a smaller generalization error than keeping them fixed. Furthermore, the latter setting exhibits singularities in the parameter space, whereas the former does not. These findings indicate that singularities play an essential role even in wide neural networks.

Sumio Watanabe

RIKEN Center for Advanced Intelligence Project 1-4-1 Nihonbashi, Chuo-ku, 103-0027, Tokyo, Japan

Neural Network Generalization Hierarchical Representation Learning Neural Network Training Dynamics

Ask PDF

Abstract

Neural Network Training Dynamics

Explore similar work

Feature Learning in Wide Neural Networks under μμμP: Identifiability and Sparse-Dictionary Decomposition of the Mean-Field Limit

Feature Learning in Linear-Width Two-Layer Networks: Two vs. One Step of Gradient Descent

A Statistical Difference between Single-Layer Learning and Hierarchical Learning in Wide Neural Networks

Feature Learning in Wide Neural Networks under $μ$ P: Identifiability and Sparse-Dictionary Decomposition of the Mean-Field Limit