A self consistent theory of Gaussian Processes captures feature learning effects in finite CNNs

I Introduction

The correspondence between Gaussian Processes (GPs) and deep neural networks (DNNs) has been instrumental in advancing our understanding of these complex algorithms. Early results related randomly initialized strongly over-parameterized DNNs with GP priors . More recent results considered training using gradient flow (or noisy gradients), where DNNs map to Bayesian inference on GPs governed by the neural tangent kernel (or the NNGP kernel ). These correspondences carry over to a wide variety of architectures, going beyond fully connected networks (FCNs) to convolutional neural networks (CNNs) , recurrent neural networks (RNNs) and even attention networks . They provide us with closed analytical expressions for the outputs of strongly over-parameterized trained DNNs, which have been used to make accurate predictions for DNN learning curves .

Despite their theoretical appeal, GPs are unable to capture feature learning , which is a well-observed key property of trained DNNs. Indeed, it was noticed that as the width tends to infinity, the neural tangent kernel (NTK) tends to a constant kernel that does not evolve during training and the weights in hidden layers change infinitesimally from their initialization values. This regime of training was thus dubbed lazy training . Other studies showed that for CNNs trained on image classification tasks, the feature learning regime generally tends to outperform the lazy regime . Clearly, working in the feature learning regime is also crucial for performing transfer learning .

It is therefore desirable to have a theoretical approach to deep learning which enjoys the generality and analytical power of GPs while capturing feature learning effects in finite DNNs. Here we make several contributions towards this goal:

We show that the mean predictor of a finite DNN trained on a large data set with noisy gradients, weight decay and MSE loss, can be obtained from GP regression on a shifted target (§III). Central to our approach is a non-linear self consistent equation involving the higher cumulants of the finite DNN (at initialization) which predicts this target shift.

Using this machinery on a toy model of a two-layer linear CNN in a teacher-student setting, we derive explicit analytical predictions which are in very good agreement with experiments even well away from the GP/lazy-learning regime (large number of channels, $C$) thus accounting for strong finite-DNN corrections (§IV). Similarly strong corrections to GPs, yielding qualitative improvements in performance, are demonstrated for the quadratic two-layer fully connected model of Ref. .

We show how our framework can be used to study statistical properties of weights in hidden layers. In particular, in the CNN toy model, we identify, both analytically and numerically, a sharp transition between a feature learning phase and a lazy learning phase (§IV.4). We define the feature learning phase as the regime where the features of the teacher network leave a clear signature in the spectrum of the student’s hidden weights posterior covariance matrix. In essence, this phase transition is analogous to the transition associated with the recovery of a low-rank signal matrix from a noisy matrix taken from the Wishart ensemble, when varying the strength of the low-rank component .

Several previous papers derived leading order finite-DNN corrections to the GP results . While these results are in principle extendable to any order in perturbation theory, such high order expansions have not been studied much, perhaps due to their complexity. In contrast, we develop an analytically tractable non-perturbative approach which we find crucial for obtaining non-negligible feature learning and associated performance enhancement effects.

Previous works studied how the behavior of infinite DNNs depends on the scaling of the top layer weights with its width. In it is shown that the standard and NTK parameterizations of a neural network do not admit an infinite-width limit that can learn features, and instead suggest an alternative parameterization which can learn features in this limit. While unifying various viewpoints on infinite DNNs, this approach does not immediately lend itself to analytical analysis of the kind proposed here.

Several works show that finite width models can generalize either better or worse than their infinite width counterparts, and provide examples where the relative performance depends on the optimization details, the DNN architecture and the statistics of the data. Here we demonstrate analytically that finite DNNs outperform their GP counterparts when the latter have a prior that lacks some constraint found in the data (e.g. positive-definiteness , weight sharing ).

Deep linear networks (FCNs and CNNs) similar to our CNN toy example have been studied in the literature . These studies use different approaches and assumptions and do not discuss the target shift mechanism which applies also for non-linear CNNs. In addition, their analytical results hinge strongly on linearity whereas our approach could be useful whenever several leading cumulants of the DNN output are known.

A concurrent work derived exact expressions for the output priors of finite FCNs induced by Gaussian priors over their weights. However, these results only apply to the limited case of a prior over a single training point and only for a FCN. In contrast, our approach applies to the setting of a large training set, it is not restricted to FCNs and yields results for the posterior predictions, not the prior. Focusing on deep linear fully connected DNNs, recent work derived analytical finite-width renormalization results for the GP kernel, by sequentially integrating out the weights of the DNN, starting from the output layer and working backwards towards the input. Our analytical approach, its scope, and the models studied here differ substantially from that work.

II Preliminaries

where $\theta_{t}$ is the vector of all network parameters at time step $t$, $\gamma$ is the strength of the weight decay, $\mathcal{L}(f_{\theta})$ is the loss as a function of the DNN output $f_{\theta}$ (where we have emphasized the dependence on the parameters $\theta$), $\sigma$ is the magnitude of noise, $\eta$ is the learning rate and $\xi_{t}\sim\mathcal{N}(0,I)$. As $\eta\to 0$ this discrete-time dynamics converge to the continuous-time Langevin equation given by $\dot{\theta}\left(t\right)=-\nabla_{\theta}\left(\frac{\gamma}{2}||\theta(t)||^{2}+\mathcal{L}\left(f_{\theta}\right)\right)+2\sigma\xi\left(t\right)$ with $\left\langle\xi_{i}(t)\xi_{j}(t^{\prime})\right\rangle=\delta_{ij}\delta\left(t-t^{\prime}\right)$, so that as $t\to\infty$ the DNN parameters $\theta$ will be sampled from the equilibrium Gibbs distribution $P(\theta)$.

As shown in , the parameter distribution $P(\theta)$ induces a posterior distribution over the trained DNN outputs $P(\vec{f})$ with the following partition function:

Here $P_{0}(\vec{f})$ is the prior generated by the finite-DNN with $\theta$ drawn from ${\mathcal{N}}(0,2\sigma^{2}/\gamma)$ where the weight decay $\gamma$ may be layer-dependent, $\left\{g_{\mu}\right\}_{\mu=1}^{n}$ are the training targets and $\vec{J}$ are source terms used to calculate the statistics of $f$. We keep the loss function $\mathcal{L}$ arbitrary at this point, committing to a specific choice in the next section. As standard , to calculate the posterior mean at any of the training points or the test point ${\mathbf{x}}_{n+1}$ from this partition function one uses

III A self consistent theory for the posterior mean and covariance

In this section we show that for a large training set, the posterior mean predictor (Eq. 3) amounts to GP regression on a shifted target ($g_{\mu}\to g_{\mu}-\Delta g_{\mu}$). This shift to the target ($\Delta g_{\mu}$) is determined by solving certain self-consistent equations involving the cumulants of the prior $P_{0}(\vec{f})$. For concreteness, we focus here on the MSE loss $\mathcal{L}=\sum_{\mu=1}^{n}\left(f_{\mu}-g_{\mu}\right)^{2}$ and comment on extensions to other losses, e.g. the cross entropy, in App. C. To this end, consider first the prior of the output of a finite DNN. Using standard manipulations (see App. A), it can be expressed as follows

where $\kappa_{\mu_{1},\dots,\mu_{r}}$ is the $r$’th multivariate cumulant of $P_{0}(\vec{f})$ . The second term in the exponent is the cumulant generating function ($\mathcal{C}$) corresponding to $P_{0}$. As discussed in App. B and Ref. , for standard initialization protocols the $r$’th cumulant will scale as $1/C^{(r/2-1)}$, where $C$ controls the over-parameterization. The second ($r=2$) cumulant which is $C$-independent, describes the NNGP kernel of the finite DNN and is denoted by $K({\mathbf{x}}_{\mu_{1}},{\mathbf{x}}_{\mu_{2}})=\kappa_{\mu_{1},\mu_{2}}$.

For a DNN with finite $C$, the prior $P_{0}(\vec{f})$ will no longer be Gaussian and cumulants with $r>2$ would contribute. This renders the partition function in Eq. 2 intractable and so some approximation is needed to make progress. To this end we note that $f$ can be integrated out (see App. A.1) to yield a partition function of the form

where $\mathcal{S}(\vec{t},\vec{J})$ is the action whose exact form is given in Eq. A.14. Interestingly, the $it_{\mu}$ variables appearing above are closely related to the discrepancies $\hat{\delta}g_{\mu}$, in particular $\langle it_{\mu}\rangle=\langle\hat{\delta}g_{\mu}\rangle/\sigma^{2}$.

