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

The note began as a chalk talk at Janelia on 18 November 2024. We begin with Gaussian variables and processes, pass through the geometry of threshold excursions, and finish with a simple calculation of local and global decoding error.

§Background

Probability enters theoretical neuroscience in several distinct ways. It can represent noise in neural recordings, pose a learning problem with stochastic data or plasticity, model random synaptic weights, or provide an analytical route into large statistical systems. Here randomness serves a different purpose. It models a neural representation itself as a random code over a continuous external variable such as position.

§Gaussian variables

§Conditional Gaussian

The most useful closure property of a multivariate Gaussian is that conditioning one subset of variables on another still gives a Gaussian. Let $\boldsymbol\xi$ and $\boldsymbol\eta$ be jointly Gaussian, with means $\mathbf m_\xi$ and $\mathbf m_\eta$, and covariance

§Gaussian processes

Moving from a Gaussian variable to a Gaussian vector replaces one value with a finite indexed collection. A Gaussian process continues that move to a continuously indexed random function. It is completely determined by its mean and covariance functions,

For a stationary process with constant mean, the covariance depends only on the separation $\Delta x=y-x$, so we write $r(x,y)=r(\Delta x)$. In particular, $r(0)=r_0$ is the pointwise variance.

§Origins of a random code

Consider $N$ spatially selective neurons presynaptic to a CA1 place cell. Let $u_i(x)$ be the bounded response of input neuron $i$ over an environment $x\in[0,L]$. For clarity, suppose all inputs share one tuning shape up to translation,

with preferred positions $x_i$ that tile the environment densely and uniformly. The subthreshold input to the CA1 cell is a random weighted sum,

where the independent weights $W_i$ have zero mean and finite variance. For any finite set of locations $\{x_k\}$, the multidimensional central limit theorem gives a jointly Gaussian vector $(f(x_1),\ldots,f(x_k))$ as $N$ grows (van der Vaart, 1998; Lehmann & Casella, 1998). The limiting random function is therefore a Gaussian process. The central limit theorem supplies the Gaussian finite-dimensional distributions. The tuning curves and uniform tiling assumptions supply smoothness and stationarity.

Uniform tiling also makes the covariance translation invariant. With input density $\rho=N/L$,

Choosing $\Var(W)=L/N$ leaves a covariance $r(\Delta x)$ that does not grow with population size.

Compare the random weighted sum with the cellular threshold that selects its connected excursions in Plate 1.1.

In this model, position $x$ is supplied by the environment. The process $f(x)$ represents subthreshold CA1 activity produced by many synaptic inputs. Cellular thresholding, together with any incident inhibition, reveals firing fields wherever $f(x)$ rises above a level of order its standard deviation.

§Excursion structure

The neural response can be idealized by a rectified threshold,

Its nonzero regions are the excursion sets of $f$ above $\theta\sqrt{r_0}$. Their number, length, separation, height, and boundary slopes expose much of the geometry of the code before any further output nonlinearity is chosen. Use the fixed path in Plate 2.1 to see how the threshold changes the crossings and connected excursions.

§The Kac-Rice formula

Kac and Rice gave a way to count level crossings of a smooth function, and then of a stochastic process, without locating each root directly (Rice, 1944; Kac & Slepian, 1959).

Compare shallow and steep crossings in Plate 2.2. Each contributes one to the count.

§Expected crossing count

We now apply the formula to a stationary, mean-zero, differentiable Gaussian process. At a fixed location, $f(x)$ and $f'(x)$ are jointly Gaussian. Stationarity gives

and

Thus

The zero covariance makes $f(x)$ and $f'(x)$ independent. Taking the expectation of (2.3) at level $T=\theta\sqrt{r_0}$ gives

Every bounded excursion has one upcrossing and one downcrossing. Its expected density is therefore half the crossing density,

Let $\Phi$ be the standard normal cumulative distribution function. The expected fraction of the domain above the threshold is $1-\Phi(\theta)$. Dividing the total occupied and unoccupied lengths by the expected number of excursions gives the approximate mean excursion length $s$ and mean gap $\bar s$,

At high thresholds, the gaps are much longer than the excursions. Consecutive events are then separated by several correlation lengths, and the crossing process approaches a Poisson process with the mean above.

§Spectral density

Define the even spectral moments

Then $\omega_0=r_0$ and $\omega_2=-r''_0$. The excursion density becomes

At the mean level, $\mu(0)=(2\pi)^{-1}\sqrt{\omega_2/\omega_0}$. It has units of inverse length and supplies a geometric correlation scale.

Repeating the argument for the derivative process gives equal densities of local maxima and minima,

The ratio between turning points and mean-level crossings defines the scale-invariant irregularity parameter

Values near one indicate a regular process with roughly one maximum and one minimum between mean-level upcrossings. Small values indicate additional high-frequency structure.

Compare the normalized densities for $\omega_0$, $\omega_2$, and $\omega_4$ in Plate 2.3. Dividing each weighted spectrum by its own moment gives unit area, so the three shapes share one vertical scale even though their unnormalized moments have different units.

§Excursion shape

Counting fields is only the first step. To describe a typical field, we condition on the event that the process crosses the threshold from below (Kac & Slepian, 1959). Palm conditioning selects the observation location from the crossings of the process. A location sampled uniformly in space follows a different distribution.

§Slope at an upcrossing

Although $f'$ is Gaussian at a uniformly chosen point, it is not Gaussian at an upcrossing. Steep slopes create more crossings per unit length and are sampled more often. The derivative factor in the Kac-Rice formula expresses exactly this selection bias. From this point onward, $\theta$ denotes the threshold itself. Replace it with $\theta\sqrt{r_0}$ to restore the variance scale.

§The Slepian process

For a Gaussian process, the conditional law inside this integral is Gaussian. Its mean and covariance therefore determine it completely. Take

Bochner’s representation gives the remaining cross-covariance,

The joint covariance is therefore

With $\omega_0=r_0$ and $\omega_2=-r''(0)$, Gaussian conditioning gives

and residual covariance

The last two terms are the reduction in uncertainty caused by fixing the height and slope at the crossing. They vanish far from the origin when the covariance and its derivative decay.

Compare the conditional mean and residual process with the high-threshold parabolic limit in Plate 2.4.