# Separation of Scales and a Thermodynamic Description of Feature Learning in
Some CNNs

Inbar Seroussi111Department of Mathematics, Weizmann Institute of Science,
Rehovot, Israel Gadi Naveh222 Hebrew University, Racah Institute of Physics,
Jerusalem, 9190401, Israel Zohar Ringel333 Hebrew University, Racah Institute
of Physics, Jerusalem, 9190401, Israel

###### Abstract

Deep neural networks (DNNs) are powerful tools for compressing and distilling
information. Their scale and complexity, often involving billions of inter-
dependent parameters, render direct microscopic analysis difficult. Under such
circumstances, a common strategy is to identify slow variables that average
the erratic behavior of the fast microscopic variables. Here, we identify a
similar separation of scales occurring in fully trained finitely over-
parameterized deep convolutional neural networks (CNNs) and fully connected
networks (FCNs). Specifically, we show that DNN layers couple only through the
second moment (kernels) of their activations and pre-activations. Moreover,
the latter fluctuates in a nearly Gaussian manner. For infinite width DNNs,
these kernels are inert, while for finite ones they adapt to the data and
yield a tractable data-aware Gaussian Process. The resulting thermodynamic
theory of deep learning yields accurate predictions in various settings. In
addition, it provides new ways of analyzing and understanding DNNs in general.

##  1 Introduction

Identifying slow or relevant variables is an essential step in analyzing
large-scale non-linear systems. In the context of deep neural networks (DNNs),
these should be some combinations of the individual weights that are weakly
fluctuating and obey a closed set of equations. One potential set of such
variables is the DNNs’ outputs themselves. Indeed, in the limit of infinitely
over-parameterized DNNs these provide an elegant picture of deep learning [24,
16, 34] based on a mapping to Gaussian Processes (GPs). However, these GP
limits miss out on several qualitative aspects, such as feature learning [30,
12] and the fact that real-world DNNs are not nearly as over-parameterized as
required for the GP description to hold [24, 27, 3, 34]. Obtaining a useful
set of slow variables for describing deep learning at finite over-
parameterization is thus an important open problem in the field.

Several works provide guidelines for this search. Noting that GP limits can
have surprisingly good performance [4] and that over-parameterization is
natural to deep learning [36, 51] we are inclined to keep some elements of the
GP picture. One such element is to work in function space and study pre-
activation and outputs instead of weights whose posterior distribution becomes
complicated even in the GP limit [20, 15]. Another element is the layer-wise
composition of hidden layer kernels [50], using which one generates the output
kernel of the GP [13]. Such a layer-wise picture is also harmonious with the
idea that DNN layers should not correlate strongly, to prevent co-adaptation
[22]. Recently, it was shown that in some limited settings, making the GP
kernel "dynamical" or flexible, so that it adapts to the dataset, can account
for differences between infinite and finite DNNs [29, 34, 1, 21, 2]. Still,
the task of finding an explicit and general set of equations describing this
flexibility in DNNs remains unsolved.

In this work, we identify such slow variables and use these to derive an
effective theory for deep learning capable of capturing various finite
channel/width ($C/N$) effects in convolutional neural networks (CNNs) and
fully connected neural networks (FCNs) such as feature learning. We argue
that:

  1. 1.

For, $C,N\gg 1$ the erratic behavior of specific channels/neurons averages out
and hidden layers coupled to each other only through two “slow" variables per
layer: The second moment of the pre-activations (pre-kernel), $K^{(l)}$, and
the second moment of activations (post-kernel), $Q^{(l)}$, of the $l$th layer.
Furthermore, for mean square error (MSE) loss, FCNs in the so-called mean-
field (MF) scaling [33] (where the last layer weights are scaled down) or CNNs
with a large read-out layer fan-in behave effectively as a GP with a data-
aware kernel determined by the second moment of pre-activations in the
penultimate layer.

  2. 2.

In settings where the kernels have a large density of dominant eigenvalues,
the posterior (or trained) pre-activations fluctuate in a nearly Gaussian
manner. Following this we use a multivariate Gaussian variational
approximation for the posterior pre-activations (pre-kernels) and derive
explicit equations (equations of state), for the covariance matrices governing
these pre-activations.

  3. 3.

We identify an emergent feature learning scale (FLS) denoted by $\chi$,
proportional to the train MSE times $n^{2}$ over $C$ (or $N$). This scale
controls the difference between the finite $C,N$ output kernel ($Q_{f}$) and
its $C,N\rightarrow\infty$ limit and in this sense reflects feature learning.
Due to the $n^{2}$ factor, $\chi$ can be $O(1)$ or larger even for $C\gg 1$,
e.g. for CNN architectures (see Fig. 1 panel c). The same holds, with $C$
replaced by $N$, for FCNs in the MF scaling [33]. Unlike perturbation theory
[47, 34, 27, 35], our theory tracks all orders of $\chi$ and treats only
$1/C,1/N$ perturbatively. The separation of scales between $\chi$ and
$1/C,1/N$ is thus central to our analysis. Its manifestation is the fact that
feature learning shifts and stretches the dynamical variables in the theory
(the pre-activations) in a considerable manner yet barely spoils their
Gaussianity.

The predictions of our approach are tested on several toy and real-world
examples using direct analytical approaches and numerical solutions to the
equations of state. Our analysis takes a physics viewpoint on this complex
non-linear problem. Rigorous mathematical proofs are left as an open problem
for future research.

We note that there are several works showing evidence that the spectrum of the
empirical weight correlation matrix show various tail effects and spikes [31].
While in deeper layers we focus on pre-activations, the spectrum of input
layer weights we obtained, is Gaussian but not independent as in Ref. [15].
Hence it can produce a variety of spectral distributions for the covariance
matrix, similar to the aforementioned ones. We note a recent interesting
work[19] arguing that the test-loss depends only on the mean and variance of
hidden activations. There, however, the setting is of a fixed trained DNN and
the statistics are over the input measure rather than over the DNN parameters
as in our case. While quantitatively different, our approach is similar in
spirit to the phenomenological layer-wise Gaussian Processes put forward in a
recent work [2]. Additional approaches for finite-width include perturbative
correction around the infinite width limit to leading [47, 34] or higher
orders [17, 23, 39]. There is however mounting evidence from bounds on GP
limits [23, 24], numerical experiments [27, 34, 35, 11], as well as the
current work, that such perturbative expansions have slow convergence in
practical regimes. In contrast our EoS are useful both numerically (see Sec. 4
and 2.2) and analytically (see Sec. 2.3) and in addition allow us to model
pre-activation distributions in the wild via our pre-kernels (see Sec. 2.4).

###  1.1 Problem Statement

Our general setup consists of DNNs trained on a labeled training set of size
$n$,
$\mathcal{D}_{n}=\\{(\bm{x}_{\mu},y_{\mu})\\}_{\mu=1}^{n}=\\{X_{n},\bm{y}\\}$
with MSE loss. The input vector is $\bm{x}_{\mu}\in\mathbb{R}^{d}$, and
$y_{\mu}$ is a scalar output. We denote vectors and tensors by boldface and
use $\mu$ and $\nu$ to represent data point indices. Our training algorithm is
Gradient Descent with a vanishing learning rate, along with weight decay, and
additive white noise with variance $\sigma^{2}$ (GD+Noise [34]). When relaxed
to its equilibrium, the resulting dynamics could also be understood in terms
of Bayesian Neural Networks [44] where one aims to find the minimum mean
square error estimator under a Gaussian prior for the DNN weights, where each
sample is corrupted by Gaussian noise with variance $\sigma^{2}$.

Our theory can be applied to any finite number of convolutional, dense, or
pooling layers. To illustrate its main aspects, let us focus on an $L$-layer
fully connected model with width $N_{l}$ ($l\in[1..L-1]$),

$\displaystyle f(\bm{x})$
$\displaystyle=\sum_{j=1}^{N_{L}-1}w^{(L)}_{j}\phi\left(h^{(L-1)}_{j}(\bm{x})\right)$
(1) $\displaystyle h^{(l)}_{j}(\bm{x})$
$\displaystyle=\sum_{i=1}^{N_{l-1}}W^{(l)}_{ji}\phi\left(h^{(l-1)}_{i}(\bm{x})\right)$
$\displaystyle h^{(1)}_{i}(\bm{x})$
$\displaystyle=\bm{w}^{(1)}_{i}\cdot\bm{x}$

where $\bm{w}^{(1)}_{i}\in\mathbb{R}^{d}$ ($i\in[1,N_{1}]$),
$\bm{w}^{(L)}\in\mathbb{R}^{N_{L-1}}$, $W^{(l)}\in\mathbb{R}^{N_{l}\times
N_{l-1}}$ are the weights of the network, and the input vector
$\bm{x}\in\mathbb{R}^{d}$ such that $d=N_{0}$. The activation function,
$\phi:\mathbb{R}\rightarrow\mathbb{R}$, is applied element-wise. Below, we
focus on error function activation, as its anti-symmetry simplifies our
theory. Generalizations to ReLU, where one needs to track both kernels and
averages, could be found in Supp. Mat. (5).

The main object we analyze is the equilibrium distribution of the GD+Noise
algorithm in function space. Adopting physics notation, this distribution can
be written as
$p(\bm{f}|\mathcal{D}_{n})=e^{-\mathcal{S}}/\mathcal{Z}(\mathcal{D}_{n})$,
where $\mathcal{Z}(\mathcal{D}_{n})=\int e^{-\mathcal{S}}$ is the partition
function, and $\mathcal{S}$ is the action or negative log-posterior (see also
Methods). Taking a Bayesian perspective, this probability distribution can
also be viewed as a posterior distribution given measurements of $y_{\mu}$
having Gaussian noise with variance $\sigma^{2}$ and a prior given by a
finite-width random DNN. As shown in Supp. Mat. (1) our partition function is
governed by the following action

$\mathcal{S}=\frac{1}{2}\sum_{i=1}^{N_{1}}(\bm{h}_{i}^{(1)})^{\mathsf{T}}\left[Q^{(1)}\right]^{-1}\bm{h}_{i}^{(1)}+\frac{1}{2}\sum^{L-1}_{l=2}\sum_{i=1}^{N_{l}}(\bm{h}_{i}^{(l)})^{\mathsf{T}}\left[\tilde{Q}^{(l)}(\bm{h^{(l-1)}})\right]^{-1}\bm{h}_{i}^{(l)}\\\
+\frac{1}{2}\bm{f}^{\mathsf{T}}\left[\tilde{Q}_{f}(\bm{h}^{(L-1)})\right]^{-1}\bm{f}+\frac{1}{2\sigma^{2}}\sum_{\mu}(f_{\mu}-y_{\mu})^{2}$
(2)

where $f_{\mu}=f(\bm{x}_{\mu})$, and

$\displaystyle\tilde{Q}_{f}(\bm{h}^{(L-1)})_{\mu\nu}=$
$\displaystyle\frac{\sigma_{L}^{2}}{N_{L-1}}\sum_{j=1}^{N_{L-1}}\phi(h_{j\mu}^{(L-1)})\phi(h_{j\nu}^{(L-1)})$
(3) $\displaystyle\tilde{Q}^{(l+1)}(\bm{h}^{(l)})_{\mu\nu}$
$\displaystyle=\frac{\sigma_{l+1}^{2}}{N_{l}}\sum_{i=1}^{N_{l}}\phi(h_{i\mu}^{(l)})\phi(h_{i\nu}^{(l)})$
$\displaystyle
Q_{\mu\nu}^{(1)}=\frac{\sigma_{1}^{2}}{d}\bm{x}_{\mu}^{T}\bm{x}_{\nu}.$

where $\bm{h}^{(l)}_{i}\in\mathbb{R}^{n}$ and $\sigma_{l}^{2}$ determine the
weight variances at $n=0$ (in the GD+Noise context, these are determined by
the strength of the weight decay [34]). We comment that for rank-deficient
matrices, the inverses are regularized by including a small positive
regularizer to be taken to zero at the end of the computation.

To familiarize ourselves with the action in Eq. (2) let us see how it
reproduces the standard Gaussian Processes picture at infinite channel/width
[34, 13]. Strictly speaking, this action is highly non-linear, since the
$\tilde{Q}$’s matrix elements contain high powers of pre-activations and since
their inverse enters the action. Crucially, however, the $\tilde{Q}$’s are
width-averaged quantities. Thus, at $N_{l}\rightarrow\infty$ one may replace
them by their averages. Furthermore, upstream dependencies, wherein
$\bm{h}^{(l)}$ affects $\bm{h}^{(l-1)}$, vanish (see Supp. Mat. (1)). Roughly
speaking, this is because $\bm{h}^{(l)}$ only feels the collective effect of
all the neurons in $\bm{h}^{(l-1)}$ rendering its feedback on any specific
neuron negligible.

Having these two simplifications in mind, we begin a layer-by-layer downstream
analysis of the DNN: As there is no upstream feedback on $\bm{h}^{(1)}$, the
average of $\tilde{Q}^{(2)}(\bm{h}^{(1)})$ (denoted $Q^{(2)}$) can be carried
under the Gaussian action of the input layer alone (first term in the action).
Replacing $\tilde{Q}^{(2)}(\bm{h}^{(1)})$ by $Q^{(2)}$ in the second term in
the action, would then imply that $\bm{h}^{(2)}$ fluctuates in a Gaussian
manner with $Q^{(2)}$ as its covariance matrix. Next, the average of
$\tilde{Q}^{(3)}(\bm{h}^{(2)})$ (denoted $Q^{(3)}$) can be found based on the
now known, Gaussian statistics of $\bm{h}^{(2)}$. Repeating this process, the
final kernel ($Q^{(L)}=Q_{f}$) is found and is exactly the one obtained using
the method introduced by Cho and Saul [13]. Together with the MSE loss term
(last term in the action), we find that the outputs ($f_{\mu}$) fluctuate in a
Gaussian manner leading to the standard GP picture of infinite width trained
DNNs [34].

Here, however, our focus is at large but finite width ($N_{l}\gg 1$). In this
more complex regime, several corrections may appear: (i) The pre-activations’
average and covariance may deviate from those of a random DNN. (ii) $Q^{(l)}$,
the covariance of activations in the $l-1$ layer, would not solely determine
the covariance of pre-activations on the downstream layer $l$, as upstream
effects between $\bm{h}^{(l+1)}$ and $\bm{h}^{(l)}$ come into play. (iii)
Inter-channel (or inter-neuron in the fully-connected case) and inter-layer
correlations may appear. (iv) The fluctuations of pre-activations may deviate
from that of a Gaussian. A priori, all these corrections may play similarly
dominant roles, thereby making analysis cumbersome.

##  2 Results

###  2.1 Effective GP Description in the Feature Learning Regime

The basic analytical insight underlying this work is that these four types of
corrections scale differently with $n$, $d$, and $N_{l}$. This allows for a
controlled mean-field treatment, which differs substantially from
straightforward perturbation theory in one over the width. As shown in Supp.
Mat. (1.6), corrections of type (iii) are often much smaller than those of
type (i) and (ii). This holds generally for hidden layers when $N_{l}$’s (or
channel’s number for CNNs) are much larger than $1$. Considering the output
layer of CNNs this requires a large fan-in and FCNs this holds when using an
MF scaling [33]. Turning to correction of type (iv) in the $l$’th layer, these
are suppressed when the average $\tilde{Q}^{(l)}$ has a large density of
dominant eigenvalues — a situation relevant for when $n$ and the input
dimension are both large relative to one.

This leaves us with corrections of types (i) and (ii). Interestingly,
following these corrections to all orders leads to a tractable mean-field
picture of learning. The latter is an augmentation of the standard
correspondence between GPs and DNNs at infinite width (NNGP)[13, 26]: Pre-
activations in different layers or channels/neurons remain uncorrelated and
Gaussian. Correlations only appear between different data-points (and latent
pixels for CNNs) within the same layer and channel/neuron. We henceforth
denote the covariance of pre-activations and activations at layer $l$ (up to
normalization by the variance of the weights) by $K^{(l)}$ and $Q^{(l)}$ and
refer to these as pre-kernel and post-kernel, respectively. However in the
NNGP [26] viewpoint, $Q^{(l)}$ is simply proportional to $K^{(l)}$ and fully
determined by the upstream kernel ($Q^{(l-1)}$) whereas here $K^{(l)}$ and
$Q^{(l)}$ differ and moreover depend both on the upstream and downstream
kernels.

Specifically, for the above $L$-layers FCN with $\phi=\mathrm{erf}$ activation
function, we obtain (see also Fig. 3 and Supp. Mat. (1) for derivation and
extension to CNNs and any number of layers)

$\displaystyle\bar{\bm{f}}$
$\displaystyle=Q_{f}[Q_{f}+\sigma^{2}I_{n}]^{-1}\bm{y}$ (4)
$\displaystyle[(K^{(L-1)})^{-1}]_{\mu\nu}$
$\displaystyle=[(Q^{(L-1)})^{-1}]_{\mu\nu}-\frac{1}{N_{L-1}}\Tr{A^{(L)}\frac{\partial
Q_{f}}{\partial[K^{(L-1)}]_{\mu\nu}}}$
$\displaystyle[[K^{(l-1)}]^{-1}]_{\mu\nu}$
$\displaystyle=[[Q^{(l-1)}]^{-1}]_{\mu\nu}+\frac{2N_{l}}{N_{l-1}}\frac{\partial
D_{\mathrm{KL}}(K^{(l)}||Q^{(l)})\,\,}{\partial[K^{(l-1)}]_{\mu\nu}}\,\,\,\text{for
all $l\in[2,L-1]$}$ $\displaystyle[\Sigma^{-1}]_{ss^{\prime}}$
$\displaystyle=\frac{d}{\sigma_{1}^{2}}\delta_{ss^{\prime}}+\frac{2N_{2}}{N_{1}}\frac{\partial
D_{\mathrm{KL}}(K^{(2)}||Q^{(2)})\,\,}{\partial\Sigma_{ss^{\prime}}}$
$\displaystyle A^{(L)}$
$\displaystyle=\sigma^{-4}(\bm{y}-\bar{\bm{f}})(\bm{y}-\bar{\bm{f}})^{\mathsf{T}}-\left[Q_{f}+\sigma^{2}I_{n}\right]^{-1}$

where $[Q_{f}]_{\mu\nu}=\sigma_{L}^{2}G(K^{(L-1)})_{\mu\nu}$,
$[Q^{(l)}]_{\mu\nu}=\sigma_{l}^{2}G(K^{(l-1)})_{\mu\nu}$ and
$G(K)_{\mu\nu}=\frac{2}{\pi}\sin^{-1}\left(\frac{2K_{\mu\nu}}{\sqrt{1+2K_{\mu\nu}}\sqrt{1+2K_{\mu\nu}}}\right)$
[45] for a matrix $K\in\mathbb{R}^{n\times n}$ (equivalent expressions are
known for several common activation functions, such as ReLU [13]). Also,
$\bar{\bm{f}}$ is the average DNN output and $D_{\mathrm{KL}}(K||Q)$ is the
KL-divergence between two Gaussians with covariance matrices $K$ and $Q$. The
input layer post-kernel is $K^{(1)}=X_{n}\Sigma X_{n}^{\mathsf{T}}$, where
$\Sigma$ is the covariance matrix of input layer weights.

As their lack of dependence on width suggests, the first equation together
with the definitions of the post-kernels $Q^{(l)},Q_{f}$ are already present
in the strict GP limit ($N_{l}\rightarrow\infty$). They are, respectively, the
GP inference formula and standard kernel recursive equations of random DNNs
[13] with $\erf$ activation. The remaining equations are, to the best of our
knowledge, novel and follow the changes to the pre-kernels and post-kernels at
finite $N_{l}$. These could be solved analytically in some simple cases (see
subsection 2.3 for the case of two-layer CNN). We note that for non-anti-
symmetric activation, one will also need to track the mean of each layer’s
pre-activation (see Supp. Mat. (5)).

To get a qualitative impression of their role, one can consider the case where
the penultimate layer ($l=L-1$) is linear, in which case
$Q_{f}=\sigma_{L}^{2}K^{(L-1)}$. Consequently,
$\frac{\partial[Q_{f}]_{\mu^{\prime},\nu^{\prime}}}{\partial[K^{(L-1)}]_{\mu,\nu}}=\sigma_{L}^{2}\delta_{\mu\mu^{\prime}}\delta_{\nu\nu^{\prime}}$,
(where $\delta_{\mu\nu}$, with double index refers here to the Kronecker
delta) and thus the second equation simplifies to