To proceed analytically we adopt the saddle point (SP) approximation which, as argued in App. A.4, relies on the fact that the non-linear terms in the action comprise of a sum of many $it_{\mu}$’s. Given that this sum is dominated by collective effects coming from all data points, expanding $\mathcal{S}(\vec{t},\vec{J})$ around the saddle point yields terms with increasingly negative powers of $n$.

For the training points $\mu\in\left\{1,\dots,n\right\}$, taking the saddle point approximation amounts to setting $\evaluated{\partial_{it_{\mu}}\mathcal{S}\left(\vec{t},\vec{J}\right)}_{\vec{J}=\vec{0}}=0$. This yields a set of equations that has precisely the form of Eq. 5, but where the target is shifted as $g_{\nu}\to g_{\nu}-\Delta g_{\nu}$ and the target shift is determined self consistently by

Equation 7 is thus an implicit equation for $\Delta g_{\nu}$ involving all training points, and it holds for the training set and the test point $\nu\in\left\{1,\dots,n+1\right\}$. Once solved, either analytically or numerically, one calculates the predictions on the test point via

Equation 5 with $g_{\nu}\to g_{\nu}-\Delta g_{\nu}$ along with Eqs. 7 and 8 are the first main result of this paper. Viewed as an algorithm, the procedure to predict the finite DNN’s output on a test point ${\mathbf{x}}_{*}$ is as follows: we shift the target in Eq. 5 as $g\to g-\Delta g$ with $\Delta g$ as in Eq. 7, arriving at a closed equation for the average discrepancies $\langle\hat{\delta}g_{\mu}\rangle$ on the training set. For some models, the cumulants $\kappa_{\nu,\mu_{2},\dots,\mu_{r}}$ can be computed for any order $r$ and it can be possible to sum the entire series, while for other models several leading cumulants might already give a reasonable approximation due to their $1/C^{r/2-1}$ scaling. The resulting coupled non-linear equations can then be solved numerically, to obtain $\Delta g_{\mu}$ from which predictions on the test point are calculated using Eq. 8.

Notwithstanding, solving such equations analytically is challenging and one of our main goals here is to provide concrete analytical insights. Thus, in §IV.2 we propose an additional approximation wherein to leading order we replace all summations over data-points with integrals over the measure from which the data-set is drawn. This approximation, taken in some cases beyond leading order as in Ref. , will yield analytically tractable equations which we solve for two simple toy models, one of a linear CNN and the other of a non-linear FCN.

The SP approximation can be extended to compute the predictor variance by expanding the action $\mathcal{S}$ to quadratic order in $it_{\mu}$ around the SP value (see App. A.3). Due to the saddle-point being an extremum this leads to $\mathcal{S}\approx\mathcal{S}_{\rm{SP}}+\frac{1}{2}t_{\mu}A^{-1}_{\mu\nu}t_{\nu}$. This leaves the previous SP approximation for the posterior mean on the training set unaffected (since the mean and maximizer of a Gaussian coincide), but is necessary to get sensible results for the posterior covariance. Empirically, in the toy models we considered in §IV we find that the finite DNN corrections to the variance are much less pronounced than those for the mean. Using the standard Gaussian integration formula, one finds that $A_{\mu\nu}$ is the covariance matrix of $it_{\mu}$. Performing such an expansion one finds

where the $it_{\mu}$ on the r.h.s. are those of the saddle point. This gives an expression for the posterior covariance matrix on the training set:

where the r.h.s. coincides with the posterior covariance of a GP with a kernel equal to $K+\Delta K$ . The variance on the test point is given by (repeating indices are summed over the training set)

IV Two toy models

Here we define a teacher-student toy model showing several qualitative real-world aspects of feature learning and analyze it via our self-consistent shifted target approach. Concretely, we consider the simplest student CNN $f({\mathbf{x}})$, having one hidden layer with linear activation, and a corresponding teacher CNN, $g({\mathbf{x}})$

Despite its simplicity, this model distils several key differences between feature learning models and lazy learning or GP models. Due to the lack of pooling layers, the GP associated with the student fails to take advantage of the weight sharing property of the underlying CNN . In fact, here it coincides with a GP of a fully-connected DNN which is quite inappropriate for the task. We thus expect that the finite network will have good performance already for $n=C^{*}(N+S)$ whereas the GP will need $n$ of order of the dimension ($NS$) to learn well . Thus, for $N+S\ll NS$ there should be a broad regime in the value of $n$ where the finite network substantially outperforms the corresponding GP. We later show (§IV.4) that this performance boost over GP is due to feature learning, as one may expect.

Conveniently, the cumulants of the student DNN of any order can be worked out exactly. Assuming $\gamma$ and $\sigma^{2}$ of the noisy GD training are chosen such thatGenerically this requires $C$ dependent and layer dependent weight decay. $a_{i,c}\sim\mathcal{N}\left(0,\sigma_{a}^{2}/CN\right),\quad{\mathbf{w}}_{c}\sim\mathcal{N}\left(\bm{0},\frac{\sigma_{w}^{2}}{S}I_{S}\right)$ (and similarly for the teacher DNN) the covariance function for the associated GP reads

Denoting $\lambda:=\frac{\sigma_{a}^{2}}{N}\frac{\sigma_{w}^{2}}{S}$, the even cumulant of arbitrary order $2m$ is (see App. F):

IV.2 Self consistent equation in the limit of a large training set

In §III our description of the self consistent equations was for a finite and fixed training set. Further analytical insight can be gained if we consider the limit of a large training set, known in the GP literature as the Equivalent Kernel (EK) limit . For a short review of this topic, see App. D. In essence, in the EK limit we replace the discrete sums over a specific draw of training set, as in Eqs. 5, 7, 8, with integrals over the entire input distribution $\mu({\mathbf{x}})$. Given a kernel that admits a spectral decomposition in terms of its eigenvalues and eigenfunctions: $K\left({\mathbf{x}},{\mathbf{x}}^{\prime}\right)=\sum_{s}\lambda_{s}\psi_{s}\left({\mathbf{x}}\right)\psi_{s}\left({\mathbf{x}}^{\prime}\right)$, the standard result for the GP posterior mean at a test point is approximated by

In the context of our theory, the EK limit allows for a derivation of a simple analytical form for the self consistent equations. As shown in App. E.1 in our toy CNN both $\Delta g$ and $\hat{\delta}g$ become linear in the target. Thus the self-consistent equations can be reduced to a single equation governing the proportionality factor ($\alpha$) between $\hat{\delta}g$ and $g$ ($\hat{\delta}g=\alpha g$). Thus starting from the general self consistent equations, 5, 7, 8, taking their EK limit, and plugging in the general cumulant for our toy model (14) we arrive at the following equation for $\alpha$

IV.3 Numerical verification

In this section we numerically verify the predictions of the self consistent theory of Sec. §IV.2, by training linear shallow student CNNs on a teacher with $C^{*}=1$ as in Eq. 12, using noisy gradients as in Eq. 1, and averaging their outputs across noise realizations and across dynamics after reaching equilibrium.

For simplicity we used $N=S$ and $n\in\left\{62,200,650\right\},S\in\left\{15,30,60\right\}$ so that $n\propto S^{1.7}$. The latter scaling places us in the poorly performing regime of the associated GP while allowing good performance of the CNN. Indeed, as aforementioned, the GP here requires $n$ of the scale of $\lambda^{-1}=NS=O(S^{2})$ for good performance , while the CNN requires $n$ of scale of the number of parameters ($C(N+S)=O(S)$).

The results are shown in Fig 1 where we compare the theoretical predictions given by the solutions of the self consistent equation (16) to the empirical values of $\alpha$ obtained by training actual CNNs and averaging their outputs across the ensemble.

IV.4 Feature learning phase transition in the CNN model

At this point there is evidence that our self-consistent shifted target approach works well within the feature learning regime of the toy model. Indeed GP is sub-optimal here, since it does not represent the CNN’s weight sharing present in the teacher network. Weight sharing is intimately tied with feature learning in the first layer, since it aggregates the information coming from all convolutional windows to refine a single set of repeating convolution-filters. Empirically, we observed a large performance gap of finite $C$ CNNs over the infinite-$C$ (GP) limit, which was also observed previously in more realistic settings . Taken together with the existence of a clear feature in the teacher, a natural explanation for this performance gap is that feature learning, which is completely absent in GPs, plays a major role in the behavior of finite $C$ CNNs.

