A theory of feature learning in neural networks

We have broken a decades-old bottleneck in machine learning by providing an analytic theory of feature learning in deep linear-width neural networks (width scaling as the input dimension and model size and data count scaling proportionaly). This is a regime where the network is highly expressive, yet it remains constrained parameter-wise with respect to the amount of data it is trained from. The model is thus forced to be "smart" in order to solve the task, i.e., to learn task-relevant features, which is the crux of deep learning.

More precisely, we study the supervised learning of a Multi-Layer Perceptron (MLP) in the matched teacher–student framework: d-dimensional inputs x are fed to a teacher network of depth L that produces noisy outputs y; a student MLP with the same architecture observes n training pairs and attempts to learn the teacher rule. The teacher’s weights are drawn independently from a generic prior and are unknown to the student. We analyse the limit where both the input dimension d and the width of each hidden layer grow proportionally, while the sample size n scales as the square of d. In this regime, the number of data points is proportional to the number of trainable parameters, and learning is controlled by a sampling ratio α. We then characterise the Bayes-optimal performance of the student network trained on such data.

Previous theoretical approaches have either focused on very narrow networks (committee-type models), which have limited expressive power, or on ultra-wide limits (NTK and random-feature models), where feature learning does not occur: they are kernel machines. Others have studied special cases (such as quadratic neural networks) that also collapse to kernel machines. While these simplifications provide valuable insights, they fail to capture the combined impact of proportional width, depth and nonlinearity that gives rise to strong feature learning and rich learning transitions.

Our paper develops a new analytical toolbox that blends the replica method with Harish Chandra–Itzykson–Zuber (HCIZ) integrals, which allows to study MLPs whose matrix degrees of freedom have a high rank. Using this framework we derive replica-symmetric formulas for the free entropy and explicit saddle-point equations (order-parameter equations) that, for the first time, give Bayes-optimal learning curves for deep neural networks in the strong feature learning regime.

The theory predicts two qualitatively distinct regimes: a universal phase (a prior-independent, effectively polynomial approximation similar to a kernel-like learning), and a specialization phase in which individual hidden units align to teacher features. Crucially, specialization is inhomogeneous — it typically starts in inner layers (closer to the input) and propagates outwards, and it can occur unevenly across neurons inside the same layer, producing a rich sequence of learning transitions.

We validate our theoretical predictions through extensive numerical experiments. The results show excellent quantitative agreement with the theory and clearly reveal metastability, partial specialization, and the progressive, layer-by-layer emergence of feature learning. We also find that practical training and sampling procedures are often trapped in these metastable states and would require exponentially long computation to reach the fully specialized, Bayes-optimal solution.

With the tools we provide, based on statistical mechanics of spin glasses and random matrix theory, 𝘆𝗼𝘂 𝗰𝗮𝗻 𝘁𝗿𝗮𝗰𝗸 𝗮𝘁 𝗲𝗮𝗰𝗵 𝗹𝗮𝘆𝗲𝗿 𝘄𝗵𝗲𝘁𝗵𝗲𝗿 𝘁𝗵𝗲 𝗻𝗲𝘁𝘄𝗼𝗿𝗸 𝗵𝗮𝘀 𝗹𝗲𝗮𝗿𝗻𝘁 𝘁𝗵𝗲 𝗵𝗶𝗱𝗱𝗲𝗻 𝗳𝗲𝗮𝘁𝘂𝗿𝗲𝘀 𝗼𝗿 𝗻𝗼𝘁 𝗮𝗹𝗺𝗼𝘀𝘁 𝗮𝘁 𝘁𝗵𝗲 𝗹𝗲𝘃𝗲𝗹 𝗼𝗳 𝗶𝗻𝗱𝗶𝘃𝗶𝗱𝘂𝗮𝗹 𝘄𝗲𝗶𝗴𝗵𝘁𝘀!

E.g., in this picture we compare theory (left column) vs. experiments (right). It shows whether the trainable weights in a given layer --each indexed by two neurons, thus forming a 2d plane-- have specialized (i.e., learnt hidden features of the teacher, yellow) or not yet (blue), as a function of the training set size (increasing from top to bottom). Being able to obtain such a detailed picture of feature learning in a deep model raises the ceiling on tractable complexity and practical relevance, multiplying by orders of magnitude the parameter count handled at full statistical physics resolution.

This brief overview highlights only a part of our results on feature learning in deep neural networks. For the full analysis and many additional insights, see the paper at https://arxiv.org/abs/2510.24616, or get in touch with us if you’d like to know more!