$\displaystyle\sigma_{L}^{2}Q_{f}^{-1}$
$\displaystyle=(Q^{(L-1)})^{-1}-\frac{\sigma_{L}^{2}}{N_{L-1}}\left(\bm{\delta}\bm{\delta}^{\mathsf{T}}-[Q_{f}+\sigma^{2}I_{n}]^{-1}\right),$
(5)

where $\bm{\delta}=(\bm{y}-\bar{\bm{f}})/\sigma^{2}$. We note in passing that
even for a non-linear penultimate layer, a similar term will arise from the
expansion of $Q_{f}$ in $K^{(L-1)}$ to linear order. From the above form,
several insights can be drawn.

First, we argue that the above equation implies that the trained DNN is more
susceptible to changes along $\bm{\delta}$ than the DNN at
$N_{L-1}\rightarrow\infty$. Noting how $Q_{f}^{-1}$ enters the action (Eq. 2),
it controls the stiffness associated with fluctuations in ${\bm{f}}$. Hence
$Q_{f}^{-1}$ makes fluctuations in the direction of $\bm{\delta}$ more likely
than they are according to $(Q^{(L-1)})^{-1}$. Since $\bm{\delta}$ measures
the discrepancy in train predictions, this effect reduces the discrepancy by
making the DNN more responsive in these directions than it is at
$N_{L-1}\rightarrow\infty$. The second term, proportional to
$\left[Q_{f}+\sigma^{2}I_{n}\right]^{-1}$ amounts to a negligible reduction in
fluctuations along eigenvectors of $Q_{f}$ corresponding to eigenvalues which
are larger than $\sigma^{2}$.

Using Eq. (5) one can also identify the aforementioned emergent feature
learning scale (or FLS) namely,
$\chi=N_{L-1}^{-1}\bm{\delta}^{\mathsf{T}}Q^{(L-1)}\bm{\delta}$. This scale
represents the magnitude of the leading term when one Taylor expands $Q_{f}$
in $1/N_{L-1}$. When $\chi=O(1)$ or larger there is a significant change in
the eigenvalues of $Q_{f}$ compared to $Q^{(L-1)}$ which indicates feature
learning. On the other hand, when this quantity is small, we are closer to the
GP regime (see Supp. Mat. (1.6)). To asses the scaling $\chi$, one can
consider the common situation where $\bm{\delta}$ has some non-negligible
overlap with dominant eigenvectors of $Q^{(L-1)}$ whose eigenvalues are on the
scale $\lambda$. Here we find
$\chi\approx\lambda\cdot\text{MSE}/\sigma^{2}\cdot n/N_{L-1}$, where MSE
denotes the mean train MSE which enters here via $||\bm{\delta}||^{2}/n$. Due
to its explicit $n$ dependency, and for $\lambda=O(n)$ at large $n$ [14] —
$\chi$ maybe $O(1)$ even at very large $N_{L-1}$ and/or when the average MSE
is rather small.

Figure 1 panel (c) shows the value of $N_{L-1}$ (or $C_{L-1}$) at which
$\chi=1$ (i.e. $N_{L-1}$ or $C_{L-1}$ at which feature learning becomes a
dominant effect) as a function of $n$ for several DNNs we study. The scale
separation, demonstrated there by the fact that $\chi$ can be $O(1)$ in
regions where $1/N_{l}$’s is negligible, is central to our analytical
approach.

This scale $\chi$ is also the reason that naive perturbation theory in
$1/N_{l}$ fails at large $n$ [34, 27, 35, 17], as it treats $\chi$ and
$O(1/N_{L-1})$ on the same footing, since they both have a single negative
power of $N_{L-1}$. In contrast, our EoS treat the FLS non-perturbatively.

Last we stress that the EoS provide us with a concrete effective GP
description for the entire DNN as well as its hidden layers. A priori one
would expect that the normality of pre-activations, a large $C,N$ trait, will
be lost at finite $C,N$. Yet we find that pre-activations remain Gaussian and
accommodate strong feature learning effects while maintaining accurate
predictions. This unexpectedly simple behavior opens various reverse
engineering possibilities wherein one infers the effective kernels from
experiments and uses their spectrum and eigenvectors to rationalize about the
DNN (see also Fig. 5).

###  2.2 Numerical Demonstration: 3 Layer FCN

Next, we test the agreement between the above results and statements and
actual trained DNNs, starting from the 3-layer FCN defined in Eq. (1) with
$L=3$. We focus here on a student-teacher setting with $n=512$ or $1024$
training data points drawn from iid Gaussian distributions with unit variance
along each input dimension. The target was generated by a randomly drawn
teacher FCN of the same type only with $N_{1}=N_{2}=1$. The student was
trained using an analog scaling to the MF scaling [33], wherein the output
layer weights are scaled down by a factor of $1/\sqrt{N_{l}}$. Whereas for the
CNNs discussed below, this choice of scaling was not required, for FCNs we
found it necessary for getting any appreciable feature learning at
$N_{1}=N_{2}=N\gg 1$ (see Fig. 1 panel (c)).

Figure 1: Theory versus experiment. Panels (a) and (b) compare fully trained
3-layer FCNs in the MF scaling with our theoretical predictions. Panel (a)
studies the most strongly fluctuating weight mode in the first layer. A QQ-
plot is shown, of its empirical distribution against a normal one, along with
histograms together with Gaussian fits (insets). As the scale of the problem
increases, these fluctuations become more and more Gaussian. Panel (b) shows
the experimental values and EoS predictions each normalized by their values at
$N_{l}\rightarrow\infty$ with fixed $\sigma_{l}$ (i.e. the GP value). Here we
used $n=1024$ and $d=128$. Panel (c) plots the $N$ (or channel number $C$) at
which the feature learning scale ($\chi$) is $1$. This dictates the crossover
between weak and strong feature learning in our theory. The color gradient
provides a qualitative separation between large and small $N,C$ and thus
correlates with the adequateness of our mean-field decoupling for the hidden
layers. Evidently, FCNs with standard scaling ($[Q_{f}]_{\mu\mu}=O(1)$ at
$n=0$) require $N_{l}=O(1)$ for appreciable feature learning. However, FCN
with an MF scaling, as well as CNNs in standard scaling, can exhibit strong
feature learning well within the regime of our mean-field decoupling.

As described in the methodology section 4, we trained $20$ FCNs using the
GD+Noise algorithm until they reached equilibrium. We use these trained FCNs
to calculate various average quantities under our partition function (Eq.
(2)). Specifically, we focused on: (i) The normalized train-loss on the scale
of $\sigma^{2}$, namely $\text{MSE}/(\sigma^{2}\sum_{\mu}y^{2}_{\mu})$ (ii)
The eigenvalues ($\lambda_{i},i\in[1..d]$) of the average $\Sigma$, where we
average over neurons, training seeds, and training time (the latter within the
equilibrium region). (iii) The normalized overlap ($\alpha$) between the
discrepancy in prediction on the training set times the target, namely
$\alpha=\sigma^{-2}\sum_{\mu=1}^{n}(\bar{f}_{\mu}-y_{\mu})y_{\mu}/(\sum_{\mu}y^{2}_{\mu})$.
We then used a JAX-based [8] numerical solver for the EoS and compared it with
the experiment.

As the results of Fig. 1 panel (b) show, the predictions of our EoS for all
these three quantities converged well as we increased $N$. Furthermore, they
do so in a region where they differ considerably from their associated GP
limit. Indeed, as shown in the Supp. Mat. (3) the top $\Sigma$ eigenvalue came
out 2-3 times larger than it is in the GP limit. The associated eigenvector
corresponded to the first layer weights of the teacher ($\bm{w}^{*}$). The
rest of the eigenvalues remained at their GP limit values. Put together, this
is a clear sign of strong feature learning.

Notably, however, this notion of feature learning does not involve
compression. Indeed, since $\Sigma$ has the same variance as in the GP limit
for directions perpendicular to $\bm{w}^{*}$, it does not compress the input
by projecting it solely on the label relevant direction ($\bm{w}^{*}$).
Instead, it exaggerates the fluctuation of student weights along, $\bm{w}^{*}$
thereby making it statistically more likely that $h^{(1)}_{c\mu}$ and
$h^{(1)}_{c\nu}$ with opposite sign of $\bm{w}^{*}\cdot\bm{x}_{\mu}$ and
$\bm{w}^{*}\cdot\bm{x}_{\nu}$ will be further apart in the space of pre-
activations.

Next, we study how $\chi$ behaves as a function of $n$ and $C$ (or $N$) for
different architectures. Fig. 1 shows the value of $N$ (or $C$) at which
$\chi=1$. As $\chi$ contains a single inverse power of $C$ at ten times this
value, $\chi$ would be $0.1$ and thus indicate only minor feature learning
effects in our EoS. As $N,C$ diminish from this latter value, our EoS yield
increasingly stronger feature learning effects. We find that both for CNNs in
the standard scaling and for FCNs with MF scaling, the crossover to feature
learning happens well within the validity region of our mean-field decoupling
(i.e. large $N$ or $C$). In contrast, FCN with standard scaling shows this
crossover when $N=O(1)$, which is outside the scope of our theory. In this
aspect, we comment that there is evidence that FCNs with standard scale are
inferior to those with mean-field scaling [49] and perform similarly to GPs
[26].

###  2.3 Analytical Solution of the EoS - Two Layer CNN

Having tested our EoS numerically, we turn to show they lend themselves, in
simple settings, to a fully analytical calculation. Amongst other things, this
will flesh out the non-perturbative nature of our results. To this end, we
consider a simple non-linear CNN with 2 layers. Though bounds have been
derived [10, 9], we are not aware of any analytical predictions for the
performance of finite non-linear 2-layer DNNs, let alone CNNs. It is therefore
a natural first application of our approach. Specifically, we consider

$\displaystyle f(\bm{x})$
$\displaystyle=\sum_{i=1}^{N}\sum_{c=1}^{C}a_{ic}\erf\left(\bm{w}_{c}\cdot\bm{x}_{i}\right)$
(6)

where $\bm{x}\in\mathbb{R}^{d}$ with $d=NS$ and
$\bm{w}_{c},\bm{x}_{i}\in\mathbb{R}^{S}$. The vector $\bm{x}_{i}$ is given by
the $iS,..,(i+1)S-1$ coordinates of $\bm{x}$. The dataset consists of
$\\{\bm{x}_{\mu}\\}_{\mu=1}^{n}$ i.i.d. samples, each sample $\bm{x}_{\mu}$ is
a centered Gaussian vector with covariance $I_{d}$. We choose a linear target
of the form $y_{\mu}=\sum_{i}a^{*}_{i}(\bm{w}^{*}\cdot\bm{x}_{\mu,i})$ where
$a^{*}_{i}\sim{\mathcal{N}}(0,1/N)$ and $w^{*}_{s}\sim{\mathcal{N}}(0,1/S)$.
This choice is not crucial but does simplify considerably the GP inference
part of the computation. We train this DNN using GD+noise and tune weight-
decay and gradient noise such that, without any data,
$a_{ic}\sim{\mathcal{N}}(0,\sigma_{\text{a}}^{2}(NC)^{-1})$ and
$[\bm{w}_{c}]_{s}\sim{\mathcal{N}}(0,\sigma_{\text{w}}^{2}S^{-1})$.

The equations of state are given by (See Supp. Mat. (3))

$\displaystyle\bar{\bm{f}}$
$\displaystyle=Q_{f}[\sigma^{2}I_{n}+Q_{f}]^{-1}\bm{y}$ (7)
$\displaystyle[Q_{f}]_{\mu\nu}$
$\displaystyle=\frac{\sigma_{\text{a}}^{2}}{N}\sum_{i}G(X_{n}\Sigma
X_{n}^{\mathsf{T}})_{\mu i,\nu i}$ $\displaystyle[\Sigma^{-1}]_{ss^{\prime}}$
$\displaystyle=\frac{S}{\sigma_{\text{w}}^{2}}\delta_{ss^{\prime}}-\frac{1}{C}\Tr{(\bm{\delta}\bm{\delta}^{\mathsf{T}}-{K_{f}}^{-1})\frac{\partial
Q_{f}}{\partial\Sigma_{ss^{\prime}}})}$

Here we denote the discrepancy from the target by
$\bm{\delta}=(\bm{y}-\bar{\bm{f}})/\sigma^{2}$. The above equations for
$\Sigma_{ss^{\prime}}$ and ${\delta}_{\mu}$ could be solved numerically and
compared with DNN training experiments. The results are shown in Fig. (2) in
solid lines and match empirical values well.

To obtain fully analytical results, we proceed with several approximations for
large $n$. First, we approximate the spectrum of the matrix $[Q_{f}]_{\mu\nu}$
based on its continuum kernel version $Q_{f}(\bm{x},\bm{x}^{\prime})$. This is
closely related to the equivalent kernel [38] approximation, which we adopt
here along with its leading order correction [14]. Similarly, we use large $n$
to replace the double summation
$\sum_{\nu\mu}{\delta}_{\mu}{\delta}_{\nu}[Q_{f}]_{\mu\nu}$ by two integrals
over the measure from which $\bm{x}_{\mu}$ are drawn ($d\mu$). See Supp. Mat.
(3.1) for further details and a discussion of the fully connected case
($N=1$).

The latter approximation also fleshes out the importance of the FLS ($\chi$).
Technically, the two summations provide for the $n^{2}$ scaling and
$\delta_{\mu}$ is related to the MSE via
$t_{\mu}=(y_{\mu}-\bar{f}_{\mu})/\sigma^{2}$ (see also Supp. Mat. (3)). The
FLS controls the deviation between the pre-kernel and post-kernel of the
penultimate layer with our EoS. Hence, in particular, it implies deviations
from GP where these are not equal.

Following our approximations, the equations acquire the full rotation symmetry
of the data-set measure, which amount to an independent orthogonal
transformation of each $\bm{x}_{i}$. Furthermore, as shown in Supp. Mat. (3),
at large $S$, $\bar{f}(\bm{x})$ (the continuum function representing
$f_{\mu}$) is linear given a linear target, $y(\bm{x})$, regardless of
$\Sigma$ and hence so is ${\delta}(\bm{x})$. The above symmetry then implies
that ${\delta}(\bm{x})$ is only a function of $\bm{w}^{*}\cdot\bm{x}_{i}$ and
furthermore takes the simple form ${\delta}(\bm{x})=\alpha y(\bm{x})$. The
quantity $\alpha$ thus measures the overlap between the discrepancy in
predictions (${\bm{t}}$) and the target. Following this, the EoS are reduced
to a non-linear equation in a single variable $\alpha$

$\displaystyle\sigma^{2}\alpha$
$\displaystyle=1-\frac{q_{\mathrm{train}}S\lambda_{\infty}l_{*}}{S\lambda_{\infty}l_{*}+\sigma^{2}/n}$
(8) $\displaystyle l_{*}$ $\displaystyle=\frac{1}{S}[1-\chi_{2}]^{-1}$
$\displaystyle\chi_{2}$ $\displaystyle=C^{-1}\alpha^{2}n^{2}\lambda_{\infty}$

where $\lambda_{\infty}$ is the dominant eigenvalue of
$Q_{f}(\bm{x},\bm{x}^{\prime})$ associated with a linear function in the limit
$C\rightarrow\infty$, $l_{*}$ is the eigenvalue of $\Sigma$ associated with
$\bm{w}^{*}$ (the remaining eigenvalues are inert and equal to $1/S$),
$\chi_{2}$ is the FLS, and we assumed
$||\bm{a}^{*}||^{2}=||\bm{w}^{*}||^{2}=\sigma_{\text{a}}^{2}=\sigma_{\text{w}}^{2}=1$
(see Supp. Mat. (3). for more generic expressions). The quantity
$q_{\mathrm{train}}$ is exactly one in the equivalent kernel limit, and its
perturbative correction (in $1/n$) can be found in Supp. Mat. (3).

Solving the above equation for $\alpha$, one obtains $l_{*}$ and hence
$\Sigma$ and also $Q_{f}$ (via Eq. (7)). Using the obtained $Q_{f}$ one can
calculate the DNN’s predictions on the test-set. The effect of the FLS is
evident in the second equation where it controls the deviations from the GP
limit. Here we also recall that $\alpha^{2}$ is the train MSE over
$\sigma^{4}$, thus $\chi_{2}$ as defined above contains the MSE factor
mentioned in the introduction.

To test the theoretical predictions, we trained two such CNNs, with
$n=\\{800,1600\\};S=64;N=20$ and varying channel number. Fig. 2, left panel,
shows the empirical test-set values for $\alpha$ (dots) compared with a
numerical solution of the equations of state (solid lines) and their
analytical solution (dashed lines). For the latter, we obtained $\Sigma$
analytically and performed the resulting GP inference with $Q_{f}$
numerically. The inset tracks the train-set results which, in this case, are
fully analytical and involve no numerical GP inference. Both predictions match
empirical values quite well, even in the regime where test root MSE is roughly
half that of a Gaussian Process ($C\rightarrow\infty$). The right panel shows
the input layer weights, dotted with $\bm{w}^{*}/|\bm{w}^{*}|$ and with a
normalized random vector. These remain Gaussian up to minor statistical noise.
Further details can be found in the methods section.

To emphasize the non-perturbative nature of our Eq. (111), let us assume for
the sake of negation that they agree with first order perturbation theory in
$1/C$ (as in Refs. [34, 27, 35, 47]). If so, we may replace $\alpha$ in the
above expression for $\chi_{2}$ by its GP value, as it already contains one
negative power of $C$ and hence receives no further corrections at that order.
Numerics show this value $\alpha_{GP}=0.558$ for $n=1600$. Plugging this in,
one obtains $l_{*}=\frac{2}{S}[1-633.2/C]^{-1}$. Clearly, this logic leads to
a contradiction unless $C\gg 633.2$. In contrast, our theory provides highly
accurate predictions for $n=1600,C=320$ and $C=640$ well away from where
$\frac{2}{S}[1-633.2/C]^{-1}$ admits a perturbation theory in $1/C$. In Supp.
Mat. (6.2) we report additional results on $l_{*}$ over its GP value.

Figure 2: Theory and experiment comparison for the 2-layer non-linear CNN.
Left. Discrepancy measured along target direction normalized by that of the
corresponding GP ($\alpha/\alpha_{C\rightarrow\infty}$). Dots denote empirical
values for the test-set, and dashed and solid lines are theoretical
predictions via an approximate-analytical and exact-numerical solution of the
equations of state, respectively. Train-set comparisons are shown in the
inset. Right. Statistics of $\bm{w}\cdot\bm{w}^{*}/|\bm{w}^{*}|$ of the
trained CNN (blue) compared with $\bm{w}$ projected on a normalized random
vector (green) for the $n=1600$ experiment with $C=640$. Dashed lines show fit
to a Gaussian.

###  2.4 Extensions to Deeper CNNs and Subsets of Real-World Data-sets

For truly deep CNNs and real-world datasets, obtaining fully analytical
predictions for DNN performance is a challenging task even in the
$C\rightarrow\infty$ limit. Still, the EoS could be solved numerically and
compared with experimental values. Furthermore, the quantities which underlie
them could be examined and reasoned upon. We do so here in two richer
settings, a 3-layer CNN trained with a teacher CNN and the Myrtle-5 CNN [42]
trained on a subset of CIFAR-10.

Our first setting extends that of the previous subsection by having an extra
activated layer and a non-linear target function. Specifically, we consider a
student CNN defined by

$\displaystyle f(\bm{x})$
$\displaystyle=\sum_{j=0}^{N-1}\sum_{c^{\prime}=1}^{C_{2}}a_{c^{\prime}j}\phi\left(h^{(2)}_{c^{\prime}j}(\bm{x})\right)$
(9) $\displaystyle h^{(2)}_{c^{\prime}j}(\bm{x})$
$\displaystyle=\sum_{i=0}^{S_{1}-1}\sum_{c=1}^{C_{1}}v_{c^{\prime}ci}\phi\left(h^{(1)}_{cji}(\bm{x})\right)$
$\displaystyle h^{(1)}_{cji}(\bm{x})$
$\displaystyle=\bm{w}_{c}\cdot\bm{x}_{i+jS_{1}}$