As shown in App. G, using our field-theory or function-space formulation, we find that to leading order in $1/C$ the ensemble average of the empirical covariance matrix, for a teacher with a single feature ${\mathbf{w}}^{*}$, is

A first conclusion that could be drawn here, is that given access to an ensemble of such trained CNNs, feature learning happens for any finite $C$ as a statistical property. We turn to discuss the more common setting where one wishes to use the features learned by a specific randomly chosen CNN from this ensemble.

To this end, we follow Ref. and model $\Sigma_{W}$ as a Wishart matrix with a rank-one perturbation. The variance of the matrix and details of the rank one perturbation are then determined by the above equation. Consequently the eigenvalue distribution is expected to follow a spiked Marchenko-Pastur (MP), which was studied extensively in . To test this modeling assumption, for each snapshot of training time (after reaching equilibrium) and noise realization we compute $\Sigma_{W}$’s eigenvalues and aggregate these across the ensemble. In Fig. 2 we plot the resulting empirical spectral distribution for varying values of $C$ while keeping $S$ fixed. Note that, differently from the usual spiked-MP model, varying $C$ here changes both the distribution of the MP bulk (which is determined by the ratio $S/C$) as well as the strength of the low-rank perturbation.

Our main finding is a phase transition between two regimes which becomes sharp as one takes $n,S\rightarrow\infty$. In the regime of large $C$ the eigenvalue distribution of $\Sigma_{W}$ is indistinguishable from the MP distribution, whereas in the regime of small $C$ an outlier eigenvalue $\lambda_{m}$ departs from the support of the bulk MP distribution and the associated top eigenvector has a non-zero overlap with ${\mathbf{w}}^{*}$, see Fig. 2. We refer to the latter as the feature-learning regime, since the feature ${\mathbf{w}}^{*}$ is manifested in the spectrum of the students weights, whereas the former is the non-feature learning regime. We use the quantity $\mathcal{Q}\equiv{\mathbf{w}}^{*\mathsf{T}}\Sigma_{W}{\mathbf{w}}^{*}$ as a surrogate for $\lambda_{m}$, as it is valid on both sides of the transition. Having established the correspondence to the MP plus low rank model, we can use the results of to find the exact location of the phase transition, which occurs at the critical value $C_{\rm{crit}}$ given by

where we assumed for simplicity $N=S$ so that $\lambda=S^{-2}$.

IV.5 Two-layer FCN with average pooling and quadratic activations

Similarly to the previous toy model, here too the GP associated with the student at large $M$ (and finite $\sigma^{2}$) overlooks a qualitative feature of the finite DNN — the fact that the first term in $f({\mathbf{x}})$ is non-negative. Interestingly, this feature provides a strong performance boost in the $\sigma^{2}\rightarrow 0$ limit compared to the associated GP. Namely the DNN, even at large $M$, performs well for $n>2d$ whereas the associated GP is expected to work well only for $n=O(d^{2})$ .

We wish to solve for the predictions of this model with our self consistent GP based approach. As shown in App. I, the cumulants of this model can be obtained from the following cumulant generating function

The associated GP kernel is given by $K({\mathbf{x}}_{\mu},{\mathbf{x}}_{\nu})=2M^{-1}\sigma_{w}^{4}({\mathbf{x}}_{\mu}\cdot{\mathbf{x}}_{\nu})^{2}$. Following this, the target shift equation, at the saddle point level, appears as

In App. I, we solve these equations numerically for $\sigma^{2}=10^{-5}$ and show that our approach captures the correct $n=2d$ threshold value. An analytic solution of these equations at low $\sigma^{2}$ using EK or other continuum approximations is left for future work (see Refs. for potential approaches). As a first step towards this goal, in App. I we consider the simpler case of $\sigma^{2}=1$ and derive the asymptotics of the learning curves which deviate strongly from those of GP for $M\ll d$.

V Discussion

In this work we presented a correspondence between ensembles of finite DNNs trained with noisy gradients and GPs trained on a shifted target. The shift in the target can be found by solving a set of self consistent equations for which we give a general form. We found explicit expressions for these equations for the case of a 2-layer linear CNN and a non-linear FCN, and solved them analytically and numerically. For the former model, we performed numerical experiments on CNNs that agree well with our theory both in the GP regime and well away from it, i.e. for small number of channels $C$, thus accounting for strong finite $C$ effects. For the latter model, the numerical solution of these equations capture a remarkable and subtle effect in these DNNs which the GP approach completely overlooks — the $n=2d$ threshold value.

Considering feature learning in the CNN model, we found that averaging over ensembles of such networks always leads to a form of feature learning. Namely, the teacher always leaves a signature on the statistics of the student’s weights. However, feature learning is usually considered at the level of a single DNN instance rather than an ensemble of DNNs. Focusing on this case, we show numerically that the eigenvalues of $\Sigma_{W}$, the student hidden weights covariance matrix, follow a Marchenko–Pastur distribution plus a rank-1 perturbation. We then use our approach to derive the critical number of channels $C_{\rm{crit}}$ below which the student is in a feature learning regime.

There are many directions for future research. Our toy models where chosen to be as simple as possible in order to demonstrate the essence of our theory on problems where lazy learning grossly under-performs finite-DNNs. Even within this setting, various extensions are interesting to consider such as adding more features to the teacher CNN (e.g. biases or a subset of linear functions which are more favorable), studying linear CNNs with overlapping convolutional windows, or deeper linear CNNs. As for non-linear CNNs, we believe it is possible to find the exact cumulants of any order for a variety of toy CNNs involving, for example, quadratic activation functions. For other cases it may be useful to develop methods for characterizing and approximating the cumulants.

More generally, we advocated here a physics-style methodology using approximations, self-consistency checks, and experimental tests. As DNNs are very complex experimental systems, we believe this mode of research is both appropriate and necessary. Nonetheless we hope the insights gained by our approach would help generate a richer and more relevant set of toy models on which mathematical proofs could be made.

References

Appendix A Derivation of the target shift equations

where unless explicitly written otherwise, summations over $\mu$ run from $1$ to $n+1$ (i.e. include the test point). Here we commit to the MSE loss which facilitates the derivation, and in App. C we give an alternative derivation that may also be applied to other losses such as cross-entropy. Our goal in this appendix is to establish that the target shift equations are in fact saddle point equations of the partition function A.2 following some transformations on the variables of integration. To this end, consider the cumulant generating function of $P_{0}\left(\vec{f}\right)$ given by

where the sum over the cumulant tensors $\kappa_{\mu_{1},\dots,\mu_{r}}$ does not include $r=1$ since our DNN priors are assumed to have zero mean. Notably one can re-express $P_{0}\left(\vec{f}\right)$ as the inverse Fourier transform of $e^{\mathcal{C}(\vec{t})}$:

where for clarity we do not keep track of multiplicative $\pi$ factors that have no effect on moments of $f_{\mu}$. As the term in the exponent (the action) is quadratic in $\left\{f_{\mu}\right\}_{\mu=1}^{n}$ and linear in $f_{n+1}$ these can be integrated out to yield an equivalent partition function phrased solely in terms of $t_{1},...,t_{n+1}$:

The identification $t_{n+1}=-iJ_{n+1}$ arises from the delta function:

Recall that $\left\langle f_{\mu}\right\rangle=\partial_{J_{\mu}}\evaluated{\log Z\left(\vec{J}\right)}_{\vec{J}=\vec{0}}$ and notice that the first term in Eq. A.8 (the cumulant generating function, $\mathcal{C}$) depends on $J_{n+1}$ and not on $\left\{J_{\mu}\right\}_{\mu=1}^{n}$ whereas the rest of the action depends on $\left\{J_{\mu}\right\}_{\mu=1}^{n}$ and not on $J_{n+1}$. Thus, for training points $\left\langle f_{\mu}\right\rangle$ amounts to the average of $g_{\mu}-i\sigma^{2}t_{\mu}$, and so we identify

where $\langle\cdots\rangle$ denotes an expectation value using $Z(\vec{J}=\vec{0})$. We comment that the above relation holds also for any (non-mixed) cumulants of $it_{\mu}$ and $\hat{\delta}g_{\mu}/\sigma^{2}$ except the covariance, where a constant difference appears due to the $O(J^{2})$ term in the action, namely

