A solvable model of feature learning in a nonlinear network

Many theoretical accounts attribute much of the success of deep neural networks to feature learning, in which training changes the features represented by a network’s hidden layers. A theory of deep learning needs to explain which features a nonlinear network learns and how its weights represent them. We do not yet have a general theory that answers these questions.

We study this question as part of PIRAMID’s Advancements in Learning Theory team, which uses statistical physics to build foundations for mechanistic interpretability and AI safety. Our progress and plans describe a theory agenda centred on hidden weight distributions and covariances that can record learned structure. The mean field sequence introduces this viewpoint by treating each neuron as one member of a population whose distribution must be self consistent. The model below gives a solvable example in which we can derive the hidden weight distribution and its covariance.

Theorists have made progress by simplifying the task or the network. Rubin and collaborators give one of the clearest solvable examples in an equilibrium Bayesian model. They study a nonlinear network with two trainable layers in a teacher and student setting. They use statistical mechanics to derive the posterior over its weights. In that posterior, a hidden layer that starts isotropic acquires populations of weights aligned with one useful direction.

In this model, explicit feature learning is a change in the weight posterior. At the transition, the posterior over one scalar overlap develops new wells at finite overlap, and the hidden weight covariance gains a rank one term along the teacher direction. Rubin and collaborators connect the representation change to a first order transition in generalization and to grokking. The network can transfer from an easy linear target component to a harder cubic component.

We explain the feature learning result and work through its derivation. We begin with a cavity reduction, which isolates one neuron against the residual left by the others. The neuron’s $D$ dimensional posterior depends on one scalar overlap. We close the calculation with a self consistency equation that connects the learned weight distribution to the network’s prediction residual. We use the cavity calculation here because it exposes the population description in a simple model. We want to carry the representation and self consistency equations beyond the cavity construction.

Noa Rubin, Inbar Seroussi, and Zohar Ringel use the model to study grokking as a first order feature learning transition (Rubin, Seroussi & Ringel, 2024). A later paper by Rubin and collaborators extends the same mechanism across several network scaling regimes (Rubin et al., 2025). We follow the overlap transition from the first paper and use the adaptive kernel relation developed across both papers.

§The target depends on one direction

Let the input be isotropic Gaussian,

and let $\mathbf e\in\RR^D$ be an unknown unit vector. The teacher depends only on the scalar projection

Its output combines the first and third probabilists’ Hermite polynomials,

The two modes are orthogonal under the standard Gaussian measure, and both use the same $\mathbf e$. The coefficient $\epsilon$ may be small.

The student is a network with two layers and an odd nonlinearity,

Each hidden weight $\mathbf w_i$ can rotate. We measure its alignment with the teacher by the overlap

Plate 1.1 shows the task and the scalar order parameter.

The linear part contains enough information to estimate the teacher direction. Consider

Gaussian orthogonality gives

The squared estimation error is $O(D/P)$, so $P=O(D)$ examples identify $\mathbf e$ to order one accuracy. The linear term supplies the direction. Without it, this derivation does not give the network a route to the cubic target. Another mechanism could still learn the target.

After $O(D)$ samples, the data contain the direction. We need to explain when training puts that direction into the hidden weight distribution and how the new distribution changes prediction.

§A Bayesian model exposes the hidden representation

The tractable model treats training as Bayesian inference. Put independent priors on the hidden and readout weights,

and use a Gaussian likelihood with variance $T$ on a training set of size $P$,

Standard scaling takes $s_a^2=O(1/N)$. The posterior balances a prior cost in $D$ dimensions against the fit reward from $P$ examples.

The hidden weight distribution provides a direct measure of feature learning. Feature learning appears when that distribution changes from isotropic to aligned with a direction selected by the target. Equilibrated Langevin training samples this posterior. Gradient descent can follow another path and can take a different time to reach equilibrium, so its dynamics need a separate argument.

§Score one neuron against the residual

Remove one neuron from the network and call the residual left by the other $N-1$ neurons

For a candidate hidden weight $\mathbf w$, define its feature vector on the training set by

The selected neuron contributes $a\boldsymbol\phi$. Conditional on $\mathbf w$, its readout coefficient is Gaussian, so we can integrate it exactly. Introduce

Completing the square in $a$ gives

Equation (3.1) is exact for a fixed cavity residual. The prior keeps $\mathbf w$ small and isotropic. The factor $L^{-1/2}$ accounts for the cost of fitting a readout coefficient, and the exponential rewards a feature that correlates with the residual. The cavity step holds the residual fixed while it scores one neuron. The final self consistency step returns the neuron to the network and requires the residual and adaptive kernel to agree.

The fit reward depends on $S^2$, so the signs $\mathbf w$ and $-\mathbf w$ remain equivalent. The mean weight can stay at zero while an even moment such as $\EE[\rho^2]$ records direction learning.

§Reduce the hidden weight posterior to one overlap

Before the transition, the residual lies mainly in the same two Hermite modes as the teacher. Write its population form as

For one hidden weight, let

The pair $(g,z)$ is jointly Gaussian, with covariance $\rho$. For a smooth function $F$, the higher order Stein identity gives