where $\bm{w}_{c},\bm{x}_{i+jS_{1}}\in\mathbb{R}^{S_{0}}$,
$a\in\mathbb{R}^{C_{2}\times N}$, $v\in\mathbb{R}^{C_{2}\times C_{1}\times
S_{1}}$, and the input vector $\bm{x}\in\mathbb{R}^{d}$ with $d=NS_{1}S_{0}$,
and the activation function, $\phi:\mathbb{R}\rightarrow\mathbb{R}$, is
applied element-wise. See Fig. 3 for illustration. Similarly to before, the
regression target ($y_{\mu}$) is generated by a random teacher CNN
($y_{\mu}=f^{*}(\bm{x}_{\mu})$) having the same architecture as the student,
only with $C_{1}=C_{2}=1$. In addition, we chose
$S_{0}=50,S_{1}=30,N=2,C_{1}=C_{2}=100$ for the student. Further details are
found in the methods section 4.

Figure 3: (a) An illustration of a 3 layers CNN architecture. (b) Learning as
described by the effective theory, for $C_{1}=C_{2}=C$. Left: for
$C\rightarrow\infty$, weights of the first layer ($\bm{w}_{c}$), pre-
activations ($\bm{h}^{(2)}$), and the output ($\bm{f}$) fluctuate according to
a Gaussian distribution with post-kernels
$\sigma^{2}_{w}I_{S_{0}}/S_{0},Q^{(2)}$, $Q_{f}$, respectively. Right: for
$C\gg 1$ but finite, each layer receives an effective Gaussian action from the
previous layer with kernels $Q^{(2)},Q_{f}$ and "back-propagation" terms which
try to align $\tilde{Q}^{(2)}(\bm{h}^{(1)}),\tilde{Q}_{f}(\bm{h}^{(2)})$ with
the second moment of $\bm{h}^{(2)}$, and $\bm{f}$ respectively. At large $n$
and input fan-in, these combine to yield an approximately Gaussian
distribution for the weights and pre-activation with pre-kernels
$\Sigma,K^{(2)}$, and the DNN outputs ($\bm{f}$), which amount to a Gaussian
Process Regression (GPR) with $Q_{f}$.

As the first test of our theory, we examine the fluctuations of pre-
activations in the input and middle layers of the trained student CNN and
check their normality. Specifically, for the input weights $\bm{w}_{c}$ we
obtain the histogram (over channels, equilibrium samples, and seeds) of
$\bm{w}_{c}\cdot\bm{w}^{*}/|\bm{w}^{*}|$, where $\bm{w}^{*}$ is the teacher
input weight and the histogram of $\bm{w}_{c}\cdot\bm{w}_{r}/|\bm{w}_{r}|$
where $\bm{w}_{r}$ is a random vector. Teacher overlap has a variance of
$0.254$ here, whereas random overlap variance, averaged over choices of
$\bm{w}_{r}$’s, was $0.039$ with an std of $0.0043$. For the hidden layer, we
obtain the histogram of
$\bm{h}^{(2)}_{c^{\prime}}\cdot\bm{h}^{(2),*}/|\bm{h}^{(1),*}|$ where
$\bm{h}^{(2),*}$ are the teacher’s pre-activations as well as
$\bm{h}^{(2)}_{c^{\prime}}\cdot\bm{h}^{(2)}_{r}/|\bm{h}^{(2)}_{r}|$ where
$\bm{h}^{(2)}_{r}$ is the pre-activation of a different randomly chosen
teacher. Teacher overlap variance here was, $64.4$ whereas average student
variance was, $2.3$ with an std of $0.12$. Fig. 4 shows the associated
histograms along with their fit to a Gaussian. The large and consistent
differences in the variance of the fluctuations between teacher directions and
random directions show that we are deep in the feature learning regime.
Remarkably, the fluctuations remain almost perfectly Gaussian. The larger
variance along teacher directions implies that by drawing DNNs from the
trained DNN ensemble and diagonalizing their empirical covariance matrices,
one is more likely to find dominant eigenvalues along these teacher
directions.

Figure 4: Pre-activation statistics in the 3-layer student-teacher setting.
Left. Histogram of student input weight vector, dotted with the normalized
teacher weight (blue) and normalized random weight (green). Right. Histogram
of student hidden layer pre-activations, dotted with the normalized teacher
pre-activations (blue) and normalized pre-activations of a random teacher
(green). Dots are empirical values and dashed lines are Gaussian fits. Insets:
2d histograms along the same vectors before (left) and after (right) training.
Within our framework, these variances are determined by
$\bm{v}^{\mathsf{T}}K^{(l)}\bm{v}$ with $\bm{v}$ being either random unit
vector or $\bm{h}^{(l)}$ of the single channel teacher. Remarkably, despite
strong changes to the kernels and various non-linearity in the action, the
pre-activation is almost perfectly Gaussian.

We turn to verify the EoS and rationalize on the behavior pre-kernels. To this
end, we average the empirical pre-activations, over channels and training
seeds, to obtain an estimator for the pre-kernel and post-kernel of $h^{(2)}$
(i.e. $Q^{(2)}$ and $K^{(2)}$) and that of the input weights ($\Sigma$). We
then obtain $\partial D_{\text{KL}}(K^{(2)}||Q^{(2)})/\partial\Sigma$
analytically using the 3rd equation from Eqs. 4, plugging in the empirical
$\Sigma$. Finally, we compare the empirical $\Sigma$ with that obtained from
the last equation from Eqs. (4). Fig. 5 left panel plots the eigenvalues of
$\Sigma$, our predictions ($\Sigma_{\text{pred}}$), and the post-kernel of the
input layer which is simply $\sigma_{\text{w}}^{-2}S_{0}I_{S_{0}}$, showing a
good match between the first two. Figure (5) right panel plots the eigenvalues
of $Q^{(2)}$ as predicted from $K^{(2)}$, compared with its empirical value.

Figure 5: Theoretical, empirical, and GP spectrum of kernels for the 3-layer
non-linear student-teacher CNN setting. Left: the leading 50 eigenvalues
making up the spectrum of $\Sigma$ for the empirical $\Sigma$ (blue), for the
predicted $\Sigma$ based on the equations of state (red), and for
$C\rightarrow\infty$ (green). Right: the leading 10 eigenvalues of the post-
kernel of the hidden layer ($Q^{(2)}$) again for the empirical, predicted, and
GP post kernels. The inset shows the next 90 trailing eigenvalues and reveals
a gap separating the leading $S_{1}$ eigenvalues from the rest. The dimension
of $Q^{(2)}$ here is $Nn=2000$, we focus on large eigenvalues as these
dominate the predictions. The outlier in the left panel is aligned with the
teacher weight vector $\bm{w}^{*}$. The 30 (or $N$) outliers on the right
panel are again a feature learning effect and represent the linear feature
proportional to $\bm{w}^{*}$ in each of the $N$ latent pixels of the hidden
layer.

Next, we trained the myrtle-5 CNN[42], capable of good performance and
containing both pooling layers and ReLU activations, with $C=256$ on a subset
of CIFAR-10 ($n=2048$). As shown in Fig. 6, pre-activations show a strong
deviation of trained DNNs from non-trained DNNs or DNNs at infinite
channel/width, and at the same time show quite a good fit to Gaussian in most
cases. This opens the possibility of reverse engineering the pre-kernels
governing this trained network and using them to rationalize about the DNN,
for instance by identifying their dominant eigenvectors.

The 2nd layer (as well as the input layer, (see Supp. Mat. (6.3))) show
deviations from Gaussianity in the leading eigenvalue. This is expected since
the kernels of these layers show quite a dilute dominant spectrum whereas VGA
requires a contribution from many adjacent modes (see Supp. Mat. (1.3)).
Interestingly, despite this non-Gaussianity in the leading eigenvalue of
layers 1 and 2, Gaussainity is restored in the downstream layers 3 and 4.
Correlations across layers and across channels within the same layer are very
weak (largely on the order of $10^{-3}$) and fully consistent with the mean-
field decoupling underlying this work. Further technical details are found in
Supp. Mat. (6.3).

Figure 6: Pre-activation statistics for the myrtle-5 CNN with ReLU activation
and $C=256$, trained on $n=2048$ CIFAR-10 images. Panels (a)+(b): Statistics
of pre-activations of untrained (grey) and trained (blue) Myrtle-5 nets, in
the 2nd layer. The histograms are aggregations over channels and
initialization seeds of the pre-activations, projected on various eigenvectors
of the corresponding covariance matrix. Panel (a) shows a projection on the
1st (leading) eigenvector, while panel (b) is the same for the 10th
eigenvector. Panels (c)+(d) are the same as (a)+(b) but for the 4th layer. The
fits to a Gaussian (dashed lines) provide an accuracy measure for our
variational Gaussian approximation. Notice that the equilibrium distribution,
even when it is near Gaussian, can differ from the initial one in either or
both its mean and variance, a signature of feature learning. Generally, we
find that Gaussianity increases with the depth of the layer (i.e. downstream
layers are more Gaussian than upstream layers), and with the index of the
eigenvector, we project on (e.g. here projection on the 10th eigenvector is
more Gaussian than on the 1st eigenvector). More details can be found in Supp.
Mat. (6.3), where we also show that correlations are small across different
layers and across channels within each layer, thus justifying neglecting these
in our theory.

##  3 Discussion

In this work, we presented what is, to the best of our knowledge, a novel
mean-field framework for analyzing finite deep non-linear neural networks in
the feature learning regime. Central to our analysis was a series of mean-
field approximations, revealing that pre-activations are weakly correlated
between layers and follow a Gaussian distribution within each layer with a
pre-kernel $K^{(l)}$. Using the latter together with the post-kernel $Q^{(l)}$
induced by the upstream layer, explicit equations-of-state (EoS) governing the
statistics of the hidden layers were given. These enabled us to derive, for
the first time, analytical predictions for the performance of non-linear CNNs
and deep non-linear FCNs in the feature learning regime. We further note that
our EoS generalize straightforwardly to combined CNN-FCN architectures,
pooling layers, and models with multiple outputs. Our theory can also be
viewed from a Bayesian perspective, the GP process represented by the equation
we find is a good approximation to the true posterior distribution generated
by a large but finite width Bayesian neural network.

Various aspects of this work invite further study. Empirically, it would be
interesting to better characterize the scope of models for which Langevin
dynamics (or potentially ensemble-averaged NTK dynamics) leads to Gaussian
pre-activations and overall GP-like behavior. Probing the “feature-learning-
load“ of each layer, by experimentally measuring the differences between the
kernels $Q^{(l)}$ and $K^{(l)}$, may also provide insights on generalization,
transfer learning, and pruning, thus complementing other diagnostic tools
suggested recently [32]. For instance, transferring a layer with a small
feature learning load may provide little benefit, and pruning a channel having
a large overlap with a leading eigenvalue of $K^{(l)}-Q^{(l)}$ may be harmful.

From the theory side, it is desirable to develop analytical techniques for
solving the EoS as well as guarantees regarding the existence and uniqueness
of solutions. In particular, exploring the possibility of spontaneous symmetry
breaking of internal symmetries such as weight inversion. Providing a
mathematical underpinning for the approximations involved here may lend itself
to developing performance bounds on the Langevin algorithm and Bayesian neural
networks[41, 15]. Similarly, one can consider using the empirical effective
kernel ($Q_{f}$) as a starting point to develop GP-based bounds [40] on
performance. Last, it is interesting to explore the approach to equilibrium of
the training dynamics and adapt the approximations carried here to the NTK
setting [24, 17].

##  4 Methods

###  4.1 Mean Field Action

Here we present the main ingredients of our theory, leading to the EoS we
find. Further details can be found in the Supp. Mat.

####  4.1.1 Decoupling of Layers and Neurons

First, we provide decoupling of Eq. (2) into layer-wise neuron-wise terms,
wherein each of the terms depends on the upstream and downstream layers only
through channel-averaged second moments of activations and pre-activations
(pre-kernel and post-kernel). Further details are found in Supp. Mat. (1).

Consider the non-linear terms in the action 2 which couple the different
layers. This coupling is mediated through the channel/width-averaged
quantities: indeed $\bm{h}^{(1)}$ depends on $\bm{h}^{(2)}$ through the
channel/width averaged square term in $\bm{h}^{(2)}$, $\bm{h}^{(2)}$ depends
on $\bm{h}^{(1)}$ through the average of
$\phi(\bm{h}^{(1)})\phi(\bm{h}^{(1)})$, and $\bm{h}^{(3)}$ depends on
$\bm{h}^{(2)}$ through the average of $\phi(\bm{h}^{(2)})\phi(\bm{h}^{(2)})$
and so forth. For $N_{l}\gg 1$ we expect these to be weakly fluctuating and
well approximated by their mean-field values. This behavior propagates till
the output layer, and in particular implies that the outputs $\bm{f}$
fluctuate in a Gaussian manner, as previously conjectured [35]. As for the
dependency of $\bm{h}^{(L-1)}$ on the $\bm{f}$ variables, it is not through a
channel/width averaged quantity. However, we find that in various scenarios,
such as FCNs with MF scaling or CNNs with large $N$, the fluctuations of
$\bm{f}$ are suppressed enabling us to replace $\bm{f}$ by its average (see
Supp. Mat. (1.7) and (2.2)). Following this, we obtain our mean-field action,

$\mathcal{S}_{\mathrm{MF}}=\sum^{L-1}_{l=1}\mathcal{S}^{(l)}_{\mathrm{MF}}+\mathcal{S}_{f,\mathrm{MF}}.$
(10)

This allows us to define the $\langle...\rangle_{\mathrm{MF}}$ with respect to
the distribution:
$\pi_{\mathrm{MF}}(\\{\bm{h}^{(l)}\\}^{L-1}_{l=1},\bm{f})=e^{-\mathcal{S}_{\mathrm{MF}}}/\mathcal{Z}_{\mathrm{MF}}$,
and the partition function
$\mathcal{Z}_{\mathrm{MF}}=\int\prod_{l=1}^{L-1}d\bm{h}^{(l)}d\bm{f}\pi_{\mathrm{MF}}(\\{\bm{h}^{(l)}\\}^{L-1}_{l=1},\bm{f})$.
Such that

$\mathcal{S}_{f,\mathrm{MF}}=\frac{1}{2\sigma^{2}}\sum_{\mu}(f_{\mu}-y_{\mu})^{2}+\sum_{\mu\nu}\frac{1}{2}f_{\mu}[Q_{f}]^{-1}_{\mu\nu}f_{\nu}$
(11)
$\mathcal{S}^{(l)}_{\mathrm{MF}}=\frac{1}{2}N_{l+1}\sum_{\mu\nu}A^{(l+1)}_{\mu\nu}\tilde{Q}_{\mu\nu}^{(l+1)}(\bm{h}^{(l)})+\sum_{\mu\nu
j}\frac{1}{2}h_{\mu j}^{(l)}[Q^{(l)}]_{\mu\nu}^{-1}h_{\nu
j}^{(l)}\,\,\,\,\,\,\text{for }l\in[1,L-1]$ (12)

where

$A^{(L)}=\sigma^{-4}(\bm{y}-\bar{\bm{f}})(\bm{y}-\bar{\bm{f}})^{\mathsf{T}}-\left[Q_{f}+\sigma^{2}I_{n}\right]^{-1},$
(13)

and

$A^{(l)}=[Q^{(l)}]^{-1}\left(I_{n}-\langle\bm{h}_{j}^{(l)}(\bm{h}_{j}^{(l)})^{T}\rangle_{\mathrm{MF}}[Q^{(l)}]^{-1}\right).$
(14)

The post-kernels are defined self-consistently as
$Q_{f}=\langle\tilde{Q}_{f}\rangle_{\mathrm{MF}}$, and
$Q^{(l)}=\langle\tilde{Q}^{(l)}\rangle_{\mathrm{MF}}$ for $l\in[1,L-1]$.

Notably, any coupling between the different layers is only through static
mean-field quantities, namely the pre-kernels and-post kernels. In addition,
all neuron-neuron couplings (and similarly, channel-channel couplings for
CNNs) have been removed.

####  4.1.2 Intra-Layer Decoupling

Despite the simplified inter-layer coupling and intra-layer neuron coupling,
the mean-field actions are still non-quadratic for all layers but the output
layer. This non-linearity couples all the $\bm{h}^{(l)}$ variables for the
same neuron (channel in the CNN case) in a way that is roughly all-to-all in
the data-point index. In atomic and nuclear physics, similar circumstances are
well described by self-consistent Hartree-Fock approximations [28, 18, 5, 37].
In our setting, this approximation is directly analogous to a variational
Gaussian approximation (VGA). In Supp. mat. (4) we argue that in the typical
case where the diagonal of $K^{(l)}$ is much larger than the off-diagonal
elements, the VGA is well controlled. Technically, we do so by showing, order
by order in perturbation theory, that the diagrams accounted for by the VGA
approximation dominate all other perturbation theory diagrams. In Supp. Mat.
(3) we also establish this using different means for $S_{0}\gg 1$ for the
specific case of two-layer CNN with a single activated layer. We further
comment that the VGA is exact for deep linear DNNs.

Accordingly, we now look for the Gaussian distribution, governed by a kernel
$K^{(l)}$ which is the closest to the above non-quadratic action. In models
with many hidden layers, this leads to the following "inverse kernel shift"
behavior for $1<l<L-2$

$\displaystyle[[K^{(l-1)}]^{-1}]_{\mu\nu}$
$\displaystyle=[[Q^{(l-1)}]^{-1}]_{\mu\nu}+\frac{2N_{l}}{N_{l-1}}\frac{\partial
D_{\mathrm{KL}}(K^{(l)}||Q^{(l)})\,\,}{\partial[K^{(l-1)}]_{\mu\nu}}$ (15)

where $l$ denotes a layer index and $D_{\mathrm{KL}}(A||B)$ is the
Kullback–Leibler (KL) divergence between two multivariate Gaussians with
covariance matrices $A$ and $B$. As shown in Supp. Mat. (1), for antisymmetric
activation functions, the derivative of the KL-divergence is explicitly given
by

$\Tr[[Q^{(l)}]^{-1}(K^{(l)}-Q^{(l)})[Q^{(l)}]^{-1}(\partial Q^{(l)}/\partial
K^{(l-1)}_{i\mu,j\nu})].$ (16)

For non-anti-symmetric ones, see Supp. Mat. (5).

###  4.2 Experimental Details

Hyperparameters. For the 2-layer CNN experiments, we used $S=64$, $N=20$, and
varying channel number. The training parameters (noise and weight-decay) were
tuned such that $\sigma^{2}=0.1$ and weight variance of $2.0$ over fan-in, for
both layers at $n=0$. The target was drawn once for all experiments using
i.i.d. Gaussian centered random $a^{*}_{i}$ and $w^{*}_{s}$ with variances
$1/N$ and $1/S$ respectively.

For the 3-layer CNN experiments, we took $S_{1}=50,S_{0}=30,N=2$. The training
parameters (noise and weight-decay) were scaled such that $\sigma^{2}=0.005$
and weight variance of $2.0$ over fan-in for the inputs and hidden layer with
no training data (at initialization). The weight variance of the read-out
layer was $15$ over the fan-in. The target was drawn again once for all
experiments from a teacher CNN with $C=1$.

For all the myrtle-5 experiments, we used $n=2048,C=256$ and ReLU activation.
The training parameters (noise and weight-decay) were scaled such that
$\sigma^{2}=0.005$ and weight variance of $2.0$ over fan-in for all layers
with no training data (at initialization).

For all the FCN experiments, we used equal width ($N_{1}=N_{2}$) and weight
decay corresponding to variance,
$\sigma_{\text{w}}^{2}=\sigma_{\text{a}}^{2}=2$ (with no training data) in the
regular scaling. For the MF scaling, we took $\sigma_{\text{a}}^{2}=2/N_{2}$.
The target was drawn again once for all experiments from a teacher CNN with
$N_{1}=N_{2}=1$. Specifically when calculating the emergent scale, we used
$\sigma_{\text{a}}^{2}=2/256$ independent of $N_{2}$.

Equilibrium sampling. To obtain weakly correlated samples from the equilibrium
distribution of the trained CNNs we used the following procedure. For the 2
and 3-layer CNNs, we used an adaptive learning rate scheduler: For the first
$100$ epochs we used a learning rate $lr_{0}/10$, then we crank up the
learning rate to $lr_{0}$. As of epoch $5000$, every $1000$ epoch we estimate
the fluctuations of the train-loss and check for spikes - events in which the
train-loss was 5-times larger than the standard deviation in the past $500$
epochs. If a spike is observed, the learning rate is reduced by a factor of
$0.7$. This continues until $50000$ epochs pass without any events. Then the
learning rate is reduced again by a factor of two and remains fixed. Samples
from these final stages were treated as equilibrium samples. We further
checked that (i) different initialization seeds trained with this protocol
reached the same train-loss statistics. (ii) No further reduction in train-
loss occurred after the final learning rate reduction. For several runs, we
also verified that increasing the last reduction of learning rate by an
additional factor of $2$ did not have any appreciable effect on the loss. The
initial $lr_{0}$ was $\sim 1e-4$ (w.r.t. a standard mean reduction MSE loss)
and the final learning rate was typically $\sim 1e-5$. The runs terminated at
epoch $300000$.