To make contact with GPs it is beneficial to expand $\mathcal{C}(t_{1},\dots,t_{n},-iJ_{n+1})$ in terms of its cumulants, and split the second cumulant, describing the DNNs’ NNGP kernel, from the rest. Namely, using Einstein summation

Writing Eq. A.8 in this fashion gives the action:

A.2 Saddle point equation for the mean predictor

Having arrived at the action A.14, we can readily derive the saddle point equations for the training points by setting:

This corresponds to treating the variables $\left\{it_{\mu}\right\}_{\mu=1}^{n}$ as non-fluctuating quantities, i.e. replacing them with their mean value: $it_{\mu}\to\left\langle it_{\mu}\right\rangle$. Performing this for the training set $\nu\in\left\{1,\dots,n\right\}$ yields

where this target shift is related to $\mathcal{C}$ of Eq. A.13 by

Finally, we get the expression for the mean predictor at the test point by setting $\left\langle f_{*}\right\rangle=\evaluated{\partial_{J_{*}}\log Z\left(\vec{J}\right)}_{\vec{J}=\vec{0}}$ and plugging in the SP values for $it_{\mu}$ on the training set from Eq. A.16 and the target shift $\Delta g$ from Eq. A.17. This gives

A.3 Posterior covariance

The posterior covariance on the test point is important for determining the average MSE loss on the test-set, as the latter involves the MSE of the mean-predictor plus the posterior covariance. Concretely, we wish to calculate $\partial^{2}_{J_{*}}\log(Z)$, express it as an expectation value w.r.t $Z$ and calculate this using $Z$ with the self-consistent target shift. To this end, note that generally if $Z(J)=\int e^{-\mathcal{S}(J)}$ then

where an Einstein summation over $\nu=1,\dots,n$ is implicit. Further recalling that on the training set we have

where here $\Delta g$ is the full quantity without any SP approximations

where we can unpack the last two terms as

which is the standard posterior covariance of a GP .

The expressions in Eqs. A.25, A.26 are exact, but to evaluate them in a more compact form we approximate the $t_{\mu}$ distribution as a Gaussian centered around the SP value.

A.3.2 Posterior covariance on the training set

Our target shift approach at the saddle-point level allows a computation of the fluctuations of $it_{\mu}$ using the standard procedure of expanding the action at the saddle point to quadratic order in $t_{\mu}$. Due to the saddle-point being an extremum this leads to $\mathcal{S}\approx\mathcal{S}_{\rm{saddle}}+\frac{1}{2}t_{\mu}A^{-1}_{\mu\nu}t_{\nu}$ and thus using the standard Gaussian integration formula, one finds that $A_{\mu\nu}$ is the covariance matrix of $it_{\mu}$. Performing such an expansion on the action of Eq. A.14 one finds

where the $it_{\mu}$ on the r.h.s. are those of the saddle-point. Recalling Eq. A.11 we have

and the r.h.s. coincides with the posterior covariance of a GP with a kernel equal to $K+\Delta K$.

A.4 A criterion for the saddle-point regime

Saddle point approximations are commonly used in statistics and physics and often rely on having partition functions of the form $Z=\int dte^{-n\mathcal{S}(t)}$ where $n$ is a large number and $\mathcal{S}$ is order $1$ ($O(1)$). In our settings we cannot simply extract such a large factor from the action and make it $O(1)$. Nonetheless, we argue that expanding the action to quadratic order around the saddle point is still a good approximation at large $n$, with $n$ being the training set size. Concretely we give the following two consistency criteria based on comparing the saddle point results with their leading order beyond-saddle-point corrections. The first is given by the latter correction to the mean predictor over the scale of the saddle point prediction

where an Einstein summation over the training-set is implicit and the derivatives are evaluated at the saddle point value. This criterion can be calculated for any specific model to verify the appropriateness of the saddle point approach. We further provide a simpler criterion

which however relies on heuristic assumptions. The main purpose of this heuristic criterion is to provide a qualitative explanation for why we expect the first criterion to be small in many interesting large $n$ settings.

To this end we first obtain the leading (beyond quadratic) correction to the mean. Consider the partition function in terms of $it$ and its expansion around the saddle point. As $P_{0}[f]$ is effectively bounded (by the Gaussian tails of the finite set of weights), the corresponding characteristic function ($e^{\mathcal{C}(t_{1},\dots,t_{n})}$) is well defined over the entire complex plane. Given this, one can deform the integration contour, along each dimension ($\int_{-\infty}^{\infty}dt_{\mu}$), which originally laid on the real axis, to $\int_{-\infty+t_{\rm{SP},\mu}}^{\infty+t_{\rm{SP},\mu}}dt_{\mu}$ where $t_{\rm{SP}}$ is purely imaginary and equals $-i\hat{\delta}g/\sigma^{2}$ so it crosses the saddle point (see also Ref. ). Next we expand the action in the deviation from the SP value: $\delta t_{\mu}=t_{\mu}-t_{\rm{SP},\mu}$ to obtain

where an Einstein summation over the training set is implicit and where we denoted for $m\geq 3$

Next we consider first order perturbation theory in the cubic term and calculate the correction to the mean of $it_{\mu}$ or equivalently $\hat{\delta}g_{\mu}/\sigma^{2}$.

Next we turn to study the scaling of this correction with $n$. To this end we first consider a single derivative of $\Delta g$ ($\partial_{it_{\nu}}\Delta g_{\mu}$). Note that $\Delta g_{\mu}$, by its definition, includes contributions from at least $n^{3}$ different $it$’s. In many cases, one expects that the value of this sum will be dominated by some finite fraction of the training set rather than by a vanishing fraction. This assumption is in fact implicit in our EK treatment where we replaced all $\sum_{\mu}$ with $n\int$. Given so, the derivative $\partial_{it_{\nu}}\Delta g_{\mu}$, which can be viewed as the sensitivity to changing $it_{\nu}$, is expected to go as one over the size of that fraction of the training set, namely as $1/n$. Under this collectivity assumption we expect the scaling

Making a similar collectivity assumption on higher derivatives yields

Generally we expect $\Delta g\approx 0$ at strong over-parameterization (as non-linear effects are suppressed by $C^{-1}$) and $\Delta g\sim O(g)$ at good performances (as this implies good performance on the training set). Thus we generally expect $\Delta g=O(g)=O(1)$ and hence large $n(\hat{\delta}g/\sigma^{2})^{2}$ controls the magnitude of the corrections. Considering the $\sigma^{2}\rightarrow 0$ limit, we note in passing that $\hat{\delta}g/\sigma^{2}$ typically remains finite. For instance it is simply $K^{-1}g$ for a Gaussian Process.

Considering the linear CNN model of the main text, we estimate the above heuristic criterion for $n=650$ and $C=8$ where $\Delta g=O(g)$ and $\hat{\delta}g\approx 0.1g$. This then gives $(6.5O(g)^{2})^{-1}$ as the small factor dominating the correction. As we choose $O(g)=3$ in that experiment, we find that the correction is roughly $1/60$. As the discrepancy is $0.1g=O(0.3)$ we expect roughly a $5\%$ relative error in predicting the discrepancy.

Appendix B Review of the Edgeworth expansion

In this section we give a review of the Edgeworth expansion, starting from the simplest case of a scalar valued RV and then moving on vector valued RVs so we can write down the expansion for the output of a generic neural network on a fixed set of inputs.

Consider scalar valued continuous iid RVs $\{Z_{i}\}$ and assume WLOG $\left\langle Z_{i}\right\rangle=0,\hskip 5.0pt\left\langle Z_{i}^{2}\right\rangle=1$, with higher cumulants $\kappa_{r}^{Z}$ for $r\geq 3$. Now consider their normalized sum $Y_{N}=\frac{1}{\sqrt{N}}\sum_{i=1}^{N}Z_{i}$. Recall that cumulants are additive, i.e. if $Z_{1},Z_{2}$ are independent RVs then $\kappa_{r}(Z_{1}+Z_{2})=\kappa_{r}(Z_{1})+\kappa_{r}(Z_{2})$ and that the $r$-th cumulant is homogeneous of degree $r$, i.e. if $c$ is any constant, then $\kappa_{r}(cZ)=c^{r}\kappa_{r}(Z)$. Combining additivity and homogeneity of cumulants we have a relation between the cumulants of $Z$ and $Y$

Now, let $\varphi(y):=(2\pi)^{-1/2}e^{-y^{2}/2}$ be the PDF of the standard normal distribution. The characteristic function of $Y$ is given by the Fourier transform of its PDF $P(y)$ and is expressed via its cumulants

