A simple mean field model of feature learning

2025/10/16 by Göring, Niclas, Mingard, Chris, Nam, Yoonsoo +1
#FOS: Computer and information sciences #Machine Learning (cs.LG)

paper · doi:10.48550/arxiv.2510.15174

Abstract

Feature learning (FL), where neural networks adapt their internal representations during training, remains poorly understood. Using methods from statistical physics, we derive a tractable, self-consistent mean-field (MF) theory for the Bayesian posterior of two-layer non-linear networks trained with stochastic gradient Langevin dynamics (SGLD). At infinite width, this theory reduces to kernel ridge regression, but at finite width it predicts a symmetry breaking phase transition where networks abruptly align with target functions. While the basic MF theory provides theoretical insight into the emergence of FL in the finite-width regime, semi-quantitatively predicting the onset of FL with noise or sample size, it substantially underestimates the improvements in generalisation after the transition. We trace this discrepancy to a key mechanism absent from the plain MF description: self-reinforcing input feature selection. Incorporating this mechanism into the MF theory allows us to quantitatively match the learning curves of SGLD-trained networks and provides mechanistic insight into FL.

Citations

• Microscopic and collective signatures of feature learning in neural networks
• Dynamical Decoupling of Generalization and Overfitting in Large Two-Layer Networks
• Adaptive kernel predictors from feature-learning infinite limits of neural networks
• From Kernels to Features: A Multi-Scale Adaptive Theory of Feature Learning
• Disentangling Rich Dynamics from Feature Learning: A Framework for Independent Measurements
• Critical feature learning in deep neural networks
• Predictive power of a Bayesian effective action for fully-connected one hidden layer neural networks in the proportional limit
• Grokking as a First Order Phase Transition in Two Layer Networks
• Local Kernel Renormalization as a mechanism for feature learning in overparametrized Convolutional Neural Networks
• Dynamics of Finite Width Kernel and Prediction Fluctuations in Mean Field Neural Networks
• Learning Single-Index Models with Shallow Neural Networks
• Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit
• Neural Networks can Learn Representations with Gradient Descent
• Learning sparse features can lead to overfitting in neural networks
• A Theoretical Analysis on Feature Learning in Neural Networks: Emergence from Inputs and Advantage over Fixed Features
• Self-Consistent Dynamical Field Theory of Kernel Evolution in Wide Neural Networks
• The high-dimensional asymptotics of first order methods with random data
• Learning Parities with Neural Networks
• Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit
• Deep Neural Networks as Gaussian Processes
• Consistency and fluctuations for stochastic gradient Langevin dynamics
• Representation Learning: A Review and New Perspectives
• Representation Learning: A Review and New Perspectives

Related

• From Kernels to Features: A Multi-Scale Adaptive Theory of Feature Learning
• Microscopic and collective signatures of feature learning in neural networks
• The Geometry of Statistical Feature Learning in Mean-Field Langevin Dynamics
• Three Mechanisms of Feature Learning in a Linear Network
• Critical feature learning in deep neural networks