For the myrtle-5 CNN trained on CIFAR-10, we first ran several runs for
$300000$ epochs using the above procedure and examined those that reached the
lowest train-loss. We then generated a fixed scheduler based on those more
successful instances, running up to $4e5$ epochs. We again verified that
further lowering the final learning rate has no appreciable effect on the
training loss and that different seeds reach similar final train-loss. This
ensures that we are indeed sampling from a valid equilibrium distribution.

For the 3-layer CNN and Myrtle-5, we found that auto-correlation times of pre-
activations change considerably between the layers. While the read-out layer
typically had an auto-correlation time of the order of $10^{3}$ epochs (at the
lowest learning rates) the auto-correlation times for the input layers could
reach $\sim 10^{6}$ or larger values. To overcome this issue, when analyzing
pre-activations of these deeper DNNs we took an ensemble containing $98$ and
$234$ different initialization seeds for the 3-layer CNN and Myrtle-5
respectively.

For the 3-layer FCN we used a fixed scheduler which starts at $1/2$ the
maximal stable learning rate and reduces the learning rate by factors of $2$
at $100,1e5,1e6,3e6$ epochs and by a factor of $4$ at $4e6,5e6$ epochs (factor
of $128$ in total). Equilibrium sampling was done between $6e6-7e6$ epochs.

Numerical solution of the equations of state. For the 2-layer CNN, the
equations of state were solved using Newton-Krylov method[25], which does not
require explicit gradients. To facilitate convergence, we adopted an annealing
procedure: For $C\sim 1000$, we obtain the solution using a GP initial value
($x_{0}$) for $\Sigma$. The optimization outcome was then used as $x_{0}$ for
the next lower value of $C$. Using 12 CPU cores, this optimization took
several hours. After obtaining $\Sigma_{ss^{\prime}}$ as a function of $C$,
the resulting kernel $[Q_{f}]_{\mu\nu}$ was used in standard GP inference to
obtain $f$ on the test-set. For the 3-layer FCN we used a more efficient JAX-
based code to generate the kernels and kernel derivatives involved in the EoS,
but otherwise followed the same procedure. Optimization took between several
minutes to a few hours on one Titan-X GPU, depending on parameters.

## Supplementary Information - Separation of Scales and a Thermodynamic
Description of Feature Learning in Some CNNs

###### Contents

  1. 1 Introduction
    1. 1.1 Problem Statement
  2. 2 Results
    1. 2.1 Effective GP Description in the Feature Learning Regime
    2. 2.2 Numerical Demonstration: 3 Layer FCN
    3. 2.3 Analytical Solution of the EoS - Two Layer CNN
    4. 2.4 Extensions to Deeper CNNs and Subsets of Real-World Data-sets
  3. 3 Discussion
  4. 4 Methods
    1. 4.1 Mean Field Action
      1. 4.1.1 Decoupling of Layers and Neurons
      2. 4.1.2 Intra-Layer Decoupling
    2. 4.2 Experimental Details
  5. 5 Derivation of the Mean-field Equations for a Fully Connected Network
    1. 5.1 Model Definition and Main Players
    2. 5.2 Mean-Field Decoupling
      1. 5.2.1 Mean-field Decoupling - Hidden Layers
      2. 5.2.2 Mean-Field Decoupling - Output Layer
    3. 5.3 Variational Gaussian approximation
      1. 5.3.1 $l\in[2,L-1]$
      2. 5.3.2 $l=1$
    4. 5.4 Equations of State (EoS)
      1. 5.4.1 Extension to test points
      2. 5.4.2 Auxiliary field correlation
    5. 5.5 An Emergent Scale
    6. 5.6 Estimating Corrections to Mean-field Results
  6. 6 Mean-Field Equations for CNNs
    1. 6.1 An Emergent Scale in CNNs
    2. 6.2 Estimating Mean-field corrections - CNNs
  7. 7 Toy Example - One Hidden Layer
    1. 7.1 Continuum Limit of Summations for CNNs and Fully-Connected DNNs
  8. 8 Validity of the Variational Gaussian Approximation
  9. 9 Variational Gaussian Approximation for ReLU Activation
  10. 10 Further details on the numerical experiments
    1. 10.1 FCN experiments
    2. 10.2 2-layer CNN experiment
    3. 10.3 Myrtle-5 CNN on subsets of CIFAR-10 experiment

##  5 Derivation of the Mean-field Equations for a Fully Connected Network

In this section, we derive the equations of state for deep fully connected NNs
with a finite number of layers, $L$. In Sec. 6 we provide a sketch of the
derivation for CNN architecture. The analysis can also be extended to other
architectures, including pooling layers and skip connections.

###  5.1 Model Definition and Main Players

The model is composed of a $L$ layer NN having $L-1$ activated hidden layers,
and one linear readout layer. Specifically, we consider

$\displaystyle f(\bm{x})$
$\displaystyle=\sum_{j=0}^{N_{L-1}-1}w^{(L)}_{j}\phi\left(h_{j}^{(L-1)}(\bm{x})\right)$
(17) $\displaystyle h_{j}^{(l+1)}(\bm{x})$
$\displaystyle=\sum_{i=0}^{N_{l}-1}W^{(l+1)}_{ji}\phi\left(h_{i}^{(l)}(\bm{x})\right)$
(18) $\displaystyle h_{i}^{(1)}(\bm{x})$
$\displaystyle=\sum^{d-1}_{j=0}{W}^{(1)}_{ij}{x}_{j}$ (19)

where $l\in[1,L-1]$, the weights of the network are
$\bm{w}^{(L)}\in\mathbb{R}^{N_{L-1}}$, ${W}^{(l)}\in\mathbb{R}^{N_{l-1}\times
N_{l}}$, such that $N_{0}=d$ and $N_{L}=1$, and $\bm{x}\in\mathbb{R}^{d}$ is
the input vector. Bold letters denote vector or tensor quantities. The
function $\phi:\mathbb{R}\rightarrow\mathbb{R}$ is the activation function
applied element-wise. In subsection 5.3, when deriving the equation of states,
we take $\phi=\mathrm{erf}$ for concreteness, other activation functions can
also be considered. See section 9 for extension to ReLU. The training data is
denoted by $\mathcal{D}_{n}=\left\\{\bm{x}_{\mu},y_{\mu}\right\\}_{\mu=0}^{n}$
and we use a square loss
$\mathcal{L}=\sum_{\mu=1}^{n}\left(y_{\mu}-f(\bm{x}_{\mu})\right)^{2}$. The
feature matrix, $X_{n}$, represents all the input samples,
$\\{\bm{x}_{\mu}\\}_{\mu=1}^{n}$.

Our main object of interest is the following equilibrium distribution of the
Langevin dynamics algorithm with noise strength $\sigma^{2}$ and weight decay
written in function space (i.e. in terms of the DNNs outputs)

$\displaystyle p(\bm{\bm{f}}|\mathcal{D}_{n})$ $\displaystyle\propto
p(\bm{f}|X_{n})\mathrm{exp}\left(-\frac{\mathcal{L}}{2\sigma^{2}}\right),$
(20)

In practice, we sample from this distribution using Gradient descent (GD), at
small learning rates, together with weight decay and noise on each weight
derivative. The weight decay parameters are $\sigma_{l}^{2}/N_{l-1}$ for layer
$l$. The first term on the r.h.s. is given by

$\displaystyle
p(\bm{f}|X_{n})=\left\langle\prod_{\mu}\delta\left[f(\bm{x}_{\mu})-\sum_{j=0}^{N_{L-1}-1}w^{(L)}_{j}\phi\left(h_{j}^{(L-1)}(\bm{x}_{\mu})\right)\right]\right\rangle_{\\{W^{(l)}\\}^{L-1}_{l=1},\bm{w}^{(L)}}$
(21)

where $\bm{f}=(f_{1},...,f_{n})$ is viewed now as a random variable following
the NN outputs on all different training points. The average
$\langle...\rangle_{\bm{w}}$ is over the weights $\bm{w}$ of the network at
equilibrium. The weights’ distribution at equilibrium can be obtained
explicitly. At $N_{l}\rightarrow\infty$, for $l\in[1,L]$ these distributions
(Eq. (21) and Eq. (20)) tend to a GP, however our interest here is at finite
$N_{l}$.

Eq. (20) and Eq. (21) can also be understood from a Bayesian perspective. Eq.
(20) can be viewed as the posterior distribution assuming each sample is
generated by a neural network model as in Eq. (19) and is corrupted by an
additive i.i.d. Gaussian noise with variance $\sigma^{2}$. The prior
distribution over the weights of the network is taken to be Gaussian with
variance $\sigma_{l}^{2}/N_{l-1}$.

To obtain a more explicitly "layer-wise" representation of Eq. 21, we next
condition over the pre-activations ($\bm{h}^{(l)}$) of each layer using Bayes’
formula, we obtain the following Markov representation of Eq. (21)

$\displaystyle p(\bm{f}|X_{n})$ $\displaystyle=\int
p(\bm{f}|\bm{h}^{(L-1)},X_{n})\prod^{L-2}_{l=1}p(\bm{h}^{(l+1)}|\bm{h}^{(l)},X_{n})d\bm{h}^{(l)}d\bm{h}^{(L-1)}.$
(22)

The hidden layers probabilities above are defined as follows:

$\displaystyle p(\bm{f}|\bm{h}^{(L-1)},X_{n})$
$\displaystyle=\left\langle\prod_{\mu}\delta\left(f_{\mu}-\sum_{j=0}^{N_{L}-1}w^{(L)}_{j}\phi\left(h_{j\mu}^{(L-1)}\right)\right)\right\rangle_{\bm{w}^{(L)}},$
(23) $\displaystyle p(\bm{h}^{(l+1)}|\bm{h}^{(l)},X_{n})$
$\displaystyle=\left\langle\prod_{\mu
j}\delta\left(h_{j\mu}^{(l+1)}-\sum_{i=0}^{N_{l}-1}W^{(l+1)}_{ji}\phi\left(h_{i\mu}^{(l)}\right)\right)\right\rangle_{W^{(l+1)}},$
(24) $\displaystyle p(\bm{h}^{(1)}|X_{n})$
$\displaystyle=\left\langle\prod_{\mu
i}\delta\left(h_{i\mu}^{(1)}-\sum^{d}_{j=0}W^{(1)}_{ij}{x}_{\mu,j}\right)\right\rangle_{W^{(1)}},$
(25)

where we write for short $h_{i\mu}^{(l)}=h_{i}^{(l)}(\bm{x}_{\mu})$ for
$l\in\\{1,...,L-1\\}$, the latter being the random variables describing
argument of the activation function (pre-activation) of the $l$th layer at
neuron $i$, on the $\bm{x}_{\mu}$ data-point. Later we will use this Markov
structure of the distribution to decouple the different layers.

We continue our analysis by using the Fourier identity, which replaces the
above delta functions by auxiliary fields,
$\bm{t},\\{\bm{m}^{(l)}\\}^{L-1}_{l=1}$. To this end, the probability
distribution over the network output given the input data can be written as
follows.

$\displaystyle p(\bm{f}|X_{n})$ $\displaystyle\propto\int
d\bm{t}\prod^{L-1}_{l=1}d\bm{m}^{(l)}d\bm{h}^{(l)}\exp(-\mathcal{S}\left(\bm{f},\bm{t},\\{\bm{m}^{(l)},\bm{h}^{(l)}\\}^{L-1}_{l=1}\right))$
(26)

where we adopt here physics notation, and define,
$\mathcal{S}\left(\bm{f},\bm{t},\\{\bm{m}^{(l)},\bm{h}^{(l)}\\}^{L-1}_{l=1}\right)$,
as the action associated with this distribution. Collecting all terms, the
action is defined as follows

$\displaystyle\mathcal{S}$
$\displaystyle=\sum^{L-1}_{l=1}\mathcal{S}^{(l)}+\mathcal{S}_{f}$ (27)
$\displaystyle\mathcal{S}^{(1)}$ $\displaystyle=\frac{1}{2}\sum_{\mu\nu
i}m_{\mu i}^{(1)}m_{\nu i}^{(1)}Q_{\mu\nu}^{(1)}-i\sum_{\mu j}m_{\mu
j}^{(1)}h_{\mu j}^{(1)}$ $\displaystyle\mathcal{S}^{(l)}$
$\displaystyle=\frac{1}{2}\sum_{\mu\nu i}m_{\mu i}^{(l)}m_{\nu
i}^{(l)}\tilde{Q}_{\mu\nu}^{(l)}(\bm{h}^{(l-1)})-i\sum_{\mu j}m_{\mu
j}^{(l)}h_{\mu j}^{(l)}$ $\displaystyle\mathcal{S}_{f}$
$\displaystyle=-i\sum_{\mu}t_{\mu}f_{\mu}+\frac{1}{2}\sum_{\mu\nu}t_{\mu}t_{\nu}\left[\tilde{Q}_{f}(\bm{h}^{(L-1)})\right]_{\mu\nu}+\frac{1}{2\sigma^{2}}\sum_{\mu}(f_{\mu}-y_{\mu})^{2}$

To obtain the above expression, we performed Gaussian integration over the
weights of all layers. In addition, we define the following matrices:

$\displaystyle\tilde{Q}_{f}(\bm{h}^{(L-1)})_{\mu\nu}=$
$\displaystyle\frac{\sigma_{L}^{2}}{N_{L-1}}\sum_{j=0}^{N_{L-1}-1}\phi(h_{j\mu}^{(L-1)})\phi(h_{j\nu}^{(L-1)})$
(28) $\displaystyle\tilde{Q}^{(l+1)}(\bm{h}^{(l)})_{\mu\nu}=$
$\displaystyle\frac{\sigma_{l+1}^{2}}{N_{l}}\sum_{i=0}^{N_{l}-1}\phi(h_{i\mu}^{(l)})\phi(h_{i\nu}^{(l)}),$
(29) $\displaystyle Q_{\mu\nu}^{(1)}=$
$\displaystyle\frac{\sigma_{1}^{2}}{d}\bm{x}_{\mu}^{\mathsf{T}}\bm{x}_{\nu}.$
(30)

Averages of these matrices within our mean-field theory defined below are
referred to henceforth as post-kernels. The name post-kernel represents the
fact that at large width, these matrices concentrate around the covariance
matrix of the post activation. At infinite width, these are the GP kernels,
but for finite width, we show that they concentrate around a different
average, which is defined self-consistently via our mean-field theory. The
action in Eq. (27) is associated with the partition function $\mathcal{Z}=\int
e^{-\mathcal{S}}d\bm{f}d\bm{t}\Pi^{L-1}_{l=1}d\bm{h}^{(l)}d\bm{m}^{(l)}$. We
comment that by integrating over the auxiliary fields, $\bm{t},\bm{m}^{(l)}$,
and using the following identity $\int
e^{-\bm{x}^{\mathsf{T}}A\bm{x}/2+\bm{J}\bm{x}}d^{n}\bm{x}=\sqrt{\frac{(2\pi)^{n}}{\mathrm{det}A}}e^{\bm{J}^{\mathsf{T}}A^{-1}\bm{J}/2}$,
where $\bm{J},\bm{x}\in\mathbb{R}^{n}$, and $A\in\mathbb{R}^{n\times n}$, one
obtains the equivalent form cited in the main text, containing only the pre-
activations and the outputs

$\mathcal{S}=\frac{1}{2}\sum_{\mu\nu
i}h_{i\mu}^{(1)}\left[Q^{(1)}\right]_{\mu\nu}^{-1}h_{i\nu}^{(1)}+\frac{1}{2}\sum^{L-1}_{l=2}\sum_{\mu\nu
i}h_{i\mu}^{(l)}\left[\tilde{Q}^{(l)}\right]_{\mu\nu}^{-1}h_{i\nu}^{(l)}\\\
+\frac{1}{2}f_{\mu}\left[\tilde{Q}_{f}(\bm{h}^{(L-1)})\right]_{\mu\nu}^{-1}f_{\nu}+\frac{1}{2\sigma^{2}}\sum_{\mu}(f_{\mu}-y_{\mu})^{2}$
(31)

In the following section, we derive the inter-layer mean-field decoupling. In
this context, the form of the action in the presence of the auxiliary fields
(Eq. (27)) turns out to be useful.

###  5.2 Mean-Field Decoupling

Next, we capitalize on the fact that the coupling between layers in Eq. 27 is
only through width-index averaged quantities, specifically the width-index
average of $\bm{m}^{(l)}$ and $\tilde{Q}^{(l)}(\bm{h}^{(l-1)})$. Under such
circumstances, it is natural to consider a mean-field type approximation where
one replaces

$\displaystyle\sum_{\mu\nu i}m_{\mu i}^{(l)}m_{\nu
i}^{(l)}\tilde{Q}_{\mu\nu}^{(l)}(\bm{h}^{(l-1)})$
$\displaystyle\rightarrow\sum_{\mu\nu i}\langle m_{\mu i}^{(l)}m_{\nu
i}^{(l)}\rangle_{\mathrm{MF}}\tilde{Q}_{\mu\nu}^{(l)}(\bm{h}^{(l-1)})$ (32)
$\displaystyle+\sum_{\mu\nu i}m_{\mu i}^{(l)}m_{\nu
i}^{(l)}\langle\tilde{Q}_{\mu\nu}^{(l)}(\bm{h}^{(l-1)})\rangle_{\mathrm{MF}}$

where $\langle...\rangle_{\mathrm{MF}}$ is the average over a mean-field
distribution which we construct via the above approximation in a self-
consistent manner. Importantly, following such a replacement, the fluctuations
of the different layers as well as the different neurons within each layer
become independent. This constitutes a major simplification and sets the stage
for the final approximation we carry - a variational Gaussian approximation
(VGA), to estimate the individual partition function of each neuron. This also
allows us to derive the equation of states of the GP process we find. However,
before turning to the VGA, the three subsections below provide a detailed
derivation of the mean-field decoupling together with an estimate of the
leading order correction to the mean-field action.

####  5.2.1 Mean-field Decoupling - Hidden Layers

Our first step is to understand the dependence of $\bm{h}^{(l)}$ on
$\bm{h}^{(l-1)}$ for all $l\in[2,L-1]$. Following the mean-field idea
introduced in Eq. (32) we first subtract average quantities and rewrite the
action of the $j$th neuron of the $l$th layer such that
$\mathcal{S}^{(l)}=\sum^{N_{l}-1}_{j=0}\mathcal{S}_{j}^{(l)}$:

$\mathcal{S}_{j}^{(l)}=\sum_{\mu\nu}\langle m_{\mu j}^{(l)}m_{\nu
j}^{(l)}\rangle_{\mathrm{MF}}\tilde{Q}_{\mu\nu}^{(l)}(\bm{h}^{(l-1)})+\sum_{\mu\nu}m_{\mu
j}^{(l)}m_{\nu
j}^{(l)}\langle\tilde{Q}_{\mu\nu}^{(l)}(\bm{h}^{(l-1)})\rangle_{\mathrm{MF}}\\\
+i\sum_{\mu}m_{\mu j}^{(l)}h_{\mu
j}^{(l)}+\frac{1}{2}\sum_{\mu\nu}\Delta[\bm{m}^{(l)}(\bm{m}^{(l)})^{\mathsf{T}}]_{\mu\nu}\Delta\tilde{Q}_{\mu\nu}^{(l)}(\bm{h}^{(l-1)})+\mathrm{const}$
(33)