where the last equality holds since $\kappa_{1}=0,\quad\kappa_{2}=1$ and $\hat{\varphi}(t)=e^{-\frac{t^{2}}{2}}$. From the CLT, we know that $P(y)\to\varphi(y)$ as $N\to\infty$. Taking the inverse Fourier transform $\mathcal{F}^{-1}$ has the effect of mapping $it\mapsto-\partial_{y}$ thus

where $H_{r}(y)$ is the $r$th probabilist’s Hermite polynomial, defined by

B.2 Edgeworth expansion for a vector valued random variable

where $\kappa_{i,j}$ is the matrix inverse of $\kappa^{i,j}$ and Einstein summation is used. The $r$’th order cumulant becomes a tensor with $r$ indices, e.g. the analogue of $\kappa_{4}$ is $\kappa^{i,j,k,l}$. The Hermite polynomials are now multi-variate polynomials, so that the first one is $H_{i}=\kappa_{i,j}y^{j}$ and the fourth one is

where the postscript bracket notation is simply a convenience to avoid listing explicitly all possible partitions of the indices, e.g. $\kappa_{i,j}\kappa_{k,l}=\kappa_{i,j}\kappa_{k,l}+\kappa_{i,k}\kappa_{j,l}+\kappa_{i,l}\kappa_{j,k}$

In our context we are interested in even distributions where all odd cumulants vanish, so the Edgeworth expansion reads

B.3 Edgeworth expansion for the posterior of Bayesian neural network

Consider an on-data formulation, i.e. a distribution over a vector space - the NN output evaluated on the training set and on a single test point, rather than a distribution over the whole function space:

where ${\mathbf{x}}_{*}$ is the test point. Let $\kappa_{r}$ denote the $r$th cumulant of the prior $P_{0}\left(\vec{f}\right)$ of the network over this space:

Take the baseline distribution to be Gaussian $P_{G}\left(\vec{f}\right)\propto\exp\left(-\frac{1}{2}\vec{f}^{\mathsf{T}}K^{-1}\vec{f}\right)$, around which we perform the Edgeworth expansion, thus the characteristic function of the prior reads

and the indices range over both the train set and the test point $\mu\in\left\{1,\dots,\underbrace{n+1}_{*}\right\}$. In our case, all odd cumulants vanish, thus

Introducing the data term and a source term, the partition function reads (denote $f\left({\mathbf{x}}_{\mu}\right)\equiv f_{\mu},\quad f\left({\mathbf{x}}_{*}\right)\equiv f_{*})$

Appendix C Target shift equations - alternative derivation

Here we derive our self-consistent target shift equations from a different approach which does not require the introduction of the $it_{\mu}$ integration variables by transforming to Fourier space. While this approach requires an additional assumption (see below) it also has the benefit of being extendable to any smooth loss function comprised of a sum over training points. In particular, below we derive it for both MSE loss and cross entropy loss.

To this end, we examine the Edgeworth expansion for the partition function given by Eq. B.14. By using a series of integration by parts and noting the boundary terms vanish, one can shift the action of the higher cumulants from the prior to the data dependent term

Doing so yields an equivalent viewpoint on the problem, wherein the Gaussian data term and the non-Gaussian prior appearing in Eq. B.14 are replaced in Eq. C.1 by a Gaussian prior and a non-Gaussian data term.

Next we argue that in the large $n$ limit, the non-Gaussian data-term can be expressed as a Gaussian-data term but on a shifted target. To this end we note that when $n$ is large, most combinations of derivatives in the exponents act on different data points. In such cases derivatives could simply be replaced as $\partial_{\mu_{i}}\rightarrow\sigma^{-2}\hat{\delta}g_{\mu_{i}}$, where $\hat{\delta}g_{\mu_{i}}\equiv g_{\mu_{i}}-f_{\mu_{i}}$ denotes the discrepancy on the training point $\mu_{i}$.

Consider next how $f_{\nu}$ on a particular training point ($\nu$) is affect by these derivative terms. Following the above observation, most terms in the exponent will not act on $f_{\nu}$ and a $1/n$ portion will contain a single derivative. The remaining rarer cases, where two derivatives act on the same $\nu$, are neglected. For each $f_{\nu}$ we thus replace $r-1$ derivatives in the order $r$ term in C.1 by discrepancies, leaving a single derivative operator that is multiplied by the following quantity

Note that the summation indices span only the training set, not the test point: $\mu_{1},...,\mu_{r-1}\in\left\{1,\dots,n\right\}$, whereas the free index spans also the test point $\nu\in\left\{1,\dots,n+1\right\}$.

Given a fixed $\overline{\Delta g}$, C.3 is the partition function corresponding to a GP with the train targets shifted by $\overline{\Delta g}_{\mu}$ and the test target shifted by $\overline{\Delta g}_{*}$. Following this we find that $\overline{\Delta g}$ depends on the discrepancy of the GP prediction which in turn depends on $\overline{\Delta g}$. In other words we obtain our self-consistent equation: $\overline{\Delta g}=\langle\Delta g_{\mu}\rangle_{Z\left(\vec{J};\overline{\Delta g}\right)}$.

The partition function C.3 reflects the correspondence between finite DNNs and a GP with its target shifted by $\overline{\Delta g}$. To facilitate the analytic solution of this self-consistent equation, we focus on the case $\left\langle\hat{\delta}g_{\mu}\hat{\delta}g_{\nu}\right\rangle\ll\left\langle\hat{\delta}g_{\mu}\right\rangle\left\langle\hat{\delta}g_{\nu}\right\rangle$ at least for $\mu\neq\nu$. We note that this was the case for the two toy models we studied.

Given this, the expectation value over $\Delta g$ using the GP defined by $Z\left(\vec{J};\overline{\Delta g}\right)$, which consists of products of expectation values of individual discrepancies and correlations between two discrepancies, can then be expressed using only the former. Omitting correlations within the GP expectation value, one obtains a simplified self-consistent equation involving only the average discrepancies:

with $\hat{\delta}g_{\mu}$ now understood as a number, also within C.2. Lastly, we plug the solution to these equations to find the prediction on the test point: $\left\langle f({\mathbf{x}}_{*})\right\rangle_{Z\left(\vec{J};\overline{\Delta g}\right)}$. These coincide with the self-consistent equations derived via the saddle point approximation in the main text.

Notably the above derivation did not hinge on having MSE loss. For any loss given as a sum over training points, $\mathcal{L}=\sum^{n}_{\mu}L_{\mu}(f_{\mu})$, the above derivation should hold with $\sigma^{-2}\hat{\delta}g_{\mu}$ in $\Delta g_{\nu}$ replaced by $\partial_{f_{\mu}}L_{\mu}$. In particular for the cross entropy loss where $f_{\nu,i}$ is the pre-softmax output of the DNN for class $i$ we will have

where $i$ and $j$ run over all classes, $i_{\nu}$ is the class of ${\mathbf{x}}_{\nu}$. Neatly, the above r.h.s. is again a form of discrepancy but this time in probability space. Namely it is $p_{\rm{model}}(i|{\mathbf{x}}_{\nu})-p_{\rm{data}}(i|{\mathbf{x}}_{\nu})$, where $p_{\rm{model}}$ is the distribution generated by the softmax layer, and $p_{\rm{data}}$ is the empirical distribution. Following this one can readily derive self-consistent equations for cross entropy loss and solve them numerically. Further analytical progress hinges on developing analogous of the EK approximation for cross entropy loss.

Appendix D Review of the Equivalent Kernel (EK)

In this appendix we generally follow , see also for more details. The posterior mean for GP regression

can be obtained as the function which minimizes the functional

where $||f||_{\mathcal{H}}$ is the RKHS norm corresponding to kernel $K$. Our goal is now to understand the behaviour of the minimizer of $J[f]$ as $n\to\infty$. Let the data pairs $\left({\mathbf{x}}_{\alpha},y_{\alpha}\right)$ be drawn from the probability measure $\mu({\mathbf{x}},y)$. The expectation value of the MSE is

Since the right term on the RHS of D.4 does not depend on $f$ we can ignore it when looking for the minimizer of the functional which is now replaced by

