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

Akmal Xodarev

cs.LGMay 23, 2026

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

Authors: Akmal Xodarev

Organizations: Independent Researcher · Tashkent, Uzbekistan

Abstract

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.

Explore similar work

Jul 7, 2026stat.ML

Width-Robust Learnability in Mean-Field Bayesian Neural Networks

Infinite-width limits are a standard way to reason about neural networks, but it is not automatic that the limiting learner has the same complexity-theoretic inductive bias as large finite networks. We study this question for Bayesian neural networks at the mean-field, or critical feature-learning, scaling. The central quantity is the \emph{reduced entropy}

s_\infty(y,\varepsilon)=\limsup_N -\frac{1}{N}\log π_N^0(L\le \varepsilon),

the intensive prior cost of representing a target function

y

to population mean-squared error

\varepsilon

. Our main result is a width-robust learnability theorem. At fixed depth, a family of Boolean-cube targets is learnable from polynomially many samples at infinite width if and only if it is learnable at polynomial width, if and only if its reduced entropy is polynomially bounded. Equivalently, up to polynomial slack in accuracy, the Bayesian mean-field learner generalizes exactly on the targets that can be represented by polynomial-size networks. The forward direction is proved by a form of subsampling: from the infinitely many hidden neurons in the mean-field solution, one can select polynomially many representatives and still preserve the learned function on every input simultaneously. At the critical scaling this subsampling has both an active'' component, which keeps the data-dependent low-dimensional statistics, and a lazy'' component, which resamples the entropy-dominated directions from the prior. Thus the infinite-width mean-field limit gives a clean analytic description of learning without introducing spurious width-dependent generalization power.

Dmitry Vaintrob, Kaarel Hänni

Bayesian Neural Networks Two-Layer Neural Networks Singular Learning Theory Mean-Field Limit

Ask PDF

May 18, 2026stat.ML

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 \textit{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

Two-Layer Neural Networks Feature Learning Gradient Descent Low-Data Regimes

Ask PDF

Oct 28, 2022cs.LG

A Functional-Space Mean-Field Theory of Partially-Trained Three-Layer Neural Networks

To understand the training dynamics of neural networks, prior studies have considered the mean-field limit of two-layer neural networks as the width tends to infinity, establishing theoretical guarantees for its convergence under gradient flow training as well as approximation and generalization capabilities. In this work, we study the infinite-width limit of a type of three-layer neural network where the first-layer weights are randomly sampled and untrained. To rigorously define the limiting model, we extend the mean-field theory by lifting the representation of neurons from Euclidean to functional spaces. This allows us to establish the mean-field training dynamics as a functional gradient flow with a time-varying kernel that remains positive-definite under suitable assumptions, thus proving a linear-rate convergence of its training loss. Furthermore, we define novel function spaces that contain the solutions obtained through the mean-field training dynamics and prove Rademacher complexity bounds for these spaces. Notably, our analysis applies to a range of scaling choices of the model, resulting in two distinct regimes of the mean-field limit that both exhibit feature learning through training.

Zhengdao Chen, Eric Vanden-Eijnden, Joan Bruna

Two-Layer Neural Networks Mean-Field Limit Wasserstein Gradient Flows Low-Data Regimes Finite

Ask PDF