where
$\Delta[\bm{m}^{(l)}(\bm{m}^{(l)})^{\mathsf{T}}]_{\mu\nu}=\frac{1}{N_{l}}\sum_{j}{m}_{\mu
j}^{(l)}{m}_{\nu j}^{(l)}-\frac{1}{N_{l}}\sum_{j}\langle{m}_{\mu
j}^{(l)}{m}_{\nu j}^{(l)}\rangle_{\mathrm{MF}}$, and
$\Delta\tilde{Q}_{\mu\nu}^{(l)}(\bm{h}^{(l-1)})=\tilde{Q}_{\mu\nu}^{(l)}(\bm{h}^{(l-1)})-\langle\tilde{Q}_{\mu\nu}^{(l)}(\bm{h}^{(l-1)})\rangle_{\mathrm{MF}}$.
Note that, in the mean-field decoupling
$\frac{1}{N_{l}}\sum_{j}\langle{m}_{\mu j}^{(l)}{m}_{\nu
j}^{(l)}\rangle_{\mathrm{MF}}=\langle{m}_{\mu j}^{(l)}{m}_{\nu
j}^{(l)}\rangle_{\mathrm{MF}}$ for any $j\in[0,N_{l}-1]$. The last term in Eq.
(33) (discarding the $const$) is the fluctuations around the mean-field
average. Since the sum over all neurons in the layer concentrates around its
average for large width by the law of large numbers, the order of this term is
$1/N_{l}$ smaller than all the other terms in the action. We, therefore,
neglect this term. This leaves us with a Gaussian distribution over the
auxiliary field, $\bm{m}^{(l)}$ which we can now integrate over and obtain:

$\int\prod_{j=0}^{N_{l}-1}d\bm{m}_{j}^{(l)}\exp\left(-{\sum_{\mu\nu}\langle
m_{\mu j}^{(l)}m_{\nu
i}^{(l)}\rangle_{\mathrm{MF}}\tilde{Q}_{\mu\nu}^{(l)}(\bm{h}^{(l-1)})-\sum_{\mu\nu}m_{\mu
j}^{(l)}m_{\nu
i}^{(l)}\langle\tilde{Q}_{\mu\nu}^{(l)}(\bm{h}^{(l-1)})\rangle_{\mathrm{MF}}-i\sum_{\mu}m_{\mu
j}^{(l)}h_{\mu j}^{(l)}}\right)\\\
=\exp\left(-\frac{1}{2}\sum_{\mu\nu,j}h_{\mu
j}^{(l)}[Q^{(l)}]_{\mu\nu}^{-1}h_{\nu
j}^{(l)}-\frac{1}{2}\sum_{\mu\nu,j}\langle m_{\mu j}^{(l)}m_{\nu
j}^{(l)}\rangle_{\mathrm{MF}}\tilde{Q}_{\mu\nu}^{(l)}(\bm{h}^{(l-1)})\right)\cdot$
(34)

In the second transition, we used the previously defined post-kernel
$Q^{(l)}=\langle\tilde{Q}_{\mu\nu}^{(l)}(\bm{h}^{(l-1)})\rangle_{\mathrm{MF}}$.

Following Eq. (34) for $l$ and $l-1$, we gather all the $h^{(l-1)}$ dependent
terms and obtain the mean-field action of the $l-1$th layer:

$\mathcal{S}^{(l-1)}_{\mathrm{MF}}=\frac{1}{2}N_{l}\sum_{\mu\nu}A^{(l)}_{\mu\nu}\tilde{Q}_{\mu\nu}^{(l)}(\bm{h}^{(l-1)})+\frac{1}{2}h_{\mu
j}^{(l-1)}[Q^{(l-1)}]_{\mu\nu}^{-1}h_{\nu j}^{(l-1)}$ (35)

where we denote by $A^{(l)}_{\mu\nu}=\langle{m}_{\mu j}^{(l)}{m}_{\nu
j}^{(l)}\rangle_{\mathrm{MF}}$, which due to the symmetry of the problem do
not depend on the specific neuron $j$. The above mean-field action is valid
for $l\in[2,L-1]$.

To calculate
$A^{(l)}=\langle\bm{m}^{(l)}(\bm{m}^{(l)})^{\mathsf{T}}\rangle_{\mathrm{MF}}$
for $l\in[2,L-1]$, we use standard multivariate Gaussian results (see
subsection 5.4.2) this yields the following result:

$\displaystyle A^{(l)}$
$\displaystyle=[Q^{(l)}]^{-1}\left(I_{n}-\langle\bm{h}_{j}^{(l)}(\bm{h}_{j}^{(l)})^{\mathsf{T}}\rangle_{\mathrm{MF}}[Q^{(l)}]^{-1}\right)$
(36)

where
$\langle\bm{h}_{j}^{(l)}(\bm{h}_{j}^{(l)})^{\mathsf{T}}\rangle_{\mathrm{MF}}$
is determined self-consistently.

####  5.2.2 Mean-Field Decoupling - Output Layer

We next consider the coupling between the final/output layer and the
penultimate layer. As done previously in the analysis of the hidden layers, we
rewrite the action of the top layer as

$\mathcal{S}_{f}=i\sum_{\mu}t_{\mu}f_{\mu}+\frac{1}{2\sigma^{2}}\sum_{\mu}(f_{\mu}-y_{\mu})^{2}+\frac{1}{2}\sum_{\mu\nu}\langle
t_{\mu}t_{\nu}\rangle_{\mathrm{MF}}\left[\tilde{Q}_{f}(\bm{h}^{(L-1)})\right]_{\mu\nu}\\\
+\frac{1}{2}\sum_{\mu\nu}t_{\mu}t_{\nu}\left[\langle\tilde{Q}_{f}(\bm{h}^{(L-1)})\rangle_{\mathrm{MF}}\right]_{\mu\nu}+\frac{1}{2}\sum_{\mu\nu}\Delta[t_{\mu}t_{\nu}]\left[\Delta\tilde{Q}_{f}(\bm{h}^{(L-1)})\right]_{\mu\nu}+\mathrm{const}$
(37)

where we define $\Delta[t_{\mu}t_{\nu}]=t_{\mu}t_{\nu}-\langle
t_{\mu}t_{\nu}\rangle_{\mathrm{MF}}$ and
$\Delta\tilde{Q}_{f}(\bm{h}^{(L-1)})=\tilde{Q}_{f}(\bm{h}^{(L-1)})-\langle\tilde{Q}_{f}(\bm{h}^{(L-1)})\rangle_{\mathrm{MF}}$.Unlike
the hidden layers, here there is no averaging of many neurons, and therefore
a-priori it is not clear whether the fluctuating term is small. Interestingly,
we show in subsection 5.6 that this term leads to $O(1/N_{L-1})$ corrections
in observable quantities. Following this analysis, we focus on the remaining
terms, we can now integrate over the auxiliary field $\bm{t}$ and obtain:

$\int
d\bm{t}\exp\left(-\sum_{\mu}it_{\mu}f_{\mu}-\frac{1}{2}\sum_{\mu\nu}\langle
t_{\mu}t_{\nu}\rangle_{\mathrm{MF}}\left[\tilde{Q}_{f}(\bm{h}^{(L-1)})\right]_{\mu\nu}-\frac{1}{2}\sum_{\mu\nu}t_{\mu}t_{\nu}\left[\langle\tilde{Q}_{f}(\bm{h}^{(L-1)})\rangle_{\mathrm{MF}}\right]_{\mu\nu}\right)\\\
=\exp\left(-\frac{1}{2}\sum_{\mu\nu}f_{\mu}[Q_{f}]^{-1}_{\mu\nu}f_{\nu}-\frac{1}{2}\sum_{\mu\nu}\langle
t_{\mu}t_{\nu}\rangle_{\mathrm{MF}}[\tilde{Q}_{f}(\bm{h}^{(L-1)})]_{\mu\nu}\right),$
(38)

where in the second transition we define the output post-kernel as
$Q_{f}=\langle\tilde{Q}_{f}(\bm{h}^{(L-1)})\rangle_{\mathrm{MF}}$.

Our mean-field decoupling is therefore valid at $N_{L-1}\gg 1$ but finite. We
stress that even in this regime, feature learning effects can still be of
order, $1$ as these are controlled by an emergent scale ($\chi$) containing
positive powers of $n$. See Fig. 1(c) in the main text.

For the $L-1$ layer, we obtain from Eq. (35) the following action:

$\displaystyle\mathcal{S}^{(L-1)}_{\mathrm{MF}}$
$\displaystyle=\frac{1}{2}\sum_{\mu\nu}A^{(L)}_{\mu\nu}\tilde{Q}_{f}(\bm{h}^{(L-1)})+\sum_{j\mu\nu}\frac{1}{2}h_{\mu
j}^{(L-1)}[Q^{(L-1)}]_{\mu\nu}^{-1}h_{\nu j}^{(L-1)},$ (39)

where we denote by $A^{(L)}_{\mu\nu}=\langle
t_{\mu}t_{\nu}\rangle_{\mathrm{MF}}$. Notably, the fluctuations of the hidden
layer (i.e. the $\bm{h}^{(L-1)},\bm{m}^{(L-1)}$ variables) are now decoupled
both from the input layer and the output layer. The dependency on
$\bm{h}^{(L-1)}$ is only through, $Q^{(L-1)}$ and on the $\bm{m}^{(L-1)}$ only
through $\langle t_{\mu}t_{\nu}\rangle_{\mathrm{MF}}$. A tedious yet
straightforward calculation similar to that carried in Sec. 5.4.2 reveals that

$\displaystyle A^{(L)}$
$\displaystyle=\langle\bm{t}\bm{t}^{\mathsf{T}}\rangle_{\mathrm{MF}}=-\bm{\delta}\bm{\delta}^{\mathsf{T}}+\left[Q_{f}+\sigma^{2}I_{n}\right]^{-1},$
(40)

where the mean field average of $\bm{t}$ is then,

$i\bar{\bm{t}}=\bm{\delta}=\frac{\bm{y}-\bar{\bm{f}}}{\sigma^{2}}=\left[Q_{f}+\sigma^{2}I_{n}\right]^{-1}\bm{y}.$
(41)

We can also now identify the last layer of the mean-field action:

$\mathcal{S}_{f,\mathrm{MF}}=\frac{1}{2\sigma^{2}}\sum_{\mu}(f_{\mu}-y_{\mu})^{2}+\frac{1}{2}f_{\mu}[Q_{f}]^{-1}_{\mu\nu}f_{\nu}$
(42)

Combining all the layers Eq. (35), Eq. (39), and Eq. (42), we can write the
mean-field action:

$\displaystyle\mathcal{S}_{\mathrm{MF}}=\sum^{L-1}_{l=1}\mathcal{S}^{(l)}_{\mathrm{MF}}+\mathcal{S}_{f,\mathrm{MF}}.$
(43)

This allows us to define the $\langle...\rangle_{\mathrm{MF}}$ with respect to
the distribution:

$\pi_{\mathrm{MF}}(\\{\bm{h}^{(l)}\\}^{L-1}_{l=1},\bm{f})=e^{-\mathcal{S}_{\mathrm{MF}}}/\mathcal{Z}_{\mathrm{MF}}$
(44)

and the partition function
$\mathcal{Z}_{\mathrm{MF}}=\int\prod_{l=1}^{L-1}d\bm{h}^{(l)}d\bm{f}\pi_{\mathrm{MF}}(\\{\bm{h}^{(l)}\\}^{L-1}_{l=1},\bm{f})$.
We note that though this distribution is decoupled among layers and neurons in
each layer, it is parameterized by mean-field average quantities. These
quantities are defined self consistently and, unlike the infinite width case,
create both downstream and upstream dependencies between layers. These
dependencies will be fleshed out in the next subsection.

###  5.3 Variational Gaussian approximation

Despite reducing the full DNN into a product of decoupled partition functions
per neuron and layer, the resulting actions for all but the top layer, are
still non-Gaussian. Following the justifications discussed in the main text,
we approximate the latter using the variational Gaussian approximation (VGA),
in which we search for the closest (in the KL divergence sense) Gaussian
distribution, with general covariance matrix $K^{(l)}$, to the non-Gaussian
mean-field distribution of the $l$th layer. We denote the optimal covariance
matrix found for each layer $K^{(l)}$ as the pre-kernel, since it is the
approximated covariance of the pre-activation $\bm{h}^{(l)}$. In this section,
we also assume for simplicity an antisymmetric activation function, such as
$\phi=\mathrm{erf}$. Due to $\phi$’s anti-symmetry, the DNN acquires an
internal symmetry, making each pre-activation as likely as its negative. At
large enough $N_{l}$, we do not expect spontaneous symmetry breaking of this
internal symmetry, thus one may consider a simplified version of the VGA which
involves only centered Gaussian. We note that if the activation function is
not antisymmetric as in the case of ReLU, a parameter for the mean should be
added to the VGA (see Sec 9 for more details).

For the first layer, one is free to choose either pre-activations or the
weights themselves as the variables, as these two are linear functions of one
another. Here we will use the weights, $W^{(1)}_{i:}$, (the $i$th row of the
matrix $W^{(1)}$) and denote the covariance matrix of these weights as
$\Sigma$. The pre-kernel matrix of the input layer pre-activations is then
given by $K^{(1)}=X_{n}\Sigma X^{\mathsf{T}}_{n}$.

In the next two subsections, we perform the VGA to find the effective Gaussian
kernel of the $l$th hidden layer (i.e. the pre-kernel $K^{(l)}$). Note that we
apply the VGA for each layer $l$ where our reference distribution is the mean-
field distribution we found in Sec. 5.2. The $l$th layer mean-field
distribution is defined as
$e^{-\mathcal{S}^{(l)}_{\mathrm{MF}}}\propto\prod_{j}\pi^{(l)}_{j}$ where the
distribution for each neuron $j$ and layer $l$ is $\pi^{(l)}_{j}$ and
$\mathcal{S}^{(l)}_{\mathrm{MF}}$ is defined in Eq. (35), Eq. (39), and Eq.
(42). We begin with the hidden layers, turn to discuss the input and final
layer, and then finally gather all the terms and report our main result which
is the Equations of State (EoS)

####  5.3.1 $l\in[2,L-1]$

The KL divergence between the above distribution and a Gaussian distribution
for all $j\in[0,N_{l}-1]$ with covariance $K^{(l)}$ is then

$D(\mathcal{N}(0,K^{(l)})|\pi_{j}^{(l)})=-\left\langle\frac{1}{2}\sum_{\mu\nu}h_{j\mu}^{(l)}\left[Q^{(l)}\right]_{\mu\nu}^{-1}h_{j\nu}^{(l)}\right\rangle_{K^{(l)}}\\\
-\frac{\sigma_{l+1}^{2}}{2N_{l}}\sum_{\mu\nu}A^{(l+1)}_{\mu\nu}\left\langle\phi(h_{j\mu}^{(l)})\phi(h_{j\nu}^{(l)})\right\rangle_{K^{(l)}}+\frac{1}{2}\log\det(K^{(l)})+\text{Const}.$
(45)

where $\langle...\rangle_{C}$ indicate average with respect to a centered
Gaussian distribution with covariance matrix $C$, and recall that
$A^{(l)}=\langle\bm{m}^{(l)}(\bm{m}^{(l)})^{\mathsf{T}}\rangle_{\mathrm{MF}}$.
To find the pre-kernels $K^{(l)}$, we take the derivative with respect to
$K^{(l)}$ and set it to zero, in order to find the closest Gaussian
distribution:

$\partial_{K^{(l)}}D(\mathcal{N}(0,K^{(l)})|\pi_{j}^{(l)})\\\
=-\frac{1}{2}[Q^{(l)}]^{-1}-\frac{N_{l+1}}{2N_{l}}\sum_{\mu\nu}A^{(l+1)}_{\mu\nu}\partial_{K^{(l)}}Q^{(l+1)}_{\mu\nu}+\frac{1}{2}[K^{(l)}]^{-1}=0$
(46)

In the first equality, since the optimal Gaussian distribution found is the
closest to the mean-field one, the expectation of the post-kernel
$\tilde{Q}^{(l)}$ with respect to the Gaussian distribution is a good
approximation of the mean-field expectation. This is also validated using
perturbation theory in section 8. Therefore, we have that,

$\sigma_{l+1}^{2}\left\langle\phi(h_{j\mu}^{(l)})\phi(h_{j\nu}^{(l)})\right\rangle_{K^{(l)}}=\sigma_{l+1}^{2}\left\langle\phi(h_{0\mu}^{(l)})\phi(h_{0\nu}^{(l)})\right\rangle_{K^{(l)}}\\\
\approx\left\langle\tilde{Q}^{(l+1)}_{\mu\nu}(\bm{h}^{(l)})\right\rangle_{\mathrm{MF}}=Q^{(l+1)}_{\mu\nu}$
(47)

where for the last layer $\tilde{Q}^{(L)}=\tilde{Q}_{f}$, such that
$Q^{(L)}=Q_{f}$. We note here that if $Q^{(l)}$ is not full rank, then
$K^{(l)}$ has the same support (vector space associated with non-zero
eigenvalues) as $Q^{(l)}$. We can now replace in Eq. (36)
$\langle\bm{h}_{j}^{(l)}(\bm{h}_{j}^{(l)})^{\mathsf{T}{}}\rangle_{\mathrm{MF}}$
by $K^{(l)}$ such that for $l\in[2,L-1]$:

$\displaystyle
A^{(l)}\approx[Q^{(l)}]^{-1}-[Q^{(l)}]^{-1}K^{(l)}[Q^{(l)}]^{-1}.$ (48)

Note that, this is indeed the optimal distribution since taking the second
derivative yields $-[K^{(l)}]^{-2}/2$ and since $K^{(l)}$ is positive definite
as a covariance matrix $-[K^{(l)}]^{-2}/2$ is negative definite, therefore it
is a global minimum.

####  5.3.2 $l=1$

We turn to find the pre-kernel of the input layer. Here, we find it more
convenient to work with the covariance matrix of the weights rather than the
pre-activation, as it is a more compact object having fewer indices. We follow
the same procedure and minimize the KL divergence for all $i\in[0,N_{1})$ to
find the closest Gaussian distribution with covariance $\Sigma$,

$D(\mathcal{N}(0,\Sigma)|\pi_{i}^{(1)})=-\left\langle\frac{d\left\|W^{(1)}_{i:}\right\|^{2}}{2\sigma_{1}^{2}}\right\rangle_{\Sigma}\\\
-\sigma_{2}^{2}\frac{N_{2}}{2N_{1}}\sum_{\mu\nu}A^{(2)}_{\mu\nu}\left\langle\phi(h_{i\mu}^{(1)})\phi(h_{i\nu}^{(1)})\right\rangle_{\Sigma}+\frac{1}{2}\log\det(\Sigma)+const\\\
=-\frac{d}{2\sigma_{1}^{2}}\mathrm{Tr}\left(\Sigma\right)-\frac{N_{2}}{2N_{1}}\sum_{\mu\nu}A^{(2)}_{\mu\nu}\left\langle\phi(h_{i\mu}^{(1)})\phi(h_{i\nu}^{(1)})\right\rangle_{\Sigma}+\frac{1}{2}\log\det(\Sigma)+const,$

taking the derivative with respect to $\Sigma$

$\frac{\partial
D(\mathcal{N}(0,\Sigma)|\pi_{i}^{(1)})}{\partial\Sigma}=-\frac{d}{2\sigma_{1}^{2}}I_{d}-\frac{N_{2}}{2N_{1}}\sum_{\mu\nu}A^{(2)}_{\mu\nu}\left[\partial_{\Sigma}Q^{(1)}\right]_{\mu\nu}+\frac{1}{2}\Sigma^{-1}=0.$
(49)

where again we replace the expectation within the Gaussian optimal
distribution with the mean-field expectation. By the same reasoning, we can
replace in Eq. (36) the term
$\langle\bm{h}_{j}^{(1)}(\bm{h}_{j}^{(1)})^{\mathsf{T}}\rangle_{\mathrm{MF}}$
by $K^{(2)}$, and use the result in Sec. 5.4.2, to obtain:

$\displaystyle
A^{(2)}\approx[Q^{(2)}]^{-1}-[Q^{(2)}]^{-1}K^{(2)}[Q^{(2)}]^{-1}.$ (50)

Plugging in the above expression for $A^{(2)}$ and combining with Eq. (46) we
have that:

$-\frac{d}{\sigma_{1}^{2}}I_{d}-\frac{N_{2}}{N_{1}}\sum_{\mu\nu}\left([Q^{(2)}]^{-1}-[Q^{(2)}]^{-1}K^{(2)}[Q^{(2)}]^{-1}\right)_{\mu\nu}\left(\partial_{\Sigma}Q^{(2)}\right)_{\mu\nu}+\Sigma^{-1}=0.$

We note that following the VGA, one can use results of Ref. [46] to obtain
explicit expressions for the post-kernels $Q_{f},Q^{(l)}$. Using the
definition of the pre-kernels in Eq. (47) we obtain for $\phi=\mathrm{erf}$,
for all $l\in[1,L-1]$