To proceed we project $g$ and $f$ on the eigenfunctions of the kernel with respect to $\mu({\mathbf{x}})$ which obey $\int\mu\left({\mathbf{x}}^{\prime}\right)K\left({\mathbf{x}},{\mathbf{x}}^{\prime}\right)\psi_{s}\left({\mathbf{x}}^{\prime}\right)=\lambda_{s}\psi_{s}\left({\mathbf{x}}\right)$. Assuming that the kernel is non-degenerate so that the $\psi$’s form a complete orthonormal basis, for a sufficiently well behaved target we may write $g\left({\mathbf{x}}\right)=\sum_{s}g_{s}\psi_{s}\left({\mathbf{x}}\right)$ where $g_{s}=\int g\left({\mathbf{x}}\right)\psi_{s}\left({\mathbf{x}}\right)d\mu\left({\mathbf{x}}\right)$, and similarly for $f$. Thus the functional becomes

This is easily minimized by taking the derivative w.r.t. each $f_{s}$ to yield

In the limit $n\to\infty$ we have $\sigma^{2}/n\to 0$ thus we expect that $f$ would converge to $g$. The rate of this convergence will depend on the smoothness of $g$, the kernel $K$ and the measure $\mu({\mathbf{x}},y)$. From D.7 we see that if $n\lambda_{s}\ll\sigma^{2}$ then $f_{s}$ is effectively zero. This means that we cannot obtain information about the coefficients of eigenfunctions with small eigenvalues until we get a sufficient amount of data. Plugging the result D.7 into $f\left({\mathbf{x}}\right)=\sum_{s}f_{s}\psi_{s}\left({\mathbf{x}}\right)$ and recalling $g_{s}=\int g\left({\mathbf{x}}^{\prime}\right)\psi_{s}\left({\mathbf{x}}^{\prime}\right)d\mu\left({\mathbf{x}}^{\prime}\right)$ we find

The term $h({\mathbf{x}},{\mathbf{x}}^{\prime})$ it the equivalent kernel. Notice the similarity to the vector-valued equivalent kernel weight function $\mathbf{h}\left({\mathbf{x}}_{*}\right)=\left(\mathbf{K}+\sigma^{2}I\right)^{-1}\mathbf{k}\left({\mathbf{x}}_{*}\right)$ where $\mathbf{K}$ denotes the $n\times n$ matrix of covariances between the training points with entries $K\left({\mathbf{x}}_{\mu},{\mathbf{x}}_{\nu}\right)$ and $\mathbf{k}\left({\mathbf{x}}_{\ast}\right)$ is the vector of covariances with elements $K\left({\mathbf{x}}_{\mu},{\mathbf{x}}_{\text{\textasteriskcentered}}\right)$. The difference is that in the usual discrete formulation the prediction was obtained as a linear combination of a finite number of observations $y_{i}$ with weights given by $h_{i}({\mathbf{x}})$ while here we have instead a continuous integral.

Appendix E Additional technical details for solving the self consistent equations

In this subsection we show how to arrive at Eq. 16 from the main text, which is a self consistent equation for the proportionality constant, $\alpha$, defined by $\hat{\delta}g=\alpha g$. We first show that both the shift and the discrepancy are linear in the target, and then derive the equation.

Recall that we assume a linear target with a single channel:

The fact that $f$ is always a linear function of the input (since the CNN linear) and the fact that it is proportional to $g$ at $C\to\infty$ (since the GP is linear in the target), motivates the ansatz:

Indeed we will show that this ansatz provides a solution to the non linear self consistent equations.

Notice that the target shift has a form of a geometric series. In the linear CNN toy model we are able to sum this entire series, whose first term is related to (using the notation introduced in §F):

For simplicity we can assume $\norm{{\mathbf{w}}^{*}}^{2}=1$ and $\sigma_{a}^{2}=1$, thus getting a simple proportionality constant of $\frac{6\lambda^{2}}{C}$. If we were to trade $g$ for $\hat{\delta}g$, as we have in $\Delta g$, we would get a similar result, with an extra factor of $\left(\frac{\alpha}{\sigma^{2}/n}\right)^{3}$. The factor of $6$ will cancel out with the factor of $1/(4-1)!$ appearing in the definition of $\Delta g$. Repeating this calculation for the sixth cumulant, one would arrive to the same result multiplied by a factor of $\frac{\lambda}{C}\left(\frac{\alpha}{\sigma^{2}/n}\right)^{2}$ due to the general form of the even cumulants (Eq. F.29) and the fact that there an extra two $(\sigma^{2}/n)^{-1}\hat{\delta}g$’s.

E.1.2 Self consistent equation in the EK limit

Starting from the proportionality relations $\hat{\delta}g=\alpha g$ and $\Delta g=\alpha_{\Delta}g$, we can now write the self consistent equation for the discrepancy as

Dividing both sides by $g$ we get a scalar equation

The factor $\alpha_{\Delta}$ can be calculated by noticing that $\Delta g$ has the form of a geometric series. To better understand what follows next, the reader should first go over §F. The first term in this series is related to contracting the fourth cumulant $\kappa_{4}$ with three $\hat{\delta}g$’s thus yielding a factor of $\frac{\lambda^{2}}{C}\left(\frac{\alpha}{\sigma^{2}/n}\right)^{3}$ (recall that in the EK approximation we trade $\sigma^{2}\to\sigma^{2}/n$). The ratio of two consecutive terms in this series is given by $\frac{\lambda}{C}\left(\frac{\alpha}{\sigma^{2}/n}\right)^{2}$. Using the formula for the sum of a geometric series we have

The EK approximation can be improved systematically using the field-theory approach of Ref. where the EK result is interpreted as the leading order contribution, in the large $n$ limit, to the average of the GP predictor over many data-set draws from the dataset measure. However, that work focused on the test performance whereas for $q_{\rm{train}}$ we require the performance on the training set. We briefly describe the main augmentations needed here and give the sub-leading and sub-sub-leading corrections to the EK result on the training set, enabling us to estimate $q_{\rm{train}}$ analytically within a $16.3\%$ relative error compared with the empirical value. Further systematic improvements are possible but are left for future work.

We thus consider the quantity $\sum_{\mu}\varphi({\mathbf{x}}_{\mu})f({\mathbf{x}}_{\mu})$ where ${\mathbf{x}}_{\mu}$ is drawn from the training set, $f({\mathbf{x}}_{\mu})$ is the predictive mean of the GP on that specific training set, and $\varphi({\mathbf{x}}_{\mu})$ is some function which we will later take to be the target function ($\varphi({\mathbf{x}})=g({\mathbf{x}})$). We wish to calculate the average of this quantity over all training set draws of size $n$. We begin by adding a source term of the form $J\sum_{\mu}\varphi({\mathbf{x}}_{\mu})f({\mathbf{x}}_{\mu})$ to the action and notice a similar term appearing in the GP action ($-\sum_{\mu}(f({\mathbf{x}}_{\mu})-g({\mathbf{x}}_{\mu}))^{2}$) due to the MSE loss. Examining this extra term one notices that it can be absorbed as a $J$ dependent shift to the target on training set ($g({\mathbf{x}}_{\mu})\rightarrow g({\mathbf{x}}_{\mu})+\frac{J\sigma^{2}}{2}\varphi({\mathbf{x}}_{\mu})$) following which the analysis of Ref. carries through straightforwardly. Doing so, the general result for the leading EK term and sub-leading correction are

Turning to the specific linear CNN toy model and carrying the above expansion up to an additional term leads to

Considering for instance $n=200,\sigma^{2}=1.0,N=30$ and $S=30$, we find $\alpha_{\rm{EK}}=0.818$ and so

recalling that $q_{\rm{train}}=\frac{\lambda+\sigma^{2}/n}{\lambda}(1-\alpha_{\rm{train}})$ we have

whereas the empirical value here is $2.8995$.

Appendix F Cumulants for a two-layer linear CNN

In this section we explicitly derive the leading (fourth and sixth) cumulants of the toy model of §IV.1, and arrive at the general formula for the even cumulant of arbitrary order.

For a general activation, we have in our setting for a 2-layer CNN

Averaging over the last layer weights gives

So this will always make two pairs out of four $\phi$’s, each with the same $i,c$ indices. Notice that, regardless of the input indices, for different channels $c\neq c^{\prime}$ we have

where in the last line we separated the diagonal and off-diagonal terms in the channel indices. So

Putting it all together, the off-diagonal terms in the channel indices cancel and we are left with

where in the last line we introduced a short-hand notation to compactly keep track of the combinations of the indices.

F.1.2 Fourth cumulant for linear CNN

