Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

A simple mean field model of feature learning

Niclas Göring, Chris Mingard, Yoonsoo Nam, Ard Louis

TL;DR

This work addresses why finite-width neural networks learn features more efficiently than fixed kernels. It develops a self-consistent mean-field (MF) theory for the Bayesian posterior of a two-layer network trained with stochastic gradient Langevin dynamics (SGLD), revealing a symmetry-breaking onset of feature learning (FL) at finite width, with infinite width reducing to kernel ridge regression. However, plain MF underestimates post-transition generalization because it misses self-reinforcing input feature selection (IFS). The authors introduce MF-ARD, a minimal ARD-based MF extension that endows coordinates with learnable variances, enabling coordinate-wise feature selection and sparsification (WTA-like), thereby quantitatively matching SGLD learning curves and revealing how ARD removes the ambient dimensionality from the FL onset. The framework shows FL comprises a two-stage process: a data-driven phase transition, followed by IFS that concentrates mass on task-relevant coordinates and yields sharp generalization gains. Overall, MF-ARD provides a tractable, mechanistic explanation for finite-width feature learning and its advantage over fixed kernels on sparse tasks.

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

A simple mean field model of feature learning

TL;DR

This work addresses why finite-width neural networks learn features more efficiently than fixed kernels. It develops a self-consistent mean-field (MF) theory for the Bayesian posterior of a two-layer network trained with stochastic gradient Langevin dynamics (SGLD), revealing a symmetry-breaking onset of feature learning (FL) at finite width, with infinite width reducing to kernel ridge regression. However, plain MF underestimates post-transition generalization because it misses self-reinforcing input feature selection (IFS). The authors introduce MF-ARD, a minimal ARD-based MF extension that endows coordinates with learnable variances, enabling coordinate-wise feature selection and sparsification (WTA-like), thereby quantitatively matching SGLD learning curves and revealing how ARD removes the ambient dimensionality from the FL onset. The framework shows FL comprises a two-stage process: a data-driven phase transition, followed by IFS that concentrates mass on task-relevant coordinates and yields sharp generalization gains. Overall, MF-ARD provides a tractable, mechanistic explanation for finite-width feature learning and its advantage over fixed kernels on sparse tasks.

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: \textit{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.
Paper Structure (60 sections, 14 theorems, 164 equations, 7 figures, 4 tables, 1 algorithm)

This paper contains 60 sections, 14 theorems, 164 equations, 7 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related work
  3. Theory: A hierarchy of models (SGLD $\rightarrow$ MF $\rightarrow$ NNGP)
  4. SGLD posterior: Fully interacting theory
  5. FL mechanism as complex neuron-neuron interaction
  6. MF Theory: Removing off-diagonal couplings
  7. NNGP limit: No interactions
  8. No FL mechanism
  9. Problem: Simple MF does not capture a central FL mechanism
  10. Stage 1: The Onset of Feature Learning as a Phase Transition
  11. Stage 2: Specialization via IFS (what simple MF misses)
  12. Why plain MF misses specialization
  13. A minimal MF model of feature learning
  14. Model - ARD as the smallest change that enables IFS
  15. How ARD induces IFS
  16. ...and 45 more sections

Key Result

Theorem 4.1

There exists an outer iteration $t_0=\mathcal{O}(1)$ of the fixed point algorithm and a constant $\varepsilon_0>0$ (independent of $d$ such that we get $\varepsilon_0$-level symmetry breaking towards $S$: $\min_{j \in S} \langle w_j^{2} \rangle_{p_{\textrm{ARD}}}^{t_0} \;-\; \max_{j \notin S} \langl

Figures (7)

  • Figure 1: a) Generalisation error vs. training set size $P$ for 2-layer $N=512$ width RELU networks trained with SGLD on $k$-sparse parity compared to the predicted error of the 3 theories presented in the paper (training details in \ref{['sec_traindetails']}). b) Hierarchy of theories presented in this paper.
  • Figure 2: SGLD Posterior: The model is a fully interacting system where every neuron (row ${\bm{w}}_i$) is coupled to every other neuron (${\bm{w}}_{i'}$) through the interaction kernel $G$. This results in a complex, high-dimensional posterior that is computationally intractable. MF theory: The pairwise interactions are replaced by an effective field that represents the average influence of all other neurons. Each neuron is now treated as an independent sample from a single, shared, data-dependent distribution, $p_{\textrm{MF}}({\bm{w}})$. NNGP Limit: All interactions are removed. Each neuron is an independent and identically distributed sample drawn from the fixed prior distribution.
  • Figure 3: $k$-sparse parity target $y({\bm{x}})=\chi_S({\bm{x}})=\Pi_{j \in S}x_j$ with index $S=\{0,1,2,3\}$ in $d=35$ with a ReLU network ($N=512$, see \ref{['appx_walsh']} and \ref{['sec_traindetails']} for more details.). a) $m_S$ vs. $P$: SGLD, MF and MF-ARD exhibit a phase transition at $P_c$, NNGP grows smoothly (no FL). b) Weight matrix plotted for the SGLD-trained NN: After the phase transition, only a very small set of neurons (9 out of 512) pick up a high correlation with the target ($m_S$ is color coded) and only for these neurons there is a strong symmetry breaking where the norm of the first $j=0,1,2,3$ coordinates is much larger than for the rest $d-k$ coordinates (coded as the height in the plot).
  • Figure 4: Setup as in \ref{['fig_ms']}. a) Coordinate-wise weight marginals $p(w_j)$: For SGLD, $p(w_j)$ has high variance and strong non-Gaussianity on input relevant coordinates $j=0,1,2,3$ (symmetry breaking with IFS). Plain MF also shows symmetry breaking for coordinates $j=0,1,2,3$, but much weaker, while MF-ARD restores the strong anisotropy. b) Distribution of specialization $\tilde{m}{=}aJ_S({\bm{w}})$ before vs. after phase transition: SGLD develops heavy tails after the transition (sparsification), MF remains narrow, MF-ARD produces the heavy-tailed specialization. Together, MF-ARD qualitatively recovers the two mechanisms, IFS and sparsification, that MF misses.
  • Figure 5: Generalisation phase diagram (Test MSE) vs. dataset size $P$ and noise $\kappa$ for $k$-sparse parity target $y({\bm{x}})=\chi_S({\bm{x}})$ with $S=\{0,1,2,3\}$ in $d=35$. a) SGLD, b) MF, c) MF-ARD: All panels show a transition from high to low test error as $P$ increases or $\kappa$ decreases. However, the plain MF theory strongly underestimates the sharpness of the transition while the MF-ARD theory largely reproduces the shape/location of this boundary. Training details are in \ref{['sec_traindetails']}.
  • ...and 2 more figures

Theorems & Definitions (29)

  • Theorem 4.1
  • proof
  • Theorem D.1
  • proof
  • Theorem D.2
  • proof
  • Lemma D.3
  • proof
  • Lemma D.4
  • proof
  • ...and 19 more