$\displaystyle[Q^{(l+1)}]_{\mu\nu}$
$\displaystyle=\sigma_{l+1}^{2}\frac{2}{\pi}\sin^{-1}\left(\frac{2[K^{(l)}]_{\mu\nu}}{\sqrt{1+2[K^{(l)}]_{\mu\mu}}\sqrt{1+2[K^{(l)}]_{\nu\nu}}}\right),$
(51)

where $K^{(1)}=X_{n}\Sigma X^{\mathsf{T}}_{n}$ and $Q^{(L)}=Q_{f}$.

###  5.4 Equations of State (EoS)

Collecting all the results above, we obtain the following closed set of
equations determining all pre-kernel and post-kernel as well as the average
output of the Langevin algorithm, $\bar{\bm{f}}$: namely

$\displaystyle\bar{\bm{f}}$
$\displaystyle=Q_{f}[\sigma^{2}I_{n}+Q_{f}]^{-1}\bm{y}$ (52)
$\displaystyle{}[Q^{(l)}]_{\mu\nu}$
$\displaystyle=\sigma_{l}^{2}\frac{2}{\pi}\sin^{-1}\left(\frac{2[K^{(l-1)}]_{\mu\nu}}{\sqrt{1+2[K^{(l-1)}]_{\mu,\mu}}\sqrt{1+2[K^{(l-1)}]_{\nu,\nu}}}\right)$
$\displaystyle{}[[K^{(l-1)}]^{-1}]_{\mu\nu}$
$\displaystyle=[[Q^{(l-1)}]^{-1}]_{\mu\nu}+\frac{N_{l}}{N_{l-1}}\mathrm{Tr}\left\\{A^{(l)}\frac{\partial
Q^{(l)}}{\partial[K^{(l-1)}]_{\mu\nu}}\right\\}\,\,\,\text{for all
$l\in[2,L]$}$ (53) $\displaystyle[\Sigma^{-1}]_{ss^{\prime}}$
$\displaystyle=\frac{d}{\sigma_{1}^{2}}\delta_{ss^{\prime}}+\frac{N_{2}}{N_{1}}\mathrm{Tr}\left[A^{(2)}\partial_{\Sigma_{ss^{\prime}}}Q^{(2)}\right]$
(54) $\displaystyle A^{(l)}$
$\displaystyle=[Q^{(l)}]^{-1}-[Q^{(l)}]^{-1}K^{(l)}[Q^{(l)}]^{-1}\,\,\,\text{for
all $l\in[2,L-1]$}$ $\displaystyle A^{(L)}$
$\displaystyle=-(\bm{y}-\bar{\bm{f}})(\bm{y}-\bar{\bm{f}})^{\mathsf{T}}\sigma^{-4}+\left[Q_{f}+\sigma^{2}I_{n}\right]^{-1}$

where $N_{L}=1$, $Q^{(L)}=Q_{f}$, and $K^{(1)}=X_{n}\Sigma
X^{\mathsf{T}}_{n}$. While not explicitly apparent, the above expression does
converge to the GP limit as $N=N_{l}\rightarrow\infty$ for all $l\in[2,L-1]$.
Indeed, for very large $N$, $K^{(L-1)}=Q^{(L-1)}+O(1/N)$. Consequently, the
term $[[Q^{(l)}]^{-1}(K^{(l)}-Q^{(l)})[Q^{(l)}]^{-1}]$ is $O(1/N)$. The term
on the r.h.s. thus vanishes as $1/N$.

We note that the second term in Eq. (53) and Eq. (54) has a more profound
meaning in terms of the information transfer between the pre-kernel and post-
kernel. Looking at the KL divergence between two centered multidimensional
Gaussian with kernel $K^{(l)}$ and $Q^{(l)}$ of the same dimension $m=Nn$

$D_{\mathrm{KL}}(K^{(l)}||Q^{(l)})=\frac{1}{2}\left(\mathrm{Tr}\left([Q^{(l)}]^{-1}K^{(l)}\right)-m+\ln\left(\frac{\det
Q^{(l)}}{\det K^{(l)}}\right)\right).$ (55)

Taking the derivative with respect to $K^{(l-1)}$ for $l\in[3,L-1]$ and with
respect to $\Sigma$ for $l=2$, we then obtain that:

$\partial_{K^{(l-1)}}D_{\mathrm{KL}}(K^{(l)}||Q^{(l)})\\\
=\frac{1}{2}\mathrm{Tr}\left(-[Q^{(l)}]^{-1}\partial_{K^{(l-1)}}Q^{(l)}[Q^{(l)}]^{-1}K^{(l)}+[Q^{(l)}]^{-1}\partial_{K^{(l-1)}}Q^{(l)}\right)\\\
=\frac{1}{2}\mathrm{Tr}\left[\left([Q^{(l)}]^{-1}(Q^{(l)}-K^{(l)})[Q^{(l)}]^{-1}\right)\partial_{K^{(l-1)}}Q^{(l)}\right]=\frac{1}{2}\mathrm{Tr}\left[A^{(l)}\partial_{K^{(l-1)}}Q^{(l)}\right],$
(56)

where in the second transition we rearranged terms and used the cyclical
property of the trace. Substituting this relation leads to Eq. (4) provided in
the main text.

####  5.4.1 Extension to test points

The main focus of this work was on predicting the behavior of finite DNNs on
the training set. Nonetheless the results could be extended to include test
points in the following straightforward manner: Add an additional “training
point” ($f_{*}$) whose MSE term looks like
$(f_{*}-y_{*})^{2}/2\tilde{\sigma}^{2}$ where $\tilde{\sigma}^{2}$ is some
constant (rather than $(f_{\mu}-y_{\mu})^{2}/2\sigma^{2}$), repeat the
derivation, and then once the EoS are obtained take $1/\tilde{\sigma}^{2}$ to
zero. Alternatively, as done in Sec. 7, one can use large $n$ considerations
to extend the kernel found by the EoS at finite $n$ to a continuum kernel. A
more comprehensive analytical and numerical study of generalization and the
EoS is left for future work.

####  5.4.2 Auxiliary field correlation

To evaluate the second moment correlation of $\bm{m}^{(l)}$ for $l\in[2,L-1]$
under the mean-field action, we follow a standard field theory technique. We
define the partition function and introduce a source field
$J\in\mathbb{R}^{N_{l}\times n}$ ,

$\displaystyle\mathcal{Z}_{\bm{J}}[X_{n};A^{(l+1)},Q^{(l)}]=\prod_{j}\mathcal{Z}_{J_{j}}$
(57)

where

$\displaystyle\mathcal{Z}_{J_{j}}=\int
e^{-\frac{\sigma_{l+1}^{2}}{N_{l}}\sum_{\mu\nu}A_{\mu\nu}^{(l+1)}\phi(h_{j\mu}^{(l)})\phi(h_{j\nu}^{(l)})-\frac{1}{2}(\bm{J}_{j}+\bm{h}_{j}^{(l)})^{\mathsf{T}}[Q^{(l)}]^{-1}(\bm{J}_{j}+\bm{h}_{j}^{(l)})}d\bm{h}_{j}^{(l)}.$
(58)

The second moment correlation then can be found by taking twice the derivative
of the following free entropy and taking the source field to zero:

$\langle\bm{m}^{(l)}(\bm{m}^{(l)})^{\mathsf{T}}\rangle_{\mathrm{MF}}=-\frac{1}{N_{l}}\sum_{j}\triangle_{\bm{J}_{j}}\log\left(\mathcal{Z}_{\bm{J}}\right)_{\bm{J}=0},$
(59)

the vectors $\bm{h}_{j}^{(l)},\bm{J}_{j}\in\mathbb{R}^{n}$ refer to the $j$th
column of, $\bm{h}^{(l)},\bm{J}$ respectively. Note that, this partition
function is decoupled in the neurons:

$\triangle_{J}\log\left(\mathcal{Z}_{\bm{J}}\right)=\sum_{kj}\nabla_{J_{j}}^{2}\log\left(\mathcal{Z}_{J_{k}}\right)=\sum_{j}\nabla_{J_{j}}\left(\mathcal{Z}_{J_{j}}^{-1}\nabla_{J_{j}}\mathcal{Z}_{J_{j}}\right)\\\
=-\frac{1}{2}\sum_{j}\nabla_{J_{j}}\left([Q^{(l)}]^{-1}(\bm{J}_{j}+\langle\bm{h}_{j}^{(l)}\rangle_{\mathrm{MF}})+(\bm{J}_{j}+\langle\bm{h}_{j}^{(l)}\rangle_{\mathrm{MF}})^{\mathsf{T}}[Q^{(l)}]^{-1}\right)\\\
=-\frac{1}{2}\sum_{j}\left([Q^{(l)}]^{-1}\left(I_{n}+\nabla_{J_{j}}\langle\bm{h}_{j}^{(l)}\rangle_{\mathrm{MF},J}\right)+\left(I_{n}+\nabla_{J_{j}}\langle\bm{h}_{j}^{(l)}\rangle_{\mathrm{MF},J}\right)^{\mathsf{T}}Q^{(l)}]^{-1}\right),$
(60)

where we define

$\nabla_{J_{j}}\langle\bm{h}_{j}^{(l)}\rangle_{\mathrm{MF},J}|_{\bm{J}=0}\\\
=\nabla_{J_{j}}\left(\mathcal{Z}_{J_{j}}^{-1}\int
e^{-\frac{\sigma_{l+1}^{2}}{2N_{l}}\sum_{\mu\nu}A_{\mu\nu}^{(l+1)}\phi\left(h_{j\mu}^{(l)}\right)\phi\left(h_{j\nu}^{(l)}\right)-\frac{1}{2}(\bm{J}_{j}+\bm{h}_{j}^{(l)})^{\mathsf{T}}[Q^{(l)}]^{-1}(\bm{J}_{j}+\bm{h}_{j}^{(l)})}\bm{h}_{j}^{(l)}d\bm{h}_{j}^{(l)}\right)|_{\bm{J}=0}\\\
=-\langle\bm{h}_{j}^{(l)}\rangle_{\mathrm{MF}}\mathcal{Z}_{j}^{-1}\nabla_{J_{j}}\mathcal{Z}_{J_{j}}|_{J=0}-\left(\langle\bm{h}_{j}^{(l)}(\bm{h}_{j}^{(l)})^{\mathsf{T}}\rangle_{\mathrm{MF}}[Q^{(1\l)}]^{-1}\right)\\\
=\langle\bm{h}_{j}^{(l)}\rangle_{\mathrm{MF}}\langle(\bm{h}_{j}^{(l)})^{\mathsf{T}}\rangle_{\mathrm{MF}}-\left(\langle\bm{h}_{j}^{(l)}(\bm{h}_{j}^{(l)})^{\mathsf{T}}\rangle_{\mathrm{MF}}[Q^{(l)}]^{-1}\right).$

Note that, by symmetry for antisymmetric activation function:
$\langle\bm{h}_{j}^{(l)}\rangle_{\mathrm{MF}}=\langle\bm{h}_{j}^{(l)}\rangle_{\mathrm{MF},\bm{J}=0}=0$,
therefore,

$\nabla_{\bm{J}_{j}}\langle\bm{h}_{j}^{(l)}\rangle_{\mathrm{MF},J}|_{\bm{J}=0}=-\langle\bm{h}_{j}^{(l)}(\bm{h}_{j}^{(l)})^{\mathsf{T}}\rangle_{\mathrm{MF}}[Q^{(l)}]^{-1}.$

Plugging this expression back in Eq. (60), and using the fact that all
matrices are symmetric, we have that

$\triangle_{{\bm{J}}}\log\left(\mathcal{Z}_{\bm{J}}\right)=N_{l}[Q^{(l)}]^{-1}\left(\langle\bm{h}_{j}^{(l)}(\bm{h}_{j}^{(l)})^{\mathsf{T}}\rangle_{\mathrm{MF}}[Q^{(l)}]^{-1}-I_{n}\right)$

Plugging back in Eq. (59), we have that

$\langle\bm{m}^{(l)}(\bm{m}^{(l)})^{\mathsf{T}}\rangle_{\mathrm{MF}}=[Q^{(l)}]^{-1}\left(I_{n}-\langle\bm{h}_{j}^{(l)}(\bm{h}_{j}^{(l)})^{\mathsf{T}}\rangle_{\mathrm{MF}}[Q^{(l)}]^{-1}\right).$

###  5.5 An Emergent Scale

Here we identify a quantity ($\chi$) whose scale characterizes the amount of
feature learning in the DNN and its deviations from the GP limit. More
specifically, when $\chi$ becomes comparable to $1$, strong feature learning
effects appear and perturbation theory in $1/N_{l}$ becomes impractical. Since
$\chi$ would consist of a non-trivial combination of factors, we refer to it
as an emergent scale. To define this scale, we work within our EoS. and ask
when $Q_{f}$ changes in a noticeable manner as we lower $N_{l}$ from the
$N_{l}\rightarrow\infty$ at fixed $\sigma_{l}^{2}$ (the GP limit). More
technically, we next solve the EoS using perturbation theory in $N^{-1}_{l}$
and estimate the magnitude of the leading $O(1/N_{l})$ terms we obtain.

For simplicity, we focus on a setting where the penultimate layer is linear,
due to this linearity we obtain $Q_{f}=\sigma_{L}^{2}K^{(L-1)}$ and so:

$[Q_{f}]_{\mu\nu}^{-1}=\sigma_{L}^{-2}[Q^{(L-1)}]_{\mu\nu}^{-1}-\frac{1}{N_{L-1}}\left(\delta_{\mu}\delta_{\nu}-\left[Q_{f}+\sigma^{2}I_{n}\right]_{\mu\nu}^{-1}\right)$

Following the aforementioned perturbation theory in $1/N_{l}$, we perform a
first-order Taylor expansion of $Q_{f}$ yielding

$Q_{f}=\sigma_{L}^{2}\left([Q^{(L-1)}]^{-1}-\frac{\sigma_{L}^{2}}{N_{L-1}}\left(Q^{(L-1)}\bm{\delta}\bm{\delta}^{\mathsf{T}}-Q^{(L-1)}\left[Q_{f}+\sigma^{2}I_{n}\right]^{-1}\right)\right)^{-1}\\\
=\sigma_{L}^{2}Q^{(L-1)}+\frac{\sigma_{L}^{2}}{N_{L-1}}Q^{(L-1)}\bm{\delta}\bm{\delta}^{\mathsf{T}}Q^{(L-1)}-\frac{\sigma_{L}^{2}}{N_{L-1}}Q^{(L-1)}\left[Q_{f}+\sigma^{2}I_{n}\right]^{-1}Q^{(L-1)}+O(N_{L-1}^{-2})$
(61)

To evaluate the magnitude of the first term we multiply by $\bm{\delta}$ from
both sides to obtain

$\bm{\delta}^{\mathsf{T}}Q_{f}\bm{\delta}=\sigma_{L}^{2}\bm{\delta}^{\mathsf{T}}Q^{(L-1)}\bm{\delta}+\frac{\sigma_{L}^{2}}{N_{L-1}}\bm{\delta}^{\mathsf{T}}Q^{(L-1)}\left(\bm{\delta}\bm{\delta}^{\mathsf{T}}-\left[Q_{f}+\sigma^{2}I_{n}\right]^{-1}\right)Q^{(L-1)}\bm{\delta}+O(N_{L-1}^{-2})$
(62)

We define
$\chi=\frac{1}{N_{L-1}}\bm{\delta}^{\mathsf{T}}Q^{(L-1)}\bm{\delta}$, and
find:

$\frac{1}{N_{L-1}\sigma_{L}^{2}}\bm{\delta}^{\mathsf{T}}Q_{f}\bm{\delta}=\chi+\chi^{2}+\frac{1}{N_{L-1}^{2}}\bm{\delta}^{\mathsf{T}}Q^{(L-1)}\left[Q_{f}+\sigma^{2}I_{n}\right]^{-1}Q^{(L-1)}\bm{\delta}+O(N_{L-1}^{-3})$
(63)

As we argue below, the last term on the r.h.s is small compared to the first
two. Putting it aside and recalling that the first term on the right-hand side
is the zeroth order term in $1/N_{l}$, we find that $\chi$ controls the ratio
between the zeroth order contribute and the first order perturbative
correction. Hence, once $\chi$ becomes order $1$, first-order perturbation in
$1/N_{l}$ becomes inaccurate. However, it can be further checked that a
second-order perturbation will contain a $\chi^{3}$ contribution, and hence
this is not just a problem in first-order perturbation theory. Rather, it is
that low order perturbation theory becomes inaccurate.

Next, we argue that having non-negligible $\chi$ also implies feature
learning, in the sense that $Q_{f}$ changes from its $N_{l}\rightarrow\infty$
value. Indeed, the quantity we are examining
($\bm{\delta}^{\mathsf{T}}Q_{f}\bm{\delta}$) involves both the discrepancy
($\bm{\delta}$) and the kernel $Q_{f}$. Thus, a-priory may change just because
$\bm{\delta}$ changes. However, $\bm{\delta}$ is a function of $Q_{f}$ via
$\bm{\delta}=[Q_{f}+\sigma^{2}I]^{-1}\bm{y}$. Thus, it cannot undergo any
change if $Q_{f}$ remains inert. Thus, we conclude that a change to
$\bm{\delta}^{\mathsf{T}}Q_{f}\bm{\delta}$ must come from a change in $Q_{f}$
and hence, by our definition, from feature learning. This combined with the
previous paragraph also shows that strong feature learning effects are beyond
the practical reach of straightforward perturbation theory.

Following some assumptions on the support of $\bm{\delta}$, the emergent scale
can be presented in a more explicit manner. Indeed, assuming $\bm{\delta}$ has
support mainly on $Q^{(L-1)}$ leading eigenvalues which are order
$\lambda_{\max}$, we can estimate $\chi\sim
n^{2}\mathrm{MSE}\lambda_{\max}/N_{L-1}$, where we also use
$\left\|\bm{\delta}\right\|^{2}=n\mathrm{MSE}$. Last we note that $\chi$ is
also a relevant scale for nonlinear penultimate layers since even in that case
$Q_{f}$ can be expanded in $K^{(L-1)}$ and the linear term in this expansion
will yield a correction to $\bm{\delta}^{\mathsf{T}}Q_{f}\bm{\delta}$
proportional to the emergent scale.

We turn to estimate the magnitude of the term we neglected namely

$\displaystyle\frac{1}{N_{L-1}^{2}}\bm{\delta}^{\mathsf{T}}Q^{(L-1)}\left[Q_{f}+\sigma^{2}I_{n}\right]^{-1}Q^{(L-1)}\bm{\delta}$
(64)

Noting $Q_{f}=\sigma_{L}^{2}K^{(L-1)}$ and that, following our EoS.,
$K^{(L-1)}=Q^{(L-1)}+O(1/N_{l})$, we can rewrite the above term up to
$1/N_{l}^{2}$ corrections as

$\displaystyle\frac{1}{N_{L-1}^{2}}\bm{\delta}^{\mathsf{T}}Q^{(L-1)}\left[\sigma_{L}^{2}Q^{(L-1)}+\sigma^{2}I_{n}\right]^{-1}Q^{(L-1)}\bm{\delta}$
(65)

following this one can establish that

$\displaystyle\frac{1}{N_{L-1}^{2}}\bm{\delta}^{\mathsf{T}}Q^{(L-1)}\left[\sigma_{L}^{2}Q^{(L-1)}+\sigma^{2}I_{n}\right]^{-1}Q^{(L-1)}\bm{\delta}\leq\frac{1}{N_{L-1}^{2}\sigma_{L}^{2}}\bm{\delta}^{\mathsf{T}}Q^{(L-1)}\bm{\delta}=\frac{\chi}{N_{L-1}}$
(66)

which is indeed negligible compared to $\chi$ at large $N_{L-1}$.

###  5.6 Estimating Corrections to Mean-field Results

In the mean-field derivation, when focusing on the two last layers, we
neglected the term

$\displaystyle\frac{1}{2}\sum_{\nu\mu}\Delta[tt]_{\mu\nu}[\Delta\tilde{Q}_{f}(\bm{h}^{(L-1)})]_{\mu\nu}$
(67)

where $\Delta[tt]_{\mu\nu}=t_{\mu}t_{\nu}-\langle
t_{\mu}t_{\nu}\rangle_{\mathrm{MF}}$ and
$\Delta\tilde{Q}_{f}(\bm{h}^{(L-1)})=\tilde{Q}_{f}(\bm{h}^{(L-1)})-{Q}_{f}$.
However, unlike in the case of hidden layers, where the above two "$\Delta$"
terms were both clearly a sum over many independent random variables (in the
mean-field picture), here only the 2nd $\Delta$ is of that form. Hence, it is
a priory less clear what makes this term negligible.