Notice that the 2nd and 3rd terms have $\left(ij\right)\left(ij\right)$ while the first term has $\left(ii\right)\left(jj\right)$. The latter will cancel out with the $\left\langle\phi_{i,c}^{\mu}\phi_{i,c}^{\nu}\right\rangle_{{\mathbf{w}}}\left\langle\phi_{j,c}^{\mu^{\prime}}\phi_{j,c}^{\nu^{\prime}}\right\rangle_{{\mathbf{w}}}$ terms. Thus

Denote $\lambda:=\frac{\sigma_{a}^{2}}{N}\frac{\sigma_{w}^{2}}{S}$ The fourth cumulant is

F.2 Sixth cumulant and above

The even moments in terms of cumulants for a vector valued RV with zero odd moments and cumulants are (see ):

where the moments are on the l.h.s. (indices with no commas) and the cumulants are on the r.h.s. (indices are separated with commas). Thus, the sixth cumulant is

In the linear case, the analogue of $\kappa^{\mu_{1},\mu_{2},\mu_{3},\mu_{4}}\kappa^{\mu_{5},\mu_{6}}$ is ($15$ such pairings, where only the numbers "move", not the $i,j,k$)

and the analogue of $\kappa^{\mu_{1},\mu_{2}}\kappa^{\mu_{3},\mu_{4}}\kappa^{\mu_{5},\mu_{6}}$ is

Below, we found the 6th moment for a linear CNN to be

Notice that for every blue term we have exactly $6$ red terms, so all of the colored terms will exactly cancel out and only the uncolored terms will survive. There are $8$ such uncolored terms for each one of the $15$ pairings, thus we will ultimately have $120$ such pairs, thus the sixth cumulant is

where the $\left[120\right]$ stands for the number of ways to pair the numbers $\left\{1,...,6\right\}$ into the form $\left(\bullet_{i}\bullet_{j}\right)\left(\bullet_{i}\bullet_{k}\right)\left(\bullet_{j}\bullet_{k}\right)$.

We can thus identify a pattern which we conjecture to hold for any even cumulant of arbitrary order $2m$:

where the indices $i_{1},\dots,i_{m}$ obey the following:

Each index appears exactly twice in each summand.

Each index cannot be paired with itself, i.e. $\left(\bullet_{i_{1}},\bullet_{i_{1}}\right)$ is not allowed.

Appendix G Feature learning phase transition

Although our main focus was on the statistics of the DNN outputs, our function-space formalism can also be used to characterize the statistics of the weights of the intermediate hidden layers. Here we focus on the linear CNN toy model given in the main text, where the learnable parameters of the student are given by $\theta=\left\{w_{c,s},a_{i,c}\right\}$. Consider first a prior distribution in output space, where throughout this section we denote: $\vec{f}\equiv\left(f_{1},\dots,f_{n}\right)$, i.e. the vector of outputs on the training set alone (without the test point). Since we are interested in the statistics of the hidden weights, we will introduce an appropriate source term in weight space $J_{c,s}$

where $z_{\theta,\mu}$ is the of output of the CNN parameterized by $\theta$ on the $\mu$’th training point. Given some loss function $\mathcal{L}$, the posterior is given by

The posterior mean of the hidden weights is thus

and the posterior covariance can be extracted from taking the second derivative, namely

Our next task is to rewrite these expectation values over weights under the posterior as expectation values of DNN training outputs ($f({\mathbf{x}}_{\mu})$) under the posterior. To this end we write down the kernel of this simple CNN such that it depends on the source terms:

We can now write the second mixed derivatives of $K_{J}$ to leading order in $J$ as

Next we take the large $C$ limit and thus have a posterior of the form $P[\vec{f},J]=P_{0}[\vec{f},J]e^{-\mathcal{L}/\sigma^{2}}$ where $P_{0}[\vec{f}]$ contains only $K_{J}^{-1}$ and none of the higher cumulants. Having the derivatives of $K_{J}^{-1}$ w.r.t. $J$ we can proceed in analyzing the derivatives of the log-partition function for the posterior w.r.t $J$. In particular the covariance matrix of the weights averaged over the different channels is

The above result is one of the two main points of this appendix: we established a mapping between expectation values over outputs and expectation values over hidden weights. Such a mapping can in principle be extended to any DNN. On the technical level, it requires the ability to calculate the cumulants as a function of the source terms, $J$. As we argue below, it may very well be that unlike in the main text, only a few cumulants are needed here.

To estimate the above expectation values we use the EK limit, where the sums over the training set are replaced by integrals over the measure $\mu({\mathbf{x}})$, the $f$’s are replaced as $f\left({\mathbf{x}}_{\mu}\right)\to\frac{\lambda}{\lambda+\sigma^{2}/n}g\left({\mathbf{x}}\right)$ and we assume the input distribution is normalized as $\int d\mu\left({\mathbf{x}}\right)x_{i}x_{j}=\delta_{ij}$. Following this we find

Comparing this to our earlier result for the covariance Eq. G.4 we get

Multiplying by $S=1/\sigma_{w}^{2}$ and recalling that $\lambda=1/NS$ we get

Repeating similar steps while also taking into account diagonal fluctuations yields another factor of $\left(\frac{1}{\lambda}+\frac{n}{\sigma^{2}}\right)^{-1}$ on the diagonal, thus arriving at the result as it appears in the main text:

The above results capture the leading order correction in $1/C$ to the weights covariance matrix. However the careful reader may be wary of the fact that the results in the main text require $1/C$ corrections to all orders and so it is potentially inadequate to use such a low order expansion deep in the feature learning regime, as we do in the main text. Here we note that not all DNN quantities need to have the same dependence on $C$. In particular it was shown in Ref. , that the weight’s low order statistics is only weakly affected by finite-width corrections whereas the output covariance matrix is strongly affected by these. We conjecture that this is the case here and that only the cumulative effect of many weights, as reflected in the output of the DNN, requires strong $1/C$ corrections.

This conjecture can be verified analytically by repeating the above procedure on the full prior (i.e. the one that contains all cumulants), obtaining the operator in terms of $f$’s corresponding the weight’s covariance matrix, and calculating its average with respect to the saddle point theory. We leave this for future work.

G.2 A surrogate quantity for the outlier

Since we used moderate $S$ values in our simulations (to maintain a reasonable compute time), we aggregated the eigenvalues of many instances of $\Sigma_{W}$ across training time and across noise realizations. Although the empirical histogram of the spectrum of $\Sigma_{W}$ agrees very well with the theoretical MP distribution (solid smooth curves in Fig. 2A), there is a substantial difference between the two at the right edge of the support $\lambda_{+}$, where the empirical histogram has a tail due to finite size effects. Thus it is hard to characterize the phase transition using the largest eigenvalue $\lambda_{\max}$ averaged across realizations. Instead, we use the quantity $\mathcal{Q}\equiv{\mathbf{w}}^{*\mathsf{T}}\Sigma_{W}{\mathbf{w}}^{*}$ as a surrogate which coincides with $\lambda_{\max}$ for $C\ll C_{\rm{crit}}$ but behaves sensibly on both sides of $C_{\rm{crit}}$, thus allowing to characterize the phase transition.

Appendix H Further details on the numerical experiments

In our experiments, we used the following hyper-parameter values. Learning rates of $\eta=10^{-6},3\cdot 10^{-7}$ which yield results with no appreciable difference in almost all cases, when we scale the amount of statistics collected (training epochs after reaching equilibrium) so that both $\eta$ values have the same amount of re-scaled training time: we used $10$ training seeds for $\eta=10^{-6}$ and 30 for $\eta=3\cdot 10^{-7}$. We used a gradient noise level of $\sigma^{2}=1.0$, but also checked for $\sigma^{2}\in\{0.1,0.01\}$ and got qualitatively similar results to those reported in the main text.

In the main text and here we do not show error bars for $\alpha$ as these are too small to be appreciated visually. They are smaller than the mean values by approximately two orders of magnitude. The error bars were found by computing the empirical standard deviation of $\alpha$ across training dynamics and training seeds.

H.2 Convergence of the training protocol to GP

In Fig. 4 we plot the MSE between the outputs of the trained CNNs and the predictions of the corresponding GP. We see that as $C$ becomes large the slope of the MSE tends to $-2.0$ indicating the $O(1/C)$ scaling of the leading corrections to the GP. This illustrates where we enter the perturbative regime of GP, and we see that this happens for larger $C$ as we increase the conv-kernel size $S$, since this also increases the input dimension $d=NS$. Thus it takes larger $C$ to enter the highly over-parameterized regime.

Appendix I Quadratic fully connected network