The identity turns the $k$th teacher mode into the $k$th power of the overlap. Gaussian inputs give the clean $\rho^k$ factor. With non Gaussian inputs, higher moments can mix the teacher modes, so the calculation needs a different moment expansion.

Setting $F=\operatorname{erf}$ yields the two correlations needed here,

The resulting correlation between the feature and residual is

The feature variance is

For large $P$, the empirical quantities concentrate as $S\simeq Pm$ and $Q\simeq Pv$.

The components of $\mathbf w$ orthogonal to $\mathbf e$ also concentrate, so

After these approximations, Equation (3.1) gives the scalar distribution

with

The first term is the prior cost of selecting one direction in $D$ dimensions. The last term rewards a feature that fits the residual. Since $m(\rho)$ begins at first order in $\rho$, enough data can make weights with finite overlap competitive.

The curvature test locates a spinodal boundary. The first order transition studied by Rubin, Seroussi, and Ringel is a global statement. Wells at finite overlap can reach the same depth as the centered well before the center loses local stability (Rubin, Seroussi & Ringel, 2024). New minima at finite overlap first appear, then reach the same depth as the centered minimum, and later make the center unstable. The middle event is the first order transition.

The 2024 paper combines the sample count and the remaining scales into an effective interaction,

Move the control in Plate 4.1 to follow the conditional action as $u$ increases. The adjacent hidden population uses the same overlap coordinate and shows how posterior mass moves between the wells.

§The hidden covariance records the learned direction

Let

Before alignment, $R=O(D^{-1})$. Once the posterior assigns finite mass near $\rho=\pm\rho_*$, $R$ becomes $O(1)$. Symmetry still gives $\EE[\rho]=0$.

Decompose a hidden weight into parts parallel and orthogonal to the teacher,

The orthogonal directions remain equivalent. To leading order at large $D$,

The covariance receives a rank one correction along $\mathbf e$. Plate 5.1 shows why the second moment is the right measure. A symmetric pair of aligned populations can have zero mean while its covariance records the learned direction.

The rank one term is a direct measure of explicit feature learning. A scalar rescaling changes every direction by the same factor, while the adaptive covariance changes the target direction. Rubin and collaborators retain this directional information across several network scaling regimes, including regimes where a scalar rescaling predicts some mean outputs well (Rubin et al., 2025).

§Self consistency closes the feature learning loop

The one neuron posterior starts from coefficients $(b,c)$ in the residual. The whole network determines those coefficients through its current prediction.

For a candidate residual, the overlap posterior defines an adaptive kernel,

For that kernel, the posterior mean residual on the training set is

Projecting $\overline{\mathbf r}$ onto $H_1$ and $H_3$ gives new coefficients $(b,c)$. A solution is a fixed point of this map. Plate 6.1 shows the full loop.

The calculation for one neuron becomes a network theory only after solving this loop. The residual specifies what remains to be learned. The hidden weight distribution changes the kernel, and the new kernel produces a new residual. We want to carry the requirement that representation and prediction error agree beyond this model.

§A sample complexity corollary

The learned direction changes the kernel power assigned to the cubic teacher mode. Before alignment, a random hidden direction has $\rho=O(D^{-1/2})$, so the cubic kernel power is $O(D^{-3})$. After alignment, a finite part of the posterior has $\rho=O(1)$, so the same channel has order one strength.

The linear target mode reveals $\mathbf e$ at $P=O(D)$. The adaptive network can reuse the learned direction for the cubic function of the same coordinate, while a fixed kernel still needs $P=O(D^3)$. The exponent three comes from squaring $h_3\propto\rho^3$ when a random overlap has size $\rho=O(D^{-1/2})$. It does not come from counting every cubic monomial in $D$ variables.

Plate 7.1 compares the two sample scales.

The later adaptive theory predicts learning of the nonlinear component at $P=O(D)$, while the NNGP and scalar rescaling descriptions retain the $O(D^3)$ requirement (Rubin et al., 2025). The earlier grokking paper observes the same change after the feature learning transition (Rubin, Seroussi & Ringel, 2024).

§What we are working on

The calculation above tells us when a hidden layer learns a useful direction, but it leaves open how the layer distributes what it has learned. If $\rho_i$ is neuron $i$’s overlap with the teacher direction, then its contribution to a degree-$k$ feature is proportional to $a_i\rho_i^k$. The same total feature can come from a few strongly aligned neurons or from many neurons whose individual contributions are small. A small aligned population stores the feature in specialists. A broad population stores the feature across many units.

Both populations can compute the same function while having different hidden covariance and unit selectivity. They can also respond differently when neurons are removed. A theory that predicts only the output cannot distinguish them. We want a theory that begins with the posterior over weights and predicts how many neurons participate and how strongly each contributes.

Our current work asks what controls the division between distributed and specialized representations. We study simple models in which training can store the same function in different ways, and we use the posterior over weights to track how many neurons participate. A theory of representation is necessary for interpretability because a faithful account must describe the function a model computes and how its neurons represent that function.

§References