In this section, we study the effect of this term in perturbation theory and
when it can be neglected.

Specifically, we treat the above term as a perturbation over the mean-field
limit and calculate its leading order effect on the mean-field average of,
$(\bm{f}-\bm{y})/\sigma^{2}$ which coincides with the mean-field average of
$\bm{t}$ which is directly related to the MSE. To expose this matter in its
simplest form, we shall assume that the penultimate layer is linear.
Consequently $Q_{f}=\sigma^{2}_{L}K^{(L-1)}$ and in addition, VGA becomes
exact (see Eq. 33 and definition of $Q_{f}(\bm{h}^{(L-1)})$).

Turning to an action formulation, we focus on the following mean-field action
of the $\bm{h}^{(L-1)},\bm{f}$ and $\bm{t}$ variables together with the
perturbation namely,

$\displaystyle\mathcal{S}_{f}=\mathcal{S}_{f,0}+\Delta\mathcal{S}_{f}$ (68)

where

$\displaystyle\mathcal{S}_{f,0}$
$\displaystyle=i\sum_{\mu}t_{\mu}f_{\mu}+\frac{1}{2\sigma^{2}}\sum_{\mu}(f_{\mu}-y_{\mu})^{2}+\frac{1}{2}\sum_{\mu\nu}\langle
t_{\mu}t_{\nu}\rangle_{\mathrm{MF}}\left[\tilde{Q}_{f}(\bm{h}^{(L-1)})\right]_{\mu\nu}$
(69)
$\displaystyle+\frac{1}{2}\sum_{\mu\nu}t_{\mu}t_{\nu}\left[\langle\tilde{Q}_{f}(\bm{h}^{(L-1)})\rangle_{\mathrm{MF}}\right]_{\mu\nu}$
$\displaystyle\Delta\mathcal{S}_{f}$
$\displaystyle=\frac{1}{2}\sum_{\mu\nu}\Delta[t_{\mu}t_{\nu}]\left[\Delta\tilde{Q}_{f}(\bm{h}^{(L-1)})\right]_{\mu\nu}.$
(70)

Perturbation in the $\Delta\mathcal{S}_{f}$ term yields the following zeroth
and first-order contributions

$\displaystyle\langle t_{\nu}\rangle$ $\displaystyle=\langle
t_{\nu}\rangle_{\mathrm{MF}}$ (71) $\displaystyle+\frac{1}{2}\sum_{a,b}\langle
t_{\nu}\Delta[t_{a}t_{b}]\Delta\tilde{Q}_{f}(\bm{h}^{(L-1)})]_{ab}\rangle_{\mathrm{MF,con}}+\text{Higher
order terms}$

where $\langle...\rangle_{\mathrm{MF},\mathrm{con}}$ means connected average
(i.e. one in which, following Wick’s theorem, the perturbation is never
contracted/averaged just with itself), and
$K_{f}=\left[Q_{f}+\sigma^{2}I_{n}\right]$. Clearly, the first order
contribution is strictly zero, since by definition
$\langle[\Delta\tilde{Q}_{f}(\bm{h}^{(L-1)})]_{ab}\rangle_{\mathrm{MF}}=0$ and
due to the mean-field decoupling under $\langle..\rangle_{\rm{MF}}$, $h$’s can
only contract with themselves. We thus turn to the first non-zero contribution
which occurs in second order namely

$\displaystyle\frac{1}{2!4}\langle
t_{\nu}\sum_{abde}\Delta[tt]_{ab}\Delta[tt]_{de}[\Delta\tilde{Q}_{f}(\bm{h}^{(L-1)})]_{ab}[\Delta\tilde{Q}_{f}(\bm{h}^{(L-1)})]_{de}\rangle_{\mathrm{MF,con}}.$
(72)

Noting that we are restricted to connected averages, the fact that within the
mean-field picture $\bm{h}^{(L-1)}$’s and $\bm{t}$’s are uncorrelated as well
as ${h}^{(L-1)}_{i}$’s between different neurons, and the fact that
$\tilde{Q}_{f}=\sigma_{L}^{2}N_{L-1}^{-1}\sum_{j=1}^{N_{L-1}}\bm{h}^{(L-1)}_{j}(\bm{h}_{j}^{(L-1)})^{\mathsf{T}}$
we obtain the following correction to $\bar{t}_{\mu}$

$\displaystyle\frac{2}{N_{L-1}}[K_{f}^{-1}Q_{f}K_{f}^{-1}Q_{f}\bar{\bm{t}}]_{\nu}$
(73)
$\displaystyle+\frac{1}{N_{L-1}}[K_{f}^{-1}Q_{f}\bar{\bm{t}}]_{\nu}\mathrm{Tr}[K_{f}^{-1}Q_{f}]$

Recalling that $K_{f}=Q_{f}+\sigma^{2}I_{n}$, the first term in the correction
to $\bar{t}_{\nu}$ is upper bounded by $\frac{2}{N_{L-1}}\bar{t}_{\nu}$ and
hence negligible. The first piece of the second term is similarly upper
bounded by $\frac{1}{N_{L-1}}\bar{t}_{\nu}$, however it also contains a factor
of
$\mathrm{Tr}[K_{f}^{-1}Q_{f}]=\sum_{\lambda}\sigma_{L}^{2}\lambda/(\sigma_{L}^{2}\lambda+\sigma^{2})$,
where $\sigma_{L}^{2}\lambda$’s are eigenvalues of $Q_{f}$.

Let us discuss how these factors behave in various settings. The first
settings are similar to that in the main text, only with standard rather than
mean-field scaling for the top layer weight decay. Here, the
$\mathrm{Tr}[K_{f}^{-1}Q_{f}]$ term at $\sigma^{2}\rightarrow 0$ is equal to
$n$. Similarly, our upper bound for the first factor becomes exact. Together,
this leads to an over $n/N_{L-1}$ ratio between the correction to the mean-
field result and the mean-field result for $\bar{t}_{\mu}$. Hence, when $n$
becomes larger than, $N_{L-1}$ we expect the mean-field treatment to fail.

Turning to nonzero $\sigma^{2}$, the spectrum of $Q_{f}$ starts to play a
role. Indeed, since the spectrum of $Q_{f}$ typically decays quickly, [6]
$\mathrm{Tr}[K_{f}^{-1}Q_{f}]$ becomes proportional to the number of $Q_{f}$
eigenvalues which are larger than $\sigma^{2}$. In addition, upper bounding
$[K_{f}^{-1}Q_{f}\bar{\bm{t}}]_{\nu}$ by $\bar{t}_{\nu}$ may be a very loose
bound if $\bar{t}$ has support on $Q_{f}$ eigenvalues which are smaller than
$\sigma^{2}$. Hence, estimating the magnitude of these corrections requires a
careful case-by-case inspection. Some preliminary numerical checks we perform
in teacher-student settings similar to those in the main text for
$\sigma^{2}\leq 0.01$ showed that the corrections to mean-field became large
as soon as feature learning began. We believe that our theory still applies to
FCNs away from student-teacher settings, in particular in cases where the
teacher induces a target that has support mainly on sub-leading eigenvalues of
$Q_{f}$. However, we leave this matter for future study. We comment that CNNs
are different in that aspect– as shown later below, they appear to obey our
mean-field treatment within the standard scaling.

Last, we consider nonzero $\sigma^{2}$ with mean-field scaling. Here one can
capitalize from the fact that the mean-field scaling, scales down $Q_{f}$,
thereby making both $[K_{f}^{-1}Q_{f}\bar{\bm{t}}]_{\nu}$ and
$\mathrm{Tr}[K_{f}^{-1}Q_{f}]$ smaller. To address this issue most cleanly,
let us define a GP limit to the mean-field scaling wherein we fix
$\sigma^{2}_{l}$ at order $1$, $\sigma_{L}^{2}$ at order $\epsilon$
(independent from $N_{L-1}$), and take $N_{l}\rightarrow\infty$. Here we find
that the eigenvalues of $Q_{f}$ are those of the standard GP limit, with an
additional factor of $\epsilon$. Similarly to the emergent scale section, we
can now ask when the mean-field correction becomes noticeable as we reduce
$N_{l}$. Assuming $\chi$ is finite but not yet of order $1$, $Q_{f}$ is not
expected to change drastically. Hence, to obtain an estimate of the mean-field
correction, we may use the $N_{l}\rightarrow\infty$ values for $Q_{f}$ and
$\bar{t}_{\mu}$. At small, $\epsilon$ this implies that
$\mathrm{Tr}[K_{f}^{-1}Q_{f}]\approx\mathrm{Tr}[Q_{f}]/\sigma^{2}$ which is of
the order of $n\epsilon/\sigma^{2}$. Similarly,
$[K_{f}^{-1}Q_{f}\bar{\bm{t}}]_{\nu}$ can be estimated by
$\epsilon\lambda_{\max}/\sigma^{2}$, where $\lambda_{\max}$ is the dominant
eigenvalue of $Q_{f}$ on which $\bm{t}$ has support. Putting this together
with the ${N_{L-1}^{-1}}$ factor we find that
$n\lambda_{\max}\epsilon^{2}/(N_{L-1}\sigma^{4})\ll 1$ is required which is a
much less stringent requirement. Indeed, by taking $\epsilon$ to zero, one can
make this as small as needed.

A key point is that making $\epsilon$ smaller makes $\chi$ bigger. Hence, our
mean-field treatment works within the feature learning regime. Indeed, the
emergent scale is defined as
$\bm{\delta}^{\mathsf{T}}Q^{(L-1)}{\bm{\delta}}/N_{L-1}$ and hence does not
involve $\epsilon$ directly. Decreasing $\epsilon$ can only make
${\bm{\delta}}$ larger since
${\bm{\delta}}=[Q_{f}+\sigma^{2}I_{n}]^{-1}{\bm{y}}$. Thus, decreasing
$\epsilon$ increases feature learning while making mean-field corrections more
negligible. On this note, we comment that in our FCN experiments we found,
similarly to Ref. [49], that mean-field scaled FCNs had better test
performance compared to those which used standard scaling similar to

##  6 Mean-Field Equations for CNNs

In this section, we provide a sketch of the derivation of the equations of
state for deep CNNs highlighting the differences between CNN and FCN. For
simplicity, we follow here the derivation of a three-layer CNN. The
generalization to any number of layers is straightforward. The model we
consider is a three-layer CNN having two activated convolutional layers and
one linear readout layer. Specifically, we consider

$\displaystyle f(\bm{x})$
$\displaystyle=\sum_{j=0}^{N-1}\sum_{c^{\prime}=1}^{C_{2}}a_{c^{\prime}j}\phi\left(h^{(2)}_{c^{\prime}j}(\bm{x})\right)$
(74) $\displaystyle h^{(2)}_{c^{\prime}j}(\bm{x})$
$\displaystyle=\sum_{i=0}^{S_{1}-1}\sum_{c=1}^{C_{1}}v_{c^{\prime}ci}\phi\left(\bm{w}_{c}\cdot\bm{x}_{i+jS_{1}}\right)$

where $\bm{w}_{c},\bm{x}_{i+jS_{1}}\in\mathbb{R}^{S_{0}}$,
$\bm{a}\in\mathbb{R}^{C_{2}\times N}$, $\bm{v}\in\mathbb{R}^{C_{2}\times
C_{1}\times S_{1}}$, and the input vector $\bm{x}\in\mathbb{R}^{d}$ with
$d=S_{0}S_{1}N$. Bold letters mark vector or tensor quantities. Similar to the
FCN case, we analyze the following distribution over the pre-activations

$\displaystyle p(\bm{\bm{f}}|\mathcal{D}_{n})$ $\displaystyle\propto
p(\bm{f}|X_{n})\mathrm{exp}\left(-\frac{1}{2\sigma^{2}}\sum_{\mu=1}^{n}\left(y_{\mu}-f(\bm{x}_{\mu})\right)^{2}\right),$
(75)

where the matrix $X_{n}$ represents all the samples,
$\\{\bm{x}_{\mu}\\}^{n}_{\mu=1}$, and the first term on the r.h.s. is given by

$\displaystyle p(\bm{f}|X_{n})$
$\displaystyle=\left\langle\prod_{\mu}\delta\left[f(\bm{x}_{\mu})-\sum_{j=0}^{N-1}\sum_{c^{\prime}=1}^{C_{2}}a_{c^{\prime}j}\phi\left(h^{(2)}_{c^{\prime}j}(\bm{x}_{\mu})\right)\right]\right\rangle_{\bm{a,v,w}}$
(76)

where $\bm{f}=(f_{1},...,f_{n})$ is viewed now as a random variable following
the CNN outputs on all different training points, and
$\langle...\rangle_{\bm{a,v,w}}$ denote average over the weights $\bm{a,v,w}$.
The above probability can be viewed as the prior induced on $\bm{f}$ by a
finite random DNN with i.i.d. Gaussian weights and variance
$\sigma_{\text{w}}^{2}/S_{0},\sigma_{\text{v}}^{2}/(S_{1}C_{1})$ and
$\sigma_{\text{a}}^{2}/(NC_{2})$ respectively for each layer. At
$C_{1},C_{2}\rightarrow\infty$ such priors tend to a GP, however our interest
here is at finite $C_{1},C_{2}$.

By conditioning on the pre-activation output ($\bm{h}^{(2)}$), Eq. (76) can be
re-written as

$\displaystyle p(\bm{f}|X_{n})$ $\displaystyle=\int
p(\bm{f}|\bm{h}^{(2)},X_{n})p(\bm{h}^{(2)}|\bm{w},X_{n})p(\bm{w})\mathcal{D}\bm{h}^{(2)}\mathcal{D}\bm{w}$
(77)

where the hidden layers probabilities used above are

$\displaystyle p(\bm{f}|\bm{h}^{(2)},X_{n})$
$\displaystyle=\left\langle\prod_{\mu}\delta\left(f_{\mu}-\sum_{j=0}^{N-1}\sum_{c^{\prime}=1}^{C_{2}}a_{c^{\prime}j}\phi\left(h^{(2)}_{c^{\prime}j\mu}\right)\right)\right\rangle_{\bm{a}},$
(78) $\displaystyle p(\bm{h}^{(2)}|\bm{w},X_{n})$
$\displaystyle=\left\langle\prod_{\mu
c^{\prime}j}\delta\left(h^{(1)}_{c^{\prime}j\mu}-\sum_{i=0}^{S_{1}-1}\sum_{c=1}^{C_{1}}v_{c^{\prime}ci}\phi\left(\bm{w}_{c}\cdot\bm{x}_{\mu,i+jS_{1}}\right)\right)\right\rangle_{\bm{v}},$

the tensor $\bm{h}^{(2)}$ consists of all $h^{(2)}_{ic\mu}$ the latter being
the random variables describing outputs of the second layer at hidden pixel
$i$, channel $c$ on the $\bm{x}_{\mu}$ data-point.

Similar to the FCN, we continue our analysis by using the Fourier identity,
which replaces the delta function with auxiliary fields $\bm{t},\bm{m}$. The
resulting action ($Z=\int e^{-\mathcal{S}}$) is then,

$\displaystyle\mathcal{S}$
$\displaystyle=\frac{S_{0}}{2\sigma_{\text{w}}^{2}}\sum_{c}\norm{\bm{w}_{c}}^{2}-i\sum_{\mu
c^{\prime}j}m_{\mu
c^{\prime}j}h^{(2)}_{c^{\prime}j}(\bm{x}_{\mu})+\frac{1}{2}\sum_{\mu\nu
j_{1}j_{2}c^{\prime}}m_{\mu c^{\prime}j_{1}}m_{\nu
c^{\prime}j_{2}}\tilde{Q}^{(2)}(\bm{w})_{\mu\nu j_{1}j_{2}}$ (79)
$\displaystyle-i\sum_{\mu}t_{\mu}f_{\mu}+\frac{1}{2}\sum_{\mu\nu}t_{\mu}t_{\nu}\left[\tilde{Q}_{f}(\bm{h}^{(2)})\right]_{\mu\nu}+\frac{1}{2\sigma^{2}}\sum_{\mu}(f_{\mu}-y_{\mu})^{2}$

Where the “channel” post-kernels for CNN contain also summation overstrides
and are defined as:

$\displaystyle\tilde{Q}_{f}(\bm{h}^{(2)})_{\mu\nu}$
$\displaystyle=\frac{\sigma_{\text{a}}^{2}}{C_{2}N}\sum_{jc^{\prime}}^{N,C_{2}}\phi\left(h^{(2)}_{c^{\prime}j\mu}\right)\phi\left(h^{(2)}_{c^{\prime}j\nu}\right)$
(80) $\displaystyle\tilde{Q}^{(2)}(\bm{w})_{\mu j_{1},\nu j_{2}}$
$\displaystyle=\frac{\sigma_{\text{v}}^{2}}{C_{1}S_{1}}\sum^{S_{1},C_{1}}_{ic}\phi\left(\bm{w}_{c}\cdot\bm{x}_{\mu,i+j_{1}S_{1}}\right)\phi\left(\bm{w}_{c}\cdot\bm{x}_{\nu,i+j_{2}S_{1}}\right).$

We comment that by averaging over the auxiliary fields, $\bm{t},\bm{m}$, one
obtains the following equivalent form, containing only the pre-activations and
the outputs

$\mathcal{S}=\frac{S_{0}}{2\sigma_{\text{w}}^{2}}\sum_{c}\norm{\bm{w}_{c}}^{2}+\frac{1}{2}\sum_{c^{\prime}\mu\nu
jj^{\prime}}h^{(2)}_{c^{\prime}j\mu}\left[\tilde{Q}^{(2)}(\bm{w})\right]^{-1}_{j\mu,j^{\prime}\nu}{h}^{(2)}_{c^{\prime}j^{\prime}\nu}\\\
+\frac{1}{2}f_{\mu}\left[\tilde{Q}_{f}(\bm{h}^{(2)})\right]^{-1}_{\mu,\nu}f_{\nu}+\frac{1}{2\sigma^{2}}\sum_{\mu}(f_{\mu}-y_{\mu})^{2}$
(81)

We continue with the auxiliary variables and derive the inter-layer mean-field
decoupling. Where for CNN the number of channels $C_{1},C_{2}$ plays the role
of width in FCN. As for FCN, we note that in Eq. (79), the hidden layer and
the output layer depend on their respective upstream layers only through the
"channel" post-kernels. Performing our mean-field decoupling as in Sec. 5.2 we
obtain the resulting action

$\displaystyle\mathcal{S}_{\mathrm{MF}}$
$\displaystyle=\mathcal{S}^{(1)}_{\mathrm{MF}}+\mathcal{S}^{(2)}_{\mathrm{MF}}+\mathcal{S}_{f,\mathrm{MF}}$
(82) $\displaystyle\mathcal{S}^{(1)}_{\mathrm{MF}}$
$\displaystyle=\frac{S_{0}}{2\sigma_{\text{w}}^{2}}\sum_{c}\norm{\bm{w}_{c}}^{2}+\frac{1}{2}\sum_{\mu\nu
j_{1}j_{2}c^{\prime}}\langle{\bm{m}\bm{m}^{\mathsf{T}}}\rangle_{\mathrm{MF},\mu
j_{1}\nu j_{2}}\tilde{Q}^{(2)}(\bm{w})_{\mu j_{1}\nu j_{2}}$
$\displaystyle\mathcal{S}^{(2)}_{\mathrm{MF}}$
$\displaystyle=\frac{1}{2}\sum_{\mu\nu}\langle
t_{\mu}t_{\nu}\rangle_{\mathrm{MF}}\left[\tilde{Q}_{f}(\bm{h}^{(2)})\right]_{\mu\nu}+\frac{1}{2}\sum_{\mu
c^{\prime}j_{1},j_{2}}h^{(2)}_{c^{\prime}j_{1}}(\bm{x}_{\mu})[{Q}^{(2)}]_{\mu
j_{1}\nu j_{2}}^{-1}h^{(2)}_{c^{\prime}j_{2}}(\bm{x}_{\mu})$
$\displaystyle\mathcal{S}_{f,\mathrm{MF}}$
$\displaystyle=\frac{1}{2\sigma^{2}}\sum_{\mu}(f_{\mu}-y_{\mu})^{2}+\frac{1}{2}\sum_{\mu\nu}f_{\mu}[Q_{f}]_{\mu\nu}^{-1}f_{\nu},$