At large $M$ and for $w_{m,i}$ drawn from ${\mathcal{N}}(0,\sigma_{w}^{2}/M)$, the student generates a GP prior. It is shown below that the GP kernel is simply $K({\mathbf{x}},{\mathbf{x}}^{\prime})=\frac{2\sigma_{w}^{4}}{M}({\mathbf{x}}\cdot{\mathbf{x}}^{\prime})^{2}$. As such it is proportional to the kernel of the above DNN with an additional linear read-out layer. The above model can be written as $\sum_{ij}x_{i}[P_{ij}-\sigma_{w}^{2}\delta_{ij}]x_{j}$ where $P_{ij}$ is a positive semi-definite matrix. The eigenvalues of the matrix appearing within the brackets are therefore larger than $-\sigma_{w}^{2}$ whereas no similar restriction occurs for DNNs with a linear read-out layer. This extra restriction is completely missed by the GP approximation and, as discussed in Ref. , leads to strong performance improvements compared to what one expects from the GP or equivalently the DNN with the linear readout layer. Here we demonstrate that our self-consistent approach at the saddle-point level captures this effects

We consider training this DNN on $n$ train points $\left\{{\mathbf{x}}_{\mu}\right\}_{\mu=1}^{n}$ using noisy GD training with weight decay $\gamma=M\sigma^{2}/\sigma_{w}^{2}$. We wish to solve for the predictions of this model with our shifted target approach. To this end, we first derive the cumulants associated with the effective Bayesian prior ($P_{0}(\vec{f})$) here. Equivalently stated, obtain the cumulants of the equilibrium distribution of $\vec{f}$ following training with no data, only a weight decay term. This latter distribution is given by

To obtain the cumulants, we calculate the cumulant generating function of this distribution given by

Taylor expanding this last expression is straightforward. For instance up to third order is gives

from which the cumulants can be directly inferred, in particular the associated GP kernel given by

Following this, the target shift equation, at the saddle point level, becomes

Figure 5 shows the numerical results for the test MSE as obtained by solving the above equations for $\hat{\delta}g$ on the training set, taking $\nu=*$ in these equation together with the self-consistent $\hat{\delta}g_{\mu}$ to find the mean-predictor, and taking the average MSE of the latter over the test set. Both test and train data sets were random points sampled uniformly from a $d$ dimensional hypersphere of radius one. The test dataset contained $100$ points and the figure shows the test MSE as a function of $n/d$ where $d=20$, $\sigma^{2}_{w}=1$, $\sigma^{2}=2.76\cdot 10^{-6}$, $M=4d$, and $w^{*}_{i}$ drawn from $\mathcal{N}(0,1)$. The non-linear equations were solved using the Newton-Krylov algorithm together with gradual annealing from $\sigma^{2}=1$ down to the above values. The figure shows the median over $60$ data sets. Remarkably, our self-consistent approach yields the expected threshold values of $n/d=2$ separating good and poor performance. Discerning whether this is a threshold or a smooth cross-over in the large $d$ limit is left for future work.

Turning to analytics, one can again employ the EK approximation as done for the CNN. However taking $\sigma^{2}$ to zero invalidates the EK approximation and requires a more advance treatment as in Ref. . We thus leave an EK type analysis of the self-consistent equation at $\sigma^{2}=0$ for future work and instead focus on the simpler case of finite $\sigma^{2}$ where analytical predictions can again be derived in similar fashion to our treatment of the CNN.

To simplify things further, we also commit to the distribution $[{\mathbf{x}}_{\mu}]_{i}\sim\mathcal{N}(0,1/d)$. In this setting $K({\mathbf{x}},{\mathbf{x}}^{\prime})$ has two distinct eigenvalues w.r.t. to this measure, the larger one ($\lambda_{0}=2M^{-1}\sigma_{w}^{4}\left(\frac{2}{d^{2}}+\frac{1}{d}\right)$) associated with $f({\mathbf{x}})=||{\mathbf{x}}||^{2}$ and a smaller one ($\lambda_{2}=2M^{-1}\sigma_{w}^{4}\frac{2}{d^{2}}$) associated with $x_{i}x_{j}$ (with $i\neq j$) and $\sum_{i}a_{i}x_{i}^{2}$ (with $\sum_{i=1}^{d}a_{i}=0$) eigenfunctions.

Next we argue that provided the discrepancy is of the following form

then within the EK limit the target shift is also of the form of the r.h.s. with $\alpha_{\Delta}$ and $\beta_{\Delta}$ and the target shift equations reduce to two coupled non-linear equations for $\alpha$ and $\beta$. Following the EK approximation, we replace all $\sum_{\mu}$ in the target shift equation with $n\int d\mu_{x}$ and obtain

Next we note that the $i\neq j$ element of the matrix $\int d\mu_{x}g({\mathbf{x}}){\mathbf{x}}{\mathbf{x}}^{\mathsf{T}}$ is given by

taking this together with the simpler term ($\int d\mu({\mathbf{x}})\frac{\beta}{\alpha}\sum_{i}x_{i}x_{i}{\mathbf{x}}{\mathbf{x}}^{\mathsf{T}}=\frac{\beta(d+2)}{\alpha d^{2}}I$)

Consider the matrix $({\mathbf{w}}_{*}{\mathbf{w}}_{*}^{\mathsf{T}}+bI)$, appearing in the above denominator with $b=\left(\frac{\beta\sigma_{w}^{2}(d+2)}{2\alpha}-\sigma_{w}^{2}\right)$, and note that

Plugging this equation into a Taylor expansion of the denominator of Eq. I.10 one finds that all the resulting terms are of the desired form of a linear superposition of $g({\mathbf{x}})$ and $\norm{{\mathbf{x}}}^{2}$. Considering the first term on the r.h.s. of Eq. I.10, $\beta\norm{x}^{2}$ is already an eigenfunction of the kernel whereas $g({\mathbf{x}})$ can be re-written as

so that the first two terms on the r.h.s. are $\lambda_{2}$ eigenfunctions and the last one is a $\lambda_{0}$ eigenfunctions. Summing these different contributions along with the aforementioned Taylor expansion, one finds that $\Delta g({\mathbf{x}})$ is indeed a linear superposition of $g({\mathbf{x}})$ and $\norm{{\mathbf{x}}}^{2}$.

Next we wish to write down the saddle-point equations for $\alpha$ and $\beta$. For simplicity we focus on the case where $g({\mathbf{x}})$ is chosen orthogonal to $\norm{{\mathbf{x}}}^{2}$ under $d\mu({\mathbf{x}})$, namely $\norm{{\mathbf{w}}_{*}}^{2}=d\sigma_{w}^{2}$. Under this choice the self-consistent equations become

where the constant $b$ was defined above and $c=\frac{n\lambda_{2}}{\sigma_{w}^{2}\sigma^{2}}$.

Next we perform several straightforward algebraic manipulations with the aim of extracting their asymptotic behavior at large $n$. Noting that $c\sigma_{w}^{2}=\frac{n}{\sigma^{2}}\lambda_{2}$, $c\norm{{\mathbf{w}}_{*}}^{2}=2d\frac{n}{\sigma^{2}}\lambda_{2}$, and $\alpha cb=-\frac{n}{\sigma^{2}}(\alpha\lambda_{2}-2\beta\lambda_{0})$ we have

noting that $d\lambda_{2}=2(\lambda_{0}-\lambda_{2})$ we find

The first equation above is linear in $\beta$ and yields in the large $d$ limit

which when placed in the second equation yields

At large $n$, we expect $\alpha$ and $\beta$ to go to zero. Accordingly to find the asymptotic decay to zero, one can approximate $\alpha(1-\alpha)\approx\alpha$, and similarly $\alpha/(1-\alpha)\approx\alpha$. This along with the large $d$ limit simplifies the equations to a quadratic equation in $\beta$

which for $\sigma^{2}/n\ll\lambda_{0}$ simplifies further into

We thus find that both $\alpha$ and $\beta$ are of the order of $\frac{\sigma^{2}/n}{2\lambda_{0}}=n^{-1}\frac{Md\sigma^{2}}{4\sigma_{w}^{4}}$. Hence $n$ scaling as $Md$ ensures good performance. This could have been anticipated as for small yet finite $\sigma^{2}$ each $n$ can be seen as a soft constrained on the parameters of the DNN and since the DNN contains $Md$ parameters $n=O(Md)$ should provide enough data to fix the student’s parameters close to the teacher’s.