where the post-kernel can now be defined as the mean-field average of the
“channel” post-kernel via the above mean-field action distribution:

$\displaystyle[Q_{f}]_{\mu\nu}$
$\displaystyle=\left\langle\tilde{Q}_{f}(\bm{h}^{(2)})_{\mu\nu}\right\rangle_{\mathrm{MF}}$
(83) $\displaystyle[Q^{(2)}]_{\mu j,\nu j^{\prime}}$
$\displaystyle=\left\langle\tilde{Q}^{(2)}(\bm{w})_{\mu j_{1},\nu
j_{2}}\right\rangle_{\mathrm{MF}}.$

The correlation functions similar to the FCN are then,

$\langle{\bm{m}\bm{m}^{\mathsf{T}}}\rangle_{\mathrm{MF}}=\left\langle
C_{2}^{-1}\sum^{C_{2}}_{c^{\prime},j=1}\bm{m}_{c^{\prime}}\bm{m}^{\mathsf{T}}_{c^{\prime}}\right\rangle_{\mathrm{MF}}=[Q^{(2)}]^{-1}\left(I_{n}-\langle\bm{h}_{1}^{(2)}(\bm{h}_{1}^{(2)})^{\mathsf{T}}\rangle_{\mathrm{MF}}[Q^{(2)}]^{-1}\right)$
(84)

The last layer of auxiliary field correlation is

$\displaystyle\langle
t_{\mu}t_{\nu}\rangle_{\mathrm{MF}}=-\delta_{\mu}\delta_{\nu}+\left[(Q_{f}+\sigma^{2}I_{n})^{-1}\right]_{\mu\nu}$
(85)

this is easily derive from Eq. (82). The average of $\bm{t}$ using the mean-
field action is then,

$\displaystyle
i\bar{\bm{t}}=\bm{\delta}=\frac{\bm{y}-\bar{\bm{f}}}{\sigma^{2}}=\left(Q_{f}+\sigma^{2}I_{n}\right)^{-1}\bm{y},$
(86)

where we denote by bar the $\langle...\rangle_{\mathrm{MF}}$ corresponding to
averages with respect to the mean-field distribution. Note that, in the main
text, we denote $\tilde{Q}^{(2)}(\bm{w})$ by $\tilde{Q}^{(2)}(\bm{h}^{(1)})$.
These are the same quantities, the connection being
${h}_{cij}^{(1)}=\bm{w}_{c}\cdot\bm{x}_{i+jS_{1}}$.

Despite reducing the full system into decoupled systems per channel and layer,
the resulting action for all but the top layer is still non-Gaussian.
Following the justifications discussed in the main text and for the FCN, we
approximate the latter using the variational Gaussian approximation (VGA).
Assuming for simplicity an antisymmetric activation function such as $\erf$,
the CNN has an internal symmetry, making each pre-activation positive output
as likely as its negative. Also, at large enough $C_{1},C_{2}$ we do not
expect spontaneous symmetry breaking, thus our VGA will involve only a
centered Gaussian. Specifically, we denote the optimal covariance of
$\bm{h}^{(2)}$, $K^{(2)}_{\mu j\nu j^{\prime}}$ as the pre-kernel of the
second layer and the optimal covariance of $\bm{w}$, $\Sigma_{ss^{\prime}}$,
is a connected to the first layer pre-kernel, $K^{(1)}=X_{n}\Sigma
X_{n}^{\mathsf{T}}$.

Following the analysis in subsection. 5.3 with the above modification for CNN,
we obtain the following closed set of equations determining all pre-kernels as
well as the outputs namely

$\displaystyle\bar{\bm{f}}$
$\displaystyle=Q_{f}[\sigma^{2}I_{n}+Q_{f}]^{-1}\bm{y}$ (87)
$\displaystyle[Q_{f}]_{\mu\nu}$
$\displaystyle=\sigma_{\text{a}}^{2}\frac{1}{N}\sum_{j}\frac{2}{\pi}\sin^{-1}\left(\frac{2[K^{(2)}]_{\mu
j,\nu j}}{\sqrt{1+2[K^{(2)}]_{\mu j,\mu j}}\sqrt{1+2[K^{(2)}]_{\nu j,\nu
j}}}\right)$ $\displaystyle[Q^{(2)}]_{\mu j,\nu j^{\prime}}$
$\displaystyle=\sigma_{\text{v}}^{2}\frac{1}{S_{1}}\sum_{i}\frac{2}{\pi}\sin^{-1}\left(\frac{2\bm{x}_{\mu
ji}\Sigma\bm{x}_{\nu j^{\prime}i}}{\sqrt{1+2\bm{x}_{\mu ji}\Sigma\bm{x}_{\mu
ji}}\sqrt{1+2\bm{x}_{\nu j^{\prime}i}\Sigma\bm{x}_{\nu j^{\prime}i}}}\right)$
$\displaystyle[(K^{(2)})^{-1}]_{\mu j,\nu j^{\prime}}$
$\displaystyle=[(Q^{(2)})^{-1}]_{\mu j,\nu
j^{\prime}}-\frac{1}{C_{2}}\Tr{A^{(3)}\frac{\partial
Q_{f}}{\partial[K^{(2)}]_{\mu j,\nu j^{\prime}}}}$
$\displaystyle[\Sigma^{-1}]_{ss^{\prime}}$
$\displaystyle=\frac{S_{0}}{\sigma_{\text{w}}^{2}}\delta_{ss^{\prime}}-\frac{C_{2}}{C_{1}}\Tr\left[[(Q^{(2)})^{-1}(K^{(2)}-Q^{(2)})(Q^{(2)})^{-1}]\partial_{\Sigma_{ss^{\prime}}}Q^{(2)}\right]$
$\displaystyle A^{(3)}$
$\displaystyle=(\bm{y}-\bar{\bm{f}})(\bm{y}-\bar{\bm{f}})^{\mathsf{T}}\sigma^{-4}-\left[Q_{f}+\sigma^{2}I_{n}\right]^{-1}$

###  6.1 An Emergent Scale in CNNs

Following the FCN case, we again examine the EoS for the penultimate layer
assuming that it is linear, solve them using perturbation theory in $1/C_{l}$,
and use the magnitude of the correction to estimate the scale at which feature
learning becomes important. For our CNNs we have that
$[Q_{f}]_{\mu\nu}=\sigma_{L}^{2}N^{-1}\sum_{j}K^{(L-1)}_{\mu j,\nu j}$ which
yields the following equation of state for $K^{(L-1)}$

$[K^{(L-1)}]_{\mu i,\nu j}^{-1}=[Q^{(L-1)}]_{\mu i,\nu
j}^{-1}-\frac{\delta_{ij}}{C_{L-1}N}\left(\delta_{\mu}\delta_{\nu}-\left[Q_{f}+\sigma^{2}I_{n}\right]_{\mu\nu}^{-1}\right)$

Next, we perform a leading order perturbation theory in $1/C_{l}$ (or
equivalently in the second term on the r.h.s) yielding

$[K^{(L-1)}]_{\mu i,\nu j}=[Q^{(L-1)}]_{\mu i,\nu
j}+\frac{1}{C_{L-1}N}\sum_{kab}[Q^{(L-1)}]_{\mu
i,ak}\delta_{a}\delta_{b}[Q^{(L-1)}]_{bk,\nu j}\\\
-\frac{1}{C_{L-1}N}\sum_{kab}[Q^{(L-1)}]_{\mu
i,ak}[Q_{f}+\sigma^{2}I_{n}]^{-1}_{ab}[Q^{(L-1)}]_{bk,\nu
j}+O(1/(C_{L-1}^{2}))$ (88)

As justified in the FCN case, we focus on the second term on the r.h.s. and
look again at $\bm{\delta}^{\mathsf{T}}Q_{f}\bm{\delta}$ yielding,

$\displaystyle\bm{\delta}^{\mathsf{T}}Q_{f}\bm{\delta}$
$\displaystyle=\frac{\sigma^{2}_{L}}{N}\sum_{i\mu\nu}\delta_{\mu}Q^{(L-1)}_{\mu
i,\nu
i}\delta_{\nu}+\frac{\sigma_{L}^{2}}{C_{L-1}N^{2}}\sum_{ikab\mu\nu}\delta_{\mu}[Q^{(L-1)}]_{\mu
i,ak}\delta_{a}\delta_{b}[Q^{(L-1)}]_{bk,\nu i}\delta_{\nu}$ (89)

We define the ratio of the second to the first term as the emergent scale.
Notably for $N=1$ it coincides with the definition for FCNs. In addition to
considering the 2-layer CNN studied in the main text and estimating it using
the same approximations, it results in the same scale.

###  6.2 Estimating Mean-field corrections - CNNs

In the FCN case, we found that an additional ingredient, on top of large
$N_{l}$, is needed for our mean-field decoupling to hold - either mean-field
scaling or a target with support only on weak $Q_{f}$ eigenvalues. For the
CNNs we have studied, and quite possibly for a much larger family of CNNs,
this additional ingredient comes naturally from the read-out layer averages
over $N$ latent pixels in the penultimate layers. Since these are expected to
be somewhat independent, one can hope that summing over these terms is similar
to increasing the number of channels by a factor of $N$ (the number of pixels
in the penultimate layer). Here we establish this more concretely. Similar to
the FCN we calculate the average of the discrepancy $t_{\mu}$ up to the second
order:

$\langle t_{\mu}\rangle=\langle
t_{\mu}\rangle_{\mathrm{MF}}+\frac{1}{2!4}\langle
t_{\mu}\sum_{abde}\Delta[tt]_{ab}\Delta[tt]_{de}[\Delta\tilde{Q}_{f}(\bm{h}^{(L-1)})]_{ab}[\Delta\tilde{Q}_{f}(\bm{h}^{(L-1)})]_{de}\rangle_{\mathrm{MF},\mathrm{con}}\\\
+\text{Higher order terms}$ (90)

For CNN second order term can be further simplified by using Wick theorem for
pre-activations of the last layer when again we consider the case of linear
activation function

$\frac{\sigma_{L}^{4}}{8N^{2}C_{L-1}}\sum_{ji}\sum_{abde}\langle
t_{\mu}\left(t_{a}t_{b}-\langle
t_{a}t_{b}\rangle_{\mathrm{MF}}\right)\left(t_{d}t_{e}-\langle
t_{d}t_{e}\rangle_{\mathrm{MF}}\right)\left(K_{ajei}^{(L-1)}K_{bjdi}^{(L-1)}+K_{ajdi}^{(L-1)}K_{bjei}^{(L-1)}\right)\rangle_{\mathrm{MF},\mathrm{con}}$
(91)

A similar derivation to 5.6 yields the following correction to $\bar{t}_{\mu}$
($(f_{\mu}-y_{\mu})/\sigma^{2}$)

$\displaystyle\frac{2\sigma_{L}^{4}}{C_{L-1}N^{2}}\sum_{abdeij}[K_{f}^{-1}]_{\mu
a}K^{(L-1)}_{ai,bj}[K_{f}^{-1}]_{bd}K^{(L-1)}_{di,ej}\bar{t}_{e}$ (92)
$\displaystyle+\frac{\sigma_{L}^{4}}{C_{L-1}N^{2}}\sum_{ij}\left(\sum_{ab}[K_{f}^{-1}]_{\mu
a}K^{(L-1)}_{ai,bj}\bar{t}_{b}\right)\left(\sum_{ab}[K_{f}^{-1}]_{ab}K^{(L-1)}_{ia,jb}\right)$

Unlike in the FCN case, $K^{(L-1)}\in\mathbb{R}^{Nn\times Nn}$ is not
proportional to
$[K_{f}]_{\mu\nu}=\sigma^{2}\delta_{\mu\nu}+\sigma^{2}_{L}N^{-1}\sum_{j}K_{\mu
j,\nu j}^{(L-1)}$ where $\mu,\nu\in[1,n]$, in the $\sigma^{2}\rightarrow 0$
limit. In particular, for $N>1$, the two matrices have different dimensions.
Hence, we cannot cancel them together just yet. To obtain an order of
magnitude estimate of the second (and more dominant) term in this
perturbation, we turn to a different route and compare the mean-field value of
the norm, $\sum_{\mu}|\langle t_{\mu}\rangle|^{2}$ which is
$\sum_{\mu}\delta^{2}_{\mu}$ with its leading correction. Following similar
arguments to the FCN section, we find that the correction is dominated by

$\displaystyle\frac{\sigma_{L}^{4}}{C_{L-1}N^{2}}\sum_{ij}\left(\sum_{\mu
ab}\bar{t}_{\mu}[K_{f}^{-1}]_{\mu
a}K^{(L-1)}_{ai,bj}\bar{t}_{b}\right)\left(\sum_{ab}[K_{f}^{-1}]_{ab}K^{(L-1)}_{ia,jb}\right)$
(93)

To simplify this expression we next note that for data-sets in which for each
$\bm{x}_{\mu}$ there exists a "symmetry-partner" point wherein all coordinates
in the fan-in of the $i$’th latent pixels are flipped - the second, trace-
like, term must vanish for $i\neq j$. This is due to the fact that $K_{f}$ is
invariant under the action of the associated symmetry, whereas
$K^{(L-1)}_{*i,*j}$ receives a minor sign whenever $i\neq j$. As
$n\rightarrow\infty$ we expect this symmetry to be approximately realized as
it is a symmetry of the underlying measure from which $\bm{x}_{\mu}$ are
drawn. Following this, we remove $i\neq j$ terms from the summation.

Next we notice that due to the approximate translation symmetry of the
dataset, at large $n$

$\displaystyle\sum_{ab}[K_{f}^{-1}]_{ab}K^{(L-1)}_{ia,ib}$ (94)

becomes independent of $i$. We thus replace it by its average and perform the
remaining summation over, $i$ which now involves only the first term to obtain

$\displaystyle\frac{1}{C_{L-1}N}\left(\sum_{\mu
ab}\bar{t}_{\mu}[K_{f}^{-1}]_{\mu
a}[Q_{f}]_{ab}\bar{t}_{b}\right)\left(\sum_{ab}[K_{f}^{-1}]_{ab}[Q_{f}]_{ab}\right)$
(95)

Recall that $[Q_{f}]_{\mu\nu}=\sigma_{L}^{2}N^{-1}\sum_{j}K_{\mu j,\nu
j}^{(L-1)}$. This resulting expression is very similar to its FCN version
(with standard scaling) with one crucial difference, which is the appearance
of the aforementioned $1/N$ factor. More specifically, the first summation is
smaller than the zeroth term ($\sum_{\mu}t_{\mu}^{2}$). The scale controlling
the mean-field decoupling is therefore
$\frac{1}{C_{L-1}N}\Tr[K_{f}^{-1}Q_{f}]$. Thus, at large $N$, we can have a
reliable mean-field decoupling even when $n=C_{L-1}$.

##  7 Toy Example - One Hidden Layer

In this section, we provide a detailed calculation of the toy example
presented in the main text (subsection 2.3). We consider $\bm{x}_{\mu}$’s as a
center Gaussian vector of size $NS$ with random i.i.d. entries with zero mean
and variance $1$. We choose a target of the form
$y(\bm{x})=\sum_{i}a^{*}_{i}\bm{w}^{*}\cdot\bm{x}_{i}$. The activation
function here is $\phi=\mathrm{erf}$. Specifically, we consider a student
network of the form

$\displaystyle f(\bm{x}_{\mu})$
$\displaystyle=\sum_{i=1}^{N}\sum_{c=1}^{C}a_{ic}\erf\left(\bm{w}_{c}\cdot\bm{x}_{\mu,i}\right).$
(96)

requiring enough data-points to resolve the target sets $n>N+S$ (number of
target parameters) while staying within the over-parameterized regime implies
$n<CNS$. We further consider the large-scale “thermodynamic” limit, where
$C,S,N,n\gg 1$.

Similarly to the 3 layer case, the equations of states here are given by,

$\displaystyle[\Sigma^{-1}]_{ss^{\prime}}$
$\displaystyle=\frac{S}{\sigma_{\text{w}}^{2}}\delta_{ss^{\prime}}+\frac{1}{C}\sum_{\mu\nu}A_{\mu\nu}^{(2)}\frac{\partial[Q_{f}]_{\mu\nu}}{\partial\Sigma_{ss^{\prime}}}$
(97) $\displaystyle[Q_{f}]_{\mu\nu}$
$\displaystyle=\frac{2\sigma_{\text{a}}^{2}}{\pi
N}\sum_{i}\sin^{-1}\left(\frac{2\bm{x}_{\mu,i}\Sigma\bm{x}_{\mu,i}}{\sqrt{1+2\bm{x}_{\mu,i}\Sigma\bm{x}_{\mu,i}}\sqrt{1+2\bm{x}_{\nu,i}\Sigma\bm{x}_{\nu,i}}}\right)$
$\displaystyle{\bm{\delta}}$
$\displaystyle=[Q_{f}+\sigma^{2}I_{n}]^{-1}\bm{y}$ $\displaystyle
A_{\mu\nu}^{(2)}$
$\displaystyle=-{\delta}_{\mu}{\delta}_{\nu}+\left[(Q_{f}+\sigma^{2}I_{n})^{-1}\right]_{\mu\nu}$

These can be viewed as non-linear equations in the $S(S-1)/2$ variables making
up the symmetric matrix $\Sigma$. Having these variables determines the
$t_{\mu}$ variables directly.

We begin with approximating the GP inference appearing in the last equation.
This can be represented as ${\bm{\delta}}=[\bm{y}-\bar{\bm{f}}]/\sigma^{2}$
where, $\bar{\bm{f}}=Q_{f}[Q_{f}+\sigma^{2}I_{n}]^{-1}\bm{y}$ following the
standard GP prediction formula (for the training set). At large $n$,
$\bar{\bm{f}}$ can be approximated using the equivalence kernel (EK)
approximation [38] together with its perturbative corrections [14, 34] (non-
perturbative approaches in $1/n$ could also be considered [7, 43]). Within
this approximation scheme, one considers
$[Q_{f}]_{\mu\nu}=Q_{f}(\bm{x}_{\mu},\bm{x}_{\nu})$ as the continuum operator
$Q_{f}(\bm{x},\bm{y})$, diagonalized on the data-set measure ($d\mu$), leading
to the following formula for $\bar{f}(\bm{x})$ in the strict EK limit

$\displaystyle\bar{f}(\bm{x})$
$\displaystyle=\sum_{\lambda}\frac{\lambda}{\lambda+\sigma^{2}/n}y_{\lambda}\phi_{\lambda}(\bm{x}),$
(98)

where $\lambda,\phi_{\lambda}(\bm{x})$ are the eigenvalues and eigenfunctions
of $Q_{f}$ and $y_{\lambda}=\int d\mu_{x}y(\bm{x})\phi_{\lambda}(\bm{x})$. To
obtain an explicit formula for,
$\bar{t}(\bm{x})=[y(\bm{x})-\bar{f}(\bm{x})]/\sigma^{2}$, we proceed by
solving the eigenvalue problem. To this end, we first consider the kernel
action on a general linear function ($\bm{w^{\prime}}\cdot\bm{z}_{j}$), where
$\bm{z}_{j},\bm{w^{\prime}}$ are vectors of size $S$ for all $j$, and
$\bm{z}_{j}$ is drawn from the dataset measure,

$\displaystyle\int
d\mu_{z}Q_{f}(\bm{x},\bm{z})\bm{w}^{\prime}\cdot\bm{z}_{j}=$
$\displaystyle\int d\mu_{z}\frac{\sigma_{\text{a}}^{2}}{N}\sum_{i}\int
d\bm{w}p(\bm{w})\erf(\bm{w}\cdot\bm{x}_{i})\erf(\bm{w}\cdot\bm{z}_{i})(\bm{w}^{\prime}\cdot\bm{z}_{j}).$
(99)

where $p(\bm{w})$ is a centered Gaussian with covariance matrix $\Sigma$. In
the second transition, we undo the kernel integral, where $d\mu_{z}$ is a
Gaussian measure. We then exchange the order of integration and sum over the
weights and the data. We now do the integration over the data,

$\displaystyle\int
d\mu_{z}\erf(\bm{w}\cdot\bm{z}_{i})\bm{w}^{\prime}\cdot\bm{z}_{j}$
$\displaystyle=\delta_{ij}\sqrt{\frac{4}{\pi}}\frac{\bm{w}^{\prime}\cdot\bm{w}}{\sqrt{1+2\norm{\bm{w}}^{2}}}$